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

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

+
+              /plugin marketplace update litellm
+            
+
+

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

One-time marketplace setup

-

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

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

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

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

Date: Thu, 16 Jul 2026 14:12:29 +0300 Subject: [PATCH 019/319] fix(claude-code): fix prettier formatting and remove unused import --- .../src/components/claude_code_plugins/helpers.test.ts | 2 +- .../src/components/claude_code_plugins/helpers.ts | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index a04134a63d8..1d07ca09df5 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -19,7 +19,7 @@ import { isValidSubPath, buildMarketplaceSettingsSnippet, } from "./helpers"; -import { MarketplacePluginEntry, PluginSource } from "./types"; +import { MarketplacePluginEntry } from "./types"; describe("buildMarketplaceSettingsSnippet", () => { it("nests the url under a source object so Claude Code accepts the marketplace", () => { diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 85b4ef357f7..5b3d4f20c1e 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -202,8 +202,7 @@ export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => * Generate install command for Claude Code CLI. * Installs the named plugin from the "litellm" marketplace registered in settings.json. */ -export const formatInstallCommand = (plugin: { name: string }): string => - `/plugin install ${plugin.name}@litellm`; +export const formatInstallCommand = (plugin: { name: string }): string => `/plugin install ${plugin.name}@litellm`; /** * Extract unique categories from plugins list From 141281b1d658e192a2b188ddf03d19f6eab58814 Mon Sep 17 00:00:00 2001 From: ozolam Date: Thu, 16 Jul 2026 14:15:43 +0300 Subject: [PATCH 020/319] ci: retrigger CI From 346c065fe0c7516b563996d2cee7061ce276fc3a Mon Sep 17 00:00:00 2001 From: ozolam Date: Tue, 11 Aug 2026 21:34:41 +0300 Subject: [PATCH 021/319] fix(claude-code): restrict marketplace catalog mutations to proxy admins --- .../claude_code_marketplace.py | 24 +++++++ .../test_claude_code_marketplace.py | 71 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 46ee9b0911d..4730372511a 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -28,6 +28,7 @@ from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, @@ -221,6 +222,18 @@ def _name_conflict_error(name: str) -> HTTPException: ) +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + """Catalog mutations are restricted to proxy admins: marketplace.json is served + unauthenticated and any registered/updated entry is immediately installable by + every user, so a non-admin key must never be able to add or overwrite one. + """ + if not is_proxy_admin(user_api_key_dict): + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins may modify the Claude Code plugin marketplace."}, + ) + + @router.post( "/claude-code/plugins", tags=["Claude Code Marketplace"], @@ -271,6 +284,8 @@ async def register_plugin( from prisma.errors import UniqueViolationError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() if not re.match(r"^[a-z0-9-]+$", request.name): @@ -468,6 +483,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Update an existing plugin in the LiteLLM marketplace. @@ -509,6 +525,8 @@ async def update_plugin( from prisma.errors import PrismaError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() _validate_plugin_source(request.source) @@ -570,6 +588,8 @@ async def enable_plugin( - plugin_name: The name of the plugin to enable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -615,6 +635,8 @@ async def disable_plugin( - plugin_name: The name of the plugin to disable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -660,6 +682,8 @@ async def delete_plugin( - plugin_name: The name of the plugin to delete """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 18e0f2cb559..d1ccd86044c 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -18,6 +18,9 @@ from litellm.types.proxy.claude_code_endpoints import ( UpdatePluginRequest, ) from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + delete_plugin, + disable_plugin, + enable_plugin, get_marketplace, register_plugin, update_plugin, @@ -72,6 +75,12 @@ _USER = UserAPIKeyAuth( user_id="test-user", ) +_NON_ADMIN_USER = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-5678", + user_id="regular-user", +) + _GIT_SUBDIR_SOURCE = { "source": "git-subdir", "url": "https://github.com/org/monorepo.git", @@ -151,6 +160,7 @@ async def test_update_plugin_replaces_existing_source(): response = await update_plugin( plugin_name=name, request=UpdatePluginRequest(source=new_source, version="2.0.0", description="updated"), + user_api_key_dict=_USER, ) assert response.status == "success" @@ -170,6 +180,7 @@ async def test_update_plugin_not_found(): await update_plugin( plugin_name="does-not-exist", request=UpdatePluginRequest(source=_GIT_SUBDIR_SOURCE), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 404 @@ -213,6 +224,7 @@ async def test_update_plugin_db_error_maps_to_structured_500(): await update_plugin( plugin_name=name, request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 500 @@ -341,3 +353,62 @@ async def test_register_plugin_unknown_source_type(): assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_register_plugin_rejects_non_admin(): + """A non-admin key cannot add an entry to the marketplace catalog.""" + request = RegisterPluginRequest(name="attacker-plugin", source=_GIT_SUBDIR_SOURCE) + + with pytest.raises(HTTPException) as exc_info: + await register_plugin(request=request, user_api_key_dict=_NON_ADMIN_USER) + + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert await table.find_unique(where={"name": "attacker-plugin"}) is None + + +@pytest.mark.asyncio +async def test_update_plugin_rejects_non_admin_overwrite(): + """A non-admin key cannot overwrite an existing plugin's source.""" + name = "trusted-plugin" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + malicious_source = {"source": "github", "repo": "attacker/malicious-repo"} + with pytest.raises(HTTPException) as exc_info: + await update_plugin( + plugin_name=name, + request=UpdatePluginRequest(source=malicious_source), + user_api_key_dict=_NON_ADMIN_USER, + ) + + assert exc_info.value.status_code == 403 + + stored = await _read_stored_manifest(name) + assert stored["source"] == _GIT_SUBDIR_SOURCE + + +@pytest.mark.asyncio +async def test_enable_disable_delete_plugin_reject_non_admin(): + """Non-admin keys cannot enable, disable, or delete catalog entries.""" + name = "trusted-plugin-2" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + for coro in ( + enable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + disable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + delete_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + ): + with pytest.raises(HTTPException) as exc_info: + await coro + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert (await table.find_unique(where={"name": name})).enabled is True From 168ee2d9fdbd63cb3f2f4fd6de665595cbcbd8b8 Mon Sep 17 00:00:00 2001 From: ozolam Date: Tue, 11 Aug 2026 22:38:25 +0300 Subject: [PATCH 022/319] fix(claude-code): use Annotated dependency style to stay within B008 lint budget --- .../claude_code_marketplace.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 4730372511a..09ba2c93ea0 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -6,21 +6,21 @@ Plugins are stored as metadata + git source references in LiteLLM database. Actual plugin files are hosted on GitHub/GitLab/Bitbucket. Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery -/claude-code/plugins - POST - Register a new plugin (create-only) -/claude-code/plugins - GET - List plugins (admin) -/claude-code/plugins/{name} - GET - Get plugin details -/claude-code/plugins/{name} - PUT - Update an existing plugin -/claude-code/plugins/{name}/enable - POST - Enable a plugin -/claude-code/plugins/{name}/disable - POST - Disable a plugin -/claude-code/plugins/{name} - DELETE - Delete a plugin +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) +/claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) +/claude-code/plugins - GET - List plugins (any authenticated key) +/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key) +/claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only) +/claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only) +/claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only) +/claude-code/plugins/{name} - DELETE - Delete a plugin (proxy admin only) """ import json import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Final, Protocol, TypedDict +from typing import Annotated, Final, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -255,6 +255,8 @@ async def register_plugin( the same name already exists it returns 409 Conflict; use PUT /claude-code/plugins/{plugin_name} to update an existing plugin. + Requires a proxy admin API key. + Parameters: - name: Plugin name (kebab-case) - source: Git source reference (github, url, or git-subdir format) @@ -483,7 +485,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an existing plugin in the LiteLLM marketplace. @@ -497,6 +499,8 @@ async def update_plugin( Returns 404 if no plugin with the given name exists; use POST /claude-code/plugins to create a new plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: Name of the plugin to update (path parameter) - source: Git source reference (github, url, or git-subdir format) @@ -584,6 +588,8 @@ async def enable_plugin( """ Enable a disabled plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to enable """ @@ -631,6 +637,8 @@ async def disable_plugin( """ Disable a plugin without deleting it. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to disable """ @@ -678,6 +686,8 @@ async def delete_plugin( """ Delete a plugin from the marketplace. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to delete """ From 3db1759d04ba8d888a08dbe8308cc75a348141d8 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 19:45:32 +0000 Subject: [PATCH 023/319] fix(bedrock): stop emitting an empty assistant delta after the finish_reason chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/chat/invoke_handler.py | 8 +++- .../llms/bedrock/chat/test_invoke_handler.py | 46 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 8d2b3dae71b..595884ae630 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -561,6 +561,10 @@ class AWSEventStreamDecoder: elif "usage" in chunk_data: usage = converse_config._transform_usage(chunk_data.get("usage", {})) + carries_message_content: Final = any( + key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason") + ) + model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: trace: Final = chunk_data.get("trace") @@ -571,8 +575,8 @@ class AWSEventStreamDecoder: finish_reason=finish_reason, index=0, # Always 0 - Bedrock never returns multiple choices delta=Delta( - content=text, - role="assistant", + content=text if carries_message_content else None, + role="assistant" if carries_message_content else None, tool_calls=[tool_use] if tool_use else None, provider_specific_fields=(provider_specific_fields if provider_specific_fields else None), thinking_blocks=thinking_blocks, diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index ee50b9db015..9783976db4a 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -1,3 +1,4 @@ +import datetime import os import sys from unittest.mock import AsyncMock, MagicMock @@ -8,6 +9,8 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, @@ -293,3 +296,46 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) + +@pytest.mark.asyncio +async def test_converse_stream_ends_on_finish_reason_chunk(): + """The usage-only metadata event Bedrock sends after messageStop must not reach the caller as an extra + assistant delta following the finish_reason chunk.""" + model = "anthropic.claude-sonnet-4-6" + events = ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "delta": {"text": "Hello"}}, + {"contentBlockIndex": 0, "delta": {"text": " world"}}, + {"contentBlockIndex": 0}, + {"stopReason": "end_turn"}, + {"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, "metrics": {"latencyMs": 100}}, + ) + + async def bedrock_stream(): + decoder = AWSEventStreamDecoder(model=model) + for event in events: + yield decoder._chunk_parser(chunk_data=event) + + wrapper = CustomStreamWrapper( + completion_stream=bedrock_stream(), + model=model, + custom_llm_provider="bedrock", + logging_obj=LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="1234", + function_id="1234", + ), + ) + + chunks = [chunk async for chunk in wrapper] + + assert [choice.finish_reason for chunk in chunks for choice in chunk.choices].count("stop") == 1 + assert chunks[-1].choices[0].finish_reason == "stop", ( + f"stream must end on the finish_reason chunk, got trailing {chunks[-1].model_dump(exclude_none=True)}" + ) + assert any(getattr(chunk, "usage", None) is not None for chunk in wrapper.chunks) + From 0e3f52a4c0b071e8f297decce690a6ba2b6615ff Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 13 Aug 2026 19:18:22 -0500 Subject: [PATCH 024/319] feat(model_prices): add gemini-3.1-flash-lite-image Register Nano Banana 2 Lite on the unprefixed, gemini/, and vertex_ai/ keys so completion_cost and pass-through spend tracking no longer treat the model as unmapped --- ...odel_prices_and_context_window_backup.json | 95 +++++++ model_prices_and_context_window.json | 95 +++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...ini_3_1_flash_lite_image_model_metadata.py | 242 ++++++++++++++++++ tests/test_litellm/test_utils.py | 2 + 5 files changed, 435 insertions(+) create mode 100644 tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1eb72c887b5..ca67d5d6844 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1eb72c887b5..ca67d5d6844 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3aa41e18f1e..36fc98a1f09 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1546,6 +1546,7 @@ def test_service_tier_fallback_pricing(): [ "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", ], ) def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py new file mode 100644 index 00000000000..aa6f03a47ff --- /dev/null +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py @@ -0,0 +1,242 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm import completion_cost +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +UNPREFIXED = "gemini-3.1-flash-lite-image" +GEMINI = "gemini/gemini-3.1-flash-lite-image" +VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" +ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) + +INPUT_COST = 2.5e-07 +INPUT_COST_BATCHES = 1.25e-07 +OUTPUT_TEXT_COST = 1.5e-06 +OUTPUT_TEXT_COST_BATCHES = 7.5e-07 +OUTPUT_IMAGE_TOKEN_COST = 3e-05 +OUTPUT_COST_PER_1K_IMAGE = 0.0336 +INPUT_COST_PER_IMAGE = 0.00028 +CACHE_READ_COST = 2.5e-08 +MAX_INPUT_TOKENS = 65536 +MAX_OUTPUT_TOKENS = 4096 +TOKENS_PER_1K_IMAGE = 1120 + + +def _load(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_gemini_3_1_flash_lite_image_is_registered(model: str): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["mode"] == "image_generation" + assert info["input_cost_per_token"] == INPUT_COST + assert info["input_cost_per_token_batches"] == INPUT_COST_BATCHES + assert info["output_cost_per_token"] == OUTPUT_TEXT_COST + assert info["output_cost_per_token_batches"] == OUTPUT_TEXT_COST_BATCHES + assert info["output_cost_per_image"] == OUTPUT_COST_PER_1K_IMAGE + assert info["output_cost_per_image_token"] == OUTPUT_IMAGE_TOKEN_COST + assert info["max_input_tokens"] == MAX_INPUT_TOKENS + assert info["max_output_tokens"] == MAX_OUTPUT_TOKENS + assert info["max_tokens"] == MAX_OUTPUT_TOKENS + assert info["supports_reasoning"] is False + assert info["supports_response_schema"] is False + assert info["supports_vision"] is True + for field in ("supports_web_search", "search_context_cost_per_query", "web_search_billing_unit"): + assert field not in info + + +def test_gemini_3_1_flash_lite_image_provider_specific_fields(): + cost_map = _load(MAIN_PATH) + + unprefixed = cost_map[UNPREFIXED] + assert unprefixed["litellm_provider"] == "vertex_ai-language-models" + assert unprefixed["cache_read_input_token_cost"] == CACHE_READ_COST + assert unprefixed["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert unprefixed["supports_function_calling"] is False + assert unprefixed["supports_prompt_caching"] is True + assert unprefixed["supports_pdf_input"] is True + assert unprefixed["supports_video_input"] is True + assert unprefixed["supported_modalities"] == ["text", "image", "video"] + + gemini = cost_map[GEMINI] + assert gemini["litellm_provider"] == "gemini" + assert gemini["supports_function_calling"] is True + assert gemini["supports_prompt_caching"] is False + assert "cache_read_input_token_cost" not in gemini + assert gemini["supported_modalities"] == ["text", "image"] + assert gemini["supported_output_modalities"] == ["text", "image"] + assert gemini["rpm"] == 1000 + assert gemini["tpm"] == 4000000 + assert gemini["input_cost_per_image"] == INPUT_COST_PER_IMAGE + + vertex = cost_map[VERTEX] + assert vertex["litellm_provider"] == "vertex_ai-language-models" + assert vertex["cache_read_input_token_cost"] == CACHE_READ_COST + assert vertex["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert vertex["supports_function_calling"] is False + assert vertex["supports_prompt_caching"] is True + + +def test_one_k_image_price_matches_official_token_math(): + assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == OUTPUT_COST_PER_1K_IMAGE + assert TOKENS_PER_1K_IMAGE * INPUT_COST == INPUT_COST_PER_IMAGE + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_gemini_prefix_routes_to_gemini(): + routed_model, provider, _, _ = get_llm_provider(model=GEMINI) + assert routed_model == UNPREFIXED + assert provider == "gemini" + + +def test_vertex_prefix_routes_to_vertex(): + routed_model, provider, _, _ = get_llm_provider(model=VERTEX) + assert routed_model == UNPREFIXED + assert provider == "vertex_ai" + + +def test_text_token_cost(local_model_cost_map): + prompt_cost, text_completion_cost = cost_per_token(model=GEMINI, prompt_tokens=1000, completion_tokens=500) + assert prompt_cost == pytest.approx(1000 * INPUT_COST) + assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) + + +def test_completion_cost_bills_one_k_image(local_model_cost_map): + response = ModelResponse() + response.model = UNPREFIXED + response.usage = Usage( + prompt_tokens=7, + completion_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=7 + TOKENS_PER_1K_IMAGE, + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0), + ) + billed = completion_cost( + completion_response=response, + model=UNPREFIXED, + custom_llm_provider="vertex_ai", + ) + expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST + assert billed == pytest.approx(expected) + + +def test_image_tokens_are_not_billed_as_text(local_model_cost_map): + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + _, image_completion_cost = generic_cost_per_token( + model=UNPREFIXED, + usage=usage, + custom_llm_provider="vertex_ai", + ) + + expected_completion_cost = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST + bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST + assert image_completion_cost > bugged_text_only_cost * 2 + assert image_completion_cost == pytest.approx(expected_completion_cost) + + +def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = gemini_image_generation_cost_calculator(model=GEMINI, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + assert cost != OUTPUT_COST_PER_1K_IMAGE + + +def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + + +def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..ada053ee38c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4312,11 +4312,13 @@ class TestVertexEmbeddingEncodingFormat: "vertex_ai/gemini-3-pro-image-preview", "vertex_ai/gemini-3.1-flash-image", "vertex_ai/gemini-3.1-flash-image-preview", + "vertex_ai/gemini-3.1-flash-lite-image", "gemini/gemini-2.5-flash-image", "gemini/gemini-3-pro-image", "gemini/gemini-3-pro-image-preview", "gemini/gemini-3.1-flash-image", "gemini/gemini-3.1-flash-image-preview", + "gemini/gemini-3.1-flash-lite-image", ], ) def test_gemini_image_models_do_not_support_reasoning( From b6ee13803d1c21bc0b07dac6e6b9ad2ede7ad7aa Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Sun, 9 Aug 2026 22:59:25 +0800 Subject: [PATCH 025/319] fix(responses-bridge): preserve reasoning input items as reasoning_content --- .../transformation.py | 164 +++++++++++++++++- .../test_reasoning_input_item_preservation.py | 147 ++++++++++++++++ 2 files changed, 310 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4892e3b348c..5d3ed0477e3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -557,7 +557,108 @@ class LiteLLMCompletionResponsesConfig: continue messages.extend(chat_completion_messages) - return messages + return LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages(messages) + + @staticmethod + def _merge_reasoning_only_assistant_messages( + messages: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ], + ) -> list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage + ]: + """ + Responses API emits prior-turn reasoning as its own ``reasoning`` input + item, which becomes a standalone assistant message with + ``content=None`` + ``reasoning_content``. Chat-completions providers + (e.g. DeepSeek V4, Kimi K2.6) expect the chain-of-thought on the + assistant message that carries the answer or tool calls. This pass + merges standalone reasoning-only assistant messages into the + immediately following assistant message. + + If the reasoning item is not followed by an assistant message (e.g. a + stateless chain replays ``reasoning`` + ``user``), the standalone + reasoning message is preserved so the reasoning is still passed back. + """ + + def _role(msg: Any) -> str: + if isinstance(msg, dict): + return str(msg.get("role") or "") + return str(getattr(msg, "role", "") or "") + + def _reasoning_text(msg: Any) -> str | None: + if isinstance(msg, dict): + value = msg.get("reasoning_content") + else: + value = getattr(msg, "reasoning_content", None) + return value if isinstance(value, str) and value else None + + def _content(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("content") + return getattr(msg, "content", None) + + def _tool_calls(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("tool_calls") + return getattr(msg, "tool_calls", None) + + merged: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ] = [] + pending_reasoning: list[str] = [] + + for msg in messages: + if ( + _role(msg) == "assistant" + and _content(msg) is None + and not _tool_calls(msg) + and _reasoning_text(msg) is not None + ): + pending_reasoning.append(_reasoning_text(msg) or "") + continue + + if pending_reasoning and _role(msg) == "assistant": + combined = "\n".join(pending_reasoning) + existing = _reasoning_text(msg) + if existing: + combined = existing + "\n" + combined + if isinstance(msg, dict): + msg["reasoning_content"] = combined + else: + setattr(msg, "reasoning_content", combined) + pending_reasoning = [] + elif pending_reasoning: + # Not followed by an assistant message — keep the reasoning + # standalone instead of dropping it. + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + pending_reasoning = [] + + merged.append(msg) + + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + + return merged @staticmethod def _merged_trailing_assistant_message( @@ -1026,6 +1127,25 @@ class LiteLLMCompletionResponsesConfig: return LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( function_call=input_item ) + elif input_item.get("type") == "reasoning": + # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. + # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this + # to be replayed as `reasoning_content` on an assistant message, not as + # visible `content` (prompt pollution) and not dropped (DeepSeek V4 + # rejects multi-turn requests with a missing `reasoning_content`). + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + if not reasoning_text: + # No plaintext reasoning is available (e.g. encrypted_content only). + # Chat-completions providers cannot consume opaque encrypted blobs, + # so skip the item instead of polluting the prompt. + return [] + return [ + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=reasoning_text, + ) + ] else: content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content @@ -1041,6 +1161,48 @@ class LiteLLMCompletionResponsesConfig: ) ] + @staticmethod + def _extract_reasoning_text_from_input_item(input_item: Mapping[str, object]) -> str | None: + """ + Extract plaintext reasoning from a ResponseReasoningItemParam. + + Handles: + - content as a string + - content as a list of blocks (output_text / summary_text / text) + - summary as a list of summary_text blocks (fallback) + + Returns None when only opaque forms (e.g. encrypted_content) are present. + """ + content: Final[object] = input_item.get("content") + if isinstance(content, str) and content.strip(): + return content + if isinstance(content, list): + text_parts: list[str] = [] + for block in content: + if not isinstance(block, Mapping): + continue + block_type = block.get("type") + if block_type in ("encrypted_content", "redacted_thinking"): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + + summary: Final[object] = input_item.get("summary") + if isinstance(summary, list): + text_parts = [] + for block in summary: + if not isinstance(block, Mapping): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + return None + @staticmethod def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py new file mode 100644 index 00000000000..5fcd4df3ff8 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -0,0 +1,147 @@ +""" +Unit tests for preserving prior-turn ``reasoning`` input items when the +Responses API is bridged to chat completions. + +Without this handling, a ``ResponseReasoningItemParam`` falls through to the +generic message branch, polluting the prompt as visible assistant ``content`` +or being silently dropped. Chat-completions providers such as DeepSeek V4 and +Kimi K2.6 require the chain-of-thought to be replayed as ``reasoning_content`` +on an assistant message. +""" + +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +def _transform_item(item): + return LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=item + ) + + +def _transform_input(input_items): + return LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=input_items + ) + + +class TestReasoningInputItemHandler: + """Reasoning input items map to assistant ``reasoning_content``.""" + + def test_reasoning_item_with_output_text_content(self): + """Standard Responses-API reasoning item with output_text blocks.""" + item = { + "type": "reasoning", + "id": "rs_abc", + "summary": [], + "content": [{"type": "output_text", "text": "step 1: think about X"}], + } + messages = _transform_item(item) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "step 1: think about X" + + def test_reasoning_item_with_string_content(self): + """Variant: reasoning content as a plain string.""" + item = {"type": "reasoning", "id": "rs_1", "content": "step 1: ..."} + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "step 1: ..." + + def test_reasoning_item_with_summary_only(self): + """SDK form: reasoning carried in summary list, no content.""" + item = { + "type": "reasoning", + "id": "rs_2", + "summary": [{"type": "summary_text", "text": "..."}], + } + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "..." + + def test_reasoning_item_with_encrypted_content_only_dropped(self): + """Opaque encrypted reasoning cannot be forwarded to chat completions.""" + item = {"type": "reasoning", "id": "rs_3", "encrypted_content": "opaque-blob"} + assert _transform_item(item) == [] + + def test_reasoning_item_empty_dropped(self): + """Reasoning item with neither content nor summary drops cleanly.""" + assert _transform_item({"type": "reasoning", "id": "rs_4"}) == [] + + +class TestReasoningInputItemMerging: + """Standalone reasoning messages merge into the following assistant turn.""" + + def test_reasoning_merged_into_following_assistant_message(self): + """Reasoning + assistant answer become one assistant message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret reasoning"}], + }, + {"type": "message", "role": "assistant", "content": "The answer."}, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "secret reasoning" + + def test_reasoning_preserved_when_followed_by_user_message(self): + """Stateless chain: reasoning + user prompt keeps the reasoning turn.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret BLUEBERRY"}], + }, + {"role": "user", "content": "What is the secret word?"}, + ] + ) + assert len(messages) == 2 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "secret BLUEBERRY" + assert messages[1]["role"] == "user" + + def test_reasoning_merged_into_function_call_assistant(self): + """Reasoning + function_call becomes one assistant tool-call message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "I should look this up"}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": '{"cwe": "79"}', + }, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["reasoning_content"] == "I should look this up" + assert len(messages[0]["tool_calls"]) == 1 + + +class TestNonReasoningInputItemUnchanged: + """Non-reasoning items still flow through the existing branches.""" + + def test_user_message_unchanged(self): + item = {"role": "user", "content": "hello"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "user" + + def test_assistant_message_unchanged(self): + item = {"role": "assistant", "content": "hi"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "assistant" + assert out[0]["content"] == "hi" From 3a77556dc14660e88a7d20f54e9762c39f24b749 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:45:25 +0800 Subject: [PATCH 026/319] fix(responses-bridge): preserve reasoning merge order when assistant already has reasoning_content --- .../transformation.py | 2 +- .../test_reasoning_input_item_preservation.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5d3ed0477e3..0604c3636ff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -628,7 +628,7 @@ class LiteLLMCompletionResponsesConfig: combined = "\n".join(pending_reasoning) existing = _reasoning_text(msg) if existing: - combined = existing + "\n" + combined + combined = combined + "\n" + existing if isinstance(msg, dict): msg["reasoning_content"] = combined else: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 5fcd4df3ff8..ecc024b7d04 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -129,6 +129,18 @@ class TestReasoningInputItemMerging: assert messages[0]["reasoning_content"] == "I should look this up" assert len(messages[0]["tool_calls"]) == 1 + def test_reasoning_merged_into_assistant_with_existing_reasoning_content(self): + """Old reasoning precedes existing reasoning on the target assistant turn.""" + messages = LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages( + [ + {"role": "assistant", "content": None, "reasoning_content": "old reasoning"}, + {"role": "assistant", "content": "The answer.", "reasoning_content": "new reasoning"}, + ] + ) + assert len(messages) == 1 + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "old reasoning\nnew reasoning" + class TestNonReasoningInputItemUnchanged: """Non-reasoning items still flow through the existing branches.""" From 5911124f1dbba1e9c58f3b53619c3f875752a20f Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:58:47 +0800 Subject: [PATCH 027/319] fix(responses-bridge): satisfy ruff strict-rule budget in reasoning merge --- .../transformation.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0604c3636ff..2c506d4a4c7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -584,24 +584,24 @@ class LiteLLMCompletionResponsesConfig: reasoning message is preserved so the reasoning is still passed back. """ - def _role(msg: Any) -> str: + def _role(msg: object) -> str: if isinstance(msg, dict): return str(msg.get("role") or "") return str(getattr(msg, "role", "") or "") - def _reasoning_text(msg: Any) -> str | None: + def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): value = msg.get("reasoning_content") else: value = getattr(msg, "reasoning_content", None) return value if isinstance(value, str) and value else None - def _content(msg: Any) -> Any: + def _content(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("content") return getattr(msg, "content", None) - def _tool_calls(msg: Any) -> Any: + def _tool_calls(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("tool_calls") return getattr(msg, "tool_calls", None) @@ -632,31 +632,35 @@ class LiteLLMCompletionResponsesConfig: if isinstance(msg, dict): msg["reasoning_content"] = combined else: - setattr(msg, "reasoning_content", combined) + setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) pending_reasoning = [] merged.append(msg) - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) return merged From 438c1850fe1223feec1e2e6e5b48f0a6c15a1328 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:09:26 +0800 Subject: [PATCH 028/319] fix(responses-bridge): satisfy type-discipline budget in reasoning merge --- .../transformation.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2c506d4a4c7..e3e62ab3c55 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -561,13 +561,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _merge_reasoning_only_assistant_messages( - messages: list[ + messages: list[ # mutable-ok: input sequence AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ], - ) -> list[ + ) -> list[ # mutable-ok: fresh merged list AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: """ @@ -591,9 +591,9 @@ class LiteLLMCompletionResponsesConfig: def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): - value = msg.get("reasoning_content") + value = msg.get("reasoning_content") # rebind-ok: branch lookup else: - value = getattr(msg, "reasoning_content", None) + value = getattr(msg, "reasoning_content", None) # rebind-ok: branch lookup return value if isinstance(value, str) and value else None def _content(msg: object) -> object | None: @@ -606,13 +606,13 @@ class LiteLLMCompletionResponsesConfig: return msg.get("tool_calls") return getattr(msg, "tool_calls", None) - merged: list[ + merged: list[ # mutable-ok: accumulator # rebind-ok: accumulator AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage - ] = [] - pending_reasoning: list[str] = [] + ] = [] # mutable-ok: accumulator + pending_reasoning: list[str] = [] # mutable-ok: accumulator # rebind-ok: accumulator for msg in messages: if ( @@ -633,11 +633,11 @@ class LiteLLMCompletionResponsesConfig: msg["reasoning_content"] = combined else: setattr(msg, "reasoning_content", combined) # noqa: B010 - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - merged.extend( + merged.extend( # mutable-ok: append reasoning messages [ ChatCompletionResponseMessage( role="assistant", @@ -647,11 +647,11 @@ class LiteLLMCompletionResponsesConfig: for text in pending_reasoning ] ) - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator merged.append(msg) - merged.extend( + merged.extend( # mutable-ok: append trailing reasoning [ ChatCompletionResponseMessage( role="assistant", @@ -1137,13 +1137,15 @@ class LiteLLMCompletionResponsesConfig: # to be replayed as `reasoning_content` on an assistant message, not as # visible `content` (prompt pollution) and not dropped (DeepSeek V4 # rejects multi-turn requests with a missing `reasoning_content`). - reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item( # rebind-ok: extraction result + input_item + ) if not reasoning_text: # No plaintext reasoning is available (e.g. encrypted_content only). # Chat-completions providers cannot consume opaque encrypted blobs, # so skip the item instead of polluting the prompt. - return [] - return [ + return [] # mutable-ok: empty drop result + return [ # mutable-ok: single message result ChatCompletionResponseMessage( role="assistant", content=None, @@ -1181,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(content, str) and content.strip(): return content if isinstance(content, list): - text_parts: list[str] = [] + text_parts: list[str] = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in content: if not isinstance(block, Mapping): continue @@ -1196,7 +1198,7 @@ class LiteLLMCompletionResponsesConfig: summary: Final[object] = input_item.get("summary") if isinstance(summary, list): - text_parts = [] + text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in summary: if not isinstance(block, Mapping): continue From de95372dfbd7bbba8c478815340dd49c1b21da11 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:24:24 +0800 Subject: [PATCH 029/319] fix(responses-bridge): type-safe reasoning_content assignment in merge pass --- .../litellm_completion_transformation/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e3e62ab3c55..d7b6b8c7b8f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -630,7 +630,7 @@ class LiteLLMCompletionResponsesConfig: if existing: combined = combined + "\n" + existing if isinstance(msg, dict): - msg["reasoning_content"] = combined + cast(dict[str, Any], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier else: setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] # mutable-ok: reset accumulator From 2d4e6afe1c7d6d3233a18668d53fafe4cffa50b3 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Tue, 18 Aug 2026 21:11:29 +0800 Subject: [PATCH 030/319] fix(guardrails): inspect responses reasoning content and summary text --- litellm/proxy/guardrails/_content_utils.py | 59 +++++++++++++------ .../transformation.py | 6 +- .../proxy/guardrails/test_content_utils.py | 54 +++++++++++++++++ 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6ed6f0013df..ae92adcb1ee 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -8,7 +8,7 @@ skip the other shapes — these helpers normalise that so every hook sees every text fragment. """ -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from typing import Any, Final # Call types whose body carries free-form chat / prompt text that @@ -33,7 +33,9 @@ def is_text_content_call_type(call_type: str) -> bool: return call_type in TEXT_CONTENT_CALL_TYPES -TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "output_text"}) +TEXT_PART_TYPES: Final[frozenset[str]] = frozenset( + {"text", "input_text", "output_text", "summary_text", "reasoning_text"} +) # Responses-API item types whose ``output`` field carries user/tool text # that guardrails should inspect. ``function_call_output`` is the @@ -42,6 +44,16 @@ TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "outpu _OUTPUT_ITEM_TYPES: Final[frozenset[str]] = frozenset({"function_call_output", "custom_tool_call_output"}) +def _part_text(part: Mapping[str, object]) -> str | None: + """Return non-empty plaintext from any content part that carries ``text``.""" + if not isinstance(part, dict): + return None + text = part.get("text") + if isinstance(text, str) and text: + return text + return None + + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or multimodal list). Non-text parts (images, audio, …) are skipped.""" @@ -58,10 +70,9 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: continue if not isinstance(part, dict): continue - if part.get("type") in TEXT_PART_TYPES: - text = part.get("text") - if isinstance(text, str) and text: - yield text + text = _part_text(part) + if text is not None: + yield text def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: @@ -75,8 +86,23 @@ def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: if isinstance(item, str): messages.append({"role": "user", "content": item}) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: + if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) + elif item.get("type") == "reasoning": + if "content" in item: + messages.append( + { # mutable-ok: append reasoning content + "role": item.get("role") or "assistant", + "content": item["content"], + } + ) + if isinstance(item.get("summary"), list): + messages.append( + { # mutable-ok: append reasoning summary + "role": item.get("role") or "assistant", + "content": item["summary"], + } + ) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: @@ -126,12 +152,7 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: if isinstance(part, str) and part: visited += 1 new_parts.append(visit(part)) - elif ( - isinstance(part, dict) - and part.get("type") in TEXT_PART_TYPES - and isinstance(part.get("text"), str) - and part["text"] - ): + elif isinstance(part, dict) and _part_text(part) is not None: visited += 1 new_parts.append({**part, "text": visit(part["text"])}) else: @@ -158,10 +179,14 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: visited += 1 input_value[idx] = visit(item) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: - if isinstance(item.get("text"), str) and item["text"]: - visited += 1 - input_value[idx] = {**item, "text": visit(item["text"])} + if _part_text(item) is not None: + visited += 1 + input_value[idx] = {**item, "text": visit(item["text"])} # mutable-ok: rewrite text part in place + elif item.get("type") == "reasoning": + if "content" in item: + item["content"] = _rewrite_content(item["content"]) + if isinstance(item.get("summary"), list): + item["summary"] = _rewrite_content(item["summary"]) elif "content" in item: item["content"] = _rewrite_content(item["content"]) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d7b6b8c7b8f..c5f40242bfd 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -638,7 +638,7 @@ class LiteLLMCompletionResponsesConfig: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. merged.extend( # mutable-ok: append reasoning messages - [ + [ # mutable-ok: append reasoning messages ChatCompletionResponseMessage( role="assistant", content=None, @@ -652,7 +652,7 @@ class LiteLLMCompletionResponsesConfig: merged.append(msg) merged.extend( # mutable-ok: append trailing reasoning - [ + [ # mutable-ok: append trailing reasoning ChatCompletionResponseMessage( role="assistant", content=None, @@ -1196,6 +1196,8 @@ class LiteLLMCompletionResponsesConfig: if text_parts: return "\n".join(text_parts) + # Guardrail traversal in litellm/proxy/guardrails/_content_utils.py + # inspects and rewrites these summary blocks before they are forwarded. summary: Final[object] = input_item.get("summary") if isinstance(summary, list): text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 3dfb98c12ea..d9e079c6d92 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -149,6 +149,22 @@ def test_iter_message_text_responses_api_tool_call_taxonomy(): assert list(iter_message_text(data)) == ["hello", "sunny"] +def test_iter_message_text_inspects_reasoning_content_and_summary(): + """VERIA: reasoning items forwarded as ``reasoning_content`` must be + inspected, including ``summary`` blocks the bridge reads as a fallback.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "content secret"}], + "summary": [{"type": "summary_text", "text": "summary secret"}], + } + ] + } + assert list(iter_message_text(data)) == ["content secret", "summary secret"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -308,6 +324,27 @@ def test_walk_user_text_redacts_mixed_list_input(): assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}} +def test_walk_user_text_redacts_reasoning_content_and_summary(): + """VERIA: in-place redaction must cover both plaintext shapes the bridge + forwards from a reasoning item.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "AKIAEXAMPLE content"}], + "summary": [{"type": "summary_text", "text": "AKIAEXAMPLE summary"}], + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 2 + item = data["input"][0] + assert item["content"][0]["text"] == "[REDACTED] content" + assert item["summary"][0]["text"] == "[REDACTED] summary" + assert item["id"] == "rs_1" + + # ── build_inspection_messages ───────────────────────────────────────────────── @@ -462,6 +499,23 @@ def test_build_inspection_messages_empty_data(): assert build_inspection_messages({"input": ""}) == [] +def test_build_inspection_messages_includes_reasoning_summary(): + """VERIA: remote guardrail APIs must see reasoning summaries even when + the reasoning item has no ``content`` field.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "secret summary"}], + } + ] + } + assert build_inspection_messages(data) == [ + {"role": "assistant", "content": "secret summary"} + ] + + # ── has_non_string_content ──────────────────────────────────────────────────── From 34e692c903c9d75b52065b4093c7d80b7eb2e00b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:16:22 -0700 Subject: [PATCH 031/319] fix(batches): decode model-encoded output file id so completed batches book spend Adds e2e coverage for batches terminal state and cost write-back, failure paths, per-backend file content downloads, and two-gateway routing (LIT-5730). --- litellm/batches/batch_utils.py | 40 +- tests/e2e/batches/COVERAGE.md | 87 +++- tests/e2e/batches/batch_client.py | 29 +- tests/e2e/batches/capabilities.py | 10 + tests/e2e/batches/test_batches_e2e.py | 471 +++++++++++++++++- .../llm_nonconversational.yaml | 10 + .../test_litellm/batches/test_batch_utils.py | 37 ++ 7 files changed, 646 insertions(+), 38 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0cf22d82ca6..6eb13d2cba7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -296,6 +296,32 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +def _provider_output_file_id(output_file_id: str) -> str: + """ + Resolve the file id the provider actually knows: unified ids yield their embedded + llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_original_file_id, + ) + + unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id) + if not unified_file_id: + return get_original_file_id(output_file_id) + try: + extracted: Final = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", + output_file_id, + e, + ) + return output_file_id + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", extracted) + return extracted + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -311,23 +337,11 @@ async def _fetch_batch_output_file_content( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + file_id: Final = _provider_output_file_id(batch.output_file_id) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca48204962a..f02d4eb4fe4 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -1,9 +1,11 @@ # Batches Test Coverage Matrix Live e2e coverage of the Batches API over a real proxy, real provider keys, and -real cost. Synchronous tier only: a batch's completion window is 24h, so these -tests never wait for `completed`. They assert the proxy accepts, routes, retrieves, -cancels, and lists a batch; everything created is deleted on teardown. +real cost. Mostly synchronous tier: a batch's completion window is 24h, so the +lifecycle matrix never waits for `completed`. It asserts the proxy accepts, routes, +retrieves, cancels, and lists a batch; everything created is deleted on teardown. +The exception is `TestBatchTerminalState`, which covers the completed state and +cost write-back via a cross-run marker baton (design below). ## Provider x operation @@ -12,19 +14,26 @@ row per supported (provider, scenario) pair, so there are no skipped cells in th parametrized run. The batches suite never skips: missing provider creds or upstream failures are hard test failures (see `tests/e2e/CLAUDE.md`). -| Provider | create | retrieve | cancel | list | file backing | -|-----------|--------|----------|--------|------|--------------| -| OpenAI | yes | yes | yes | yes | OpenAI Files | -| Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | -| Bedrock | yes (unified only) | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | +| Provider | create | retrieve | cancel | list | content download | file backing | +|-----------|--------|----------|--------|------|------------------|--------------| +| OpenAI | yes | yes | yes | yes | yes (lifecycle + terminal output) | OpenAI Files | +| Azure | yes | yes | yes | yes | yes (byte-verbatim) | Azure Files | +| Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | +| Bedrock | yes (unified only) | yes | no (limited upstream) | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off -(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. +(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix; +flipping those gates is tracked in LIT-4774 and deliberately not part of this suite. Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only); `model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no model-less passthrough path. +`GET /v1/files/{id}/content` is exercised for the unified upload path per backend in +`test_unified_file_content_downloads`. Azure stores the JSONL verbatim, so its download +is asserted byte-equal to the upload. Vertex (GCS) and Bedrock (S3) transform lines at +upload time, so those assert a 200 with non-empty parseable JSON lines instead. Gemini +(non-Vertex) raises `NotImplementedError` for file content and has no cell here. + ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`) Each create-capable provider runs all four. The test asserts the returned file id @@ -71,11 +80,59 @@ File delete asserts `object=="file"` and `deleted==True`. | `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, per-backend content download, failure paths, second-hop routing, terminal state + cost | + +## Failure paths + +`TestBatchFailurePaths` pins the customer-facing error contracts. A malformed input +file is a 400 at upload naming the bad content. A JSONL line whose url contradicts +the batch endpoint passes create (providers validate asynchronously) and drives the +batch to `failed` with structured `errors.data` (code/line/message), a null +`output_file_id`, and a $0 spend row keyed `{batch_id}_batch_cost` (LIT-4852: a +failed batch books $0 instead of crashing cost tracking). Cancelling that failed +batch is a 409 naming the terminal status. A file id encoded for one deployment wins +over a conflicting `model` param on create: the batch routes and re-encodes by the +file's embedded model (foreign-id precedence). + +## Second hop (two chained gateways) + +`TestBatchSecondHop` registers a `litellm_proxy/` deployment pointing at +the proxy's own base URL with a freshly minted virtual key, so unified upload and +create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: +`target_model_names` is rewritten to the inner deployment on the second hop and the +nested managed ids round-trip retrieve. This self-chaining only needs the proxy to +reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. + +## Terminal state + cost write-back (cross-run marker baton) + +The 24h completion window rules out submit-and-wait inside one run, so +`TestBatchTerminalState` amortizes across runs. Each run submits a 1-line marker +batch (stable metadata key/value plus a per-run field) and deliberately never +cancels or deletes it or its input file: the marker is the baton the next run picks +up (OpenAI files expire on their own after ~30 days). Polling is list-only, up to 5 +minutes, because retrieving a non-terminal batch books a $0 spend row whose +request_id then blocks the later real-cost row (`skip_duplicates`); the single +retrieve happens only once a completed marker exists. The assertion target is the +newest completed marker from ANY run: run-scoped deployment names mean the list +re-encodes prior-run batches under new encoded ids, so their spend keys are fresh +and a prior-run marker is billable by this run. On the 6h stage cadence the full +assertions are therefore deterministic from run 2 onward. On a cold start (no +completed marker within the poll budget) the test passes on the submission +assertions alone: a documented vacuous pass, not a skip. Markers aged past the 24h +window (25h-73h band, within the newest 100-item list page) must be terminal. + +The cost assertion is the LIT-5730 headline: retrieving a completed model-encoded +batch must write a positive spend row with call_type `aretrieve_batch` and token +usage. Before the fix in `litellm/batches/batch_utils.py`, the retrieve endpoint +re-encoded the response's `output_file_id` in place before the queued logging +worker ran, the worker sent that encoded id to OpenAI, got a 404, and the spend row +never landed. ## Out of scope (intentionally) -Driving a batch to `completed`, cost tracking on completion, and the DB write-back -are not covered here; the 24h window makes them unfit for a synchronous gate. That -logic belongs in a DI-stubbed proxy integration test under `tests/test_litellm/proxy/` -where the provider client is injected to return `completed` deterministically. +Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a +terminal DB status short-circuits retrieve for those ids, so the terminal-state cell +uses the encoded path; poller timing does not fit an e2e gate and belongs in a +DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Bedrock +cancel/list stay gated pending LIT-4774. Gemini (non-Vertex) file content raises +`NotImplementedError` upstream and is not a coverage cell. diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..21a56f3398f 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -44,6 +44,17 @@ class FileList(BaseModel): data: list[FileObject] = [] +class BatchErrorItem(BaseModel): + code: str | None = None + line: int | None = None + message: str | None = None + + +class BatchErrorList(BaseModel): + object: str | None = None + data: list[BatchErrorItem] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -51,6 +62,9 @@ class BatchObject(BaseModel): endpoint: str | None = None input_file_id: str | None = None output_file_id: str | None = None + error_file_id: str | None = None + errors: BatchErrorList | None = None + metadata: dict[str, str] | None = None completion_window: str | None = None created_at: int | None = None model: str | None = None @@ -72,12 +86,18 @@ class BatchCreateBody(BaseModel): endpoint: str = "/v1/chat/completions" completion_window: str = "24h" model: str | None = None + metadata: dict[str, str] | None = None class ModelQuery(BaseModel): model: str | None = None +class BatchListQuery(BaseModel): + model: str | None = None + limit: int | None = None + + def is_model_access_denied(resp: StreamingResponse) -> bool: """True if the proxy rejected the call because the key may not access the model.""" return resp.status_code == 403 and "key_model_access_denied" in resp.body @@ -168,12 +188,17 @@ class BatchClient: ) def list_batches( - self, *, key: str, provider: str | None = None + self, + *, + key: str, + provider: str | None = None, + model: str | None = None, + limit: int | None = None, ) -> Result[BatchList]: return self.proxy.transport.get( _batches_path(provider), headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=BatchListQuery(model=model, limit=limit), response_type=BatchList, ) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3988fb5e7e1..ce1f68184a7 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -210,6 +210,16 @@ def is_model_encoded_id(id_str: str) -> bool: return False +def decoded_model_from_id(id_str: str) -> str | None: + """Deployment name embedded in a model-encoded file/batch id, or None.""" + for prefix in ("file-", "batch_"): + if id_str.startswith(prefix): + decoded = _b64_decode(id_str[len(prefix) :]) + if decoded.startswith("litellm:") and ";model," in decoded: + return decoded.split(";model,", 1)[1].split(";")[0] + return None + + def matches_id_shape(shape: IdShape, id_str: str) -> bool: if shape == "managed": return is_managed_id(id_str) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..75b3a6cd758 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -1,11 +1,12 @@ """Live e2e for the Batches API across every provider LiteLLM supports. -Synchronous tier only: a batch's completion window is 24h, so these never wait for -"completed". Each case uploads a tiny JSONL, creates the batch through one of the -four routing scenarios, asserts it was accepted (non-terminal status) and routed to -the right provider, then retrieves / cancels / lists where the provider supports it. -Everything created is deleted on teardown. Completion + cost tracking are out of -scope here (see COVERAGE.md). +Mostly synchronous tier: a batch's completion window is 24h, so the lifecycle +matrix never waits for "completed". Each case uploads a tiny JSONL, creates the +batch through one of the four routing scenarios, asserts it was accepted +(non-terminal status) and routed to the right provider, then retrieves / cancels / +lists where the provider supports it. Everything created is deleted on teardown. +The exception is TestBatchTerminalState, which carries completed-state + cost +write-back coverage via a cross-run marker baton (design in COVERAGE.md). Routing signal: for provider_fallback the raw batch id discriminates the provider; for the encoded/unified/model_param scenarios the proxy re-encodes the id, so the @@ -22,8 +23,9 @@ from datetime import datetime, timedelta, timezone from typing import Callable import pytest +from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import PROXY_BASE_URL, unique_marker from batch_client import ( UPLOAD_FILENAME, @@ -40,9 +42,12 @@ from capabilities import ( CAPABILITIES, FILE_ID_SHAPE, OPENAI_BATCH_MODEL, + PROVIDERS, Capability, + Provider, batch_model_name, coverage_cells_for_lifecycle, + decoded_model_from_id, is_managed_id, matches_id_shape, raw_id_matches_provider, @@ -475,9 +480,22 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" +FILE_CONTENT_CELLS = { + "azure": "llm.files.azure_openai.content.nonstream.works", + "vertex_ai": "llm.files.vertex.content.nonstream.works", + "bedrock": "llm.files.bedrock.content.nonstream.works", +} +BYTE_FIDELITY_CONTENT_PROVIDERS = frozenset({"azure"}) + class TestBatchFileContent: - """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes.""" + """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes. + + Azure stores the upload verbatim, so its download is asserted byte-equal. + Vertex (GCS) and Bedrock (S3) transform each JSONL line into the provider's + request format at upload time, so their downloads assert 200 plus non-empty + parseable JSON lines instead of byte equality. + """ @pytest.mark.covers( "llm.files.openai.content.nonstream.works", @@ -521,6 +539,62 @@ class TestBatchFileContent: "downloaded file content must match the uploaded JSONL bytes" ) + @pytest.mark.parametrize( + "provider", + [ + pytest.param( + p, + id=p.name, + marks=pytest.mark.covers( + FILE_CONTENT_CELLS[p.name], exercised_on=["files"] + ), + ) + for p in PROVIDERS + if p.name in FILE_CONTENT_CELLS + ], + ) + def test_unified_file_content_downloads( + self, + provider: Provider, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + payload = render_jsonl(provider.raw_model) + file = unwrap( + client.upload_file( + content=payload, + form=FileUploadForm(purpose="batch", target_model_names=provider.model), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider=provider.name) + assert is_managed_id(file.id), ( + f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"{provider.name}: file content must be 200, " + f"got {downloaded.status_code}: {downloaded.body[:300]}" + ) + body = downloaded.body.strip() + assert body, f"{provider.name}: file content download returned an empty body" + if provider.name in BYTE_FIDELITY_CONTENT_PROVIDERS: + assert body == payload.decode().strip(), ( + f"{provider.name}: downloaded content must match the uploaded JSONL bytes" + ) + else: + for line in body.splitlines(): + assert json.loads(line), ( + f"{provider.name}: content line is not JSON: {line[:200]}" + ) + class TestOpenAIFiles: """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. @@ -866,3 +940,384 @@ class TestHostedVllmBatch: f"hosted_vllm batch has non-transitional status {batch.status!r}" ) assert_batch_object(batch) + + +BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"}) +FAILED_BATCH_POLL_SECONDS = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS = 5.0 + +AZURE_BATCH_RAW_MODEL = next(p.raw_model for p in PROVIDERS if p.name == "azure") + + +def _mismatched_endpoint_jsonl(model: str) -> bytes: + line = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model, "input": "ping"}, + } + return (json.dumps(line) + "\n").encode() + + +def _poll_until_terminal(client: BatchClient, batch_id: str, key: str) -> BatchObject: + deadline = time.monotonic() + FAILED_BATCH_POLL_SECONDS + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + while fetched.status not in BATCH_TERMINAL_STATUSES and time.monotonic() < deadline: + time.sleep(FAILED_BATCH_POLL_INTERVAL_SECONDS) + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + return fetched + + +class TestBatchFailurePaths: + """Customer-facing failure contracts for /v1/batches. + + A malformed input file is rejected at upload with a 400 naming the bad + content. A JSONL line whose url contradicts the batch endpoint is accepted + at create (providers validate asynchronously) and drives the batch to + "failed" with structured per-line errors, a null output_file_id, and a + zero-cost spend row (LIT-4852: a failed batch must book $0, not crash cost + tracking). Cancelling that already-failed batch returns a 409 naming the + terminal status. A file id encoded for one deployment wins over a + conflicting model param on create: the batch routes (and re-encodes) by the + file's embedded model, pinning that precedence. + """ + + @pytest.mark.covers( + "llm.batches.openai.malformed_jsonl.nonstream.works", + exercised_on=["files"], + ) + def test_malformed_jsonl_upload_rejected( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + result = client.upload_file( + content=b"this is not json\n", + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=resources.key(), + ) + match result: + case UnknownApiError(status_code=400, body=body): + assert "json" in body.lower(), ( + f"400 must name the malformed JSONL so users can fix the file, got: {body[:300]}" + ) + case _: + pytest.fail(f"malformed JSONL upload must be rejected with a 400, got: {result}") + + @pytest.mark.covers( + "llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works", + "llm.batches.openai.cancel_terminal.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_endpoint_mismatch_fails_batch_and_cancel_conflicts( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=_mismatched_endpoint_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + + fetched = _poll_until_terminal(client, batch.id, key) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch must fail, got {fetched.status!r}" + ) + assert fetched.output_file_id is None, ( + f"failed batch must have no output file, got {fetched.output_file_id!r}" + ) + assert fetched.errors is not None and fetched.errors.data, ( + "failed batch must surface structured errors so users can fix the JSONL" + ) + first_error = fetched.errors.data[0] + assert first_error.message, "batch error item has no message" + assert first_error.code, "batch error item has no code" + + rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost") + assert rows, ( + f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)" + ) + assert all((row.spend or 0) == 0 for row in rows), ( + f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}" + ) + assert rows[0].call_type == "aretrieve_batch", ( + f"batch cost row call_type={rows[0].call_type!r}" + ) + + conflict = client.cancel_batch(batch.id, key=key) + match conflict: + case UnknownApiError(status_code=409, body=body): + assert "failed" in body.lower(), ( + f"409 must name the terminal status blocking the cancel, got: {body[:300]}" + ) + case _: + pytest.fail(f"cancel of a failed batch must return a 409 conflict, got: {conflict}") + + @pytest.mark.covers( + "llm.batches.openai.foreign_file_id.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_foreign_encoded_file_id_routes_by_file_model( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(AZURE_BATCH_RAW_MODEL), + form=FileUploadForm(purpose="batch"), + model=AZURE_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( + f"upload did not encode the azure deployment into the file id: {file.id!r}" + ) + + created = client.create_batch( + body=BatchCreateBody(input_file_id=file.id, model=OPENAI_BATCH_MODEL), key=key + ) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( + "create with a foreign encoded file id must route by the file's embedded model, " + f"but the batch id encodes {decoded_model_from_id(batch.id)!r} " + f"(model param was {OPENAI_BATCH_MODEL!r})" + ) + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "retrieved foreign-file batch has no status" + + +class TestBatchSecondHop: + """Two-proxy batch routing: a litellm_proxy deployment chained to the gateway + itself (LIT-5347, PR #36240). + + The hop deployment's litellm_params point litellm_proxy/ at this + gateway's own base URL with a freshly minted virtual key, so the unified + upload and batch create traverse gateway -> gateway -> OpenAI. The regression + this pins: target_model_names must be rewritten to the inner deployment on + the second hop and the nested managed ids must round-trip retrieve. + """ + + @pytest.mark.covers( + "llm.batches.openai.second_hop.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_create_and_retrieve_via_chained_gateway( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + hop_name = batch_model_name("openai-batch-hop") + model_id = client.create_model( + hop_name, + LiteLLMParamsBody( + model=f"litellm_proxy/{OPENAI_BATCH_MODEL}", + api_base=PROXY_BASE_URL, + api_key=key, + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch", target_model_names=hop_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert is_managed_id(file.id), ( + f"second-hop unified upload must return a managed file id, got {file.id!r}" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert is_managed_id(batch.id), ( + f"second-hop create must return a managed batch id, got {batch.id!r}" + ) + assert batch.status in CREATED_BATCH_STATUSES, ( + f"second-hop batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) + + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "second-hop retrieve returned no status" + + +class BatchOutputBody(BaseModel): + choices: list[object] = [] + + +class BatchOutputResponse(BaseModel): + status_code: int | None = None + body: BatchOutputBody | None = None + + +class BatchOutputLine(BaseModel): + response: BatchOutputResponse + + +TERMINAL_MARKER_KEY = "litellm_e2e_suite" +TERMINAL_MARKER_VALUE = "batches-terminal-baton" +TERMINAL_POLL_SECONDS = 300.0 +TERMINAL_POLL_INTERVAL_SECONDS = 10.0 +TERMINAL_LIST_LIMIT = 100 +TERMINAL_BAND_MIN_AGE_SECONDS = 25 * 3600 +TERMINAL_BAND_MAX_AGE_SECONDS = 73 * 3600 + + +def _marker_batches(client: BatchClient, key: str) -> list[BatchObject]: + listed = unwrap( + client.list_batches(key=key, model=OPENAI_BATCH_MODEL, limit=TERMINAL_LIST_LIMIT) + ) + return [ + b + for b in listed.data + if (b.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE + ] + + +def _await_completed_marker( + client: BatchClient, key: str +) -> tuple[BatchObject | None, list[BatchObject]]: + deadline = time.monotonic() + TERMINAL_POLL_SECONDS + while True: + markers = _marker_batches(client, key) + completed = max( + (b for b in markers if b.status == "completed"), + key=lambda b: b.created_at or 0, + default=None, + ) + if completed is not None or time.monotonic() >= deadline: + return completed, markers + time.sleep(TERMINAL_POLL_INTERVAL_SECONDS) + + +def _assert_aged_markers_terminal(markers: list[BatchObject]) -> None: + now = time.time() + stuck = [ + b + for b in markers + if b.created_at is not None + and TERMINAL_BAND_MIN_AGE_SECONDS <= now - b.created_at <= TERMINAL_BAND_MAX_AGE_SECONDS + and b.status not in BATCH_TERMINAL_STATUSES + ] + assert not stuck, ( + "marker batches past their 24h completion window must be terminal; stuck: " + f"{[(b.id, b.status, b.created_at) for b in stuck]}" + ) + + +class TestBatchTerminalState: + """Terminal state + cost write-back via a cross-run marker baton. + + Each run submits a 1-line marker batch (stable metadata key/value plus a + per-run field) and never cancels or deletes it: the marker is the baton the + next run picks up. Polling is list-only for up to 5 minutes because a + retrieve of a non-terminal batch books a $0 spend row whose request_id then + blocks the real-cost row (skip_duplicates); the single retrieve happens only + once a completed marker exists. The assertion target is the newest completed + marker from ANY run, so on the 6h stage cadence the full assertions are + deterministic from run 2 onward. On a cold start (no marker has ever + completed within the poll budget) the test passes on the submission + assertions alone: that is a documented vacuous pass, not a skip, and this + run's marker becomes the next run's target. Markers aged past OpenAI's 24h + completion window (25h-73h band, within the newest list page) must be + terminal. The cost assertion is the LIT-5730 headline: retrieving a + completed model-encoded batch must write a positive spend row keyed + {batch_id}_batch_cost; before the fix the logging worker fetched the + re-encoded output_file_id, 404d, and the row never landed. + """ + + @pytest.mark.covers( + "llm.batches.openai.terminal_state.nonstream.works", + "llm.batches.openai.terminal_state.nonstream.cost_logged", + exercised_on=["batches", "files"], + ) + def test_completed_batch_downloads_output_and_books_cost( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + created = client.create_batch( + body=BatchCreateBody( + input_file_id=file.id, + metadata={ + TERMINAL_MARKER_KEY: TERMINAL_MARKER_VALUE, + "run": unique_marker(), + }, + ), + key=key, + ) + require_successful_call(created) + submitted = BatchObject.model_validate_json(created.body) + assert submitted.status in CREATED_BATCH_STATUSES, ( + f"marker batch has non-transitional status {submitted.status!r}" + ) + assert (submitted.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE, ( + f"create dropped the marker metadata: {submitted.metadata!r}" + ) + + completed, markers = _await_completed_marker(client, key) + _assert_aged_markers_terminal(markers) + if completed is None: + return + + fetched = retrieve_batch(client, completed.id, key=key, provider=None) + assert fetched.status == "completed", ( + f"listed-completed marker retrieved as {fetched.status!r}" + ) + assert fetched.output_file_id, "completed batch has no output_file_id" + + downloaded = client.proxy.transport.download( + f"/v1/files/{fetched.output_file_id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"output content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + first_line = BatchOutputLine.model_validate_json(downloaded.body.strip().splitlines()[0]) + assert first_line.response.status_code == 200, ( + f"batch output line reports failure: {downloaded.body[:400]}" + ) + assert first_line.response.body is not None and first_line.response.body.choices, ( + "batch output line has no choices" + ) + + rows = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced = [row for row in rows if (row.spend or 0) > 0] + assert priced, ( + f"completed batch {fetched.id} wrote no positive-cost spend row under " + f"request_id {fetched.id}_batch_cost; cost write-back is broken (LIT-5730)" + ) + cost_row = priced[0] + assert cost_row.call_type == "aretrieve_batch", ( + f"batch cost row call_type={cost_row.call_type!r}" + ) + assert (cost_row.total_tokens or 0) > 0, ( + f"batch cost row has no token usage: {cost_row.total_tokens!r}" + ) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..3de462683d2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -24,6 +24,13 @@ - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"} - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"} - {id: llm.batches.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.18 / LIT-4778", rationale: "Missing input_file_id and invalid batch id rejected"} +- {id: llm.batches.openai.terminal_state.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "A batch actually reaches completed and its output file downloads through GET /v1/files/{id}/content with per-line provider responses"} +- {id: llm.batches.openai.terminal_state.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_batches_e2e.py / LIT-5730", fail_before_fix: proven, rationale: "Retrieving a completed model-encoded batch writes a positive spend row keyed {batch_id}_batch_cost (pins LIT-4852/LIT-5666; before the fix the logging worker 404d fetching the re-encoded output_file_id and the row was never written)"} +- {id: llm.batches.openai.malformed_jsonl.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Uploading a non-JSON batch file is rejected with a 400 naming the bad line"} +- {id: llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "JSONL line url that contradicts the batch endpoint drives the batch to failed with structured errors, retrieve stays clean, and the terminal retrieve books a zero-cost spend row (LIT-4852)"} +- {id: llm.batches.openai.cancel_terminal.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Cancelling an already-terminal batch returns a 409 conflict naming the terminal status"} +- {id: llm.batches.openai.foreign_file_id.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Create with one deployment's encoded file id and a conflicting model param routes by the file's embedded model; the returned batch id pins that precedence"} +- {id: llm.batches.openai.second_hop.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5347", rationale: "A litellm_proxy deployment chained to the gateway itself preserves target_model_names through nested unified ids; upload, create, and retrieve work over the two-hop chain (PR #36240)"} - {id: llm.files.openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai_files_endpoints/files_endpoints.py:46", rationale: "File upload returns OpenAIFileObject"} - {id: llm.files.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.16 / LIT-4778", rationale: "File upload without purpose rejected"} - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} @@ -36,6 +43,9 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} +- {id: llm.files.azure_openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on an Azure unified file returns the uploaded JSONL bytes verbatim"} +- {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} +- {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index ebe093c591c..08cdf945b80 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -800,6 +800,43 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke assert captured["custom_llm_provider"] == "vertex_ai" +@pytest.mark.asyncio +async def test_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch): + import litellm.files.main as files_main + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model") + + await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai") + + assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR" + assert captured["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai") + + assert captured["file_id"] == "file-abc123" + + def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens): return { "request": { From 42cffe93c821ce60a8a6f24d95cd01d995e52f32 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 19 Aug 2026 16:45:52 -0700 Subject: [PATCH 032/319] Add moonshot/kimi-k3 to model prices and context window map Pricing per https://platform.kimi.ai/docs/pricing/chat-k3: - $3.00/M input (cache miss), $0.30/M cache read, $15.00/M output - 1,048,576 context window; max_completion_tokens settable up to 1,048,576 - Supports reasoning (reasoning_effort low/high/max), tool calling, structured output, vision and video input Co-Authored-By: Claude Fable 5 --- .../model_prices_and_context_window_backup.json | 17 +++++++++++++++++ model_prices_and_context_window.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0eca17272d..a7c9825ee7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0eca17272d..a7c9825ee7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", From f8b31f493a62a7b43a2effced84c8a9557929ffd Mon Sep 17 00:00:00 2001 From: 72004 Date: Fri, 21 Aug 2026 02:05:05 +0500 Subject: [PATCH 033/319] fix: don't retire a completed batch from cost recovery while output_file_id is still lagging --- .../openai_files_endpoints/common_utils.py | 21 +++++++++- .../test_files_common_utils.py | 42 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b9af01e9aea..f8896771077 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1288,6 +1288,25 @@ def batch_cost_poller_is_active() -> bool: return False +def _completed_batch_safe_to_retire(response) -> bool: + """Whether a "completed" batch may be retired from cost recovery. + + ``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's + cost-recovery poller, so setting it retires the batch permanently. A batch can + reach ``status="completed"`` while ``output_file_id`` is still ``None`` (the + provider response briefly lags before the output id populates). Retiring in that + window loses the spend record forever. Retire only once we can prove there is + nothing left to recover: the output file has actually arrived, or the provider + reports no successful request lines. When counts are unknown, stay eligible so + the next poller pass revisits it. (#37713) + """ + if getattr(response, "output_file_id", None) is not None: + return True + request_counts = getattr(response, "request_counts", None) + completed = getattr(request_counts, "completed", None) + return completed == 0 + + async def update_batch_in_database( batch_id: str, unified_batch_id: str | Literal[False], @@ -1369,7 +1388,7 @@ async def update_batch_in_database( } poller_owns: Final = batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting - if db_status == "complete" and not poller_owns: + if db_status == "complete" and not poller_owns and _completed_batch_safe_to_retire(response): update_data["batch_processed"] = True try: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 6ffb7daaa2d..eb6596e274c 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -431,3 +431,45 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") assert data == {"batch_id": "unified-batch-id"} + + +from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, +) + + +def _completed_batch(output_file_id, completed=None) -> LiteLLMBatch: + kwargs = dict( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id=output_file_id, + error_file_id=None, + ) + if completed is not None: + kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0} + return LiteLLMBatch(**kwargs) + + +class TestCompletedBatchSafeToRetire: + """A completed batch is only safe to retire from cost recovery once its output + file has arrived or the provider proves no successful lines (#37713).""" + + def test_output_file_present_is_safe(self): + assert _completed_batch_safe_to_retire(_completed_batch("file-out")) is True + + def test_no_output_and_no_successful_lines_is_safe(self): + # Every request line errored -> nothing left to recover. + assert _completed_batch_safe_to_retire(_completed_batch(None, completed=0)) is True + + def test_no_output_but_successful_lines_is_not_safe(self): + # The bug: output_file_id is lagging; retiring here loses the spend record. + assert _completed_batch_safe_to_retire(_completed_batch(None, completed=5)) is False + + def test_no_output_and_unknown_counts_is_not_safe(self): + # Counts unknown -> stay eligible so the next poller pass revisits it. + assert _completed_batch_safe_to_retire(_completed_batch(None)) is False From 67d16a499dc208f69ddab20e17648b503acc07ec Mon Sep 17 00:00:00 2001 From: 72004 Date: Fri, 21 Aug 2026 02:13:11 +0500 Subject: [PATCH 034/319] Type the batch-retire helpers and rename test helper to avoid shadowing existing _completed_batch --- litellm/proxy/openai_files_endpoints/common_utils.py | 2 +- .../openai_files_endpoint/test_files_common_utils.py | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f8896771077..2e8ae6af7a9 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1288,7 +1288,7 @@ def batch_cost_poller_is_active() -> bool: return False -def _completed_batch_safe_to_retire(response) -> bool: +def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: """Whether a "completed" batch may be retired from cost recovery. ``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index eb6596e274c..3de9e61463f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -438,7 +438,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) -def _completed_batch(output_file_id, completed=None) -> LiteLLMBatch: +def _completed_batch_for_retire( + output_file_id: str | None, completed: int | None = None +) -> LiteLLMBatch: kwargs = dict( id="batch-1", completion_window="24h", @@ -460,16 +462,16 @@ class TestCompletedBatchSafeToRetire: file has arrived or the provider proves no successful lines (#37713).""" def test_output_file_present_is_safe(self): - assert _completed_batch_safe_to_retire(_completed_batch("file-out")) is True + assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True def test_no_output_and_no_successful_lines_is_safe(self): # Every request line errored -> nothing left to recover. - assert _completed_batch_safe_to_retire(_completed_batch(None, completed=0)) is True + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True def test_no_output_but_successful_lines_is_not_safe(self): # The bug: output_file_id is lagging; retiring here loses the spend record. - assert _completed_batch_safe_to_retire(_completed_batch(None, completed=5)) is False + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False def test_no_output_and_unknown_counts_is_not_safe(self): # Counts unknown -> stay eligible so the next poller pass revisits it. - assert _completed_batch_safe_to_retire(_completed_batch(None)) is False + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False From 7d9e3756980135699a43f5c3c3d892a87b7ec842 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 21 Aug 2026 17:28:12 +1000 Subject: [PATCH 035/319] fix(scx-ai): use the published scx.ai rates and the scx_ai docs url Applies the review suggestions. The cost map now carries the rates published on https://scx.ai/pricing, GLM-5.2 at 0.61 in, 0.22 cached, 1.98 out and Qwen3.8-Max at 1.65 in, 0.21 cached, 4.99 out per million tokens, and cites that page as the source rather than a third party gateway. The provider link is corrected to https://docs.litellm.ai/docs/providers/scx_ai to match the page that shipped as scx_ai.md. Both the primary files and their backup mirrors are updated. --- .../model_prices_and_context_window_backup.json | 14 +++++++------- litellm/provider_endpoints_support_backup.json | 2 +- model_prices_and_context_window.json | 14 +++++++------- provider_endpoints_support.json | 2 +- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 222f4dd4db6..6c86c56d52a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36725,15 +36725,15 @@ "source": "https://cloud.sambanova.ai/plans/pricing" }, "scx-ai/GLM-5.2": { - "cache_read_input_token_cost": 1.375e-07, - "input_cost_per_token": 5.5e-07, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 6.1e-07, "litellm_provider": "scx-ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.9255e-06, - "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", + "output_cost_per_token": 1.98e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -36743,14 +36743,14 @@ }, "scx-ai/Qwen3.8-Max": { "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 1.815e-06, + "input_cost_per_token": 1.65e-06, "litellm_provider": "scx-ai", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5.4461e-06, - "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "output_cost_per_token": 4.99e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index c1928c34349..86c14fb4cd8 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2029,7 +2029,7 @@ }, "scx-ai": { "display_name": "SCX.ai (`scx-ai`)", - "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "url": "https://docs.litellm.ai/docs/providers/scx_ai", "endpoints": { "chat_completions": true, "messages": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 222f4dd4db6..6c86c56d52a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36725,15 +36725,15 @@ "source": "https://cloud.sambanova.ai/plans/pricing" }, "scx-ai/GLM-5.2": { - "cache_read_input_token_cost": 1.375e-07, - "input_cost_per_token": 5.5e-07, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 6.1e-07, "litellm_provider": "scx-ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.9255e-06, - "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", + "output_cost_per_token": 1.98e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -36743,14 +36743,14 @@ }, "scx-ai/Qwen3.8-Max": { "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 1.815e-06, + "input_cost_per_token": 1.65e-06, "litellm_provider": "scx-ai", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5.4461e-06, - "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "output_cost_per_token": 4.99e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 3da0ec7d6b4..1d8d374c2c4 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2263,7 +2263,7 @@ }, "scx-ai": { "display_name": "SCX.ai (`scx-ai`)", - "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "url": "https://docs.litellm.ai/docs/providers/scx_ai", "endpoints": { "chat_completions": true, "messages": false, From 86efa2bcfde555e519140f6d9721e9570aada2d2 Mon Sep 17 00:00:00 2001 From: longwind48 Date: Fri, 21 Aug 2026 17:27:28 +0800 Subject: [PATCH 036/319] feat(bedrock): serve gpt-5.6 cross-region inference profiles on bedrock runtime GPT-5.6 Sol, Terra and Luna reached the bedrock-runtime data plane on 2026-08-17, separately from the existing bedrock-mantle path. On runtime they are served only through cross-region inference profiles, so bedrock/us.openai.gpt-5.6-* had no cost map entry and fell through to the Invoke route, which rewrites the token cap to max_tokens and is rejected as unsupported_parameter on both /v1/chat/completions and /v1/responses. Register the Geo and Global profiles as bedrock_converse so routing reaches Converse, which AWS documents and serves for these models, and price each profile from its own published rate table. No bare key: the control plane reports inferenceTypesSupported INFERENCE_PROFILE with no on-demand throughput, so a bare id is not invocable. Declare the published cache-read and cache-write rates. Bedrock rejects an explicit cachePoint block for these models, so supports_prompt_caching stays off, but it caches long prefixes implicitly and reports the cache tokens in usage either way. Without the cost fields a cache-read turn bills only its uncached tokens: measured against live Bedrock, a 15609-token cached prefix came to $0.000176 instead of $0.00876095. Clients that resend a long prefix every turn are the worst affected. Reasoning stays unadvertised. Converse rejects the Anthropic-shaped thinking block LiteLLM sends for reasoning_effort; the shape these models accept is additionalModelRequestFields {"reasoning": {"effort": ...}}, which needs a transform change tracked by #34105. Advertising it without that change is what made the earlier attempt in #37307 fail. --- ...odel_prices_and_context_window_backup.json | 150 ++++++++++ model_prices_and_context_window.json | 150 ++++++++++ ..._cross_region_inference_profile_mapping.py | 258 +++++++++++++++++- 3 files changed, 557 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 91c10d13e8e..be8e6da5f59 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-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", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "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, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 91c10d13e8e..be8e6da5f59 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-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", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "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, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 3a27f3ed002..22aba59fb5d 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,13 +1,132 @@ """Test Bedrock cross-region inference profile model mapping""" +import json import os import sys +from functools import lru_cache +from pathlib import Path +from typing import NamedTuple + +import pytest sys.path.insert(0, os.path.abspath("../../../..")) +import litellm +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.utils import _get_model_info_helper from litellm.cost_calculator import completion_cost -from litellm.types.utils import ModelResponse, Usage, Choices, Message +from litellm.types.utils import ( + Choices, + Message, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Resolve models against this checkout's cost map instead of the network-fetched + ``main`` copy, which lags this branch until merge.""" + original_converse_models = set(litellm.bedrock_converse_models) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + try: + litellm.bedrock_converse_models.update( + key + for key, value in litellm.model_cost.items() + if isinstance(value, dict) + and value.get("litellm_provider") == "bedrock_converse" + ) + yield + finally: + litellm.bedrock_converse_models.clear() + litellm.bedrock_converse_models.update(original_converse_models) + litellm.get_model_info.cache_clear() + + +class GptProfile(NamedTuple): + model_id: str + input_cost: float + input_cost_above_272k: float + cache_write: float + cache_write_above_272k: float + cache_read: float + cache_read_above_272k: float + output_cost: float + output_cost_above_272k: float + + +GPT_5_6_PROFILES = [ + GptProfile( + model_id="us.openai.gpt-5.6-sol", + input_cost=5.5e-06, input_cost_above_272k=1.1e-05, + cache_write=6.875e-06, cache_write_above_272k=1.375e-05, + cache_read=5.5e-07, cache_read_above_272k=1.1e-06, + output_cost=3.3e-05, output_cost_above_272k=4.95e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-sol", + input_cost=5e-06, input_cost_above_272k=1e-05, + cache_write=6.25e-06, cache_write_above_272k=1.25e-05, + cache_read=5e-07, cache_read_above_272k=1e-06, + output_cost=3e-05, output_cost_above_272k=4.5e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-terra", + input_cost=2.2e-06, input_cost_above_272k=4.4e-06, + cache_write=2.75e-06, cache_write_above_272k=5.5e-06, + cache_read=2.2e-07, cache_read_above_272k=4.4e-07, + output_cost=1.32e-05, output_cost_above_272k=1.98e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-terra", + input_cost=2e-06, input_cost_above_272k=4e-06, + cache_write=2.5e-06, cache_write_above_272k=5e-06, + cache_read=2e-07, cache_read_above_272k=4e-07, + output_cost=1.2e-05, output_cost_above_272k=1.8e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-luna", + input_cost=2.2e-07, input_cost_above_272k=4.4e-07, + cache_write=2.75e-07, cache_write_above_272k=5.5e-07, + cache_read=2.2e-08, cache_read_above_272k=4.4e-08, + output_cost=1.32e-06, output_cost_above_272k=1.98e-06, + ), + GptProfile( + model_id="global.openai.gpt-5.6-luna", + input_cost=2e-07, input_cost_above_272k=4e-07, + cache_write=2.5e-07, cache_write_above_272k=5e-07, + cache_read=2e-08, cache_read_above_272k=4e-08, + output_cost=1.2e-06, output_cost_above_272k=1.8e-06, + ), +] + + +@lru_cache(maxsize=1) +def _packaged_cost_map(): + """The map litellm actually resolves against, for fields ModelInfoBase drops.""" + path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" + return json.loads(path.read_text()) + + +def _bedrock_response(model, usage): + return ModelResponse( + id="test", + created=1234567890, + model=model, + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="OK", role="assistant"), + ) + ], + usage=usage, + ) def test_bedrock_cross_region_inference_profile_mapping(): @@ -52,3 +171,140 @@ def test_proxy_cost_calculation_scenario(): ) expected_cost = (100 * 8e-07) + (50 * 4e-06) assert cost == expected_cost + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_map): + """GPT-5.6 is served by Converse on bedrock-runtime, never by Invoke.""" + assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): + """Geo and Global profiles carry their own published rates, per context tier.""" + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["litellm_provider"] == "bedrock_converse" + assert model_info["mode"] == "chat" + assert model_info["max_input_tokens"] == 1000000 + assert model_info["input_cost_per_token"] == profile.input_cost + assert ( + model_info["input_cost_per_token_above_272k_tokens"] + == profile.input_cost_above_272k + ) + assert model_info["output_cost_per_token"] == profile.output_cost + assert ( + model_info["output_cost_per_token_above_272k_tokens"] + == profile.output_cost_above_272k + ) + assert model_info["cache_creation_input_token_cost"] == profile.cache_write + assert ( + model_info["cache_creation_input_token_cost_above_272k_tokens"] + == profile.cache_write_above_272k + ) + assert model_info["cache_read_input_token_cost"] == profile.cache_read + assert ( + model_info["cache_read_input_token_cost_above_272k_tokens"] + == profile.cache_read_above_272k + ) + + +def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): + """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" + response = _bedrock_response( + "bedrock/us.openai.gpt-5.6-sol", + Usage(prompt_tokens=300000, completion_tokens=1000, total_tokens=301000), + ) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx((300000 * 1.1e-05) + (1000 * 4.95e-05), rel=1e-9) + + +def test_bedrock_gpt_5_6_bills_cache_read_tokens(local_model_cost_map): + """Bedrock caches long prefixes implicitly and reports them, so a cache-read turn + must be billed at the cache rate rather than dropped to zero.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=15609), + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 5.5e-07) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + # Without cache_read_input_token_cost the cached prefix bills at zero. + assert cost > (15611 * 5.5e-06) * 0.1 + + +def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): + """The write side of the same cache cycle is billed at the 30m cache-write rate.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + cache_creation_input_tokens=15609, + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 6.875e-06) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( + profile, local_model_cost_map +): + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["supports_function_calling"] is True + assert model_info["supports_tool_choice"] is True + assert model_info["supports_vision"] is True + + # Bedrock rejects an explicit cachePoint block for these models, so the flag that + # offers caller-driven caching stays off even though the cache rates are declared. + assert not model_info.get("supports_prompt_caching") + + # ModelInfoBase drops these two, so they are read from the map litellm resolves. + raw = _packaged_cost_map()[profile.model_id] + assert raw["supported_modalities"] == ["text", "image"] + assert raw["supported_output_modalities"] == ["text"] + # No bedrock_converse entry declares supported_endpoints; these models are reachable + # on chat completions and on the Responses API without it. + assert "supported_endpoints" not in raw + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_offers_tools_but_not_reasoning(profile, local_model_cost_map): + """Converse rejects the Anthropic-shaped thinking block LiteLLM emits for + reasoning_effort, so neither reasoning param may be offered yet, while the tool + params these models do accept must be.""" + supported = AmazonConverseConfig().get_supported_openai_params( + model=f"bedrock/{profile.model_id}" + ) + + assert "tools" in supported + assert "tool_choice" in supported + assert "reasoning_effort" not in supported + assert "thinking" not in supported From 13d4074492aa03b4d35a62fc8ffb8de2ef40e8dc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 21 Aug 2026 04:41:21 -0700 Subject: [PATCH 037/319] test(mcp): retire the last file of the dead tests/litellm mirror tests/litellm/ was a second mirror beside tests/test_litellm/ that no workflow, Makefile target, or CircleCI job ever named. Its other 33 files were reconciled during August 2026; this one stayed behind under a ci-coverage-allowlist entry asking a later pass to decide which of its five orphan behaviours still hold. They no longer hold as written: 25 of its 32 cases fail against today's code, because the file froze on the day it stopped being collected and the endpoints kept moving. Three of the five are already covered by the live twin, and better. test_get_request_base_url_xff_trust_gate parametrizes the trust gate in both directions, including the exact untrusted-caller case the orphan asserted, and the standard and legacy protected-resource shapes are both exercised through use_standard_pattern. The other two were the only tests anywhere for validate_trusted_redirect_uri under that same gate, so they are ported rather than dropped, rebuilt on the live file's request-mock conventions. Both directions are load-bearing: forcing is_request_from_trusted_proxy to True fails the untrusted case, forcing it to False fails the trusted one. 313 tests pass in the live file, up from 311. Dropping the dead file clears one zero-assert TQ001 violation, so its ceiling ratchets down with it. --- .github/ci-coverage-allowlist.yml | 10 - test-quality-budget.json | 2 +- .../mcp_server/test_discoverable_endpoints.py | 1268 ----------------- .../mcp_server/test_discoverable_endpoints.py | 49 + 4 files changed, 50 insertions(+), 1279 deletions(-) delete mode 100644 tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index ff8fa864d4a..918589f84d1 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -48,16 +48,6 @@ test_paths: choice it informed is settled paths: - tests/code_coverage_tests/test_aio_http_image_conversion.py - - reason: >- - The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its - other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging - their bodies into the live file of the same name. This one cannot follow either route yet: - its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no - counterpart while 25 assertions fail against today's code, so what survives that rewrite - is a judgement about the endpoints, not a merge. Revisit by deciding which of the five - behaviours still hold - paths: - - tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py - reason: >- No job invokes this suite and its files mix pure transformation tests with ones driving live vendor vector stores, so assigning them needs a per-file decision diff --git a/test-quality-budget.json b/test-quality-budget.json index 1613c8c75cb..91ae881c83a 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -1,6 +1,6 @@ { "TQ001": { - "limit": 750 + "limit": 746 }, "TQ002": { "limit": 742 diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py deleted file mode 100644 index 2a8768df722..00000000000 --- a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ /dev/null @@ -1,1268 +0,0 @@ -"""Tests for MCP OAuth discoverable endpoints""" - -import pytest -from fastapi import HTTPException -from unittest.mock import AsyncMock, MagicMock, patch - -TRUSTED_PROXY_IP = "10.0.0.5" -TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] - - -def set_request_from_trusted_proxy(mock_request): - mock_request.client = MagicMock() - mock_request.client.host = TRUSTED_PROXY_IP - - -@pytest.fixture -def trusted_proxy_origin_headers(): - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - ): - yield - - -@pytest.mark.asyncio -async def test_authorize_endpoint_includes_response_type(): - """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Mock the encryption functions to avoid needing a signing key - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify response is a redirect - assert response.status_code == 307 # FastAPI RedirectResponse default - - # Verify response_type is in the redirect URL - assert "response_type=code" in response.headers["location"] - assert "https://provider.com/oauth/authorize" in response.headers["location"] - assert "client_id=test_client_id" in response.headers["location"] - assert "scope=read+write" in response.headers["location"] - - -@pytest.mark.asyncio -async def test_authorize_endpoint_forwards_pkce_parameters(): - """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server (simulating Google OAuth) - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock the encryption function - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" - - # Call authorize endpoint with PKCE parameters - response = await authorize( - request=mock_request, - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - redirect_uri="http://localhost:60108/callback", - state="test_client_state", - code_challenge="x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk", - code_challenge_method="S256", - ) - - # Verify response is a redirect - assert response.status_code == 307 - - # Verify PKCE parameters are included in the redirect URL - location = response.headers["location"] - assert "https://accounts.google.com/o/oauth2/v2/auth" in location - assert "code_challenge=x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk" in location - assert "code_challenge_method=S256" in location - assert "client_id=669428968603-test.apps.googleusercontent.com" in location - assert "response_type=code" in location - - -@pytest.mark.asyncio -async def test_token_endpoint_forwards_code_verifier(): - """Test that token endpoint forwards code_verifier for PKCE flow""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "ya29.test_access_token", - "token_type": "Bearer", - "expires_in": 3599, - "scope": "openid email https://www.googleapis.com/auth/drive", - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client with AsyncMock for async methods - from unittest.mock import AsyncMock - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_async_client = MagicMock() - # Use AsyncMock for the async post method - mock_async_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_async_client - - # Call token endpoint with code_verifier - response = await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="4/test_authorization_code", - redirect_uri="http://localhost:60108/callback", - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - client_secret="GOCSPX-test_secret", - code_verifier="test_code_verifier_from_client", - ) - - # Verify that the token endpoint was called with code_verifier - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - - # Check the data parameter includes code_verifier - assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" - assert call_args[1]["data"]["code"] == "4/test_authorization_code" - assert ( - call_args[1]["data"]["client_id"] - == "669428968603-test.apps.googleusercontent.com" - ) - assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" - assert call_args[1]["data"]["grant_type"] == "authorization_code" - - # Verify response - response_data = response.body - import json - - token_data = json.loads(response_data) - assert token_data["access_token"] == "ya29.test_access_token" - assert token_data["token_type"] == "Bearer" - - -@pytest.mark.asyncio -async def test_register_client_without_mcp_server_name_returns_dummy(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_returns_existing_server_credentials(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="stored_server", - name="stored_server", - server_name="stored_server", - alias="stored_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="existing-client", - client_secret="existing-secret", - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - try: - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - assert result == { - "client_id": "stored_server", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_remote_registration_success(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="remote_server", - name="remote_server", - server_name="remote_server", - alias="remote_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - client_secret=None, - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - registration_url="https://provider.example/oauth/register", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - request_payload = { - "client_name": "Litellm Proxy", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "token_endpoint_auth_method": "client_secret_post", - } - - mock_response = MagicMock() - mock_response.json.return_value = { - "client_id": "generated-client", - "client_secret": "generated-secret", - } - mock_response.raise_for_status = MagicMock() - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - try: - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value=request_payload), - ), - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", - return_value=mock_async_client, - ), - ): - response = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - import json - - assert response.status_code == 200 - payload = json.loads(response.body.decode("utf-8")) - assert payload == mock_response.json.return_value - - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - assert call_args.args[0] == oauth2_server.registration_url - assert call_args.kwargs["headers"] == { - "Content-Type": "application/json", - "Accept": "application/json", - } - assert call_args.kwargs["json"]["redirect_uris"] == [ - "https://proxy.litellm.example/callback" - ] - assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] - assert ( - call_args.kwargs["json"]["token_endpoint_auth_method"] - == request_payload["token_endpoint_auth_method"] - ) - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses HTTPS in the redirect_uri parameter - location = response.headers["location"] - - # The redirect_uri parameter sent to the OAuth provider should use HTTPS - assert ( - "redirect_uri=https%3A%2F%2Flitellm.example.com%2Fcallback" in location - or "redirect_uri=https://litellm.example.com/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses HTTPS - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://litellm-proxy.example.com/callback" - ) - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_standard_pattern(): - """Test that oauth_protected_resource_mcp_standard returns standard MCP URL pattern (/mcp/{server_name})""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp_standard, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the standard pattern endpoint - response = await oauth_protected_resource_mcp_standard( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses standard MCP pattern: /mcp/{server_name} - assert response["resource"] == "https://litellm.example.com/mcp/test_server" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_legacy_pattern(): - """Test that oauth_protected_resource_mcp returns legacy URL pattern (/{server_name}/mcp)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the legacy pattern endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses legacy pattern: /{server_name}/mcp - assert response["resource"] == "https://litellm.example.com/test_server/mcp" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_servers"][0].startswith( - "https://litellm.example.com/" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_authorization_server_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_authorization_server_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_authorization_server_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_endpoint"].startswith("https://litellm.example.com/") - assert response["token_endpoint"].startswith("https://litellm.example.com/") - assert response["registration_endpoint"].startswith("https://litellm.example.com/") - assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_register_client_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that register_client uses X-Forwarded-Proto for redirect_uris""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://proxy.litellm.example/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - # Verify the redirect_uris use HTTPS - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy: - # Internal: http://localhost:8888/github/mcp - # External: https://proxy.example.com/github/mcp - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses the forwarded host and scheme - location = response.headers["location"] - - # The redirect_uri parameter should use the external URL - assert ( - "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" - in location - or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy without port in host - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses the external URL - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://proxy.example.com/github/mcp/callback" - ) - - -@pytest.mark.parametrize( - "base_url,x_forwarded_proto,x_forwarded_host,x_forwarded_port,expected_url", - [ - # Case 1: No forwarded headers - use original URL as-is (no trailing slash) - ( - "http://localhost:4000/", - None, - None, - None, - "http://localhost:4000", - ), - # Case 2: Only X-Forwarded-Proto - change scheme only - ( - "http://localhost:4000/", - "https", - None, - None, - "https://localhost:4000", - ), - # Case 3: X-Forwarded-Proto + X-Forwarded-Host - change scheme and host - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - None, - "https://proxy.example.com", - ), - # Case 4: X-Forwarded-Host with port included in host header - ( - "http://localhost:4000/", - "https", - "proxy.example.com:8080", - None, - "https://proxy.example.com:8080", - ), - # Case 5: X-Forwarded-Host + X-Forwarded-Port as separate headers - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - "8443", - "https://proxy.example.com:8443", - ), - # Case 6: Only X-Forwarded-Host without proto - use original scheme - ( - "http://localhost:4000/", - None, - "proxy.example.com", - None, - "http://proxy.example.com", - ), - # Case 7: Only X-Forwarded-Port without host - preserves original port if present - # (This is safer behavior - X-Forwarded-Port alone is unusual) - ( - "http://localhost:4000/", - None, - None, - "8443", - "http://localhost:4000", # Original port preserved when already present - ), - # Case 8: Complex internal URL with path (path is preserved) - ( - "http://localhost:8888/github/mcp", - "https", - "proxy.example.com", - None, - "https://proxy.example.com/github/mcp", - ), - # Case 9: IPv6 address in X-Forwarded-Host (should not treat :: as port separator) - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]", - None, - "https://[2001:db8::1]", - ), - # Case 10: IPv6 address with port - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]:8080", - None, - "https://[2001:db8::1]:8080", - ), - # Case 11: X-Forwarded-Host already has port, X-Forwarded-Port also provided (host wins) - ( - "http://localhost:4000/", - "https", - "proxy.example.com:9000", - "8443", - "https://proxy.example.com:9000", - ), - # Case 12: Standard proxy setup (most common case) - ( - "http://127.0.0.1:8888/", - "https", - "chatproxy.company.com", - None, - "https://chatproxy.company.com", - ), - # Case 13: Internal URL already has port, X-Forwarded-Port does NOT override - # (safer behavior - preserves original port when X-Forwarded-Host not provided) - ( - "http://localhost:4000/", - None, - None, - "443", - "http://localhost:4000", # Original port preserved - ), - # Case 14: Original URL with existing port in netloc, X-Forwarded-Host replaces it - ( - "http://internal.local:8888/", - "https", - "external.com", - None, - "https://external.com", - ), - ], -) -def test_get_request_base_url_comprehensive( - base_url, - x_forwarded_proto, - x_forwarded_host, - x_forwarded_port, - expected_url, - trusted_proxy_origin_headers, -): - """Comprehensive test for get_request_base_url with various header combinations""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = base_url - set_request_from_trusted_proxy(mock_request) - - # Build headers dict - headers = {} - if x_forwarded_proto: - headers["X-Forwarded-Proto"] = x_forwarded_proto - if x_forwarded_host: - headers["X-Forwarded-Host"] = x_forwarded_host - if x_forwarded_port: - headers["X-Forwarded-Port"] = x_forwarded_port - - # Mock headers.get() to return our test values - def mock_get(header_name, default=None): - return headers.get(header_name, default) - - mock_request.headers.get = mock_get - - # Test the function - result = get_request_base_url(mock_request) - - # Verify result - assert result == expected_url, ( - f"Expected '{expected_url}' but got '{result}'\n" - f"Input: base_url={base_url}, " - f"X-Forwarded-Proto={x_forwarded_proto}, " - f"X-Forwarded-Host={x_forwarded_host}, " - f"X-Forwarded-Port={x_forwarded_port}" - ) - - -def test_get_request_base_url_ignores_forwarded_headers_from_untrusted_client(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - "X-Forwarded-Port": "443", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ): - assert get_request_base_url(mock_request) == "https://gateway.example.com/mcp" - - -def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with ( - patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ), - pytest.raises(HTTPException), - ): - validate_trusted_redirect_uri( - mock_request, - "https://attacker.example.com/callback", - ) - - -def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy( - trusted_proxy_origin_headers, -): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:4000/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - validate_trusted_redirect_uri( - mock_request, - "https://proxy.example.com/callback", - ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b4d3782ba43..442bfe8a090 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2957,6 +2957,55 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection(caplog, monk assert "X-Forwarded-Host" in msg +@pytest.mark.parametrize( + "direct_ip,expect_accepted", + [ + ("10.0.0.7", True), + ("203.0.113.5", False), + ], +) +def test_validate_trusted_redirect_uri_follows_the_xff_trust_gate(direct_ip, expect_accepted, monkeypatch): + try: + from fastapi import HTTPException, Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP oauth_utils not available") + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("MCP_TRUSTED_REDIRECT_ORIGINS", raising=False) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = direct_ip + + headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + mock_request.headers.__contains__ = lambda self_, name: name in headers + + redirect_uri = "https://proxy.example.com/callback" + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + + with patch("litellm.proxy.proxy_server.general_settings", general_settings, create=True): + if expect_accepted: + validate_trusted_redirect_uri(mock_request, redirect_uri) + return + with pytest.raises(HTTPException) as exc_info: + validate_trusted_redirect_uri(mock_request, redirect_uri) + + assert exc_info.value.status_code == 400 + assert "proxy.example.com" in str(exc_info.value.detail) + + @pytest.mark.parametrize( "bad_value", [ From f9f8320972f6589dd5aac0877a1bc89c4f200028 Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 21 Aug 2026 11:47:04 -0400 Subject: [PATCH 038/319] fix(files): list unscoped managed files Read owner-scoped managed rows directly when no provider or model is supplied, avoiding an unauthenticated OpenAI fallback. Refs #35362 --- .../proxy/hooks/managed_files.py | 19 ++++-- litellm/llms/base_llm/files/transformation.py | 6 +- .../openai_files_endpoints/files_endpoints.py | 36 ++++++----- .../proxy/test_managed_files_hook.py | 33 +++++++++++ .../test_files_endpoint.py | 59 +++++++++++++++++++ 5 files changed, 130 insertions(+), 23 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index c986e835e4f..9b62284072d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1365,12 +1365,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def afile_list( self, - purpose: Optional[OpenAIFilesPurpose], + purpose: str | None, litellm_parent_otel_span: Optional[Span], + user_api_key_dict: UserAPIKeyAuth, **data: Dict, - ) -> List[OpenAIFileObject]: - """Handled in files_endpoints.py""" - return [] + ) -> Dict[str, object]: + owner_filter: Final = build_owner_filter(user_api_key_dict) + if owner_filter is None: + return build_list_page([]) + + rows: Final = await _managed_file_table(self.prisma_client).find_many(where=owner_filter) + files: Final = [ + parsed_file_object.model_copy(update={"id": row.unified_file_id}) + for row in rows + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None + and (purpose is None or parsed_file_object.purpose == purpose) + ] + return build_list_page(files) def _is_batch_polling_enabled(self) -> bool: """ diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 174be93448b..7c19326b627 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -13,7 +13,6 @@ from litellm.types.llms.openai import ( FileContentRequest, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, - OpenAIFilesPurpose, ) from litellm.types.utils import LlmProviders, ModelResponse @@ -240,10 +239,11 @@ class BaseFileEndpoints(ABC): @abstractmethod async def afile_list( self, - purpose: OpenAIFilesPurpose | None, + purpose: str | None, litellm_parent_otel_span: Span | None, + user_api_key_dict: UserAPIKeyAuth, **data: dict, - ) -> list[OpenAIFileObject]: + ) -> dict[str, object]: pass @abstractmethod diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 37cfd9d073d..a482cc54748 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1488,24 +1488,28 @@ async def list_files( or get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) or await get_custom_llm_provider_from_request_body(request=request) - or "openai" ) + managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") + if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints): + response = await managed_files_obj.afile_list( + purpose=purpose, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + user_api_key_dict=user_api_key_dict, + ) + else: + resolved_custom_llm_provider: Final = custom_llm_provider or "openai" + apply_team_provider_credentials( + data=data, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=resolved_custom_llm_provider, + ) - # No model/target_model_names pinned: resolve upstream credentials from - # the team's deployment for this provider so the call is authenticated - # against the team's own account (e.g. the team's openai deployment). - apply_team_provider_credentials( - data=data, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) - - response = await litellm.afile_list( - custom_llm_provider=custom_llm_provider, - purpose=purpose, - **data, - ) + response = await litellm.afile_list( + custom_llm_provider=resolved_custom_llm_provider, + purpose=purpose, + **data, + ) if response is None: raise HTTPException( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index fcd03e77aa2..b39d2ef8559 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -190,6 +190,39 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie assert files[0].purpose == raw_provider_object.purpose +@pytest.mark.asyncio +async def test_afile_list_returns_owner_scoped_managed_files(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object=_make_file_object("file-provider-id").model_dump(), + unified_file_id="unified-file-id", + ), + MagicMock( + file_object=_make_file_object("file-other-purpose").model_copy( + update={"purpose": "batch"} + ).model_dump(), + unified_file_id="unified-other-purpose", + ), + ] + ) + + response = await managed_files.afile_list( + purpose="batch_output", + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + ) + + managed_files.prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"created_by": "test-user"} + ) + assert [file.id for file in response["data"]] == ["unified-file-id"] + assert response["first_id"] == "unified-file-id" + assert response["last_id"] == "unified-file-id" + assert response["has_more"] is False + + @pytest.mark.asyncio async def test_parse_managed_file_object_warning_omits_rejected_values(caplog): from litellm_enterprise.proxy.hooks.managed_files import ( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bf9323cdc6a..e6101d3edd8 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -2468,6 +2468,65 @@ def test_list_files_without_target_model_names_uses_team_openai_deployment( proxy_logging_obj.post_call_failure_hook.assert_not_called() +def test_unscoped_list_files_uses_managed_file_store( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.proxy._types import LitellmUserRoles + + managed_file = OpenAIFileObject( + id="unified-file-id", + object="file", + bytes=100, + created_at=1700000000, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + managed_files = mocker.MagicMock(spec=BaseFileEndpoints) + managed_files.afile_list = mocker.AsyncMock( + return_value={ + "object": "list", + "data": [managed_file], + "first_id": managed_file.id, + "last_id": managed_file.id, + "has_more": False, + } + ) + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=None) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + provider_list = mocker.patch.object(litellm, "afile_list", new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert response.json()["data"][0]["id"] == "unified-file-id" + managed_files.afile_list.assert_awaited_once() + assert managed_files.afile_list.await_args.kwargs["user_api_key_dict"].user_id == "test-user" + provider_list.assert_not_awaited() + proxy_logging_obj.post_call_failure_hook.assert_not_called() + + def test_list_files_restricted_team_does_not_leak_global_openai_credentials( mocker: MockerFixture, monkeypatch ): From 7da34e8aed341b3368b14810d91052ebccb97117 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 21 Aug 2026 09:47:52 -0700 Subject: [PATCH 039/319] fix(proxy): make per-model budgets track spend, enforce, and report the same counter (#37736) Per-model budgets were three separate things pretending to be one. The enforcement check, the post-call increment and the info endpoints each derived their own cache key, so a budget could refuse traffic at 429 while /key/info reported zero usage, and a Bedrock model id never matched a budget keyed on the bare family name. /user/new echoed a model_max_budget back and stored an empty dict, and nothing enforced a user-scoped per-model budget at all. One owner now builds the counter key from the configured budget model, and enforcement, the increment and the info endpoints all read it. Bedrock ids resolve through the model-cost map. Auth carries the user's budget onto the token on every branch that reaches the spend hook, including JWT and auto-registration. Native passthrough attaches the three budget metadata keys its StandardLoggingUserAPIKeyMetadata does not carry, so /anthropic/... and /bedrock/... traffic is counted and capped like /v1/chat/completions. The dashboard gains the per-model budget editor it never had, on the key create, key edit and internal-user edit forms. It is read-only without an enterprise license, matching the write gate the proxy already enforces, and an untouched budget is left out of an update so an unrelated edit cannot trip that gate. The editor hydrates from either BudgetConfig spelling, since model_max_budget is a plain dict that the proxy stores exactly as the client sent it, and it carries through the fields it does not model. Without both, editing one model would drop another model row entirely and silently discard its tpm_limit and rpm_limit. /user/info refreshes its local copy of the user field by field after a save, so model_max_budget joins that list. Left out, a saved cap read back as the old one when the form was reopened, and clearing the row to recover would then wipe the value that had actually persisted. A zero-dollar cap is the strictest limit expressible, not the absence of one, so it is enforced rather than skipped on falsiness, spend exactly at the cap is refused the way every sibling budget check already refuses it, and a counter that was never written reads as zero spend rather than as unknown. The usage endpoints read every counter in one batched lookup, so a large model_max_budget cannot fan out into one concurrent cache call per configured model. Every auth path honours the same zero-cost skip flag, so none of them can refuse a free request that another serves. The custom-auth helper gains the flag it never had, which also changes its pre-existing key and end-user checks. The compaction summary gate checks the user scope alongside the key and end-user ones. This file propagates all three budgets into the summary subrequest, so enforcing only two let compaction increment a counter it could not be refused by. Custom auth attaches the user's budget to the token unconditionally, since the post-call spend hook reads it there: gating the attach on the same condition as enforcement left the counter uncharged whenever the request was not itself enforceable. An entry that will not validate is skipped rather than raised on, so one malformed scope cannot abort every other scope's increment or turn a config typo into a 500. The edit forms re-seed the budget editor when a different key or user is loaded. Its rows are seeded once and cannot re-read their own value prop, so without this a save wrote the previously loaded record's budgets onto the current one. Only the built-in provider pass-through routes carry the budget metadata. get_model_from_request deliberately resolves no model for a user-defined pass-through, since its body is forwarded verbatim and names an upstream model, so attaching there would charge a counter nothing on that route can refuse. --- .../context_management/editors/compact.py | 32 +- litellm/proxy/_types.py | 6 + litellm/proxy/auth/auth_utils.py | 4 +- litellm/proxy/auth/user_api_key_auth.py | 140 ++- .../proxy/hooks/model_max_budget_limiter.py | 584 +++++---- litellm/proxy/litellm_pre_call_utils.py | 2 + .../internal_user_endpoints.py | 21 +- .../key_management_endpoints.py | 84 +- .../pass_through_endpoints.py | 17 + ...test_unit_test_max_model_budget_limiter.py | 1054 ++++++++++++++--- .../test_user_api_key_auth.py | 670 ++++++++++- .../context_management/test_compact.py | 76 ++ .../test_internal_user_endpoints.py | 75 ++ .../test_key_management_endpoints.py | 104 +- .../test_pass_through_endpoints.py | 840 +++++-------- .../users/_components/BulkEditUsers.tsx | 3 + .../users/_components/user_edit_view.test.tsx | 121 +- .../users/_components/user_edit_view.tsx | 25 + .../user_info_view.integration.test.tsx | 44 +- .../_components/view_users/user_info_view.tsx | 8 + .../ModelMaxBudgetEditor.integration.test.tsx | 69 ++ .../ModelMaxBudgetEditor.test.ts | 140 +++ .../key_team_helpers/ModelMaxBudgetEditor.tsx | 233 ++++ .../components/key_team_helpers/key_list.tsx | 4 +- .../modelMaxBudgetPayload.test.ts | 71 ++ .../key_team_helpers/modelMaxBudgetPayload.ts | 44 + .../useModelMaxBudgetField.ts | 34 + .../key_team_helpers/useSeededState.ts | 26 + .../src/components/networking.tsx | 3 + .../organisms/createKeyPayload.test.ts | 25 + .../components/organisms/createKeyPayload.ts | 3 + .../organisms/create_key_button.tsx | 19 + .../templates/key_edit_view.test.tsx | 89 ++ .../components/templates/key_edit_view.tsx | 15 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 35 files changed, 3656 insertions(+), 1041 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.test.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/modelMaxBudgetPayload.test.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/modelMaxBudgetPayload.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/useModelMaxBudgetField.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/useSeededState.ts diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index a7c462a8fb0..2a87afb5990 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -56,7 +56,7 @@ from ..result import PolyfillResult # so the summary's spend is attributed to the same scopes. The list mirrors the # fields populated by # ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``. -# ``user_api_key_model_max_budget`` / ``user_api_key_end_user_model_max_budget`` +# The three ``*_model_max_budget`` fields # are what ``_PROXY_VirtualKeyModelMaxBudgetLimiter`` reads post-call to update # the per-model spend caches, so without them the summary spend would never # count against the caller's model budget. ``user_api_key_end_user_id`` / @@ -76,6 +76,7 @@ _PROPAGATED_METADATA_KEYS: Final = ( "user_api_key_end_user_id", "user_api_end_user_max_budget", "user_api_key_model_max_budget", + "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", "litellm_call_id", "litellm_parent_otel_span", @@ -317,10 +318,14 @@ async def _check_summary_model_budget( The summary subrequest never passes back through ``user_api_key_auth``, so without this gate a caller whose ``model_max_budget`` for ``context_management_summary_model`` is exhausted could keep consuming that - model via compaction. Mirrors the ``model_max_budget`` / - ``end_user_model_max_budget`` enforcement that ``user_api_key_auth`` runs for - the client-requested model. Returns True outside the proxy or when no + model via compaction. Mirrors the per-model budget enforcement that + ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no per-model budget is configured. + + All three scopes are checked because the summary's spend is charged to all + three: this file propagates the key, user and end-user budgets into the + subrequest's metadata, so enforcing only two of them would let compaction + increment a counter it can never be refused by. """ if user_api_key_auth is None: return True @@ -347,6 +352,25 @@ async def _check_summary_model_budget( ) return False + user_model_max_budget: Final = getattr(user_api_key_auth, "user_model_max_budget", None) + user_id: Final = getattr(user_api_key_auth, "user_id", None) + if isinstance(user_model_max_budget, dict) and user_model_max_budget and user_id is not None: + try: + await model_max_budget_limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model=summary_model, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the key and end-user scopes do + verbose_logger.warning( + "compact_20260112: unexpected error during user model-budget check for summary_model=%s; denying: %s", + summary_model, + e, + ) + return False + end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00cbd13cfdc..e51a3138d2d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2805,6 +2805,10 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob user_email: str | None = None user_spend: float | None = None user_max_budget: float | None = None + # Values stay `object` rather than BudgetConfig: this is the raw JSON column, + # and validating it here would make one malformed row fail auth outright. + # resolve_model_budget validates the single entry a request actually needs. + user_model_max_budget: dict[str, object] | None = None request_route: str | None = None is_session_token: bool = False # Server-only marker set exclusively by the MCP gateway admission path @@ -2982,6 +2986,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): sso_user_id: str | None = None teams: list[str] = [] # Just team IDs, not full team objects object_permission: LiteLLM_ObjectPermissionTable | None = None + model_max_budget: dict | None = None + model_max_budget_usage: dict | None = None from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ce662ee0374..d04a71535ef 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1801,7 +1801,7 @@ def _format_model_candidates( return candidates -def _request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool: +def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool: """Whether FastAPI resolved this request to a user-defined pass-through handler. Reads the marker set by ``create_pass_through_route`` off the dispatched endpoint @@ -1842,7 +1842,7 @@ def get_model_from_request( and does not carry the marker. Built-in provider passthrough routes (``/vertex_ai``, ``/gemini``, ...) are separate handlers and keep model enforcement. """ - if _request_dispatched_to_pass_through_endpoint(request): + if request_dispatched_to_pass_through_endpoint(request): return None candidates: Final = _extract_model_candidates_from_request( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99592d44f9b..fe4f1ee4ae5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -11,6 +11,7 @@ import asyncio import fnmatch import re import secrets +from collections.abc import Mapping from datetime import datetime, timezone from typing import Any, Final, NamedTuple, Protocol, Union, cast @@ -186,6 +187,62 @@ class _KeyModelBudgetLimiter(Protocol): async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ... +class _UserModelBudgetLimiter(Protocol): + async def is_user_within_model_budget( + self, user_id: str, user_model_max_budget: Mapping[str, object], model: str + ) -> bool: ... + + +async def _read_user_model_max_budget( + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: object, + proxy_logging_obj: ProxyLogging, +) -> dict | None: + """The user row's `model_max_budget`, or None when the row cannot be read. + + A user whose row is missing must not be refused: this is a budget lookup, + and the main auth path likewise treats an unreadable user as no user. + """ + if user_id is None or prisma_client is None: + return None + try: + user_obj: Final = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, # pyright: ignore[reportArgumentType] # Span is a runtime union, not usable in an annotation here + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # mirrors the main path's tolerance + verbose_logger.debug("Unable to read user for the per-model budget check: %s", e) + return None + return getattr(user_obj, "model_max_budget", None) + + +async def _check_user_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _UserModelBudgetLimiter, + models: list[str], +) -> None: + """Enforce the internal user's own `model_max_budget` across the request's models. + + Separate from the key check: a user's per-model budget caps every key they + own, so a caller cannot escape it by minting another key. + """ + user_model_max_budget: Final = valid_token.user_model_max_budget + if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_user_within_model_budget( + user_id=valid_token.user_id, + user_model_max_budget=user_model_max_budget, + model=model_name, + ) + + async def _check_key_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _KeyModelBudgetLimiter, @@ -1390,6 +1447,7 @@ async def _user_api_key_auth_builder( end_user_id=end_user_id, user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), + user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), team_member_rpm_limit=( team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None ), @@ -1427,6 +1485,13 @@ async def _user_api_key_auth_builder( if auto_registered is not None: auto_registered.jwt_claims = jwt_claims auto_registered.user_email = user_email + # The auto-registered token is built from the new key's + # columns, which carry no user budget. Carry over the + # already-loaded user row rather than re-reading it, or + # the budget check below has nothing to enforce. + auto_registered.user_model_max_budget = ( + user_object.model_max_budget if user_object is not None else None + ) valid_token = auto_registered api_key = valid_token.token or "" @@ -1458,6 +1523,28 @@ async def _user_api_key_auth_builder( valid_token.project_metadata = _jwt_project_obj.metadata valid_token.project_alias = _jwt_project_obj.project_alias + # JWT auth returns here rather than falling through to the + # virtual-key checks below, so the user's per-model budget + # has to be enforced on this path too. Without it the + # post-call increment still charges the counter and nothing + # ever reads it, which is worse than not tracking at all. + # Guarded by the same flag the virtual-key path uses, or a + # zero-cost model would be refused here and allowed there, + # while the log above claims all budget checks were skipped. + if not skip_budget_checks: + await _check_user_model_budget( + valid_token=cast(UserAPIKeyAuth, valid_token), + model_max_budget_limiter=model_max_budget_limiter, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + ) + ), + ) + return cast(UserAPIKeyAuth, valid_token) #### ELSE #### @@ -1811,6 +1898,12 @@ async def _user_api_key_auth_builder( ) user_obj = None + if user_obj is not None: + # The joint verification-token view carries the key's columns only, so the + # user's own per-model budget reaches enforcement and the post-call + # increment through the row fetched here. + valid_token.user_model_max_budget = user_obj.model_max_budget + if ( user_obj is not None and isinstance(user_obj.metadata, dict) @@ -1974,6 +2067,14 @@ async def _user_api_key_auth_builder( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # Check 5a. Internal user model_max_budget + if current_models: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=current_models, + ) + # Check 5b. End-user model max budget end_user_mmb: Final = valid_token.end_user_model_max_budget if ( @@ -2757,6 +2858,7 @@ async def _return_user_api_key_auth_obj( user_email=user_obj.user_email, user_spend=getattr(user_obj, "spend", None), user_max_budget=getattr(user_obj, "max_budget", None), + user_model_max_budget=getattr(user_obj, "model_max_budget", None), ) if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj): user_api_key_kwargs.update( @@ -3020,10 +3122,21 @@ async def _run_post_custom_auth_checks( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # A zero-cost model cannot move any counter, so refusing it means refusing on + # spend some other model accrued. The JWT and virtual-key paths already skip + # every budget check for these; this path did not, so the same request could + # be refused under custom auth and served under the other two. + skip_budget_checks: Final = ( + _is_model_cost_zero(model=current_model, llm_router=llm_router) + if current_model is not None and llm_router is not None + else False + ) + # 3. Check key-level model_max_budget max_budget_per_model: Final = valid_token.model_max_budget if ( - max_budget_per_model is not None + not skip_budget_checks + and max_budget_per_model is not None and isinstance(max_budget_per_model, dict) and len(max_budget_per_model) > 0 and current_models @@ -3050,10 +3163,33 @@ async def _run_post_custom_auth_checks( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # 3b. Attach and check the internal user's model_max_budget. + # Custom auth builds its own token, so unlike the main path nothing has + # loaded the user row yet. The attach is unconditional because the post-call + # spend hook reads this field off the token: gating it on the same condition + # as enforcement would leave the user's counter uncharged whenever this + # request was not itself enforceable, which is the untracked-spend bug this + # PR exists to fix. + user_budget: Final = await _read_user_model_max_budget( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + valid_token.user_model_max_budget = user_budget # rebind-ok: the spend hook reads it off this token + if not skip_budget_checks and current_models: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=current_models, + ) + # 4. Check end-user model_max_budget end_user_mmb: Final = valid_token.end_user_model_max_budget if ( - end_user_mmb is not None + not skip_budget_checks + and end_user_mmb is not None and isinstance(end_user_mmb, dict) and len(end_user_mmb) > 0 and current_models diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 215969ef899..c5d10b2749b 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,21 +1,253 @@ import json +import time +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ( - BudgetConfig, - GenericBudgetConfigType, - StandardLoggingPayload, -) +from litellm.types.utils import BudgetConfig, StandardLoggingPayload VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend" +USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend" + +_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( + { + Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX, + } +) + +_LEGACY_REQUEST_MODEL_SCOPES: Final = frozenset({Litellm_EntityType.KEY, Litellm_EntityType.END_USER}) + +_PROCESS_STARTED_AT: Final = time.monotonic() + +_BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( + { + Litellm_EntityType.KEY: "virtual_key_budget_start_time", + Litellm_EntityType.USER: "user_model_budget_start_time", + Litellm_EntityType.END_USER: "end_user_budget_start_time", + } +) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelBudget: + """The `model_max_budget` entry a request resolved to. + + ``budget_model`` is the key as the operator configured it, not the model + name on the request. Every counter is keyed on it so enforcement, the + post-call increment and the `/key/info` + `/user/info` usage reads cannot + disagree about which counter a request belongs to. + """ + + budget_model: str + budget_config: BudgetConfig + + +def model_budget_spend_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + budget_model: str, + budget_duration: str | None, +) -> str: + """Sole owner of the per-model spend counter key, shared by its writer and all of its readers.""" + return f"{_SPEND_CACHE_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" + + +def _legacy_request_model_spend_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + model: str, + resolved: ResolvedModelBudget, +) -> str | None: + """The counter this request was billed to before the budget model owned the key, or None. + + Upgrading proxies carry live counters keyed on the model as REQUESTED + (`openai/gpt-4`) rather than as configured (`gpt-4`), and those were the + counters the previous version enforced on. Nothing writes that spelling once + this version is running, so the pre-upgrade and post-upgrade counters hold + disjoint halves of one window and adding them is the window's real spend. + + Only the key and end-user scopes ever had one. The user scope is introduced + by this change, so it has no counter to carry. + + The carry stops one budget window after start-up, because a legacy counter + belongs to a window that was already open when this process replaced the one + writing it. Past that point the lookup could only ever miss. + """ + budget_duration: Final = resolved.budget_config.budget_duration + if entity_type not in _LEGACY_REQUEST_MODEL_SCOPES or budget_duration is None: + return None + if time.monotonic() - _PROCESS_STARTED_AT >= duration_in_seconds(budget_duration): + return None + return model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=model, + budget_duration=budget_duration, + ) + + +def model_budget_start_time_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + budget_model: str, + budget_duration: str | None, +) -> str: + """Window start for one (entity, budget model) pair. + + Scoped per budget model because an entity may budget two models over + different periods, and a shared start time lets the shorter period restart + the longer one's window. + """ + return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" + + +def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None: + """Find the `model_max_budget` entry that governs `model`, or None.""" + for candidate in _budget_model_candidates(model): + raw_budget_config = model_max_budget.get(candidate) + if raw_budget_config is None: + continue + if (budget_config := _usable_budget_config(raw_budget_config)) is None: + # An entry that will not validate cannot be keyed, so it cannot be + # enforced or incremented. Skip to the next candidate rather than + # raising: raising would abort every other scope's increment and turn + # a config typo into a 500, and stopping here would let one malformed + # specific entry disable a perfectly good bare-family budget beside + # it. The candidate chain already falls through an ABSENT entry, and + # an unparseable one is indistinguishable from absent to enforcement. + # `validate_model_max_budget` rejects these on the write path, so + # reaching here means config.yaml or a direct DB edit. + verbose_proxy_logger.warning( + "Ignoring unusable model_max_budget entry for %s; it cannot be enforced or tracked", + candidate, + ) + continue + return ResolvedModelBudget(budget_model=candidate, budget_config=budget_config) + return None + + +def _budget_model_candidates(model: str) -> tuple[str, ...]: + """Names a budget may be configured under for a request on `model`, most specific first. + + Beyond the model as sent, a budget may be keyed on the model without its + ``{custom_llm_provider}/`` prefix (``gpt-4o`` governs ``openai/gpt-4o``), on + the Bedrock base model (``anthropic.claude-opus-4-8`` governs the + cross-region ``us.anthropic.claude-opus-4-8``), or on the bare family name + that Bedrock id shares with its direct-provider twin (``claude-opus-4-8``). + """ + return tuple(dict.fromkeys((model, model.split("/")[-1], *_bedrock_candidates(model)))) + + +def _bedrock_candidates(model: str) -> tuple[str, ...]: + """Bedrock-only candidates, empty unless litellm prices `model` as a Bedrock model. + + Gating on the cost map rather than on a vendor allowlist is what makes + splitting the leading dotted segment safe: most dotted model ids are not + Bedrock ids at all (``azure/gpt-4.1``, ``gpt-image-1.5``), and splitting one + of those would produce a garbage candidate. + """ + base_model: Final = get_bedrock_base_model(model) + cost_entry: Final = litellm.model_cost.get(base_model) + if not isinstance(cost_entry, dict) or not str(cost_entry.get("litellm_provider", "")).startswith("bedrock"): + return () + _, _, without_vendor = base_model.partition(".") + return (base_model, without_vendor) if without_vendor else (base_model,) + + +async def build_model_max_budget_usage( + entity_type: Litellm_EntityType, + entity_id: str | None, + model_max_budget: Mapping[str, object] | None, + cache: DualCache | None, +) -> dict[str, dict[str, object]]: + """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it. + + `cache` must be the DualCache the limiter writes the counters to; callers + read it off the limiter rather than re-deriving it, so a scope that is being + blocked can never report zero usage. + """ + if cache is None or entity_id is None or not model_max_budget: + return {} + + budgets: Final = tuple( + (budget_model, budget_config) + for budget_model, raw_budget_config in model_max_budget.items() + for budget_config in (_usable_budget_config(raw_budget_config),) + if budget_config is not None + ) + if not budgets: + return {} + spend_keys: Final = tuple( + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ) + batched: Final = await cache.async_batch_get_cache( + keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here + ) + # async_batch_get_cache returns None if it fails internally, and its result is + # index-aligned with `keys` otherwise. An unusable result reads as a miss, + # which is what a never-written counter already reads as. + current_spends: Final = ( + tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets) + ) + return { + budget_model: { + "current_spend": round(_as_spend(current_spend), 4), + "budget_limit": budget_config.max_budget, + "time_period": budget_config.budget_duration, + } + for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True) + } + + +def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: + try: + budget_config: Final = BudgetConfig.model_validate(raw_budget_config) + if budget_config.budget_duration is None: + return None + duration_in_seconds(budget_config.budget_duration) + except Exception: # noqa: BLE001 # a malformed entry must not fail the whole report + return None + return budget_config + + +def _as_spend(current_spend: object) -> float: + try: + return float(current_spend or 0.0) # pyright: ignore[reportArgumentType] # non-numeric falls to the except + except (TypeError, ValueError): + return 0.0 + + +def _resolve_entity_model_budgets( + model: str, + entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], +) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]: + """Drop the scopes that do not budget `model`, keeping only what can be incremented.""" + return tuple( + (entity_type, entity_id, resolved) + for entity_type, entity_id, model_max_budget in entity_budgets + if entity_id is not None and isinstance(model_max_budget, Mapping) and model_max_budget + for resolved in (resolve_model_budget(model=model, model_max_budget=model_max_budget),) + if resolved is not None and resolved.budget_config.budget_duration is not None + ) class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): @@ -41,47 +273,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Raises: BudgetExceededError: If the user_api_key_dict has exceeded the model budget """ - _model_max_budget: Final = user_api_key_dict.model_max_budget - internal_model_max_budget: Final[GenericBudgetConfigType] = {} - - for _model, _budget_info in _model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - - verbose_proxy_logger.debug( - "internal_model_max_budget %s", - json.dumps(internal_model_max_budget, indent=4, default=str), + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.KEY, + entity_id=user_api_key_dict.token, + model_max_budget=user_api_key_dict.model_max_budget, + model=model, + exceeded_message=( + f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, " + f"exceeded budget for model={model}" + ), ) - # check if current model is in internal_model_max_budget - _current_model_budget_info: Final = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if _current_model_budget_info is None: - verbose_proxy_logger.debug("Model %s not found in internal_model_max_budget", model) - return True - - # check if current model is within budget - if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: - _current_spend: Final = await self._get_virtual_key_spend_for_model( - user_api_key_hash=user_api_key_dict.token, - model=model, - key_budget_config=_current_model_budget_info, - ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): - raise litellm.BudgetExceededError( - message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", - current_cost=_current_spend, - max_budget=_current_model_budget_info.max_budget, - entity_type=Litellm_EntityType.KEY.value, - entity_id=user_api_key_dict.token, - ) - - return True - async def get_fallback_model_within_budget( self, user_api_key_dict: UserAPIKeyAuth, @@ -96,10 +298,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): continue return None + async def is_user_within_model_budget( + self, + user_id: str, + user_model_max_budget: Mapping[str, object], + model: str, + ) -> bool: + """ + Check if the internal user is within the model budget + + Raises: + BudgetExceededError: If the user has exceeded the model budget + """ + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + model=model, + exceeded_message=f"LiteLLM User: {user_id}, exceeded budget for model={model}", + ) + async def is_end_user_within_model_budget( self, end_user_id: str, - end_user_model_max_budget: dict, + end_user_model_max_budget: Mapping[str, object], model: str, ) -> bool: """ @@ -108,116 +330,81 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Raises: BudgetExceededError: If the end_user has exceeded the model budget """ - internal_model_max_budget: Final[GenericBudgetConfigType] = {} - - for _model, _budget_info in end_user_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - - verbose_proxy_logger.debug( - "end_user internal_model_max_budget %s", - json.dumps(internal_model_max_budget, indent=4, default=str), + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.END_USER, + entity_id=end_user_id, + model_max_budget=end_user_model_max_budget, + model=model, + exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", ) - # check if current model is in internal_model_max_budget - _current_model_budget_info: Final = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if _current_model_budget_info is None: - verbose_proxy_logger.debug("Model %s not found in end_user_model_max_budget", model) + async def _is_entity_within_model_budget( + self, + entity_type: Litellm_EntityType, + entity_id: str | None, + model_max_budget: Mapping[str, object] | None, + model: str, + exceeded_message: str, + ) -> bool: + if not model_max_budget: + return True + resolved: Final = resolve_model_budget(model=model, model_max_budget=model_max_budget) + if resolved is None: + verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) return True - # check if current model is within budget - if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: - _current_spend: Final = await self._get_end_user_spend_for_model( - end_user_id=end_user_id, - model=model, - key_budget_config=_current_model_budget_info, - ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): - raise litellm.BudgetExceededError( - message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", - current_cost=_current_spend, - max_budget=_current_model_budget_info.max_budget, - entity_type=Litellm_EntityType.END_USER.value, - entity_id=end_user_id, - ) + max_budget: Final = resolved.budget_config.max_budget + if max_budget is None or max_budget < 0: + return True + current_spend: Final = await self._get_spend_for_model_budget( + entity_type=entity_type, + entity_id=entity_id, + model=model, + resolved=resolved, + ) + if current_spend >= max_budget: + raise litellm.BudgetExceededError( + message=exceeded_message, + current_cost=current_spend, + max_budget=max_budget, + entity_type=entity_type.value, + entity_id=entity_id, + ) return True - async def _get_end_user_spend_for_model( + async def _get_spend_for_model_budget( self, - end_user_id: str, + entity_type: Litellm_EntityType, + entity_id: str | None, model: str, - key_budget_config: BudgetConfig, - ) -> float | None: - # 1. model: directly look up `model` - end_user_model_spend_cache_key = ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" - ) - _current_spend = await self.dual_cache.async_get_cache( - key=end_user_model_spend_cache_key, - ) + resolved: ResolvedModelBudget, + ) -> float: + """Spend charged to this budget in the current window, legacy counter included. - if _current_spend is None: - # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=end_user_model_spend_cache_key, - ) - return _current_spend - - async def _get_virtual_key_spend_for_model( - self, - user_api_key_hash: str | None, - model: str, - key_budget_config: BudgetConfig, - ) -> float | None: + A counter that was never written is zero spend, not unknown spend. The + distinction only shows up at a zero-dollar cap, where skipping the + comparison would let the strictest possible limit admit every request. """ - Get the current spend for a virtual key for a model - - Lookup model in this order: - 1. model: directly look up `model` - 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - """ - - # 1. model: directly look up `model` - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" + spend_key: Final = model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, ) - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, + legacy_spend_key: Final = _legacy_request_model_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + model=model, + resolved=resolved, ) + current_spend: Final = _as_spend(await self._cached_spend(spend_key)) + if legacy_spend_key is None or legacy_spend_key == spend_key: + return current_spend + return current_spend + _as_spend(await self._cached_spend(legacy_spend_key)) - if _current_spend is None: - # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - # if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview - virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - return _current_spend - - def _get_request_model_budget_config( - self, model: str, internal_model_max_budget: GenericBudgetConfigType - ) -> BudgetConfig | None: - """ - Get the budget config for the request model - - 1. Check if `model` is in `internal_model_max_budget` - 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` - """ - return internal_model_max_budget.get(model, None) or internal_model_max_budget.get( - self._get_model_without_custom_llm_provider(model), None - ) - - def _get_model_without_custom_llm_provider(self, model: str) -> str: - if "/" in model: - return model.split("/")[-1] - return model + async def _cached_spend(self, spend_key: str) -> float | None: + return await self.dual_cache.async_get_cache(key=spend_key) async def async_filter_deployments( self, @@ -245,80 +432,63 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): _litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {} _metadata: Final[dict] = _litellm_params.get("metadata", {}) or {} - user_api_key_model_max_budget: Final[dict | None] = _metadata.get("user_api_key_model_max_budget", None) - user_api_key_end_user_model_max_budget: Final[dict | None] = _metadata.get( - "user_api_key_end_user_model_max_budget", None - ) - if (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) and ( - user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0 - ): - verbose_proxy_logger.debug( - "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty." - ) - return + payload_metadata: Final = standard_logging_payload.get("metadata") or {} - response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) # Use model_group (the user-facing model alias, e.g. "gpt-4o") when - # available. The enforcement path (is_key_within_model_budget) receives - # the model name from request_data["model"] which is the model group - # alias, so the spend tracking cache key must use the same name. - # Falling back to the deployment-level "model" field preserves - # behaviour for non-proxy or non-router deployments where model_group - # is None. + # available. The enforcement path receives the model name from + # request_data["model"] which is the model group alias, so the spend + # tracking cache key must resolve from the same name. Falling back to + # the deployment-level "model" field preserves behaviour for non-proxy + # or non-router deployments where model_group is None. model: Final = standard_logging_payload.get("model_group") or standard_logging_payload.get("model") - virtual_key: Final = standard_logging_payload.get("metadata", {}).get("user_api_key_hash") - end_user_id = standard_logging_payload.get("end_user") or standard_logging_payload.get("metadata", {}).get( - "user_api_key_end_user_id" - ) - if model is None: return - if ( - virtual_key is not None - and user_api_key_model_max_budget is not None - and len(user_api_key_model_max_budget) > 0 - ): - internal_model_max_budget: GenericBudgetConfigType = {} - for _model, _budget_info in user_api_key_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if key_budget_config is not None and key_budget_config.budget_duration: - virtual_spend_key: Final = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" - ) - virtual_start_time_key: Final = f"virtual_key_budget_start_time:{virtual_key}" - await self._increment_spend_for_key( - budget_config=key_budget_config, - spend_key=virtual_spend_key, - start_time_key=virtual_start_time_key, - response_cost=response_cost, - ) + response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + entity_budgets: Final = ( + ( + Litellm_EntityType.KEY, + payload_metadata.get("user_api_key_hash"), + _metadata.get("user_api_key_model_max_budget"), + ), + ( + Litellm_EntityType.USER, + payload_metadata.get("user_api_key_user_id"), + _metadata.get("user_api_key_user_model_max_budget"), + ), + ( + Litellm_EntityType.END_USER, + standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"), + _metadata.get("user_api_key_end_user_model_max_budget"), + ), + ) - if ( - end_user_id is not None - and user_api_key_end_user_model_max_budget is not None - and len(user_api_key_end_user_model_max_budget) > 0 - ): - internal_model_max_budget: GenericBudgetConfigType = {} - for _model, _budget_info in user_api_key_end_user_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget + resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets) + if not resolved_budgets: + verbose_proxy_logger.debug( + "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " + "no key, user or end-user model_max_budget covers model=%s", + model, + ) + return + + for entity_type, entity_id, resolved in resolved_budgets: + await self._increment_spend_for_key( + budget_config=resolved.budget_config, + spend_key=model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, + ), + start_time_key=model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, + ), + response_cost=response_cost, ) - if key_budget_config is not None and key_budget_config.budget_duration: - end_user_spend_key: Final = ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" - ) - end_user_start_time_key: Final = f"end_user_budget_start_time:{end_user_id}" - await self._increment_spend_for_key( - budget_config=key_budget_config, - spend_key=end_user_spend_key, - start_time_key=end_user_start_time_key, - response_cost=response_cost, - ) if self.dual_cache.redis_cache is not None: await self._push_in_memory_increments_to_redis() diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2ec5c34958c..1541b8acfdc 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1943,6 +1943,8 @@ async def add_litellm_data_to_request( # Follow same pattern as team and API key budgets data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget + user_model_budget: Final = user_api_key_dict.user_model_max_budget + data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9c725c54d08..c2f5b8eeb8b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, user_object_permission_id_cache_key, ) +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( DailySpendRecord, @@ -817,6 +818,7 @@ def _build_user_info_response( keys: list[LiteLLM_VerificationToken] | None, team_list: list[TeamListResponseObject], teams_1: list[TeamListResponseObject] | None, + model_max_budget_usage: dict[str, dict[str, object]] | None = None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -830,6 +832,8 @@ def _build_user_info_response( if isinstance(_user_info, dict): _user_info.pop("password", None) _user_info["metadata"] = _redact_scim_enterprise_metadata(_user_info.get("metadata")) + if model_max_budget_usage is not None: + _user_info["model_max_budget_usage"] = model_max_budget_usage return UserInfoResponse( user_id=user_id, @@ -864,7 +868,7 @@ async def user_info( --header 'Authorization: Bearer sk-1234' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: user_id = _normalize_user_info_user_id(request=request, user_id=user_id) @@ -910,6 +914,12 @@ async def user_info( keys=keys, team_list=team_list, teams_1=teams_1, + model_max_budget_usage=await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=getattr(user_info, "model_max_budget", None), + cache=model_max_budget_limiter.dual_cache, + ), ) return response_data @@ -1007,7 +1017,7 @@ async def user_info_v2( --header 'Authorization: Bearer sk-1234' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: if prisma_client is None: @@ -1062,6 +1072,13 @@ async def user_info_v2( sso_user_id=user_data.get("sso_user_id"), teams=user_data.get("teams") or [], object_permission=user_data.get("object_permission"), + model_max_budget=user_data.get("model_max_budget"), + model_max_budget_usage=await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_data.get("user_id", user_id), + model_max_budget=user_data.get("model_max_budget"), + cache=model_max_budget_limiter.dual_cache, + ), ) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_info_v2(): Exception occured - %s", e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 71218d6114b..bf42aeeec05 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -29,6 +29,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.caching.dual_cache import DualCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -47,7 +48,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s rotate_sso_identity_assertions_master_key, ) from litellm.proxy._types import * -from litellm.proxy._types import LiteLLM_VerificationToken, hash_token +from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken, hash_token from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, can_team_access_model, @@ -73,9 +74,7 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks -from litellm.proxy.hooks.model_max_budget_limiter import ( - VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, -) +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, @@ -3511,62 +3510,17 @@ async def delete_key_fn( raise handle_exception_on_proxy(e) -async def _get_model_max_budget_current_spend( - api_key_hash: str, - model: str, - budget_config: BudgetConfig, - user_api_key_cache: UserApiKeyCache, -) -> float: - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" - ) - current_spend: float | None = await user_api_key_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - if current_spend is None: - model_without_prefix: Final = model.split("/")[-1] if "/" in model else model - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" - f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}" - ) - current_spend = await user_api_key_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - try: - return float(current_spend or 0.0) - except (TypeError, ValueError): - return 0.0 - - async def _build_model_max_budget_usage( api_key_hash: str, model_max_budget: Mapping[str, Mapping[str, object]], - user_api_key_cache: UserApiKeyCache | None, + user_api_key_cache: DualCache | None, ) -> dict[str, dict[str, object]]: - if user_api_key_cache is None or not model_max_budget: - return {} - - result: Final[dict[str, dict[str, object]]] = {} - for model, budget_info in model_max_budget.items(): - try: - budget_config = BudgetConfig.model_validate(budget_info) - if budget_config.budget_duration is None: - continue - duration_in_seconds(budget_config.budget_duration) - except Exception: # noqa: BLE001 - continue - spend = await _get_model_max_budget_current_spend( - api_key_hash=api_key_hash, - model=model, - budget_config=budget_config, - user_api_key_cache=user_api_key_cache, - ) - result[model] = { - "current_spend": round(spend, 4), - "budget_limit": budget_config.max_budget, - "time_period": budget_config.budget_duration, - } - return result + return await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=api_key_hash, + model_max_budget=model_max_budget, + cache=user_api_key_cache, + ) @router.post( @@ -3596,7 +3550,10 @@ async def info_key_fn_v2( -d {"keys": ["sk-1", "sk-2", "sk-3"]} ``` """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + model_max_budget_limiter, + prisma_client, + ) try: if prisma_client is None: @@ -3648,7 +3605,7 @@ async def info_key_fn_v2( k_dict["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=k_token_hash, model_max_budget=model_max_budget, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=model_max_budget_limiter.dual_cache, ) filtered_key_info.append(k_dict) @@ -3707,7 +3664,10 @@ async def info_key_fn( -H "Authorization: Bearer sk-test-example-key-123" ``` """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + model_max_budget_limiter, + prisma_client, + ) try: if prisma_client is None: @@ -3760,7 +3720,7 @@ async def info_key_fn( key_info["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=key_token_hash, model_max_budget=model_max_budget, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=model_max_budget_limiter.dual_cache, ) # Attach object_permission if object_permission_id is set @@ -3953,6 +3913,10 @@ async def generate_key_helper_fn( } if teams is not None: user_data["teams"] = teams + if model_max_budget: + # Only when supplied: the SSO and default-key callers reach this with the + # empty default, and writing that would clear an existing user's budgets. + user_data["model_max_budget"] = model_max_budget_json key_data: Final = { "token": token, "key_alias": key_alias, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d45421489e7..1915a853983 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -64,6 +64,7 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, @@ -568,6 +569,22 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key + # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this + # the post-call increment finds nothing and every passthrough request goes untracked and + # unenforced. Set after the client merge so a request body cannot supply its own budget. + # + # Only for the built-in provider routes. `get_model_from_request` returns + # None for a user-defined pass-through, deliberately: its body is forwarded + # verbatim, so `model` there names an UPSTREAM model rather than a + # LiteLLM-managed one. Enforcement is therefore skipped on those routes, and + # charging a counter anyway would track spend that nothing can refuse, and + # would attribute it to a budget the operator scoped to a LiteLLM model that + # merely shares the name. + if not request_dispatched_to_pass_through_endpoint(request): + _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget + _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget + _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget _metadata.update( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 55459721906..fcdb3c9246c 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -2,16 +2,21 @@ import os import sys from unittest.mock import AsyncMock, patch -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system-path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system-path import pytest import litellm from litellm.caching.caching import DualCache +from datetime import datetime, timezone + +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_max_budget_limiter import ( + _budget_model_candidates, _PROXY_VirtualKeyModelMaxBudgetLimiter, + build_model_max_budget_usage, + resolve_model_budget, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -24,41 +29,95 @@ def budget_limiter(): return _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) -# Test _get_model_without_custom_llm_provider -def test_get_model_without_custom_llm_provider(budget_limiter): +# Test _budget_model_candidates +def test_budget_model_candidates(): # Test with custom provider - assert ( - budget_limiter._get_model_without_custom_llm_provider("openai/gpt-4") == "gpt-4" - ) + assert _budget_model_candidates("openai/gpt-4") == ("openai/gpt-4", "gpt-4") - # Test without custom provider - assert budget_limiter._get_model_without_custom_llm_provider("gpt-4") == "gpt-4" + # Test without custom provider: no duplicate candidate + assert _budget_model_candidates("gpt-4") == ("gpt-4",) -# Test _get_request_model_budget_config -def test_get_request_model_budget_config(budget_limiter): - internal_budget = { - "gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d"), - "claude-3": GenericBudgetInfo(budget_limit=50.0, time_period="1d"), +@pytest.mark.parametrize( + "model,expected", + [ + ( + "bedrock/anthropic.claude-opus-4-8", + ( + "bedrock/anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-8", + "claude-opus-4-8", + ), + ), + ( + "us.anthropic.claude-opus-4-8", + ( + "us.anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-8", + "claude-opus-4-8", + ), + ), + ( + "bedrock/converse/us.amazon.nova-pro-v1:0", + ( + "bedrock/converse/us.amazon.nova-pro-v1:0", + "us.amazon.nova-pro-v1:0", + "amazon.nova-pro-v1:0", + "nova-pro-v1:0", + ), + ), + ], +) +def test_budget_model_candidates_reach_the_bedrock_family_name(model, expected): + """ + Bedrock ids carry a dotted vendor segment ("anthropic.", "amazon.") on top of + the optional cross-region prefix, so a budget configured under the bare + family name would otherwise never match Bedrock traffic: no enforcement and + no spend tracking at all. + """ + assert _budget_model_candidates(model) == expected + + +@pytest.mark.parametrize( + "model", + [ + "azure/gpt-4.1", + "gpt-image-1.5", + "not-a-real-model.with.dots", + "ft:gpt-4o:acme::abc", + ], +) +def test_budget_model_candidates_never_split_a_non_bedrock_dotted_name(model): + """ + Most dotted model ids are versions, not Bedrock vendor prefixes. Splitting one + would offer a garbage candidate ("gpt-4.1" -> "1") that could collide with an + unrelated budget entry, so the split is gated on litellm pricing the model as + a Bedrock model. + """ + for candidate in _budget_model_candidates(model): + assert candidate in (model, model.split("/")[-1]) + + +# Test resolve_model_budget +def test_resolve_model_budget(): + model_max_budget = { + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 50.0, "time_period": "1d"}, } # Test direct model match - config = budget_limiter._get_request_model_budget_config( - model="gpt-4", internal_model_max_budget=internal_budget - ) - assert config.max_budget == 100.0 + resolved = resolve_model_budget(model="gpt-4", model_max_budget=model_max_budget) + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 100.0 - # Test model with provider - config = budget_limiter._get_request_model_budget_config( - model="openai/gpt-4", internal_model_max_budget=internal_budget - ) - assert config.max_budget == 100.0 + # Test model with provider: the counter is keyed on the CONFIGURED name, + # not the request name, so every reader looks it up the same way. + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 100.0 # Test non-existent model - config = budget_limiter._get_request_model_budget_config( - model="non-existent", internal_model_max_budget=internal_budget - ) - assert config is None + assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None # Test is_key_within_model_budget @@ -72,47 +131,47 @@ async def test_is_key_within_model_budget(budget_limiter): ) # Test when model is within budget - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=50.0 - ): - assert ( - await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") - is True - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): + assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True # Test when model exceeds budget - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") # Test model not in budget config - assert ( - await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") - is True + assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True + + +# Test _get_spend_for_model_budget +@pytest.mark.asyncio +async def test_get_spend_for_model_budget_reads_the_configured_model_key( + budget_limiter, +): + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, ) + model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + # openai/gpt-4 resolves to the configured "gpt-4" entry, so the lookup must + # hit the same key async_log_success_event writes. + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) -# Test _get_virtual_key_spend_for_model -@pytest.mark.asyncio -async def test_get_virtual_key_spend_for_model(budget_limiter): - budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d") + async def _spend(key): + return 50.0 if key == f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d" else None - # Mock cache get - with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0): - spend = await budget_limiter._get_virtual_key_spend_for_model( - user_api_key_hash="test-key", model="gpt-4", key_budget_config=budget_config - ) - assert spend == 50.0 - - # Test with provider prefix - spend = await budget_limiter._get_virtual_key_spend_for_model( - user_api_key_hash="test-key", + with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get: + spend = await budget_limiter._get_spend_for_model_budget( + entity_type=Litellm_EntityType.KEY, + entity_id="test-key", model="openai/gpt-4", - key_budget_config=budget_config, + resolved=resolved, ) assert spend == 50.0 + assert [call.kwargs["key"] for call in mock_get.call_args_list] == [ + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d", + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:openai/gpt-4:1d", + ] @pytest.mark.asyncio @@ -138,9 +197,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim "metadata": {"user_api_key_hash": virtual_key}, }, "litellm_params": { - "metadata": { - "user_api_key_model_max_budget": user_api_key_model_max_budget - }, + "metadata": {"user_api_key_model_max_budget": user_api_key_model_max_budget}, }, } with patch.object( @@ -148,15 +205,11 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}") assert call_kwargs["response_cost"] == 0.05 @@ -164,9 +217,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim @pytest.mark.asyncio async def test_is_end_user_within_model_budget(budget_limiter): # Test when model is within budget - with patch.object( - budget_limiter, "_get_end_user_spend_for_model", return_value=50.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): assert ( await budget_limiter.is_end_user_within_model_budget( "test-user", @@ -177,9 +228,7 @@ async def test_is_end_user_within_model_budget(budget_limiter): ) # Test when model exceeds budget - with patch.object( - budget_limiter, "_get_end_user_spend_for_model", return_value=150.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_end_user_within_model_budget( "test-user", @@ -198,25 +247,31 @@ async def test_is_end_user_within_model_budget(budget_limiter): ) -# Test _get_end_user_spend_for_model +# Test _get_spend_for_model_budget for the end-user scope @pytest.mark.asyncio -async def test_get_end_user_spend_for_model(budget_limiter): - budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d") +async def test_get_spend_for_end_user_model_budget(budget_limiter): + from litellm.proxy.hooks.model_max_budget_limiter import ( + END_USER_SPEND_CACHE_KEY_PREFIX, + ) - # Mock cache get - with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0): - spend = await budget_limiter._get_end_user_spend_for_model( - end_user_id="test-user", model="gpt-4", key_budget_config=budget_config - ) - assert spend == 50.0 + model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) - # Test with provider prefix - spend = await budget_limiter._get_end_user_spend_for_model( - end_user_id="test-user", + async def _spend(key): + return 50.0 if key == f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d" else None + + with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get: + spend = await budget_limiter._get_spend_for_model_budget( + entity_type=Litellm_EntityType.END_USER, + entity_id="test-user", model="openai/gpt-4", - key_budget_config=budget_config, + resolved=resolved, ) assert spend == 50.0 + assert [call.kwargs["key"] for call in mock_get.call_args_list] == [ + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d", + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:openai/gpt-4:1d", + ] @pytest.mark.asyncio @@ -261,16 +316,12 @@ async def test_async_log_success_event_uses_model_group_for_cache_key(budget_lim "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] # The cache key must use the model_group name, NOT the deployment name - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}") assert call_kwargs["response_cost"] == 0.10 @@ -310,15 +361,11 @@ async def test_async_log_success_event_falls_back_to_model_when_no_model_group( "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}") @pytest.mark.asyncio @@ -357,15 +404,11 @@ async def test_async_log_success_event_end_user_uses_model_group(budget_limiter) "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}" - ) + assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}") @pytest.mark.asyncio @@ -393,9 +436,7 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration( "metadata": {"user_api_key_end_user_id": end_user_id}, }, "litellm_params": { - "metadata": { - "user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget - }, + "metadata": {"user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget}, }, } with patch.object( @@ -403,15 +444,11 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration( "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}" - ) + assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}") assert call_kwargs["response_cost"] == 0.05 @@ -446,9 +483,7 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config "_push_in_memory_increments_to_redis", new_callable=AsyncMock, ) as mock_push: - await limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_push.assert_awaited_once() @@ -457,10 +492,7 @@ async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter, ): user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={}) - assert ( - await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - is None - ) + assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None @pytest.mark.asyncio @@ -472,12 +504,8 @@ async def test_get_fallback_model_within_budget_returns_first_within_budget( model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}}, budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=1.0 - ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0): + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result == "gpt-4o-mini" @@ -494,17 +522,15 @@ async def test_get_fallback_model_within_budget_skips_exhausted_fallback( budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - async def _spend_for_model(user_api_key_hash, model, key_budget_config): - return 150.0 if model == "gpt-4o-mini" else 1.0 + async def _spend_for_model(entity_type, entity_id, model, resolved): + return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0 with patch.object( budget_limiter, - "_get_virtual_key_spend_for_model", + "_get_spend_for_model_budget", side_effect=_spend_for_model, ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result == "claude-haiku" @@ -520,12 +546,8 @@ async def test_get_fallback_model_within_budget_returns_none_when_chain_exhauste }, budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0 - ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result is None @@ -554,7 +576,761 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim "_push_in_memory_increments_to_redis", new_callable=AsyncMock, ) as mock_push: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_push.assert_not_awaited() + + +def _success_kwargs( + *, + model_group, + deployment_model=None, + response_cost=0.5, + key_hash=None, + key_model_max_budget=None, + user_id=None, + user_model_max_budget=None, + end_user_id=None, + end_user_model_max_budget=None, +): + return { + "standard_logging_object": { + "response_cost": response_cost, + "model": deployment_model or model_group, + "model_group": model_group, + "end_user": end_user_id, + "metadata": { + "user_api_key_hash": key_hash, + "user_api_key_user_id": user_id, + "user_api_key_end_user_id": end_user_id, + }, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": key_model_max_budget, + "user_api_key_user_model_max_budget": user_model_max_budget, + "user_api_key_end_user_model_max_budget": end_user_model_max_budget, + }, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_model", + ["gpt-4", "openai/gpt-4"], + ids=["request_model_matches_budget_key", "request_model_carries_provider_prefix"], +) +async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request_model): + """ + The counter written post-call, the counter enforcement reads and the counter + /key/info reports must be one and the same, including when the request model + is not byte-identical to the configured budget key. + + Regression: the increment used to be keyed on the REQUEST model while + /key/info only ever looked up the CONFIGURED model, so a key could be + actively blocked at 429 while reporting current_spend 0. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_hash = "vk-hash" + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group=request_model, + response_cost=0.75, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) + assert usage == { + "gpt-4": { + "current_spend": 0.75, + "budget_limit": 1.0, + "time_period": "1d", + } + } + + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + # Still under the 1.0 limit. + assert await limiter.is_key_within_model_budget(user_api_key, request_model) is True + + await limiter.async_log_success_event( + _success_kwargs( + model_group=request_model, + response_cost=0.75, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key, request_model) + + usage_after = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) + assert usage_after["gpt-4"]["current_spend"] == 1.5 + + +@pytest.mark.asyncio +async def test_user_model_budget_is_tracked_and_enforced(): + """ + An internal user's own model_max_budget must be incremented post-call and + enforced, independently of any key-level budget. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + user_id = "user-1" + user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + + assert ( + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="openai/gpt-4", + ) + is True + ) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="openai/gpt-4", + response_cost=1.5, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + + with pytest.raises(litellm.BudgetExceededError) as exc: + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="openai/gpt-4", + ) + assert exc.value.entity_type == Litellm_EntityType.USER.value + + +@pytest.mark.asyncio +async def test_user_model_budget_counter_is_separate_from_the_key_counter(): + """ + A key budget and a user budget over the same model are two independent + counters, so one request must charge each exactly once. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=2.0, + key_hash="vk-hash", + key_model_max_budget=model_max_budget, + user_id="user-1", + user_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-hash:gpt-4:1d") == 2.0 + assert await dual_cache.async_get_cache(key="user_model_spend:user-1:gpt-4:1d") == 2.0 + + +@pytest.mark.asyncio +async def test_two_models_on_one_key_do_not_share_a_budget_window(): + """ + A key budgeting two models over different periods must own one window start + per model: a shared start lets the shorter period restart the longer one. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + model_max_budget = { + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "30d"}, + } + + start_time_keys = [] + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + for model in ("gpt-4", "claude-3"): + await limiter.async_log_success_event( + _success_kwargs( + model_group=model, + key_hash="vk-hash", + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + start_time_keys = [call.kwargs["start_time_key"] for call in mock_increment.call_args_list] + + assert start_time_keys == [ + "virtual_key_budget_start_time:vk-hash:gpt-4:1d", + "virtual_key_budget_start_time:vk-hash:claude-3:30d", + ] + assert len(set(start_time_keys)) == 2 + + +@pytest.mark.asyncio +async def test_no_increment_when_no_scope_budgets_the_model(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + key_hash="vk-hash", + key_model_max_budget={"claude-3": {"budget_limit": 1.0, "time_period": "1d"}}, + user_id="user-1", + user_model_max_budget={}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + mock_increment.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_build_model_max_budget_usage_skips_unusable_entries(): + """A malformed or period-less entry must be omitted, not crash the report.""" + dual_cache = DualCache() + await dual_cache.async_set_cache(key="virtual_key_spend:vk:gpt-4:1d", value=3.0) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk", + model_max_budget={ + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "no-period": {"budget_limit": 10.0}, + "bad-period": {"budget_limit": 10.0, "time_period": "not-a-duration"}, + }, + cache=dual_cache, + ) + assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_bedrock_traffic_charges_the_bare_family_name_budget(): + """ + The reported case: a budget configured as "claude-opus-4-8" with traffic on + "bedrock/anthropic.claude-opus-4-8". Before the fix nothing matched, so spend + was never tracked and the budget was never enforced no matter how far over it + the key went. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_hash = "vk-hash" + model_max_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="bedrock/anthropic.claude-opus-4-8", + response_cost=1.5, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) == { + "claude-opus-4-8": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "18h", + } + } + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key, "bedrock/anthropic.claude-opus-4-8") + + +@pytest.mark.asyncio +async def test_user_model_budget_window_resets_when_the_period_elapses(): + """ + A monthly user budget must start a fresh window once the period elapses, + and the window start must be scoped to that one budget model so a second + model on a shorter period cannot drag it forward. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + model_budget_spend_cache_key, + model_budget_start_time_cache_key, + ) + + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + user_id = "user-1" + user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + spend_key = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model="gpt-4", + budget_duration="1mo", + ) + start_time_key = model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model="gpt-4", + budget_duration="1mo", + ) + + kwargs = _success_kwargs( + model_group="gpt-4", + response_cost=1.5, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) + assert await dual_cache.async_get_cache(key=spend_key) == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="gpt-4", + ) + + # Age the window past its period. The next charge opens a new window rather + # than adding to the exhausted one. + elapsed = duration_in_seconds("1mo") + 60 + await dual_cache.async_set_cache( + key=start_time_key, + value=datetime.now(timezone.utc).timestamp() - elapsed, + ttl=elapsed, + ) + + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) + assert await dual_cache.async_get_cache(key=spend_key) == 1.5 + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + + +@pytest.mark.asyncio +async def test_a_zero_dollar_cap_blocks_the_model(): + """ + 0 is the operator saying "nobody may spend anything on this model", which is + the strictest cap expressible, not the absence of one. Skipping it on + falsiness turned the strictest setting into no setting at all, so the model + stayed wide open. The dashboard editor can produce this value, so it has to + mean something. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key = UserAPIKeyAuth( + token="hash-zero", + model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc: + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") + assert exc.value.max_budget == 0 + + +@pytest.mark.asyncio +async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent(): + """The usage endpoints must show the 0 too, or an operator cannot see the block they configured.""" + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-zero", + model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, + cache=DualCache(), + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_spend_exactly_at_the_cap_is_refused(): + """ + Spending the whole budget exhausts it. `>` let a caller sit exactly on the + limit and keep going, and every sibling budget check in the codebase + (RouterBudgetLimiting, the key and team budget checks) uses `>=`. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + budget = {"gpt-4": {"budget_limit": 2.0, "time_period": "1d"}} + key = UserAPIKeyAuth(token="hash-exact", model_max_budget=budget) + + await limiter.async_log_success_event( + _success_kwargs(model_group="gpt-4", response_cost=2.0, key_hash="hash-exact", key_model_max_budget=budget), + response_obj=None, + start_time=None, + end_time=None, + ) + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") + + +@pytest.mark.asyncio +async def test_usage_report_reads_every_counter_in_one_batched_lookup(): + """ + model_max_budget is caller-supplied and unbounded in size, so one cache + coroutine per configured model let a large map fan out into an unbounded + number of concurrent lookups on an endpoint anyone holding the key can call. + One batched read keeps it to a single round trip whatever the map's size. + """ + dual_cache = DualCache() + budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)} + + with ( + patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched, + patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single, + ): + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-many", + model_max_budget=budget, + cache=dual_cache, + ) + + assert batched.await_count == 1 + assert len(batched.await_args.kwargs["keys"]) == 50 + assert single.await_count == 0 + assert len(usage) == 50 + + +@pytest.mark.asyncio +async def test_usage_report_survives_a_batch_lookup_that_returns_nothing(): + """ + async_batch_get_cache swallows its own failures and returns None. Zipping + that against the budgets would raise and take the whole /key/info response + with it, so an unusable result has to read as a miss instead. + """ + dual_cache = DualCache() + with patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=None)): + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-none", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_one_malformed_scope_does_not_abort_the_other_scopes(): + """ + Every scope is resolved before any of them is incremented, so a single + unusable entry used to raise out of resolution and leave the key counter + unwritten too. The key's budget is well formed here and must still be + charged despite the user's entry being garbage. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=0.25, + key_hash="hash-mixed", + key_model_max_budget=key_budget, + user_id="user-mixed", + user_model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-mixed", + model_max_budget=key_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_an_unusable_budget_entry_is_not_enforced_instead_of_raising(): + """ + A config typo must not turn every request for that model into a 500. It + cannot be keyed, so it cannot be enforced; the write path rejects these, so + reaching here means config.yaml or a direct DB edit. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + key = UserAPIKeyAuth( + token="hash-malformed", + model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ) + + assert await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") is True + + +def test_resolve_model_budget_returns_none_for_an_unusable_entry(): + assert ( + resolve_model_budget( + model="gpt-4", + model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ) + is None + ) + + +def test_a_malformed_specific_entry_does_not_hide_a_usable_family_budget(): + """ + The candidate chain is most-specific-first and already falls through an entry + that is ABSENT. An entry that will not parse is indistinguishable from absent + as far as enforcement goes, so it has to fall through too: otherwise one bad + provider-prefixed entry silently disables the valid bare-family budget sitting + next to it, and the model goes uncapped. + """ + resolved = resolve_model_budget( + model="openai/gpt-4", + model_max_budget={ + "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "gpt-4": {"budget_limit": 7.0, "time_period": "1d"}, + }, + ) + + assert resolved is not None + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 7.0 + + +@pytest.mark.asyncio +async def test_a_malformed_specific_entry_still_enforces_the_family_budget(): + """The fall-through has to reach enforcement, not just resolution.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + budget = { + "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "gpt-4": {"budget_limit": 1.0, "time_period": "1d"}, + } + key = UserAPIKeyAuth(token="hash-fallthrough", model_max_budget=budget) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="openai/gpt-4", + response_cost=2.0, + key_hash="hash-fallthrough", + key_model_max_budget=budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="openai/gpt-4") + + +def test_documented_budget_spelling_survives_model_validate(): + """ + `budget_limit` / `time_period` are the spelling the docs, the CRUD endpoints + and the dashboard editor all use, and BudgetConfig maps them onto + `max_budget` / `budget_duration` inside its `__init__`. + + Pydantic v2 normally bypasses a custom `__init__` in `model_validate`, and + this code path validates rather than constructing. It works today, but that + is a property of the installed Pydantic rather than of anything in this + repository, so an upgrade could silently stop applying the mapping and + quietly disable every budget written in the documented spelling. Pinned here + so that becomes a red test instead of an outage. + """ + from litellm.types.utils import BudgetConfig + + validated = BudgetConfig.model_validate({"budget_limit": 5, "time_period": "1d"}) + assert validated.max_budget == 5.0 + assert validated.budget_duration == "1d" + + # Control: an unrecognised key must NOT populate max_budget, or the assertion + # above would also pass against a model that accepted anything at all. + ignored = BudgetConfig.model_validate({"bogus_limit": 5, "time_period": "1d"}) + assert ignored.max_budget is None + + +def test_resolution_accepts_both_documented_spellings(): + """The resolver is what enforcement, tracking and reporting all go through.""" + for budget in ( + {"gpt-4": {"budget_limit": 5, "time_period": "1d"}}, + {"gpt-4": {"max_budget": 5, "budget_duration": "1d"}}, + ): + resolved = resolve_model_budget(model="gpt-4", model_max_budget=budget) + assert resolved is not None, f"{budget} resolved to nothing" + assert resolved.budget_config.max_budget == 5.0 + assert resolved.budget_config.budget_duration == "1d" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity_type, prefix", + [ + (Litellm_EntityType.KEY, "virtual_key_spend"), + (Litellm_EntityType.END_USER, "end_user_model_spend"), + ], +) +async def test_a_pre_upgrade_counter_keyed_on_the_request_model_still_enforces(entity_type, prefix): + """An upgrading proxy must not hand out a second allowance for the window it is already in. + + Before the counter key moved to the configured budget model, spend for a + request on `openai/gpt-4` against a budget configured as `gpt-4` was both + written to and enforced on `{prefix}:{id}:openai/gpt-4:1d`. Reading only the + configured-model key finds that counter empty and admits another full budget + until the window expires. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key=f"{prefix}:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + if entity_type == Litellm_EntityType.KEY: + await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget), + model="openai/gpt-4", + ) + else: + await limiter.is_end_user_within_model_budget( + end_user_id="entity-1", + end_user_model_max_budget=model_max_budget, + model="openai/gpt-4", + ) + assert exc_info.value.current_cost == 25.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "legacy_spend, current_spend, expect_blocked", + [(6.0, 5.0, True), (2.0, 3.0, False)], +) +async def test_the_pre_upgrade_and_post_upgrade_counters_add_up_over_one_window( + legacy_spend, current_spend, expect_blocked +): + """The two counters hold disjoint halves of one window, so the window's spend is their sum. + + Nothing writes the request-model spelling once this version is running, so + the legacy counter is frozen at whatever the previous version charged and + the configured-model counter carries everything since. Either one alone + under-reports the window: 6 + 5 is over a cap of 10 that neither half + reaches on its own. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache( + key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=legacy_spend, ttl=86400 + ) + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=current_spend, ttl=86400) + + async def enforce(): + return await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget), + model="openai/gpt-4", + ) + + if expect_blocked: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await enforce() + assert exc_info.value.current_cost == legacy_spend + current_spend + else: + assert await enforce() is True + + +@pytest.mark.asyncio +async def test_the_configured_model_counter_is_never_counted_twice(): + """When the request names the budget exactly there is no legacy counter, only the one key. + + Both keys are `virtual_key_spend:entity-1:gpt-4:1d` here, so a lookup that + added them without noticing would charge 12 against a cap of 10 and refuse a + key that has spent 6. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=6.0, ttl=86400) + + assert ( + await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth( + token="entity-1", + model_max_budget={"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}, + ), + model="gpt-4", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_the_pre_upgrade_counter_is_no_longer_read_a_window_after_start_up(monkeypatch): + """The carry is bounded, so it cannot become a permanent second lookup on every request. + + A counter written by the previous version belongs to a window that was + already open when this process replaced it, so once a full window has passed + since start-up there is nothing left for the lookup to find. + """ + import litellm.proxy.hooks.model_max_budget_limiter as limiter_module + + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400) + user_api_key = UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget) + + # Control: within the first window since start-up the same counter blocks, + # so the assertion below cannot pass against a lookup that never worked. + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4") + + monkeypatch.setattr(limiter_module, "_PROCESS_STARTED_AT", limiter_module.time.monotonic() - 86401) + assert await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4") is True + + +@pytest.mark.asyncio +async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry(): + """The user scope is introduced by this change, so a request-model key under it is not one of ours. + + Reading one would invent a counter no previous version ever wrote, which is + the opposite of preserving one. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:openai/gpt-4:1d", value=25.0, ttl=86400) + + assert ( + await limiter.is_user_within_model_budget( + user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" + ) + is True + ) + + # Control: the same overspend under the key this scope does own must block, + # or the assertion above would pass against a scope that enforces nothing. + await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:gpt-4:1d", value=25.0, ttl=86400) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_user_within_model_budget( + user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" + ) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 58dbe3ad370..e9566254dbc 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -7,9 +7,7 @@ import sys import litellm.proxy import litellm.proxy.proxy_server -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path from typing import Dict, List, Optional from unittest.mock import MagicMock, patch, AsyncMock @@ -50,9 +48,7 @@ class Request: ), # Request with no client IP should not be allowed ], ) -def test_check_valid_ip( - allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool -): +def test_check_valid_ip(allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool): from litellm.proxy.auth.auth_utils import _check_valid_ip request = Request(client_ip) @@ -121,9 +117,7 @@ async def test_check_blocked_team(): last_refreshed_at=time.time(), ) await asyncio.sleep(1) - team_obj = LiteLLM_TeamTableCachedObj( - team_id=_team_id, blocked=False, last_refreshed_at=time.time() - ) + team_obj = LiteLLM_TeamTableCachedObj(team_id=_team_id, blocked=False, last_refreshed_at=time.time()) hashed_token = hash_token(user_key) print(f"STORING TOKEN UNDER KEY={hashed_token}") user_api_key_cache.set_cache(key=hashed_token, value=valid_token) @@ -173,9 +167,7 @@ async def test_team_object_has_object_permission_id(): request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - with patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ) as mock_common_checks: + with patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock) as mock_common_checks: mock_common_checks.return_value = True await user_api_key_auth(request=request, api_key="Bearer " + user_key) @@ -200,9 +192,7 @@ async def test_returned_user_api_key_auth(user_role, expected_role): from datetime import datetime new_obj = await _return_user_api_key_auth_obj( - user_obj=LiteLLM_UserTable( - user_role=user_role, user_id="", max_budget=None, user_email="" - ), + user_obj=LiteLLM_UserTable(user_role=user_role, user_id="", max_budget=None, user_email=""), api_key="hello-world", parent_otel_span=None, valid_token_dict={}, @@ -258,9 +248,7 @@ async def test_aaauser_personal_budgets(key_ownership): spend=20, ) - user_obj = LiteLLM_UserTable( - user_id=_user_id, spend=11, max_budget=10, user_email="" - ) + user_obj = LiteLLM_UserTable(user_id=_user_id, spend=11, max_budget=10, user_email="") user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token) user_api_key_cache.set_cache(key="{}".format(_user_id), value=user_obj) @@ -273,10 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership): test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache") - assert ( - test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) - == valid_token - ) + assert test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) == valid_token if key_ownership == "user_key": with pytest.raises(ProxyException) as exc_info: @@ -311,9 +296,7 @@ async def test_user_api_key_auth_fails_with_prohibited_params(prohibited_param): request.body = return_body try: - response = await user_api_key_auth( - request=request, api_key="Bearer " + user_key - ) + response = await user_api_key_auth(request=request, api_key="Bearer " + user_key) except Exception as e: print("error str=", str(e)) error_message = str(e.message) @@ -519,9 +502,7 @@ def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api verbose_proxy_logger.setLevel(logging.DEBUG) request = MagicMock(spec=Request) request.headers = headers - api_key = get_api_key_from_custom_header( - request=request, custom_litellm_key_header_name=custom_header_name - ) + api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name) assert api_key == expected_api_key @@ -572,9 +553,7 @@ from litellm.proxy._types import LitellmUserRoles (LitellmUserRoles.TEAM, "1234", "1234", True), ], ) -def test_allowed_route_inside_route( - user_role, auth_user_id, requested_user_id, expected_result -): +def test_allowed_route_inside_route(user_role, auth_user_id, requested_user_id, expected_result): from litellm.proxy.auth.auth_checks import allowed_route_check_inside_route from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles @@ -715,9 +694,7 @@ async def test_soft_budget_alert(): try: # Call user_api_key_auth - response = await user_api_key_auth( - request=request, api_key="Bearer " + user_key - ) + response = await user_api_key_auth(request=request, api_key="Bearer " + user_key) # Assert the request was allowed (no exception raised) assert response is not None @@ -883,9 +860,7 @@ async def test_user_api_key_auth_websocket(): mock_websocket.url = URL(url="/ws") # Mock the return value of `user_api_key_auth` when it's called within the `user_api_key_auth_websocket` function - with patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True - ) as mock_user_api_key_auth: + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: # Make the call to the WebSocket function await user_api_key_auth_websocket(mock_websocket) @@ -896,17 +871,11 @@ async def test_user_api_key_auth_websocket(): request_arg = mock_user_api_key_auth.call_args.kwargs["request"] # Verify that the request has headers set - assert hasattr( - request_arg, "headers" - ), "Request object should have headers attribute" - assert ( - "authorization" in request_arg.headers - ), "Request headers should contain authorization" + assert hasattr(request_arg, "headers"), "Request object should have headers attribute" + assert "authorization" in request_arg.headers, "Request headers should contain authorization" assert request_arg.headers["authorization"] == "Bearer some_api_key" - assert ( - mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" - ) + assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" @pytest.mark.asyncio @@ -929,9 +898,7 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): } mock_websocket.url = URL(url="/v1/realtime") - with patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True - ) as mock_user_api_key_auth: + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: await user_api_key_auth_websocket(mock_websocket) request_arg = mock_user_api_key_auth.call_args.kwargs["request"] @@ -1127,9 +1094,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): ) request._url = URL(url="/team/new") - monkeypatch.setattr( - litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True} - ) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True}) # Initialize jwt_handler with a default LiteLLM_JWTAuth so that the # virtual_key_claim_field check in user_api_key_auth doesn't fail with @@ -1158,9 +1123,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): ): try: await user_api_key_auth(request=request, api_key="Bearer fake.jwt.token") - pytest.fail( - "Expected this call to fail. Non-admin user should not access team routes." - ) + pytest.fail("Expected this call to fail. Non-admin user should not access team routes.") except ProxyException as e: print("e", e) assert "Only proxy admin can be used to generate" in str(e.message) @@ -1220,9 +1183,7 @@ async def test_user_api_key_from_query_param(): from litellm.proxy.proxy_server import hash_token, user_api_key_cache user_key = "sk-query-1234" - user_api_key_cache.set_cache( - key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)) - ) + user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))) setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -1235,9 +1196,7 @@ async def test_user_api_key_from_query_param(): "query_string": f"alt=sse&key={user_key}".encode(), } ) - request._url = URL( - url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}" - ) + request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}") async def return_body(): return b"{}" @@ -1246,3 +1205,592 @@ async def test_user_api_key_from_query_param(): valid_token = await user_api_key_auth(request=request, api_key="") assert valid_token.token == hash_token(user_key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "user_id,user_model_max_budget,expected_calls", + [ + ("u-1", {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 1), + ("u-1", {}, 0), + ("u-1", None, 0), + (None, {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 0), + ], + ids=["enforced", "empty_budget", "no_budget", "no_user_id"], +) +async def test_check_user_model_budget(user_id, user_model_max_budget, expected_calls): + """ + An internal user's model_max_budget must reach the limiter. Before this it was + stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read + by nothing, so a user-level per-model budget never blocked anything. + """ + from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget + + calls = [] + + class _Limiter: + async def is_user_within_model_budget(self, user_id, user_model_max_budget, model): + calls.append((user_id, user_model_max_budget, model)) + return True + + valid_token = UserAPIKeyAuth( + token="hash", + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=_Limiter(), + models=["gpt-4"], + ) + assert len(calls) == expected_calls + if expected_calls: + assert calls[0] == ("u-1", user_model_max_budget, "gpt-4") + + +@pytest.mark.asyncio +async def test_user_model_max_budget_is_threaded_onto_the_auth_object(): + """ + The limiter can only enforce what auth carries. Regression for the user row's + model_max_budget being dropped on the way into UserAPIKeyAuth. + """ + from datetime import datetime + + from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj + + budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + user_obj = LiteLLM_UserTable( + user_id="u-1", + max_budget=None, + spend=0.0, + user_email=None, + models=[], + model_max_budget=budget, + ) + + auth_obj = await _return_user_api_key_auth_obj( + user_obj=user_obj, + api_key="sk-1234", + parent_otel_span=None, + valid_token_dict={"token": "hash"}, + route="/chat/completions", + start_time=datetime.now(), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + assert auth_obj.user_model_max_budget == budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "over_budget,expect_refusal", + [(True, True), (False, False)], + ids=["over_budget_is_refused", "under_budget_is_served"], +) +async def test_user_model_budget_is_enforced_through_user_api_key_auth(over_budget, expect_refusal): + """ + Drive the real auth entry point, not the helper. + + The user's model_max_budget lives on the user row, and the joint + verification-token view auth builds its token from does not carry it. A test + that only exercises the helper passes while the whole path is inert, so this + one goes through user_api_key_auth with a key that has no per-model budget of + its own and asserts the USER's budget decides the outcome. + """ + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_UserTable, Litellm_EntityType, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key + from litellm.proxy.proxy_server import ( + hash_token, + model_max_budget_limiter, + user_api_key_cache, + ) + + user_id = "user-model-budget" + model = "gpt-4o" + key = "sk-user-model-budget" + hashed = hash_token(key) + user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}} + + setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "prisma_client", "present") + + await user_api_key_cache.async_set_cache( + key=hashed, + value=UserAPIKeyAuth(token=hashed, user_id=user_id, models=[], model_max_budget={}), + model_type=UserAPIKeyAuth, + ) + await model_max_budget_limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model=model, + budget_duration="1mo", + ), + value=5.0 if over_budget else 0.25, + ttl=600, + ) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return f'{{"model": "{model}"}}'.encode() + + request.body = return_body + + async def fake_get_user_object(**kwargs): + return LiteLLM_UserTable( + user_id=user_id, + max_budget=None, + spend=0.0, + user_email=None, + models=[], + model_max_budget=user_model_max_budget, + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=fake_get_user_object, + ): + if expect_refusal: + with pytest.raises(Exception) as exc: + await user_api_key_auth(request=request, api_key="Bearer " + key) + assert "budget" in str(exc.value).lower() + assert user_id in str(exc.value) + else: + result = await user_api_key_auth(request=request, api_key="Bearer " + key) + # The budget must also reach the token, or the post-call increment + # has nothing to charge and the counter never grows. + assert result.user_model_max_budget == user_model_max_budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "over_budget,expect_refusal", + [(True, True), (False, False)], + ids=["over_budget_is_refused", "under_budget_is_served"], +) +async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(over_budget, expect_refusal): + """ + JWT auth returns its own token instead of falling through to the + virtual-key budget checks, so the user's per-model budget has to be enforced + on that path explicitly. + + The dangerous shape is not "no tracking": the post-call increment charges the + JWT user's counter either way, so without this check the counter grows and + nothing ever reads it, which looks enforced and is not. + """ + from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget + from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key + from litellm.proxy.proxy_server import model_max_budget_limiter + + user_id = "jwt-user-model-budget" + model = "gpt-4o" + user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}} + + await model_max_budget_limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model=model, + budget_duration="1mo", + ), + value=5.0 if over_budget else 0.25, + ttl=600, + ) + + # The token the JWT branch builds and returns. + valid_token = UserAPIKeyAuth( + api_key=None, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + + if expect_refusal: + with pytest.raises(litellm.BudgetExceededError) as exc: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=[model], + ) + assert exc.value.entity_type == Litellm_EntityType.USER.value + else: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=[model], + ) + + +def test_jwt_path_enforces_the_user_model_budget_before_returning(): + """ + The JWT branch returns early, so the enforcement call has to sit before that + return rather than in the virtual-key block. Assert on the call graph, since + a helper-level test passes whether or not the JWT path ever calls it. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + def calls_before_each_return(node): + seen_check = [] + for child in ast.walk(node): + if isinstance(child, ast.Call): + fn = child.func + name = getattr(fn, "id", None) or getattr(fn, "attr", None) + if name == "_check_user_model_budget": + seen_check.append(child.lineno) + return seen_check + + check_lines = calls_before_each_return(tree) + assert check_lines, "_user_api_key_auth_builder never enforces the user model budget" + + jwt_returns = [ + n.lineno + for n in ast.walk(tree) + if isinstance(n, ast.Return) and isinstance(n.value, ast.Call) and getattr(n.value.func, "id", None) == "cast" + ] + assert jwt_returns, "expected the JWT branch's `return cast(UserAPIKeyAuth, valid_token)`" + assert any(check < jwt_return for check in check_lines for jwt_return in jwt_returns), ( + "the user model-budget check must run before the JWT branch returns" + ) + + +def test_every_jwt_branch_carries_the_user_model_budget(): + """ + Each JWT branch that builds or replaces `valid_token` has to put the user's + model budget on it, or the enforcement call a few lines later has nothing to + read and silently admits the request. + + The auto-register branch is the one that regressed: it REPLACES the token + built above it with a key-scoped one whose columns carry no user budget. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + ] + targets = { + t.value.id + for node in assignments + for t in node.targets + if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) + } + assert "auto_registered" in targets, ( + f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" + ) + assert "valid_token" in targets, "the virtual-key path must carry the user's model budget" + + +@pytest.mark.asyncio +async def test_user_budget_lookup_tolerates_an_unreadable_user(): + """ + `get_user_object(user_id_upsert=False)` raises a bare Exception when the row + is simply ABSENT, which is the ordinary state for a custom-auth deployment + that never writes users to the proxy DB. Refusing on that exception would + turn "no user row" into a 4xx for every such request, and a transient DB + blip into a full outage. + + The virtual-key path makes the same call and swallows the same exception + ("Unable to get user from db/cache. Setting user_obj to None"), so this is + the established contract, not a shortcut. There is also nothing to enforce: + the budget being looked up lives on the row that could not be read. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + prisma_client = MagicMock() + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(side_effect=Exception("No user table row")), + ): + budget = await _read_user_model_max_budget( + user_id="user-with-no-row", + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget is None + + +@pytest.mark.asyncio +async def test_user_budget_lookup_is_also_unenforced_when_the_database_is_down(): + """ + KNOWN LIMITATION, pinned deliberately rather than discovered later. + + `get_user_object` cannot tell "row absent" from "database unreachable": the + absent case raises inside its own try (auth_checks.py:2177) and the handler + at :2213 rewrites every exception into the same + `ValueError("User doesn't exist in db...")`. A connection error, a query + timeout and a malformed row all reach us as that one type and message. + + So tolerating the absent case, which the test above requires, unavoidably + tolerates an outage too, and a user who DOES have a per-model budget goes + unenforced while the DB is unreachable. This is pre-existing behaviour of + `get_user_object` that the virtual-key path inherits identically; it is not + introduced here. Distinguishing them needs a dedicated exception type for + the absent case and a change to both auth paths. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + db_down = ValueError("User doesn't exist in db. 'user_id'=u-1. Got error - Connection refused") + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(side_effect=db_down), + ): + budget = await _read_user_model_max_budget( + user_id="u-1", + prisma_client=MagicMock(), + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget is None + + +@pytest.mark.asyncio +async def test_user_budget_lookup_returns_the_budget_when_the_row_reads(): + """Positive control: the tolerance above must not be swallowing every result.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + stored = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_obj = MagicMock() + user_obj.model_max_budget = stored + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(return_value=user_obj), + ): + budget = await _read_user_model_max_budget( + user_id="user-1", + prisma_client=MagicMock(), + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget == stored + + +def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): + """ + `skip_budget_checks` is computed per request for zero-cost models, and the + JWT branch logs "Skipping all budget checks" when it is set. Any enforcement + call that ignores it makes the same request behave differently depending on + whether the caller used a JWT or a virtual key, and makes that log a lie. + + Structural rather than behavioural on purpose: the defect is a call site + sitting outside a guard, and driving both auth paths to a zero-cost model + would prove it for the two requests exercised rather than for every site. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: + for parent in ast.walk(node): + if not isinstance(parent, ast.If): + continue + test = parent.test + is_skip_guard = ( + isinstance(test, ast.UnaryOp) + and isinstance(test.op, ast.Not) + and isinstance(test.operand, ast.Name) + and test.operand.id == "skip_budget_checks" + ) + if is_skip_guard and any(sub is target for sub in ast.walk(parent)): + return True + return False + + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" + ] + assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" + + unguarded = [c for c in calls if not guarded_by_skip(tree, c)] + assert not unguarded, ( + f"{len(unguarded)} _check_user_model_budget call(s) run even when " + "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other" + ) + + +def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): + """ + The custom-auth helper runs its own key, user and end-user per-model budget + checks. If it does not honour the zero-cost skip that the JWT and + virtual-key paths honour, the same free request is refused under one auth + method and served under the others. + + Asserted structurally, on the same reasoning as the sibling test: the defect + is a check sitting outside a guard, and it must hold for checks added later + rather than only for whichever request a behavioural test happened to drive. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + src = textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks)) + tree = ast.parse(src) + + assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" + + budget_calls = ( + "_check_key_model_budget_with_fallback", + "_check_user_model_budget", + "is_end_user_within_model_budget", + ) + + def guarding_ifs(target: ast.AST) -> list[ast.If]: + return [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is target for sub in ast.walk(node)) + ] + + def mentions_skip(node: ast.If) -> bool: + return any(isinstance(sub, ast.Name) and sub.id == "skip_budget_checks" for sub in ast.walk(node.test)) + + for call_name in budget_calls: + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and ( + (isinstance(node.func, ast.Name) and node.func.id == call_name) + or (isinstance(node.func, ast.Attribute) and node.func.attr == call_name) + ) + ] + assert calls, f"{call_name} is no longer called here; update this invariant" + for call in calls: + assert any(mentions_skip(node) for node in guarding_ifs(call)), ( + f"{call_name} runs even for a zero-cost model, so custom auth refuses " + "requests the JWT and virtual-key paths serve" + ) + + +def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce(): + """ + The post-call spend hook reads `user_model_max_budget` off the token, so the + attach has to happen whether or not THIS request was enforceable. Gating it + on the same condition as the check leaves the user's counter uncharged for + every request with no resolvable model or a zero-cost one, which is exactly + the untracked-spend defect this PR fixes. + + Structural, because the failure is an assignment sitting inside a guard. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks))) + + attaches = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + ] + assert attaches, "custom auth no longer attaches the user budget at all" + + for attach in attaches: + enclosing_ifs = [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is attach for sub in ast.walk(node)) + ] + assert not enclosing_ifs, ( + "the user budget is attached inside a conditional, so the spend hook " + "cannot charge the user counter whenever that condition is false" + ) + + +def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): + """ + A JWT that maps to an existing virtual key resolves through the resolver + store, which builds the token from the KEY row alone and therefore carries + no user-level per-model budget. That branch sets `do_standard_jwt_auth = + False` precisely so it falls through to the shared virtual-key checks, where + the user row is loaded and its budget copied onto the token. + + Reviewed as a bypass three times, so the two halves it depends on are pinned + here: the branch must not return before the shared block, and the shared + block must copy the user row's budget onto the token. Structural on purpose, + because the claim is about control flow reaching a statement, and it has to + hold for branches added later rather than for one mocked request. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + # Half one: the shared block copies the user row's budget onto the token. + copies_user_row = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + and any(isinstance(v, ast.Attribute) and v.attr == "model_max_budget" for v in ast.walk(node.value)) + ] + assert copies_user_row, ( + "nothing copies the user row's model_max_budget onto the token, so a mapped-key " + "JWT reaches enforcement carrying the key's columns only" + ) + + # Half two: the mapped-key branch does not return before reaching it. + disables_standard_auth = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "do_standard_jwt_auth" for t in node.targets) + and isinstance(node.value, ast.Constant) + and node.value.value is False + ] + assert len(disables_standard_auth) == 1, "expected exactly one mapped-key branch" + marker = disables_standard_auth[0] + + enclosing = [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is marker for sub in node.body) + ] + assert enclosing, "could not locate the mapped-key branch body" + + returns_after = [ + node for node in ast.walk(enclosing[0]) if isinstance(node, ast.Return) and node.lineno > marker.lineno + ] + assert not returns_after, ( + "the mapped-key branch returns before the shared virtual-key checks, so the " + "user's per-model budget is never attached and never enforced" + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 6cc1d9e5add..28c82fdf528 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1202,6 +1202,8 @@ def _fake_user_api_key_auth( model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, + user_model_max_budget=None, + user_id=None, token=None, ): """Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields @@ -1220,6 +1222,8 @@ def _fake_user_api_key_auth( auth.model_max_budget = model_max_budget auth.end_user_model_max_budget = end_user_model_max_budget auth.end_user_id = end_user_id + auth.user_model_max_budget = user_model_max_budget + auth.user_id = user_id auth.token = token return auth @@ -1548,6 +1552,78 @@ async def test_summary_model_denied_when_key_over_model_budget(): assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" +async def test_summary_model_denied_when_user_over_model_budget(): + """Internal-user per-model budget is enforced for the summary subrequest too. + + This file propagates `user_api_key_user_model_max_budget` into the summary + subrequest's metadata, so its spend charges the user's counter. Enforcing + only the key and end-user scopes would let compaction increment a counter it + can never be refused by, which is the asymmetry this PR exists to remove. + """ + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + + auth = _fake_user_api_key_auth( + key_models=["all-proxy-models"], + user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}}, + user_id="user-over-budget", + token="hashed-token", + ) + + limiter = MagicMock() + limiter.is_user_within_model_budget = AsyncMock( + side_effect=litellm.BudgetExceededError( + message="over budget", current_cost=10, max_budget=5 + ) + ) + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" + + # The limiter is a mock, so it would accept any kwargs. Pin the call shape and + # check it against the real method, or a rename there would keep this test + # green while breaking compaction in production. + limiter.is_user_within_model_budget.assert_awaited_once_with( + user_id="user-over-budget", + user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}}, + model="claude-haiku-4-5", + ) + import inspect + + from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + ) + + real_params = inspect.signature( + _PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget + ).parameters + for kwarg in ("user_id", "user_model_max_budget", "model"): + assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts" + + async def test_summary_model_denied_when_end_user_over_model_budget(): """End-user per-model budget is enforced for the summary subrequest too.""" import litellm diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 8a036d7e62f..da51d513b39 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2975,6 +2975,7 @@ async def test_user_info_v2_response_shape(mocker): "updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc), "sso_user_id": None, "teams": ["team-a", "team-b"], + "model_max_budget": {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}}, } async def mock_find_unique(*args, **kwargs): @@ -3018,9 +3019,20 @@ async def test_user_info_v2_response_shape(mocker): "sso_user_id", "teams", "object_permission", + "model_max_budget", + "model_max_budget_usage", } assert set(response_dict.keys()) == expected_fields + # The dashboard's user edit form hydrates its per-model budget rows from + # these two, so dropping them makes a save replace the user's budgets. + assert response_dict["model_max_budget"] == { + "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"} + } + assert response_dict["model_max_budget_usage"] == { + "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} + } + # Verify teams is a list of strings (team IDs), not team objects assert isinstance(response.teams, list) assert all(isinstance(t, str) for t in response.teams) @@ -4150,3 +4162,66 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): assert response.object_permission.mcp_tool_permissions == { "github": ["list_issues"] } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_max_budget,expected_written", + [ + ( + {"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}}, + '{"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}}', + ), + (None, None), + ({}, None), + ], + ids=["supplied", "omitted", "empty"], +) +async def test_user_new_persists_model_max_budget( + monkeypatch, model_max_budget, expected_written +): + """ + /user/new used to echo model_max_budget back while writing {} to the user row, + so a per-model budget looked configured and was read by nothing. + + The omitted/empty cases are the other half: SSO and default-key callers reach + generate_key_helper_fn with no budget, and writing "{}" for them would clear + an existing user's budgets. + """ + from litellm.proxy.management_endpoints import key_management_endpoints + + captured = {} + + class _FakeUserRow: + models = [] + + class _FakePrisma: + async def insert_data(self, data, table_name): + if table_name == "user": + captured["user_data"] = dict(data) + return _FakeUserRow() + captured["key_data"] = dict(data) + return SimpleNamespace( + token=data.get("token"), + litellm_budget_table=None, + created_at=None, + updated_at=None, + ) + + async def get_data(self, *args, **kwargs): + return None + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False) + # model_max_budget is an enterprise feature; without this the call is rejected + # before it ever reaches the write this test is about. + monkeypatch.setattr(proxy_server, "premium_user", True, raising=False) + + await key_management_endpoints.generate_key_helper_fn( + request_type="user", + user_id="u-1", + model_max_budget=model_max_budget, + ) + + assert captured["user_data"].get("model_max_budget") == expected_written diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 069cfa01178..fff6368cfc6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13507,10 +13507,15 @@ async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.23) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_budget_test:gpt-4o:1d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_budget_test:gpt-4o:1d": 0.23}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13568,6 +13573,10 @@ async def test_info_key_fn_no_model_max_budget_skips_usage(monkeypatch): monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + mock_user_api_key_cache, + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13621,10 +13630,15 @@ async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.55) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_v2_test:gpt-4o:7d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_v2_test:gpt-4o:7d": 0.55}), + ) mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.token = test_key_token @@ -13680,10 +13694,15 @@ async def test_info_key_fn_budget_table_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=1.20) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d": 1.20}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13748,10 +13767,15 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=2.50) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d": 2.50}), + ) mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.token = test_key_token @@ -13793,8 +13817,13 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): @pytest.mark.asyncio -async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): - """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match.""" +async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): + """/key/info reads the one counter enforcement reads: the configured budget model. + + It used to probe a second, provider-stripped key because the counter was + written under the request model instead, which is what let a key report zero + usage while being blocked at 429. + """ from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import LiteLLM_VerificationToken @@ -13808,10 +13837,15 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.75]) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d": 0.75}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13847,7 +13881,22 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): assert "model_max_budget_usage" in result["info"] usage = result["info"]["model_max_budget_usage"] assert usage["openai/gpt-4o"]["current_spend"] == 0.75 - assert mock_user_api_key_cache.async_get_cache.await_count == 2 + + +async def _budget_cache(seeded): + """A real DualCache holding spend at the given LITERAL counter keys. + + The keys are spelled out in full on purpose. Seeding via + model_budget_spend_cache_key would move the seed and the read together, so + any change to the key format would still match itself and these tests could + never fail, which is the exact bug they exist to catch. + """ + from litellm.caching.caching import DualCache + + cache = DualCache() + for key, spend in seeded.items(): + await cache.async_set_cache(key, spend) + return cache @pytest.mark.asyncio @@ -13872,19 +13921,16 @@ async def test_build_model_max_budget_usage_reads_current_cache_window(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.30) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:30d": 0.30}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "30d"}}, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) + # 0.30 comes back only if the key matched virtual_key_spend:some-hash:gpt-4o:30d. assert result["gpt-4o"]["current_spend"] == 0.30 - mock_user_api_key_cache.async_get_cache.assert_awaited_once_with( - key="virtual_key_spend:some-hash:gpt-4o:30d" - ) @pytest.mark.asyncio @@ -13915,8 +13961,7 @@ async def test_build_model_max_budget_usage_skips_model_without_duration(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.10) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:1d": 0.10}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", @@ -13924,11 +13969,10 @@ async def test_build_model_max_budget_usage_skips_model_without_duration(): "gpt-4o": {"budget_limit": 1.0, "time_period": "1d"}, "gpt-3.5-turbo": {"budget_limit": 0.5}, }, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) - assert "gpt-4o" in result + assert result["gpt-4o"]["current_spend"] == 0.10 assert "gpt-3.5-turbo" not in result - assert mock_user_api_key_cache.async_get_cache.await_count == 1 @pytest.mark.asyncio @@ -13961,8 +14005,7 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.20) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-3.5-turbo:7d": 0.20}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", @@ -13970,32 +14013,35 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped(): "gpt-4o": {"max_budget": "not-a-number", "budget_duration": "1d"}, "gpt-3.5-turbo": {"budget_limit": 0.5, "time_period": "7d"}, }, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) assert "gpt-4o" not in result - assert "gpt-3.5-turbo" in result - assert mock_user_api_key_cache.async_get_cache.await_count == 1 + assert result["gpt-3.5-turbo"]["current_spend"] == 0.20 @pytest.mark.asyncio -async def test_build_model_max_budget_usage_provider_prefix_cache_fallback(): +async def test_build_model_max_budget_usage_reads_only_the_configured_model_key(): + """One lookup, at the configured budget model. + + The counter is written under the name the operator configured, so probing a + provider-stripped variant would read a key nothing writes. + """ from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.55]) + cache = await _budget_cache({"virtual_key_spend:test-hash:openai/gpt-4o:7d": 0.55}) result = await _build_model_max_budget_usage( api_key_hash="test-hash", model_max_budget={"openai/gpt-4o": {"budget_limit": 2.0, "time_period": "7d"}}, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) + # Seeded only under the configured name, so a provider-stripped probe reads 0.0. assert result["openai/gpt-4o"]["current_spend"] == 0.55 - assert mock_user_api_key_cache.async_get_cache.await_count == 2 def test_list_keys_substring_matching_param_defaults_to_false(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 090acf2dbb0..4c6ba23c88c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -15,9 +15,7 @@ from fastapi import Request, UploadFile from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, @@ -73,9 +71,7 @@ async def test_build_request_files_from_upload_file(): upload_file = UploadFile(file=file, filename="test.txt", headers=headers) upload_file.read = AsyncMock(return_value=file_content) - result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file - ) + result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(upload_file) assert result == ("test.txt", file_content, "text/plain") # Test with Starlette UploadFile @@ -87,9 +83,7 @@ async def test_build_request_files_from_upload_file(): ) starlette_file.read = AsyncMock(return_value=file_content) - result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - starlette_file - ) + result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(starlette_file) assert result == ("test2.txt", file_content, "text/plain") @@ -275,9 +269,7 @@ async def test_non_streaming_http_request_handler_multipart_with_non_empty_parse """ request = MagicMock(spec=Request) request.method = "POST" - request.headers = Headers( - {"content-type": "multipart/form-data; boundary=------------------------test"} - ) + request.headers = Headers({"content-type": "multipart/form-data; boundary=------------------------test"}) file_content = b"test file content" file = BytesIO(file_content) @@ -316,9 +308,7 @@ async def test_pass_through_request_failure_handler(): Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs. """ with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - with patch( - "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client: with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" ) as mock_processing: @@ -329,9 +319,7 @@ async def test_pass_through_request_failure_handler(): # Setup mock for httpx client mock_client = MagicMock() mock_client.client = MagicMock() - mock_client.client.request = AsyncMock( - side_effect=httpx.HTTPError("Request failed") - ) + mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) mock_get_client.return_value = mock_client # Mock headers for custom headers @@ -364,9 +352,7 @@ async def test_pass_through_request_failure_handler(): # Verify the arguments to post_call_failure_hook call_args = mock_proxy_logging.post_call_failure_hook.call_args[1] assert call_args["user_api_key_dict"] == mock_user_api_key_dict - assert isinstance( - call_args["original_exception"], TypeError - ) # Now expecting TypeError + assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError assert "traceback_str" in call_args @@ -410,27 +396,14 @@ def test_is_langfuse_route(): handler = PassThroughEndpointLogging() # Test positive cases - assert ( - handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces") - is True - ) - assert ( - handler.is_langfuse_route( - "https://proxy.example.com/langfuse/api/public/sessions" - ) - is True - ) + assert handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces") is True + assert handler.is_langfuse_route("https://proxy.example.com/langfuse/api/public/sessions") is True assert handler.is_langfuse_route("/langfuse/api/public/ingestion") is True assert handler.is_langfuse_route("http://localhost:4000/langfuse/") is True # Test negative cases - assert ( - handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False - ) - assert ( - handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages") - is False - ) + assert handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False + assert handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages") is False assert handler.is_langfuse_route("https://example.com/other") is False assert handler.is_langfuse_route("") is False @@ -447,17 +420,9 @@ def test_is_vertex_route_ignores_plain_predict_path_segment(): """ handler = PassThroughEndpointLogging() - assert ( - handler.is_vertex_route( - "https://upstream.example.com/ml/api/v1/time-series-forecast/predict" - ) - is False - ) + assert handler.is_vertex_route("https://upstream.example.com/ml/api/v1/time-series-forecast/predict") is False assert handler.is_vertex_route("https://upstream.example.com/api/v1/search") is False - assert ( - handler.is_vertex_route("https://upstream.example.com/predict/generateContent") - is False - ) + assert handler.is_vertex_route("https://upstream.example.com/predict/generateContent") is False assert ( handler.is_vertex_route( @@ -483,10 +448,7 @@ def test_is_vertex_route_ignores_plain_predict_path_segment(): ) is True ) - assert ( - handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") - is True - ) + assert handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") is True assert ( handler.is_vertex_route( @@ -545,9 +507,7 @@ async def test_custom_passthrough_predict_path_logs_via_generic_handler(): mock_vertex_handler.assert_not_called() handler._handle_logging.assert_awaited_once() - logged_object = handler._handle_logging.call_args.kwargs[ - "standard_logging_response_object" - ] + logged_object = handler._handle_logging.call_args.kwargs["standard_logging_response_object"] assert logged_object == {"response": '{"forecast": [1, 2, 3]}'} @@ -600,10 +560,7 @@ async def test_langfuse_passthrough_no_logging(): assert result is None # Verify that the passthrough_logging_payload was still set (this happens before the langfuse check) - assert ( - mock_logging_obj.model_call_details["passthrough_logging_payload"] - == passthrough_logging_payload - ) + assert mock_logging_obj.model_call_details["passthrough_logging_payload"] == passthrough_logging_payload def test_construct_target_url_with_subpath(): @@ -1051,9 +1008,7 @@ async def test_create_pass_through_route_with_cost_per_request(): # Mock the pass_through_request function to capture its call with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -1100,10 +1055,7 @@ def test_resolve_pass_through_request_timeout_precedence(): assert resolve_pass_through_request_timeout(endpoint_timeout=800) == 800.0 with patch("litellm.proxy.proxy_server.general_settings", {}): - assert ( - resolve_pass_through_request_timeout() - == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS - ) + assert resolve_pass_through_request_timeout() == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS def test_resolve_llm_passthrough_timeout_precedence(): @@ -1135,15 +1087,11 @@ async def test_pass_through_request_uses_resolved_timeout(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" ) as mock_get_client: - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda **kwargs: kwargs["data"] - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) mock_client = MagicMock() mock_client.client = MagicMock() - mock_client.client.request = AsyncMock( - side_effect=httpx.HTTPError("Request failed") - ) + mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) mock_get_client.return_value = mock_client mock_request = MagicMock(spec=Request) @@ -1181,9 +1129,7 @@ async def test_create_pass_through_route_forwards_timeout(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -1296,9 +1242,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body" ) as mock_get_response_body: # Setup mock for pre_call_hook and post_call_failure_hook - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"test": "data"} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1308,9 +1252,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {} - mock_response.aread = AsyncMock( - return_value=b'{"success": true}' - ) + mock_response.aread = AsyncMock(return_value=b'{"success": true}') mock_response.text = '{"success": true}' mock_response.raise_for_status = MagicMock() @@ -1330,9 +1272,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.url = "http://test-proxy.com/api/endpoint" - mock_request.body = AsyncMock( - return_value=b'{"message": "test request"}' - ) + mock_request.body = AsyncMock(return_value=b'{"message": "test request"}') mock_request.headers = Headers({}) mock_request.query_params = QueryParams({}) @@ -1411,9 +1351,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" ) as mock_chunk_processor: - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3", "stream": True} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3", "stream": True}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1438,9 +1376,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.url = "http://test-proxy.com/v1/messages" - mock_request.body = AsyncMock( - return_value=b'{"model": "claude-3", "stream": true}' - ) + mock_request.body = AsyncMock(return_value=b'{"model": "claude-3", "stream": true}') mock_request.headers = Headers({}) mock_request.query_params = QueryParams({}) @@ -1456,9 +1392,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): assert async_client.send.call_args.kwargs["stream"] is True mock_chunk_processor.assert_called_once() - logging_obj = mock_chunk_processor.call_args.kwargs[ - "litellm_logging_obj" - ] + logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"] assert logging_obj.stream is True assert logging_obj.model_call_details["stream"] is True @@ -1479,9 +1413,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" ) as mock_chunk_processor: - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3"} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3"}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1521,9 +1453,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): async_client.send.assert_awaited_once() mock_chunk_processor.assert_called_once() - logging_obj = mock_chunk_processor.call_args.kwargs[ - "litellm_logging_obj" - ] + logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"] assert logging_obj.stream is True assert logging_obj.model_call_details["stream"] is True @@ -1550,16 +1480,10 @@ async def test_create_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Mock existing config (empty list) - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=[] - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[]) # Create test endpoint data test_endpoint = PassThroughGenericEndpoint( @@ -1629,12 +1553,8 @@ async def test_update_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Create existing endpoint data existing_endpoint_id = "test-endpoint-123" existing_endpoints = [ @@ -1731,18 +1651,14 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist(): registry: dict = {} with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, patch("litellm.proxy.proxy_server.update_config_general_settings"), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", registry, ), ): - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=[] - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[]) # auth is not passed -> defaults to True on PassThroughGenericEndpoint endpoint = PassThroughGenericEndpoint( @@ -1757,19 +1673,12 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist(): ) assert any(value.get("auth") is True for value in registry.values()) - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/secure-passthrough", method="POST" - ) - is True - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/secure-passthrough", method="POST") is True post_request = MagicMock(spec=Request) post_request.method = "POST" - without_allowlist = UserAPIKeyAuth( - user_id="u", allowed_routes=["llm_api_routes"] - ) + without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"]) with pytest.raises(HTTPException) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( route="/secure-passthrough", @@ -1826,9 +1735,7 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist(): ] with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, patch("litellm.proxy.proxy_server.update_config_general_settings"), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", @@ -1851,19 +1758,12 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist(): user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), ) - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/edited-passthrough", method="POST" - ) - is True - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/edited-passthrough", method="POST") is True post_request = MagicMock(spec=Request) post_request.method = "POST" - without_allowlist = UserAPIKeyAuth( - user_id="u", allowed_routes=["llm_api_routes"] - ) + without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"]) with pytest.raises(HTTPException) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( route="/edited-passthrough", @@ -1905,12 +1805,8 @@ async def test_update_pass_through_endpoint_preserves_auth_false(): ] with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, - patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, + patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", registry, @@ -1937,12 +1833,7 @@ async def test_update_pass_through_endpoint_preserves_auth_false(): persisted = mock_update_config.call_args[1]["data"].field_value[0] assert persisted["auth"] is False - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/public-passthrough", method="POST" - ) - is False - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/public-passthrough", method="POST") is False @pytest.mark.asyncio @@ -1962,9 +1853,7 @@ async def test_update_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -1982,9 +1871,7 @@ async def test_update_pass_through_endpoint_not_found(): ) # Create update data - update_data = PassThroughGenericEndpoint( - path="/test/endpoint", target="http://newapi.com/v2" - ) + update_data = PassThroughGenericEndpoint(path="/test/endpoint", target="http://newapi.com/v2") # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) @@ -2023,12 +1910,8 @@ async def test_delete_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Create existing endpoint data endpoint_to_delete_id = "test-endpoint-123" other_endpoint_id = "other-endpoint-456" @@ -2106,9 +1989,7 @@ async def test_delete_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -2199,14 +2080,8 @@ async def test_get_pass_through_endpoints_includes_config_and_db(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._get_pass_through_endpoints_from_config" ) as mock_get_config: - db_objects = [ - PassThroughGenericEndpoint(**ep, is_from_config=False) - for ep in db_endpoints - ] - config_objects = [ - PassThroughGenericEndpoint(**ep, is_from_config=True) - for ep in config_endpoints - ] + db_objects = [PassThroughGenericEndpoint(**ep, is_from_config=False) for ep in db_endpoints] + config_objects = [PassThroughGenericEndpoint(**ep, is_from_config=True) for ep in config_endpoints] mock_get_db.return_value = db_objects mock_get_config.return_value = config_objects @@ -2280,13 +2155,9 @@ async def test_delete_pass_through_endpoint_empty_list(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock empty config - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) @@ -2325,9 +2196,7 @@ async def test_pass_through_request_query_params_forwarding(): ) as mock_get_response_body: # Setup mock for pre_call_hook test_body = {"name": "Azure Assistant", "model": "gpt-4o"} - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value=test_body - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body) mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} ) @@ -2336,9 +2205,7 @@ async def test_pass_through_request_query_params_forwarding(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - mock_response.aread = AsyncMock( - return_value=b'{"id": "asst_123", "object": "assistant"}' - ) + mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}') mock_response.text = '{"id": "asst_123", "object": "assistant"}' mock_response.raise_for_status = MagicMock() @@ -2360,20 +2227,12 @@ async def test_pass_through_request_query_params_forwarding(): # Create mock request with query parameters (Azure API version) mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = ( - "http://localhost:4000/azure-assistant/openai/assistants" - ) - mock_request.body = AsyncMock( - return_value=json.dumps(test_body).encode() - ) - mock_request.headers = Headers( - {"Content-Type": "application/json"} - ) + mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants" + mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode()) + mock_request.headers = Headers({"Content-Type": "application/json"}) # Create QueryParams with api-version parameter - mock_request.query_params = QueryParams( - [("api-version", "2025-01-01-preview")] - ) + mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")]) # Create mock user API key dict mock_user_api_key_dict = MagicMock() @@ -2395,9 +2254,7 @@ async def test_pass_through_request_query_params_forwarding(): # The key assertion: query parameters should be preserved and passed to the HTTP handler assert "requested_query_params" in call_kwargs - assert call_kwargs["requested_query_params"] == { - "api-version": "2025-01-01-preview" - } + assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"} assert call_kwargs.get("forward_multipart") is False # Verify the target URL is correct @@ -2447,9 +2304,7 @@ async def _run_pass_through_and_capture_wire_url( "PassThroughEndpoint client not found in in_memory_llm_clients_cache; " "get_async_httpx_client may not be caching this provider." ) - cache_dict[cache_key] = SimpleNamespace( - client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)) - ) + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) mock_request = MagicMock(spec=Request) mock_request.method = "GET" @@ -2458,18 +2313,14 @@ async def _run_pass_through_and_capture_wire_url( mock_request.body = AsyncMock(return_value=b"") mock_proxy_logging = MagicMock() - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type: data - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) mock_proxy_logging.get_proxy_hook = MagicMock(return_value=managed_files_hook) try: with ExitStack() as stack: - stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) - ) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)) if managed_files_hook is not None: stack.enter_context( patch( @@ -2637,26 +2488,16 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/allowed1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/allowed2", target="http://example.com/api2" - ), - PassThroughGenericEndpoint( - id="endpoint-3", path="/api/notallowed", target="http://example.com/api3" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"), + PassThroughGenericEndpoint(id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"), ] # Mock prisma client mock_prisma_client = MagicMock() mock_team = MagicMock() - mock_team.metadata = { - "allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"] - } - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_team.metadata = {"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2671,9 +2512,7 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter(): assert result[1].path == "/api/allowed2" # Verify database call - mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with( - where={"team_id": "test-team-123"} - ) + mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(where={"team_id": "test-team-123"}) @pytest.mark.asyncio @@ -2691,9 +2530,7 @@ async def test_filter_endpoints_by_team_allowed_routes_team_not_found(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test", target="http://example.com/api" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test", target="http://example.com/api"), ] # Mock prisma client to return None (team not found) @@ -2726,21 +2563,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_metadata(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has None metadata mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = None - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2768,21 +2599,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has metadata but no allowed_passthrough_routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"some_other_key": "some_value"} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2810,21 +2635,15 @@ async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has empty allowed_passthrough_routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"allowed_passthrough_routes": []} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2850,29 +2669,21 @@ async def test_filter_endpoints_by_team_allowed_routes_partial_match(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/openai", target="http://example.com/openai" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/openai", target="http://example.com/openai"), PassThroughGenericEndpoint( id="endpoint-2", path="/api/anthropic", target="http://example.com/anthropic", ), - PassThroughGenericEndpoint( - id="endpoint-3", path="/api/azure", target="http://example.com/azure" - ), - PassThroughGenericEndpoint( - id="endpoint-4", path="/api/cohere", target="http://example.com/cohere" - ), + PassThroughGenericEndpoint(id="endpoint-3", path="/api/azure", target="http://example.com/azure"), + PassThroughGenericEndpoint(id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"), ] # Mock prisma client with team that allows only 2 routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2904,9 +2715,7 @@ async def test_bedrock_router_passthrough_metadata_initialization(): ) # Mock ProxyBaseLLMRequestProcessing to verify it's used - with patch( - "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing" - ) as mock_processing_class: + with patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing") as mock_processing_class: # Setup mock instance mock_processor = MagicMock() mock_processing_class.return_value = mock_processor @@ -2914,12 +2723,8 @@ async def test_bedrock_router_passthrough_metadata_initialization(): # Mock successful response mock_response = MagicMock() mock_response.status_code = 200 - mock_response.aread = AsyncMock( - return_value=b'{"content": [{"text": "Hello"}]}' - ) - mock_processor.base_passthrough_process_llm_request = AsyncMock( - return_value=mock_response - ) + mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}') + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value=mock_response) # Create mock request with headers mock_request = MagicMock(spec=Request) @@ -2986,18 +2791,10 @@ async def test_bedrock_router_passthrough_metadata_initialization(): call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args[1] # These are the critical parameters that ensure metadata is properly initialized: - assert ( - call_kwargs["request"] == mock_request - ), "Request must be passed for header extraction" - assert ( - call_kwargs["user_api_key_dict"] == mock_user_api_key_dict - ), "User API key dict needed for metadata" - assert ( - call_kwargs["proxy_logging_obj"] == mock_proxy_logging - ), "Logging obj needed for hooks" - assert ( - call_kwargs["llm_router"] == mock_router - ), "Router needed for model routing" + assert call_kwargs["request"] == mock_request, "Request must be passed for header extraction" + assert call_kwargs["user_api_key_dict"] == mock_user_api_key_dict, "User API key dict needed for metadata" + assert call_kwargs["proxy_logging_obj"] == mock_proxy_logging, "Logging obj needed for hooks" + assert call_kwargs["llm_router"] == mock_router, "Router needed for model routing" assert call_kwargs["model"] == "my-bedrock-model", "Model name must be passed" # Verify response was returned @@ -3060,18 +2857,12 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata(): # Bedrock passthrough uses litellm_metadata to prevent key-level # tags from leaking into the provider payload (GH#30629). assert "litellm_metadata" in result, "litellm_metadata should be present in result" - assert ( - "headers" in result["litellm_metadata"] - ), "headers should be present in litellm_metadata" - assert isinstance( - result["litellm_metadata"]["headers"], dict - ), "headers should be a dictionary" + assert "headers" in result["litellm_metadata"], "headers should be present in litellm_metadata" + assert isinstance(result["litellm_metadata"]["headers"], dict), "headers should be a dictionary" # Verify specific headers are accessible (important for guardrails) headers = result["litellm_metadata"]["headers"] - assert ( - "user-agent" in headers or "User-Agent" in headers - ), "User-Agent header should be accessible in metadata" + assert "user-agent" in headers or "User-Agent" in headers, "User-Agent header should be accessible in metadata" # Also verify proxy_server_request has headers (original location) assert "proxy_server_request" in result @@ -3106,9 +2897,7 @@ async def test_create_pass_through_route_custom_body_url_target(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -3145,9 +2934,7 @@ async def test_create_pass_through_route_custom_body_url_target(): "retrievalQuery": {"text": "What is in the knowledge base?"}, } - setattr( - mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body - ) + setattr(mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body) await endpoint_func( request=mock_request, @@ -3185,9 +2972,7 @@ async def test_pass_through_request_non_streaming_uses_content_for_state_raw_bod mock_request.headers = Headers({"Content-Type": "application/json"}) mock_request.state = SimpleNamespace() setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) - mock_request.body = AsyncMock( - return_value=json.dumps(parsed_from_wire).encode("utf-8") - ) + mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8")) mock_user = MagicMock() mock_user.api_key = "sk-test" @@ -3260,9 +3045,7 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body(): mock_request.headers = Headers({"Content-Type": "application/json"}) mock_request.state = SimpleNamespace() setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) - mock_request.body = AsyncMock( - return_value=json.dumps(parsed_from_wire).encode("utf-8") - ) + mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8")) mock_user = MagicMock() mock_user.api_key = "sk-test" @@ -3336,9 +3119,7 @@ async def test_create_pass_through_route_no_custom_body_falls_back(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -3407,32 +3188,12 @@ def test_is_registered_pass_through_route_with_custom_root(): } with patch("litellm.proxy.utils.get_server_root_path", return_value="/proxy"): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/proxy/api/endpoint" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/api/endpoint" - ) - is True - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is True + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True with patch("litellm.proxy.utils.get_server_root_path", return_value="/"): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/api/endpoint" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/proxy/api/endpoint" - ) - is False - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is False # Clean up _registered_pass_through_routes.clear() @@ -3464,24 +3225,18 @@ def test_get_registered_pass_through_route_with_custom_root(): with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"): # Prefixed incoming route - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/litellm/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/litellm/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" assert result["headers"]["Authorization"] == "Bearer token123" # Bare incoming route (get_request_route convention) - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" with patch("litellm.proxy.utils.get_server_root_path", return_value="/"): - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" @@ -3535,12 +3290,7 @@ def test_db_registered_pass_through_route_bare_path_convention( "litellm.proxy.utils.get_server_root_path", return_value=server_root_path, ): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - incoming_route - ) - is should_match - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route(incoming_route) is should_match _registered_pass_through_routes.clear() @@ -3559,25 +3309,13 @@ def test_mapped_pass_through_routes_with_server_root_path(): with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"): # prefixed route should match mapped routes like /vertex_ai assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/litellm/vertex_ai/v1/projects/foo" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/litellm/bedrock/model/invoke" - ) + InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/vertex_ai/v1/projects/foo") is True ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/bedrock/model/invoke") is True # bare route without prefix should not match when root is set - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/vertex_ai/v1/projects/foo" - ) - is False - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/vertex_ai/v1/projects/foo") is False @pytest.mark.asyncio @@ -3594,24 +3332,18 @@ async def test_multipart_passthrough_preserves_boundary(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = httpx.Headers({"content-type": "application/json"}) - mock_response.aread = AsyncMock( - return_value=b'{"filename": "test.txt", "size": 17}' - ) + mock_response.aread = AsyncMock(return_value=b'{"filename": "test.txt", "size": 17}') mock_response.text = '{"filename": "test.txt", "size": 17}' async def mock_httpx_request(method, url, **kwargs): # Verify that files parameter is passed (not json) assert "files" in kwargs, "Files should be passed for multipart requests" - file_parts = [ - value for name, value in kwargs["files"] if name == "file" - ] + file_parts = [value for name, value in kwargs["files"] if name == "file"] assert len(file_parts) == 1, "File field should be in files" # Verify content-type is NOT in headers (httpx will set it with correct boundary) headers = kwargs.get("headers", {}) - assert ( - "content-type" not in headers - ), "content-type should be removed for multipart" + assert "content-type" not in headers, "content-type should be removed for multipart" filename, content, content_type = file_parts[0] assert filename == "test.txt" @@ -3684,9 +3416,7 @@ def test_get_response_headers_strips_server_and_date(): "connection", "keep-alive", ): - assert ( - stripped not in lowered_keys - ), f"{stripped!r} must not be forwarded by passthrough" + assert stripped not in lowered_keys, f"{stripped!r} must not be forwarded by passthrough" # Application/business headers must still pass through. lowered = {k.lower(): v for k, v in result.items()} @@ -3724,9 +3454,7 @@ class TestStaleRouteCleanupOnReload: ) stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) mock_set_env = stack.enter_context( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header" - ) + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header") ) mock_set_env.return_value = {} return stack @@ -3771,14 +3499,10 @@ class TestStaleRouteCleanupOnReload: so the registry would hold both paths instead of only ``/b``. """ with self._patches(): - await initialize_pass_through_endpoints( - [{"path": "/a", "target": "http://example.com"}] - ) + await initialize_pass_through_endpoints([{"path": "/a", "target": "http://example.com"}]) assert self._paths_in_registry() == ["/a"] - await initialize_pass_through_endpoints( - [{"path": "/b", "target": "http://example.com"}] - ) + await initialize_pass_through_endpoints([{"path": "/b", "target": "http://example.com"}]) assert self._paths_in_registry() == ["/b"] @@ -3801,12 +3525,8 @@ class TestStaleRouteCleanupOnReload: ] ) - assert InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/live-passthrough" - ) - assert InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/live-passthrough/some/subpath" - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough") + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough/some/subpath") # Regression (LIT-3538): a pre-call guardrail block on a passthrough endpoint @@ -3907,40 +3627,26 @@ async def _drive_pass_through_block(raised_exception): 400, ), ( - _FastAPIHTTPException( - status_code=400, detail={"error": "Violated moderation policy"} - ), + _FastAPIHTTPException(status_code=400, detail={"error": "Violated moderation policy"}), 400, ), ], ) -async def test_pre_call_guardrail_block_logs_warning_not_exception( - guardrail_exception, expected_code -): +async def test_pre_call_guardrail_block_logs_warning_not_exception(guardrail_exception, expected_code): status_code, logger = await _drive_pass_through_block(guardrail_exception) assert int(status_code) == expected_code - assert ( - logger.exception.call_count == 0 - ), "guardrail block must not be logged as an ERROR with a traceback" - assert ( - logger.warning.call_count == 1 - ), "guardrail block must be logged once at WARNING" + assert logger.exception.call_count == 0, "guardrail block must not be logged as an ERROR with a traceback" + assert logger.warning.call_count == 1, "guardrail block must be logged once at WARNING" @pytest.mark.asyncio async def test_non_guardrail_exception_still_logs_with_traceback(): - status_code, logger = await _drive_pass_through_block( - RuntimeError("upstream connection reset") - ) + status_code, logger = await _drive_pass_through_block(RuntimeError("upstream connection reset")) assert int(status_code) == 500 - assert ( - logger.exception.call_count == 1 - ), "a genuine failure must still be logged via verbose_proxy_logger.exception" - assert ( - logger.warning.call_count == 0 - ), "a genuine failure must not be downgraded to WARNING" + assert logger.exception.call_count == 1, "a genuine failure must still be logged via verbose_proxy_logger.exception" + assert logger.warning.call_count == 0, "a genuine failure must not be downgraded to WARNING" # Regression: generic config-based passthrough (`pass_through_request`) used to @@ -3979,9 +3685,7 @@ async def test_pass_through_request_non_streaming_upstream_error_returned_unchan ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4067,9 +3771,7 @@ async def test_pass_through_request_upstream_error_failure_hook_exception_is_swa mock_proxy_logging.post_call_failure_hook = AsyncMock( side_effect=RuntimeError("alerting integration misconfigured") ) - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4119,9 +3821,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_success_handler.return_value = None async_client = MagicMock() @@ -4149,10 +3849,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( streamed_chunks = [chunk async for chunk in response.body_iterator] await asyncio.sleep(0) - streamed_bytes = b"".join( - chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") - for chunk in streamed_chunks - ) + streamed_bytes = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks) assert streamed_bytes == upstream_content assert json.loads(streamed_bytes) == _UPSTREAM_ERROR_BODY @@ -4198,9 +3895,7 @@ async def test_pass_through_request_non_streaming_success_unchanged(): ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4244,9 +3939,7 @@ async def test_pass_through_request_internal_failure_still_raises_proxy_exceptio from litellm.proxy._types import ProxyException with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=RuntimeError("auth backend unavailable") - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=RuntimeError("auth backend unavailable")) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_request = MagicMock(spec=Request) @@ -4335,9 +4028,7 @@ def _inject_fake_passthrough_client(transport, timeout): def _enter_relay_logging_mocks(stack, parsed_body): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) @@ -4347,11 +4038,7 @@ def _enter_relay_logging_mocks(stack, parsed_body): ) ) mock_success_handler.return_value = None - stack.enter_context( - patch.object( - GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock() - ) - ) + stack.enter_context(patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock())) return mock_proxy_logging, mock_success_handler @@ -4437,10 +4124,7 @@ async def test_pass_through_request_relays_non_json_body_without_buffering(): mock_success_handler.assert_called_once() success_kwargs = mock_success_handler.call_args.kwargs assert success_kwargs["response_body"] is None - assert ( - success_kwargs["url_route"] - == "http://upstream.test/v1/messages/batches/b1/results" - ) + assert success_kwargs["url_route"] == "http://upstream.test/v1/messages/batches/b1/results" finally: cleanup() await fake_client.aclose() @@ -4518,9 +4202,7 @@ async def test_pass_through_request_upstream_error_body_stays_buffered(): ) try: with ExitStack() as stack: - mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks( - stack, {} - ) + mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(stack, {}) response = await pass_through_request( request=_relay_client_request(), @@ -4590,18 +4272,11 @@ async def test_pass_through_relay_client_disconnect_logs_partial_relay_warning(c partial_relay_warnings = [ record.getMessage() for record in caplog.records - if record.levelno == logging.WARNING - and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() + if record.levelno == logging.WARNING and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() ] assert len(partial_relay_warnings) == 1 - assert ( - "http://upstream.test/v1/messages/batches/b1/results" - in partial_relay_warnings[0] - ) - assert ( - f"{len(first_chunk)} bytes were sent to the client" - in partial_relay_warnings[0] - ) + assert "http://upstream.test/v1/messages/batches/b1/results" in partial_relay_warnings[0] + assert f"{len(first_chunk)} bytes were sent to the client" in partial_relay_warnings[0] assert upstream_stream.closed is True mock_success_handler.assert_called_once() @@ -4648,10 +4323,7 @@ async def test_pass_through_relay_full_consumption_logs_no_partial_relay_warning relayed = [chunk async for chunk in response.body_iterator] assert b"".join(relayed) == b"".join(upstream_chunks) - assert not any( - _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() - for record in caplog.records - ) + assert not any(_PARTIAL_RELAY_WARNING_MARKER in record.getMessage() for record in caplog.records) mock_success_handler.assert_called_once() finally: cleanup() @@ -4689,9 +4361,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body): logging worker would have run so the test can await them.""" from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) @@ -4708,9 +4378,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body): return mock_proxy_logging, enqueued -async def _run_upstream_reporting_passthrough( - upstream_headers, status_code=200, cost_per_request=None -): +async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200, cost_per_request=None): """Drive a generic pass-through against an upstream that reports its own cost/usage. Returns (recorded standard logging payloads, proxy logging mock).""" from litellm.proxy._types import UserAPIKeyAuth @@ -4731,9 +4399,7 @@ async def _run_upstream_reporting_passthrough( request=_relay_client_request(method="POST"), target="http://internal-api.test/v1/summarize", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-upstream-usage", team_id="team-fil" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"), cost_per_request=cost_per_request, ) for coroutine in enqueued: @@ -4802,23 +4468,17 @@ async def test_passthrough_records_upstream_reported_cost_on_error_response(): ) mock_proxy_logging.post_call_failure_hook.assert_awaited_once() - request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[ - "request_data" - ] + request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"] assert request_data["response_cost"] == 0.00021 assert request_data["combined_usage_object"] == litellm.Usage(total_tokens=930) @pytest.mark.asyncio async def test_passthrough_error_response_without_usage_headers_records_no_spend(): - _, mock_proxy_logging = await _run_upstream_reporting_passthrough( - {}, status_code=500 - ) + _, mock_proxy_logging = await _run_upstream_reporting_passthrough({}, status_code=500) mock_proxy_logging.post_call_failure_hook.assert_awaited_once() - request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[ - "request_data" - ] + request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"] assert "combined_usage_object" not in request_data @@ -4853,14 +4513,10 @@ async def test_streaming_passthrough_records_cost_and_tokens_reported_by_upstrea request=_relay_client_request(method="POST"), target="http://internal-api.test/v1/summarize", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-upstream-usage", team_id="team-fil" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"), ) assert isinstance(response, StreamingResponse) - assert [chunk async for chunk in response.body_iterator] == [ - b'data: {"delta": "hi"}\n\n' - ] + assert [chunk async for chunk in response.body_iterator] == [b'data: {"delta": "hi"}\n\n'] for coroutine in enqueued: await coroutine finally: @@ -4968,9 +4624,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", return_value=FakeUpstreamConnect(upstream_ws), ), - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER" - ) as mock_worker, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, ): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() @@ -5052,9 +4706,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", return_value=FakeUpstreamConnect(upstream_ws), ), - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER" - ) as mock_worker, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, ): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() @@ -5199,9 +4851,7 @@ async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rc "abnormal": Close(1006, "connection died"), "no_status": Close(1005, ""), }[rcvd_close] - upstream_ws = ClosingUpstreamWebSocket( - ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None) - ) + upstream_ws = ClosingUpstreamWebSocket(ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None)) websocket = _client_websocket(_pending_receive) with _patched_websocket_passthrough_environment(upstream_ws): @@ -5276,14 +4926,15 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f def _passthrough_kwargs_for_reservation( - user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None + user_api_key_dict: UserAPIKeyAuth, + parsed_body: Optional[dict] = None, + user_defined_route: bool = False, ) -> dict: mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = ( - "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" - ) + mock_request.url = "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" mock_request.headers = Headers({}) + mock_request.scope = {"endpoint": _marked_pass_through_endpoint()} if user_defined_route else {} return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( request=mock_request, @@ -5352,10 +5003,7 @@ async def test_passthrough_success_reconciles_budget_reservation(): reservation = user_api_key_dict.budget_reservation kwargs = _passthrough_kwargs_for_reservation(user_api_key_dict) - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] - is reservation - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is reservation increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) @@ -5381,9 +5029,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation(): user_api_key_dict, parsed_body={"litellm_metadata": {"user_api_key_budget_reservation": forged}}, ) - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) @@ -5391,9 +5037,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation(): assert increment_spend_counters.await_args.kwargs["budget_reservation"] is None -async def _drive_streaming_pass_through( - upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True -): +async def _drive_streaming_pass_through(upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True): """Drive pass_through_request against an upstream that stalls before its first byte. ``client_asked_for_stream`` picks which of pass_through_request's two streaming @@ -5405,22 +5049,14 @@ async def _drive_streaming_pass_through( ) with ExitStack() as stack: - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_get_client = stack.enter_context( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" - ) - ) - mock_chunk_processor = stack.enter_context( - patch.object(PassThroughStreamingHandler, "chunk_processor") + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") ) + mock_chunk_processor = stack.enter_context(patch.object(PassThroughStreamingHandler, "chunk_processor")) mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3", "stream": True} - if client_asked_for_stream - else {"model": "claude-3"} + return_value={"model": "claude-3", "stream": True} if client_asked_for_stream else {"model": "claude-3"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) @@ -5510,9 +5146,7 @@ async def test_pass_through_binary_event_stream_is_never_given_an_sse_comment(): @pytest.mark.asyncio @pytest.mark.parametrize("configured_interval, expect_ping", [(0.05, True), (None, False)]) -async def test_pass_through_route_pings_while_the_upstream_call_is_still_running( - configured_interval, expect_ping -): +async def test_pass_through_route_pings_while_the_upstream_call_is_still_running(configured_interval, expect_ping): """The upstream withholds its response headers until its first token, so the whole time-to-first-token is spent inside pass_through_request with nothing on the wire (issue #34819).""" @@ -5542,9 +5176,7 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running ) ) stack.enter_context(patch(f"{module}.pass_through_request", slow_pass_through)) - stack.enter_context( - patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval) - ) + stack.enter_context(patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval)) endpoint_func = create_pass_through_route( endpoint="/v1/messages", @@ -5571,3 +5203,147 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running assert (collected[0] == b": ping\n\n") is expect_ping assert collected[-1] in (MESSAGE_START_SSE_FRAME, MESSAGE_START_SSE_FRAME.decode()) + + +def test_passthrough_carries_the_per_model_budgets(): + """ + Native passthrough builds its logging metadata from + StandardLoggingUserAPIKeyMetadata, which has no budget field, and never calls + add_litellm_data_to_request. Without these three keys the post-call increment + exits early, so a /bedrock/... request is costed but its per-model counter is + never written: the budget reports zero forever and enforces nothing. + """ + key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_budget = {"claude-opus-4-8": {"budget_limit": 2.0, "time_period": "1mo"}} + end_user_budget = {"claude-opus-4-8": {"budget_limit": 3.0, "time_period": "1d"}} + + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth( + token="hash", + user_id="u-1", + model_max_budget=key_budget, + user_model_max_budget=user_budget, + end_user_model_max_budget=end_user_budget, + ) + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["user_api_key_model_max_budget"] == key_budget + assert metadata["user_api_key_user_model_max_budget"] == user_budget + assert metadata["user_api_key_end_user_model_max_budget"] == end_user_budget + + +def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body(): + """ + These keys decide budget enforcement, so a caller-supplied body must not be + able to raise its own cap. They are set after the client metadata merge for + the same reason user_api_key and the parent span are. + """ + key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth(token="hash", user_id="u-1", model_max_budget=key_budget), + parsed_body={ + "litellm_metadata": { + "user_api_key_model_max_budget": {"claude-opus-4-8": {"budget_limit": 999999.0, "time_period": "18h"}} + } + }, + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["user_api_key_model_max_budget"] == key_budget + + +def _marked_pass_through_endpoint(): + """An endpoint carrying the marker ``create_pass_through_route`` sets.""" + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + def _endpoint(): # pragma: no cover - identity only + return None + + setattr(_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) # noqa: B010 # name is a module constant + return _endpoint + + +def test_user_defined_passthrough_is_neither_tracked_nor_enforced(): + """ + `get_model_from_request` returns None for a user-defined pass-through on + purpose: the body is forwarded verbatim, so its `model` names an UPSTREAM + model rather than a LiteLLM-managed one, and enforcing key/team allowlists + against it would reject valid requests. Enforcement is therefore skipped + on those routes. + + Attaching the budget metadata anyway would charge a counter that nothing on + that route can refuse, and would attribute the spend to a budget the operator + scoped to a LiteLLM model that merely shares the name. Tracking and + enforcement have to agree: both on for the built-in provider routes, both off + here. + """ + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth( + token="hash", + user_id="u-1", + model_max_budget={"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}, + ), + user_defined_route=True, + ) + + metadata = kwargs["litellm_params"]["metadata"] + for field in ( + "user_api_key_model_max_budget", + "user_api_key_user_model_max_budget", + "user_api_key_end_user_model_max_budget", + ): + assert field not in metadata, f"{field} was attached on a route that never enforces it" + + +@pytest.mark.parametrize( + "handler_name", + [ + "anthropic_proxy_route", + "bedrock_proxy_route", + "gemini_proxy_route", + "cohere_proxy_route", + "vllm_proxy_route", + "mistral_proxy_route", + ], +) +def test_builtin_provider_routes_do_not_carry_the_user_defined_marker(handler_name): + """ + The budget metadata is attached only when the dispatched endpoint is NOT a + user-defined pass-through, so the built-in provider handlers must not carry + that marker or native provider spend would stop being tracked and enforced. + + These handlers DO call `create_pass_through_route` internally, and that + factory sets the marker on what it returns. But the result is awaited + immediately rather than registered, so FastAPI puts the decorated handler in + `request.scope["endpoint"]`, and that is what the marker check reads. This + test pins the distinction between calling the factory and being dispatched as + its product, which is easy to misread from a grep alone. + """ + from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + handler = getattr(llm_passthrough_endpoints, handler_name) + assert getattr(handler, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is False, ( + f"{handler_name} is marked as a user-defined pass-through, so per-model budget " + "metadata would be skipped and native provider spend would go untracked" + ) + + +def test_the_marker_check_distinguishes_the_two_route_kinds(): + """Positive control: the factory's product IS marked, so the check can discriminate.""" + from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint + from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints + + marked = MagicMock(spec=Request) + marked.scope = {"endpoint": _marked_pass_through_endpoint()} + assert request_dispatched_to_pass_through_endpoint(marked) is True + + builtin = MagicMock(spec=Request) + builtin.scope = {"endpoint": llm_passthrough_endpoints.anthropic_proxy_route} + assert request_dispatched_to_pass_through_endpoint(builtin) is False diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx index b3326df1ff8..459e3fd8c92 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx @@ -10,6 +10,7 @@ import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface BulkEditUserModalProps { open: boolean; @@ -36,6 +37,7 @@ const BulkEditUserModal: React.FC = ({ userModels, allowAllUsers = false, }) => { + const { premiumUser } = useAuthorized(); const [loading, setLoading] = useState(false); const [selectedTeams, setSelectedTeams] = useState([]); const [teamBudget, setTeamBudget] = useState(null); @@ -362,6 +364,7 @@ const BulkEditUserModal: React.FC = ({ userModels={userModels} possibleUIRoles={possibleUIRoles} isBulkEdit={true} + premiumUser={premiumUser === true} /> {loading && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx index 0ca68cc665e..8fb94ce477e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx @@ -1,4 +1,4 @@ -import { cleanup, screen, waitFor } from "@testing-library/react"; +import { cleanup, fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../../../tests/test-utils"; @@ -612,6 +612,125 @@ describe("UserEditView", () => { expect(onSubmit).not.toHaveBeenCalled(); }); + // /user/new validates model_max_budget behind an enterprise license, so a + // form that re-sends what is already stored turns an unrelated edit into a + // 400 on a proxy without one. + describe("per-model budgets", () => { + const withStoredBudgets = { + ...MOCK_USER_DATA, + user_info: { + ...MOCK_USER_DATA.user_info, + model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "30d" } }, + }, + }; + + it("should leave model_max_budget out of an edit that did not touch it", async () => { + const payload = await submittedPayload({ userData: withStoredBudgets, premiumUser: true }); + + expect(payload).not.toHaveProperty("model_max_budget"); + }); + + // The proxy stores model_max_budget as a plain dict, exactly as the client + // sent it, and BudgetConfig documents the max_budget/budget_duration + // spelling. A row hydrated from the spelling the editor does not read mounts + // with an empty cap, and every edit re-emits ALL rows, so touching one + // model's budget silently deletes another's. + it("should keep a row stored under the BudgetConfig aliases when a sibling row is edited", async () => { + const onSubmit = vi.fn(); + renderWithProviders( + , + ); + + const [aliasRow, canonicalRow] = await screen.findAllByPlaceholderText("Max spend ($)"); + expect(aliasRow).toHaveValue(5); + + fireEvent.change(canonicalRow, { target: { value: "3" } }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0].model_max_budget).toEqual({ + "gpt-4": { budget_limit: 5, time_period: "30d" }, + "gpt-3.5-turbo": { budget_limit: 3, time_period: "1h" }, + }); + }); + + // The effect already re-seeds the form on a userData change, so that change + // does happen while this component stays mounted. The editor holds its rows + // in state seeded once, so without a matching re-seed the rows on screen + // keep describing the previously loaded user and a save overwrites theirs. + it("re-seeds the editor when a different user is loaded", async () => { + const withBudget = (limit: number, id: string) => ({ + ...MOCK_USER_DATA, + user_id: id, + user_info: { + ...MOCK_USER_DATA.user_info, + model_max_budget: { "gpt-4": { budget_limit: limit, time_period: "1h" } }, + }, + }); + + const { rerender } = renderWithProviders( + , + ); + expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(5); + + rerender(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(99); + }); + + // BulkEditUsers copies a fixed field list into its payload and never reads + // model_max_budget, so an editor rendered here would take input and throw + // it away. It also has no single stored budget to diff against, since its + // userData stands in for every selected user. + it("does not offer the editor in bulk edit, where the value would be discarded", async () => { + renderWithProviders( + , + ); + + await screen.findByRole("button", { name: /save changes/i }); + expect(screen.queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument(); + }); + + it("should lock the editor when the proxy has no enterprise license", async () => { + renderWithProviders(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toBeDisabled(); + }); + + it("should leave the editor usable when the proxy has one", async () => { + renderWithProviders(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toBeEnabled(); + }); + }); + it("should send an empty-string metadata through untouched rather than as an object", async () => { const onSubmit = vi.fn(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 5612f5cd5f6..ed0c08adf38 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -2,6 +2,9 @@ import React, { useMemo, useState } from "react"; import { z } from "zod/v4"; import { all_admin_roles } from "@/utils/roles"; import BudgetDurationDropdown from "@/components/common_components/budget_duration_dropdown"; +import { ModelMaxBudget, ModelMaxBudgetField } from "@/components/key_team_helpers/ModelMaxBudgetEditor"; +import { modelMaxBudgetUpdate } from "@/components/key_team_helpers/modelMaxBudgetPayload"; +import { useSeededState } from "@/components/key_team_helpers/useSeededState"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; @@ -30,6 +33,7 @@ interface UserEditViewProps { possibleUIRoles: Record> | null; isBulkEdit?: boolean; objectPermission?: ObjectPermission | null; + premiumUser?: boolean; } const MCP_SELECTION_SHAPE = z.object({ @@ -135,9 +139,14 @@ export function UserEditView({ possibleUIRoles, isBulkEdit = false, objectPermission, + premiumUser = false, }: UserEditViewProps) { const canEditMcpPermissions = !isBulkEdit && all_admin_roles.includes(userRole || ""); const [unlimitedBudget, setUnlimitedBudget] = useState(false); + const [modelMaxBudget, setModelMaxBudget] = useSeededState( + userData.user_id, + () => userData.user_info?.model_max_budget ?? {}, + ); const schema = useMemo(() => budgetSchema(unlimitedBudget), [unlimitedBudget]); const form = useZodForm(schema, { defaultValues: toFormValues(userData, objectPermission, isBulkEdit, canEditMcpPermissions), @@ -162,9 +171,11 @@ export function UserEditView({ return; } + const modelBudgets = modelMaxBudgetUpdate(modelMaxBudget, userData.user_info?.model_max_budget); onSubmit({ ...values, ...("metadata" in values ? { metadata: metadata.value } : {}), + ...(modelBudgets !== undefined && { model_max_budget: modelBudgets }), max_budget: unlimitedBudget || values.max_budget === "" || values.max_budget === undefined ? null : values.max_budget, }); @@ -282,6 +293,20 @@ export function UserEditView({ {({ id, value, onChange }) => } + {/* Bulk edit forwards a fixed field list and has no single stored budget to + diff against, so the editor would silently discard whatever was typed. */} + {!isBulkEdit && ( + + )} + {({ ref, value, ...control }) => (