diff --git a/litellm/constants.py b/litellm/constants.py index 3ff80d8b7dd..7b40f432446 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1772,6 +1772,8 @@ RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_S RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25"))) +PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, float(os.getenv("PROXY_DB_LOOKUP_DEADLINE_SECONDS", "10"))) +PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, float(os.getenv("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", "30"))) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 7990832dc48..6cc0d9444cd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2017,8 +2017,8 @@ def _deployment_model_info( return cast(ModelInfo, registered_deployment_info) # cast-ok: router registers deployment prices under its id if litellm_logging_obj is None: return None - litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) - if litellm_params is None: + litellm_params: Final = litellm_logging_obj.litellm_params + if not litellm_params: return None return next( ( @@ -2036,7 +2036,9 @@ def _ocr_model_info( router_model_id: str | None, ) -> OCRPricing | None: deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) - litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None + litellm_params: Final = ( + litellm_logging_obj.litellm_params if custom_pricing and litellm_logging_obj is not None else None + ) if litellm_params is None: return deployment_info return _layered_ocr_pricing(litellm_params, deployment_info) diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index a9ee851d529..df644fd7f4a 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -129,7 +129,7 @@ async def list_tools_with_pagination( ) tools.extend(result.tools) - next_cursor = getattr(result, "next_cursor", None) + next_cursor = result.next_cursor if not isinstance(next_cursor, str) or not next_cursor: return tools if next_cursor in seen_cursors: diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 4ab8d9796b2..57e60fea759 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -112,7 +112,7 @@ class ArizeLogger(OpenTelemetry): if value is None or value in ("", "None"): return None try: - rate = float(value) + rate: Final = float(value) except (TypeError, ValueError): verbose_logger.warning( "ArizeLogger: %s value %r is not a number; exporting the request", diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 7e3c4cc3ce8..4c2f75bb4d7 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -21,7 +21,7 @@ from __future__ import annotations import os from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_proxy_logger @@ -35,17 +35,6 @@ else: AsyncIOScheduler = Any -class _PodLockManager(Protocol): - """The subset of PodLockManager this logger drives to serialize the export across pods.""" - - @property - def redis_cache(self) -> object: ... - - async def acquire_lock(self, cronjob_id: str) -> bool | None: ... - - async def release_lock(self, cronjob_id: str) -> None: ... - - def _parse_metrics_marker( marker: object | None, ) -> datetime | None: @@ -237,13 +226,10 @@ class MavvrikFocusLogger(FocusLogger): """Scheduler entry point — uses Mavvrik-specific pod-lock key.""" from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415 - pod_lock_manager: _PodLockManager | None = None - if proxy_logging_obj is not None: - writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) - if writer is not None: - pod_lock_manager = getattr(writer, "pod_lock_manager", None) - - if pod_lock_manager and pod_lock_manager.redis_cache: + pod_lock_manager: Final = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager if proxy_logging_obj is not None else None + ) + if pod_lock_manager is not None and pod_lock_manager.redis_cache: acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME) if not acquired: verbose_proxy_logger.debug("Mavvrik FOCUS export: unable to acquire pod lock") diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 54578323fa4..90bbd5a00d8 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1849,7 +1849,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: # Handle both Anthropic-style input and OpenAI-style function.arguments query = None - tool_args: dict | None = None # mutable-ok: the tool call's own arguments dict + tool_args: dict[str, object] | None = None # mutable-ok: the tool call's own arguments dict if "input" in tool_call and isinstance(tool_call["input"], dict): tool_args = tool_call["input"] query = tool_args.get("query") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index a2d40279c49..3afa6a913b5 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -365,7 +365,7 @@ def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object: return getattr(user_api_key_auth, "budget_reservation", None) -def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None: +def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict[str, object] | None: stamped: Final = metadata.get("user_api_key_budget_reservation") if isinstance(stamped, dict): return stamped diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0603414cabd..2cecec729c2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5191,7 +5191,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetryV2) - and getattr(callback, "callback_name", None) == callback_name + and callback.callback_name == callback_name and (serves_a_destination or not _exports_nowhere(callback.config)) ): return callback @@ -6663,7 +6663,7 @@ def get_standard_logging_object_payload( cost_breakdown=request_cost_breakdown, autorouter_savings=autorouter_savings, autorouter_savings_estimate=( - { + { # mutable-ok: spend-log JSON serialization requires plain mappings "version": 3, "status": "unknown", "reason": "pending_projection", diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py index 6331d815bdc..9688b511ea8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -5,7 +5,12 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never -from litellm.types.utils import StandardLoggingZeroCostDiagnostic, Usage +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, + StandardLoggingZeroCostDiagnostic, + Usage, +) ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" @@ -18,8 +23,8 @@ _NESTED_PRICING: Final = TypeAdapter(Mapping[str, object] | tuple[object, ...]) _MAX_PRICING_DEPTH: Final = 4 -def _audio_tokens(details: object) -> int: - audio_tokens: Final = getattr(details, "audio_tokens", None) +def _audio_tokens(details: PromptTokensDetailsWrapper | CompletionTokensDetailsWrapper | None) -> int: + audio_tokens: Final = details.audio_tokens if details is not None else None return audio_tokens if isinstance(audio_tokens, int) and audio_tokens > 0 else 0 diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 378295e1b7a..8e5d2cd0a17 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2003,11 +2003,11 @@ def strip_encrypted_reasoning_from_messages(messages: object) -> None: """ if not isinstance(messages, list): return - for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + for content in anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json _strip_encrypted_reasoning_from_blocks(content) -def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: +def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: return ( cast(list[object], content) # cast-ok: narrowed by isinstance for message in messages diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f97a274708f..fa687b585f5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1329,7 +1329,7 @@ class CustomStreamWrapper: "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, "original_chunk": cached_chunk, - "tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None), + "tool_calls": cached_choice.delta.tool_calls if cached_choice is not None else None, } completion_obj["content"] = response_obj["text"] diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index c5a71daaba5..0813e0827d2 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -48,7 +48,7 @@ def _registry_api_key(agent_litellm_params: Mapping[str, object]) -> str | None: return configured_api_key if isinstance(configured_api_key, str) else None -def _registry_headers(agent_litellm_params: Mapping[str, object]) -> dict[str, Any] | None: +def _registry_headers(agent_litellm_params: Mapping[str, object]) -> dict[str, object] | None: stored_headers: Final = agent_litellm_params.get("headers") if not isinstance(stored_headers, Mapping): return None diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a78f633f5d7..e1c727ad235 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -685,7 +685,7 @@ class AnthropicMessagesHandler(BaseTranslation): return data - def _hoisted_top_level_system_message(self, data: dict) -> AllMessageValues | None: + def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None: """Return the system message produced by translating the top-level prompt.""" system: Final = data.get("system") if not system: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 4486eb0985a..20753afee5c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -11,7 +11,6 @@ from typing import ( Final, Literal, Protocol, - cast, # noqa: TID251 # rebuilt message_delta dict spans the ContentBlockDelta/MessageBlockDelta union get_args, ) @@ -27,6 +26,7 @@ from litellm.types.llms.anthropic import ( ContentBlockDelta, ContextManagementResponse, MessageBlockDelta, + MessageDelta, StreamingContentBlockDeltaType, UsageDelta, UsageIteration, @@ -1028,26 +1028,22 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self, processed_chunk: ContentBlockDelta | MessageBlockDelta, ) -> ContentBlockDelta | MessageBlockDelta: - if processed_chunk.get("type") != "message_delta" or not self._refusal_text: + if processed_chunk["type"] != "message_delta" or not self._refusal_text: return processed_chunk - delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) # cast-ok: keys checked before use + delta: Final = processed_chunk["delta"] if delta.get("stop_reason") == "max_tokens": return processed_chunk from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( refusal_stop_details, ) - return cast( # cast-ok: rebuilt dict matches the message_delta TypedDict shape for this branch - ContentBlockDelta | MessageBlockDelta, - { # mutable-ok: fresh translation payload; never mutated after construction - **processed_chunk, - "delta": { # mutable-ok: fresh message_delta payload; never mutated after construction - **delta, - "stop_reason": "refusal", - "stop_details": refusal_stop_details(self._refusal_text), - }, - }, - ) + refusal_delta: Final[MessageDelta] = { + **delta, + "stop_reason": "refusal", + "stop_details": refusal_stop_details(self._refusal_text), + } + refusal_chunk: Final[MessageBlockDelta] = {**processed_chunk, "delta": refusal_delta} + return refusal_chunk @staticmethod def _delta_has_content(processed_chunk: Mapping[str, object]) -> bool: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 7545dff1408..89105c00428 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -37,7 +37,7 @@ def _mapping_field(container: object, key: str) -> object | None: """One key of a raw provider payload, or None when the payload is not a mapping.""" if not isinstance(container, Mapping): return None - return cast(Mapping[str, object], container).get(key) # cast-ok: raw payload, callers re-check every value + return container.get(key) def _mapping_str_field(container: object, key: str) -> str | None: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index e3d3425f8a6..6a31173a9c6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -169,7 +169,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: cls, summary: Iterable[object], encrypted_content: object, - ) -> dict[str, Any] | None: # mutable-ok: API message payload + ) -> dict[str, object] | None: # mutable-ok: API message payload """The one Anthropic block for a Responses reasoning item. The item's encrypted reasoning rides the block's opaque field (`signature`, or @@ -198,7 +198,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @classmethod def _assistant_group_to_input_items( cls, group: tuple[Mapping[str, object], ...] - ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload + ) -> tuple[dict[str, object], ...]: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") if btype in ("thinking", "redacted_thinking"): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index ad8e29ad4ac..66cebe0175d 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -994,7 +994,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _spread_text_rewrite_over_stream_events( self, - stream_events: Sequence[Any], + stream_events: Sequence[object], rewritten_text: str, guardrail_name: str, ) -> None: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index fc9c6bcc19f..9cfaaab89f0 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -7,6 +7,7 @@ Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embed Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -161,12 +162,14 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: """ Get the error class for Vercel AI Gateway errors. """ return VercelAIGatewayException( message=error_message, status_code=status_code, - headers=headers, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a50665c71df..941ec4ad419 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -286,7 +286,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ Check if the model is Gemini 3 or newer. """ - model_name = model.split("/")[-1].lower() + model_name: Final = model.split("/")[-1].lower() is_vertex_fine_tuned_model: Final = model_name.isdigit() or ( model.startswith("gemini/") and not model_name.startswith("gemini-") ) diff --git a/litellm/main.py b/litellm/main.py index 4c40d864169..ceac729d3f0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1182,7 +1182,7 @@ def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: return False -def _without_anthropic_only_tool_keys(tool: dict) -> dict: +def _without_anthropic_only_tool_keys(tool: dict[str, object]) -> dict[str, object]: kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} function: Final = tool.get("function") if not isinstance(function, dict): @@ -1193,7 +1193,7 @@ def _without_anthropic_only_tool_keys(tool: dict) -> dict: } -def _drop_anthropic_only_tool_keys(tools: list[dict] | None) -> list[dict] | None: +def _drop_anthropic_only_tool_keys(tools: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if tools is None: return None return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d2aec118d31..69a64de6e04 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "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, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "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, @@ -5546,7 +5546,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "azure", @@ -8647,6 +8647,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -8695,6 +8696,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.375e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-07, "cache_read_input_token_cost": 1.1e-08, @@ -8743,6 +8745,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, "cache_read_input_token_cost": 2.2e-07, @@ -10657,7 +10660,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -27347,9 +27350,11 @@ }, "vertex_ai/gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai", @@ -27359,6 +27364,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "regional_endpoint_uplift_multiplier": 1.1, @@ -29601,9 +29607,11 @@ }, "gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", @@ -29613,6 +29621,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -41233,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.396e-07, + "input_cost_per_token": 9.24462e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8792e-06, + "output_cost_per_token": 1.848924e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.83e-08, + "cache_read_input_token_cost": 7.70385e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -48128,6 +48137,20 @@ "/v1/realtime" ] }, + "vertex_ai/chirp_2": { + "input_cost_per_second": 0.00026667, + "litellm_provider": "vertex_ai", + "metadata": { + "calculation": "$0.016/60 seconds = $0.00026667 per second", + "original_pricing_per_minute": 0.016 + }, + "mode": "audio_transcription", + "source": "https://cloud.google.com/speech-to-text/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime" + ] + }, "vertex_ai/claude-3-5-haiku": { "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, @@ -50152,6 +50175,7 @@ ] }, "vertex_ai/google/gemma-4-26b-a4b-it-maas": { + "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 262144, @@ -57133,6 +57157,130 @@ "/v1/responses" ] }, + "us.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "us.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, "global.openai.gpt-5.6-luna": { "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, @@ -63236,6 +63384,115 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/FW-DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.1e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3": { + "cache_read_input_token_cost": 3.25e-07, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3-Flash": { + "cache_read_input_token_cost": 3.8e-08, + "input_cost_per_token": 1.88e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.25e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GPT-OSS-120B": { + "cache_read_input_token_cost": 8.2e-08, + "input_cost_per_token": 1.65e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure_ai/Cohere-command-a-plus-05-2026": { + "input_cost_per_token": 8e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/mistral-medium-3-5": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_response_schema": true, + "supports_vision": true + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -65454,13 +65711,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.4e-07, - "output_cost_per_token": 2.64e-06, - "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943717, + "max_tokens": 943717, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -65843,7 +66100,7 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.062e-07, + "input_cost_per_token": 6.562e-07, "output_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", @@ -66165,9 +66422,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 8.554e-08, + "output_cost_per_token": 1.7108e-07, + "cache_read_input_token_cost": 1.7108e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -68634,7 +68891,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68645,7 +68902,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -68811,6 +69068,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -68824,6 +69082,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 3e-07, "cache_read_input_token_cost": 1.2e-08, @@ -68872,6 +69131,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/eu/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost": 2.4e-07, @@ -69066,7 +69326,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69077,7 +69337,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -69282,6 +69542,114 @@ "mode": "embedding", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "azure/gpt-live-1": { + "input_cost_per_second": 0.000833333333333, + "litellm_provider": "azure", + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, + "azure/gpt-live-transcribe": { + "input_cost_per_second": 0.000283333333333, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "azure", + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-realtime-translate": { + "input_cost_per_second": 0.000566666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, "litellm_provider": "aihubmix", @@ -70438,6 +70806,34 @@ "output_cost_per_token": 0.0, "source": "https://docs.typesafe.ai/models" }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true + }, + "wandb/google/gemma-4-26B-A4B-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "wandb", + "max_input_tokens": 262000, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "wandb/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 1.5e-07, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index ef931827d85..45ab52690bf 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -540,7 +540,9 @@ def llm_passthrough_route( ) ## IS STREAMING REQUEST - _streaming_request_data: dict = data if isinstance(data, dict) else (json if isinstance(json, dict) else {}) + _streaming_request_data: Final[dict[str, object]] = ( + data if isinstance(data, dict) else (json if isinstance(json, dict) else {}) + ) is_streaming_request: Final = provider_config.is_streaming_request( endpoint=endpoint, request_data=_streaming_request_data, diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index e5d11271c67..77bdbd26b35 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -86,8 +86,8 @@ async def oauth_authorization_uses_gateway_credential(request: Request) -> bool: async def _opaque_bearer_is_gateway_credential(token: str) -> bool: - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - is_envelope, # noqa: PLC0415 # envelope imports bridge types + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # envelope imports bridge types + is_envelope, is_refresh_envelope, ) from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ade829a1b67..d42c1c6b879 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2634,7 +2634,9 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: } -def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict: +def _build_aggregate_authorization_server_response( + request: Request, token_exchange_available: bool +) -> dict[str, object]: """RFC 8414 metadata for the gateway as the aggregate authorization server. The issuer is ``{base}/mcp`` and must stay equal to the value the diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0c520142fb3..312dcb27d89 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3490,7 +3490,7 @@ class MCPServerManager: passthrough_server_ids: Final = [ server.server_id for server in self.get_registry().values() - if getattr(server, "auth_type", None) == MCPAuth.true_passthrough + if server.auth_type == MCPAuth.true_passthrough ] combined_servers.update(passthrough_server_ids) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 3c060752934..71e46f8df25 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -128,14 +128,15 @@ def with_mcp_proxy_identity(tool: Tool, server_id: str) -> Tool: def _mcp_proxy_identity(tool: Tool) -> MCPProxyToolIdentity: - identity: Final = (tool.meta or {}).get(_MCP_PROXY_IDENTITY_META_KEY) # mutable-ok: absent metadata default + identity: Final = None if tool.meta is None else tool.meta.get(_MCP_PROXY_IDENTITY_META_KEY) if not isinstance(identity, Mapping): raise TypeError("MCP proxy tool identity is missing") server_id: Final = identity.get("server_id") tool_name: Final = identity.get("tool_name") if not isinstance(server_id, str) or not isinstance(tool_name, str): raise TypeError("MCP proxy tool identity is invalid") - return {"server_id": server_id, "tool_name": tool_name} # mutable-ok: TypedDict identity payload + resolved: Final[MCPProxyToolIdentity] = {"server_id": server_id, "tool_name": tool_name} + return resolved def mcp_proxy_tool_id(tool: Tool) -> str: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d2fed1eb421..67950e603c0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,11 +15,11 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm from litellm._logging import verbose_proxy_logger @@ -110,7 +110,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_auth_cache_key, team_membership_reservation_cache_key, ) -from litellm.proxy.db.db_lookup_gate import db_lookup_gate +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, @@ -223,30 +223,79 @@ class _PrismaTableHolder(Protocol[RowT_co]): def table(self) -> _PrismaAuthTable[RowT_co]: ... -def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow]) -> _PrismaAuthTable[_PrismaDictableRow]: - return repo.table +class _FindOneKwargs(TypedDict): + where: ReadOnly[Required[Mapping[str, object]]] + include: ReadOnly[NotRequired[Mapping[str, object] | None]] + + +class _FindManyKwargs(TypedDict): + where: ReadOnly[NotRequired[Mapping[str, object] | None]] + include: ReadOnly[NotRequired[Mapping[str, object] | None]] + take: ReadOnly[NotRequired[int | None]] + + +class _DeadlineBoundedTable(Generic[RowT_co]): + """Every read on the wrapped table fails with ``DBLookupDeadlineExceeded`` once + ``PROXY_DB_LOOKUP_DEADLINE_SECONDS`` passes, so a stalled database fails the + request fast instead of parking it in the pod until it fills its memory.""" + + __slots__ = ("_lookup", "_table") + + def __init__(self, table: _PrismaAuthTable[RowT_co], lookup: str) -> None: + self._table: Final = table + self._lookup: Final = lookup + + async def find_unique( + self, + **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> RowT_co | None: + return await bounded_db_lookup(self._table.find_unique(**kwargs), name=self._lookup) + + async def find_first( + self, + **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> RowT_co | None: + return await bounded_db_lookup(self._table.find_first(**kwargs), name=self._lookup) + + async def find_many( + self, + **kwargs: Unpack[_FindManyKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> Sequence[RowT_co]: + return await bounded_db_lookup(self._table.find_many(**kwargs), name=self._lookup) + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> RowT_co | None: + return await self._table.update(where=where, data=data) + + async def create(self, *, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co: + return await self._table.create(data=data, include=include) + + +def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow], lookup: str) -> _PrismaAuthTable[_PrismaDictableRow]: + return _DeadlineBoundedTable(repo.table, lookup) def _jwt_key_mapping_table( repo: _PrismaTableHolder[_PrismaJWTKeyMappingRow], ) -> _PrismaAuthTable[_PrismaJWTKeyMappingRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "jwt_key_mapping") -def _model_dump_table(repo: _PrismaTableHolder[_PrismaModelDumpRow]) -> _PrismaAuthTable[_PrismaModelDumpRow]: - return repo.table +def _model_dump_table( + repo: _PrismaTableHolder[_PrismaModelDumpRow], lookup: str +) -> _PrismaAuthTable[_PrismaModelDumpRow]: + return _DeadlineBoundedTable(repo.table, lookup) def _team_table(repo: _PrismaTableHolder[_PrismaTeamRow]) -> _PrismaAuthTable[_PrismaTeamRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "team") def _vector_store_table(repo: _PrismaTableHolder[_PrismaVectorStoreRow]) -> _PrismaAuthTable[_PrismaVectorStoreRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "vector_store") def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_PrismaUserRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "user") class _VectorStorePermissionsRow(Protocol): @@ -257,7 +306,7 @@ class _VectorStorePermissionsRow(Protocol): def _object_permission_table( repo: _PrismaTableHolder[_VectorStorePermissionsRow], ) -> _PrismaAuthTable[_VectorStorePermissionsRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "object_permission") class _PrismaTagRow(Protocol): @@ -1422,7 +1471,7 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( where={"budget_id": default_budget_id} # mutable-ok: prisma where clause ) @@ -1483,7 +1532,7 @@ async def get_team_member_default_budget( return cached_budget try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( where={"budget_id": budget_id} ) except Exception: @@ -1877,7 +1926,7 @@ async def get_end_user_object( # Fetch from database try: - response: Final = await _dictable_table(EndUserRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -2286,7 +2335,7 @@ async def _fetch_team_membership_from_db( proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: _ = parent_otel_span, proxy_logging_obj - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client), "team_membership").find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -3290,7 +3339,7 @@ async def get_access_object( # Not in cache - fetch from DB try: - response: Final = await _dictable_table(AccessGroupRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(AccessGroupRepository(prisma_client), "access_group").find_unique( where={"access_group_id": access_group_id} ) @@ -3472,7 +3521,7 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await _model_dump_table(OrganizationRepository(prisma_client)).find_many( + orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( where={"organization_alias": org_alias} ) @@ -3650,10 +3699,32 @@ async def _fetch_key_object_from_db_with_reconnect( prisma_client: PrismaClient, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, + deadline_seconds: float | None = None, ) -> BaseModel | None: """ Fetch key object from DB and retry once if a DB connection error can be healed. + The gate wait, the query, the reconnect, and the retry share one deadline, so a + stalled database fails the request with ``DBLookupDeadlineExceeded`` instead of + parking it. """ + return await bounded_db_lookup( + _fetch_key_object_from_db_unbounded( + hashed_token=hashed_token, + prisma_client=prisma_client, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ), + name="key", + deadline_seconds=deadline_seconds, + ) + + +async def _fetch_key_object_from_db_unbounded( + hashed_token: str, + prisma_client: PrismaClient, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> BaseModel | None: async with db_lookup_gate.current(): try: return await prisma_client.get_data( @@ -3874,9 +3945,9 @@ async def get_object_permission( # else, check db try: - response: Final = await _dictable_table(ObjectPermissionRepository(prisma_client)).find_unique( - where={"object_permission_id": object_permission_id} - ) + response: Final = await _dictable_table( + ObjectPermissionRepository(prisma_client), "object_permission" + ).find_unique(where={"object_permission_id": object_permission_id}) if response is None: return None @@ -4008,7 +4079,9 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await _model_dump_table(OrganizationRepository(prisma_client)).find_unique(**query_kwargs) + response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( + **query_kwargs + ) except Exception: # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them @@ -4073,7 +4146,7 @@ async def get_org_object_for_request( ) except OrganizationNotFoundError: return None - except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits + except Exception as e: if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) return None @@ -5948,7 +6021,7 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row: Final = await _model_dump_table(ProjectRepository(prisma_client)).find_unique( + project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d4e3b224273..f636a173146 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -120,6 +120,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, team_membership_auth_cache_key, ) +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state @@ -735,8 +736,9 @@ async def _fetch_global_spend_with_event_coordination( """ async def _load_global_spend() -> float | None: - proxy_budget_row: Final = await prisma_client.replica_db.litellm_usertable.find_unique( - where={"user_id": LITELLM_PROXY_BUDGET_NAME} + proxy_budget_row: Final = await bounded_db_lookup( + prisma_client.replica_db.litellm_usertable.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), + name="proxy_budget", ) return float(proxy_budget_row.spend) if proxy_budget_row is not None else None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2e43c6b0d22..904070cfadd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2323,7 +2323,7 @@ class ProxyBaseLLMRequestProcessing: return fallbacks if isinstance(fallbacks, list) and fallbacks else None @staticmethod - def _resolve_fallback_models(model: str, fallbacks: list) -> list | None: + def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None: from litellm.router_utils.fallback_event_handlers import get_fallback_model_group fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 4b3ff2711b3..3c6555662e6 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -13,9 +13,12 @@ from __future__ import annotations import re from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm if TYPE_CHECKING: @@ -28,6 +31,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code" _CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" _ONE_MILLION_SUFFIX: Final = "[1m]" _ONE_MILLION_TOKENS: Final = 1_000_000 +_ALIAS_ENTRIES: Final = TypeAdapter(Mapping[object, object]) +_NO_ALIASES: Final[Mapping[str, str]] = MappingProxyType({}) def configured_display_names( @@ -152,6 +157,77 @@ class ClaudeCodeRoutingNames: ) +@dataclass(frozen=True, slots=True) +class CallerAliases: + """`own` are the caller's key and team alias maps, the names `/v1/models` lists for it. + `rewrite` are the maps `/chat/completions` rewrites its model through, in the order it + applies them: the team's, the key's in `add_litellm_data_to_request`, then the global + `model_alias_map` and the key's again in `common_processing_pre_call_logic`.""" + + own: tuple[object, ...] + rewrite: tuple[object, ...] + + +def caller_alias_maps( + key_aliases: object, + team_aliases: object, + key_team_id: str | None, + listed_team_id: str | None, +) -> CallerAliases: + """Team aliases count only when listing the team the key authenticated as.""" + if listed_team_id is not None and listed_team_id != key_team_id: + return CallerAliases((key_aliases,), (key_aliases, litellm.model_alias_map, key_aliases)) + return CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, litellm.model_alias_map, key_aliases)) + + +def _alias_map(aliases: object) -> Mapping[str, str]: + try: + entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + except ValidationError: + return _NO_ALIASES + return MappingProxyType( + {alias: target for alias, target in entries.items() if isinstance(alias, str) and isinstance(target, str)} + ) + + +def _alias_names(alias_maps: Sequence[Mapping[str, str]]) -> tuple[str, ...]: + return tuple(dict.fromkeys(alias for aliases in alias_maps for alias in aliases)) + + +def _rewrite(model_id: str, alias_maps: Sequence[Mapping[str, str]]) -> str | None: + target: Final = reduce(lambda name, aliases: aliases.get(name, name), alias_maps, model_id) + return None if target == model_id else target + + +def alias_target(model_id: str, aliases: CallerAliases, listed: Container[str] = frozenset()) -> str | None: + """The model group `/chat/completions` rewrites `model_id` to, else None. A `model_id` + already `listed` keeps its own row, so it is never rewritten.""" + if model_id in listed: + return None + return _rewrite(model_id, tuple(_alias_map(alias_map) for alias_map in aliases.rewrite)) + + +def alias_listing_entries( + entries: Sequence[tuple[str, str]], + aliases: CallerAliases, +) -> tuple[tuple[str, str], ...]: + """`entries` plus one `(alias, lookup_id)` row per key or team alias whose target is + listed. An alias colliding with a listed id keeps the listed entry.""" + maps: Final = tuple(_alias_map(alias_map) for alias_map in aliases.rewrite) + own: Final = tuple(_alias_map(alias_map) for alias_map in aliases.own) + lookup_by_response: Final = MappingProxyType(dict(entries)) + lookup_ids: Final = frozenset(lookup_by_response.values()) + targets: Final = MappingProxyType( + {alias: _rewrite(alias, maps) for alias in _alias_names(own) if alias not in lookup_by_response} + ) + added: Final = tuple( + (alias, lookup_by_response.get(target, target)) + for alias, target in targets.items() + if target is not None and (target in lookup_by_response or target in lookup_ids) + ) + return (*entries, *added) + + def claude_code_requested_group( requested: str, llm_router: Router, @@ -218,7 +294,7 @@ class TeamModelNameTranslator: @staticmethod def _response_to_lookup_map( - model_names: list[str], + model_names: Sequence[str], internal_to_public: dict[str, str], ) -> dict[str, str]: """Map each public response id to the first internal lookup id seen in @@ -235,7 +311,7 @@ class TeamModelNameTranslator: @staticmethod def listing_entries( - model_names: list[str], + model_names: Sequence[str], llm_router: Router | None, general_settings: Mapping[str, object], ) -> list[tuple[str, str]]: diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 78036237993..05a4a989152 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -486,7 +486,7 @@ class BaselineAccountingStore: async def _pages( self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None ) -> AsyncIterator[tuple[_StoredRecord, ...]]: - cursor: float | None = None + cursor: float | None = None # rebind-ok: keyset pagination advances after each complete timestamp group while page := _RECORDS.validate_python( tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from)) ): @@ -627,7 +627,7 @@ async def flush_baseline_accounting(client: PrismaClient) -> None: more_queued: Final = bool(client.baseline_accounting_transactions) try: remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5) - except (Exception, asyncio.CancelledError) as error: # noqa: BLE001 # unknown acknowledgements can be replayed safely + except (Exception, asyncio.CancelledError) as error: async with client.baseline_accounting_lock: client.baseline_accounting_transactions.extend(batch) if isinstance(error, asyncio.CancelledError): diff --git a/litellm/proxy/db/db_lookup_gate.py b/litellm/proxy/db/db_lookup_gate.py index 2fd427687bd..d2aefde929f 100644 --- a/litellm/proxy/db/db_lookup_gate.py +++ b/litellm/proxy/db/db_lookup_gate.py @@ -1,7 +1,14 @@ import asyncio -from typing import Final +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar -from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY +from litellm.constants import ( + PROXY_DB_LOOKUP_DEADLINE_SECONDS, + PROXY_DB_LOOKUP_MAX_CONCURRENCY, +) + +LookupT = TypeVar("LookupT") class LoopBoundSemaphore: @@ -20,4 +27,64 @@ class LoopBoundSemaphore: return self._semaphore +class DBLookupDeadlineExceeded(asyncio.TimeoutError): + def __init__(self, lookup: str, deadline_seconds: float) -> None: + super().__init__(f"{lookup} lookup did not answer within {deadline_seconds:g}s") + self.lookup: Final = lookup + self.deadline_seconds: Final = deadline_seconds + + +class DBLookupStallTracker: + __slots__ = ("_clock", "_last_hit") + + def __init__(self, clock: Callable[[], float] = time.monotonic) -> None: + self._clock: Final = clock + self._last_hit: float | None = None + + def record_hit(self) -> None: + self._last_hit = self._clock() + + def clear(self) -> None: + self._last_hit = None + + def stalled_within(self, window_seconds: float) -> bool: + if self._last_hit is None: + return False + return self._clock() - self._last_hit < window_seconds + + db_lookup_gate: Final = LoopBoundSemaphore(PROXY_DB_LOOKUP_MAX_CONCURRENCY) +db_lookup_stall_tracker: Final = DBLookupStallTracker() + + +def _consume_abandoned_lookup(task: asyncio.Future[LookupT]) -> None: + if not task.cancelled(): + task.exception() + + +async def bounded_db_lookup( + lookup: Awaitable[LookupT], + *, + name: str, + deadline_seconds: float | None = None, + tracker: DBLookupStallTracker = db_lookup_stall_tracker, +) -> LookupT: + timeout: Final = PROXY_DB_LOOKUP_DEADLINE_SECONDS if deadline_seconds is None else deadline_seconds + task: Final = asyncio.ensure_future(lookup) + try: + done, _ = await asyncio.wait({task}, timeout=timeout) + except asyncio.CancelledError: + task.cancel() + raise + if task not in done: + task.cancel() + task.add_done_callback(_consume_abandoned_lookup) + tracker.record_hit() + raise DBLookupDeadlineExceeded(name, timeout) + try: + return task.result() + except DBLookupDeadlineExceeded: + raise + except asyncio.TimeoutError as e: + tracker.record_hit() + raise DBLookupDeadlineExceeded(name, timeout) from e diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 38d6fb9b99d..022ca2efc9e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -11,6 +11,7 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, ) +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.secret_managers.main import str_to_bool # Bounds the __cause__/__context__ walk in find_database_service_unavailable_error_in_chain. @@ -104,7 +105,7 @@ class PrismaDBExceptionHandler: """ import prisma.engine.errors - if isinstance(e, DB_CONNECTION_ERROR_TYPES): + if isinstance(e, (*DB_CONNECTION_ERROR_TYPES, DBLookupDeadlineExceeded)): return True if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): return True diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2dd028454d6..f8e102d2682 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,7 +23,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType -from litellm.proxy.db.db_lookup_gate import db_lookup_gate +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository @@ -134,36 +134,9 @@ class SpendCounterReseed: if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: - async with db_lookup_gate.current(): - if counter_key.startswith("spend:key:"): - token: Final = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) - elif counter_key.startswith("spend:team_member:"): - suffix: Final = counter_key[len("spend:team_member:") :] - if ":" not in suffix: - return None - user_id, team_id = suffix.rsplit(":", 1) - row = await TeamMembershipRepository(prisma_client).table.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): - return None - elif counter_key.startswith("spend:org:"): - org_id: Final = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique( - where={"organization_id": org_id} - ) - elif counter_key.startswith("spend:project:"): - project_id: Final = counter_key[len("spend:project:") :] - row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) - else: - return None + row: Final = await bounded_db_lookup( + SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter" + ) except Exception: verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None @@ -171,13 +144,47 @@ class SpendCounterReseed: return None return float(getattr(row, "spend", 0.0) or 0.0) + @staticmethod + async def _counter_row(prisma_client: "PrismaClient", counter_key: str) -> object | None: + async with db_lookup_gate.current(): + if counter_key.startswith("spend:key:"): + token: Final = counter_key[len("spend:key:") :] + return await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) + if counter_key.startswith("spend:team_member:"): + suffix: Final = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + return await TeamMembershipRepository(prisma_client).table.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + if counter_key.startswith("spend:team:"): + return await TeamRepository(prisma_client).table.find_unique( + where={"team_id": counter_key[len("spend:team:") :]} + ) + if counter_key.startswith("spend:user:"): + return await UserRepository(prisma_client).table.find_unique( + where={"user_id": counter_key[len("spend:user:") :]} + ) + if counter_key.startswith("spend:org:"): + return await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": counter_key[len("spend:org:") :]} + ) + if counter_key.startswith("spend:project:"): + return await ProjectRepository(prisma_client).table.find_unique( + where={"project_id": counter_key[len("spend:project:") :]} + ) + return None + @staticmethod async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX): return None where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} try: - row: Final = await EndUserRepository(prisma_client).table.find_unique(where=where) + row: Final = await bounded_db_lookup( + EndUserRepository(prisma_client).table.find_unique(where=where), name="end_user_spend" + ) except Exception: # noqa: BLE001 # a failed floor read falls back to the cached spend, like from_db verbose_proxy_logger.exception("SpendCounterReseed.end_user_from_db: failed for %s", counter_key) return None diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 7bafad26569..c422902d30d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -120,7 +120,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> _OPTIONAL_PresidioPIIMasking, ) - explicit_filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None) + explicit_filter_scope: Final = litellm_params.presidio_filter_scope filter_scope: Final = explicit_filter_scope or ("input" if _is_mcp_only_mode(litellm_params.mode) else "both") run_input: Final = filter_scope in ("input", "both") run_output: Final = filter_scope in ("output", "both") diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 64fd59bbe44..f8801e65c82 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -16,7 +16,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger, verbose_proxy_logger -from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS +from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS from litellm.integrations.SlackAlerting.ms_teams import ( MS_TEAMS_ALERT_HEADERS, build_ms_teams_payload, @@ -44,6 +44,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -1723,7 +1724,7 @@ async def _get_health_readiness_details( # check DB if prisma_client is not None: # if db passed in, check if it's connected - db_health_status: Final = await _db_health_readiness_check() + db_status: Final = _readiness_db_status(await _db_health_readiness_check()) # A configured DB that is not reachable means the worker cannot # serve requests that depend on persisted state (keys, budgets, # spend logs). Return 503 so orchestrators take this pod out of @@ -1733,13 +1734,13 @@ async def _get_health_readiness_details( # report the DB state through the body instead. if ( response is not None - and db_health_status["status"] != "connected" + and db_status != "connected" and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() ): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", - "db": db_health_status["status"], + "db": db_status, "cache": cache_type, "litellm_version": version, "success_callbacks": success_callback_names, @@ -1816,24 +1817,32 @@ def _authorize_drain_request(request: Request) -> None: ) +def _readiness_db_status(db_health_status: DBHealthCache) -> str: + """A pod whose pre-request lookups hit their deadline inside the stall window + reports "stalled" even though the ping succeeds: the ping is a fresh + connection, the stalled lookups are the ones requests actually wait on.""" + if db_health_status["status"] != "connected": + return db_health_status["status"] + if db_lookup_stall_tracker.stalled_within(PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS): + return "stalled" + return "connected" + + async def _resolve_public_readiness_db(response: Response) -> str: """ Return the db status string for the public probe and flip the response to - 503 when a configured DB is unreachable. Mirrors the legacy values: - "Not connected" (no DB configured), "connected", "disconnected". + 503 when a configured DB is unreachable or stalled. Mirrors the legacy values: + "Not connected" (no DB configured), "connected", "disconnected", plus "stalled". """ from litellm.proxy.proxy_server import prisma_client if prisma_client is None: return "Not connected" - db_health_status: Final = await _db_health_readiness_check() - if ( - db_health_status["status"] != "connected" - and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - ): + db_status: Final = _readiness_db_status(await _db_health_readiness_check()) + if db_status != "connected" and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE - return db_health_status["status"] + return db_status @router.get( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 81894a5ff12..4dfea5cd472 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_checks import ( log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( DBSpendUpdateWriter, debitable_model_access_groups, @@ -186,8 +187,8 @@ class _ProxyDBLogger(CustomLogger): ) _metadata["error_information"] = _error_information - _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( - metadata=_metadata, + _metadata = await _ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( + metadata=_metadata, original_exception=original_exception ) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} @@ -472,6 +473,12 @@ class _ProxyDBLogger(CustomLogger): spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) + @staticmethod + async def _enrich_failure_metadata_unless_db_stalled(metadata: dict, original_exception: Exception) -> dict: + if isinstance(original_exception, DBLookupDeadlineExceeded): + return metadata + return await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) + @staticmethod async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: """ @@ -770,7 +777,7 @@ async def _reconcile_budget_reservation_before_db_update( "Failed to invalidate budget reservation counters after pre-persist reconcile failed" ) finally: - budget_reservation["finalized"] = True # rebind-ok: the counter update reads the stamp off the shared dict + budget_reservation["finalized"] = True # rebind-ok: stamps the caller's shared dict for the counter update async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 0bd4eb5a5d8..59c06a3f888 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -528,7 +528,7 @@ def _prisma_value(value: object) -> object: return list(value) if isinstance(value, tuple) else value -def member_budget_patch(source: BaseModel) -> dict[str, Any]: +def member_budget_patch(source: BaseModel) -> Mapping[str, object]: """Map the per-member limit fields a request actually set to their budget-table columns (merge-patch: a sent value updates, an explicit null clears, an absent field is left untouched).""" @@ -561,7 +561,7 @@ async def _upsert_budget_and_membership( user_id: str, existing_budget_id: str | None, user_api_key_dict: UserAPIKeyAuth, - budget_patch: dict[str, Any], + budget_patch: Mapping[str, object], team_default_budget_id: str | None = None, shared_budget_ids: frozenset[str] | None = None, ): @@ -624,9 +624,9 @@ async def _upsert_budget_and_membership( if is_shared_default and not temp_only else None ) - source: Final[Mapping[str, Any]] = source_row.model_dump() if source_row is not None else MappingProxyType({}) + source: Final[Mapping[str, object]] = source_row.model_dump() if source_row is not None else MappingProxyType({}) - create_data: Final[dict[str, Any]] = { # mutable-ok: Prisma create payloads are dict-shaped + create_data: Final[dict[str, object]] = { # mutable-ok: Prisma create payloads are dict-shaped "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", **MappingProxyType( diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index dd3f4ff1b12..ec8fd312766 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -348,7 +348,7 @@ async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _Pre data: Final = {**dumped, "user_id": user.user_id} # mutable-ok: /user/new defaults helper mutates in place data_json: Final = _JSON_OBJECT.validate_python(_update_internal_new_user_params(data, user.request)) with_permission: Final = _JSON_OBJECT.validate_python( - await _set_object_permission(data_json=data_json, prisma_client=prisma_client) # pyright: ignore[reportUnknownArgumentType] # validated by the adapter + await _set_object_permission(data_json=data_json, prisma_client=prisma_client) ) return _PreparedUser(user, _USER_ROW.validate_python(with_permission)) except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only @@ -509,7 +509,7 @@ class _TeamsData(TypedDict): def _default_member_budget_id(team: LiteLLM_TeamTable) -> str | None: metadata: Final = ( _JSON_OBJECT.validate_python( - team.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter + team.metadata # pyright: ignore[reportUnknownMemberType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter ) if team.metadata # pyright: ignore[reportUnknownMemberType] # same bare dict else None diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index 8b5b601fe8a..c7b89a6dd6c 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -204,7 +204,7 @@ def _error_message(exc: BaseException) -> str: if isinstance(exc, HTTPException) and isinstance(exc.detail, dict): return str(exc.detail.get("error", exc.detail)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # HTTPException.detail is untyped if isinstance(exc, HTTPException): - return str(exc.detail) # pyright: ignore[reportUnknownArgumentType] # HTTPException.detail is untyped + return str(exc.detail) return str(exc) or type(exc).__name__ diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 4e60c318f03..eaa03b67b40 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3902,7 +3902,7 @@ async def handle_gigachat_passthrough_router_model( is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] - data: dict[str, Any] = await _read_request_body(request=request) # Any needed for proxy pipeline + data: Final[dict[str, object]] = await _read_request_body(request=request) if user_api_key_dict is not None: auth_metadata: Final = { metadata_key: value diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 2cdeddbea30..040250637ea 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -458,7 +458,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def _is_audio_predict_response( model: str, - json_response: dict, # mutable-ok: predicate inspects the decoded provider response dictionary without mutation + json_response: Mapping[str, object], ) -> bool: return ( VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0 @@ -467,7 +467,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def _get_audio_prediction_count( - json_response: dict, # mutable-ok: counter inspects the decoded provider response dictionary without mutation + json_response: Mapping[str, object], ) -> int: predictions: Final = json_response.get("predictions") if not isinstance(predictions, list): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cb5c41705c2..f97798bb9bf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -422,6 +422,9 @@ from litellm.proxy.common_utils.model_deprecation import collect_model_deprecati from litellm.proxy.common_utils.model_listing_utils import ( ClaudeCodeRoutingNames, TeamModelNameTranslator, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_view_ids, configured_display_names, is_claude_code_client, @@ -11172,14 +11175,13 @@ async def model_list( view_aliases: Final = ( view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None ) + caller_aliases: Final = caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ) routing_names: Final = ClaudeCodeRoutingNames( llm_router, team_id or user_api_key_dict.team_id, - ( - user_api_key_dict.aliases, - user_api_key_dict.team_model_aliases, - view_aliases, - ), + (*caller_aliases.rewrite, view_aliases), ) # Validate scope parameter if provided @@ -11307,7 +11309,9 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + entries: Final = alias_listing_entries( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases + ) for response_id, lookup_id in entries: model_info = create_model_info_response( model_id=lookup_id, @@ -11391,7 +11395,8 @@ async def model_info( ) # Mirror /v1/models' visibility filter so first-occurrence resolution - # cannot land on a deployment the listing had hidden. + # cannot land on a deployment the listing had hidden. Undiscoverable + # models stay retrievable by id, they only drop out of the alias guard. blocked_names: Final = llm_router.get_fully_blocked_model_names() if llm_router is not None else set() unhealthy_names: Final = await get_hidden_unhealthy_model_names( healthy_only=healthy_only, @@ -11401,10 +11406,25 @@ async def model_info( hidden_names: Final = blocked_names | unhealthy_names if hidden_names: all_models = [m for m in all_models if m not in hidden_names] + undiscoverable_names: Final = undiscoverable_model_names( + all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id + ) internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) + aliased_model_id: Final = alias_target( + model_id, + caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ), + frozenset( + response_id + for response_id, _ in TeamModelNameTranslator.listing_entries( + tuple(m for m in all_models if m not in undiscoverable_names), llm_router, settings + ) + ), + ) resolved_model_id: Final = TeamModelNameTranslator.resolve_public_name( - model_id=model_id, + model_id=aliased_model_id or model_id, available_models=all_models, llm_router=llm_router, general_settings=settings, @@ -11434,7 +11454,8 @@ async def model_info( fallback_type=None, llm_router=llm_router, ) - return {**response, "id": internal_to_public.get(resolved_model_id, model_id)} # mutable-ok: response id differs + response_id: Final = model_id if aliased_model_id else internal_to_public.get(resolved_model_id, model_id) + return {**response, "id": response_id} # mutable-ok: response id differs def _blocked_response_usage(original_response: object | None) -> "litellm.Usage": diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ca4008dba3d..75eefb2e73b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence from enum import Enum from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args from uuid import uuid4 import fastapi @@ -49,7 +49,7 @@ if TYPE_CHECKING: router: Final = APIRouter() -_ResponseDocSchemas = dict[int | str, dict[str, Any]] # pyright: ignore[reportExplicitAny] # fastapi's responses kwarg +_ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {200: {"model": ResponsesAPIResponse}} RESPONSES_API_CREATE_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 9515fecb586..ebf8e59a72a 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -198,10 +198,6 @@ async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) -async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]: - return (await _scan_pending(prisma_client)).days - - async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun overwrites every group with the same totals.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6a20dc3ffdd..34c952270a1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2389,7 +2389,6 @@ class ProxyLogging: ) try: - # Execute guardrail pipelines before the normal callback loop if not skip_guardrails: data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e421cae0724..bd239922fd3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2254,7 +2254,7 @@ class LiteLLMCompletionResponsesConfig: ) -> Mapping[str, ResponseFunctionWebSearch]: calls: Final[dict[str, ResponseFunctionWebSearch]] = {} # mutable-ok: indexes provider-built calls for choice in chat_completion_response.choices: - provider_fields = getattr(choice.message, "provider_specific_fields", None) + provider_fields = choice.message.provider_specific_fields if not isinstance(provider_fields, Mapping): continue web_search_calls = provider_fields.get("web_search_calls") diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 64989c4cf1c..70f2a7db6da 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1360,7 +1360,7 @@ def _billed_terminal_response( return None usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict return ResponsesAPIResponse.model_construct( - **{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread + **{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportArgumentType] # same untyped dict spread ) diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index c0d0d1de8e3..02e57975626 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,7 +1,7 @@ from collections.abc import Mapping from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, Literal, NamedTuple, Protocol +from typing import Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias from uuid import uuid4 import httpx @@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.typesafe_passthr from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN -JevProbability = Annotated[float, Field(ge=0.0, le=1.0)] +JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] DEFAULT_JEV_INSTRUCTIONS: Final = _DEFAULT_JEV_INSTRUCTIONS diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index cb5c3089685..e5e40d7d6f5 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -50,6 +50,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.prompt_templates.common_utils import ( + anthropic_content_lists, encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) @@ -155,10 +156,7 @@ class EncryptedContentAffinityCheck(CustomLogger): return iter(()) return ( cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance - for message in cast(list[object], messages) # cast-ok: narrowed by isinstance - if isinstance(message, Mapping) - for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance - if isinstance(content, list) + for content in anthropic_content_lists(cast(list[object], messages)) # cast-ok: narrowed by isinstance for block in cast(list[object], content) # cast-ok: narrowed by isinstance if isinstance(block, Mapping) ) diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 076b7759c6d..94ccddc92b7 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass -from typing import Final, Generic, TypeVar +from typing import Final, Generic, TypeAlias, TypeVar from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding @@ -10,11 +10,11 @@ from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.configuration import decision as rollout_decision -RequestT = TypeVar("RequestT") -NativeT = TypeVar("NativeT") -ResultT = TypeVar("ResultT") +RequestT: Final = TypeVar("RequestT") +NativeT: Final = TypeVar("NativeT") +ResultT: Final = TypeVar("ResultT") -NativeHook = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT] +NativeHook: TypeAlias = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT] def call_hook( diff --git a/litellm/rust_bridge/response_metadata.py b/litellm/rust_bridge/response_metadata.py index 1c03515720e..ae459710b34 100644 --- a/litellm/rust_bridge/response_metadata.py +++ b/litellm/rust_bridge/response_metadata.py @@ -1,10 +1,10 @@ -from typing import TypeVar +from typing import Final, TypeVar from litellm.router_utils.add_retry_fallback_headers import ( _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # reuse the proxy's identity-preserving response metadata writer ) -ResultT = TypeVar("ResultT") +ResultT: Final = TypeVar("ResultT") def mark_rust_response(response: ResultT) -> ResultT: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 3674bb670d5..e4c41c3ee5b 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -558,7 +558,6 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict): message: str # Specifies any errors that occur during generation. -# TwelveLabs Marengo Embed types TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 064e3040054..caf88e5d517 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3860,7 +3860,7 @@ def without_server_derived_pricing(model_info: Mapping[str, Any]) -> Mapping[str ) -def echoed_cost_map_pricing_fields(model_info: Mapping[str, Any]) -> tuple[str, ...]: +def echoed_cost_map_pricing_fields(model_info: Mapping[str, object]) -> tuple[str, ...]: """Pricing fields a stored ``model_info`` blob copied from a ``/model/info`` response. Only ``litellm.get_model_info`` emits ``key`` (the resolved cost-map entry), so a stored @@ -3891,7 +3891,7 @@ def echoed_cost_map_fields( ) -def pricing_override_fields(*sources: Mapping[str, Any]) -> tuple[str, ...]: +def pricing_override_fields(*sources: Mapping[str, object]) -> tuple[str, ...]: return tuple( sorted( frozenset( diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d2aec118d31..69a64de6e04 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "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, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "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, @@ -5546,7 +5546,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "azure", @@ -8647,6 +8647,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -8695,6 +8696,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.375e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-07, "cache_read_input_token_cost": 1.1e-08, @@ -8743,6 +8745,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, "cache_read_input_token_cost": 2.2e-07, @@ -10657,7 +10660,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -27347,9 +27350,11 @@ }, "vertex_ai/gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai", @@ -27359,6 +27364,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "regional_endpoint_uplift_multiplier": 1.1, @@ -29601,9 +29607,11 @@ }, "gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", @@ -29613,6 +29621,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -41233,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.396e-07, + "input_cost_per_token": 9.24462e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8792e-06, + "output_cost_per_token": 1.848924e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.83e-08, + "cache_read_input_token_cost": 7.70385e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -48128,6 +48137,20 @@ "/v1/realtime" ] }, + "vertex_ai/chirp_2": { + "input_cost_per_second": 0.00026667, + "litellm_provider": "vertex_ai", + "metadata": { + "calculation": "$0.016/60 seconds = $0.00026667 per second", + "original_pricing_per_minute": 0.016 + }, + "mode": "audio_transcription", + "source": "https://cloud.google.com/speech-to-text/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime" + ] + }, "vertex_ai/claude-3-5-haiku": { "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, @@ -50152,6 +50175,7 @@ ] }, "vertex_ai/google/gemma-4-26b-a4b-it-maas": { + "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 262144, @@ -57133,6 +57157,130 @@ "/v1/responses" ] }, + "us.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "us.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, "global.openai.gpt-5.6-luna": { "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, @@ -63236,6 +63384,115 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/FW-DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.1e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3": { + "cache_read_input_token_cost": 3.25e-07, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3-Flash": { + "cache_read_input_token_cost": 3.8e-08, + "input_cost_per_token": 1.88e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.25e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GPT-OSS-120B": { + "cache_read_input_token_cost": 8.2e-08, + "input_cost_per_token": 1.65e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure_ai/Cohere-command-a-plus-05-2026": { + "input_cost_per_token": 8e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/mistral-medium-3-5": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_response_schema": true, + "supports_vision": true + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -65454,13 +65711,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.4e-07, - "output_cost_per_token": 2.64e-06, - "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943717, + "max_tokens": 943717, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -65843,7 +66100,7 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.062e-07, + "input_cost_per_token": 6.562e-07, "output_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", @@ -66165,9 +66422,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 8.554e-08, + "output_cost_per_token": 1.7108e-07, + "cache_read_input_token_cost": 1.7108e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -68634,7 +68891,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68645,7 +68902,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -68811,6 +69068,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -68824,6 +69082,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 3e-07, "cache_read_input_token_cost": 1.2e-08, @@ -68872,6 +69131,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/eu/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost": 2.4e-07, @@ -69066,7 +69326,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69077,7 +69337,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -69282,6 +69542,114 @@ "mode": "embedding", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "azure/gpt-live-1": { + "input_cost_per_second": 0.000833333333333, + "litellm_provider": "azure", + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, + "azure/gpt-live-transcribe": { + "input_cost_per_second": 0.000283333333333, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "azure", + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-realtime-translate": { + "input_cost_per_second": 0.000566666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, "litellm_provider": "aihubmix", @@ -70438,6 +70806,34 @@ "output_cost_per_token": 0.0, "source": "https://docs.typesafe.ai/models" }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true + }, + "wandb/google/gemma-4-26B-A4B-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "wandb", + "max_input_tokens": 262000, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "wandb/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 1.5e-07, diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index f695af3cb11..3bd98ff5b0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -17,6 +17,7 @@ - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"} - {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"} +- {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bd7f5171172..c96f4b0bef1 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1426,6 +1426,7 @@ class TeamNewBody(BaseModel): team_id: str | None = None organization_id: str | None = None metadata: TeamMetadata | None = None + model_aliases: dict[str, str] | None = None class TeamNewResponse(BaseModel): diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 057a735e5e5..93c198586f6 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -14,12 +14,15 @@ own endpoints, so no test ever holds a signing key. from __future__ import annotations from dataclasses import dataclass +from typing import Final -from e2e_http import AuthHeaders, NoBody, ProbeResult, Result +from e2e_http import AnthropicHeaders, AuthHeaders, NoBody, ProbeResult, Result from idp import Keycloak, keycloak_from_env from models import ( ChatBody, ChatResponse, + ModelsListParams, + ModelsListResponse, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -88,6 +91,17 @@ class OtherClient: response_type=ChatResponse, ) + def list_models_as(self, token: str, *, anthropic: bool = False) -> Result[ModelsListResponse]: + """GET /v1/models under `token`, in the OpenAI shape or, with `anthropic`, the + Anthropic Models API shape Claude Code reads. Both carry `data[].id`.""" + bearer: Final = self.proxy.transport.bearer(token) + return self.proxy.transport.get( + "/v1/models", + headers=AnthropicHeaders(authorization=bearer.authorization) if anthropic else bearer, + params=ModelsListParams(return_wildcard_routes=False), + response_type=ModelsListResponse, + ) + def list_users_as(self, key: str) -> Result[UserListResponse]: """GET /user/list under `key`. Admin-only, so it doubles as the master key's authorization proof: the master key (proxy admin) reads it, a diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index de06c376f76..a8aa474f04f 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -78,9 +78,35 @@ def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam: return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias) -def _ping() -> ChatBody: +@dataclass(frozen=True, slots=True) +class AliasedTeam: + identity: Identity + alias: str + target: str + + +@pytest.fixture +def aliased_team(client: OtherClient, resources: ResourceManager) -> AliasedTeam: + """An identity whose team carries a model_aliases entry, the name a managed + client such as Claude Code sends and the team rewrites to a real model group.""" + marker: Final = unique_marker() + provisioned: Final = _provision(client, resources, marker=marker) + alias: Final = f"e2e-jwt-model-alias-{marker}" + team_id: Final = client.proxy.create_team( + TeamNewBody( + team_alias=f"e2e-jwt-aliased-{marker}", + team_id=provisioned.group, + models=[CHEAP_OPENAI_MODEL], + model_aliases={alias: CHEAP_OPENAI_MODEL}, + ) + ) + resources.defer(lambda: client.proxy.delete_team(team_id)) + return AliasedTeam(identity=provisioned, alias=alias, target=CHEAP_OPENAI_MODEL) + + +def _ping(model: str = CHEAP_OPENAI_MODEL) -> ChatBody: return ChatBody( - model=CHEAP_OPENAI_MODEL, + model=model, messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], max_tokens=16, ) @@ -221,6 +247,23 @@ class TestJwtTeamHeader: f"{bound_team.team_id!r}, got {by_alias!r}" ) + @pytest.mark.covers("other.auth.jwt.team_model_alias_listed_and_routes") + @pytest.mark.parametrize("anthropic", [False, True], ids=["openai_shape", "anthropic_shape"]) + def test_team_model_alias_is_listed_by_v1_models_under_the_same_token_that_routes_it( + self, client: OtherClient, aliased_team: AliasedTeam, anthropic: bool + ) -> None: + token: Final = client.idp.access_token(aliased_team.identity) + + routed: Final = unwrap(client.proxy.chat(token, _ping(model=aliased_team.alias))) + assert routed.choices, f"precondition: /chat/completions must route the team alias, got {routed}" + + listed: Final = tuple(entry.id for entry in unwrap(client.list_models_as(token, anthropic=anthropic)).data) + assert aliased_team.alias in listed, ( + f"/v1/models must list team alias {aliased_team.alias!r} that the same token routes on " + f"/chat/completions, got {listed}" + ) + assert aliased_team.target in listed, f"the alias target {aliased_team.target!r} must stay listed, got {listed}" + @pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied") def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value( self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index 84c3a4e244a..b50d6e19458 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -115,7 +115,7 @@ def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: [ ("vertex_ai/chirp_3", True), ("chirp_3", True), - ("chirp_2", False), + ("chirp_2", True), ("gemini-live-2.5-flash", False), ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), ("vertex_ai/gemini-3.5-transcribe-live-preview", False), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index f6c7faaf28d..05199bbcb13 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -646,6 +646,115 @@ async def test_fetch_key_object_from_db_bounds_in_flight_prisma_requests(): assert prisma.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY +@pytest.fixture +def _clear_db_lookup_stall() -> Iterator[None]: + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + db_lookup_stall_tracker.clear() + yield + db_lookup_stall_tracker.clear() + + +class _StalledPrisma: + def __init__(self) -> None: + self.attempt_db_reconnect = AsyncMock(return_value=True) + self.db = MagicMock() + self.db.litellm_teamtable.find_unique = AsyncMock(side_effect=_stall_forever) + self.db.litellm_teamtable.update = AsyncMock(side_effect=_answer_slowly) + self.replica_db = self.db + + async def get_data(self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None) -> None: + await _stall_forever() + + +async def _stall_forever(**kwargs: object) -> None: + await asyncio.Event().wait() + + +async def _answer_slowly(**kwargs: object) -> Mapping[str, object]: + await asyncio.sleep(0.15) + return {"team_id": "slow-write"} + + +@pytest.mark.asyncio +async def test_fetch_key_object_from_db_fails_a_stalled_burst_within_the_deadline_without_reconnecting( + _clear_db_lookup_stall, +): + """The incident: a stalled database parked every request in the pod with liveness + and readiness green until it OOMed. Every lookup in a burst larger than the gate, + the ones queued behind it included, must fail within one deadline, must not try to + reconnect (the transport is fine, the query is slow), and must leave every gate slot + free for the next burst.""" + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + prisma: Final = _StalledPrisma() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 3 + started: Final = time.monotonic() + + results: Final = await asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"hashed-token-{i}", + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + deadline_seconds=0.2, + ) + for i in range(burst) + ), + return_exceptions=True, + ) + elapsed: Final = time.monotonic() - started + + assert len(results) == burst + assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) + assert all(PrismaDBExceptionHandler.is_database_service_unavailable_error(result) for result in results) + assert elapsed < 3 + prisma.attempt_db_reconnect.assert_not_awaited() + + recovered: Final = _InFlightCountingPrisma() + after: Final = await asyncio.wait_for( + asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"after-{i}", + prisma_client=recovered, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + ) + for i in range(PROXY_DB_LOOKUP_MAX_CONCURRENCY) + ) + ), + timeout=5, + ) + assert {r.token for r in after if r is not None} == {f"after-{i}" for i in range(PROXY_DB_LOOKUP_MAX_CONCURRENCY)} + + +@pytest.mark.asyncio +async def test_team_lookup_fails_at_the_db_lookup_deadline_while_writes_stay_unbounded(_clear_db_lookup_stall): + """Team, user, budget, and membership reads share the key lookup's deadline through + the typed table wrappers; writes do not, since a slow write must land rather than + fail the request that already passed auth.""" + from litellm.proxy.auth.auth_checks import _team_table + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + from litellm.repositories.table_repositories import TeamRepository + + prisma: Final = _StalledPrisma() + with patch( # test-quality-ok: lowers the module-level lookup deadline so the stalled-read test finishes fast + "litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05 + ): + started: Final = time.monotonic() + with pytest.raises(DBLookupDeadlineExceeded, match=r"team lookup did not answer within 0\.05s"): + await _get_team_db_check(team_id="stalled-team", prisma_client=prisma) # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + assert time.monotonic() - started < 2 + + written: Final = await _team_table(TeamRepository(prisma)).update( + where={"team_id": "slow-write"}, data={"spend": 1.0} + ) + + assert written == {"team_id": "slow-write"} + + def _fake_redis_cache(): fake_redis = MagicMock() fake_redis.async_get_cache = AsyncMock(return_value=None) @@ -6221,7 +6330,9 @@ async def test_get_org_object_for_request_serves_last_known_org_through_db_outag proxy_logging_obj=None, ) - with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists + with patch( + "litellm.proxy.proxy_server.general_settings", {} + ): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists warm = await _lookup() assert warm is not None and warm.organization_alias == "platform-org" await user_api_key_cache.async_delete_cache("org_id:org-1:with_budget") @@ -9015,20 +9126,34 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s client: Final = MagicMock(db=MagicMock(litellm_accessgrouptable=MagicMock(find_unique=reader))) client.replica_db = client.db with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: [TQ008] prove reads stay on the injected connection - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), # test-quality-ok: [TQ008] isolate the process cache + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: [TQ008] prove reads stay on the injected connection + patch( + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), # test-quality-ok: [TQ008] isolate the process cache ): if channel == "team": - assert await can_team_access_model( - model="allowed", team_object=LiteLLM_TeamTable(team_id="team-a", models=["other"], access_group_ids=["group-a"]), - llm_router=None, prisma_client=client, - ) is True + assert ( + await can_team_access_model( + model="allowed", + team_object=LiteLLM_TeamTable(team_id="team-a", models=["other"], access_group_ids=["group-a"]), + llm_router=None, + prisma_client=client, + ) + is True + ) else: - assert await can_key_call_model( - model="allowed", llm_model_list=None, - valid_token=UserAPIKeyAuth(models=["other"], access_group_ids=["group-a"]), - llm_router=None, prisma_client=client, - ) is True + assert ( + await can_key_call_model( + model="allowed", + llm_model_list=None, + valid_token=UserAPIKeyAuth(models=["other"], access_group_ids=["group-a"]), + llm_router=None, + prisma_client=client, + ) + is True + ) reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) @@ -9049,6 +9174,7 @@ def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default(): litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]), ) + def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None: assert route_skips_budget_checks(route="/v1/models") is True assert route_skips_budget_checks(route="/spend/logs") is True @@ -9166,7 +9292,9 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): return fallback_spend with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, @@ -9194,7 +9322,9 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): ), ) with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, @@ -9256,7 +9386,9 @@ async def test_team_member_budget_check_adds_temp_increase_to_live_team_default( return fallback_spend with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 46175e73b2c..07868f6c307 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4,6 +4,7 @@ import logging import os import subprocess import sys +import time from collections.abc import Mapping from contextlib import contextmanager from datetime import datetime, timedelta, timezone @@ -186,11 +187,7 @@ async def test_disable_budget_reservation_does_not_log_per_request(caplog): general_settings={"disable_budget_reservation": True}, ) - records = [ - record - for record in caplog.records - if "disable_budget_reservation is enabled" in record.message - ] + records = [record for record in caplog.records if "disable_budget_reservation is enabled" in record.message] assert records == [] assert user_api_key_auth_obj.budget_reservation is None @@ -234,9 +231,7 @@ async def test_budget_reservation_runs_when_not_disabled(): ({}, False), ], ) -async def test_fail_closed_budget_enforcement_reaches_reservation( - general_settings, expected_flag -): +async def test_fail_closed_budget_enforcement_reaches_reservation(general_settings, expected_flag): """#33923: the strict flag must be threaded into reserve_budget_for_request so a failed reservation write can reject instead of failing open.""" user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") @@ -259,10 +254,7 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( general_settings=general_settings, ) - assert ( - mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] - is expected_flag - ) + assert mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] is expected_flag @pytest.mark.asyncio @@ -274,9 +266,7 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( ({}, False), ], ) -async def test_apply_user_budget_to_team_keys_reaches_reservation( - general_settings, expected_flag -): +async def test_apply_user_budget_to_team_keys_reaches_reservation(general_settings, expected_flag): """The opt-in lives in general_settings but is consumed inside _get_budget_counters, so it has to be threaded through reserve_budget_for_request or the reservation path keeps exempting team keys while the read path enforces.""" @@ -300,9 +290,7 @@ async def test_apply_user_budget_to_team_keys_reaches_reservation( general_settings=general_settings, ) - assert ( - mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag - ) + assert mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag @pytest.mark.asyncio @@ -402,9 +390,7 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -435,9 +421,7 @@ async def test_custom_auth_enforces_key_model_access_from_file_route_header_with "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -468,9 +452,7 @@ async def test_custom_auth_honors_key_level_model_access_restriction_denied_with "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -506,9 +488,7 @@ def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) return { @@ -770,9 +750,7 @@ async def test_enterprise_custom_auth_runs_post_custom_auth_checks_when_opt_in() litellm.enable_post_custom_auth_checks = original_flag -def _assert_get_api_key_with_custom_litellm_key_header( - custom_litellm_key_header, api_key, passed_in_key -): +def _assert_get_api_key_with_custom_litellm_key_header(custom_litellm_key_header, api_key, passed_in_key): assert get_api_key( custom_litellm_key_header=custom_litellm_key_header, api_key=None, @@ -829,9 +807,7 @@ def _assert_get_api_key_with_custom_litellm_key_header( ("App:LiteLLM", None, False, False), ], ) -def test_routing_selector_matches_claim_parametrized( - selector_value, claim_value, expected, split_space_delimited -): +def test_routing_selector_matches_claim_parametrized(selector_value, claim_value, expected, split_space_delimited): assert ( _routing_selector_matches_claim( selector_value=selector_value, @@ -925,10 +901,7 @@ def test_routing_selector_matches_claim_parametrized( ], ) def test_matches_routing_override_parametrized(override, token_claims, expected): - assert ( - _matches_routing_override(token_claims=token_claims, override=override) - is expected - ) + assert _matches_routing_override(token_claims=token_claims, override=override) is expected def test_get_api_key_with_custom_litellm_key_header_bearer_prefix(): @@ -1007,12 +980,9 @@ def test_team_metadata_with_tags_flows_through_jwt_auth(): ) # Verify team_metadata is set - assert ( - user_api_key_auth.team_metadata is not None - ), "team_metadata should be populated" + assert user_api_key_auth.team_metadata is not None, "team_metadata should be populated" assert user_api_key_auth.team_metadata == team_object.metadata, ( - f"team_metadata not correctly mapped. " - f"Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" + f"team_metadata not correctly mapped. Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" ) # Specifically verify tags are present @@ -1051,9 +1021,7 @@ def test_route_checks_is_llm_api_route(): ] for route in openai_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Anthropic routes anthropic_routes = [ @@ -1062,9 +1030,7 @@ def test_route_checks_is_llm_api_route(): ] for route in anthropic_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test passthrough routes (this is the key improvement over the old route checking) passthrough_routes = [ @@ -1084,9 +1050,7 @@ def test_route_checks_is_llm_api_route(): ] for route in passthrough_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test MCP routes mcp_routes = [ @@ -1096,9 +1060,7 @@ def test_route_checks_is_llm_api_route(): ] for route in mcp_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test LiteLLM native RAG routes rag_routes = [ @@ -1108,9 +1070,7 @@ def test_route_checks_is_llm_api_route(): "/v1/rag/query", ] for route in rag_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test routes with placeholders placeholder_routes = [ @@ -1125,9 +1085,7 @@ def test_route_checks_is_llm_api_route(): ] for route in placeholder_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Azure OpenAI routes azure_routes = [ @@ -1138,9 +1096,7 @@ def test_route_checks_is_llm_api_route(): ] for route in azure_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test non-LLM routes (should return False) non_llm_routes = [ @@ -1159,9 +1115,7 @@ def test_route_checks_is_llm_api_route(): ] for route in non_llm_routes: - assert not RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should NOT be identified as LLM API route" + assert not RouteChecks.is_llm_api_route(route=route), f"Route {route} should NOT be identified as LLM API route" # Test invalid inputs invalid_inputs = [ @@ -1173,9 +1127,9 @@ def test_route_checks_is_llm_api_route(): ] for invalid_input in invalid_inputs: - assert not RouteChecks.is_llm_api_route( - route=invalid_input - ), f"Invalid input {invalid_input} should return False" + assert not RouteChecks.is_llm_api_route(route=invalid_input), ( + f"Invalid input {invalid_input} should return False" + ) @pytest.mark.asyncio @@ -1222,9 +1176,7 @@ async def test_proxy_admin_expired_key_from_cache(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() # Mock post_call_failure_hook as async function returning None (no transformation) mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) @@ -1261,9 +1213,7 @@ async def test_proxy_admin_expired_key_from_cache(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1287,36 +1237,30 @@ async def test_proxy_admin_expired_key_from_cache(): ) # Verify that ProxyException was raised with expired_key type - assert hasattr( - exc_info.value, "type" - ), "Exception should have 'type' attribute" - assert ( - exc_info.value.type == ProxyErrorTypes.expired_key - ), f"Expected expired_key error type, got {exc_info.value.type}" + assert hasattr(exc_info.value, "type"), "Exception should have 'type' attribute" + assert exc_info.value.type == ProxyErrorTypes.expired_key, ( + f"Expected expired_key error type, got {exc_info.value.type}" + ) assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED - assert "Expired Key" in str( - exc_info.value.message - ), f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + assert "Expired Key" in str(exc_info.value.message), ( + f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + ) # Verify that the param field does NOT leak the full API key (Issue #18731) # The param should be abbreviated like "sk-...XXXX" not the full plaintext key - assert ( - exc_info.value.param is not None - ), "Exception should have 'param' attribute" + assert exc_info.value.param is not None, "Exception should have 'param' attribute" assert exc_info.value.param != api_key, ( f"SECURITY: Full API key should NOT be in param field! " f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'" ) - assert exc_info.value.param.startswith( - "sk-..." - ), f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + assert exc_info.value.param.startswith("sk-..."), ( + f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + ) # Verify that cache deletion was called mock_delete_cache.assert_called_once() call_args = mock_delete_cache.call_args - assert ( - call_args[1]["hashed_token"] == hashed_key - ), "Cache deletion should be called with the hashed key" + assert call_args[1]["hashed_token"] == hashed_key, "Cache deletion should be called with the hashed key" finally: # Restore all module-level attributes so subsequent tests are not affected for attr, val in _original_values.items(): @@ -1354,9 +1298,7 @@ async def test_scim_deactivated_user_key_is_rejected(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) mock_prisma_client = MagicMock() @@ -1377,9 +1319,7 @@ async def test_scim_deactivated_user_key_is_rejected(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1446,9 +1386,7 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -1467,9 +1405,7 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1521,9 +1457,7 @@ async def test_master_key_auth_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -1542,9 +1476,7 @@ async def test_master_key_auth_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1597,9 +1529,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) mock_prisma_client = MagicMock() @@ -1620,9 +1550,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -2155,7 +2083,10 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(post_call_failure_hook=AsyncMock(return_value=None)), + ), patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), patch( "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", @@ -2220,7 +2151,9 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id(): plaintext = "sk-auto-registered-agent" token_hash = hash_token(plaintext) persisted_principal = IdentityStore._principal_from_key( - UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"), + UserAPIKeyAuth( + token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id" + ), auth_method=AuthMethod.API_KEY, credential_ref=CredentialRef(token_id=token_hash), ) @@ -2437,10 +2370,7 @@ class TestJWTOAuth2Coexistence: def test_is_jwt_detects_jwt_tokens(self): """JWT tokens have 3 dot-separated parts.""" assert JWTHandler.is_jwt("header.payload.signature") is True - assert ( - JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") - is True - ) + assert JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") is True def test_is_jwt_rejects_opaque_tokens(self): """Opaque OAuth2 tokens do not have 3 dot-separated parts.""" @@ -2549,10 +2479,7 @@ class TestJWTOAuth2Coexistence: assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "403" - assert ( - "Oauth2 token validation is only available for premium users" - in exc_info.value.message - ) + assert "Oauth2 token validation is only available for premium users" in exc_info.value.message mock_oauth2.assert_not_called() @pytest.mark.asyncio @@ -2745,9 +2672,7 @@ class TestJWTOAuth2Coexistence: assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team" assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user" assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org" - assert ( - mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" - ) + assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" assert result.org_id == "validated-org" assert result.user_email == "validated@example.com" @@ -2826,10 +2751,7 @@ class TestJWTOAuth2Coexistence: assert result.user_id == "mapped-user" assert result.user_email == "mapped@example.com" - assert ( - mock_get_user_object.call_args_list[0].kwargs["user_email"] - == "mapped@example.com" - ) + assert mock_get_user_object.call_args_list[0].kwargs["user_email"] == "mapped@example.com" @pytest.mark.asyncio async def test_mapped_virtual_key_does_not_backfill_mismatched_owner(self): @@ -2906,8 +2828,7 @@ class TestJWTOAuth2Coexistence: assert result.user_id == "other-owner" assert result.user_email is None assert all( - call.kwargs.get("user_email") != "principal@example.com" - for call in mock_get_user_object.call_args_list + call.kwargs.get("user_email") != "principal@example.com" for call in mock_get_user_object.call_args_list ) @pytest.mark.asyncio @@ -3713,9 +3634,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -3847,9 +3766,7 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -3899,9 +3816,9 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): request_data={}, ) - assert result.team_metadata == { - "guardrails": ["test-guardrail-333"] - }, f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" + assert result.team_metadata == {"guardrails": ["test-guardrail-333"]}, ( + f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" + ) finally: for k, v in _originals.items(): @@ -4226,9 +4143,7 @@ async def test_auth_flow_fallback_team_object_permission_none_when_unreadable(): # --------------------------------------------------------------------------- -def _proxy_attrs_for_centralized_checks( - user_custom_auth=None, flag=False, master_key="sk-test-master" -): +def _proxy_attrs_for_centralized_checks(user_custom_auth=None, flag=False, master_key="sk-test-master"): """Build the minimal proxy_server module attributes that _run_centralized_common_checks reads. @@ -4439,9 +4354,7 @@ async def _run_centralized_checks_with_key_end_user_budget( request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") attrs = { - **_proxy_attrs_for_centralized_checks( - user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth - ), + **_proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth), "prisma_client": prisma_client, "user_api_key_cache": user_api_key_cache if user_api_key_cache is not None else DualCache(), "proxy_logging_obj": proxy_logging_obj, @@ -4632,7 +4545,9 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the for k, v in attrs.items(): setattr(_proxy_server_mod, k, v) with ( - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test patch( # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", new_callable=AsyncMock, @@ -4665,9 +4580,7 @@ async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - attrs = _proxy_attrs_for_centralized_checks( - user_custom_auth=AsyncMock(), flag=False - ) + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock(), flag=False) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} try: for k, v in attrs.items(): @@ -5073,9 +4986,7 @@ async def test_centralized_common_checks_reserves_request_end_user_budget(): "applied_adjustment": 0.0, } ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:alice" - ) == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:end_user:alice") == pytest.approx(0.6) @pytest.mark.asyncio @@ -5090,9 +5001,7 @@ async def test_centralized_common_checks_short_circuits_when_master_key_unset(): from litellm.proxy._types import LitellmUserRoles - token = UserAPIKeyAuth( - api_key="sk-test", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER - ) + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER) request = Request(scope={"type": "http"}) request._url = URL(url="/get/config/callbacks") @@ -5893,9 +5802,7 @@ async def test_centralized_common_checks_user_http_exception_isolates_to_user_on request._url = URL(url="/chat/completions") request._body = json.dumps({"user": "alice", "model": "gpt-4o"}).encode() - fetched_team = LiteLLM_TeamTableCachedObj( - team_id="t1", max_budget=20.0, models=["gpt-4o"] - ) + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", max_budget=20.0, models=["gpt-4o"]) fetched_end_user = LiteLLM_EndUserTable(user_id="alice", blocked=False, spend=1.0) fetched_project = LiteLLM_ProjectTableCachedObj( project_id="proj-1", @@ -6024,10 +5931,46 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, ("org-pinned", None, None, "preset", None, "success", False, False, "org-pinned", "preset", (None, None, None)), ("org-view", None, None, None, 3, "success", False, False, "org-view", None, (None, None, 3)), ("org-missing", None, None, None, None, "missing", False, False, "org-missing", None, (None, None, None)), - ("org-db-failure-allowed", None, None, None, None, "db_failure", True, False, "org-db-failure-allowed", None, (None, None, None)), - ("org-db-failure-denied", None, None, None, None, "db_failure", False, True, "org-db-failure-denied", None, (None, None, None)), + ( + "org-db-failure-allowed", + None, + None, + None, + None, + "db_failure", + True, + False, + "org-db-failure-allowed", + None, + (None, None, None), + ), + ( + "org-db-failure-denied", + None, + None, + None, + None, + "db_failure", + False, + True, + "org-db-failure-denied", + None, + (None, None, None), + ), ("org-bad-row", None, None, None, None, "bad_row", False, False, "org-bad-row", None, (None, None, None)), - ("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)), + ( + "org-nobudget", + None, + None, + None, + None, + "no_budget", + False, + False, + "org-nobudget", + "acme-org", + (None, None, None), + ), ], ) async def test_centralized_common_checks_inherits_org_identity( @@ -6336,9 +6279,7 @@ async def test_user_api_key_auth_sets_end_user_id_when_builder_skips_it(): } ) request._url = URL(url="/chat/completions") - request._body = json.dumps( - {"model": "gpt-4o", "user": "alice@example.com"} - ).encode() + request._body = json.dumps({"model": "gpt-4o", "user": "alice@example.com"}).encode() attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} @@ -6382,9 +6323,7 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() import litellm.proxy.proxy_server as _proxy_server_mod - builder_token = UserAPIKeyAuth( - api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id" - ) + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id") request = Request( scope={ @@ -6394,9 +6333,7 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() } ) request._url = URL(url="/chat/completions") - request._body = json.dumps( - {"model": "gpt-4o", "user": "different-id-from-body"} - ).encode() + request._body = json.dumps({"model": "gpt-4o", "user": "different-id-from-body"}).encode() attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} @@ -6729,6 +6666,83 @@ async def _run_builder_with_key_lookup(get_key_object_mock): setattr(_proxy_server_mod, k, v) +class _StalledKeyLookupPrisma: + """A database whose connection answers the readiness ping but whose key lookups + never return, which is what the incident's locked table looked like.""" + + def __init__(self) -> None: + self.health_check = AsyncMock(return_value=True) + self.attempt_db_reconnect = AsyncMock(return_value=True) + self.db = MagicMock() + + async def get_data(self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None) -> None: + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_burst_against_a_stalled_db_fails_fast_with_503_and_turns_readiness_red(): + """The incident, end to end: N requests into a proxy whose database stalls used to + park in the pod with readiness green until it OOMed. Now every one of them fails + within the lookup deadline as a 503, and the next readiness probe takes the pod out + of rotation.""" + import httpx + from fastapi import Depends, FastAPI + + import litellm.proxy.health_endpoints._health_endpoints as health_endpoints + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + app = FastAPI() + + @app.post("/chat/completions", dependencies=[Depends(user_api_key_auth)]) + async def chat_completions() -> Mapping[str, bool]: + return {"served": True} + + app.include_router(health_endpoints.router) + app.add_exception_handler(ProxyException, _proxy_server_mod.openai_exception_handler) + + attrs = {**_proxy_attrs_for_db_lookup(), "prisma_client": _StalledKeyLookupPrisma()} + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + health_endpoints.db_health_cache = {"status": "unknown", "last_updated": datetime.now() - timedelta(seconds=60)} + db_lookup_stall_tracker.clear() + burst = 60 + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: lowers the module-level lookup deadline so the stalled burst finishes fast + "litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.2 + ), + patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), + ): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://t") as client: + started = time.monotonic() + responses = await asyncio.gather( + *( + client.post( + "/chat/completions", + json={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": f"Bearer sk-stalled-{i}"}, + ) + for i in range(burst) + ) + ) + elapsed = time.monotonic() - started + readiness = await client.get("/health/readiness") + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + db_lookup_stall_tracker.clear() + + assert len(responses) == burst + assert {r.status_code for r in responses} == {status.HTTP_503_SERVICE_UNAVAILABLE} + assert {r.json()["error"]["type"] for r in responses} == {ProxyErrorTypes.no_db_connection.value} + assert all("temporarily unreachable" in r.json()["error"]["message"] for r in responses) + assert elapsed < 5 + assert readiness.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert readiness.json()["db"] == "stalled" + + @pytest.mark.asyncio async def test_builder_returns_503_when_db_lookup_raises_infra_error(): """End-to-end: a DB infrastructure failure during the key lookup must @@ -6796,9 +6810,7 @@ def _mint_cli_session_token(monkeypatch, *, user_id="cli-admin"): models=["gpt-3.5-turbo"], max_budget=100.0, ) - return ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info, team_id="cli-team", team_alias="cli-team-alias" - ) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info, team_id="cli-team", team_alias="cli-team-alias") @pytest.mark.asyncio @@ -6848,7 +6860,7 @@ async def test_random_non_sk_token_is_rejected(monkeypatch): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): - with pytest.raises(Exception, match='LiteLLM Virtual Key expected\\.') as exc_info: + with pytest.raises(Exception, match="LiteLLM Virtual Key expected\\.") as exc_info: await user_api_key_auth( request=mock_request, api_key="Bearer not-a-real-token", @@ -6927,9 +6939,7 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa user_role=LitellmUserRoles.INTERNAL_USER.value, models=[], ) - cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info, team_id="team-abc", team_alias="my-team" - ) + cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info, team_id="team-abc", team_alias="my-team") import litellm.proxy.proxy_server as _proxy_server_mod from fastapi import Request @@ -7236,7 +7246,7 @@ async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", None), ): - with pytest.raises(Exception, match='JWT Auth is an enterprise only feature\\. You must be a') as exc_info: + with pytest.raises(Exception, match="JWT Auth is an enterprise only feature\\. You must be a") as exc_info: await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_token}", @@ -7275,13 +7285,9 @@ async def test_auth_does_not_rewrite_cached_key_object_back_into_cache(): metadata={"model_rpm_limit": {"gpt-5.4-mini": 3}}, last_refreshed_at=1000.0, ) - await key_cache.async_set_cache( - key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth - ) + await key_cache.async_set_cache(key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth) - fetch_from_db = AsyncMock( - side_effect=AssertionError("cache-hit auth must not touch the DB") - ) + fetch_from_db = AsyncMock(side_effect=AssertionError("cache-hit auth must not touch the DB")) proxy_logging_obj = MagicMock() proxy_logging_obj.internal_usage_cache = MagicMock() @@ -7328,9 +7334,7 @@ async def test_auth_does_not_rewrite_cached_key_object_back_into_cache(): assert result.token == hashed_key fetch_from_db.assert_not_called() - cached_after = await key_cache.async_get_cache( - key=hashed_key, model_type=UserAPIKeyAuth - ) + cached_after = await key_cache.async_get_cache(key=hashed_key, model_type=UserAPIKeyAuth) assert cached_after is not None assert cached_after.last_refreshed_at == 1000.0 assert cached_after.metadata == {"model_rpm_limit": {"gpt-5.4-mini": 3}} @@ -7406,7 +7410,9 @@ class TestJWTAuthUserEmail: assert result.user_email == "resolved@example.com" @pytest.mark.asyncio - @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"]) + @pytest.mark.parametrize( + "route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"] + ) @pytest.mark.parametrize("active", [False, True, None, "false", 0]) @pytest.mark.parametrize("is_admin", [False, True]) async def test_jwt_auth_rejects_deactivated_user( @@ -7481,9 +7487,7 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_within_budget_does_not_reroute(self): - valid_token = UserAPIKeyAuth( - token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} - ) + valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) limiter = AsyncMock() limiter.is_key_within_model_budget.return_value = True request_data = {"model": "gpt-4o"} @@ -7508,9 +7512,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7524,9 +7526,7 @@ class TestCheckKeyModelBudgetWithFallback: ) assert request_data["model"] == "gpt-4o-mini" - limiter.get_fallback_model_within_budget.assert_awaited_once_with( - user_api_key_dict=valid_token, model="gpt-4o" - ) + limiter.get_fallback_model_within_budget.assert_awaited_once_with(user_api_key_dict=valid_token, model="gpt-4o") # the rerouted model must be visible to a later, separate # `_read_request_body` call on the same `request` (route handlers # re-parse the body from this cache instead of reusing the dict). @@ -7535,9 +7535,7 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_raises_when_every_fallback_also_exceeded(self): - valid_token = UserAPIKeyAuth( - token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} - ) + valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) limiter = AsyncMock() original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.is_key_within_model_budget.side_effect = original_error @@ -7607,9 +7605,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7677,9 +7673,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7760,9 +7754,7 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): ) assert result == 42.5 - prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with( - where={"user_id": "litellm-proxy-budget"} - ) + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "litellm-proxy-budget"}) @pytest.mark.asyncio @@ -8138,9 +8130,7 @@ async def test_jwt_shaped_key_error_names_enable_jwt_auth_when_disabled(): Prometheus invalid-key filter and the admin UI both substring-match it. Keys that are not JWT-shaped must not pick up the hint. """ - jwt_error = await _proxy_exception_for_key( - "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9.c2lnbmF0dXJl", {}, True - ) + jwt_error = await _proxy_exception_for_key("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9.c2lnbmF0dXJl", {}, True) assert jwt_error.code == "401" assert "enable_jwt_auth" in jwt_error.message @@ -8150,9 +8140,7 @@ async def test_jwt_shaped_key_error_names_enable_jwt_auth_when_disabled(): assert "is a JWT" not in jwt_error.message opaque_error = await _proxy_exception_for_key("not-a-jwt-at-all", {}, True) - two_segment_error = await _proxy_exception_for_key( - "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9", {}, True - ) + two_segment_error = await _proxy_exception_for_key("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9", {}, True) assert "enable_jwt_auth" not in opaque_error.message assert "enable_jwt_auth" not in two_segment_error.message @@ -8181,9 +8169,7 @@ class TestLitellmReceivedAtStamping: on OTEL being configured to see a true request-arrival timestamp.""" def test_stamped_even_when_otel_is_not_configured(self, monkeypatch): - monkeypatch.setattr( - "litellm.proxy.proxy_server.open_telemetry_logger", None - ) + monkeypatch.setattr("litellm.proxy.proxy_server.open_telemetry_logger", None) request = MagicMock() request.state = SimpleNamespace() @@ -8215,7 +8201,7 @@ class TestLitellmReceivedAtStamping: _RECORDING_DDTRACE = dedent( - ''' + """ import functools import inspect @@ -8264,11 +8250,11 @@ _RECORDING_DDTRACE = dedent( tracer = _Tracer() - ''' + """ ) _DDTRACE_AUTH_PROBE = dedent( - ''' + """ import asyncio import json @@ -8307,7 +8293,7 @@ _DDTRACE_AUTH_PROBE = dedent( asyncio.run(main()) - ''' + """ ) @@ -8454,25 +8440,43 @@ async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_a @pytest.mark.asyncio -@pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) +@pytest.mark.parametrize( + "route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"] +) async def test_claude_view_normalizes_before_model_access(monkeypatch, route): from starlette.requests import Request from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access source = "foo[1m]" encoded = "claude-router-" + source.encode().hex() + "[1m]" - router = litellm.Router(model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}]) + router = litellm.Router( + model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}] + ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) data = {"model": encoded, "messages": [{"role": "user", "content": "hi"}]} request = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) token = UserAPIKeyAuth(models=[source]) - await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + await _enforce_key_and_fallback_model_access( + valid_token=token, + request_data=data, + route=route, + request=request, + llm_model_list=router.model_list, + llm_router=router, + ) assert data["model"] == source assert (await request.json())["model"] == source assert json.loads(await request.body())["model"] == source assert request.scope["parsed_body"][1]["model"] == source with pytest.raises(ProxyException): - await _enforce_key_and_fallback_model_access(valid_token=UserAPIKeyAuth(models=["other"]), request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + await _enforce_key_and_fallback_model_access( + valid_token=UserAPIKeyAuth(models=["other"]), + request_data=data, + route=route, + request=request, + llm_model_list=router.model_list, + llm_router=router, + ) @pytest.mark.asyncio @@ -8484,10 +8488,18 @@ async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer) encoded = "claude-router-666f6f" names = ("foo", "other", encoded) if layer == "literal" else ("foo", "other") alias = {encoded: "other"} - router = litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names], model_group_alias=alias if layer == "router" else None) + router = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names + ], + model_group_alias=alias if layer == "router" else None, + ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) - token = UserAPIKeyAuth(aliases=alias if layer == "key" else {}, router_settings={"model_group_alias": alias} if layer == "hierarchical" else None) + token = UserAPIKeyAuth( + aliases=alias if layer == "key" else {}, + router_settings={"model_group_alias": alias} if layer == "hierarchical" else None, + ) data = {"model": encoded} request = Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}) await _normalize_claude_model(data, token, request, "/v1/messages") diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py index 7ef03140093..97e38d5e17a 100644 --- a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -4,9 +4,14 @@ from itertools import combinations import pytest +import litellm from litellm import Router from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, ClaudeCodeRoutingNames, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_group_name, claude_code_model_id, claude_code_requested_group, @@ -22,6 +27,10 @@ def _marked(name): return f"{_encoded(name)}[1m]" +def _caller(*maps: object) -> CallerAliases: + return CallerAliases(maps, maps) + + def _row(name, limit=1000000): return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} @@ -29,15 +38,16 @@ def _row(name, limit=1000000): def _router(*names, aliases=None): return Router( model_list=[ - {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} - for name in names + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names ], model_group_alias=aliases, ) @pytest.mark.parametrize("limit", [None, 999999, 1000000]) -@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +@pytest.mark.parametrize( + "name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"] +) def test_listing_round_trips_entire_source_name(name, limit): names = frozenset({name}) view = claude_code_model_id(name, limit, names) @@ -48,7 +58,15 @@ def test_listing_round_trips_entire_source_name(name, limit): def test_collision_matrix_round_trips_without_duplicate_ids(): - universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + universe = ( + "foo", + "foo[1m]", + "claude-router-foo", + _encoded("foo"), + _encoded("foo") + "[1m]", + "claude-opus-5", + "claude-opus-5[1m]", + ) for pair in combinations(universe, 2): for visible in (pair, pair[:1], pair[1:]): names = frozenset(pair) @@ -57,18 +75,31 @@ def test_collision_matrix_round_trips_without_duplicate_ids(): assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) -@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +@pytest.mark.parametrize( + "spelling", + [ + "claude-router-foo", + "claude-router-ff", + "claude-router-66 6f6f", + "claude-router-666F6F", + "claude-router-", + _encoded("missing"), + ], +) def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): assert claude_code_group_name(spelling, frozenset({"foo"})) is None -@pytest.mark.parametrize("headers,enabled", [ - ({"user-agent": "claude-code/2.1.267"}, True), - ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), - ({"x-gateway-client": "Claude-Code"}, True), - ({"user-agent": "anthropic-sdk-python/0.40"}, False), - ({}, False), -]) +@pytest.mark.parametrize( + "headers,enabled", + [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), + ], +) def test_only_claude_code_gets_the_view(headers, enabled): rows = (_row("foo"), _row("claude-opus-5")) view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) @@ -77,12 +108,15 @@ def test_only_claude_code_gets_the_view(headers, enabled): @pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): - import litellm - encoded = _encoded("foo") alias = {encoded: "other"} monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) - router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + router = _router( + "foo", + "other", + *((encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), + aliases=alias if layer == "router" else None, + ) maps = (alias,) if layer in ("key", "team") else () names = ClaudeCodeRoutingNames(router, None, maps) assert claude_code_requested_group(encoded, router, None, maps) is None @@ -98,12 +132,113 @@ def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): assert claude_code_requested_group(_marked(source), router, None) == source +def test_team_alias_is_listed_under_its_target_metadata_and_only_when_the_target_is_accessible() -> None: + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("team-public", "model_name_team_1_abc")] + aliases = ( + {"gpt-4.1-mini": "team-public"}, + None, + {"claude-sonnet-4-5": "gpt-4.1-mini", "via-public": "team-public", "not-granted": "gpt-4.1"}, + ) + assert alias_listing_entries(entries, _caller(*aliases)) == ( + *entries, + ("claude-sonnet-4-5", "gpt-4.1-mini"), + ("via-public", "model_name_team_1_abc"), + ) + assert alias_listing_entries(entries, _caller(None, {})) == tuple(entries) + + +def test_alias_target_resolves_the_requested_alias_across_key_and_team_maps() -> None: + maps = _caller({"o": "gpt-4.1"}, {"claude-sonnet-4-5": "gpt-4.1-mini"}) + assert alias_target("claude-sonnet-4-5", maps) == "gpt-4.1-mini" + assert alias_target("gpt-4.1-mini", _caller(None, {"claude-sonnet-4-5": "gpt-4.1-mini"})) is None + + +def test_alias_colliding_with_a_listed_id_keeps_the_listed_model_at_list_and_retrieval() -> None: + entries = [("fast", "fast"), ("gpt-4.1-mini", "gpt-4.1-mini")] + maps = _caller({"fast": "gpt-4.1-mini"}) + listed = frozenset(response_id for response_id, _ in entries) + + assert alias_listing_entries(entries, maps) == tuple(entries) + assert alias_target("fast", maps, listed) is None + assert alias_target("fast", maps) == "gpt-4.1-mini" + + +def test_alias_maps_apply_in_the_order_chat_completions_applies_them() -> None: + team_then_key = _caller({"fast": "gpt-4.1-mini", "hop": "mid"}, {"fast": "gpt-4.1", "mid": "gpt-4.1"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("gpt-4.1", "gpt-4.1")] + + assert alias_target("fast", team_then_key) == "gpt-4.1-mini" + assert alias_target("hop", team_then_key) == "gpt-4.1" + assert alias_listing_entries(entries, team_then_key) == ( + *entries, + ("fast", "gpt-4.1-mini"), + ("hop", "gpt-4.1"), + ("mid", "gpt-4.1"), + ) + + +def test_one_bad_alias_entry_hides_only_itself() -> None: + aliases = {"fast": "gpt-4.1-mini", "broken": 5, 7: "gpt-4.1-mini"} + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + + assert alias_listing_entries(entries, _caller(aliases)) == (*entries, ("fast", "gpt-4.1-mini")) + assert alias_target("fast", _caller(aliases)) == "gpt-4.1-mini" + + +def test_chained_key_alias_is_listed_only_when_its_final_target_is_listable() -> None: + key_aliases = {"a": "b", "b": "hidden"} + entries = [("b", "b")] + + assert alias_listing_entries(entries, caller_alias_maps(key_aliases, None, "team-a", None)) == (*entries,) + assert alias_target("a", caller_alias_maps(key_aliases, None, "team-a", None)) == "hidden" + + +def test_team_aliases_only_apply_when_listing_the_team_the_key_authenticated_as( + monkeypatch: pytest.MonkeyPatch, +) -> None: + key_aliases, team_aliases, global_aliases = {"k": "gpt-4.1"}, {"t": "gpt-4.1-mini"}, {"g": "gpt-4.1"} + monkeypatch.setattr(litellm, "model_alias_map", global_aliases) + own_team = CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, global_aliases, key_aliases)) + assert caller_alias_maps(key_aliases, team_aliases, "team-a", None) == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-a") == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-b") == CallerAliases( + (key_aliases,), (key_aliases, global_aliases, key_aliases) + ) + + +def test_global_alias_rewrites_between_the_two_key_passes_like_chat_completions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"b": "d"}) + key_aliases = {"a": "b", "b": "c"} + entries = [("c", "c"), ("d", "d")] + maps = caller_alias_maps(key_aliases, None, "team-a", None) + + assert alias_target("a", maps) == "d" + assert alias_listing_entries(entries, maps) == (*entries, ("a", "d"), ("b", "c")) + + +def test_global_aliases_rewrite_but_are_not_listed_as_caller_rows(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"g": "gpt-4.1-mini"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + maps = caller_alias_maps({"k": "g"}, None, "team-a", None) + + assert alias_listing_entries(entries, maps) == (*entries, ("k", "gpt-4.1-mini")) + assert alias_target("g", maps) == "gpt-4.1-mini" + + def test_team_public_name_uses_the_same_scope_at_list_and_request(): - router = Router(model_list=[{ - "model_name": "model_name_team-a_id", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, - "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, - }]) - shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + router = Router( + model_list=[ + { + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + } + ] + ) + shown = claude_code_view_ids( + (_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a") + )["shared"] assert claude_code_requested_group(shown, router, "team-a") == "shared" assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/db/test_db_lookup_gate.py b/tests/test_litellm/proxy/db/test_db_lookup_gate.py new file mode 100644 index 00000000000..68903170840 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_db_lookup_gate.py @@ -0,0 +1,111 @@ +import asyncio +import time +from typing import Final + +import pytest + +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded, DBLookupStallTracker, bounded_db_lookup + + +async def _never_answers() -> None: + await asyncio.Event().wait() + + +class _FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_a_stalled_lookup_at_the_deadline_and_records_the_hit(): + tracker: Final = DBLookupStallTracker() + started: Final = time.monotonic() + + with pytest.raises(DBLookupDeadlineExceeded) as exc_info: + await bounded_db_lookup(_never_answers(), name="team", deadline_seconds=0.05, tracker=tracker) + + assert time.monotonic() - started < 2 + assert exc_info.value.lookup == "team" + assert exc_info.value.deadline_seconds == 0.05 + assert str(exc_info.value) == "team lookup did not answer within 0.05s" + assert isinstance(exc_info.value, asyncio.TimeoutError) + assert tracker.stalled_within(30) is True + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_returns_a_prompt_answer_without_recording_a_stall(): + tracker: Final = DBLookupStallTracker() + + async def answers() -> str: + return "row" + + assert await bounded_db_lookup(answers(), name="key", deadline_seconds=0.05, tracker=tracker) == "row" + assert tracker.stalled_within(30) is False + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_a_whole_stalled_burst_within_one_deadline(): + tracker: Final = DBLookupStallTracker() + burst: Final = 200 + started: Final = time.monotonic() + + results: Final = await asyncio.gather( + *( + bounded_db_lookup(_never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) + for i in range(burst) + ), + return_exceptions=True, + ) + + assert time.monotonic() - started < 2 + assert len(results) == burst + assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) + assert tracker.stalled_within(30) is True + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_at_the_deadline_even_when_the_lookup_absorbs_the_cancel(): + tracker: Final = DBLookupStallTracker() + absorbed: Final = asyncio.Event() + let_go: Final = asyncio.Event() + + async def absorbs_the_cancel() -> str: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + absorbed.set() + await let_go.wait() + return "late row" + + started: Final = time.monotonic() + with pytest.raises(DBLookupDeadlineExceeded): + await asyncio.wait_for( + bounded_db_lookup(absorbs_the_cancel(), name="key", deadline_seconds=0.05, tracker=tracker), + timeout=2, + ) + + assert time.monotonic() - started < 1 + assert tracker.stalled_within(30) is True + await asyncio.wait_for(absorbed.wait(), timeout=1) + let_go.set() + await asyncio.sleep(0) + + +def test_stall_tracker_reports_a_stall_only_inside_the_window(): + clock: Final = _FakeClock() + tracker: Final = DBLookupStallTracker(clock=clock) + + assert tracker.stalled_within(30) is False + tracker.record_hit() + assert tracker.stalled_within(30) is True + assert tracker.stalled_within(0) is False + clock.now += 29.9 + assert tracker.stalled_within(30) is True + clock.now += 0.2 + assert tracker.stalled_within(30) is False + tracker.record_hit() + tracker.clear() + assert tracker.stalled_within(30) is False diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 613ca847115..09f4d294ad0 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -774,3 +774,18 @@ def test_connection_error_answers_when_prisma_is_mocked_after_import(): with patch.dict(sys.modules, {"prisma": MagicMock()}): assert PrismaDBExceptionHandler.is_database_connection_error(Exception("x")) is False assert PrismaDBExceptionHandler.is_database_connection_error(httpx.ConnectError("refused")) is True + + +def test_db_lookup_deadline_is_a_connection_and_unavailability_error_but_never_a_transport_error(): + """A lookup that hit its deadline fails the request as a 503 and counts as a + DB outage for ``allow_requests_on_db_unavailable``, but it must not be read + as a broken transport: that would send every parked request into + ``attempt_db_reconnect`` and turn a slow database into a reconnect storm.""" + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + deadline: Final = DBLookupDeadlineExceeded("key", 10.0) + + assert PrismaDBExceptionHandler.is_database_connection_error(deadline) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(deadline) is True + assert PrismaDBExceptionHandler.is_database_transport_error(deadline) is False + assert "temporarily unreachable" in PrismaDBExceptionHandler.database_unavailable_message(deadline) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index b524902329f..b360ae05def 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -12,12 +12,14 @@ from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final +from unittest.mock import AsyncMock import pytest from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY +from litellm.proxy.db.db_lookup_gate import LoopBoundSemaphore, db_lookup_stall_tracker from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) @@ -447,6 +449,31 @@ async def test_from_db_returns_none_for_a_missing_project_row(): assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None +@pytest.mark.asyncio +async def test_from_db_deadline_covers_the_wait_for_a_gate_slot(monkeypatch: pytest.MonkeyPatch) -> None: + """A saturated gate must fail the lookup at the deadline instead of parking + the request on a gate slot outside the bounded window.""" + gate: Final = LoopBoundSemaphore(1) + monkeypatch.setattr("litellm.proxy.db.spend_counter_reseed.db_lookup_gate", gate) + monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05) + find_unique: Final = AsyncMock() + prisma: Final = SimpleNamespace( + db=SimpleNamespace(litellm_verificationtoken=SimpleNamespace(find_unique=find_unique)) + ) + db_lookup_stall_tracker.clear() + try: + async with gate.current(): + result: Final = await asyncio.wait_for( + SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:key:abc"), + timeout=1.0, + ) + assert result is None + assert db_lookup_stall_tracker.stalled_within(60.0) + find_unique.assert_not_called() + finally: + db_lookup_stall_tracker.clear() + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth 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 761cd0685f2..ee4c468a460 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2703,6 +2703,134 @@ async def test_health_readiness_details_returns_200_when_db_down_and_allow_reque assert result["db"] == "disconnected" +@pytest.fixture +def _clear_db_lookup_stall() -> Iterator[None]: + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + db_lookup_stall_tracker.clear() + yield + db_lookup_stall_tracker.clear() + + +def _connected_prisma() -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(return_value=True) + return mock_prisma + + +def _forget_db_health_cache() -> None: + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_stalled_after_a_db_lookup_deadline_hit(_clear_db_lookup_stall): + """The incident's readiness stayed green while every request sat parked on the + database: the probe's own ping is a fresh connection that answers fine. A lookup + that hit its deadline inside the stall window must take the pod out of rotation.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ): + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result == {"status": "healthy", "db": "stalled"} + + +@pytest.mark.asyncio +async def test_health_readiness_details_returns_503_stalled_after_a_db_lookup_deadline_hit(_clear_db_lookup_stall): + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import _get_health_readiness_details + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ): + result = await _get_health_readiness_details(response=response) + + assert response.status_code == 503 + assert result["db"] == "stalled" + + +@pytest.mark.asyncio +async def test_health_readiness_stays_200_with_stalled_body_when_requests_are_allowed_on_db_unavailable( + _clear_db_lookup_stall, +): + """The fail-open deployment keeps serving through a stalled database, so the pod + must stay in rotation and report the stall through the body, exactly as it does + for a disconnected one.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ), + patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ), + ): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "stalled"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hit_recorded", [False, True]) +async def test_health_readiness_reports_connected_without_a_stall_inside_the_window( + _clear_db_lookup_stall, hit_recorded: bool +): + """No deadline hit, or a window of 0 (the opt-out), keeps the ordinary connected + answer, so a healthy pod never leaves rotation over the stall check.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + if hit_recorded: + db_lookup_stall_tracker.record_hit() + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ), + patch( # test-quality-ok: lowers the module-level stall window to its opt-out value for the recorded-hit case + "litellm.proxy.health_endpoints._health_endpoints.PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", + 0.0 if hit_recorded else 30.0, + ), + ): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "connected"} + + @pytest.mark.asyncio async def test_db_health_readiness_check_bounds_hung_health_check(): """ 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 0e9c336a9eb..b5e594db701 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 @@ -5,12 +5,14 @@ from datetime import datetime from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( @@ -1679,6 +1681,92 @@ async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): assert metadata["user_api_key_team_alias"] == "my-team-alias" +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_skips_the_key_lookup_when_the_failure_is_a_db_stall(): + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="hashed_key") + request_data = { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + with ( + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key_object, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team_object, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=DBLookupDeadlineExceeded("key", 10.0), + user_api_key_dict=user_api_key_dict, + ) + + mock_get_key_object.assert_not_called() + mock_get_team_object.assert_not_called() + mock_update_database.assert_called_once() + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["status"] == "failure" + assert metadata["user_api_key"] == "hashed_key" + assert metadata["user_api_key_alias"] is None + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_still_enriches_metadata_for_a_non_stall_failure(): + """Only a DBLookupDeadlineExceeded skips the key lookup; a transport error + from the provider call must still resolve the key's alias for the failure row.""" + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="hashed_key") + request_data = { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "my-key-alias" + mock_key_obj.user_id = "my-user-id" + mock_key_obj.team_id = "my-team-id" + mock_key_obj.org_id = None + mock_key_obj.project_id = None + + with ( + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ) as mock_get_key_object, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ), + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=httpx.ConnectError("boom"), + user_api_key_dict=user_api_key_dict, + ) + + mock_get_key_object.assert_called_once() + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_alias"] == "my-key-alias" + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_enriches_missing_team_alias(): """ @@ -2035,9 +2123,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + patch( + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), # test-quality-ok: same function-body import, no injection seam + patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() diff --git a/tests/test_litellm/proxy/test_model_list_aliases.py b/tests/test_litellm/proxy/test_model_list_aliases.py new file mode 100644 index 00000000000..25a941bf2ae --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_aliases.py @@ -0,0 +1,151 @@ +""" +Tests for key and team `model_aliases` on the model listing endpoints: GET /v1/models +(`model_list`, OpenAI and Anthropic shapes) and GET /v1/models/{id} (`model_info`). +An alias the caller can complete on is listed next to its target and resolves by name. +""" + +import pytest +from starlette.requests import Request + +from litellm import Router +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _deployment(model_name: str, model: str = "openai/gpt-4.1-mini", **model_info: str | bool) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": f"{model_name}-id", **model_info}, + } + + +@pytest.fixture +def router(monkeypatch: pytest.MonkeyPatch) -> Router: + router = Router( + model_list=[ + _deployment("gpt-4.1-mini"), + _deployment("gpt-4.1", model="openai/gpt-4.1"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-chat"), + _deployment("hidden", model="anthropic/claude-sonnet-4-5", discoverable=False), + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _team_member( + team_id: str = "team1", models: list[str] | None = None, **aliases: dict[str, str] | None +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=team_id, + team_models=["gpt-4.1-mini", "model_name_team1_abc"], + models=models or ["gpt-4.1-mini", "model_name_team1_abc"], + **aliases, + ) + + +def _anthropic_request(*extra_headers: tuple[bytes, bytes]) -> Request: + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/v1/models", + "query_string": b"", + "headers": [(b"anthropic-version", b"2023-06-01"), *extra_headers], + } + ) + + +def _claude_code_request() -> Request: + return _anthropic_request((b"user-agent", b"claude-cli/2.1.267 (external, cli)")) + + +async def _v1_models(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> list[str]: + response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, request=request) + return [m["id"] for m in response["data"]] + + +@pytest.mark.asyncio +async def test_v1_models_lists_team_alias_next_to_its_target_in_both_shapes(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + assert await _v1_models(caller, request=_anthropic_request()) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + + +@pytest.mark.asyncio +async def test_claude_code_picker_lists_the_alias_under_its_own_name(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + picker_ids = await _v1_models(caller, request=_claude_code_request()) + assert any(picker_id.startswith("claude-sonnet-4-5") for picker_id in picker_ids), picker_ids + + +@pytest.mark.asyncio +async def test_v1_models_lists_key_alias_and_hides_alias_to_a_model_the_caller_cannot_list(router: Router) -> None: + caller = _team_member(aliases={"mini": "gpt-4.1-mini", "big": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "mini"] + + +@pytest.mark.asyncio +async def test_v1_models_resolves_a_team_alias_through_the_key_alias_like_chat_completions_does(router: Router) -> None: + caller = _team_member(team_model_aliases={"fast": "mid"}, aliases={"fast": "gpt-4.1", "mid": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast", "mid"] + response = await proxy_server.model_info(model_id="fast", user_api_key_dict=caller) + assert response["id"] == "fast" + + +@pytest.mark.asyncio +async def test_v1_models_skips_only_the_malformed_alias_entries(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": 5, "fast": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast"] + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_a_team_alias_to_its_target_metadata(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + response = await proxy_server.model_info(model_id="claude-sonnet-4-5", user_api_key_dict=caller) + assert response["id"] == "claude-sonnet-4-5" + assert response["owned_by"] == "openai" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_retrieves_the_listed_model_when_an_alias_collides_with_its_id(router: Router) -> None: + caller = _team_member(aliases={"team-chat": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat"] + response = await proxy_server.model_info(model_id="team-chat", user_api_key_dict=caller) + assert response["id"] == "team-chat" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_an_alias_named_like_an_undiscoverable_model_to_the_alias_target( + router: Router, +) -> None: + caller = _team_member(aliases={"hidden": "gpt-4.1-mini"}, models=["gpt-4.1-mini", "model_name_team1_abc", "hidden"]) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "hidden"] + target = await proxy_server.model_info(model_id="gpt-4.1-mini", user_api_key_dict=caller) + response = await proxy_server.model_info(model_id="hidden", user_api_key_dict=caller) + assert response == {**target, "id": "hidden"} + + +@pytest.mark.asyncio +async def test_v1_models_by_id_keeps_the_alias_as_id_when_it_targets_a_team_scoped_model(router: Router) -> None: + caller = _team_member(team_model_aliases={"chat": "team-chat"}) + + assert "chat" in await _v1_models(caller) + response = await proxy_server.model_info(model_id="chat", user_api_key_dict=caller) + assert response["id"] == "chat"