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/156] 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/156] 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/156] 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/156] 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/156] 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/156] 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 abf7dab0c2822edf8c3b2bc78618e62e5e6941f8 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 31 Jul 2026 21:14:18 +0000 Subject: [PATCH 007/156] 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 008/156] 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 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 009/156] 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 010/156] 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 011/156] 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 012/156] 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 013/156] 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 014/156] 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 015/156] 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 4e37425a782b7bf09d18e42ef9f072f6108155e4 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 21 Aug 2026 15:26:42 -0700 Subject: [PATCH 016/156] ci: ban row-rewriting DML from prisma migrations --- .../check_migrations_no_data_rewrites.py | 318 ++++++++++++++++++ .../test_check_migrations_no_data_rewrites.py | 234 +++++++++++++ 2 files changed, 552 insertions(+) create mode 100644 tests/code_coverage_tests/check_migrations_no_data_rewrites.py create mode 100644 tests/test_litellm/test_check_migrations_no_data_rewrites.py diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..fc0eafb3b16 --- /dev/null +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -0,0 +1,318 @@ +#!/usr/bin/env python3 +"""Ban row-rewriting DML from Prisma migrations. + +Migrations run synchronously at proxy boot, before the process serves traffic, so +anything whose cost scales with existing table size turns into downtime. A single +`UPDATE` with no batching over a spend-log-sized table is minutes of unavailability +plus a doubled heap that plain autovacuum will not give back. + +Flagged, per statement, by its leading keyword: + + UPDATE rewrites every matching row, and `WHERE` does not bound the scan + DELETE same scan, and the dead tuples outlive the migration + INSERT only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded + by the literal row list and passes + WITH a CTE-led statement containing any of the above + +Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a +statement's leading keyword, so they pass. + +Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this +repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise +hide. + +Add a column and let the application populate it, or run the rewrite as an opt-in +batched job outside boot. When a rewrite is genuinely bounded and must ship inside +the migration, put `-- data-migration-ok: ` on the statement, naming what +bounds it. The reason is required. + +`GRANDFATHERED` freezes the violations that predate this check. Prisma records a +checksum for every applied migration and this repo treats applied files as +immutable, so those two cannot take an inline marker. The set is closed; a new +migration belongs nowhere in it. +""" + +from __future__ import annotations + +import re +import sys +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" + +GRANDFATHERED = frozenset( + { + "20260817000000_shadow_eval_multi_key", + "20260818224500_add_shadow_eval_stopped_by", + } +) + +MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTILINE) +DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") +FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +STATEMENT = re.compile(r"[^;]+") + +REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) + +STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( + { + "INSERT", + "SELECT", + "WITH", + "ALTER", + "CREATE", + "DROP", + "TRUNCATE", + "COMMENT", + "GRANT", + "REVOKE", + "COPY", + "SET", + "PERFORM", + "RAISE", + "RETURN", + "EXECUTE", + "CALL", + "REINDEX", + "REFRESH", + "VACUUM", + "ANALYZE", + } +) + +GUIDANCE = """ +Migrations apply at proxy boot, before it serves traffic, so a statement whose cost +scales with table size is downtime. Add the column and let the application backfill +it, or move the rewrite to a batched job outside boot. + +If the rewrite is genuinely bounded and has to ship in the migration, mark the +statement with the bound spelled out: + + -- data-migration-ok: + UPDATE ... +""" + + +@dataclass(frozen=True, slots=True) +class Violation: + migration: str + line: int + keyword: str + + def render(self) -> str: + location = f"{MIGRATIONS_DIR.relative_to(REPO_ROOT)}/{self.migration}/migration.sql" + return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot" + + +def blank(text: str) -> str: + return "".join(character if character == "\n" else " " for character in text) + + +def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...]]: + """Blank comments and quoted text, keeping offsets, and locate dollar-quoted bodies.""" + chunks: list[str] = [] + bodies: list[tuple[int, int]] = [] + index = 0 + length = len(sql) + + while index < length: + pair = sql[index : index + 2] + + if pair == "--": + stop = sql.find("\n", index) + stop = length if stop == -1 else stop + chunks.append(blank(sql[index:stop])) + index = stop + continue + + if pair == "/*": + stop = skip_block_comment(sql, index) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + character = sql[index] + + if character in "'\"": + stop = skip_quoted(sql, index, character) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + if character == "$": + tag = DOLLAR_TAG.match(sql, index) + if tag is not None: + closing = sql.find(tag.group(), tag.end()) + body_end = length if closing == -1 else closing + stop = length if closing == -1 else closing + len(tag.group()) + bodies.append((tag.end(), body_end)) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + chunks.append(character) + index += 1 + + return "".join(chunks), tuple(bodies) + + +def skip_block_comment(sql: str, start: int) -> int: + depth = 1 + index = start + 2 + while index < len(sql) and depth > 0: + pair = sql[index : index + 2] + if pair == "/*": + depth += 1 + index += 2 + elif pair == "*/": + depth -= 1 + index += 2 + else: + index += 1 + return index + + +def skip_quoted(sql: str, start: int, quote: str) -> int: + index = start + 1 + while index < len(sql): + if sql[index] != quote: + index += 1 + elif sql[index + 1 : index + 2] == quote: + index += 2 + else: + return index + 1 + return len(sql) + + +def strip_parens(statement: str) -> str: + """Blank parenthesised groups in place, so an `IF EXISTS (SELECT ...)` guard does not + stand in for the statement it guards.""" + chunks: list[str] = [] + depth = 0 + + for character in statement: + if character == "(": + depth += 1 + chunks.append(" ") + elif character == ")": + depth = max(depth - 1, 0) + chunks.append(" ") + elif depth > 0 and character != "\n": + chunks.append(" ") + else: + chunks.append(character) + + return "".join(chunks) + + +def leading_keyword(statement: str) -> re.Match[str] | None: + """The statement's own keyword, looking past PL/pgSQL block syntax such as + `BEGIN`, `IF ... THEN` and `END`.""" + return next( + (word for word in FIRST_WORD.finditer(statement) if word.group().upper() in STATEMENT_KEYWORDS), + None, + ) + + +def offending_keyword(statement: str) -> str | None: + word = leading_keyword(strip_parens(statement)) + if word is None: + return None + + keyword = word.group().upper() + + if keyword in REWRITES_ROWS: + return keyword + + if keyword == "INSERT": + return "INSERT ... SELECT" if contains(statement, "SELECT") else None + + if keyword == "WITH": + nested = next((name for name in sorted(REWRITES_ROWS) if contains(statement, name)), None) + if nested is not None: + return f"WITH ... {nested}" + if contains(statement, "INSERT") and contains(statement, "SELECT"): + return "WITH ... INSERT ... SELECT" + + return None + + +def contains(statement: str, keyword: str) -> bool: + return re.search(rf"\b{keyword}\b", statement, re.IGNORECASE) is not None + + +def exempt_lines(sql: str) -> frozenset[int]: + return frozenset(sql.count("\n", 0, match.start()) + 1 for match in MARKER.finditer(sql)) + + +def scan(sql: str, migration: str, exempt: frozenset[int], offset: int = 0) -> Iterator[Violation]: + masked, bodies = mask(sql) + + for match in STATEMENT.finditer(masked): + keyword = offending_keyword(match.group()) + if keyword is None: + continue + first = line_of(sql, offset + keyword_start(match)) + last = line_of(sql, offset + match.end()) + if any(line in exempt for line in range(first - 1, last + 1)): + continue + yield Violation(migration, first, keyword) + + for start, end in bodies: + yield from scan(sql[start:end], migration, exempt, offset + start) + + +def keyword_start(statement: re.Match[str]) -> int: + word = leading_keyword(strip_parens(statement.group())) + return statement.start() + (0 if word is None else word.start()) + + +def line_of(sql: str, offset: int) -> int: + return sql.count("\n", 0, offset) + 1 + + +def scan_migration(directory: Path) -> tuple[Violation, ...]: + sql = (directory / "migration.sql").read_text(encoding="utf-8") + return tuple(scan(sql, directory.name, exempt_lines(sql))) + + +def stale_grandfathers(found: Mapping[str, tuple[Violation, ...]]) -> tuple[str, ...]: + clean = (name for name in GRANDFATHERED & found.keys() if not found[name]) + missing = GRANDFATHERED - found.keys() + return tuple(sorted((*clean, *missing))) + + +def main() -> int: + if not MIGRATIONS_DIR.is_dir(): + print(f"migrations directory not found: {MIGRATIONS_DIR}", file=sys.stderr) + return 2 + + directories = tuple(sorted(path for path in MIGRATIONS_DIR.iterdir() if (path / "migration.sql").is_file())) + found = {directory.name: scan_migration(directory) for directory in directories} + violations = tuple( + violation for name, results in found.items() if name not in GRANDFATHERED for violation in results + ) + + for violation in violations: + print(violation.render()) + + stale = stale_grandfathers(found) + for name in stale: + print(f"{name}: listed in GRANDFATHERED but no longer violates; remove it from the set") + + if violations: + print(GUIDANCE, file=sys.stderr) + print(f"{len(violations)} data-rewriting statement(s) in migrations.", file=sys.stderr) + + if violations or stale: + return 1 + + print(f"No data-rewriting statements in {len(directories)} migrations.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..40b55d741bc --- /dev/null +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -0,0 +1,234 @@ +"""Tests for tests/code_coverage_tests/check_migrations_no_data_rewrites.py. + +The checker reads migration.sql as SQL rather than as text, so the cases that matter +are the ones a grep would get wrong: `ON DELETE CASCADE` in a foreign key (60-odd +occurrences in the shipped migrations), an `UPDATE` inside a string literal or a +comment, and an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for +conditional DDL. +""" + +import importlib.util +import sys +from pathlib import Path + +_CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py" +_SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH) +assert _SPEC is not None and _SPEC.loader is not None +checker = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = checker +_SPEC.loader.exec_module(checker) + + +def _scan(tmp_path: Path, sql: str) -> tuple: + directory = tmp_path / "20260101000000_fixture" + directory.mkdir(exist_ok=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + return checker.scan_migration(directory) + + +def _keywords(tmp_path: Path, sql: str) -> tuple: + return tuple(violation.keyword for violation in _scan(tmp_path, sql)) + + +class TestRowRewritesAreFlagged: + def test_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'DELETE FROM "Foo" WHERE "a" IS NULL;') == ("DELETE",) + + def test_update_without_trailing_semicolon_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1') == ("UPDATE",) + + def test_lowercase_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'update "Foo" set "a" = 1;') == ("UPDATE",) + + def test_merge_is_flagged(self, tmp_path): + sql = 'MERGE INTO "Foo" t USING "Bar" s ON t."id" = s."id" WHEN MATCHED THEN UPDATE SET "a" = s."a";' + assert _keywords(tmp_path, sql) == ("MERGE",) + + def test_every_offending_statement_is_reported(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("UPDATE", "DELETE") + + def test_the_incident_migration_is_flagged(self, tmp_path): + sql = ( + 'UPDATE "LiteLLM_SpendLogs"\n' + ' SET "created_at" = "endTime",\n' + ' "updated_at" = "endTime"\n' + ' WHERE "created_at" > "endTime" + interval \'1 hour\';\n' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestSchemaStatementsPass: + def test_on_delete_cascade_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE CASCADE ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_on_delete_set_null_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE SET NULL ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_add_column_with_default_passes(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;' + assert _keywords(tmp_path, sql) == () + + def test_drop_table_passes(self, tmp_path): + assert _keywords(tmp_path, 'DROP TABLE IF EXISTS "Foo";') == () + + def test_empty_file_passes(self, tmp_path): + assert _keywords(tmp_path, "") == () + + def test_only_comments_passes(self, tmp_path): + assert _keywords(tmp_path, "-- nothing to do here\n") == () + + +class TestInsert: + def test_insert_values_is_bounded_and_passes(self, tmp_path): + assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == () + + def test_insert_select_scans_and_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar";') == ("INSERT ... SELECT",) + + +class TestCommonTableExpressions: + def test_cte_led_update_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) UPDATE "Foo" SET "a" = 1 FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... UPDATE",) + + def test_cte_led_delete_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) DELETE FROM "Foo" USING batch;' + assert _keywords(tmp_path, sql) == ("WITH ... DELETE",) + + def test_cte_led_insert_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") SELECT "id" FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_read_only_cte_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + + +class TestDollarQuotedBlocks: + def test_update_inside_do_block_is_flagged(self, tmp_path): + sql = 'DO $$\nBEGIN\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_conditional_ddl_do_block_passes(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'x') THEN\n" + ' ALTER TABLE "Foo" DROP CONSTRAINT "x";\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_tagged_dollar_quote_is_scanned(self, tmp_path): + sql = 'DO $body$\nBEGIN\n DELETE FROM "Foo";\nEND $body$;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_semicolons_inside_do_block_do_not_split_outer_statements(self, tmp_path): + sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + +class TestQuotingAndComments: + def test_update_inside_string_literal_passes(self, tmp_path): + sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'UPDATE nothing';" + assert _keywords(tmp_path, sql) == () + + def test_escaped_quote_inside_string_does_not_leak(self, tmp_path): + sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'it''s fine';\n" + assert _keywords(tmp_path, sql) == () + + def test_update_inside_line_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '-- UPDATE "Foo" SET "a" = 1;\nDROP TABLE "Bar";') == () + + def test_update_inside_block_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '/* UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";') == () + + def test_nested_block_comment_passes(self, tmp_path): + sql = '/* outer /* UPDATE "Foo" SET "a" = 1; */ still comment */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + + def test_update_inside_quoted_identifier_passes(self, tmp_path): + assert _keywords(tmp_path, 'ALTER TABLE "UPDATE Foo" ADD COLUMN "b" TEXT;') == () + + def test_positional_parameter_is_not_a_dollar_quote(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nUPDATE "Foo" SET "b" = $1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestEscapeHatch: + def test_marker_with_reason_exempts_the_statement(self, tmp_path): + sql = '-- data-migration-ok: one row per tenant, at most a few hundred\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_marker_without_reason_does_not_exempt(self, tmp_path): + assert _keywords(tmp_path, '-- data-migration-ok:\nUPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_marker_exempts_only_its_own_statement(self, tmp_path): + sql = ( + "-- data-migration-ok: bounded to in-flight jobs\n" + 'UPDATE "Foo" SET "a" = 1;\n' + 'UPDATE "Bar" SET "b" = 2;\n' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_marker_works_inside_a_do_block(self, tmp_path): + sql = 'DO $$\nBEGIN\n -- data-migration-ok: single row\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + def test_marker_below_the_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestReporting: + def test_line_number_points_at_the_statement_keyword(self, tmp_path): + sql = '-- CreateIndex\nCREATE INDEX "i" ON "Foo"("a");\n\nUPDATE "Foo" SET "a" = 1;' + assert _scan(tmp_path, sql)[0].line == 4 + + def test_render_names_the_migration_and_line(self, tmp_path): + violation = _scan(tmp_path, '\n\nDELETE FROM "Foo";')[0] + rendered = violation.render() + assert "20260101000000_fixture/migration.sql:3" in rendered + assert "DELETE" in rendered + + +class TestGrandfathering: + def test_every_grandfathered_migration_still_violates(self): + for name in sorted(checker.GRANDFATHERED): + directory = checker.MIGRATIONS_DIR / name + assert directory.is_dir(), f"{name} no longer exists; drop it from GRANDFATHERED" + assert checker.scan_migration(directory), f"{name} is clean; drop it from GRANDFATHERED" + + def test_stale_entry_is_reported_when_a_migration_stops_violating(self): + found = {name: () for name in checker.GRANDFATHERED} + assert checker.stale_grandfathers(found) == tuple(sorted(checker.GRANDFATHERED)) + + def test_missing_entry_is_reported(self): + assert checker.stale_grandfathers({}) == tuple(sorted(checker.GRANDFATHERED)) + + def test_no_stale_entries_against_the_real_tree(self): + found = { + path.name: checker.scan_migration(path) + for path in checker.MIGRATIONS_DIR.iterdir() + if (path / "migration.sql").is_file() + } + assert checker.stale_grandfathers(found) == () + + +class TestShippedMigrations: + def test_the_repo_is_clean(self): + assert checker.main() == 0 From 6c7dfbd2498a9b17d5f1afb57434cc9166fdbdcf Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 21 Aug 2026 15:27:56 -0700 Subject: [PATCH 017/156] ci: run the migration data-rewrite check in code quality --- .github/workflows/test-code-quality.yml | 3 +++ CLAUDE.md | 2 ++ 2 files changed, 5 insertions(+) diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 8f62837d29a..d0ac0b6fdee 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -128,6 +128,9 @@ jobs: - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py + - name: check_migrations_no_data_rewrites + run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py + - name: memory_test run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py diff --git a/CLAUDE.md b/CLAUDE.md index b3383b4a895..4a661f6effe 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -79,6 +79,8 @@ Do not put names of customers or customer company names in code, PR descriptions CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. Add the column and let the application populate it, or run the rewrite as an opt-in batched job outside boot. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this; when a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` + Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): - Composition over inheritance From 777eb8af107bf03eba4120a510d15bb1180a6af8 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 21 Aug 2026 15:37:05 -0700 Subject: [PATCH 018/156] test: close the sql-lexing gaps found by mutation testing Drops the doubled-quote branch in skip_quoted, which masked the same span either way and so could not be covered, and orders the failure report before the guidance text. --- .../check_migrations_no_data_rewrites.py | 18 +++----- .../test_check_migrations_no_data_rewrites.py | 45 ++++++++++++++++--- 2 files changed, 46 insertions(+), 17 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index fc0eafb3b16..61845ac521f 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -10,6 +10,7 @@ Flagged, per statement, by its leading keyword: UPDATE rewrites every matching row, and `WHERE` does not bound the scan DELETE same scan, and the dead tuples outlive the migration + MERGE both of the above in one statement INSERT only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded by the literal row list and passes WITH a CTE-led statement containing any of the above @@ -176,15 +177,10 @@ def skip_block_comment(sql: str, start: int) -> int: def skip_quoted(sql: str, start: int, quote: str) -> int: - index = start + 1 - while index < len(sql): - if sql[index] != quote: - index += 1 - elif sql[index + 1 : index + 2] == quote: - index += 2 - else: - return index + 1 - return len(sql) + """One quoted run, up to and including its closing quote. A doubled quote needs no + special case: closing on the first and reopening on the second masks the same span.""" + stop = sql.find(quote, start + 1) + return len(sql) if stop == -1 else stop + 1 def strip_parens(statement: str) -> str: @@ -304,8 +300,8 @@ def main() -> int: print(f"{name}: listed in GRANDFATHERED but no longer violates; remove it from the set") if violations: - print(GUIDANCE, file=sys.stderr) - print(f"{len(violations)} data-rewriting statement(s) in migrations.", file=sys.stderr) + print(f"\n{len(violations)} data-rewriting statement(s) in migrations.") + print(GUIDANCE) if violations or stale: return 1 diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 40b55d741bc..a4841f7a699 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -132,10 +132,36 @@ class TestDollarQuotedBlocks: ) assert _keywords(tmp_path, sql) == () + def test_guarded_update_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_guard_with_a_nested_call_still_flags_the_update(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo" WHERE lower("a") = \'x\' UNION SELECT 1) THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_tagged_dollar_quote_is_scanned(self, tmp_path): sql = 'DO $body$\nBEGIN\n DELETE FROM "Foo";\nEND $body$;' assert _keywords(tmp_path, sql) == ("DELETE",) + def test_tagged_dollar_quote_holds_an_apostrophe(self, tmp_path): + sql = 'INSERT INTO "Foo" ("t") VALUES ($body$don\'t$body$);\nUPDATE "Bar" SET "b" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_semicolons_inside_do_block_do_not_split_outer_statements(self, tmp_path): sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;' assert _keywords(tmp_path, sql) == () @@ -143,7 +169,7 @@ class TestDollarQuotedBlocks: class TestQuotingAndComments: def test_update_inside_string_literal_passes(self, tmp_path): - sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'UPDATE nothing';" + sql = 'ALTER TABLE "Foo" ADD COLUMN "note" TEXT NOT NULL DEFAULT \'UPDATE nothing\';' assert _keywords(tmp_path, sql) == () def test_escaped_quote_inside_string_does_not_leak(self, tmp_path): @@ -160,9 +186,20 @@ class TestQuotingAndComments: sql = '/* outer /* UPDATE "Foo" SET "a" = 1; */ still comment */\nDROP TABLE "Bar";' assert _keywords(tmp_path, sql) == () + def test_nested_block_comment_masks_past_the_inner_close(self, tmp_path): + sql = '/* outer /* inner */ UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + def test_update_inside_quoted_identifier_passes(self, tmp_path): assert _keywords(tmp_path, 'ALTER TABLE "UPDATE Foo" ADD COLUMN "b" TEXT;') == () + def test_select_in_a_quoted_identifier_does_not_make_an_insert_a_rewrite(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "SELECT Foo" ("id") VALUES (\'a\');') == () + + def test_update_in_a_quoted_identifier_does_not_make_a_cte_a_rewrite(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "UPDATE Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + def test_positional_parameter_is_not_a_dollar_quote(self, tmp_path): sql = 'ALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nUPDATE "Foo" SET "b" = $1;' assert _keywords(tmp_path, sql) == ("UPDATE",) @@ -177,11 +214,7 @@ class TestEscapeHatch: assert _keywords(tmp_path, '-- data-migration-ok:\nUPDATE "Foo" SET "a" = 1;') == ("UPDATE",) def test_marker_exempts_only_its_own_statement(self, tmp_path): - sql = ( - "-- data-migration-ok: bounded to in-flight jobs\n" - 'UPDATE "Foo" SET "a" = 1;\n' - 'UPDATE "Bar" SET "b" = 2;\n' - ) + sql = '-- data-migration-ok: bounded to in-flight jobs\nUPDATE "Foo" SET "a" = 1;\nUPDATE "Bar" SET "b" = 2;\n' assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 3 From 7e6d303e382b386e5d2d562a069e91494196778d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:13:47 -0700 Subject: [PATCH 019/156] fix: count migration lines against the whole file, scan EXECUTE'd sql, allow bounded inserts scan() recursed into a dollar-quoted body with the sliced text but kept absolute offsets, so line_of counted newlines in the slice against a position past its end. Any DO $$ block below the first line reported a wrong line, which also misaligned the -- data-migration-ok: markers: an unrelated marker earlier in the file could exempt a rewrite inside a block, and a marker sitting right above one failed to. Line numbers now always count against the whole migration text. EXECUTE was treated as harmless while its quoted SQL was masked, so a rewrite handed over as a string walked through the gate. The literal an EXECUTE runs is now scanned like a dollar-quoted body. INSERT was classified by searching the whole statement for SELECT, so a bounded INSERT ... VALUES holding a scalar subquery, or led by a helper CTE, was flagged as INSERT ... SELECT. A top-level VALUES now bounds the insert, and a VALUES buried in a subquery still does not. --- .../check_migrations_no_data_rewrites.py | 62 ++++++++-- .../test_check_migrations_no_data_rewrites.py | 108 ++++++++++++++++++ 2 files changed, 158 insertions(+), 12 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 61845ac521f..8ebdc43fac2 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -12,7 +12,8 @@ Flagged, per statement, by its leading keyword: DELETE same scan, and the dead tuples outlive the migration MERGE both of the above in one statement INSERT only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded - by the literal row list and passes + by the literal row list and passes, scalar subqueries in that list + included WITH a CTE-led statement containing any of the above Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a @@ -20,7 +21,12 @@ statement's leading keyword, so they pass. Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise -hide. +hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the +same to Postgres whether it is spelled out or handed over as a string. + +Line numbers always count against the whole migration file, however deeply the +statement is nested, so a reported line points at the statement and the markers +below line up with the statements they exempt. Add a column and let the application populate it, or run the rewrite as an opt-in batched job outside boot. When a rewrite is genuinely bounded and must ship inside @@ -112,10 +118,12 @@ def blank(text: str) -> str: return "".join(character if character == "\n" else " " for character in text) -def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...]]: - """Blank comments and quoted text, keeping offsets, and locate dollar-quoted bodies.""" +def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...], tuple[tuple[int, int], ...]]: + """Blank comments and quoted text, keeping offsets, and locate the spans that can still + hold SQL: dollar-quoted bodies, and the single-quoted literals `EXECUTE` runs.""" chunks: list[str] = [] bodies: list[tuple[int, int]] = [] + literals: list[tuple[int, int]] = [] index = 0 length = len(sql) @@ -139,6 +147,9 @@ def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...]]: if character in "'\"": stop = skip_quoted(sql, index, character) + if character == "'": + closed = sql[stop - 1 : stop] == character + literals.append((index + 1, max(index + 1, stop - 1 if closed else stop))) chunks.append(blank(sql[index:stop])) index = stop continue @@ -157,7 +168,7 @@ def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...]]: chunks.append(character) index += 1 - return "".join(chunks), tuple(bodies) + return "".join(chunks), tuple(bodies), tuple(literals) def skip_block_comment(sql: str, start: int) -> int: @@ -224,18 +235,31 @@ def offending_keyword(statement: str) -> str | None: return keyword if keyword == "INSERT": - return "INSERT ... SELECT" if contains(statement, "SELECT") else None + return "INSERT ... SELECT" if draws_rows_from_a_select(statement) else None if keyword == "WITH": nested = next((name for name in sorted(REWRITES_ROWS) if contains(statement, name)), None) if nested is not None: return f"WITH ... {nested}" - if contains(statement, "INSERT") and contains(statement, "SELECT"): + if contains(statement, "INSERT") and draws_rows_from_a_select(statement): return "WITH ... INSERT ... SELECT" return None +def draws_rows_from_a_select(statement: str) -> bool: + """Whether an `INSERT` takes its rows from a query rather than a literal list. A + top-level `VALUES` bounds the insert to the rows written out there, so the scalar + subqueries and helper CTEs that sit in parentheses around it do not make it a + rewrite.""" + return contains(statement, "SELECT") and not contains(strip_parens(statement), "VALUES") + + +def leads_with(statement: str, keyword: str) -> bool: + word = leading_keyword(strip_parens(statement)) + return word is not None and word.group().upper() == keyword + + def contains(statement: str, keyword: str) -> bool: return re.search(rf"\b{keyword}\b", statement, re.IGNORECASE) is not None @@ -244,21 +268,35 @@ def exempt_lines(sql: str) -> frozenset[int]: return frozenset(sql.count("\n", 0, match.start()) + 1 for match in MARKER.finditer(sql)) -def scan(sql: str, migration: str, exempt: frozenset[int], offset: int = 0) -> Iterator[Violation]: - masked, bodies = mask(sql) +def scan(sql: str, migration: str, exempt: frozenset[int]) -> Iterator[Violation]: + yield from scan_region(sql, sql, migration, exempt, 0) + + +def scan_region( + document: str, region: str, migration: str, exempt: frozenset[int], offset: int +) -> Iterator[Violation]: + """Violations in one region of `document`, whose text begins at `offset`. Lines are + always counted against the whole document, so a statement nested in a dollar-quoted + body reports its real file line and lines up with the markers read from that file.""" + masked, bodies, literals = mask(region) for match in STATEMENT.finditer(masked): + if leads_with(match.group(), "EXECUTE"): + for start, end in literals: + if match.start() <= start and end <= match.end(): + yield from scan_region(document, region[start:end], migration, exempt, offset + start) + continue keyword = offending_keyword(match.group()) if keyword is None: continue - first = line_of(sql, offset + keyword_start(match)) - last = line_of(sql, offset + match.end()) + first = line_of(document, offset + keyword_start(match)) + last = line_of(document, offset + match.end()) if any(line in exempt for line in range(first - 1, last + 1)): continue yield Violation(migration, first, keyword) for start, end in bodies: - yield from scan(sql[start:end], migration, exempt, offset + start) + yield from scan_region(document, region[start:end], migration, exempt, offset + start) def keyword_start(statement: re.Match[str]) -> int: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index a4841f7a699..9c6871cf557 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -97,6 +97,22 @@ class TestInsert: def test_insert_select_scans_and_is_flagged(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar";') == ("INSERT ... SELECT",) + def test_insert_values_with_a_scalar_subquery_passes(self, tmp_path): + sql = 'INSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT max("id")::text FROM "Bar"));' + assert _keywords(tmp_path, sql) == () + + def test_insert_values_with_a_scalar_subquery_per_row_passes(self, tmp_path): + sql = ( + 'INSERT INTO "Config" ("k", "v") VALUES\n' + " ('a', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'a')),\n" + " ('b', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'b'));" + ) + assert _keywords(tmp_path, sql) == () + + def test_values_inside_a_subquery_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM (VALUES (1), (2)) AS "v"("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -115,6 +131,10 @@ class TestCommonTableExpressions: sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;' assert _keywords(tmp_path, sql) == () + def test_cte_led_insert_values_is_bounded_and_passes(self, tmp_path): + sql = 'WITH latest AS (SELECT max("id") AS "id" FROM "Bar")\nINSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT "id"::text FROM latest));' + assert _keywords(tmp_path, sql) == () + class TestDollarQuotedBlocks: def test_update_inside_do_block_is_flagged(self, tmp_path): @@ -166,6 +186,32 @@ class TestDollarQuotedBlocks: sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;' assert _keywords(tmp_path, sql) == () + def test_line_number_inside_a_do_block_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_line_number_inside_a_nested_body_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- CreateIndex\n" + 'CREATE INDEX "i" ON "Foo"("a");\n' + "\n" + "DO $outer$\n" + "BEGIN\n" + " EXECUTE $inner$\n" + ' UPDATE "Foo" SET "a" = 1\n' + " $inner$;\n" + "END $outer$;" + ) + assert _scan(tmp_path, sql)[0].line == 7 + class TestQuotingAndComments: def test_update_inside_string_literal_passes(self, tmp_path): @@ -226,6 +272,68 @@ class TestEscapeHatch: sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_marker_inside_a_do_block_below_the_first_line_exempts(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_marker_above_a_do_block_does_not_exempt_a_rewrite_inside_it(self, tmp_path): + sql = ( + "-- data-migration-ok: bounded, this belongs to the insert below\n" + "INSERT INTO \"Config\" (\"k\") VALUES ('x');\n" + "\n" + 'DO $$ BEGIN UPDATE "Foo" SET "b" = 1; END $$;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + +class TestDynamicSql: + def test_execute_of_a_quoted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_execute_of_a_quoted_delete_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\"';\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_execute_of_a_formatted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE format('UPDATE %I SET \"a\" = 1', 'Foo');\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_execute_of_a_dollar_quoted_update_is_flagged(self, tmp_path): + sql = 'DO $outer$\nBEGIN\n EXECUTE $q$UPDATE "Foo" SET "a" = 1$q$;\nEND $outer$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_doubled_quote_inside_executed_sql_does_not_hide_the_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = date_trunc(''day'', \"t\")';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_execute_of_ddl_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_execute_of_a_read_only_query_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_executed_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n -- data-migration-ok: one row\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_literal_that_is_not_executed_is_still_inert(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('UPDATE \"Bar\" SET \"a\" = 1');" + assert _keywords(tmp_path, sql) == () + class TestReporting: def test_line_number_points_at_the_statement_keyword(self, tmp_path): From e7dea842c36be1585cc93c56473a087ed129e23d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:29:26 -0700 Subject: [PATCH 020/156] fix: scan sql held in a variable, and bound inserts by their own row source --- .../check_migrations_no_data_rewrites.py | 30 ++++++---- .../test_check_migrations_no_data_rewrites.py | 58 +++++++++++++++++++ 2 files changed, 77 insertions(+), 11 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 8ebdc43fac2..1e67617cead 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -11,9 +11,10 @@ Flagged, per statement, by its leading keyword: UPDATE rewrites every matching row, and `WHERE` does not bound the scan DELETE same scan, and the dead tuples outlive the migration MERGE both of the above in one statement - INSERT only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded - by the literal row list and passes, scalar subqueries in that list - included + INSERT only when it draws rows from a `SELECT`; an insert whose row source is + a leading `VALUES` is bounded by the rows spelled out there and passes, + scalar subqueries in that list included, while a `VALUES` reached + through a subquery or a set operation bounds nothing WITH a CTE-led statement containing any of the above Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a @@ -22,7 +23,9 @@ statement's leading keyword, so they pass. Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the -same to Postgres whether it is spelled out or handed over as a string. +same to Postgres whether it is spelled out or handed over as a string, and so is a +literal assigned to a variable with `:=`, which is where an `EXECUTE` further down +the body gets its statement from. Line numbers always count against the whole migration file, however deeply the statement is nested, so a reported line points at the statement and the markers @@ -248,11 +251,17 @@ def offending_keyword(statement: str) -> str | None: def draws_rows_from_a_select(statement: str) -> bool: - """Whether an `INSERT` takes its rows from a query rather than a literal list. A - top-level `VALUES` bounds the insert to the rows written out there, so the scalar - subqueries and helper CTEs that sit in parentheses around it do not make it a - rewrite.""" - return contains(statement, "SELECT") and not contains(strip_parens(statement), "VALUES") + """Whether an `INSERT` takes its rows from a query rather than a literal list. Only a + `SELECT` the insert is built on counts, so the scalar subqueries and helper CTEs that + sit in parentheses around a `VALUES` list do not make it a rewrite, while one reached + through a set operation does.""" + return contains(strip_parens(statement), "SELECT") + + +def hands_off_sql(statement: str) -> bool: + """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs + one outright, and an assignment parks one in a variable for an `EXECUTE` further down.""" + return leads_with(statement, "EXECUTE") or ":=" in statement def leads_with(statement: str, keyword: str) -> bool: @@ -281,11 +290,10 @@ def scan_region( masked, bodies, literals = mask(region) for match in STATEMENT.finditer(masked): - if leads_with(match.group(), "EXECUTE"): + if hands_off_sql(match.group()): for start, end in literals: if match.start() <= start and end <= match.end(): yield from scan_region(document, region[start:end], migration, exempt, offset + start) - continue keyword = offending_keyword(match.group()) if keyword is None: continue diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 9c6871cf557..2aa61b50949 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -113,6 +113,18 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM (VALUES (1), (2)) AS "v"("id");' assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + def test_values_after_a_set_operation_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" UNION ALL VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_values_after_an_except_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" EXCEPT VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_select_term_after_a_values_list_is_still_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1), (2) UNION ALL SELECT "id" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -272,6 +284,11 @@ class TestEscapeHatch: sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_marker_written_below_its_statement_leaves_that_statement_flagged(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + def test_marker_inside_a_do_block_below_the_first_line_exempts(self, tmp_path): sql = ( "-- AlterTable\n" @@ -334,6 +351,47 @@ class TestDynamicSql: sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('UPDATE \"Bar\" SET \"a\" = 1');" assert _keywords(tmp_path, sql) == () + def test_a_rewrite_declared_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_rewrite_assigned_in_the_body_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_ddl_assigned_to_a_variable_passes(self, tmp_path): + sql = "DO $$\nDECLARE\n stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n EXECUTE stmt;\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_held_in_a_variable(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := 'UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + class TestReporting: def test_line_number_points_at_the_statement_keyword(self, tmp_path): From c4d9a1ac6c501da875636d72a7cf12dd676fa952 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:40:44 -0700 Subject: [PATCH 021/156] fix: keep a marker trailing a statement from exempting the next one --- .../check_migrations_no_data_rewrites.py | 46 ++++++++++++++----- .../test_check_migrations_no_data_rewrites.py | 23 ++++++++++ 2 files changed, 58 insertions(+), 11 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 1e67617cead..558bce5ba6a 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -33,8 +33,10 @@ below line up with the statements they exempt. Add a column and let the application populate it, or run the rewrite as an opt-in batched job outside boot. When a rewrite is genuinely bounded and must ship inside -the migration, put `-- data-migration-ok: ` on the statement, naming what -bounds it. The reason is required. +the migration, put `-- data-migration-ok: ` on the statement or on the line +above it, naming what bounds it. The reason is required. A marker sharing a line +with the statement it follows exempts that statement alone, so the next statement +down is still checked rather than picking the marker up as its own. `GRANDFATHERED` freezes the violations that predate this check. Prisma records a checksum for every applied migration and this repo treats applied files as @@ -117,6 +119,19 @@ class Violation: return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot" +@dataclass(frozen=True, slots=True) +class Markers: + lines: frozenset[int] + standalone: frozenset[int] + + def exempt(self, first: int, last: int) -> bool: + """Whether a statement spanning `first` to `last` carries a marker. A marker alone on + its line speaks for the statement below it, which is how one written above a rewrite + exempts it. A marker sharing its line with the statement it follows speaks for that + statement only, so the next statement down does not inherit the exemption.""" + return any(line in self.lines for line in range(first, last + 1)) or first - 1 in self.standalone + + def blank(text: str) -> str: return "".join(character if character == "\n" else " " for character in text) @@ -273,16 +288,25 @@ def contains(statement: str, keyword: str) -> bool: return re.search(rf"\b{keyword}\b", statement, re.IGNORECASE) is not None -def exempt_lines(sql: str) -> frozenset[int]: - return frozenset(sql.count("\n", 0, match.start()) + 1 for match in MARKER.finditer(sql)) +def read_markers(sql: str) -> Markers: + lines: set[int] = set() + standalone: set[int] = set() + + for match in MARKER.finditer(sql): + line = sql.count("\n", 0, match.start()) + 1 + lines.add(line) + if not sql[sql.rfind("\n", 0, match.start()) + 1 : match.start()].strip(): + standalone.add(line) + + return Markers(frozenset(lines), frozenset(standalone)) -def scan(sql: str, migration: str, exempt: frozenset[int]) -> Iterator[Violation]: - yield from scan_region(sql, sql, migration, exempt, 0) +def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]: + yield from scan_region(sql, sql, migration, markers, 0) def scan_region( - document: str, region: str, migration: str, exempt: frozenset[int], offset: int + document: str, region: str, migration: str, markers: Markers, offset: int ) -> Iterator[Violation]: """Violations in one region of `document`, whose text begins at `offset`. Lines are always counted against the whole document, so a statement nested in a dollar-quoted @@ -293,18 +317,18 @@ def scan_region( if hands_off_sql(match.group()): for start, end in literals: if match.start() <= start and end <= match.end(): - yield from scan_region(document, region[start:end], migration, exempt, offset + start) + yield from scan_region(document, region[start:end], migration, markers, offset + start) keyword = offending_keyword(match.group()) if keyword is None: continue first = line_of(document, offset + keyword_start(match)) last = line_of(document, offset + match.end()) - if any(line in exempt for line in range(first - 1, last + 1)): + if markers.exempt(first, last): continue yield Violation(migration, first, keyword) for start, end in bodies: - yield from scan_region(document, region[start:end], migration, exempt, offset + start) + yield from scan_region(document, region[start:end], migration, markers, offset + start) def keyword_start(statement: re.Match[str]) -> int: @@ -318,7 +342,7 @@ def line_of(sql: str, offset: int) -> int: def scan_migration(directory: Path) -> tuple[Violation, ...]: sql = (directory / "migration.sql").read_text(encoding="utf-8") - return tuple(scan(sql, directory.name, exempt_lines(sql))) + return tuple(scan(sql, directory.name, read_markers(sql))) def stale_grandfathers(found: Mapping[str, tuple[Violation, ...]]) -> tuple[str, ...]: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 2aa61b50949..6befc895b4b 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -302,6 +302,29 @@ class TestEscapeHatch: ) assert _keywords(tmp_path, sql) == () + def test_a_marker_trailing_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + def test_a_marker_trailing_a_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + + def test_a_marker_trailing_a_multiline_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = ( + 'UPDATE "Foo"\n' + ' SET "a" = 1; -- data-migration-ok: one row\n' + 'UPDATE "Bar" SET "b" = 2;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_marker_alone_between_two_statements_belongs_to_the_one_below_it(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + def test_marker_above_a_do_block_does_not_exempt_a_rewrite_inside_it(self, tmp_path): sql = ( "-- data-migration-ok: bounded, this belongs to the insert below\n" From 622d9c598d0dccc17cc97fefbe69f385fdba0ced Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:56:02 -0700 Subject: [PATCH 022/156] fix: read dynamic SQL through the statement that hands it off A marker on an EXECUTE now covers the SQL that EXECUTE runs, so it goes where the migration reads rather than inside the string. A literal whose first line sat below its EXECUTE was missing the marker entirely, and the documented placement failed CI. A literal assigned with := counts as SQL only when an EXECUTE in the same body runs that variable by name. An error message naming a DELETE the application handles is text, and the only way to silence it before was a marker claiming a bounded data migration that was not there at all. --- .../check_migrations_no_data_rewrites.py | 49 +++++++++---- .../test_check_migrations_no_data_rewrites.py | 71 +++++++++++++++++++ 2 files changed, 107 insertions(+), 13 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 558bce5ba6a..2d5bb6ec726 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -24,8 +24,9 @@ Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the same to Postgres whether it is spelled out or handed over as a string, and so is a -literal assigned to a variable with `:=`, which is where an `EXECUTE` further down -the body gets its statement from. +literal assigned with `:=` to a variable some `EXECUTE` in the same body then runs +by name. A literal nothing runs is text, however much it reads like a statement, +so an error message naming a `DELETE` the application handles stays a message. Line numbers always count against the whole migration file, however deeply the statement is nested, so a reported line points at the statement and the markers @@ -36,7 +37,9 @@ batched job outside boot. When a rewrite is genuinely bounded and must ship insi the migration, put `-- data-migration-ok: ` on the statement or on the line above it, naming what bounds it. The reason is required. A marker sharing a line with the statement it follows exempts that statement alone, so the next statement -down is still checked rather than picking the marker up as its own. +down is still checked rather than picking the marker up as its own. A marker on an +`EXECUTE` or on the assignment feeding one covers the SQL that statement hands off, +so it goes where the migration reads rather than inside the string. `GRANDFATHERED` freezes the violations that predate this check. Prisma records a checksum for every applied migration and this repo treats applied files as @@ -66,6 +69,7 @@ MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTIL DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") STATEMENT = re.compile(r"[^;]+") +RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -273,10 +277,27 @@ def draws_rows_from_a_select(statement: str) -> bool: return contains(strip_parens(statement), "SELECT") -def hands_off_sql(statement: str) -> bool: - """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs - one outright, and an assignment parks one in a variable for an `EXECUTE` further down.""" - return leads_with(statement, "EXECUTE") or ":=" in statement +def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: + """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs one + outright. An assignment parks one in a variable, which counts only when something further + down runs that variable by name, since a string the migration never executes is text.""" + return leads_with(statement, "EXECUTE") or bool(assigned_names(statement) & executed) + + +def assigned_names(statement: str) -> frozenset[str]: + """The candidate variable names an assignment writes to, taken as every word ahead of the + `:=`. A declaration carries its type and sometimes a leading `DECLARE` alongside the name, + and none of that is worth parsing when the only question is which name is executed.""" + head, separator, _ = statement.partition(":=") + if not separator: + return frozenset() + return frozenset(word.group().lower() for word in FIRST_WORD.finditer(head)) + + +def executed_names(masked: str) -> frozenset[str]: + """The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps + an `EXECUTE` written inside a comment or a string from counting.""" + return frozenset(match.group(1).lower() for match in RUN_BY_NAME.finditer(masked)) def leads_with(statement: str, keyword: str) -> bool: @@ -312,18 +333,20 @@ def scan_region( always counted against the whole document, so a statement nested in a dollar-quoted body reports its real file line and lines up with the markers read from that file.""" masked, bodies, literals = mask(region) + executed = executed_names(masked) for match in STATEMENT.finditer(masked): - if hands_off_sql(match.group()): + first = line_of(document, offset + keyword_start(match)) + last = line_of(document, offset + match.end()) + exempt = markers.exempt(first, last) + + if hands_off_sql(match.group(), executed) and not exempt: for start, end in literals: if match.start() <= start and end <= match.end(): yield from scan_region(document, region[start:end], migration, markers, offset + start) + keyword = offending_keyword(match.group()) - if keyword is None: - continue - first = line_of(document, offset + keyword_start(match)) - last = line_of(document, offset + match.end()) - if markers.exempt(first, last): + if keyword is None or exempt: continue yield Violation(migration, first, keyword) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 6befc895b4b..926202f1d0e 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -415,6 +415,77 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == () + def test_a_marker_exempts_an_execute_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " EXECUTE '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_assignment_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_an_unmarked_execute_whose_sql_starts_on_a_later_line_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE '\n" + " UPDATE \"Foo\" SET \"a\" = 1';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_message_assigned_but_never_executed_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped, the application backfills them';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_notice_naming_a_delete_it_never_runs_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " note text := 'DELETE FROM legacy rows is handled by the application';\n" + "BEGIN\n" + " RAISE NOTICE '%', note;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped';\n" + " stmt text := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 4 + class TestReporting: def test_line_number_points_at_the_statement_keyword(self, tmp_path): From 3cca3f540286fd45514856f43eac95432c1dbd5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:04:16 -0700 Subject: [PATCH 023/156] fix: scan a DO body written in single quotes DO takes its body as a string literal, and dollar quoting is a convenience rather than a requirement. A migration spelling the body in single quotes got its rewrite through untouched, since nothing was reading that literal as SQL. It is ordinary syntax rather than an attempt to hide anything, so the miss was reachable by accident. The module docstring now also records where concatenated dynamic SQL stops being readable, which is a keyword split across fragments that do not hold it. Every fragment is scanned, so the shapes people actually write are all still caught. --- .../check_migrations_no_data_rewrites.py | 23 +++++++++++---- .../test_check_migrations_no_data_rewrites.py | 29 +++++++++++++++++++ 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 2d5bb6ec726..bf405a8aa31 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -25,8 +25,17 @@ repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the same to Postgres whether it is spelled out or handed over as a string, and so is a literal assigned with `:=` to a variable some `EXECUTE` in the same body then runs -by name. A literal nothing runs is text, however much it reads like a statement, -so an error message naming a `DELETE` the application handles stays a message. +by name, and so is the body of a `DO` written in single quotes rather than dollar +quotes. A literal nothing runs is text, however much it reads like a statement, so +an error message naming a `DELETE` the application handles stays a message. + +Each literal is read on its own, so a keyword built by concatenating fragments that +do not contain it (`'UPD' || 'ATE ...'`) is not caught. Every fragment is scanned, +so a concatenation is caught wherever the keyword survives whole in one of them, +which covers `'UPDATE ' || quote_ident(t)` and the rest of the readable shapes. The +gap needs a keyword deliberately split down the middle, and this check is a guard +against a rewrite reaching a boot unnoticed, not a defence against someone hiding +one on purpose. Line numbers always count against the whole migration file, however deeply the statement is nested, so a reported line points at the statement and the markers @@ -91,6 +100,7 @@ STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( "RAISE", "RETURN", "EXECUTE", + "DO", "CALL", "REINDEX", "REFRESH", @@ -279,9 +289,12 @@ def draws_rows_from_a_select(statement: str) -> bool: def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs one - outright. An assignment parks one in a variable, which counts only when something further - down runs that variable by name, since a string the migration never executes is text.""" - return leads_with(statement, "EXECUTE") or bool(assigned_names(statement) & executed) + outright, and so does `DO`, whose body is a string wherever it is not dollar-quoted. An + assignment parks one in a variable, which counts only when something further down runs + that variable by name, since a string the migration never executes is text.""" + if leads_with(statement, "EXECUTE") or leads_with(statement, "DO"): + return True + return bool(assigned_names(statement) & executed) def assigned_names(statement: str) -> frozenset[str]: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 926202f1d0e..c394eb30c37 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -472,6 +472,35 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == () + def test_a_do_body_in_single_quotes_is_scanned(self, tmp_path): + sql = "DO 'BEGIN UPDATE \"Foo\" SET \"a\" = 1; END';" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_a_quoted_do_body_with_a_language_clause_is_scanned(self, tmp_path): + sql = "DO LANGUAGE plpgsql 'BEGIN DELETE FROM \"Foo\"; END';" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_quoted_do_body_holding_only_ddl_passes(self, tmp_path): + sql = "DO 'BEGIN ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT; END';" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_quoted_do_body(self, tmp_path): + sql = ( + "-- data-migration-ok: one config row, keyed by its primary key\n" + "DO 'BEGIN UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''; END';" + ) + assert _keywords(tmp_path, sql) == () + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE ' || quote_ident('Foo') || ' SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_later_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'WITH x AS (SELECT 1) ' || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path): sql = ( "DO $$\n" From 3d16e327c6023a708f0ac8791ebe7400d0e46456 Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 21 Aug 2026 18:17:03 -0700 Subject: [PATCH 024/156] fix: judge an EXPLAIN-wrapped statement on the statement itself EXPLAIN ANALYZE runs the statement it wraps rather than only planning it, but ANALYZE sits in the keyword set, so it stood in for the keyword underneath and a rewrite left under one reached boot unflagged. --- .../check_migrations_no_data_rewrites.py | 34 ++++++++++++++---- .../test_check_migrations_no_data_rewrites.py | 35 +++++++++++++++++++ 2 files changed, 63 insertions(+), 6 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index bf405a8aa31..572c20805a0 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -20,6 +20,12 @@ Flagged, per statement, by its leading keyword: Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a statement's leading keyword, so they pass. +A statement wrapped in `EXPLAIN` is judged on the statement itself, because the +`ANALYZE` form runs it rather than only planning it, and a rewrite left under one +rewrites the table on the way to printing its timings. Explaining a rewrite without +`ANALYZE` is flagged too: nothing here needs the plan of a statement it is being +told not to run at boot, and a marker is a cheap answer if one ever does. + Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the @@ -79,6 +85,8 @@ DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") STATEMENT = re.compile(r"[^;]+") RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) +EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) +NOT_A_NEWLINE = re.compile(r"[^\n]") REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -247,17 +255,31 @@ def strip_parens(statement: str) -> str: return "".join(chunks) +def strip_explain(statement: str) -> str: + """Blank an `EXPLAIN` written with bare options, since the `ANALYZE` among them would + otherwise stand in for the keyword of the statement being explained. That statement is + the one worth reading: `EXPLAIN ANALYZE` runs it rather than only planning it, so a + rewrite underneath rewrites the table for real. The parenthesised option list needs + nothing here, already being blanked as a group.""" + return EXPLAIN_OPTIONS.sub(lambda match: NOT_A_NEWLINE.sub(" ", match.group()), statement) + + def leading_keyword(statement: str) -> re.Match[str] | None: - """The statement's own keyword, looking past PL/pgSQL block syntax such as - `BEGIN`, `IF ... THEN` and `END`.""" + """The statement's own keyword, looking past what wraps it: a parenthesised guard, + PL/pgSQL block syntax such as `BEGIN`, `IF ... THEN` and `END`, and an `EXPLAIN`. + Offsets survive both strips, so the match still points into `statement` itself.""" return next( - (word for word in FIRST_WORD.finditer(statement) if word.group().upper() in STATEMENT_KEYWORDS), + ( + word + for word in FIRST_WORD.finditer(strip_explain(strip_parens(statement))) + if word.group().upper() in STATEMENT_KEYWORDS + ), None, ) def offending_keyword(statement: str) -> str | None: - word = leading_keyword(strip_parens(statement)) + word = leading_keyword(statement) if word is None: return None @@ -314,7 +336,7 @@ def executed_names(masked: str) -> frozenset[str]: def leads_with(statement: str, keyword: str) -> bool: - word = leading_keyword(strip_parens(statement)) + word = leading_keyword(statement) return word is not None and word.group().upper() == keyword @@ -368,7 +390,7 @@ def scan_region( def keyword_start(statement: re.Match[str]) -> int: - word = leading_keyword(strip_parens(statement.group())) + word = leading_keyword(statement.group()) return statement.start() + (0 if word is None else word.start()) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index c394eb30c37..fb56d032617 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -516,6 +516,41 @@ class TestDynamicSql: assert _scan(tmp_path, sql)[0].line == 4 +class TestExplain: + def test_explain_analyze_over_an_update_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_explain_analyze_verbose_over_a_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE VERBOSE DELETE FROM "Foo";') == ("DELETE",) + + def test_explain_with_a_parenthesised_analyze_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN (ANALYZE, BUFFERS) UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_explain_analyze_over_an_insert_select_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE INSERT INTO "Foo" SELECT "a" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_explain_analyze_over_a_select_passes(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE SELECT * FROM "Foo";') == () + + def test_a_marker_exempts_an_explained_rewrite(self, tmp_path): + sql = '-- data-migration-ok: one config row\nEXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_an_analyze_of_its_own_passes(self, tmp_path): + assert _keywords(tmp_path, 'ANALYZE "Foo";') == () + + def test_a_vacuum_analyze_passes(self, tmp_path): + assert _keywords(tmp_path, 'VACUUM ANALYZE "Foo";') == () + + def test_an_explained_rewrite_inside_a_block_reports_its_line(self, tmp_path): + sql = 'DO $$\nBEGIN\n EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + class TestReporting: def test_line_number_points_at_the_statement_keyword(self, tmp_path): sql = '-- CreateIndex\nCREATE INDEX "i" ON "Foo"("a");\n\nUPDATE "Foo" SET "a" = 1;' From aabaa5151b0f6de6201f50f044ef99ae057251b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:43:04 -0700 Subject: [PATCH 025/156] docs: say where a marker goes for a dollar-quoted dynamic payload A dollar-quoted payload is read as its own region rather than as a handed-off string, so the marker belongs on the rewrite inside it. Pin that placement, and pin that a marker on a DO block header never covers the block's body. --- .../check_migrations_no_data_rewrites.py | 11 +-- .../test_check_migrations_no_data_rewrites.py | 71 +++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 572c20805a0..cc541a70102 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -53,8 +53,12 @@ the migration, put `-- data-migration-ok: ` on the statement or on the l above it, naming what bounds it. The reason is required. A marker sharing a line with the statement it follows exempts that statement alone, so the next statement down is still checked rather than picking the marker up as its own. A marker on an -`EXECUTE` or on the assignment feeding one covers the SQL that statement hands off, -so it goes where the migration reads rather than inside the string. +`EXECUTE` or on the assignment feeding one covers the single-quoted SQL that +statement hands off, so it goes where the migration reads rather than inside the +string. A dollar-quoted payload is not a string to this check but a region read like +any other body, so a rewrite inside one takes its marker on the rewrite itself. That +placement is deliberate rather than an oversight: a marker covering a whole body +would let one written for a `DO` block silence a rewrite added to that block later. `GRANDFATHERED` freezes the violations that predate this check. Prisma records a checksum for every applied migration and this repo treats applied files as @@ -86,7 +90,6 @@ FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") STATEMENT = re.compile(r"[^;]+") RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) -NOT_A_NEWLINE = re.compile(r"[^\n]") REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -261,7 +264,7 @@ def strip_explain(statement: str) -> str: the one worth reading: `EXPLAIN ANALYZE` runs it rather than only planning it, so a rewrite underneath rewrites the table for real. The parenthesised option list needs nothing here, already being blanked as a group.""" - return EXPLAIN_OPTIONS.sub(lambda match: NOT_A_NEWLINE.sub(" ", match.group()), statement) + return EXPLAIN_OPTIONS.sub(lambda match: blank(match.group()), statement) def leading_keyword(statement: str) -> re.Match[str] | None: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index fb56d032617..86216f59174 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -335,6 +335,42 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 4 + def test_marker_directly_above_a_do_block_does_not_exempt_its_body(self, tmp_path): + sql = ( + "-- data-migration-ok: seeding two default rows\n" + "DO $$\n" + "BEGIN\n" + ' INSERT INTO "Foo" ("a") VALUES (1);\n' + ' UPDATE "Foo" SET "a" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_marker_on_the_do_line_does_not_exempt_its_body(self, tmp_path): + sql = ( + "DO $$ -- data-migration-ok: bounded to one row\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marked_rewrite_does_not_exempt_a_later_one_in_the_same_block(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + ' UPDATE "Bar" SET "b" = 2;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + class TestDynamicSql: def test_execute_of_a_quoted_update_is_flagged(self, tmp_path): @@ -515,6 +551,41 @@ class TestDynamicSql: assert _keywords(tmp_path, sql) == ("DELETE",) assert _scan(tmp_path, sql)[0].line == 4 + def test_a_marker_on_an_execute_covers_its_single_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE ' -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " ';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_on_an_execute_does_not_reach_into_a_dollar_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$ -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marker_inside_a_dollar_quoted_payload_exempts_its_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + class TestExplain: def test_explain_analyze_over_an_update_is_flagged(self, tmp_path): From b573679384282e8a9ffa265a5af0b47e61a8bf89 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 21 Aug 2026 18:44:46 -0700 Subject: [PATCH 026/156] fix(proxy): keep every value of a repeated form key in get_form_data get_form_data collapsed the FormData multidict with dict(form) before the loop that rebuilds `foo[]` arrays ever ran, so a request sending timestamp_granularities[]=word and timestamp_granularities[]=segment reached the provider as ["segment"] with the first value silently dropped. Read the multidict with multi_items() instead. The test could not catch it because its mock was a plain dict carrying the same key twice, which Python collapses exactly the way the bug did. Every request.form mock that fed get_form_data now returns real FormData. --- .../proxy/common_utils/http_parsing_utils.py | 6 +- tests/test_litellm/ocr/test_ocr_file_input.py | 3 +- .../common_utils/test_http_parsing_utils.py | 74 ++++++++----------- .../test_llm_pass_through_endpoints.py | 3 +- 4 files changed, 38 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 1e4344a71f4..4cb55f6966e 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -274,10 +274,10 @@ async def get_form_data(request: Request) -> dict[str, Any]: Handles when OpenAI SDKs pass form keys as `timestamp_granularities[]="word"` instead of `timestamp_granularities=["word", "sentence"]` """ form: Final = await request.form() - form_data: Final = dict(form) parsed_form_data: Final[dict[str, Any]] = {} - for key, value in form_data.items(): - # OpenAI SDKs pass form keys as `timestamp_granularities[]="word"` instead of `timestamp_granularities=["word", "sentence"]` + # multi_items(), not dict(form): a dict drops every value but the last of a repeated key, + # which is the whole array this function exists to rebuild + for key, value in form.multi_items(): if key.endswith("[]"): clean_key = key[:-2] parsed_form_data.setdefault(clean_key, []).append(value) diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index e6216d7c580..feb98d14c03 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock import orjson import pytest +from starlette.datastructures import FormData from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type @@ -470,7 +471,7 @@ class TestProxySecurityGuard: mock_request = MagicMock() mock_request.headers = {"content-type": "multipart/form-data; boundary=---"} - mock_request.form = AsyncMock(return_value=mock_form) + mock_request.form = AsyncMock(return_value=FormData(mock_form)) result = await self._parse_multipart(mock_request) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 869d228d5a4..98af1fb7ce2 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -7,6 +7,7 @@ import orjson import pytest from fastapi import Request from fastapi.testclient import TestClient +from starlette.datastructures import FormData sys.path.insert( 0, os.path.abspath("../../../..") @@ -73,7 +74,7 @@ async def test_form_data_parsing(): test_data = {"name": "test_user", "message": "hello world"} # Mock the form method to return the test data as an awaitable - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -124,7 +125,7 @@ async def test_form_data_with_json_metadata(): } # Mock the form method to return the test data as an awaitable - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -165,7 +166,7 @@ async def test_form_data_with_invalid_json_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -188,7 +189,7 @@ async def test_form_data_without_metadata(): test_data = {"model": "whisper-1", "file": "audio.mp3", "language": "en"} # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -219,7 +220,7 @@ async def test_form_data_with_empty_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -254,7 +255,7 @@ async def test_form_data_with_dict_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -285,7 +286,7 @@ async def test_form_data_with_none_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -500,33 +501,29 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch): @pytest.mark.asyncio async def test_get_form_data(): """ - Test that get_form_data correctly handles form data with array notation. - Tests audio transcription parameters as a specific example. + A repeated `foo[]` key is how the OpenAI SDKs send a list, so every value has to + survive. `FormData`, not a dict: a dict cannot even hold the duplicate key. """ - # Create a mock request with transcription form data mock_request = MagicMock() + mock_request.form = AsyncMock( + return_value=FormData( + [ + ("file", "file_object"), + ("model", "gpt-4o-transcribe"), + ("include[]", "logprobs"), + ("language", "en"), + ("prompt", "Transcribe this audio file"), + ("response_format", "json"), + ("stream", "false"), + ("temperature", "0.2"), + ("timestamp_granularities[]", "word"), + ("timestamp_granularities[]", "segment"), + ] + ) + ) - # Create mock form data with array notation for timestamp_granularities - mock_form_data = { - "file": "file_object", # In a real request this would be an UploadFile - "model": "gpt-4o-transcribe", - "include[]": "logprobs", # Array notation - "language": "en", - "prompt": "Transcribe this audio file", - "response_format": "json", - "stream": "false", - "temperature": "0.2", - "timestamp_granularities[]": "word", # First array item - "timestamp_granularities[]": "segment", # Second array item (would overwrite in dict, but handled by the function) - } - - # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=mock_form_data) - - # Call the function being tested result = await get_form_data(mock_request) - # Verify regular form fields are preserved assert result["file"] == "file_object" assert result["model"] == "gpt-4o-transcribe" assert result["language"] == "en" @@ -534,17 +531,8 @@ async def test_get_form_data(): assert result["response_format"] == "json" assert result["stream"] == "false" assert result["temperature"] == "0.2" - - # Verify array fields are correctly parsed - assert "include" in result - assert isinstance(result["include"], list) - assert "logprobs" in result["include"] - - assert "timestamp_granularities" in result - assert isinstance(result["timestamp_granularities"], list) - # Note: In a real MultiDict, both values would be present - # But in our mock dictionary the second value overwrites the first - assert "segment" in result["timestamp_granularities"] + assert result["include"] == ["logprobs"] + assert result["timestamp_granularities"] == ["word", "segment"] def test_get_tags_from_request_body_with_metadata_tags(): @@ -958,7 +946,7 @@ class TestReadRequestBodyNonCanonicalContentType: mock_request = MagicMock() mock_request.body = AsyncMock(return_value=orjson.dumps(payload)) - mock_request.form = AsyncMock(return_value={}) + mock_request.form = AsyncMock(return_value=FormData({})) mock_request.headers = {"content-type": content_type} mock_request.scope = {} @@ -969,7 +957,7 @@ class TestReadRequestBodyNonCanonicalContentType: @pytest.mark.asyncio async def test_real_form_post_still_parsed_as_form(self): mock_request = MagicMock() - mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.form = AsyncMock(return_value=FormData({"k": "v"})) mock_request.body = AsyncMock(return_value=b"") mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} @@ -1025,7 +1013,7 @@ class TestGetRequestBody: mock_request = MagicMock() mock_request.method = "POST" mock_request.headers = {"content-type": "multipart/form-data; boundary=x"} - mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.form = AsyncMock(return_value=FormData({"k": "v"})) mock_request.scope = {} result = await get_request_body(mock_request) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 6568f6aeacf..88141fd1c90 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -13,6 +13,7 @@ import httpx import pytest from fastapi import HTTPException, Request, Response from fastapi.testclient import TestClient +from starlette.datastructures import FormData sys.path.insert( 0, os.path.abspath("../../../..") @@ -1384,7 +1385,7 @@ async def test_is_streaming_request_fn(): mock_request = Mock() mock_request.method = "POST" mock_request.headers = {"content-type": "multipart/form-data"} - mock_request.form = AsyncMock(return_value={"stream": "true"}) + mock_request.form = AsyncMock(return_value=FormData({"stream": "true"})) assert await is_streaming_request_fn(mock_request) is True From b7f8016002c080f64ca76a1136fc3fe103a5ee75 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 21 Aug 2026 18:44:57 -0700 Subject: [PATCH 027/156] test: gate the test suite on F601, B023, B025 and F632 Four more ruff rules for code the test suite runs but never checks. F601 is the one that paid: the duplicate key it flagged in a get_form_data fixture was the mock reproducing the production bug fixed in the previous commit. B025 removed two unreachable handlers, one of them a pytest.skip shadowed by an earlier `pass`, so an upstream Vertex flake reported green having asserted nothing. F632 turned an `is ""` identity check, which passes only on CPython interning, into the `== ""` it meant. B023 fixed three closures over loop variables, all latent today but one iteration-order change away from checking the last case N times. --- ruff-tests.toml | 15 ++++++++++++++ tests/code_coverage_tests/bedrock_pricing.py | 2 +- tests/load_tests/test_langsmith_load_test.py | 5 ----- .../test_amazing_vertex_completion.py | 10 ++-------- .../logging_callback_tests/test_spend_logs.py | 1 - .../test_key_generate_prisma.py | 20 ++++++++----------- .../test_ollama_completion_transformation.py | 2 +- .../llms/watsonx/test_watsonx_common_utils.py | 6 +----- tests/test_litellm/test_utils.py | 2 -- 9 files changed, 28 insertions(+), 35 deletions(-) diff --git a/ruff-tests.toml b/ruff-tests.toml index de0931f5e69..8f21ca31a6f 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -36,6 +36,17 @@ # `re.search`, so a `.` copied out of an error message is a wildcard and the block # accepts messages the author never meant to accept. Mark a real regex raw, wrap a # literal message in `re.escape`, and the pattern says which one it is +# F601 the same key literal twice in one dict. Python keeps the last value, so the +# first is dropped before the test ever runs, and a fixture that looks like it +# covers two cases covers one +# B023 a closure over a loop variable. Every closure sees the last iteration's value, +# so a per-case callback built in a loop checks the last case N times. Bind the +# value as a parameter instead +# B025 an `except` for a type an earlier `except` already catches. The second handler +# is unreachable, so the recovery or skip written there never happens +# F632 `is` against a literal. It compares identity, so it passes only where CPython +# happens to intern the value and stops meaning what it says the moment the +# value is built at runtime # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -58,4 +69,8 @@ lint.select = [ "PLR0133", "PLW0127", "RUF043", + "F601", + "B023", + "B025", + "F632", ] diff --git a/tests/code_coverage_tests/bedrock_pricing.py b/tests/code_coverage_tests/bedrock_pricing.py index b2c9e78b06c..e6219109f6d 100644 --- a/tests/code_coverage_tests/bedrock_pricing.py +++ b/tests/code_coverage_tests/bedrock_pricing.py @@ -95,7 +95,7 @@ def get_bedrock_pricing(url, providers): else: # General logic for other providers section = soup.find( - "h2", text=lambda t: t and provider.lower() in t.lower() + "h2", text=lambda t, needle=provider.lower(): t and needle in t.lower() ) if not section: pricing_data[provider] = "Provider section not found" diff --git a/tests/load_tests/test_langsmith_load_test.py b/tests/load_tests/test_langsmith_load_test.py index cf9fe526b74..40b976541a5 100644 --- a/tests/load_tests/test_langsmith_load_test.py +++ b/tests/load_tests/test_langsmith_load_test.py @@ -66,11 +66,6 @@ def test_langsmith_logging_async(): except Exception as e: pytest.fail(f"An exception occurred - {e}") - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"An exception occurred - {e}") - async def make_async_calls(metadata=None, **completion_kwargs): total_tasks = 300 diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index a52b5975f6e..53b3b2d6071 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -4207,13 +4207,7 @@ def test_gemini_google_maps_tool_simple(): ) print(f"Response: {response.model_dump_json(indent=4)}") assert response.choices[0].message.content is not None - except (litellm.RateLimitError, litellm.InternalServerError): - # Transient Vertex-side failures (rate limiting, 500 INTERNAL from the - # Google Maps grounding backend) are not LiteLLM bugs — don't fail CI. - pass - except litellm.InternalServerError: - pytest.skip( - "Google Maps Platform returned a transient 500 (upstream flake); skipping." - ) + except (litellm.RateLimitError, litellm.InternalServerError) as e: + pytest.skip(f"Transient Vertex-side failure, not a LiteLLM bug: {e}") except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/logging_callback_tests/test_spend_logs.py b/tests/logging_callback_tests/test_spend_logs.py index 709aa81f421..3aa0b3ebd90 100644 --- a/tests/logging_callback_tests/test_spend_logs.py +++ b/tests/logging_callback_tests/test_spend_logs.py @@ -148,7 +148,6 @@ def test_spend_logs_payload(model_id: Optional[str]): "completion_start_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 954146), "max_tokens": 10, "extra_body": {}, - "custom_llm_provider": "azure", "input": [ {"role": "system", "content": "you are a helpful assistant.\n"}, {"role": "user", "content": "bom dia"}, diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index efedc156429..4037107c474 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3327,16 +3327,17 @@ async def test_team_access_groups(prisma_client): request._url = URL(url="/chat/completions") + def body_reader(requested_model: str): + async def return_body() -> bytes: + return f'{{"model": "{requested_model}"}}'.encode() + + return return_body + for model in ["gpt-4o", "gemini-pro-vision"]: # Expect these to pass - async def return_body(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - request.body = return_body + request.body = body_reader(model) # use generated key to auth in print( @@ -3346,14 +3347,9 @@ async def test_team_access_groups(prisma_client): for model in ["gpt-4", "gpt-4o-mini", "gemini-experimental"]: # Expect these to fail - async def return_body_2(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - request.body = return_body_2 + request.body = body_reader(model) # use generated key to auth in print( diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index dd59cdcac1c..e746c1bfd6b 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -480,7 +480,7 @@ class TestOllamaTextCompletionResponseIterator: assert isinstance(result, ModelResponseStream) assert result.choices and result.choices[0].delta is not None assert result.choices[0].delta.content == None - assert getattr(result.choices[0].delta, "reasoning_content", None) is "" + assert getattr(result.choices[0].delta, "reasoning_content", None) == "" def test_chunk_parser_done_chunk(self): """Test that done chunks work correctly.""" diff --git a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py b/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py index be74dc40eda..0f3a6bae1ff 100644 --- a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py +++ b/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py @@ -129,11 +129,7 @@ class TestGenerateIAMToken: mock_client.reset_mock() mock_cache.reset_mock() - # Configure mock to return values based on env_keys - def get_secret_side_effect(key): - return env_keys.get(key) - - mock_get_secret_str.side_effect = get_secret_side_effect + mock_get_secret_str.side_effect = env_keys.get mock_response = MagicMock() mock_response.json.return_value = { diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index bd23ca11fbe..a01540c51e8 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -864,7 +864,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_character_above_128k_tokens": {"type": "number"}, "input_cost_per_image": {"type": "number"}, "input_cost_per_image_above_128k_tokens": {"type": "number"}, - "input_cost_per_image_token": {"type": "number"}, "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, @@ -1008,7 +1007,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): }, "bedrock_converse_supports_strict_tools": {"type": "boolean"}, "tpm": {"type": "number"}, - "provider_specific_entry": {"type": "object"}, "supported_endpoints": { "type": "array", "items": { From 64267ebd28ec35fa297a164b3c2405d6b5abc2bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:58:34 -0700 Subject: [PATCH 028/156] fix: flag an INSERT whose rows come from a parenthesised query Postgres takes the row source parenthesised, so `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table at boot. Reading only the unparenthesised text let it through: 777eb8af10 caught it, then e7dea842c3 traded it away to stop a VALUES list joined to a query by a set operation from bounding nothing. Read the top level first so set operations still count, then fall back to the whole statement when no top-level VALUES bounds the insert. `TABLE t` is a row source as much as a `SELECT` is, and it was passing too --- .../check_migrations_no_data_rewrites.py | 42 +++++++++++++------ .../test_check_migrations_no_data_rewrites.py | 34 +++++++++++++++ 2 files changed, 63 insertions(+), 13 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index cc541a70102..a8975e2529e 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -11,10 +11,12 @@ Flagged, per statement, by its leading keyword: UPDATE rewrites every matching row, and `WHERE` does not bound the scan DELETE same scan, and the dead tuples outlive the migration MERGE both of the above in one statement - INSERT only when it draws rows from a `SELECT`; an insert whose row source is - a leading `VALUES` is bounded by the rows spelled out there and passes, - scalar subqueries in that list included, while a `VALUES` reached - through a subquery or a set operation bounds nothing + INSERT only when its rows come from a query rather than a literal `VALUES` + list. The query counts wherever it sits, since Postgres takes it + parenthesised, and `TABLE t` is one as much as a `SELECT` is. An + insert bounded by a leading `VALUES` passes, scalar subqueries in that + list included, while a `VALUES` reached through a subquery or joined + to a query by a set operation bounds nothing WITH a CTE-led statement containing any of the above Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a @@ -292,24 +294,38 @@ def offending_keyword(statement: str) -> str | None: return keyword if keyword == "INSERT": - return "INSERT ... SELECT" if draws_rows_from_a_select(statement) else None + source = row_source_keyword(statement) + return None if source is None else f"INSERT ... {source}" if keyword == "WITH": nested = next((name for name in sorted(REWRITES_ROWS) if contains(statement, name)), None) if nested is not None: return f"WITH ... {nested}" - if contains(statement, "INSERT") and draws_rows_from_a_select(statement): - return "WITH ... INSERT ... SELECT" + if contains(statement, "INSERT"): + source = row_source_keyword(statement) + if source is not None: + return f"WITH ... INSERT ... {source}" return None -def draws_rows_from_a_select(statement: str) -> bool: - """Whether an `INSERT` takes its rows from a query rather than a literal list. Only a - `SELECT` the insert is built on counts, so the scalar subqueries and helper CTEs that - sit in parentheses around a `VALUES` list do not make it a rewrite, while one reached - through a set operation does.""" - return contains(strip_parens(statement), "SELECT") +def row_source_keyword(statement: str) -> str | None: + """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list + does. A query outside every parenthesis is the row source outright, including one a set + operation joins to a `VALUES` list. Failing that, a `VALUES` outside every parenthesis + is itself the row source, so the scalar subqueries and helper CTEs nested within that + list do not make the insert a rewrite. Failing both, the rows come from a parenthesised + query, which Postgres accepts and which reading only the unparenthesised text would let + through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" + outer = strip_parens(statement) + joined = row_source_in(outer) + if joined is not None: + return joined + return None if contains(outer, "VALUES") else row_source_in(statement) + + +def row_source_in(text: str) -> str | None: + return next((word for word in ("SELECT", "TABLE") if contains(text, word)), None) def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 86216f59174..15603218a27 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -125,6 +125,32 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") VALUES (1), (2) UNION ALL SELECT "id" FROM "Bar";' assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + def test_a_parenthesised_select_row_source_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_without_a_column_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_spanning_lines_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id")\n(\n SELECT "id" FROM "Bar"\n);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_over_a_values_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT * FROM (VALUES (1), (2)) AS "v"("id"));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_set_operation_over_parenthesised_selects_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT 1) UNION (SELECT 2);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_row_source_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) + + def test_a_table_named_in_the_insert_target_does_not_flag_it(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "audit table" ("id") VALUES (1);') == () + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -139,6 +165,14 @@ class TestCommonTableExpressions: sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") SELECT "id" FROM batch;' assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + def test_cte_led_insert_from_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") (SELECT "id" FROM batch);' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_cte_led_insert_into_a_values_list_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT max("id") FROM "Bar") INSERT INTO "Foo" ("id") VALUES (1);' + assert _keywords(tmp_path, sql) == () + def test_read_only_cte_passes(self, tmp_path): sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;' assert _keywords(tmp_path, sql) == () From 73af0e96921d6e0b3c855c2618facb7f2ea01cb8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:17:27 -0700 Subject: [PATCH 029/156] fix: catch two more row sources the gate let through A parenthesised query term joined to a top-level VALUES list sat behind strip_parens, so an insert reading `VALUES (1) UNION ALL (SELECT ...)` copied a whole table past the gate. A VALUES list now bounds an insert only while no set operation sits beside it at that same level. PL/pgSQL also parks dynamic SQL in a variable through a query's INTO and through the bare `=` it takes as the assignment operator, and assigned_names read neither, so a rewrite handed to a later EXECUTE went unseen. A bare `=` counts only where the words ahead of it make it an assignment rather than a test. --- .../check_migrations_no_data_rewrites.py | 77 ++++++++++--- .../test_check_migrations_no_data_rewrites.py | 102 ++++++++++++++++++ 2 files changed, 162 insertions(+), 17 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index a8975e2529e..912943265bd 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -32,9 +32,10 @@ Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the same to Postgres whether it is spelled out or handed over as a string, and so is a -literal assigned with `:=` to a variable some `EXECUTE` in the same body then runs -by name, and so is the body of a `DO` written in single quotes rather than dollar -quotes. A literal nothing runs is text, however much it reads like a statement, so +literal parked in a variable some `EXECUTE` in the same body then runs by name, +however it got there: an assignment with `:=`, the bare `=` PL/pgSQL takes as the +same operator, or a query returning it through `INTO`. So is the body of a `DO` +written in single quotes rather than dollar quotes. A literal nothing runs is text, however much it reads like a statement, so an error message naming a `DELETE` the application handles stays a message. Each literal is read on its own, so a keyword built by concatenating fragments that @@ -91,10 +92,17 @@ DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") STATEMENT = re.compile(r"[^;]+") RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) +INTO_TARGETS = re.compile( + r"\bINTO[ \t]+(?:STRICT[ \t]+)?" + r"([A-Za-z_][A-Za-z0-9_]*(?:[ \t]*,[ \t]*[A-Za-z_][A-Za-z0-9_]*)*)", + re.IGNORECASE, +) EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) +JOINS_QUERIES = ("UNION", "INTERSECT", "EXCEPT") + STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( { "INSERT", @@ -122,6 +130,10 @@ STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( } ) +GUARDS_A_CONDITION = frozenset({"IF", "ELSIF", "ELSEIF", "CASE", "WHEN", "WHILE", "EXIT", "ASSERT"}) + +OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -311,17 +323,21 @@ def offending_keyword(statement: str) -> str | None: def row_source_keyword(statement: str) -> str | None: """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list - does. A query outside every parenthesis is the row source outright, including one a set - operation joins to a `VALUES` list. Failing that, a `VALUES` outside every parenthesis - is itself the row source, so the scalar subqueries and helper CTEs nested within that - list do not make the insert a rewrite. Failing both, the rows come from a parenthesised - query, which Postgres accepts and which reading only the unparenthesised text would let - through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" + does. A query outside every parenthesis is the row source outright. Failing that, a + `VALUES` outside every parenthesis is itself the row source, so the scalar subqueries + and helper CTEs nested within that list do not make the insert a rewrite, though only + while no set operation sits beside it at that same level: one that does joins the list + to a second query term, and that term is the row source however deeply it is + parenthesised. Failing both, the rows come from a parenthesised query, which Postgres + accepts and which reading only the unparenthesised text would let through: + `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: return joined - return None if contains(outer, "VALUES") else row_source_in(statement) + if contains(outer, "VALUES") and not any(contains(outer, word) for word in JOINS_QUERIES): + return None + return row_source_in(statement) def row_source_in(text: str) -> str | None: @@ -339,13 +355,40 @@ def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: def assigned_names(statement: str) -> frozenset[str]: - """The candidate variable names an assignment writes to, taken as every word ahead of the - `:=`. A declaration carries its type and sometimes a leading `DECLARE` alongside the name, - and none of that is worth parsing when the only question is which name is executed.""" - head, separator, _ = statement.partition(":=") - if not separator: - return frozenset() - return frozenset(word.group().lower() for word in FIRST_WORD.finditer(head)) + """The candidate variable names a statement writes to, taken as every word ahead of the + assignment operator. A declaration carries its type and sometimes a leading `DECLARE` + alongside the name, and none of that is worth parsing when the only question is which + name is executed. PL/pgSQL spells that operator `:=` and takes a bare `=` as the same + thing, so both count, the second only where `assigns_rather_than_compares` reads it as + an assignment. A query assigns through the target list after its `INTO` instead, which + is how a rewrite reaches a variable with neither operator appearing at all.""" + names: set[str] = set() + + head, operator, _ = statement.partition(":=") + assigns = bool(operator) + if not assigns: + head, operator, _ = statement.partition("=") + assigns = bool(operator) and assigns_rather_than_compares(head) + if assigns: + names.update(word.group().lower() for word in FIRST_WORD.finditer(head)) + + for targets in INTO_TARGETS.finditer(statement): + names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) + + return frozenset(names) + + +def assigns_rather_than_compares(head: str) -> bool: + """Whether the bare `=` this text runs up to writes a variable or tests one. Only the + words ahead of it tell the two apart: an assignment is reached with a name and perhaps a + type, while a comparison is reached either through a statement carrying its own keyword + or through a word that guards a condition. Those words stop counting once something + opens a block after them, since a `THEN` ends the condition its `IF` began and the + assignment that follows on the same line is an assignment like any other.""" + words = [word.group().upper() for word in FIRST_WORD.finditer(head)] + opened = max((index for index, word in enumerate(words) if word in OPENS_A_BLOCK), default=-1) + reached = set(words[opened + 1 :]) + return not (reached & STATEMENT_KEYWORDS) and not (reached & GUARDS_A_CONDITION) def executed_names(masked: str) -> frozenset[str]: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 15603218a27..f70d49ddb83 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -145,6 +145,22 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") (SELECT 1) UNION (SELECT 2);' assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + def test_a_values_list_joined_to_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_excepting_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) EXCEPT (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_table_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_set_operation_inside_a_values_list_does_not_flag_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES ((SELECT 1 UNION SELECT 2 LIMIT 1));' + assert _keywords(tmp_path, sql) == () + def test_a_table_row_source_is_flagged(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) @@ -469,6 +485,92 @@ class TestDynamicSql: assert _keywords(tmp_path, sql) == ("DELETE",) assert _scan(tmp_path, sql)[0].line == 5 + def test_a_rewrite_selected_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1' INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_strict_target_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_assigned_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_assigned_with_a_bare_equals_after_then_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " IF true THEN stmt = 'DELETE FROM \"Foo\"'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_declared_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text;\n" + "BEGIN\n" + " SELECT 'UPDATE of legacy rows is skipped' INTO msg;\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_comparing_an_executed_variable_does_not_flag_the_comparison(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'DELETE FROM \"Foo\"' THEN RAISE NOTICE 'never'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_ddl_assigned_to_a_variable_passes(self, tmp_path): sql = "DO $$\nDECLARE\n stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n EXECUTE stmt;\nEND $$;" assert _keywords(tmp_path, sql) == () From d6f9bce4bb7c0783cb994f3ef165c182069557dc Mon Sep 17 00:00:00 2001 From: milan Date: Sat, 22 Aug 2026 04:05:02 +0000 Subject: [PATCH 030/156] fix(a2a): normalize agent card protocolBinding casing before transport match Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/a2a_protocol/card_resolver.py | 28 ++++++++++++ litellm/a2a_protocol/main.py | 13 +++--- tests/test_litellm/a2a_protocol/test_main.py | 45 +++++++++++++++++--- 3 files changed, 76 insertions(+), 10 deletions(-) diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index d14d892256b..8070d88b761 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger @@ -48,6 +49,33 @@ def is_localhost_or_internal_url(url: str | None) -> bool: return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS) +_CANONICAL_PROTOCOL_BINDINGS: Final = MappingProxyType( + { + "jsonrpc": "JSONRPC", + "http+json": "HTTP+JSON", + "grpc": "GRPC", + } +) + + +def normalize_agent_card_protocol_bindings(agent_card: "AgentCard") -> "AgentCard": + """ + Canonicalize protocolBinding casing on the card's supported interfaces. + + Some A2A servers (e.g. LangGraph Platform) serve agent cards with lowercase + bindings like "jsonrpc", but a2a-sdk's ClientFactory matches bindings + case-sensitively against its uppercase TransportProtocol constants and fails + with "no compatible transports found." for spec-adjacent casings. + """ + interfaces: Final = getattr(agent_card, "supported_interfaces", None) or () + for interface in interfaces: + binding: str = getattr(interface, "protocol_binding", "") or "" + canonical = _CANONICAL_PROTOCOL_BINDINGS.get(binding.lower()) + if canonical is not None and binding != canonical: + interface.protocol_binding = canonical + return agent_card + + def get_agent_card_url(agent_card: "AgentCard") -> str | None: """Return the agent endpoint URL from the resolved SDK card.""" url: Final = getattr(agent_card, "url", None) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 1c6ebf0b95c..2d4267cc0db 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -73,6 +73,7 @@ except ImportError: from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, get_agent_card_url, + normalize_agent_card_protocol_bindings, ) from litellm.a2a_protocol.exception_mapping_utils import ( handle_a2a_localhost_retry, @@ -782,13 +783,17 @@ async def create_a2a_client( if extra_headers: verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) + resolver: Final = A2ACardResolver(httpx_client=httpx_client, base_url=base_url) + agent_card: Final = normalize_agent_card_protocol_bindings( + await resolver.get_agent_card(http_kwargs={"headers": extra_headers} if extra_headers else None) + ) + a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall] - base_url, + agent_card, client_config=ClientConfig( # pyright: ignore[reportOptionalCall] httpx_client=httpx_client, streaming=streaming, ), - resolver_http_kwargs={"headers": extra_headers} if extra_headers else None, ) # Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse # the configured httpx client and this agent's headers without excavating @@ -799,9 +804,7 @@ async def create_a2a_client( if extra_headers else None ) - agent_card: Final = getattr(a2a_client, "_card", None) - if agent_card is not None: - a2a_client._litellm_agent_card = agent_card + a2a_client._litellm_agent_card = agent_card verbose_logger.info("A2A client created for %s", base_url) diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 08f6b9f25bb..29aa3aaabd7 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -176,10 +176,32 @@ _AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"} _AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"} +_V1_RPC_REPLY = { + "jsonrpc": "2.0", + "id": "reply", + "result": {"message": {"messageId": "reply-1", "role": "ROLE_AGENT", "parts": [{"text": "pong"}]}}, +} + + +_LOWERCASE_BINDING_CARD = { + "name": "langgraph-agent", + "version": "1.0.0", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [], + "supportedInterfaces": [ + {"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"} + ], +} + + class _RequestRecorder: """Records the headers httpx put on the wire, per outbound request.""" - def __init__(self): + def __init__(self, card=_AGENT_CARD, rpc_reply=_RPC_REPLY): + self.card = card + self.rpc_reply = rpc_reply self.card_requests = [] self.rpc_requests = [] self.client = None @@ -188,23 +210,23 @@ class _RequestRecorder: headers = {k.lower(): v for k, v in request.headers.items()} if request.method == "GET": self.card_requests.append(headers) - return httpx.Response(200, json=_AGENT_CARD) + return httpx.Response(200, json=self.card) self.rpc_requests.append(headers) - return httpx.Response(200, json=_RPC_REPLY) + return httpx.Response(200, json=self.rpc_reply) def _a2a_client_cache_key(timeout: float) -> str: return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider -async def _seed_shared_a2a_client() -> _RequestRecorder: +async def _seed_shared_a2a_client(card=_AGENT_CARD, rpc_reply=_RPC_REPLY) -> _RequestRecorder: """Put the one A2A client the cache will hand out behind a mock transport. Seeding has to happen on the test's own event loop, because the client cache keys on it. The injected client is a real httpx.AsyncClient, so the merge of per-request headers over client defaults, which is what these tests are about, stays real. """ - recorder = _RequestRecorder() + recorder = _RequestRecorder(card=card, rpc_reply=rpc_reply) handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) owned_client = handler.client handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) @@ -311,6 +333,19 @@ async def test_streaming_send_carries_only_its_own_caller_headers(isolated_clien assert received["b"]["x-tenant"] == "tenant-b" +@pytest.mark.asyncio +async def test_lowercase_protocol_binding_in_agent_card_still_gets_a_client(isolated_client_cache): + """LangGraph Platform serves cards with protocolBinding "jsonrpc"; a2a-sdk matches + bindings case-sensitively, so without normalization client creation raises + ValueError("no compatible transports found.").""" + await _seed_shared_a2a_client(card=_LOWERCASE_BINDING_CARD, rpc_reply=_V1_RPC_REPLY) + + a2a_client = await create_a2a_client(base_url="http://127.0.0.1:9") + response = await _send_message(a2a_client, _send_request("lc")) + + assert type(response.root.result).__name__ == "Message" + + @pytest.mark.asyncio async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cache): """Agent cards can sit behind the same auth as the agent, so the card fetch must stay From 329202004ef15d77e79b4c7c234c208567c22108 Mon Sep 17 00:00:00 2001 From: milan Date: Sat, 22 Aug 2026 04:17:16 +0000 Subject: [PATCH 031/156] test(a2a): stub card resolver in create_a2a_client unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_agent_header_isolation.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py index 5193ad989dc..30095796a7f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py @@ -245,13 +245,21 @@ def _fake_get_async_httpx_client_factory(captured_calls: list): return _fake_get_async_httpx_client -async def _fake_create_client(base_url, client_config=None, **kwargs): +async def _fake_create_client(agent_card, client_config=None, **kwargs): client = MagicMock() if client_config is not None: client._litellm_httpx_client = client_config.httpx_client return client +def _fake_card_resolver(httpx_client, base_url, **kwargs): + resolver = MagicMock() + card = MagicMock() + card.supported_interfaces = () + resolver.get_agent_card = AsyncMock(return_value=card) + return resolver + + @pytest.mark.asyncio async def test_create_a2a_client_leaves_the_shared_client_untouched(): """ @@ -276,6 +284,10 @@ async def test_create_a2a_client_leaves_the_shared_client_untouched(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client( base_url="http://agent-a:9999", @@ -321,6 +333,10 @@ async def test_create_a2a_client_default_timeout_matches_constant(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client(base_url="http://127.0.0.1:9") @@ -352,6 +368,10 @@ async def test_create_a2a_client_explicit_timeout_overrides_default(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client(base_url="http://127.0.0.1:9", timeout=42.5) From 1f57e7ea19f7889907956683dd523a88cdece74e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 08:49:02 -0700 Subject: [PATCH 032/156] fix: stop reading an INSERT target table as an assignment target Scanning a statement's own literals for SQL only makes sense when the name before INTO is a variable the body later executes. INSERT INTO names a table there, so an insert into a table sharing a variable's name was flagged for whatever its column values happened to spell. An INSERT that really does assign reaches INTO through RETURNING, which the preceding word separates. Also names the scope boundary in the module docstring: the ban is on row-rewriting DML, not on everything whose cost scales with table size. --- .../check_migrations_no_data_rewrites.py | 20 +++++++++++++++ .../test_check_migrations_no_data_rewrites.py | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 912943265bd..cbc24fe3660 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -6,6 +6,13 @@ anything whose cost scales with existing table size turns into downtime. A singl `UPDATE` with no batching over a spend-log-sized table is minutes of unavailability plus a doubled heap that plain autovacuum will not give back. +What is banned is the row-rewriting DML behind that, not everything whose cost +scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is +not binary coercible, and a volatile `DEFAULT` on a new column all read the whole +table and all pass. That is deliberate: a rule wide enough to reach them fires on +most ordinary migrations, and a marker everyone adds by reflex stops carrying +information. The outage this was written for was a backfill. + Flagged, per statement, by its leading keyword: UPDATE rewrites every matching row, and `WHERE` does not bound the scan @@ -97,6 +104,7 @@ INTO_TARGETS = re.compile( r"([A-Za-z_][A-Za-z0-9_]*(?:[ \t]*,[ \t]*[A-Za-z_][A-Za-z0-9_]*)*)", re.IGNORECASE, ) +PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -373,11 +381,23 @@ def assigned_names(statement: str) -> frozenset[str]: names.update(word.group().lower() for word in FIRST_WORD.finditer(head)) for targets in INTO_TARGETS.finditer(statement): + if names_a_table(statement[: targets.start()]): + continue names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) return frozenset(names) +def names_a_table(before: str) -> bool: + """Whether the `INTO` this text runs up to introduces a table rather than a query's + target list. `INSERT INTO` is the one that does, and reading its table as somewhere a + string was parked would have an insert scanned for the SQL its own literals spell out. + An `INSERT` that really does assign reaches its `INTO` through a `RETURNING` list, so + the word immediately before is what separates the two.""" + word = PRECEDING_WORD.search(before) + return word is not None and word.group(1).upper() == "INSERT" + + def assigns_rather_than_compares(head: str) -> bool: """Whether the bare `=` this text runs up to writes a variable or tests one. Only the words ahead of it tell the two apart: an assignment is reached with a name and perhaps a diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index f70d49ddb83..7ed8b9978a3 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -547,6 +547,31 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_an_insert_target_table_is_not_read_as_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " audit text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " INSERT INTO audit (note) VALUES ('DELETE FROM \"Foo\" is left to the app');\n" + " EXECUTE audit;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_returned_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " INSERT INTO \"Log\" (\"sql\") VALUES ('UPDATE \"Foo\" SET \"a\" = 1')\n" + " RETURNING \"sql\" INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): sql = ( "DO $$\n" From 813d3a991fabb7d75a314733e9fb88b0d5b94515 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 09:55:18 -0700 Subject: [PATCH 033/156] fix: read every assignment operator, not only the statement's first A bare = was found with partition, so a comparison earlier on the line took the one slot and the assignment after it went unread. INTO targets and the name an EXECUTE runs are also allowed to sit on the next line now. --- .../check_migrations_no_data_rewrites.py | 55 ++++++----- .../test_check_migrations_no_data_rewrites.py | 91 +++++++++++++++++++ 2 files changed, 124 insertions(+), 22 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index cbc24fe3660..9e6df46899b 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -98,12 +98,13 @@ MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTIL DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") STATEMENT = re.compile(r"[^;]+") -RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) +RUN_BY_NAME = re.compile(r"\bEXECUTE\s+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) INTO_TARGETS = re.compile( - r"\bINTO[ \t]+(?:STRICT[ \t]+)?" - r"([A-Za-z_][A-Za-z0-9_]*(?:[ \t]*,[ \t]*[A-Za-z_][A-Za-z0-9_]*)*)", + r"\bINTO\s+(?:STRICT\s+)?" + r"([A-Za-z_][A-Za-z0-9_]*(?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)", re.IGNORECASE, ) +ASSIGNS = re.compile(r":=|(?!:=])=(?!=)") PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) @@ -368,17 +369,17 @@ def assigned_names(statement: str) -> frozenset[str]: alongside the name, and none of that is worth parsing when the only question is which name is executed. PL/pgSQL spells that operator `:=` and takes a bare `=` as the same thing, so both count, the second only where `assigns_rather_than_compares` reads it as - an assignment. A query assigns through the target list after its `INTO` instead, which - is how a rewrite reaches a variable with neither operator appearing at all.""" + an assignment. Every operator in the statement is read rather than only the first, since + a comparison earlier on the line would otherwise claim the one slot and hide the + assignment after it: `IF n = 1 THEN stmt = '...'` writes `stmt` at its second `=`. A + query assigns through the target list after its `INTO` instead, which is how a rewrite + reaches a variable with neither operator appearing at all.""" names: set[str] = set() - head, operator, _ = statement.partition(":=") - assigns = bool(operator) - if not assigns: - head, operator, _ = statement.partition("=") - assigns = bool(operator) and assigns_rather_than_compares(head) - if assigns: - names.update(word.group().lower() for word in FIRST_WORD.finditer(head)) + for operator in ASSIGNS.finditer(statement): + reached = reached_words(statement[: operator.start()]) + if operator.group() == ":=" or assigns_rather_than_compares(reached): + names.update(word.lower() for word in reached) for targets in INTO_TARGETS.finditer(statement): if names_a_table(statement[: targets.start()]): @@ -398,22 +399,32 @@ def names_a_table(before: str) -> bool: return word is not None and word.group(1).upper() == "INSERT" -def assigns_rather_than_compares(head: str) -> bool: - """Whether the bare `=` this text runs up to writes a variable or tests one. Only the - words ahead of it tell the two apart: an assignment is reached with a name and perhaps a - type, while a comparison is reached either through a statement carrying its own keyword - or through a word that guards a condition. Those words stop counting once something - opens a block after them, since a `THEN` ends the condition its `IF` began and the - assignment that follows on the same line is an assignment like any other.""" +def reached_words(head: str) -> tuple[str, ...]: + """The words an assignment operator is reached through, which is everything since the last + word to open a block. A `THEN` ends the condition its `IF` began, so nothing ahead of it + describes what follows, and neither the name being written nor the keywords that would + mark a comparison ever sit further back than that.""" words = [word.group().upper() for word in FIRST_WORD.finditer(head)] opened = max((index for index, word in enumerate(words) if word in OPENS_A_BLOCK), default=-1) - reached = set(words[opened + 1 :]) - return not (reached & STATEMENT_KEYWORDS) and not (reached & GUARDS_A_CONDITION) + return tuple(words[opened + 1 :]) + + +def assigns_rather_than_compares(reached: tuple[str, ...]) -> bool: + """Whether a bare `=` reached through these words writes a variable or tests one. They are + all that tells the two apart: an assignment is reached with a name and perhaps a type, + while a comparison is reached either through a statement carrying its own keyword or + through a word that guards a condition.""" + words = set(reached) + return not (words & STATEMENT_KEYWORDS) and not (words & GUARDS_A_CONDITION) def executed_names(masked: str) -> frozenset[str]: """The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps - an `EXECUTE` written inside a comment or a string from counting.""" + an `EXECUTE` written inside a comment or a string from counting. Masking blanks a literal + in place rather than removing it, so `EXECUTE '...'` can leave the word after it looking + like the name being run. Reaching that word means crossing no semicolon, which leaves only + the syntax `INTO`, `USING` and `END`, and no assignment ever writes one of those, so the + stray name has nothing to match.""" return frozenset(match.group(1).lower() for match in RUN_BY_NAME.finditer(masked)) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 7ed8b9978a3..fe53e1ac823 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -572,6 +572,97 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_rewrite_assigned_past_an_earlier_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " IF total = 1 THEN stmt = 'DELETE FROM \"Foo\" WHERE true'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_rewrite_assigned_past_a_loop_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 3;\n" + "BEGIN\n" + " WHILE total >= 1 LOOP stmt = 'UPDATE \"Foo\" SET \"a\" = 1'; END LOOP;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1'\n" + " INTO\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_strict_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_executed_by_a_name_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE\n" + " stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_compared_against_is_not_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'UPDATE \"Foo\" SET \"a\" = 1' THEN\n" + " RAISE NOTICE 'the application owns that one';\n" + " END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_passed_to_execute_as_a_parameter_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE 'INSERT INTO \"Log\" (\"sql\") VALUES ($1)' USING stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): sql = ( "DO $$\n" From fdcf867dbfe1fdab304e3c690eb0e50365b254a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 10:38:55 -0700 Subject: [PATCH 034/156] fix: read the one assignment a statement holds, and the loop that walks a query Reading every operator let a comparison beside an assignment look like one. `ok := n = 1 AND stmt = ''` registered `stmt` as written, which collided with the `EXECUTE stmt` further down and flagged a block that rewrites nothing. A statement holds one assignment at most, so the search now stops at the first operator that reads as one: everything after it is the expression being assigned, where an `=` only ever compares. Nine shapes were flagged this way, a cast, a `coalesce`, a `format`, a named-argument arrow and the rest, and all of them are valid PL/pgSQL that leaves the table untouched. `INTO` and `USING` no longer count as names an `EXECUTE` runs. Masking blanks a literal in place, so `EXECUTE '' INTO n` left `INTO` looking like the name being run, and an ordinary query reaching the same word collided with it. The docstring claiming that collision was impossible was wrong, and both words are now dropped instead. A loop is a fourth way a literal reaches a variable. `FOR stmt IN SELECT '' LOOP EXECUTE stmt` empties the table and the gate passed it, so the target of a `FOR` or a `FOREACH` is read as assigned too. Reading each statement once rather than once per operator also drops the cost of a statement with thousands of them from seconds to milliseconds. --- .../check_migrations_no_data_rewrites.py | 103 +++++++++----- .../test_check_migrations_no_data_rewrites.py | 134 ++++++++++++++++++ 2 files changed, 198 insertions(+), 39 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 9e6df46899b..dcb81593e3e 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -41,8 +41,9 @@ hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads t same to Postgres whether it is spelled out or handed over as a string, and so is a literal parked in a variable some `EXECUTE` in the same body then runs by name, however it got there: an assignment with `:=`, the bare `=` PL/pgSQL takes as the -same operator, or a query returning it through `INTO`. So is the body of a `DO` -written in single quotes rather than dollar quotes. A literal nothing runs is text, however much it reads like a statement, so +same operator, a query returning it through `INTO`, or a loop walking the query it +came out of. So is the body of a `DO` written in single quotes rather than dollar +quotes. A literal nothing runs is text, however much it reads like a statement, so an error message naming a `DELETE` the application handles stays a message. Each literal is read on its own, so a keyword built by concatenating fragments that @@ -104,7 +105,8 @@ INTO_TARGETS = re.compile( r"([A-Za-z_][A-Za-z0-9_]*(?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)", re.IGNORECASE, ) -ASSIGNS = re.compile(r":=|(?!:=])=(?!=)") +LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE) +WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?!:=])=(?![=>])") PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) @@ -143,6 +145,8 @@ GUARDS_A_CONDITION = frozenset({"IF", "ELSIF", "ELSEIF", "CASE", "WHEN", "WHILE" OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) +NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -364,28 +368,21 @@ def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: def assigned_names(statement: str) -> frozenset[str]: - """The candidate variable names a statement writes to, taken as every word ahead of the - assignment operator. A declaration carries its type and sometimes a leading `DECLARE` - alongside the name, and none of that is worth parsing when the only question is which - name is executed. PL/pgSQL spells that operator `:=` and takes a bare `=` as the same - thing, so both count, the second only where `assigns_rather_than_compares` reads it as - an assignment. Every operator in the statement is read rather than only the first, since - a comparison earlier on the line would otherwise claim the one slot and hide the - assignment after it: `IF n = 1 THEN stmt = '...'` writes `stmt` at its second `=`. A - query assigns through the target list after its `INTO` instead, which is how a rewrite - reaches a variable with neither operator appearing at all.""" - names: set[str] = set() - - for operator in ASSIGNS.finditer(statement): - reached = reached_words(statement[: operator.start()]) - if operator.group() == ":=" or assigns_rather_than_compares(reached): - names.update(word.lower() for word in reached) + """The candidate variable names a statement writes to. An assignment is read as every + word ahead of its operator, since a declaration carries its type and sometimes a leading + `DECLARE` alongside the name, and none of that is worth parsing when the only question + is which name is executed. A query assigns through the target list after its `INTO` + instead, and a loop through the variable it walks its query with, which is how a rewrite + reaches a variable with no operator appearing at all.""" + names = {word.lower() for word in assignment_reach(statement)} for targets in INTO_TARGETS.finditer(statement): if names_a_table(statement[: targets.start()]): continue names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) + names.update(loop.group(1).lower() for loop in LOOP_TARGET.finditer(statement)) + return frozenset(names) @@ -399,33 +396,61 @@ def names_a_table(before: str) -> bool: return word is not None and word.group(1).upper() == "INSERT" -def reached_words(head: str) -> tuple[str, ...]: - """The words an assignment operator is reached through, which is everything since the last - word to open a block. A `THEN` ends the condition its `IF` began, so nothing ahead of it - describes what follows, and neither the name being written nor the keywords that would - mark a comparison ever sit further back than that.""" - words = [word.group().upper() for word in FIRST_WORD.finditer(head)] - opened = max((index for index, word in enumerate(words) if word in OPENS_A_BLOCK), default=-1) - return tuple(words[opened + 1 :]) +def assignment_reach(statement: str) -> tuple[str, ...]: + """The words the statement's assignment is reached through, empty where it holds none. + PL/pgSQL spells the operator `:=` and takes a bare `=` as the same thing, so both count, + the second only where none of the words reached so far `marks_a_comparison`. The search + stops at the first operator that reads as an assignment, because a statement holds one + at most and everything after it is the expression being assigned, where an `=` only ever + compares: that is what keeps `ok := stmt = ''` from reading as a write to `stmt`. + What comes before can still be a comparison the assignment sits behind, as in + `IF n = 1 THEN stmt = ''`, and a word opening a block ends what it is reached + through, since nothing ahead of the `THEN` describes what follows it.""" + reached: list[str] = [] + compares = False + + for token in WORD_OR_ASSIGN.finditer(statement): + word = token.group().upper() + + if word == ":=": + return tuple(reached) + + if word == "=": + if not compares: + return tuple(reached) + continue + + if word in OPENS_A_BLOCK: + reached.clear() + compares = False + continue + + reached.append(word) + compares = compares or marks_a_comparison(word) + + return () -def assigns_rather_than_compares(reached: tuple[str, ...]) -> bool: - """Whether a bare `=` reached through these words writes a variable or tests one. They are - all that tells the two apart: an assignment is reached with a name and perhaps a type, - while a comparison is reached either through a statement carrying its own keyword or - through a word that guards a condition.""" - words = set(reached) - return not (words & STATEMENT_KEYWORDS) and not (words & GUARDS_A_CONDITION) +def marks_a_comparison(word: str) -> bool: + """Whether reaching a bare `=` through this word means the operator tests a variable + rather than writing one. These are all that tell the two apart: an assignment is reached + with a name and perhaps a type, while a comparison is reached either through a statement + carrying its own keyword or through a word that guards a condition.""" + return word in STATEMENT_KEYWORDS or word in GUARDS_A_CONDITION def executed_names(masked: str) -> frozenset[str]: """The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps an `EXECUTE` written inside a comment or a string from counting. Masking blanks a literal - in place rather than removing it, so `EXECUTE '...'` can leave the word after it looking - like the name being run. Reaching that word means crossing no semicolon, which leaves only - the syntax `INTO`, `USING` and `END`, and no assignment ever writes one of those, so the - stray name has nothing to match.""" - return frozenset(match.group(1).lower() for match in RUN_BY_NAME.finditer(masked)) + in place rather than removing it, so `EXECUTE '...'` leaves whatever follows the literal + looking like the name being run. Only `INTO` and `USING` can sit there, since the syntax + allows nothing else between an `EXECUTE` and the semicolon ending it, and neither is ever + a variable, so both are dropped rather than left to collide with a query reaching one.""" + return frozenset( + match.group(1).lower() + for match in RUN_BY_NAME.finditer(masked) + if match.group(1).upper() not in NEVER_A_VARIABLE + ) def leads_with(statement: str, keyword: str) -> bool: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index fe53e1ac823..b93e8422933 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -663,6 +663,140 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == () + def test_a_rewrite_compared_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total = 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " RAISE NOTICE 'purge script? %', ok;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_named_as_an_argument_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := probe_match(subject => stmt, wanted => 'DELETE FROM \"Foo\"');\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_after_a_wider_comparison_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total >= 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_assigned_through_a_case_expression_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " stmt := CASE WHEN total = 1 THEN 'DELETE FROM \"Foo\"' ELSE 'SELECT 1' END;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_query_reaching_into_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\"'\n" + " INTO total;\n" + " SELECT (CASE WHEN total > 0 THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\"\n" + " WHERE \"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_reaching_using_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\" WHERE \"a\" = $1'\n" + " USING 'k1';\n" + " SELECT (CASE WHEN true THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\" x JOIN \"Foo\" y USING (\"a\")\n" + " WHERE x.\"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_walked_by_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOR stmt IN SELECT 'DELETE FROM \"Foo\"' LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_walked_by_a_foreach_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOREACH stmt IN ARRAY ARRAY['UPDATE \"Foo\" SET \"a\" = 1'] LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_loop_over_a_query_running_nothing_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " rec record;\n" + "BEGIN\n" + " FOR rec IN SELECT \"a\" FROM \"Foo\" LOOP\n" + " RAISE NOTICE 'the DELETE FROM \"Foo\" path is the application''s: %', rec;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): sql = ( "DO $$\n" From 6d6a2fcfb892ff69f968b0ec8407a9d08618f84b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 10:57:00 -0700 Subject: [PATCH 035/156] fix: match a marker to the statement it is written against, not to its line --- .../check_migrations_no_data_rewrites.py | 78 +++++++++++++------ .../test_check_migrations_no_data_rewrites.py | 13 ++++ 2 files changed, 67 insertions(+), 24 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index dcb81593e3e..7b0029c73cb 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -172,16 +172,38 @@ class Violation: @dataclass(frozen=True, slots=True) -class Markers: - lines: frozenset[int] - standalone: frozenset[int] +class Marker: + start: int + end: int + standalone: bool - def exempt(self, first: int, last: int) -> bool: - """Whether a statement spanning `first` to `last` carries a marker. A marker alone on - its line speaks for the statement below it, which is how one written above a rewrite - exempts it. A marker sharing its line with the statement it follows speaks for that - statement only, so the next statement down does not inherit the exemption.""" - return any(line in self.lines for line in range(first, last + 1)) or first - 1 in self.standalone + +@dataclass(frozen=True, slots=True) +class Markers: + sql: str + written: tuple[Marker, ...] + + def exempt(self, start: int, end: int) -> bool: + """Whether the statement spanning `start` to `end` carries a marker.""" + return any(self.speaks_for(marker, start, end) for marker in self.written) + + def speaks_for(self, marker: Marker, start: int, end: int) -> bool: + """Whether a marker is written against this statement. One alone on its line speaks for + the statement below it, which is how a marker written above a rewrite exempts it, and one + sharing its line with code speaks for the statement it follows. Either is matched by where + it sits rather than by the line it lands on, so a second statement sharing that line does + not inherit the exemption. A marker inside a statement speaks for it whichever kind it is, + which is how one on the opening line of a long statement still covers the whole of it.""" + if start <= marker.start < end: + return True + if marker.standalone: + return self.only_separators(marker.end, start) + return self.only_separators(end, marker.start) + + def only_separators(self, start: int, end: int) -> bool: + """Whether nothing but statement separators lie between two points, which is what makes a + marker and a statement adjacent whatever whitespace and line breaks sit between them.""" + return start <= end and not self.sql[start:end].strip(" \t\r\n;") def blank(text: str) -> str: @@ -463,16 +485,17 @@ def contains(statement: str, keyword: str) -> bool: def read_markers(sql: str) -> Markers: - lines: set[int] = set() - standalone: set[int] = set() + return Markers( + sql, + tuple( + Marker(match.start(), match.end(), alone_on_its_line(sql, match.start())) + for match in MARKER.finditer(sql) + ), + ) - for match in MARKER.finditer(sql): - line = sql.count("\n", 0, match.start()) + 1 - lines.add(line) - if not sql[sql.rfind("\n", 0, match.start()) + 1 : match.start()].strip(): - standalone.add(line) - return Markers(frozenset(lines), frozenset(standalone)) +def alone_on_its_line(sql: str, start: int) -> bool: + return not sql[sql.rfind("\n", 0, start) + 1 : start].strip() def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]: @@ -482,16 +505,16 @@ def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]: def scan_region( document: str, region: str, migration: str, markers: Markers, offset: int ) -> Iterator[Violation]: - """Violations in one region of `document`, whose text begins at `offset`. Lines are - always counted against the whole document, so a statement nested in a dollar-quoted - body reports its real file line and lines up with the markers read from that file.""" + """Violations in one region of `document`, whose text begins at `offset`. Positions are + always counted against the whole document, so a statement nested in a dollar-quoted body + reports its real file line and lines up with the markers read from that file.""" masked, bodies, literals = mask(region) executed = executed_names(masked) for match in STATEMENT.finditer(masked): - first = line_of(document, offset + keyword_start(match)) - last = line_of(document, offset + match.end()) - exempt = markers.exempt(first, last) + start = offset + statement_start(match) + end = offset + match.end() + exempt = markers.exempt(start, end) if hands_off_sql(match.group(), executed) and not exempt: for start, end in literals: @@ -501,12 +524,19 @@ def scan_region( keyword = offending_keyword(match.group()) if keyword is None or exempt: continue - yield Violation(migration, first, keyword) + yield Violation(migration, line_of(document, offset + keyword_start(match)), keyword) for start, end in bodies: yield from scan_region(document, region[start:end], migration, markers, offset + start) +def statement_start(statement: re.Match[str]) -> int: + """Where the statement's own text begins, past the whitespace and blanked comments it picked + up from whatever sat between it and the statement before it, one of which can be a marker.""" + text = statement.group() + return statement.start() + len(text) - len(text.lstrip()) + + def keyword_start(statement: re.Match[str]) -> int: word = leading_keyword(statement.group()) return statement.start() + (0 if word is None else word.start()) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index b93e8422933..1573a1b5706 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -375,6 +375,19 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 1 + def test_a_trailing_marker_exempts_only_the_statement_it_follows(self, tmp_path): + sql = 'DELETE FROM "Foo" WHERE "a" = 1; UPDATE "Bar" SET "b" = 2; -- data-migration-ok: one row' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_above_a_shared_line_exempts_only_the_first_statement_on_it(self, tmp_path): + sql = '-- data-migration-ok: one row\nUPDATE "Foo" SET "a" = 1; DELETE FROM "Bar" WHERE "b" = 2;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_on_the_opening_line_of_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" -- data-migration-ok: one row\n SET "a" = 1;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 3 + def test_marker_above_a_do_block_does_not_exempt_a_rewrite_inside_it(self, tmp_path): sql = ( "-- data-migration-ok: bounded, this belongs to the insert below\n" From dee93e2d4842d62531b17eeeb9ec9b37bc30508f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:12:26 -0700 Subject: [PATCH 036/156] fix: keep a marker on its own line bound to the statement directly below it --- .../check_migrations_no_data_rewrites.py | 10 ++++++++-- .../test_check_migrations_no_data_rewrites.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 7b0029c73cb..3d57691cea4 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -197,12 +197,18 @@ class Markers: if start <= marker.start < end: return True if marker.standalone: - return self.only_separators(marker.end, start) + return self.on_the_line_below(marker.end, start) return self.only_separators(end, marker.start) + def on_the_line_below(self, start: int, end: int) -> bool: + """Whether a marker on its own line is written directly above the statement, which means + one line break and nothing else that carries meaning. A blank line between the two leaves + the marker reading as a note about the file rather than a bound on what follows it.""" + return self.only_separators(start, end) and self.sql[start:end].count("\n") == 1 + def only_separators(self, start: int, end: int) -> bool: """Whether nothing but statement separators lie between two points, which is what makes a - marker and a statement adjacent whatever whitespace and line breaks sit between them.""" + marker and the statement it follows adjacent however they are laid out.""" return start <= end and not self.sql[start:end].strip(" \t\r\n;") diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 1573a1b5706..5413f8cdc62 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -375,6 +375,11 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 1 + def test_a_marker_a_blank_line_above_a_statement_does_not_exempt_it(self, tmp_path): + sql = '-- data-migration-ok: one row\n\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + def test_a_trailing_marker_exempts_only_the_statement_it_follows(self, tmp_path): sql = 'DELETE FROM "Foo" WHERE "a" = 1; UPDATE "Bar" SET "b" = 2; -- data-migration-ok: one row' assert _keywords(tmp_path, sql) == ("DELETE",) @@ -410,6 +415,11 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 5 + def test_marker_directly_above_a_one_line_do_block_does_not_exempt_its_body(self, tmp_path): + sql = '-- data-migration-ok: seeding one default row\nDO $$ BEGIN UPDATE "Foo" SET "a" = 1; END $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + def test_marker_on_the_do_line_does_not_exempt_its_body(self, tmp_path): sql = ( "DO $$ -- data-migration-ok: bounded to one row\n" From 766f72f1d7c1dc7b7b1d1d93f8dde46090f76f9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:19:42 -0700 Subject: [PATCH 037/156] fix(anthropic): round-trip thinking blocks to OpenAI backends The experimental /v1/messages adapters lost prior-turn reasoning three different ways once the request left for an OpenAI-shaped backend. On the Responses path, thinking blocks were flattened into output_text inside the assistant message, so the model read its own private reasoning back as visible prose and no reasoning item was ever sent. They now become Responses reasoning input items, grouped by signature so summary parts that arrived as one item go back as one item. The response direction stops hardcoding signature=None and carries the reasoning item id, which is what lets the next turn regroup them; the streaming wrapper emits the matching signature_delta. On the chat completions path the adapter attached thinking_blocks but never set reasoning_content, so Moonshot and DeepSeek substituted a single-space placeholder and other providers sent nothing. It is now derived from the thinking blocks. With use_chat_completions_url_for_anthropic_messages and a model that itself bridges to /v1/responses, the assistant message was dropped whole: reasoning, text, and all. That branch now emits the reasoning items and the message content alongside the tool calls. Fixes #24985 --- .../transformation.py | 32 +++++- .../prompt_templates/common_utils.py | 32 ++++++ .../adapters/transformation.py | 4 + .../responses_adapters/streaming_iterator.py | 13 ++- .../responses_adapters/transformation.py | 107 ++++++++++++------ ...responses_transformation_transformation.py | 87 ++++++++++++++ ...al_pass_through_adapters_transformation.py | 42 +++++++ ...t_responses_adapters_streaming_iterator.py | 14 ++- .../test_responses_adapters_transformation.py | 90 ++++++++++++++- 9 files changed, 379 insertions(+), 42 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 6103b1bf484..ea54bd83e12 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -21,6 +21,9 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + responses_reasoning_item_from_thinking_blocks, +) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, @@ -85,6 +88,21 @@ def _get_reasoning_items( return [] +def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: # mutable-ok: API message payload + """Reasoning input items for an assistant message. + + Stored reasoning items win because they carry an id the Responses API minted; thinking + blocks are the fallback for turns that arrived over another API surface. + """ + stored: Final = [_reasoning_item_to_response_input(r_item) for r_item in _get_reasoning_items(msg)] + if stored: + return stored + from_thinking: Final = responses_reasoning_item_from_thinking_blocks( + cast(Iterable[Mapping[str, Any]], msg.get("thinking_blocks") or ()) # cast-ok: untyped client thinking blocks + ) + return [from_thinking] if from_thinking is not None else [] + + def _build_reasoning_item( item_id: str, encrypted_content: str | None, @@ -372,8 +390,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) ) elif role == "assistant" and tool_calls and isinstance(tool_calls, list): - for r_item in _get_reasoning_items(msg): - input_items.append(_reasoning_item_to_response_input(r_item)) + input_items.extend(_reasoning_input_items(msg)) + if content: + input_items.append( + { + "type": "message", + "role": "assistant", + "content": self._convert_content_to_responses_format(content, "assistant"), + } + ) for tool_call in tool_calls: function = tool_call.get("function") custom = tool_call.get("custom") @@ -400,8 +425,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): raise ValueError(f"tool call not supported: {tool_call}") elif content is not None: if role == "assistant": - for r_item in _get_reasoning_items(msg): - input_items.append(_reasoning_item_to_response_input(r_item)) + input_items.extend(_reasoning_input_items(msg)) input_items.append( { "type": "message", diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2db5776047b..4211fbe610a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1549,6 +1549,38 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content +def reasoning_content_from_thinking_blocks( + thinking_blocks: Iterable[Mapping[str, Any]], +) -> str: + """Flatten Anthropic thinking blocks into the `reasoning_content` string chat models expect. + + Redacted blocks carry no readable text, so they contribute nothing. + """ + return "\n".join( + text + for block in thinking_blocks + if block.get("type") == "thinking" and (text := str(block.get("thinking") or "")) + ) + + +def responses_reasoning_item_from_thinking_blocks( + thinking_blocks: Iterable[Mapping[str, Any]], +) -> dict[str, Any] | None: # mutable-ok: API message payload + """Build a Responses API `reasoning` input item from Anthropic thinking blocks. + + The item carries no `id`: the Responses API rejects an empty one and 404s on any id it + did not mint itself, while an item without an id is always accepted. + """ + summary: Final[list[dict[str, Any]]] = [ # mutable-ok: API message payload + {"type": "summary_text", "text": text} # mutable-ok: API message payload + for block in thinking_blocks + if block.get("type") == "thinking" and (text := str(block.get("thinking") or "")) + ] + if not summary: + return None + return {"type": "reasoning", "summary": summary} # mutable-ok: API message payload + + def _parse_content_for_reasoning( message_text: str | None, ) -> tuple[str | None, str | None]: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 34c2d837127..7c89da81fe6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -64,6 +64,7 @@ from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingCho from litellm.litellm_core_utils.prompt_templates.common_utils import ( parse_tool_call_arguments, + reasoning_content_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -592,6 +593,9 @@ class LiteLLMAnthropicMessagesAdapter: assistant_message["tool_calls"] = tool_calls if len(thinking_blocks) > 0: assistant_message["thinking_blocks"] = thinking_blocks + reasoning_content = reasoning_content_from_thinking_blocks(thinking_blocks) + if reasoning_content: + assistant_message["reasoning_content"] = reasoning_content new_messages.append(assistant_message) return new_messages diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index e2ad9c9c6d3..286edb24b9b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -152,7 +152,7 @@ class AnthropicResponsesStreamWrapper: if block_idx < 0: if not delta: return - block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""}) + block_idx = self._open_block(item_id, {"type": "thinking", "thinking": "", "signature": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -189,6 +189,17 @@ class AnthropicResponsesStreamWrapper: block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: return + done_item_type: Final = ( + getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) if item else None + ) + if done_item_type == "reasoning" and item_id: + self._chunk_queue.append( + { + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "signature_delta", "signature": item_id}, + } + ) self._chunk_queue.append( { "type": "content_block_stop", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 25d729d8606..b1061f86d0b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,12 +6,14 @@ path used for OpenAI and Azure models. """ import json -from collections.abc import Iterable +from collections.abc import Iterable, Mapping +from itertools import groupby from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + responses_reasoning_item_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -100,6 +102,58 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(block, dict) and block.get("type") == "text" and (text := block.get("text")) # pyright: ignore[reportUnnecessaryIsInstance] # untrusted client payload ] + @staticmethod + def _summary_part_text(part: object) -> str: + if isinstance(part, Mapping): + mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json + return str(mapping.get("text", "")) + return str(getattr(part, "text", "")) + + @classmethod + def _thinking_blocks_from_reasoning_item( + cls, + item_id: str | None, + summary: Iterable[object], + ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload + """Anthropic thinking blocks for one Responses reasoning item. + + The reasoning item id rides along as the signature so that a follow-up turn can + regroup the summary parts into the single reasoning item they came from. + """ + return tuple( + AnthropicResponseContentBlockThinking( + type="thinking", + thinking=text, + signature=item_id or None, + ).model_dump() + for part in summary + if (text := cls._summary_part_text(part)) + ) + + @staticmethod + def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, Any]]) -> str: + """Group consecutive thinking blocks sharing a signature; keep every other block alone.""" + index, block = indexed_block + signature: Final = block.get("signature") or "" + return f"thinking:{signature}" if block.get("type") == "thinking" else f"block:{index}" + + @classmethod + def _assistant_group_to_input_item( + cls, group: tuple[Mapping[str, Any], ...] + ) -> dict[str, Any] | None: # mutable-ok: API message payload + first: Final = group[0] + btype: Final = first.get("type") + if btype == "thinking": + return responses_reasoning_item_from_thinking_blocks(group) + if btype == "tool_use": + return { # mutable-ok: API message payload + "type": "function_call", + "call_id": first.get("id", ""), + "name": first.get("name", ""), + "arguments": json.dumps(first.get("input", {})), + } + return None + def translate_messages_to_responses_input( self, messages: list[AllAnthropicPassThroughMessageValues], @@ -113,6 +167,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: user image -> message(role=user, input_image) user tool_result -> function_call_output assistant text -> message(role=assistant, output_text) + assistant thinking -> reasoning assistant tool_use -> function_call """ input_items: Final[list[dict[str, Any]]] = [] @@ -233,27 +288,17 @@ class LiteLLMAnthropicToResponsesAPIAdapter: } ) elif isinstance(content, list): - asst_parts: list[dict[str, Any]] = [] - for block in content: - if not isinstance(block, dict): - continue - btype = block.get("type") - if btype == "text": - asst_parts.append({"type": "output_text", "text": block.get("text", "")}) - elif btype == "tool_use": - # tool_use becomes a top-level function_call item - input_items.append( - { - "type": "function_call", - "call_id": block.get("id", ""), - "name": block.get("name", ""), - "arguments": json.dumps(block.get("input", {})), - } - ) - elif btype == "thinking": - thinking_text = block.get("thinking", "") - if thinking_text: - asst_parts.append({"type": "output_text", "text": thinking_text}) + blocks = tuple(block for block in content if isinstance(block, dict)) + input_items.extend( + item + for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key) + if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None + ) + asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload + {"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload + for block in blocks + if block.get("type") == "text" + ] if asst_parts: input_items.append( { @@ -514,16 +559,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for item in response.output: if isinstance(item, ResponseReasoningItem): - for summary in item.summary: - text = getattr(summary, "text", "") - if text: - content.append( - AnthropicResponseContentBlockThinking( - type="thinking", - thinking=text, - signature=None, - ).model_dump() - ) + content.extend(self._thinking_blocks_from_reasoning_item(item.id, item.summary)) elif isinstance(item, ResponseOutputMessage): for part in item.content: @@ -555,6 +591,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content.append( AnthropicResponseContentBlockText(type="text", text=part.get("text", "")).model_dump() ) + elif item_type == "reasoning": + content.extend( + self._thinking_blocks_from_reasoning_item( + item.get("id"), + cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json + ) + ) elif item_type == "function_call": try: input_data = json.loads(item.get("arguments", "{}")) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 4ff92aaf87d..124dc67b4fd 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3762,3 +3762,90 @@ def test_response_incomplete_stream_event_without_details_defaults_to_length(): result = iterator.chunk_parser(chunk) assert result.choices[0].finish_reason == "length" + + +def test_assistant_message_with_tool_calls_keeps_its_content(): + """Regression for https://github.com/BerriAI/litellm/issues/24985. + + An assistant turn that both answered and called a tool used to lose its whole message: + the branch handling tool_calls emitted the calls and dropped the text. + """ + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": "Let me look that up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Denver"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "88F"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + assistant_message = next( + item for item in input_items if item.get("type") == "message" and item.get("role") == "assistant" + ) + assert assistant_message["content"] == [{"type": "output_text", "text": "Let me look that up."}] + assert [item.get("type") for item in input_items] == [ + "message", + "message", + "function_call", + "function_call_output", + ] + + +def test_assistant_thinking_blocks_become_a_reasoning_input_item(): + """Thinking blocks are how an Anthropic-shaped turn carries reasoning into this bridge.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": "Denver is sunny.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "sig1"}, + {"type": "redacted_thinking", "data": "REDACTED"}, + ], + }, + {"role": "user", "content": "Why?"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_item = next(item for item in input_items if item.get("type") == "reasoning") + assert reasoning_item["summary"] == [{"type": "summary_text", "text": "August in Denver is dry."}] + assert "id" not in reasoning_item + + +def test_stored_reasoning_items_win_over_thinking_blocks(): + """A minted reasoning id beats a re-derived one, so the two must not both be sent.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + { + "role": "assistant", + "content": "Denver is sunny.", + "reasoning_items": [ + { + "type": "reasoning", + "id": "rs_real", + "summary": [{"type": "summary_text", "text": "August in Denver is dry."}], + } + ], + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "rs_real"} + ], + }, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_items = [item for item in input_items if item.get("type") == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["id"] == "rs_real" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e4dacc308dc..d0169963962 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -359,6 +359,48 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert result[1]["tool_calls"][0]["id"] == "toolu_01234" +def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): + """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. + + Without it Moonshot and DeepSeek fill in a single-space placeholder and the model gets + a blank where its own prior reasoning belongs. + """ + + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "text", "text": "Which city is best for a picnic?"}], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "Denver is dry in August.", "signature": "sig1"}, + {"type": "thinking", "thinking": "San Francisco is foggy.", "signature": "sig2"}, + {"type": "redacted_thinking", "data": "REDACTED"}, + {"type": "text", "text": "Denver."}, + ], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert result[1]["reasoning_content"] == "Denver is dry in August.\nSan Francisco is foggy." + assert result[1]["content"] == "Denver." + + +def test_translate_anthropic_messages_to_openai_sets_no_reasoning_content_without_thinking(): + anthropic_messages = [ + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[{"type": "text", "text": "Denver."}], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert "reasoning_content" not in result[0] + + def test_translate_anthropic_messages_to_openai_tool_message_placement(): """Test that tool result messages are placed before user messages in the conversation order.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 8b591fcd7da..12a6ab4324f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -139,14 +139,26 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 0), ("content_block_delta", 0), ("content_block_delta", 0), + ("content_block_delta", 0), ("content_block_stop", 0), ("content_block_start", 1), ("content_block_delta", 1), ("content_block_stop", 1), ] - assert chunks[1]["content_block"] == {"type": "thinking", "thinking": ""} + assert chunks[1]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" + def test_reasoning_item_id_is_streamed_as_the_thinking_signature(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weighing options"])) + + signature_deltas = [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] + assert [(c["index"], c["delta"]["signature"]) for c in signature_deltas] == [(0, "rs_1")] + + def test_no_signature_delta_without_a_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=[])) + + assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] + class TestToolUseBlockClosedExactlyOnce: """Regression for https://github.com/BerriAI/litellm/issues/37273. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 964f4b9f68b..6479c43ee7a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -486,8 +486,8 @@ class TestTranslateMessagesToResponsesInput: } ] - def test_assistant_thinking_block_becomes_output_text(self): - """Assistant thinking block text is included as output_text.""" + def test_assistant_thinking_block_becomes_reasoning_item(self): + """Assistant thinking block becomes a reasoning item, never visible assistant prose.""" messages = [ { "role": "assistant", @@ -495,7 +495,65 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [{"type": "output_text", "text": "Let me reason step by step."}] + assert result == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Let me reason step by step."}], + } + ] + + def test_reasoning_item_carries_no_id(self): + """A fabricated reasoning id 404s upstream, so the item must go out without one.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "Private reasoning.", "signature": "rs_abc123"}], + } + ] + result = _translate_messages(messages) + assert "id" not in result[0] + + def test_thinking_blocks_sharing_a_signature_become_one_reasoning_item(self): + """Summary parts of one upstream reasoning item are regrouped by their signature.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "First part.", "signature": "rs_abc123"}, + {"type": "thinking", "thinking": "Second part.", "signature": "rs_abc123"}, + {"type": "thinking", "thinking": "A later item.", "signature": "rs_def456"}, + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "reasoning", + "summary": [ + {"type": "summary_text", "text": "First part."}, + {"type": "summary_text", "text": "Second part."}, + ], + }, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "A later item."}], + }, + ] + + def test_thinking_and_text_stay_separate(self): + """The visible answer stays the only thing in the assistant message.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "The user wants Denver."}, + {"type": "text", "text": "Denver is the best pick."}, + ], + } + ] + result = _translate_messages(messages) + assert [item["type"] for item in result] == ["reasoning", "message"] + assert result[1]["content"] == [{"type": "output_text", "text": "Denver is the best pick."}] def test_assistant_empty_thinking_block_skipped(self): """Assistant thinking block with empty thinking text is skipped.""" @@ -1094,7 +1152,7 @@ def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMo return item -def _make_reasoning_item(summaries: List[str]) -> MagicMock: +def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> MagicMock: """Build a mock ResponseReasoningItem.""" from openai.types.responses import ResponseReasoningItem # type: ignore[import] @@ -1105,6 +1163,7 @@ def _make_reasoning_item(summaries: List[str]) -> MagicMock: summary_mocks.append(s) item = MagicMock(spec=ResponseReasoningItem) + item.id = item_id item.summary = summary_mocks return item @@ -1178,6 +1237,29 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [] + def test_reasoning_item_id_becomes_thinking_signature(self): + """The reasoning item id rides back as the signature so the next turn can regroup it.""" + reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert [block["signature"] for block in result["content"]] == ["rs_abc123", "rs_abc123"] + + def test_dict_reasoning_item_becomes_thinking_block(self): + """A reasoning item arriving as a plain dict is kept, not dropped.""" + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_dict_1", + "summary": [{"type": "summary_text", "text": "Weighing the options."}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + {"type": "thinking", "thinking": "Weighing the options.", "signature": "rs_dict_1"} + ] + def test_usage_mapped_correctly(self): """Input/output tokens from ResponseAPIUsage are mapped to AnthropicUsage.""" response = _make_mock_response( From e61baa6d0160991f7f7cd4fd1e361536cf4f80a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:25:31 -0700 Subject: [PATCH 038/156] fix: read one quoted run as one literal, and stop before bind values --- .../check_migrations_no_data_rewrites.py | 32 ++++++++++++++--- .../test_check_migrations_no_data_rewrites.py | 34 +++++++++++++++++++ 2 files changed, 61 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 3d57691cea4..b4907425838 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -147,6 +147,8 @@ OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) +BIND_VALUES = re.compile(r"\bUSING\b", re.IGNORECASE) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -286,10 +288,20 @@ def skip_block_comment(sql: str, start: int) -> int: def skip_quoted(sql: str, start: int, quote: str) -> int: - """One quoted run, up to and including its closing quote. A doubled quote needs no - special case: closing on the first and reopening on the second masks the same span.""" - stop = sql.find(quote, start + 1) - return len(sql) if stop == -1 else stop + 1 + """One quoted run, up to and including its closing quote. A doubled quote is an escaped + quote sitting inside the run rather than the end of it. Closing on the first and reopening + on the second would mask the same span, which is why this looked like it needed no special + case, but the run is also handed on whole as one literal, and splitting it there offers the + tail of a string to be read as SQL in its own right.""" + index = start + 1 + while True: + stop = sql.find(quote, index) + if stop == -1: + return len(sql) + if sql[stop + 1 : stop + 2] == quote: + index = stop + 2 + continue + return stop + 1 def strip_parens(statement: str) -> str: @@ -523,8 +535,9 @@ def scan_region( exempt = markers.exempt(start, end) if hands_off_sql(match.group(), executed) and not exempt: + commands_end = match.start() + bind_values_start(match.group()) for start, end in literals: - if match.start() <= start and end <= match.end(): + if match.start() <= start and end <= commands_end: yield from scan_region(document, region[start:end], migration, markers, offset + start) keyword = offending_keyword(match.group()) @@ -536,6 +549,15 @@ def scan_region( yield from scan_region(document, region[start:end], migration, markers, offset + start) +def bind_values_start(statement: str) -> int: + """Where a statement stops handing commands to the server and starts listing bind values. + The expressions after `USING` are values substituted into the command, never commands in + their own right, so one that merely spells out a rewrite is not running it. Read off the + masked text, so a `USING` written inside the command string is not mistaken for this one.""" + keyword = BIND_VALUES.search(statement) + return len(statement) if keyword is None else keyword.start() + + def statement_start(statement: re.Match[str]) -> int: """Where the statement's own text begins, past the whitespace and blanked comments it picked up from whatever sat between it and the statement before it, one of which can be a marker.""" diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 5413f8cdc62..21fafcd205f 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -467,6 +467,40 @@ class TestDynamicSql: sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = date_trunc(''day'', \"t\")';\nEND $$;" assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_rewrite_quoted_as_data_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''UPDATE \"Foo\" SET \"a\" = 1''';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_doubled_quote_does_not_split_the_literal_it_sits_in(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('a''UPDATE \"Bar\" SET \"a\" = 1''b');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_following_a_doubled_quote_in_the_same_payload_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''x''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_in_a_later_command_before_bind_values_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1; DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_bind_value_naming_a_rewrite_is_not_run(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'INSERT INTO \"Audit\" (\"note\") VALUES ($1)'" + " USING 'DELETE FROM \"Foo\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_executed_with_bind_values_is_still_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_using_written_inside_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" USING \"Bar\"" + " WHERE \"Foo\".\"a\" = \"Bar\".\"a\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + def test_execute_of_ddl_passes(self, tmp_path): sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" assert _keywords(tmp_path, sql) == () From 94caab7302ba471647bfba51bb1e5f8ad4cc5222 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:46:18 -0700 Subject: [PATCH 039/156] fix(interactions): release the reservation for creates nothing will poll, and let OTEL see the settled cost Two review findings, both in the handoff between the create's success callback and the background poll task. The callback deferred its budget reservation release for any interactions response with no usage, but the scheduler only starts a poll task when the status is in_progress and an id is present. A create that came back terminal without usage therefore matched the callback's test, got no poll task, and left its reservation open forever: the pre-call estimate stayed added to the key, user, team and org spend counters, and the key began refusing traffic against budget it had never spent. The two conditions now come from one shared gate so they cannot drift apart again. Settling a background interaction re-runs the success handlers for a second result on the same request, and OTEL dedupes span emission on a marker held in that request's metadata. The in-progress create claimed the marker, so the completion, the only event carrying usage and cost, was dropped as a duplicate by OTEL and by every integration deriving from it. Clearing the success-scoped markers alongside the existing dedup flag lets the cost span through, leaving failure and guardrail markers untouched. --- .../interactions/background_cost_polling.py | 13 ++- litellm/litellm_core_utils/litellm_logging.py | 29 +++++- .../proxy/hooks/proxy_track_cost_callback.py | 5 +- .../test_litellm_logging.py | 38 +++++++ .../hooks/test_proxy_track_cost_callback.py | 99 +++++++++++++++++++ 5 files changed, 180 insertions(+), 4 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index ccf4c2dd853..262d8e3260d 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -162,6 +162,17 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction") +def is_pollable_background_interaction(response: InteractionsAPIResponse) -> bool: + """ + The single gate deciding whether a create's response gets a poll task. + The proxy's success callback defers releasing the budget reservation for + exactly these responses, on the promise that a poll task will settle them, + so a response one site accepts and the other refuses strands its + reservation on the spend counters with nothing left to reconcile it. + """ + return response.status == "in_progress" and bool(response.id) + + @dataclass(frozen=True, slots=True) class _ActiveBackgroundPoll: task: "asyncio.Task[None]" @@ -188,7 +199,7 @@ def maybe_schedule_background_interaction_cost_polling( return None if not isinstance(response, InteractionsAPIResponse): return None - if response.status != "in_progress" or not response.id: + if not is_pollable_background_interaction(response): return None logging_obj = create_kwargs.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e7d7213b47c..275803c5aed 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2200,12 +2200,37 @@ class Logging(LiteLLMLoggingBaseClass): 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 + the dedup flags 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) + self._reset_success_emission_dedupe() await self.async_success_handler(result=result) + def _reset_success_emission_dedupe(self) -> None: + """ + Success callbacks dedupe per request, because the sync and async + handlers both fire on some paths and would otherwise report one call + twice. A settled background interaction is a genuinely second success + event on the same request, so every such marker has to be cleared or + the completion, the only event that carries usage and cost, is + discarded as a duplicate of the in-progress create. + """ + self.model_call_details.pop("has_logged_async_success", None) + litellm_params = self.model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + metadata = litellm_params.get("metadata") + if not isinstance(metadata, dict): + return + otel_internal = metadata.get("_otel_internal") + if not isinstance(otel_internal, dict): + return + spans_logged = otel_internal.get("spans_logged") + if not isinstance(spans_logged, dict): + return + for scope in [key for key in spans_logged if isinstance(key, tuple) and key[-1:] == ("success",)]: + del spans_logged[scope] + 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 55d3c6c6ed5..46cf62ede1c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -478,9 +478,12 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: def _is_unbilled_in_progress_interaction(completion_response: object) -> bool: + from litellm.interactions.background_cost_polling import is_pollable_background_interaction from litellm.types.interactions import InteractionsAPIResponse - return isinstance(completion_response, InteractionsAPIResponse) and completion_response.usage is None + if not isinstance(completion_response, InteractionsAPIResponse): + return False + return completion_response.usage is None and is_pollable_background_interaction(completion_response) def _should_track_cost_callback( 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 8a682517f64..7c90d31261b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4644,6 +4644,44 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 +@pytest.mark.asyncio +async def test_background_interaction_completion_lets_otel_emit_the_cost_span(): + """ + OTEL, and every integration that derives from it, dedupes span emission on + a marker kept in the request's own metadata. The in-progress create claims + that marker, so without clearing it the settled completion, the only event + carrying usage and cost, is discarded as a duplicate and every + OTEL-family backend shows the interaction as a span with no cost at all. + """ + import datetime as dt + + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + from litellm.types.interactions import InteractionsAPIResponse + + otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) + 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 otel._emit_once(logging_obj.model_call_details, "success") is True + assert otel._emit_once(logging_obj.model_call_details, "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 otel._emit_once(logging_obj.model_call_details, "success") is True + + @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 642db7e2d5d..93e0cbec596 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 @@ -859,6 +859,105 @@ async def test_track_cost_callback_releases_reservation_for_in_progress_interact mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status", + ["completed", "failed", "cancelled", "incomplete", "requires_action", "budget_exceeded"], +) +async def test_track_cost_callback_releases_reservation_for_unpollable_interaction(status): + """ + Only an in-progress create gets a poll task, so a create that comes back + terminal with no usage has nobody left to reconcile its reservation. The + callback must release it there and then, or the pre-call estimate stays + added to the key, user, team and org spend counters and starts refusing + traffic against budget that was never actually spent. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + terminal_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + 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=terminal_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_reservation_for_interaction_without_an_id(): + """ + The scheduler also refuses a response with no id, since it has nothing to + poll for, so the callback must not defer to a poll task that will never + exist. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + idless_response = InteractionsAPIResponse( + id="", + 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=idless_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + + +@pytest.mark.parametrize( + "status", + ["in_progress", "completed", "failed", "cancelled", "incomplete", "requires_action"], +) +@pytest.mark.parametrize("interaction_id", ["interactions/bg-abc", ""]) +def test_callback_defers_exactly_the_interactions_the_scheduler_polls(status, interaction_id): + """ + Pins the invariant the two modules share: the callback may only hold a + budget reservation open for a response the scheduler will actually poll. + Any drift between the two gates leaks reservations onto live spend + counters, so assert they agree rather than restating either condition. + """ + from litellm.interactions.background_cost_polling import is_pollable_background_interaction + from litellm.proxy.hooks.proxy_track_cost_callback import _is_unbilled_in_progress_interaction + from litellm.types.interactions import InteractionsAPIResponse + + response = InteractionsAPIResponse( + id=interaction_id, + model="gemini-3-flash-preview", + status=status, + ) + + assert _is_unbilled_in_progress_interaction(response) is is_pollable_background_interaction(response) + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): """ From 923da852fa1b245d5edc30bd9f563c52d7220cb7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 12:03:49 -0700 Subject: [PATCH 040/156] fix: read a loop body as its own statement, not as part of the header A `FOR ... LOOP` header carries no semicolon of its own, so the first statement of the loop body is written into the same semicolon-delimited run. Reading the pair as one statement let the header's row source stand in as the keyword for both, which hid whatever the loop repeats: a plain `UPDATE` in a query-driven loop went unreported, and so did an `EXECUTE` of one. That is the shape a row-by-row backfill takes, and it is the shape this gate exists to stop. --- .../check_migrations_no_data_rewrites.py | 43 ++++-- .../test_check_migrations_no_data_rewrites.py | 134 ++++++++++++++++++ 2 files changed, 162 insertions(+), 15 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index b4907425838..c82d52167b4 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -106,6 +106,7 @@ INTO_TARGETS = re.compile( re.IGNORECASE, ) LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE) +LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTALL) WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?!:=])=(?![=>])") PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) @@ -530,25 +531,37 @@ def scan_region( executed = executed_names(masked) for match in STATEMENT.finditer(masked): - start = offset + statement_start(match) - end = offset + match.end() - exempt = markers.exempt(start, end) + exempt = markers.exempt(offset + statement_start(match), offset + match.end()) - if hands_off_sql(match.group(), executed) and not exempt: - commands_end = match.start() + bind_values_start(match.group()) - for start, end in literals: - if match.start() <= start and end <= commands_end: - yield from scan_region(document, region[start:end], migration, markers, offset + start) + for clause, base in clauses(match.group(), match.start()): + if hands_off_sql(clause, executed) and not exempt: + commands_end = base + bind_values_start(clause) + for start, end in literals: + if base <= start and end <= commands_end: + yield from scan_region(document, region[start:end], migration, markers, offset + start) - keyword = offending_keyword(match.group()) - if keyword is None or exempt: - continue - yield Violation(migration, line_of(document, offset + keyword_start(match)), keyword) + keyword = offending_keyword(clause) + if keyword is None or exempt: + continue + yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) for start, end in bodies: yield from scan_region(document, region[start:end], migration, markers, offset + start) +def clauses(statement: str, start: int) -> Iterator[tuple[str, int]]: + """The statements written inside one semicolon-delimited run, each with where it begins. A + `FOR ... LOOP` header takes no semicolon of its own, so the first statement of the loop body + is written into the same run, and reading the pair as one statement lets the header's row + source stand in as the keyword for both. That hides the statement the loop repeats, which is + the shape a row-by-row backfill takes. Splitting after each header, nested ones included, + reads the header and the body as the separate statements Postgres runs them as.""" + edges = (0, *(header.end() for header in LOOP_HEADER.finditer(statement)), len(statement)) + for opens, closes in zip(edges, edges[1:]): + if opens < closes: + yield statement[opens:closes], start + opens + + def bind_values_start(statement: str) -> int: """Where a statement stops handing commands to the server and starts listing bind values. The expressions after `USING` are values substituted into the command, never commands in @@ -565,9 +578,9 @@ def statement_start(statement: re.Match[str]) -> int: return statement.start() + len(text) - len(text.lstrip()) -def keyword_start(statement: re.Match[str]) -> int: - word = leading_keyword(statement.group()) - return statement.start() + (0 if word is None else word.start()) +def keyword_start(clause: str, base: int) -> int: + word = leading_keyword(clause) + return base + (0 if word is None else word.start()) def line_of(sql: str, offset: int) -> int: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 21fafcd205f..970bc9f5264 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -275,6 +275,140 @@ class TestDollarQuotedBlocks: assert _scan(tmp_path, sql)[0].line == 7 +class TestLoopBodies: + def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_delete_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' DELETE FROM "Foo" WHERE "id" = r."id";\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_the_loop_query_does_not_hide_the_body(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT a."id" FROM "A" a JOIN "B" b USING ("id") LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_executed_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_nested_under_a_guard_inside_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' IF r."id" > 0 THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_inside_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP FOR b IN SELECT "id" FROM "B" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_supplying_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP\n' + ' FOR b IN UPDATE "Foo" SET "x" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_loop_running_only_ddl_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' CREATE INDEX "i" ON "Foo"("a");\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_loop_over_a_rewrite_returning_rows_is_flagged_once(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN UPDATE "Foo" SET "a" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_on_a_loop_exempts_the_rewrite_it_repeats(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + " -- data-migration-ok: one row\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_select_for_update_lock_is_not_read_as_a_loop(self, tmp_path): + sql = 'DO $$\nBEGIN\n PERFORM 1 FROM "Foo" FOR UPDATE;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + class TestQuotingAndComments: def test_update_inside_string_literal_passes(self, tmp_path): sql = 'ALTER TABLE "Foo" ADD COLUMN "note" TEXT NOT NULL DEFAULT \'UPDATE nothing\';' From d4608af27bf9362b3b2b310e6deef61e312af71f Mon Sep 17 00:00:00 2001 From: tin Date: Sat, 22 Aug 2026 19:30:54 +0000 Subject: [PATCH 041/156] fix(proxy): skip health checks for strategy routers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/health_check.py | 22 +++------- .../proxy/test_health_check_max_tokens.py | 44 +++++++++++-------- 2 files changed, 31 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index f9d408fb7de..4d70406caf8 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -18,6 +18,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS, ) +from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model ILLEGAL_DISPLAY_PARAMS: Final = [ "messages", @@ -182,30 +183,17 @@ async def run_with_timeout(task, timeout): return {"error": "Timeout exceeded", "exception": timeout_exception} -def _is_semantic_auto_router_deployment(litellm_params: dict) -> bool: - """ - True for semantic auto_router deployments (auto_router/) that are not - sub-strategies (complexity_router, adaptive_router, quality_router). - - These are meta-routers that select among real LLM deployments at request time; - they have no LLM endpoint to health-check. - """ +def _is_strategy_router_deployment(litellm_params: dict) -> bool: + """True for strategy-router deployments.""" model: Final[object] = litellm_params.get("model", "") - if not isinstance(model, str): - return False - if not model.startswith("auto_router/"): - return False - for sub_strategy in ("complexity_router", "adaptive_router", "quality_router"): - if model.startswith(f"auto_router/{sub_strategy}"): - return False - return True + return isinstance(model, str) and classify_strategy_router_model(model) is not None async def _run_model_health_check(model: dict): litellm_params = model["litellm_params"] model_info: Final = model.get("model_info", {}) - if _is_semantic_auto_router_deployment(litellm_params): + if _is_strategy_router_deployment(litellm_params): return {} mode: Final = _resolve_health_check_mode( diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 5a606d5f74e..80c6f2bba25 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -5,7 +5,7 @@ import pytest from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers from litellm.proxy import health_check as hc_module from litellm.proxy.health_check import ( - _is_semantic_auto_router_deployment, + _is_strategy_router_deployment, _resolve_health_check_max_tokens, _resolve_health_check_mode, _update_litellm_params_for_health_check, @@ -495,33 +495,22 @@ def test_autodetected_embedding_skips_reasoning_effort(): assert "max_tokens" not in updated -# --------------------------------------------------------------------------- -# auto_router (semantic router) deployments must be skipped by health checks. -# -# These are meta-routers that select among real LLM deployments at request -# time. They have no LLM endpoint to probe. Before this fix, the health check -# passed model="auto_router/router_1" to get_llm_provider(), which raised -# BadRequestError: "Unmapped LLM provider for this endpoint" because -# auto_router is not a real LLM provider. -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize( "model, expected", [ ("auto_router/router_1", True), ("auto_router/my_router", True), - ("auto_router/complexity_router", False), - ("auto_router/adaptive_router", False), - ("auto_router/quality_router", False), - ("auto_router/adaptive_router/subpath", False), + ("auto_router/complexity_router", True), + ("auto_router/adaptive_router", True), + ("auto_router/quality_router", True), + ("auto_router/adaptive_router/subpath", True), ("gpt-4", False), ("openai/gpt-4", False), ("bedrock/claude", False), ], ) -def test_is_semantic_auto_router_deployment(model, expected): - assert _is_semantic_auto_router_deployment({"model": model}) == expected +def test_is_strategy_router_deployment(model, expected): + assert _is_strategy_router_deployment({"model": model}) == expected @pytest.mark.asyncio @@ -543,3 +532,22 @@ async def test_run_model_health_check_skips_auto_router_deployment(): fake_ahealth_check.assert_not_called() assert result == {} + + +@pytest.mark.asyncio +async def test_run_model_health_check_skips_complexity_router_deployment(): + fake_ahealth_check = AsyncMock(return_value={}) + model = { + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"simple": "gpt-4o-mini"}}, + "complexity_router_default_model": "gpt-4o-mini", + }, + "model_info": {}, + } + + with patch.object(hc_module.litellm, "ahealth_check", fake_ahealth_check): + result = await hc_module._run_model_health_check(model) + + fake_ahealth_check.assert_not_called() + assert result == {} From 32bf1aba2958a8601d78ba7f86eb97d58a1adf85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:27:20 -0700 Subject: [PATCH 042/156] fix(anthropic): stop signing replayed thinking blocks and strip reasoning_content A reasoning item id is not an Anthropic signature. Passing it off as one got the block replayed to Anthropic and Bedrock as if it were real, and every backend that verifies signatures rejected the turn. Thinking blocks now come back unsigned, and the streaming path no longer emits a signature_delta for them. Azure AI Foundry, Fireworks, and vLLM reject unknown message fields, so they now strip reasoning_content alongside thinking_blocks the way Mistral already did. The thinking-block helpers take ChatCompletionThinkingBlock and ChatCompletionRedactedThinkingBlock instead of loose mappings. --- .../transformation.py | 8 +-- .../prompt_templates/common_utils.py | 34 ++++++++----- .../responses_adapters/streaming_iterator.py | 11 ---- .../responses_adapters/transformation.py | 25 ++++++---- litellm/llms/azure_ai/chat/transformation.py | 10 +++- .../llms/fireworks_ai/chat/transformation.py | 1 + .../llms/hosted_vllm/chat/transformation.py | 5 +- ...t_responses_adapters_streaming_iterator.py | 10 +--- .../test_responses_adapters_transformation.py | 50 +++++++++++++------ .../chat/test_azure_ai_transformation.py | 2 + .../test_fireworks_ai_chat_transformation.py | 2 + .../test_hosted_vllm_chat_transformation.py | 2 + 12 files changed, 96 insertions(+), 64 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index ea54bd83e12..3e55c3c637e 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -97,10 +97,10 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: stored: Final = [_reasoning_item_to_response_input(r_item) for r_item in _get_reasoning_items(msg)] if stored: return stored - from_thinking: Final = responses_reasoning_item_from_thinking_blocks( - cast(Iterable[Mapping[str, Any]], msg.get("thinking_blocks") or ()) # cast-ok: untyped client thinking blocks - ) - return [from_thinking] if from_thinking is not None else [] + raw_blocks: Final = msg.get("thinking_blocks") or () + blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json + from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) + return [] if from_thinking is None else [dict(from_thinking)] def _build_reasoning_item( diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 4211fbe610a..e0ba2e13c6e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -28,8 +28,12 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionFileObject, ChatCompletionImageObject, + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ChatCompletionRedactedThinkingBlock, ChatCompletionResponseMessage, ChatCompletionTextObject, + ChatCompletionThinkingBlock, ChatCompletionToolParam, ChatCompletionUserMessage, ) @@ -1549,36 +1553,42 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content +def _readable_thinking_text( + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, +) -> str: + """The text a chat model can read back, empty for redacted blocks and malformed ones.""" + if block.get("type") != "thinking": + return "" + thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag + return str(thinking or "") + + def reasoning_content_from_thinking_blocks( - thinking_blocks: Iterable[Mapping[str, Any]], + thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], ) -> str: """Flatten Anthropic thinking blocks into the `reasoning_content` string chat models expect. Redacted blocks carry no readable text, so they contribute nothing. """ - return "\n".join( - text - for block in thinking_blocks - if block.get("type") == "thinking" and (text := str(block.get("thinking") or "")) - ) + return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block))) def responses_reasoning_item_from_thinking_blocks( - thinking_blocks: Iterable[Mapping[str, Any]], -) -> dict[str, Any] | None: # mutable-ok: API message payload + thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], +) -> ChatCompletionReasoningItem | None: """Build a Responses API `reasoning` input item from Anthropic thinking blocks. The item carries no `id`: the Responses API rejects an empty one and 404s on any id it did not mint itself, while an item without an id is always accepted. """ - summary: Final[list[dict[str, Any]]] = [ # mutable-ok: API message payload - {"type": "summary_text", "text": text} # mutable-ok: API message payload + summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload + ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text) for block in thinking_blocks - if block.get("type") == "thinking" and (text := str(block.get("thinking") or "")) + if (text := _readable_thinking_text(block)) ] if not summary: return None - return {"type": "reasoning", "summary": summary} # mutable-ok: API message payload + return ChatCompletionReasoningItem(type="reasoning", summary=summary) def _parse_content_for_reasoning( diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 286edb24b9b..5577d4a9c2d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -189,17 +189,6 @@ class AnthropicResponsesStreamWrapper: block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: return - done_item_type: Final = ( - getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) if item else None - ) - if done_item_type == "reasoning" and item_id: - self._chunk_queue.append( - { - "type": "content_block_delta", - "index": block_idx, - "delta": {"type": "signature_delta", "signature": item_id}, - } - ) self._chunk_queue.append( { "type": "content_block_stop", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index b1061f86d0b..1a6b0498a52 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -38,7 +38,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, AnthropicUsage, ) -from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import ( + ChatCompletionThinkingBlock, + ResponseAPIUsage, + ResponsesAPIResponse, +) class LiteLLMAnthropicToResponsesAPIAdapter: @@ -112,19 +116,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @classmethod def _thinking_blocks_from_reasoning_item( cls, - item_id: str | None, summary: Iterable[object], ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload """Anthropic thinking blocks for one Responses reasoning item. - The reasoning item id rides along as the signature so that a follow-up turn can - regroup the summary parts into the single reasoning item they came from. + The signature stays empty: only Anthropic can sign a thinking block, and a stand-in + value would be replayed as a real one and rejected by every backend that verifies it. """ return tuple( AnthropicResponseContentBlockThinking( type="thinking", thinking=text, - signature=item_id or None, + signature=None, ).model_dump() for part in summary if (text := cls._summary_part_text(part)) @@ -132,10 +135,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, Any]]) -> str: - """Group consecutive thinking blocks sharing a signature; keep every other block alone.""" + """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block - signature: Final = block.get("signature") or "" - return f"thinking:{signature}" if block.get("type") == "thinking" else f"block:{index}" + return "thinking" if block.get("type") == "thinking" else f"block:{index}" @classmethod def _assistant_group_to_input_item( @@ -144,7 +146,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: first: Final = group[0] btype: Final = first.get("type") if btype == "thinking": - return responses_reasoning_item_from_thinking_blocks(group) + blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload + reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) + return None if reasoning_item is None else dict(reasoning_item) if btype == "tool_use": return { # mutable-ok: API message payload "type": "function_call", @@ -559,7 +563,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for item in response.output: if isinstance(item, ResponseReasoningItem): - content.extend(self._thinking_blocks_from_reasoning_item(item.id, item.summary)) + content.extend(self._thinking_blocks_from_reasoning_item(item.summary)) elif isinstance(item, ResponseOutputMessage): for part in item.content: @@ -594,7 +598,6 @@ class LiteLLMAnthropicToResponsesAPIAdapter: elif item_type == "reasoning": content.extend( self._thinking_blocks_from_reasoning_item( - item.get("id"), cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json ) ) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index bc8ea31ea8c..9e7161120cc 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -30,7 +30,12 @@ class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" -NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ("thinking_blocks", "provider_specific_fields", "cache_control") +NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( + "thinking_blocks", + "reasoning_content", + "provider_specific_fields", + "cache_control", +) class AzureAIStudioConfig(OpenAIConfig): @@ -173,7 +178,8 @@ class AzureAIStudioConfig(OpenAIConfig): """ - Azure AI Studio doesn't support content as a list. This handles: 1. Strips message fields that are not part of the OpenAI chat-completions - schema (thinking_blocks, provider_specific_fields, cache_control). + schema (thinking_blocks, reasoning_content, provider_specific_fields, + cache_control). Azure AI Foundry backends set additionalProperties=false and reject these with "Extra inputs are not permitted", which breaks multi-turn Anthropic-format clients that echo thinking blocks back as history. diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index e64237da978..4e9731ef485 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -504,6 +504,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): m = cast(dict, message) m.pop("provider_specific_fields", None) m.pop("thinking_blocks", None) + m.pop("reasoning_content", None) return messages diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 46a2320b655..29dc485732f 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -164,12 +164,13 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): """ Support translating: - video files from file_id or file_data to video_url - - thinking_blocks on assistant messages are removed, and content lists - are converted to strings for vLLM compatibility + - thinking_blocks and reasoning_content on assistant messages are removed, + and content lists are converted to strings for vLLM compatibility """ for message in messages: if message["role"] == "assistant": message.pop("thinking_blocks", None) + message.pop("reasoning_content", None) existing_content = message.get("content") if isinstance(existing_content, list): text_parts = [] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 12a6ab4324f..aebbed88c70 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -139,7 +139,6 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 0), ("content_block_delta", 0), ("content_block_delta", 0), - ("content_block_delta", 0), ("content_block_stop", 0), ("content_block_start", 1), ("content_block_delta", 1), @@ -148,15 +147,10 @@ class TestReasoningItemWithoutSummaryText: assert chunks[1]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" - def test_reasoning_item_id_is_streamed_as_the_thinking_signature(self): + def test_the_reasoning_item_id_is_never_streamed_as_a_signature(self): + """A stand-in signature would be replayed as a real one, so none is ever sent.""" chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weighing options"])) - signature_deltas = [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] - assert [(c["index"], c["delta"]["signature"]) for c in signature_deltas] == [(0, "rs_1")] - - def test_no_signature_delta_without_a_thinking_block(self): - chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=[])) - assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 6479c43ee7a..790bcd269e0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -513,15 +513,14 @@ class TestTranslateMessagesToResponsesInput: result = _translate_messages(messages) assert "id" not in result[0] - def test_thinking_blocks_sharing_a_signature_become_one_reasoning_item(self): - """Summary parts of one upstream reasoning item are regrouped by their signature.""" + def test_consecutive_thinking_blocks_become_one_reasoning_item(self): + """Summary parts of one upstream reasoning item are regrouped into that item.""" messages = [ { "role": "assistant", "content": [ - {"type": "thinking", "thinking": "First part.", "signature": "rs_abc123"}, - {"type": "thinking", "thinking": "Second part.", "signature": "rs_abc123"}, - {"type": "thinking", "thinking": "A later item.", "signature": "rs_def456"}, + {"type": "thinking", "thinking": "First part."}, + {"type": "thinking", "thinking": "Second part."}, ], } ] @@ -533,13 +532,26 @@ class TestTranslateMessagesToResponsesInput: {"type": "summary_text", "text": "First part."}, {"type": "summary_text", "text": "Second part."}, ], - }, - { - "type": "reasoning", - "summary": [{"type": "summary_text", "text": "A later item."}], - }, + } ] + def test_a_tool_call_splits_the_reasoning_items_around_it(self): + """Thinking on either side of a tool call belongs to two different reasoning items.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Before the call."}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Denver"}}, + {"type": "thinking", "thinking": "After the call."}, + ], + } + ] + result = _translate_messages(messages) + assert [item["type"] for item in result] == ["reasoning", "function_call", "reasoning"] + assert result[0]["summary"] == [{"type": "summary_text", "text": "Before the call."}] + assert result[2]["summary"] == [{"type": "summary_text", "text": "After the call."}] + def test_thinking_and_text_stay_separate(self): """The visible answer stays the only thing in the assistant message.""" messages = [ @@ -1237,12 +1249,12 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [] - def test_reasoning_item_id_becomes_thinking_signature(self): - """The reasoning item id rides back as the signature so the next turn can regroup it.""" + def test_reasoning_item_id_never_becomes_a_thinking_signature(self): + """Only Anthropic can sign a thinking block, so a stand-in signature is never invented.""" reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") response = _make_mock_response(output=[reasoning]) result: Any = _ADAPTER.translate_response(response) - assert [block["signature"] for block in result["content"]] == ["rs_abc123", "rs_abc123"] + assert [block["signature"] for block in result["content"]] == [None, None] def test_dict_reasoning_item_becomes_thinking_block(self): """A reasoning item arriving as a plain dict is kept, not dropped.""" @@ -1257,9 +1269,19 @@ class TestTranslateResponse: ) result: Any = _ADAPTER.translate_response(response) assert result["content"] == [ - {"type": "thinking", "thinking": "Weighing the options.", "signature": "rs_dict_1"} + {"type": "thinking", "thinking": "Weighing the options.", "signature": None} ] + def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self): + """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _drop_unsignable_thinking_blocks, + ) + + response = _make_mock_response(output=[_make_reasoning_item(["Part one."], item_id="rs_abc123")]) + result: Any = _ADAPTER.translate_response(response) + assert _drop_unsignable_thinking_blocks(result["content"]) == [] + def test_usage_mapped_correctly(self): """Input/output tokens from ResponseAPIUsage are mapped to AnthropicUsage.""" response = _make_mock_response( diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 0fd9a381a5a..d4fcbc823a6 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -300,6 +300,7 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): "cache_control": {"type": "ephemeral"}, } ], + "reasoning_content": "The user wants me to read a file.", "provider_specific_fields": {"thought_signature": "sig-top"}, "tool_calls": [ { @@ -327,6 +328,7 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): transformed_messages = request["messages"] assert not _find_key_anywhere(transformed_messages, "thinking_blocks") + assert not _find_key_anywhere(transformed_messages, "reasoning_content") assert not _find_key_anywhere(transformed_messages, "provider_specific_fields") assert not _find_key_anywhere(transformed_messages, "cache_control") diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index e728fc4bc40..63a749dab84 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -473,12 +473,14 @@ def test_transform_messages_helper_strips_thinking_blocks(): "thinking_blocks": [ {"type": "thinking", "thinking": "internal", "signature": ""} ], + "reasoning_content": "internal", }, ] out = config._transform_messages_helper( messages, model="accounts/fireworks/models/glm-5p1", litellm_params={} ) assert "thinking_blocks" not in out[1] + assert "reasoning_content" not in out[1] assert out[1]["content"] == "I can help." diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index e316cd14dd4..82b05601a85 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -200,6 +200,7 @@ def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content(): "signature": "abc123", } ], + "reasoning_content": "Let me reason about this...", }, { "role": "user", @@ -218,6 +219,7 @@ def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content(): assert isinstance(assistant_msg["content"], str) assert assistant_msg["content"] == "Here is my answer." assert "thinking_blocks" not in assistant_msg + assert "reasoning_content" not in assistant_msg def test_hosted_vllm_thinking_blocks_with_list_content(): From 12c4652fc424800d14c6d942144064a9a09a09b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:43:51 -0700 Subject: [PATCH 043/156] fix: read bind values from the USING the command expression has closed A `JOIN ... USING` inside a subquery that helps build an EXECUTE's command was taken for the start of its bind values, so anything written after it went unscanned and a rewrite there was never reported. Only a `USING` with the parentheses closed can be the bind-values clause. --- .../check_migrations_no_data_rewrites.py | 11 +++++++--- .../test_check_migrations_no_data_rewrites.py | 21 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index c82d52167b4..b52289b4fe2 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -566,9 +566,14 @@ def bind_values_start(statement: str) -> int: """Where a statement stops handing commands to the server and starts listing bind values. The expressions after `USING` are values substituted into the command, never commands in their own right, so one that merely spells out a rewrite is not running it. Read off the - masked text, so a `USING` written inside the command string is not mistaken for this one.""" - keyword = BIND_VALUES.search(statement) - return len(statement) if keyword is None else keyword.start() + masked text, so a `USING` written inside the command string is not mistaken for this one, + and only once the parentheses have closed, so that the `USING` of a `JOIN` in a subquery + that helps build the command does not cut the command short and hide the rest of it.""" + for keyword in BIND_VALUES.finditer(statement): + preceding = statement[: keyword.start()] + if preceding.count("(") == preceding.count(")"): + return keyword.start() + return len(statement) def statement_start(statement: re.Match[str]) -> int: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 970bc9f5264..866a0d6bfe6 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -635,6 +635,27 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == ("DELETE",) + def test_a_join_using_in_a_subquery_building_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_join_using_does_not_take_the_place_of_the_real_bind_values(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = $1' USING 2;\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_bind_value_naming_a_rewrite_after_a_subquery_join_is_not_run(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) USING 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_execute_of_ddl_passes(self, tmp_path): sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" assert _keywords(tmp_path, sql) == () From 6befeb8a17f93bd65b028fdd0b8c4c3087c92a1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:48:32 -0700 Subject: [PATCH 044/156] fix(interactions): stop the cost poll loop instead of spinning on a non-positive interval --- litellm/interactions/background_cost_polling.py | 2 +- .../interactions/test_background_cost_polling.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 262d8e3260d..f2957d08840 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -74,7 +74,7 @@ async def _fetch_interaction(context: BackgroundInteractionPollContext) -> Inter def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: elapsed = 0.0 interval = initial - while elapsed + interval <= timeout: + while interval > 0 and elapsed + interval <= timeout: yield interval elapsed += interval interval = min(interval * 2, maximum) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 64436ebe4a7..7908e0e8f17 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -1,11 +1,13 @@ import asyncio import time +from itertools import islice from typing import Optional import pytest from litellm.interactions.background_cost_polling import ( _SETTLED_KEY, + _poll_intervals, BackgroundInteractionPollContext, maybe_schedule_background_interaction_cost_polling, maybe_settle_background_interaction_before_delete, @@ -92,6 +94,17 @@ def _fetch_sequence(*responses): return fetch, calls +@pytest.mark.parametrize( + "initial, maximum", + [(0.0, 0.002), (0.001, 0.0), (-1.0, 0.002), (0.0, 0.0)], +) +def test_poll_intervals_stops_instead_of_looping_on_a_non_positive_interval(initial, maximum): + intervals = list(islice(_poll_intervals(initial=initial, maximum=maximum, timeout=3600.0), 10)) + + assert len(intervals) < 10 + assert all(interval > 0 for interval in intervals) + + @pytest.mark.asyncio async def test_poller_bills_once_when_interaction_completes(): logging_obj = _logging_obj() From 5317a5ab5019c1d40c2ae4994876e2ed6a81b581 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:53:18 -0700 Subject: [PATCH 045/156] test: cover the reverse bridge on an assistant message that precedes its function_call --- ..._tool_output_order_preserved_for_gemini.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py index 3a1c77d1dab..3dec1d571c4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py @@ -114,3 +114,46 @@ def test_assistant_message_after_tool_call_is_folded_into_it(): tool_call_idx = next(i for i, m in enumerate(msgs) if isinstance(m, dict) and m.get("tool_calls")) assert msgs[tool_call_idx].get("role") == "assistant" assert msgs[tool_call_idx + 1].get("role") == "tool" + + +def test_assistant_message_before_function_call_keeps_one_assistant_turn(): + """The chat->responses bridge emits an assistant message ahead of its function_call. + + Round-tripping that order back to chat must fold both into a single assistant + turn, so the tool result still follows the message that made the call. + """ + msgs = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=[ + { + "role": "user", + "type": "message", + "content": [{"type": "input_text", "text": "What is the weather?"}], + }, + { + "role": "assistant", + "type": "message", + "content": [{"type": "output_text", "text": "Let me check."}], + }, + { + "type": "function_call", + "name": "get_weather", + "call_id": "call_1", + "arguments": "{}", + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "sunny", + }, + ] + ) + + assistant_msgs = [m for m in msgs if isinstance(m, dict) and m.get("role") == "assistant"] + assert len(assistant_msgs) == 1 + assistant = assistant_msgs[0] + assert assistant["content"] == [{"type": "text", "text": "Let me check."}] + assert [tc["function"]["name"] for tc in assistant["tool_calls"]] == ["get_weather"] + + assistant_idx = msgs.index(assistant) + assert msgs[assistant_idx + 1].get("role") == "tool" + assert msgs[assistant_idx + 1].get("tool_call_id") == "call_1" From c55b400f41054d35652380ea2ed1bada6b986527 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:55:30 -0700 Subject: [PATCH 046/156] fix(databricks): bill cached tokens at cache rates and add missing Claude pricing Databricks cost calculation multiplied every prompt token by the input rate, so a cache read cost the same as an uncached token. Route it through generic_cost_per_token, which already understands cache reads and cache writes, and add the cache rates the registry was missing. Adds Claude Opus 4.7, Opus 4.8, Opus 5, Sonnet 5 and Fable 5 on Databricks. --- litellm/llms/databricks/cost_calculator.py | 61 ++--- ...odel_prices_and_context_window_backup.json | 243 +++++++++++++++++- model_prices_and_context_window.json | 243 +++++++++++++++++- .../test_databricks_cost_calculator.py | 90 +++++++ 4 files changed, 579 insertions(+), 58 deletions(-) create mode 100644 tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 05647883ebf..64166e6fc11 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -3,10 +3,31 @@ Helper util for handling databricks-specific cost calculation - e.g.: handling 'dbrx-instruct-*' """ +from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import Usage -from litellm.utils import get_model_info + +_LEGACY_ENDPOINT_NAMES: Final = MappingProxyType( + { + "dbrx-instruct": "databricks-dbrx-instruct", + "meta-llama-3.1-70b-instruct": "databricks-meta-llama-3-1-70b-instruct", + "meta-llama-3.1-405b-instruct": "databricks-meta-llama-3-1-405b-instruct", + "mixtral-8x7b-instruct-v0.1": "databricks-mixtral-8x7b-instruct", + "bge-large-en": "databricks-bge-large-en", + "gte-large-en": "databricks-gte-large-en", + "llama-2-70b-chat": "databricks-llama-2-70b-chat", + } +) + + +def _registry_key(model: str) -> str: + name: Final = model.removeprefix("databricks/") + return next( + (key for prefix, key in _LEGACY_ENDPOINT_NAMES.items() if name.startswith(prefix)), + name, + ) def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: @@ -20,36 +41,8 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ - base_model = model - if model.startswith("databricks/dbrx-instruct") or model.startswith("dbrx-instruct"): - base_model = "databricks-dbrx-instruct" - elif model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith("meta-llama-3.1-70b-instruct"): - base_model = "databricks-meta-llama-3-1-70b-instruct" - elif model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( - "meta-llama-3.1-405b-instruct" - ): - base_model = "databricks-meta-llama-3-1-405b-instruct" - elif ( - model.startswith("databricks/mixtral-8x7b-instruct-v0.1") - or model.startswith("mixtral-8x7b-instruct-v0.1") - or model.startswith("databricks/mixtral-8x7b-instruct-v0.1") - or model.startswith("mixtral-8x7b-instruct-v0.1") - ): - base_model = "databricks-mixtral-8x7b-instruct" - elif model.startswith("databricks/bge-large-en") or model.startswith("bge-large-en"): - base_model = "databricks-bge-large-en" - elif model.startswith("databricks/gte-large-en") or model.startswith("gte-large-en"): - base_model = "databricks-gte-large-en" - elif model.startswith("databricks/llama-2-70b-chat") or model.startswith("llama-2-70b-chat"): - base_model = "databricks-llama-2-70b-chat" - ## GET MODEL INFO - model_info: Final = get_model_info(model=base_model, custom_llm_provider="databricks") - - ## CALCULATE INPUT COST - - prompt_cost: Final[float] = usage["prompt_tokens"] * model_info["input_cost_per_token"] - - ## CALCULATE OUTPUT COST - completion_cost: Final = usage["completion_tokens"] * model_info["output_cost_per_token"] - - return prompt_cost, completion_cost + return generic_cost_per_token( + model=_registry_key(model), + usage=usage, + custom_llm_provider="databricks", + ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3af7d9e5019..0bbb4b7e377 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14566,6 +14566,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14581,10 +14583,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-fable-5": { + "cache_creation_input_token_cost": 1.2500075e-05, + "cache_read_input_token_cost": 1.0000060000000001e-06, + "input_cost_per_token": 1.000006e-05, + "input_dbu_cost_per_token": 0.000142858, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 5.0000020000000004e-05, + "output_dbu_cost_per_token": 0.000714286, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": true + }, "databricks/databricks-claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.2500250000000002e-06, + "cache_read_input_token_cost": 1.0000200000000002e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14600,10 +14633,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { + "cache_creation_input_token_cost": 1.8750025e-05, + "cache_read_input_token_cost": 1.500002e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14619,10 +14655,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { + "cache_creation_input_token_cost": 1.8750025e-05, + "cache_read_input_token_cost": 1.500002e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14638,10 +14677,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14657,11 +14699,14 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14677,10 +14722,93 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-opus-4-7": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 2048, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-4-8": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-5": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-claude-sonnet-4": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14696,10 +14824,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14715,10 +14846,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14734,10 +14868,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14753,10 +14890,40 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-sonnet-5": { + "cache_creation_input_token_cost": 2.4999625e-06, + "cache_read_input_token_cost": 1.9999700000000004e-07, + "input_cost_per_token": 1.9999700000000004e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Claude Sonnet 5 DBU rates are Anthropic's introductory launch pricing, in effect through 2026-08-31, after which the standard Sonnet 5 rates (equal to Sonnet 4.5 / 4.6) take effect." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-gemini-2-5-flash": { + "cache_creation_input_token_cost": 3.0001999999999996e-07, + "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, "litellm_provider": "databricks", @@ -14771,9 +14938,12 @@ "output_dbu_cost_per_token": 3.5714e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-2-5-pro": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14788,9 +14958,12 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-flash-lite": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14805,9 +14978,12 @@ "output_dbu_cost_per_token": 2.6786e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14822,9 +14998,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { + "cache_creation_input_token_cost": 6.250300000000001e-07, + "cache_read_input_token_cost": 6.250300000000001e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -14839,9 +15018,12 @@ "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14856,6 +15038,7 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { @@ -14874,6 +15057,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14886,9 +15071,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14901,9 +15089,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-max": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14916,9 +15107,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { + "cache_creation_input_token_cost": 2.4997000000000006e-07, + "cache_read_input_token_cost": 2.4997000000000005e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -14931,9 +15125,12 @@ "mode": "chat", "output_cost_per_token": 1.99997e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14946,9 +15143,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14961,9 +15161,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14976,9 +15179,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14991,9 +15197,12 @@ "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-mini": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.499800000000001e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15006,9 +15215,12 @@ "mode": "chat", "output_cost_per_token": 4.50002e-06, "output_dbu_cost_per_token": 6.4286e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-nano": { + "cache_creation_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost": 1.9999000000000003e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15021,9 +15233,12 @@ "mode": "chat", "output_cost_per_token": 1.24999e-06, "output_dbu_cost_per_token": 1.7857e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { + "cache_creation_input_token_cost": 2.4997000000000006e-07, + "cache_read_input_token_cost": 2.4997000000000005e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15036,9 +15251,12 @@ "mode": "chat", "output_cost_per_token": 1.9999700000000004e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-nano": { + "cache_creation_input_token_cost": 4.998e-08, + "cache_read_input_token_cost": 4.998000000000001e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15051,7 +15269,8 @@ "mode": "chat", "output_cost_per_token": 3.9998000000000007e-07, "output_dbu_cost_per_token": 5.714000000000001e-06, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { "input_cost_per_token": 1.5000999999999998e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3af7d9e5019..0bbb4b7e377 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14566,6 +14566,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14581,10 +14583,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-fable-5": { + "cache_creation_input_token_cost": 1.2500075e-05, + "cache_read_input_token_cost": 1.0000060000000001e-06, + "input_cost_per_token": 1.000006e-05, + "input_dbu_cost_per_token": 0.000142858, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 5.0000020000000004e-05, + "output_dbu_cost_per_token": 0.000714286, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": true + }, "databricks/databricks-claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.2500250000000002e-06, + "cache_read_input_token_cost": 1.0000200000000002e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14600,10 +14633,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { + "cache_creation_input_token_cost": 1.8750025e-05, + "cache_read_input_token_cost": 1.500002e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14619,10 +14655,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { + "cache_creation_input_token_cost": 1.8750025e-05, + "cache_read_input_token_cost": 1.500002e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14638,10 +14677,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14657,11 +14699,14 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14677,10 +14722,93 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-opus-4-7": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 2048, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-4-8": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-5": { + "cache_creation_input_token_cost": 6.2500375e-06, + "cache_read_input_token_cost": 5.000030000000001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-claude-sonnet-4": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14696,10 +14824,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14715,10 +14846,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14734,10 +14868,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14753,10 +14890,40 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-sonnet-5": { + "cache_creation_input_token_cost": 2.4999625e-06, + "cache_read_input_token_cost": 1.9999700000000004e-07, + "input_cost_per_token": 1.9999700000000004e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Claude Sonnet 5 DBU rates are Anthropic's introductory launch pricing, in effect through 2026-08-31, after which the standard Sonnet 5 rates (equal to Sonnet 4.5 / 4.6) take effect." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-gemini-2-5-flash": { + "cache_creation_input_token_cost": 3.0001999999999996e-07, + "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, "litellm_provider": "databricks", @@ -14771,9 +14938,12 @@ "output_dbu_cost_per_token": 3.5714e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-2-5-pro": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14788,9 +14958,12 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-flash-lite": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14805,9 +14978,12 @@ "output_dbu_cost_per_token": 2.6786e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14822,9 +14998,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { + "cache_creation_input_token_cost": 6.250300000000001e-07, + "cache_read_input_token_cost": 6.250300000000001e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -14839,9 +15018,12 @@ "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14856,6 +15038,7 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { @@ -14874,6 +15057,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14886,9 +15071,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14901,9 +15089,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-max": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14916,9 +15107,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { + "cache_creation_input_token_cost": 2.4997000000000006e-07, + "cache_read_input_token_cost": 2.4997000000000005e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -14931,9 +15125,12 @@ "mode": "chat", "output_cost_per_token": 1.99997e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14946,9 +15143,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14961,9 +15161,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { + "cache_creation_input_token_cost": 1.7500000000000002e-06, + "cache_read_input_token_cost": 1.7500000000000002e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14976,9 +15179,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14991,9 +15197,12 @@ "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-mini": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.499800000000001e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15006,9 +15215,12 @@ "mode": "chat", "output_cost_per_token": 4.50002e-06, "output_dbu_cost_per_token": 6.4286e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-nano": { + "cache_creation_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost": 1.9999000000000003e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15021,9 +15233,12 @@ "mode": "chat", "output_cost_per_token": 1.24999e-06, "output_dbu_cost_per_token": 1.7857e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { + "cache_creation_input_token_cost": 2.4997000000000006e-07, + "cache_read_input_token_cost": 2.4997000000000005e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15036,9 +15251,12 @@ "mode": "chat", "output_cost_per_token": 1.9999700000000004e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-nano": { + "cache_creation_input_token_cost": 4.998e-08, + "cache_read_input_token_cost": 4.998000000000001e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15051,7 +15269,8 @@ "mode": "chat", "output_cost_per_token": 3.9998000000000007e-07, "output_dbu_cost_per_token": 5.714000000000001e-06, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { "input_cost_per_token": 1.5000999999999998e-07, diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py new file mode 100644 index 00000000000..e2be7812c1b --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -0,0 +1,90 @@ +from typing import Final + +import pytest + +import litellm +from litellm.llms.databricks.cost_calculator import cost_per_token +from litellm.types.utils import Usage + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def _model_info(model: str) -> dict: + return litellm.get_model_info(model=model, custom_llm_provider="databricks") + + +@pytest.mark.parametrize( + "model", + [ + "databricks/databricks-claude-opus-4-8", + "databricks/databricks-claude-opus-5", + "databricks/databricks-claude-sonnet-5", + ], +) +def test_cached_tokens_bill_at_cache_rates(local_model_cost_map, model): + info: Final = _model_info(model) + usage: Final = Usage( + prompt_tokens=11000, + completion_tokens=500, + total_tokens=11500, + cache_creation_input_tokens=2000, + cache_read_input_tokens=8000, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx( + 1000 * info["input_cost_per_token"] + + 2000 * info["cache_creation_input_token_cost"] + + 8000 * info["cache_read_input_token_cost"] + ) + assert completion_cost == pytest.approx(500 * info["output_cost_per_token"]) + assert prompt_cost < 11000 * info["input_cost_per_token"] + + +def test_uncached_request_bills_every_prompt_token_at_the_input_rate(local_model_cost_map): + model: Final = "databricks/databricks-claude-sonnet-5" + info: Final = _model_info(model) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx(1000 * info["input_cost_per_token"]) + assert completion_cost == pytest.approx(200 * info["output_cost_per_token"]) + + +def test_legacy_endpoint_names_still_resolve(local_model_cost_map): + info: Final = _model_info("databricks/databricks-mixtral-8x7b-instruct") + usage: Final = Usage(prompt_tokens=100, completion_tokens=100, total_tokens=200) + + prompt_cost, completion_cost = cost_per_token(model="databricks/mixtral-8x7b-instruct-v0.1", usage=usage) + + assert prompt_cost == pytest.approx(100 * info["input_cost_per_token"]) + assert completion_cost == pytest.approx(100 * info["output_cost_per_token"]) + + +@pytest.mark.parametrize( + "model", + [ + "databricks/databricks-claude-opus-4-7", + "databricks/databricks-claude-opus-4-8", + "databricks/databricks-claude-opus-5", + "databricks/databricks-claude-sonnet-5", + "databricks/databricks-claude-fable-5", + ], +) +def test_new_models_carry_cache_pricing(local_model_cost_map, model): + info: Final = _model_info(model) + + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert info["cache_creation_input_token_cost"] == pytest.approx(1.25 * info["input_cost_per_token"], rel=1e-4) + assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"], rel=1e-4) + assert info["supports_prompt_caching"] is True From a46f919e8d746b97b9d19fb6f6be0fdb2b37b626 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:05:10 -0700 Subject: [PATCH 047/156] fix: read a set-operated insert's row source term by term An `INSERT` whose `VALUES` list holds a scalar subquery was reported as a rewrite whenever that list was not the plain top-level one: joined to another term by `UNION`, `INTERSECT` or `EXCEPT`, or written inside parentheses, which Postgres accepts. Both shapes insert a fixed handful of rows, so the gate was rejecting migrations that do nothing wrong. A set operation is now split into its terms and each is read on its own, since the insert is a rewrite when any one term is a query. A row source kept in parentheses is read on its own terms too. The operators are found outside every parenthesis, so a set operation written inside a `VALUES` list does not cut the list in half. --- .../check_migrations_no_data_rewrites.py | 60 ++++++++++++++++--- .../test_check_migrations_no_data_rewrites.py | 33 ++++++++++ 2 files changed, 86 insertions(+), 7 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index b52289b4fe2..1cb61586060 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -115,6 +115,8 @@ REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) JOINS_QUERIES = ("UNION", "INTERSECT", "EXCEPT") +SET_OPERATION = re.compile(rf"\b(?:{'|'.join(JOINS_QUERIES)})\b", re.IGNORECASE) + STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( { "INSERT", @@ -378,20 +380,64 @@ def offending_keyword(statement: str) -> str | None: def row_source_keyword(statement: str) -> str | None: """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list does. A query outside every parenthesis is the row source outright. Failing that, a - `VALUES` outside every parenthesis is itself the row source, so the scalar subqueries - and helper CTEs nested within that list do not make the insert a rewrite, though only - while no set operation sits beside it at that same level: one that does joins the list - to a second query term, and that term is the row source however deeply it is - parenthesised. Failing both, the rows come from a parenthesised query, which Postgres + set operation at that same level joins several terms, and the insert is a rewrite when + any one of them is a query, so each term is read on its own rather than the statement + read whole. Failing that, a `VALUES` outside every parenthesis is itself the row source, + so the scalar subqueries and helper CTEs nested within that list do not make the insert + a rewrite. Failing all three, the rows come from a parenthesised group, which Postgres accepts and which reading only the unparenthesised text would let through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: return joined - if contains(outer, "VALUES") and not any(contains(outer, word) for word in JOINS_QUERIES): + if SET_OPERATION.search(outer): + sources = (row_source_keyword(term) for term in set_operation_terms(statement, outer)) + return next((source for source in sources if source is not None), None) + if contains(outer, "VALUES"): return None - return row_source_in(statement) + wrapped = parenthesised_row_source(statement) + return row_source_in(statement) if wrapped is None else row_source_keyword(wrapped) + + +def set_operation_terms(statement: str, outer: str) -> Iterator[str]: + """The terms a top-level set operation joins. The operators are read from the text outside + every parenthesis, which `strip_parens` blanks in place rather than removing, so their + offsets are offsets into the statement itself and each term comes back from the original + text with its own parentheses intact. Reading them at that level is what keeps a set + operation written inside a `VALUES` list from cutting the list in half. An `ALL` or a + `DISTINCT` stays at the head of the term that follows, where it names no row source and + so reads as nothing.""" + edges = [0] + for operation in SET_OPERATION.finditer(outer): + edges += [operation.start(), operation.end()] + edges.append(len(statement)) + + for opens, closes in zip(edges[::2], edges[1::2]): + yield statement[opens:closes] + + +def parenthesised_row_source(statement: str) -> str | None: + """What the last group of parentheses closed at the statement's outermost level holds, + which is where an `INSERT` keeps a row source it has wrapped, the column list before it + being a group of its own. Postgres takes `INSERT INTO "t" ("a") (SELECT ...)` and + `... (VALUES (1))` alike, so reading the wrapped text on its own terms is what stops a + scalar subquery nested inside a wrapped `VALUES` list standing in for the rows.""" + depth = 0 + opens = None + wrapped = None + + for index, character in enumerate(statement): + if character == "(": + if depth == 0: + opens = index + depth += 1 + elif character == ")": + depth = max(depth - 1, 0) + if depth == 0 and opens is not None: + wrapped = statement[opens + 1 : index] + + return wrapped def row_source_in(text: str) -> str | None: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 866a0d6bfe6..c6c41a40c1a 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -161,6 +161,39 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") VALUES ((SELECT 1 UNION SELECT 2 LIMIT 1));' assert _keywords(tmp_path, sql) == () + def test_a_scalar_subquery_in_a_set_operated_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"))' + " UNION ALL VALUES (2);" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_a_parenthesised_values_list_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")));' + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_set_operated_parenthesised_values_lists_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + " UNION ALL (VALUES (2));" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_set_operation_inside_a_values_list_does_not_split_the_terms(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"' + ' UNION SELECT max("id") FROM "Bar")) UNION ALL VALUES (2);' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_term_written_before_a_values_term_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") UNION ALL (VALUES (2));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_term_beside_parenthesised_values_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + def test_a_table_row_source_is_flagged(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) From 6c23f5ffba4f146f68e77a92c8eb23b6a88d904a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:12:02 -0700 Subject: [PATCH 048/156] fix(databricks): price Sonnet 5 at standard rates past the introductory window The introductory DBU rates run through 2026-08-31 and pricing carries no expiry date, so a static introductory entry would undercharge by a third from September 1 and let spend outrun enforced budgets. Ship the standard rates, which match Sonnet 4.5 and 4.6, and keep the introductory numbers in the entry notes. Also give the new cost calculator tests full type annotations. --- ...odel_prices_and_context_window_backup.json | 14 +++++------ model_prices_and_context_window.json | 14 +++++------ .../test_databricks_cost_calculator.py | 25 +++++++++++++------ 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0bbb4b7e377..e296e828b92 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14895,20 +14895,20 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 2.4999625e-06, - "cache_read_input_token_cost": 1.9999700000000004e-07, - "input_cost_per_token": 1.9999700000000004e-06, - "input_dbu_cost_per_token": 2.8571e-05, + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Claude Sonnet 5 DBU rates are Anthropic's introductory launch pricing, in effect through 2026-08-31, after which the standard Sonnet 5 rates (equal to Sonnet 4.5 / 4.6) take effect." + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Anthropic's introductory launch rates (28.571 input / 142.857 output DBU) run through 2026-08-31. The standard rates listed here, equal to Sonnet 4.5 / 4.6, are used instead because pricing carries no expiry date, and undercharging past the window would let spend outrun enforced budgets." }, "mode": "chat", - "output_cost_per_token": 9.999990000000002e-06, - "output_dbu_cost_per_token": 0.000142857, + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0bbb4b7e377..e296e828b92 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14895,20 +14895,20 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 2.4999625e-06, - "cache_read_input_token_cost": 1.9999700000000004e-07, - "input_cost_per_token": 1.9999700000000004e-06, - "input_dbu_cost_per_token": 2.8571e-05, + "cache_creation_input_token_cost": 3.7499875e-06, + "cache_read_input_token_cost": 2.9999900000000006e-07, + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Claude Sonnet 5 DBU rates are Anthropic's introductory launch pricing, in effect through 2026-08-31, after which the standard Sonnet 5 rates (equal to Sonnet 4.5 / 4.6) take effect." + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Anthropic's introductory launch rates (28.571 input / 142.857 output DBU) run through 2026-08-31. The standard rates listed here, equal to Sonnet 4.5 / 4.6, are used instead because pricing carries no expiry date, and undercharging past the window would let spend outrun enforced budgets." }, "mode": "chat", - "output_cost_per_token": 9.999990000000002e-06, - "output_dbu_cost_per_token": 0.000142857, + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index e2be7812c1b..9353292c668 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -1,14 +1,15 @@ +from collections.abc import Iterator from typing import Final import pytest import litellm from litellm.llms.databricks.cost_calculator import cost_per_token -from litellm.types.utils import Usage +from litellm.types.utils import ModelInfo, Usage @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: 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() @@ -16,7 +17,7 @@ def local_model_cost_map(monkeypatch): litellm.get_model_info.cache_clear() -def _model_info(model: str) -> dict: +def _model_info(model: str) -> ModelInfo: return litellm.get_model_info(model=model, custom_llm_provider="databricks") @@ -28,7 +29,7 @@ def _model_info(model: str) -> dict: "databricks/databricks-claude-sonnet-5", ], ) -def test_cached_tokens_bill_at_cache_rates(local_model_cost_map, model): +def test_cached_tokens_bill_at_cache_rates(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) usage: Final = Usage( prompt_tokens=11000, @@ -49,7 +50,7 @@ def test_cached_tokens_bill_at_cache_rates(local_model_cost_map, model): assert prompt_cost < 11000 * info["input_cost_per_token"] -def test_uncached_request_bills_every_prompt_token_at_the_input_rate(local_model_cost_map): +def test_uncached_request_bills_every_prompt_token_at_the_input_rate(local_model_cost_map: None) -> None: model: Final = "databricks/databricks-claude-sonnet-5" info: Final = _model_info(model) usage: Final = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200) @@ -60,7 +61,7 @@ def test_uncached_request_bills_every_prompt_token_at_the_input_rate(local_model assert completion_cost == pytest.approx(200 * info["output_cost_per_token"]) -def test_legacy_endpoint_names_still_resolve(local_model_cost_map): +def test_legacy_endpoint_names_still_resolve(local_model_cost_map: None) -> None: info: Final = _model_info("databricks/databricks-mixtral-8x7b-instruct") usage: Final = Usage(prompt_tokens=100, completion_tokens=100, total_tokens=200) @@ -80,7 +81,7 @@ def test_legacy_endpoint_names_still_resolve(local_model_cost_map): "databricks/databricks-claude-fable-5", ], ) -def test_new_models_carry_cache_pricing(local_model_cost_map, model): +def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) assert info["input_cost_per_token"] > 0 @@ -88,3 +89,13 @@ def test_new_models_carry_cache_pricing(local_model_cost_map, model): assert info["cache_creation_input_token_cost"] == pytest.approx(1.25 * info["input_cost_per_token"], rel=1e-4) assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"], rel=1e-4) assert info["supports_prompt_caching"] is True + + +def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: + sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") + sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") + + assert sonnet_5["input_cost_per_token"] == pytest.approx(sonnet_4_6["input_cost_per_token"]) + assert sonnet_5["output_cost_per_token"] == pytest.approx(sonnet_4_6["output_cost_per_token"]) + assert sonnet_5["cache_creation_input_token_cost"] == pytest.approx(sonnet_4_6["cache_creation_input_token_cost"]) + assert sonnet_5["cache_read_input_token_cost"] == pytest.approx(sonnet_4_6["cache_read_input_token_cost"]) From 7e59f8c2095fb86d9d8186988ca4fc9e9d832623 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:20:24 -0700 Subject: [PATCH 049/156] fix: read every parenthesised group for the row source, not the last Taking the last group at the statement's outermost level assumed the row source was written there, and an insert is allowed to carry more after it: `(SELECT ...) ON CONFLICT ("id") DO NOTHING` ends on the conflict target and `... RETURNING ("id")` on the returning list, so the query supplying the rows was never reached and a full table copy passed the gate. Each group is now read on its own terms and the first to name a row source is the answer, since the others are the column list and the clauses an insert may carry, none of which names one. --- .../check_migrations_no_data_rewrites.py | 29 ++++++++++--------- .../test_check_migrations_no_data_rewrites.py | 18 ++++++++++++ 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 1cb61586060..17bc90bd0d6 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -386,7 +386,10 @@ def row_source_keyword(statement: str) -> str | None: so the scalar subqueries and helper CTEs nested within that list do not make the insert a rewrite. Failing all three, the rows come from a parenthesised group, which Postgres accepts and which reading only the unparenthesised text would let through: - `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" + `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table. Each group at that level is + read on its own terms and the first to name a row source is the answer, since the ones + around it are the column list, the conflict target and the rest of the clauses an insert + is allowed to carry, and any of those can be the last group written.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: @@ -396,8 +399,11 @@ def row_source_keyword(statement: str) -> str | None: return next((source for source in sources if source is not None), None) if contains(outer, "VALUES"): return None - wrapped = parenthesised_row_source(statement) - return row_source_in(statement) if wrapped is None else row_source_keyword(wrapped) + groups = list(parenthesised_groups(statement)) + if not groups: + return row_source_in(statement) + sources = (row_source_keyword(group) for group in groups) + return next((source for source in sources if source is not None), None) def set_operation_terms(statement: str, outer: str) -> Iterator[str]: @@ -417,15 +423,14 @@ def set_operation_terms(statement: str, outer: str) -> Iterator[str]: yield statement[opens:closes] -def parenthesised_row_source(statement: str) -> str | None: - """What the last group of parentheses closed at the statement's outermost level holds, - which is where an `INSERT` keeps a row source it has wrapped, the column list before it - being a group of its own. Postgres takes `INSERT INTO "t" ("a") (SELECT ...)` and - `... (VALUES (1))` alike, so reading the wrapped text on its own terms is what stops a - scalar subquery nested inside a wrapped `VALUES` list standing in for the rows.""" +def parenthesised_groups(statement: str) -> Iterator[str]: + """What each group of parentheses closed at the statement's outermost level holds, in the + order they are written. One of them is where an `INSERT` keeps a row source it has + wrapped, since Postgres takes `INSERT INTO "t" ("a") (SELECT ...)` and `... (VALUES (1))` + alike, and reading a group on its own terms is what stops a scalar subquery nested inside + a wrapped `VALUES` list standing in for the rows.""" depth = 0 opens = None - wrapped = None for index, character in enumerate(statement): if character == "(": @@ -435,9 +440,7 @@ def parenthesised_row_source(statement: str) -> str | None: elif character == ")": depth = max(depth - 1, 0) if depth == 0 and opens is not None: - wrapped = statement[opens + 1 : index] - - return wrapped + yield statement[opens + 1 : index] def row_source_in(text: str) -> str | None: diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index c6c41a40c1a..b78b34e52c7 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -186,6 +186,24 @@ class TestInsert: ) assert _keywords(tmp_path, sql) == () + def test_a_conflict_target_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar")' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_returning_list_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING ("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_conflict_target_beside_a_bounded_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == () + def test_a_query_term_written_before_a_values_term_is_still_the_row_source(self, tmp_path): sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") UNION ALL (VALUES (2));' assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) From 8b566a7f0aa3e1c935c0d84e646379a505d6904d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:28:12 -0700 Subject: [PATCH 050/156] fix(interactions): stop two settlement paths from pinning the budget reservation Both leave a background interaction's pre-call reservation open, so the serving process keeps refusing traffic on the key at the estimated cost while its recorded spend stays near zero. A raise from the completion event propagated out with the settlement gate already claimed, and nothing retries a claim that is set, so the reservation was never released. Billing now releases it on the way out. `requires_action` was missing from the terminal set. It is terminal for the interaction it names: the API has no operation that resumes one, and a caller answers a tool request by creating a new interaction whose `previous_interaction_id` points at it. A function-calling background create that stopped there was polled until the 3600s timeout, losing the tokens it had already spent producing the tool request and holding its reservation open for that whole window. --- .../interactions/background_cost_polling.py | 28 ++++++- .../test_background_cost_polling.py | 78 +++++++++++++++++++ 2 files changed, 103 insertions(+), 3 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index f2957d08840..5d7dca80b73 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -10,6 +10,14 @@ 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. +``requires_action`` is terminal for the interaction it names. The API has no +operation that resumes one: a caller answers a tool request by creating a new +interaction whose ``previous_interaction_id`` points at it, and that new +interaction bills itself. The paused interaction keeps the tokens it already +spent producing the tool request, so it is billed and settled where it stops +rather than polled until the timeout, which would both lose that usage and +hold its budget reservation open for the whole timeout window. + 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 @@ -39,7 +47,7 @@ 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"}) +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"}) @dataclass(frozen=True, slots=True) @@ -125,7 +133,7 @@ async def poll_and_log_background_interaction_cost( 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) + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) else: await _release_open_budget_reservation(logging_obj=context.logging_obj) return @@ -162,6 +170,20 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction") +async def _bill_settled_interaction(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> None: + """ + Claiming the settlement makes the claimer solely responsible for the + reservation, and no one retries a claim that is already set. A billing + failure here must therefore release the reservation on its way out, or it + stays pinned at the estimated cost until the whole poll times out. + """ + try: + await logging_obj.async_log_background_interaction_completion(result=response) + except Exception: + await _release_open_budget_reservation(logging_obj=logging_obj) + raise + + def is_pollable_background_interaction(response: InteractionsAPIResponse) -> bool: """ The single gate deciding whether a create's response gets a poll task. @@ -247,6 +269,6 @@ async def maybe_settle_background_interaction_before_delete( 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) + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) return await _release_open_budget_reservation(logging_obj=context.logging_obj) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 7908e0e8f17..759bdc71dfc 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -59,6 +59,10 @@ def _logging_obj_with_reservation(reservation: dict) -> LitellmLogging: return _logging_obj(litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}) +async def _raise_on_billing(result: InteractionsAPIResponse) -> None: + raise RuntimeError("cost calculation failed for a settled background interaction") + + def _context(logging_obj: LitellmLogging, timeout_seconds: float = 1.0) -> BackgroundInteractionPollContext: return BackgroundInteractionPollContext( interaction_id="interactions/bg-abc", @@ -120,6 +124,33 @@ async def test_poller_bills_once_when_interaction_completes(): assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 +@pytest.mark.asyncio +async def test_poller_bills_an_interaction_paused_for_a_tool_result(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("requires_action", 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_does_not_pin_the_budget_for_an_interaction_paused_for_a_tool_result(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("requires_action", 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_stops_without_billing_on_terminal_status_without_usage(): logging_obj = _logging_obj() @@ -170,6 +201,19 @@ async def test_poller_releases_budget_reservation_on_timeout_give_up(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_when_billing_raises(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + logging_obj.async_log_background_interaction_completion = _raise_on_billing + + with pytest.raises(RuntimeError): + 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_leaves_reservation_reconciliation_to_the_completion_event(): reservation = _reservation() @@ -248,6 +292,22 @@ def _register_poll(logging_obj: LitellmLogging, poll_fetch=None) -> asyncio.Task return task +@pytest.mark.asyncio +async def test_delete_settlement_bills_an_interaction_paused_for_a_tool_result(): + logging_obj = _logging_obj() + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("requires_action", 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 + await asyncio.wait_for(task, timeout=5) + + @pytest.mark.asyncio async def test_delete_settlement_bills_pending_background_interaction(): logging_obj = _logging_obj() @@ -299,6 +359,24 @@ async def test_delete_settlement_releases_reservation_when_prefetch_fails(): await asyncio.wait_for(task, timeout=5) +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_billing_raises(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + logging_obj.async_log_background_interaction_completion = _raise_on_billing + + with pytest.raises(RuntimeError): + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + 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)) From da3dcb139d207c9cdff006cd2fcd022007782f43 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:37:37 -0700 Subject: [PATCH 051/156] fix(databricks): charge the input rate for cache tokens on models with no cache pricing The shared cost calculator treats a missing cache rate as free, so routing Databricks through it billed cached tokens at zero on the 14 entries that publish no cache pricing. On a 10,000 token prompt with 8,000 cache reads that is $0.0010000 against the correct $0.0050001, a fivefold undercharge. Those entries now declare cache rates equal to their input rate, which is what a model with no caching discount should charge, and a test pins every priced Databricks entry to declaring cache rates so no future entry can regress into it. Also repoints the provider-neutral generalization test off an id the new Opus 5 entry now shadows, adds backup-to-main parity tests for the five new entries, pins that Databricks Claude is never auto-injected with cache control despite reporting caching support, and trims the Sonnet 5 pricing note, which is served on an unauthenticated route. --- ...odel_prices_and_context_window_backup.json | 30 ++++++++++- model_prices_and_context_window.json | 30 ++++++++++- .../test_anthropic_cache_control_hook.py | 8 +++ .../test_fallback_generalizations.py | 4 +- .../test_databricks_cost_calculator.py | 50 +++++++++++++++++++ 5 files changed, 119 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e296e828b92..f1e330cb4ff 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14551,6 +14551,8 @@ ] }, "databricks/databricks-bge-large-en": { + "cache_creation_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.0003e-07, "input_dbu_cost_per_token": 1.429e-06, "litellm_provider": "databricks", @@ -14904,7 +14906,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Anthropic's introductory launch rates (28.571 input / 142.857 output DBU) run through 2026-08-31. The standard rates listed here, equal to Sonnet 4.5 / 4.6, are used instead because pricing carries no expiry date, and undercharging past the window would let spend outrun enforced budgets." + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, @@ -15042,6 +15044,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15273,6 +15277,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15288,6 +15294,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-oss-20b": { + "cache_creation_input_token_cost": 7e-08, + "cache_read_input_token_cost": 7e-08, "input_cost_per_token": 7e-08, "input_dbu_cost_per_token": 1e-06, "litellm_provider": "databricks", @@ -15303,6 +15311,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { + "cache_creation_input_token_cost": 1.2999000000000001e-07, + "cache_read_input_token_cost": 1.2999000000000001e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15318,6 +15328,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-llama-2-70b-chat": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15334,6 +15346,8 @@ "supports_tool_choice": true }, "databricks/databricks-llama-4-maverick": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15350,6 +15364,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-405b-instruct": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 5.00003e-06, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -15366,6 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15381,6 +15399,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-meta-llama-3-3-70b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15397,6 +15417,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-70b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15413,6 +15435,8 @@ "supports_tool_choice": true }, "databricks/databricks-mixtral-8x7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15429,6 +15453,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-30b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15445,6 +15471,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e296e828b92..f1e330cb4ff 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14551,6 +14551,8 @@ ] }, "databricks/databricks-bge-large-en": { + "cache_creation_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.0003e-07, "input_dbu_cost_per_token": 1.429e-06, "litellm_provider": "databricks", @@ -14904,7 +14906,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Anthropic's introductory launch rates (28.571 input / 142.857 output DBU) run through 2026-08-31. The standard rates listed here, equal to Sonnet 4.5 / 4.6, are used instead because pricing carries no expiry date, and undercharging past the window would let spend outrun enforced budgets." + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, @@ -15042,6 +15044,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15273,6 +15277,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15288,6 +15294,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-oss-20b": { + "cache_creation_input_token_cost": 7e-08, + "cache_read_input_token_cost": 7e-08, "input_cost_per_token": 7e-08, "input_dbu_cost_per_token": 1e-06, "litellm_provider": "databricks", @@ -15303,6 +15311,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { + "cache_creation_input_token_cost": 1.2999000000000001e-07, + "cache_read_input_token_cost": 1.2999000000000001e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15318,6 +15328,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-llama-2-70b-chat": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15334,6 +15346,8 @@ "supports_tool_choice": true }, "databricks/databricks-llama-4-maverick": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15350,6 +15364,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-405b-instruct": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 5.00003e-06, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -15366,6 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { + "cache_creation_input_token_cost": 1.5000999999999998e-07, + "cache_read_input_token_cost": 1.5000999999999998e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15381,6 +15399,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-meta-llama-3-3-70b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15397,6 +15417,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-70b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15413,6 +15435,8 @@ "supports_tool_choice": true }, "databricks/databricks-mixtral-8x7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15429,6 +15453,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-30b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15445,6 +15471,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index b6e063a6d94..9281a80dfa5 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1599,6 +1599,14 @@ class TestEnableAnthropicPromptCaching: assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True assert self._points(model=model, provider=provider) == [] + def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch): + from litellm.utils import supports_prompt_caching + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + model = "databricks/databricks-claude-sonnet-4-5" + assert supports_prompt_caching(model=model, custom_llm_provider="databricks") is True + assert self._points(model=model, provider="databricks") == [] + def test_model_without_caching_support_not_injected(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) assert self._points(model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == [] diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 882429fd7cd..0587628e2fe 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -378,7 +378,7 @@ def test_shipped_rules_flag_unmapped_fable_as_always_on_thinking(shipped_cost_ma "model,provider", [ ("claude-opus-4-9@20260101", "vertex_ai"), - ("databricks-claude-opus-5-1", "databricks"), + ("databricks-claude-haiku-5-1", "databricks"), ], ) def test_shipped_rules_are_provider_neutral_for_unmapped_ids(shipped_cost_map, model, provider): @@ -388,6 +388,8 @@ def test_shipped_rules_are_provider_neutral_for_unmapped_ids(shipped_cost_map, m assert info["supports_adaptive_thinking"] is True assert info["supports_mid_conversation_system"] is True assert info["supports_function_calling"] is True + assert not info.get("input_cost_per_token") + assert not info.get("output_cost_per_token") @pytest.mark.parametrize( diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 9353292c668..102a861fe1a 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -1,4 +1,6 @@ +import json from collections.abc import Iterator +from pathlib import Path from typing import Final import pytest @@ -7,6 +9,17 @@ import litellm from litellm.llms.databricks.cost_calculator import cost_per_token from litellm.types.utils import ModelInfo, Usage +REPO_ROOT: Final = Path(__file__).parents[4] +MAIN_PRICES: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PRICES: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" +NEW_MODELS: Final = ( + "databricks/databricks-claude-opus-4-7", + "databricks/databricks-claude-opus-4-8", + "databricks/databricks-claude-opus-5", + "databricks/databricks-claude-sonnet-5", + "databricks/databricks-claude-fable-5", +) + @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: @@ -91,6 +104,43 @@ def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) assert info["supports_prompt_caching"] is True +def test_every_priced_databricks_model_declares_cache_rates(local_model_cost_map: None) -> None: + undeclared: Final = [ + model + for model, info in litellm.model_cost.items() + if model.startswith("databricks/") + and info.get("input_cost_per_token") is not None + and info.get("cache_read_input_token_cost") is None + ] + + assert undeclared == [] + + +def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( + local_model_cost_map: None, +) -> None: + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + info: Final = _model_info(model) + usage: Final = Usage( + prompt_tokens=10000, + completion_tokens=100, + total_tokens=10100, + cache_read_input_tokens=8000, + ) + + prompt_cost, _ = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx(10000 * info["input_cost_per_token"]) + + +@pytest.mark.parametrize("model", NEW_MODELS) +def test_backup_price_map_matches_main(model: str) -> None: + main_cost: Final = json.loads(MAIN_PRICES.read_text()) + backup_cost: Final = json.loads(BACKUP_PRICES.read_text()) + + assert backup_cost.get(model) == main_cost.get(model) + + def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") From e487e470c732c447050a2a40d1ff2f270b1f6397 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:51:17 -0700 Subject: [PATCH 052/156] test(databricks): pin the cache-control test to the in-repo cost map The new supports_prompt_caching assertion reads a capability this branch adds to the registry, so it only holds against the bundled map. CI leaves LITELLM_LOCAL_MODEL_COST_MAP unset and fetches main's copy, which lags the branch until merge, so the test failed there while passing locally. Use the local_model_cost_map fixture the repo already provides, matching what the other two test files in this change do. --- .../integrations/test_anthropic_cache_control_hook.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 9281a80dfa5..cf3318b32fe 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1599,7 +1599,7 @@ class TestEnableAnthropicPromptCaching: assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True assert self._points(model=model, provider=provider) == [] - def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch): + def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch, local_model_cost_map): from litellm.utils import supports_prompt_caching monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) From 68bda9499533545065299c4995f74214d0c92e05 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:53:12 -0700 Subject: [PATCH 053/156] chore: mark the new API payload literals for the type-discipline budget --- .../litellm_responses_transformation/transformation.py | 7 ++++--- .../responses_adapters/streaming_iterator.py | 5 ++++- .../responses_adapters/transformation.py | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 3e55c3c637e..cf227cfa7d1 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -94,13 +94,14 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: Stored reasoning items win because they carry an id the Responses API minted; thinking blocks are the fallback for turns that arrived over another API surface. """ - stored: Final = [_reasoning_item_to_response_input(r_item) for r_item in _get_reasoning_items(msg)] + items: Final = _get_reasoning_items(msg) + stored: Final = [_reasoning_item_to_response_input(item) for item in items] # mutable-ok: API message payload if stored: return stored raw_blocks: Final = msg.get("thinking_blocks") or () blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return [] if from_thinking is None else [dict(from_thinking)] + return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload def _build_reasoning_item( @@ -393,7 +394,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): input_items.extend(_reasoning_input_items(msg)) if content: input_items.append( - { + { # mutable-ok: API message payload "type": "message", "role": "assistant", "content": self._convert_content_to_responses_format(content, "assistant"), diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 5577d4a9c2d..292d2622c7f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -152,7 +152,10 @@ class AnthropicResponsesStreamWrapper: if block_idx < 0: if not delta: return - block_idx = self._open_block(item_id, {"type": "thinking", "thinking": "", "signature": ""}) + block_idx = self._open_block( + item_id, + {"type": "thinking", "thinking": "", "signature": ""}, # mutable-ok: API message payload + ) self._chunk_queue.append( { "type": "content_block_delta", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 1a6b0498a52..23a9a60d810 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -148,13 +148,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if btype == "thinking": blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return None if reasoning_item is None else dict(reasoning_item) + return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload if btype == "tool_use": return { # mutable-ok: API message payload "type": "function_call", "call_id": first.get("id", ""), "name": first.get("name", ""), - "arguments": json.dumps(first.get("input", {})), + "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload } return None From 0ab8ef60bf86f4fbc61359488c0d967f75b01a9b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 22 Aug 2026 22:54:08 +0000 Subject: [PATCH 054/156] fix(interactions): stop the unpollable-create path from firing a false cost-tracking alert The tightened gate correctly stopped deferring the reservation release for InteractionsAPIResponses the scheduler will not poll (terminal status, or in_progress without an id), but the response then fell through into the generic 'Cost tracking failed' raise and the failed_tracking_alert path. A create returning failed, cancelled, requires_action, incomplete, budget_exceeded, or an id-less in_progress without usage therefore released its reservation as intended and, in the same breath, alerted operators for a legitimate no-usage response, both creating noise and masking real cost-tracking failures. The two gates are now nested under a single 'unbilled interaction response' outer check, so any InteractionsAPIResponse with no usage takes either the defer path (pollable, polling on) or the release-and-return path, and none of them fall through to the generic failure raise. The two regression tests also now assert failed_tracking_alert is not called, closing the observation gap the report flagged. --- litellm/proxy/hooks/proxy_track_cost_callback.py | 16 ++++++++++++---- .../hooks/test_proxy_track_cost_callback.py | 9 ++++++++- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 46cf62ede1c..70677521d6d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -319,8 +319,10 @@ class _ProxyDBLogger(CustomLogger): elif budget_reservation is not None: await _release_budget_reservation(budget_reservation=budget_reservation) else: - if _is_unbilled_in_progress_interaction(completion_response): - if BACKGROUND_INTERACTION_COST_POLLING_ENABLED: + if _is_unbilled_interaction_response(completion_response): + if BACKGROUND_INTERACTION_COST_POLLING_ENABLED and _is_unbilled_in_progress_interaction( + completion_response + ): 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" @@ -328,8 +330,8 @@ class _ProxyDBLogger(CustomLogger): return await _release_budget_reservation(budget_reservation=budget_reservation) verbose_proxy_logger.debug( - "Background interaction cost polling is disabled; released the budget " - "reservation for an in-progress interaction that will not be billed" + "Released the budget reservation for an interaction create with no usage " + "that no poll task will settle" ) return await _release_budget_reservation(budget_reservation=budget_reservation) @@ -477,6 +479,12 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +def _is_unbilled_interaction_response(completion_response: object) -> bool: + from litellm.types.interactions import InteractionsAPIResponse + + return isinstance(completion_response, InteractionsAPIResponse) and completion_response.usage is None + + def _is_unbilled_in_progress_interaction(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import is_pollable_background_interaction from litellm.types.interactions import InteractionsAPIResponse 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 93e0cbec596..fd79f9ea3b3 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 @@ -871,6 +871,10 @@ async def test_track_cost_callback_releases_reservation_for_unpollable_interacti callback must release it there and then, or the pre-call estimate stays added to the key, user, team and org spend counters and starts refusing traffic against budget that was never actually spent. + + A no-usage terminal create is also not a cost-tracking failure, so the + callback must not fire ``failed_tracking_alert``: doing so would flood + operators with false alerts and mask real cost-tracking failures. """ from litellm.types.interactions import InteractionsAPIResponse @@ -897,6 +901,7 @@ async def test_track_cost_callback_releases_reservation_for_unpollable_interacti ) assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() @pytest.mark.asyncio @@ -904,7 +909,8 @@ async def test_track_cost_callback_releases_reservation_for_interaction_without_ """ The scheduler also refuses a response with no id, since it has nothing to poll for, so the callback must not defer to a poll task that will never - exist. + exist, and it must not fire ``failed_tracking_alert`` for what is a + legitimate no-usage response rather than a cost-tracking failure. """ from litellm.types.interactions import InteractionsAPIResponse @@ -931,6 +937,7 @@ async def test_track_cost_callback_releases_reservation_for_interaction_without_ ) assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() @pytest.mark.parametrize( From d46b0bddd7f39628e4782b2839ca4ca26291ef51 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 22 Aug 2026 22:58:04 +0000 Subject: [PATCH 055/156] fix: drop null summary text and emit reasoning-only assistant turns --- .../litellm_responses_transformation/transformation.py | 2 ++ .../responses_adapters/transformation.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index cf227cfa7d1..13dd41f9b86 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -434,6 +434,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "content": self._convert_content_to_responses_format(content, cast(str, role)), } ) + elif role == "assistant": + input_items.extend(_reasoning_input_items(msg)) return input_items, instructions diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 23a9a60d810..6d47d0de19f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -110,8 +110,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def _summary_part_text(part: object) -> str: if isinstance(part, Mapping): mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json - return str(mapping.get("text", "")) - return str(getattr(part, "text", "")) + return str(mapping.get("text") or "") + return str(getattr(part, "text", None) or "") @classmethod def _thinking_blocks_from_reasoning_item( From 444a0232728fe6b2ca826c28b44d052d2c59960e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:02:08 -0700 Subject: [PATCH 056/156] chore(databricks): write pricing at the derivation's precision The figures came out of a float product of the DBU rate and the dollar rate, so the map carried representation tails such as 2.9999900000000006e-07 where the derivation only means 2.99999e-07. Trim the 44 values this change adds or edits to the shortest literal that round-trips to the same figure, leaving every pre-existing value alone. The largest move is under 1e-15 relative, so no billed amount changes. --- ...odel_prices_and_context_window_backup.json | 88 +++++++++---------- model_prices_and_context_window.json | 88 +++++++++---------- 2 files changed, 88 insertions(+), 88 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f1e330cb4ff..8c80714c089 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14569,7 +14569,7 @@ }, "databricks/databricks-claude-3-7-sonnet": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14591,7 +14591,7 @@ }, "databricks/databricks-claude-fable-5": { "cache_creation_input_token_cost": 1.2500075e-05, - "cache_read_input_token_cost": 1.0000060000000001e-06, + "cache_read_input_token_cost": 1.000006e-06, "input_cost_per_token": 1.000006e-05, "input_dbu_cost_per_token": 0.000142858, "litellm_provider": "databricks", @@ -14602,7 +14602,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 5.0000020000000004e-05, + "output_cost_per_token": 5.000002e-05, "output_dbu_cost_per_token": 0.000714286, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14618,8 +14618,8 @@ "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { - "cache_creation_input_token_cost": 1.2500250000000002e-06, - "cache_read_input_token_cost": 1.0000200000000002e-07, + "cache_creation_input_token_cost": 1.250025e-06, + "cache_read_input_token_cost": 1.00002e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14685,7 +14685,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14708,7 +14708,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14730,7 +14730,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14741,7 +14741,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14756,7 +14756,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14767,7 +14767,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14783,7 +14783,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14794,7 +14794,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14810,7 +14810,7 @@ }, "databricks/databricks-claude-sonnet-4": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14832,7 +14832,7 @@ }, "databricks/databricks-claude-sonnet-4-1": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14854,7 +14854,7 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14876,7 +14876,7 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14898,8 +14898,8 @@ }, "databricks/databricks-claude-sonnet-5": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "cache_read_input_token_cost": 2.99999e-07, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -14909,7 +14909,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token": 1.500002e-05, "output_dbu_cost_per_token": 0.000214286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14924,7 +14924,7 @@ "supports_vision": true }, "databricks/databricks-gemini-2-5-flash": { - "cache_creation_input_token_cost": 3.0001999999999996e-07, + "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, @@ -15004,8 +15004,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { - "cache_creation_input_token_cost": 6.250300000000001e-07, - "cache_read_input_token_cost": 6.250300000000001e-08, + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -15044,8 +15044,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15115,8 +15115,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { - "cache_creation_input_token_cost": 2.4997000000000006e-07, - "cache_read_input_token_cost": 2.4997000000000005e-08, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.4997e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15133,8 +15133,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15151,8 +15151,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15169,8 +15169,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15206,7 +15206,7 @@ }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, - "cache_read_input_token_cost": 7.499800000000001e-08, + "cache_read_input_token_cost": 7.4998e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15224,7 +15224,7 @@ }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, - "cache_read_input_token_cost": 1.9999000000000003e-08, + "cache_read_input_token_cost": 1.9999e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15241,8 +15241,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { - "cache_creation_input_token_cost": 2.4997000000000006e-07, - "cache_read_input_token_cost": 2.4997000000000005e-08, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.4997e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15260,7 +15260,7 @@ }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, - "cache_read_input_token_cost": 4.998000000000001e-09, + "cache_read_input_token_cost": 4.998e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15277,8 +15277,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15311,8 +15311,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { - "cache_creation_input_token_cost": 1.2999000000000001e-07, - "cache_read_input_token_cost": 1.2999000000000001e-07, + "cache_creation_input_token_cost": 1.2999e-07, + "cache_read_input_token_cost": 1.2999e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15382,8 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f1e330cb4ff..8c80714c089 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14569,7 +14569,7 @@ }, "databricks/databricks-claude-3-7-sonnet": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14591,7 +14591,7 @@ }, "databricks/databricks-claude-fable-5": { "cache_creation_input_token_cost": 1.2500075e-05, - "cache_read_input_token_cost": 1.0000060000000001e-06, + "cache_read_input_token_cost": 1.000006e-06, "input_cost_per_token": 1.000006e-05, "input_dbu_cost_per_token": 0.000142858, "litellm_provider": "databricks", @@ -14602,7 +14602,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 5.0000020000000004e-05, + "output_cost_per_token": 5.000002e-05, "output_dbu_cost_per_token": 0.000714286, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14618,8 +14618,8 @@ "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { - "cache_creation_input_token_cost": 1.2500250000000002e-06, - "cache_read_input_token_cost": 1.0000200000000002e-07, + "cache_creation_input_token_cost": 1.250025e-06, + "cache_read_input_token_cost": 1.00002e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14685,7 +14685,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14708,7 +14708,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14730,7 +14730,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14741,7 +14741,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14756,7 +14756,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14767,7 +14767,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14783,7 +14783,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.000030000000001e-07, + "cache_read_input_token_cost": 5.00003e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14794,7 +14794,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, + "output_cost_per_token": 2.500001e-05, "output_dbu_cost_per_token": 0.000357143, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14810,7 +14810,7 @@ }, "databricks/databricks-claude-sonnet-4": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14832,7 +14832,7 @@ }, "databricks/databricks-claude-sonnet-4-1": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14854,7 +14854,7 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14876,7 +14876,7 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, + "cache_read_input_token_cost": 2.99999e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14898,8 +14898,8 @@ }, "databricks/databricks-claude-sonnet-5": { "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.9999900000000006e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "cache_read_input_token_cost": 2.99999e-07, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -14909,7 +14909,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token": 1.500002e-05, "output_dbu_cost_per_token": 0.000214286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", @@ -14924,7 +14924,7 @@ "supports_vision": true }, "databricks/databricks-gemini-2-5-flash": { - "cache_creation_input_token_cost": 3.0001999999999996e-07, + "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, @@ -15004,8 +15004,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { - "cache_creation_input_token_cost": 6.250300000000001e-07, - "cache_read_input_token_cost": 6.250300000000001e-08, + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -15044,8 +15044,8 @@ "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15115,8 +15115,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { - "cache_creation_input_token_cost": 2.4997000000000006e-07, - "cache_read_input_token_cost": 2.4997000000000005e-08, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.4997e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15133,8 +15133,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15151,8 +15151,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15169,8 +15169,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { - "cache_creation_input_token_cost": 1.7500000000000002e-06, - "cache_read_input_token_cost": 1.7500000000000002e-07, + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -15206,7 +15206,7 @@ }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, - "cache_read_input_token_cost": 7.499800000000001e-08, + "cache_read_input_token_cost": 7.4998e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15224,7 +15224,7 @@ }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, - "cache_read_input_token_cost": 1.9999000000000003e-08, + "cache_read_input_token_cost": 1.9999e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15241,8 +15241,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { - "cache_creation_input_token_cost": 2.4997000000000006e-07, - "cache_read_input_token_cost": 2.4997000000000005e-08, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.4997e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15260,7 +15260,7 @@ }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, - "cache_read_input_token_cost": 4.998000000000001e-09, + "cache_read_input_token_cost": 4.998e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15277,8 +15277,8 @@ "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15311,8 +15311,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { - "cache_creation_input_token_cost": 1.2999000000000001e-07, - "cache_read_input_token_cost": 1.2999000000000001e-07, + "cache_creation_input_token_cost": 1.2999e-07, + "cache_read_input_token_cost": 1.2999e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15382,8 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { - "cache_creation_input_token_cost": 1.5000999999999998e-07, - "cache_read_input_token_cost": 1.5000999999999998e-07, + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", From b460254428f5c0ff87a1d9e9b5d30e87536cd1d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:06:52 -0700 Subject: [PATCH 057/156] fix(interactions): price the settled background body against its own deployment The poll fetches the terminal interaction through its own client call, which stamps a response cost computed by a throwaway logging object holding none of the original request's deployment context: no model_info, no router model_id, no deployment litellm_params. Carrying that cost into the settlement event billed custom-priced deployments at the wrong rate, and it also satisfied the "already calculated" shortcut in _response_cost_calculator's caller, so the settlement never repriced and never built a cost breakdown. The zeros stamped by the usage-less create survived into the spend log row and the OTEL span. Dropping the imported cost before re-emitting makes the settlement price the settled body itself, against the deployment that served the create. --- litellm/litellm_core_utils/litellm_logging.py | 13 ++++++ .../test_litellm_logging.py | 43 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 275803c5aed..853df91cb89 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2202,7 +2202,20 @@ class Logging(LiteLLMLoggingBaseClass): ``in_progress`` response (no usage, so no cost was tracked); clearing the dedup flags lets the completed result flow through cost calculation and spend tracking exactly once, spanning create to completion. + + The poll fetched this body through its own client call, which priced it + against a throwaway logging object holding none of this request's + deployment context: no ``model_info``, no router ``model_id``, no + deployment ``litellm_params``. Keeping that price would bill a + custom-priced deployment at the wrong rate, and it would also satisfy + the "already calculated" shortcut and skip repricing here, leaving the + cost breakdown at the zeros the usage-less create stamped and writing + those zeros to the spend log. Dropping it makes this event price the + settled body itself, against the deployment that served the create. """ + settled_hidden_params: Final = getattr(result, "_hidden_params", None) + if isinstance(settled_hidden_params, dict): + settled_hidden_params.pop("response_cost", None) self._reset_success_emission_dedupe() await self.async_success_handler(result=result) 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 7c90d31261b..a7db8ae7a28 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4644,6 +4644,49 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 +@pytest.mark.asyncio +async def test_background_interaction_completion_prices_the_settled_body_itself(): + """ + The poll fetches the settled body through its own client call, which + prices it against a throwaway logging object holding none of this + request's deployment context. Adopting that price would bill a + custom-priced deployment at the wrong rate, and it would also satisfy the + "already calculated" shortcut and skip repricing, leaving the breakdown at + the zeros the usage-less create stamped and writing those to the spend log. + """ + 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(), + ) + + completed = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + completed._hidden_params = {"response_cost": 99.0} + + await logging_obj.async_log_background_interaction_completion(result=completed) + + response_cost = logging_obj.model_call_details["response_cost"] + assert response_cost != 99.0 + assert response_cost > 0 + + cost_breakdown = logging_obj.model_call_details["standard_logging_object"]["cost_breakdown"] + assert cost_breakdown["total_cost"] == response_cost + assert cost_breakdown["input_cost"] > 0 + assert cost_breakdown["output_cost"] > 0 + + @pytest.mark.asyncio async def test_background_interaction_completion_lets_otel_emit_the_cost_span(): """ From 521dc973c302e7672ed8cbf1d8607a7cb1ba72a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:15:14 -0700 Subject: [PATCH 058/156] fix(interactions): keep alerting when an interaction that produced output has no usage Silencing the cost-tracking alert for every usage-less interaction response went one status too far. An interaction that stopped at failed, cancelled, incomplete or budget_exceeded genuinely has nothing to charge for, so alerting on it is noise. completed and requires_action are different: both mean the model produced output, so a usage block is always expected, and one arriving without it means the charge for real work was lost. That is precisely the case failed_tracking_alert exists to surface, and swallowing it would let an operator's interactions bill nothing with no signal that anything was wrong. The status knowledge lives next to the other status predicates rather than in the proxy callback. The reservation is still released on both paths, since suppressing the alert was never what freed it. --- .../interactions/background_cost_polling.py | 17 +++++++ .../proxy/hooks/proxy_track_cost_callback.py | 5 +- .../hooks/test_proxy_track_cost_callback.py | 50 +++++++++++++++++-- 3 files changed, 67 insertions(+), 5 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 5d7dca80b73..279d2d6b32e 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -49,6 +49,8 @@ if TYPE_CHECKING: _TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"}) +_STATUSES_THAT_PRODUCED_OUTPUT = frozenset({"completed", "requires_action"}) + @dataclass(frozen=True, slots=True) class BackgroundInteractionPollContext: @@ -195,6 +197,21 @@ def is_pollable_background_interaction(response: InteractionsAPIResponse) -> boo return response.status == "in_progress" and bool(response.id) +def missing_usage_is_expected(response: InteractionsAPIResponse) -> bool: + """ + Whether a response arriving with no usage block is a normal outcome rather + than lost billing data. An interaction that is still running, or that + stopped at ``failed``, ``cancelled``, ``incomplete`` or ``budget_exceeded``, + has nothing to charge for and should not raise a cost-tracking alarm. + + ``completed`` and ``requires_action`` both mean the model produced output, + so a usage block is always expected with them. If one arrives without it + the charge for real work has been lost, which is precisely what the + proxy's cost-tracking alert exists to surface. + """ + return response.status not in _STATUSES_THAT_PRODUCED_OUTPUT + + @dataclass(frozen=True, slots=True) class _ActiveBackgroundPoll: task: "asyncio.Task[None]" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 70677521d6d..6abfca1d3a0 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -480,9 +480,12 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: def _is_unbilled_interaction_response(completion_response: object) -> bool: + from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse - return isinstance(completion_response, InteractionsAPIResponse) and completion_response.usage is None + if not isinstance(completion_response, InteractionsAPIResponse): + return False + return completion_response.usage is None and missing_usage_is_expected(completion_response) def _is_unbilled_in_progress_interaction(completion_response: object) -> bool: 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 fd79f9ea3b3..d840611fc7b 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 @@ -862,7 +862,7 @@ async def test_track_cost_callback_releases_reservation_for_in_progress_interact @pytest.mark.asyncio @pytest.mark.parametrize( "status", - ["completed", "failed", "cancelled", "incomplete", "requires_action", "budget_exceeded"], + ["failed", "cancelled", "incomplete", "budget_exceeded"], ) async def test_track_cost_callback_releases_reservation_for_unpollable_interaction(status): """ @@ -872,9 +872,10 @@ async def test_track_cost_callback_releases_reservation_for_unpollable_interacti added to the key, user, team and org spend counters and starts refusing traffic against budget that was never actually spent. - A no-usage terminal create is also not a cost-tracking failure, so the - callback must not fire ``failed_tracking_alert``: doing so would flood - operators with false alerts and mask real cost-tracking failures. + None of these statuses produced output, so their missing usage is a normal + outcome rather than a cost-tracking failure, and the callback must not fire + ``failed_tracking_alert``: doing so would flood operators with false alerts + and mask real cost-tracking failures. """ from litellm.types.interactions import InteractionsAPIResponse @@ -904,6 +905,47 @@ async def test_track_cost_callback_releases_reservation_for_unpollable_interacti mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "requires_action"]) +async def test_track_cost_callback_alerts_when_an_interaction_that_produced_output_has_no_usage(status): + """ + ``completed`` and ``requires_action`` both mean the model produced output, + so a usage block is always expected with them. One arriving without it + means the charge for real work was lost, which is exactly what the + cost-tracking alert is for: silencing it here would let an operator's + interactions bill nothing with no signal that anything went wrong. + + The reservation still has to be released, since suppressing the alert was + never what freed it. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + usageless_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + 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=usageless_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_called_once() + + @pytest.mark.asyncio async def test_track_cost_callback_releases_reservation_for_interaction_without_an_id(): """ From 8a1fe281fdbc0d63c272febd634741cc14005295 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:29:35 -0700 Subject: [PATCH 059/156] fix(responses-adapter): skip null reasoning summary text and keep thinking-only assistant turns --- .../transformation.py | 4 +++- .../responses_adapters/transformation.py | 4 ++-- ...responses_transformation_transformation.py | 22 +++++++++++++++++++ .../test_responses_adapters_transformation.py | 14 ++++++++++++ 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index cf227cfa7d1..b94e91b3034 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -428,12 +428,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if role == "assistant": input_items.extend(_reasoning_input_items(msg)) input_items.append( - { + { # mutable-ok: API message payload "type": "message", "role": role, "content": self._convert_content_to_responses_format(content, cast(str, role)), } ) + elif role == "assistant": + input_items.extend(_reasoning_input_items(msg)) return input_items, instructions diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 23a9a60d810..0b38123e787 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -110,8 +110,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def _summary_part_text(part: object) -> str: if isinstance(part, Mapping): mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json - return str(mapping.get("text", "")) - return str(getattr(part, "text", "")) + return str(mapping.get("text") or "") + return str(getattr(part, "text", "") or "") @classmethod def _thinking_blocks_from_reasoning_item( diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 124dc67b4fd..1fb74b2b7bf 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3824,6 +3824,28 @@ def test_assistant_thinking_blocks_become_a_reasoning_input_item(): assert "id" not in reasoning_item +def test_thinking_only_assistant_turn_still_sends_its_reasoning(): + """An assistant turn can be pure reasoning, with no visible text and no tool call.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "sig1"} + ], + }, + {"role": "user", "content": "Why?"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_items = [item for item in input_items if item.get("type") == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["summary"] == [{"type": "summary_text", "text": "August in Denver is dry."}] + + def test_stored_reasoning_items_win_over_thinking_blocks(): """A minted reasoning id beats a re-derived one, so the two must not both be sent.""" handler = LiteLLMResponsesTransformationHandler() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 790bcd269e0..8225e7cff39 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -1249,6 +1249,20 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [] + def test_null_summary_text_skipped_rather_than_stringified(self): + """A summary part whose text is null must not reach the client as the word "None".""" + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_null_1", + "summary": [{"type": "summary_text", "text": None}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [] + def test_reasoning_item_id_never_becomes_a_thinking_signature(self): """Only Anthropic can sign a thinking block, so a stand-in signature is never invented.""" reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") From 0a869b72268685b27d5a0be4742221dc8f24d6db Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:42:52 -0700 Subject: [PATCH 060/156] fix(databricks): derive cache rates from published cache DBU Cache rates were derived as ratios of the dollar input rate (1.25x write, 0.1x read) while input and output derive from the published DBU table times $0.070. Databricks publishes cache write and cache read DBU per model, and those are not exact multiples of the input DBU, so the two rules disagreed by up to 0.1 percent. Rewrites 43 cache literals across 31 entries to published_cache_DBU x $0.070. Skips databricks-gemini-2-5-pro and databricks-gemini-2-5-flash, whose input and output rates predate the current table by a 1.25x increase; their cache rates stay tied to their own input rate so each entry remains internally consistent. Replaces the ratio assertions with a test pinning the absolute published DBU figures for the five new models, and adds a test pinning the older-vintage exception. Corrects the metadata note on the five new entries, which claimed the reference-only *_dbu_cost_per_token fields drive cost calculation. --- ...odel_prices_and_context_window_backup.json | 98 +++++++++---------- model_prices_and_context_window.json | 98 +++++++++---------- .../test_databricks_cost_calculator.py | 81 ++++++++++----- 3 files changed, 153 insertions(+), 124 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8c80714c089..8c8b909493a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14568,8 +14568,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14590,8 +14590,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.2500075e-05, - "cache_read_input_token_cost": 1.000006e-06, + "cache_creation_input_token_cost": 1.250004e-05, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.000006e-05, "input_dbu_cost_per_token": 0.000142858, "litellm_provider": "databricks", @@ -14599,7 +14599,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 5.000002e-05, @@ -14614,12 +14614,12 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true, + "supports_vision": false, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { - "cache_creation_input_token_cost": 1.250025e-06, - "cache_read_input_token_cost": 1.00002e-07, + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14640,8 +14640,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { - "cache_creation_input_token_cost": 1.8750025e-05, - "cache_read_input_token_cost": 1.500002e-06, + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14662,8 +14662,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { - "cache_creation_input_token_cost": 1.8750025e-05, - "cache_read_input_token_cost": 1.500002e-06, + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14684,8 +14684,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14707,8 +14707,8 @@ "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14729,8 +14729,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-7": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14738,7 +14738,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14755,8 +14755,8 @@ "supports_vision": true }, "databricks/databricks-claude-opus-4-8": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14764,7 +14764,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14782,8 +14782,8 @@ "supports_vision": true }, "databricks/databricks-claude-opus-5": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14791,7 +14791,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14809,8 +14809,8 @@ "supports_vision": true }, "databricks/databricks-claude-sonnet-4": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14831,8 +14831,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14853,8 +14853,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14875,8 +14875,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14897,8 +14897,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14906,7 +14906,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", "output_cost_per_token": 1.500002e-05, @@ -14965,7 +14965,7 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, - "cache_read_input_token_cost": 3.1248e-08, + "cache_read_input_token_cost": 3.122e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14985,7 +14985,7 @@ }, "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15005,7 +15005,7 @@ }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, - "cache_read_input_token_cost": 6.2503e-08, + "cache_read_input_token_cost": 6.251e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -15025,7 +15025,7 @@ }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15062,7 +15062,7 @@ }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15080,7 +15080,7 @@ }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15098,7 +15098,7 @@ }, "databricks/databricks-gpt-5-1-codex-max": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15116,7 +15116,7 @@ }, "databricks/databricks-gpt-5-1-codex-mini": { "cache_creation_input_token_cost": 2.4997e-07, - "cache_read_input_token_cost": 2.4997e-08, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15188,7 +15188,7 @@ }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15206,7 +15206,7 @@ }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, - "cache_read_input_token_cost": 7.4998e-08, + "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15224,7 +15224,7 @@ }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, - "cache_read_input_token_cost": 1.9999e-08, + "cache_read_input_token_cost": 2.002e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15242,7 +15242,7 @@ }, "databricks/databricks-gpt-5-mini": { "cache_creation_input_token_cost": 2.4997e-07, - "cache_read_input_token_cost": 2.4997e-08, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15260,7 +15260,7 @@ }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, - "cache_read_input_token_cost": 4.998e-09, + "cache_read_input_token_cost": 4.97e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8c80714c089..8c8b909493a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14568,8 +14568,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14590,8 +14590,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.2500075e-05, - "cache_read_input_token_cost": 1.000006e-06, + "cache_creation_input_token_cost": 1.250004e-05, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.000006e-05, "input_dbu_cost_per_token": 0.000142858, "litellm_provider": "databricks", @@ -14599,7 +14599,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 5.000002e-05, @@ -14614,12 +14614,12 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true, + "supports_vision": false, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { - "cache_creation_input_token_cost": 1.250025e-06, - "cache_read_input_token_cost": 1.00002e-07, + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14640,8 +14640,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { - "cache_creation_input_token_cost": 1.8750025e-05, - "cache_read_input_token_cost": 1.500002e-06, + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14662,8 +14662,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { - "cache_creation_input_token_cost": 1.8750025e-05, - "cache_read_input_token_cost": 1.500002e-06, + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14684,8 +14684,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14707,8 +14707,8 @@ "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14729,8 +14729,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-7": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14738,7 +14738,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14755,8 +14755,8 @@ "supports_vision": true }, "databricks/databricks-claude-opus-4-8": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14764,7 +14764,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14782,8 +14782,8 @@ "supports_vision": true }, "databricks/databricks-claude-opus-5": { - "cache_creation_input_token_cost": 6.2500375e-06, - "cache_read_input_token_cost": 5.00003e-07, + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14791,7 +14791,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", "output_cost_per_token": 2.500001e-05, @@ -14809,8 +14809,8 @@ "supports_vision": true }, "databricks/databricks-claude-sonnet-4": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14831,8 +14831,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14853,8 +14853,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14875,8 +14875,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14897,8 +14897,8 @@ "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.7499875e-06, - "cache_read_input_token_cost": 2.99999e-07, + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14906,7 +14906,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation. Introductory launch rates of 28.571 input / 142.857 output DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." }, "mode": "chat", "output_cost_per_token": 1.500002e-05, @@ -14965,7 +14965,7 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, - "cache_read_input_token_cost": 3.1248e-08, + "cache_read_input_token_cost": 3.122e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14985,7 +14985,7 @@ }, "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15005,7 +15005,7 @@ }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, - "cache_read_input_token_cost": 6.2503e-08, + "cache_read_input_token_cost": 6.251e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -15025,7 +15025,7 @@ }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15062,7 +15062,7 @@ }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15080,7 +15080,7 @@ }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15098,7 +15098,7 @@ }, "databricks/databricks-gpt-5-1-codex-max": { "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.24999e-07, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -15116,7 +15116,7 @@ }, "databricks/databricks-gpt-5-1-codex-mini": { "cache_creation_input_token_cost": 2.4997e-07, - "cache_read_input_token_cost": 2.4997e-08, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15188,7 +15188,7 @@ }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, - "cache_read_input_token_cost": 2.49998e-07, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -15206,7 +15206,7 @@ }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, - "cache_read_input_token_cost": 7.4998e-08, + "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15224,7 +15224,7 @@ }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, - "cache_read_input_token_cost": 1.9999e-08, + "cache_read_input_token_cost": 2.002e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15242,7 +15242,7 @@ }, "databricks/databricks-gpt-5-mini": { "cache_creation_input_token_cost": 2.4997e-07, - "cache_read_input_token_cost": 2.4997e-08, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15260,7 +15260,7 @@ }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, - "cache_read_input_token_cost": 4.998e-09, + "cache_read_input_token_cost": 4.97e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 102a861fe1a..0f1652b1b2a 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from decimal import Decimal from pathlib import Path from typing import Final @@ -20,20 +20,30 @@ NEW_MODELS: Final = ( "databricks/databricks-claude-fable-5", ) - -@pytest.fixture -def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - yield - litellm.get_model_info.cache_clear() +DOLLARS_PER_DBU: Final = Decimal("0.070") +PRICE_FIELDS: Final = ( + "input_cost_per_token", + "output_cost_per_token", + "cache_creation_input_token_cost", + "cache_read_input_token_cost", +) +PUBLISHED_DBU_PER_MILLION: Final = { + "databricks/databricks-claude-opus-4-7": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-8": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-5": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"), +} def _model_info(model: str) -> ModelInfo: return litellm.get_model_info(model=model, custom_llm_provider="databricks") +def _dollars_per_token(dbu_per_million: str) -> float: + return float(Decimal(dbu_per_million) * DOLLARS_PER_DBU / Decimal(10) ** 6) + + @pytest.mark.parametrize( "model", [ @@ -84,23 +94,22 @@ def test_legacy_endpoint_names_still_resolve(local_model_cost_map: None) -> None assert completion_cost == pytest.approx(100 * info["output_cost_per_token"]) -@pytest.mark.parametrize( - "model", - [ - "databricks/databricks-claude-opus-4-7", - "databricks/databricks-claude-opus-4-8", - "databricks/databricks-claude-opus-5", - "databricks/databricks-claude-sonnet-5", - "databricks/databricks-claude-fable-5", - ], -) +@pytest.mark.parametrize("model", NEW_MODELS) +def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + + for field, dbu_per_million in zip(PRICE_FIELDS, PUBLISHED_DBU_PER_MILLION[model]): + assert info[field] == _dollars_per_token(dbu_per_million), field + + +@pytest.mark.parametrize("model", NEW_MODELS) def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) assert info["input_cost_per_token"] > 0 assert info["output_cost_per_token"] > 0 - assert info["cache_creation_input_token_cost"] == pytest.approx(1.25 * info["input_cost_per_token"], rel=1e-4) - assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"], rel=1e-4) + assert info["cache_creation_input_token_cost"] > info["input_cost_per_token"] + assert info["cache_read_input_token_cost"] < info["input_cost_per_token"] assert info["supports_prompt_caching"] is True @@ -130,7 +139,10 @@ def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( prompt_cost, _ = cost_per_token(model=model, usage=usage) + assert info["cache_read_input_token_cost"] == info["input_cost_per_token"] + assert info["cache_creation_input_token_cost"] == info["input_cost_per_token"] assert prompt_cost == pytest.approx(10000 * info["input_cost_per_token"]) + assert prompt_cost > 8000 * info["input_cost_per_token"] @pytest.mark.parametrize("model", NEW_MODELS) @@ -138,14 +150,31 @@ def test_backup_price_map_matches_main(model: str) -> None: main_cost: Final = json.loads(MAIN_PRICES.read_text()) backup_cost: Final = json.loads(BACKUP_PRICES.read_text()) - assert backup_cost.get(model) == main_cost.get(model) + assert model in main_cost + assert model in backup_cost + assert backup_cost[model] == main_cost[model] def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") - assert sonnet_5["input_cost_per_token"] == pytest.approx(sonnet_4_6["input_cost_per_token"]) - assert sonnet_5["output_cost_per_token"] == pytest.approx(sonnet_4_6["output_cost_per_token"]) - assert sonnet_5["cache_creation_input_token_cost"] == pytest.approx(sonnet_4_6["cache_creation_input_token_cost"]) - assert sonnet_5["cache_read_input_token_cost"] == pytest.approx(sonnet_4_6["cache_read_input_token_cost"]) + for field in PRICE_FIELDS: + assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field + + +@pytest.mark.parametrize( + "model", + [ + "databricks/databricks-gemini-2-5-pro", + "databricks/databricks-gemini-2-5-flash", + ], +) +def test_entries_priced_at_an_older_vintage_keep_cache_rates_tied_to_their_own_input( + local_model_cost_map: None, + model: str, +) -> None: + info: Final = _model_info(model) + + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) + assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) From 112224a7fe51b78587d7a2e09d7b9ce8527ec448 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:30:10 -0700 Subject: [PATCH 061/156] fix(databricks): enable vision on fable-5 and pin every cache rate to published DBU databricks-claude-fable-5 was the only fable-5 entry in the registry declaring supports_vision false, and the only one of the five new entries to do so. Only the five new models were pinned against the published DBU table, so the 26 cache literals added to pre-existing entries were checked by nothing independent. Extend the table to all 33 entries carrying cache rates and assert both cache fields against it for the 31 that take the published rates, leaving the two older-vintage gemini-2-5 entries to their existing guard. Also widen the cache-declaration guard to both cache fields, and replace the single-model equals-input assertion with one that covers all 14 entries publishing no cache rates. --- ...odel_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../test_databricks_cost_calculator.py | 80 +++++++++++++++---- 3 files changed, 68 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8c8b909493a..1dc84f06aa7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14614,7 +14614,7 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": false, + "supports_vision": true, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8c8b909493a..1dc84f06aa7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14614,7 +14614,7 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": false, + "supports_vision": true, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 0f1652b1b2a..81517e23b53 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -28,12 +28,45 @@ PRICE_FIELDS: Final = ( "cache_read_input_token_cost", ) PUBLISHED_DBU_PER_MILLION: Final = { - "databricks/databricks-claude-opus-4-7": ("71.429", "357.143", "89.286", "7.143"), - "databricks/databricks-claude-opus-4-8": ("71.429", "357.143", "89.286", "7.143"), - "databricks/databricks-claude-opus-5": ("71.429", "357.143", "89.286", "7.143"), - "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), "databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"), + "databricks/databricks-claude-opus-5": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-8": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-7": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-6": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-5": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-1": ("214.286", "1071.429", "267.857", "21.429"), + "databricks/databricks-claude-opus-4": ("214.286", "1071.429", "267.857", "21.429"), + "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-6": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-1": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-3-7-sonnet": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-haiku-4-5": ("14.286", "71.429", "17.857", "1.429"), + "databricks/databricks-gpt-5": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1-codex-max": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1-codex-mini": ("3.571", "28.571", "3.571", "0.357"), + "databricks/databricks-gpt-5-mini": ("3.571", "28.571", "3.571", "0.357"), + "databricks/databricks-gpt-5-nano": ("0.714", "5.714", "0.714", "0.071"), + "databricks/databricks-gpt-5-2": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-2-codex": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-3-codex": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-4": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gpt-5-4-mini": ("10.714", "64.286", "10.714", "1.071"), + "databricks/databricks-gpt-5-4-nano": ("2.857", "17.857", "2.857", "0.286"), + "databricks/databricks-gemini-3-1-pro": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gemini-3-pro": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gemini-3-flash": ("8.929", "53.571", "8.929", "0.893"), + "databricks/databricks-gemini-3-1-flash-lite": ("4.464", "26.786", "4.464", "0.446"), + "databricks/databricks-gemini-2-5-pro": ("22.321", "178.571", "22.321", "2.232"), + "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), } +OLDER_VINTAGE_MODELS: Final = ( + "databricks/databricks-gemini-2-5-pro", + "databricks/databricks-gemini-2-5-flash", +) +CACHE_FIELDS: Final = ("cache_creation_input_token_cost", "cache_read_input_token_cost") def _model_info(model: str) -> ModelInfo: @@ -102,6 +135,15 @@ def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, mod assert info[field] == _dollars_per_token(dbu_per_million), field +@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(OLDER_VINTAGE_MODELS))) +def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] + + for field, dbu_per_million in zip(CACHE_FIELDS, cache_dbu_per_million): + assert info[field] == _dollars_per_token(dbu_per_million), field + + @pytest.mark.parametrize("model", NEW_MODELS) def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) @@ -119,7 +161,7 @@ def test_every_priced_databricks_model_declares_cache_rates(local_model_cost_map for model, info in litellm.model_cost.items() if model.startswith("databricks/") and info.get("input_cost_per_token") is not None - and info.get("cache_read_input_token_cost") is None + and any(info.get(field) is None for field in CACHE_FIELDS) ] assert undeclared == [] @@ -139,12 +181,28 @@ def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( prompt_cost, _ = cost_per_token(model=model, usage=usage) - assert info["cache_read_input_token_cost"] == info["input_cost_per_token"] - assert info["cache_creation_input_token_cost"] == info["input_cost_per_token"] assert prompt_cost == pytest.approx(10000 * info["input_cost_per_token"]) assert prompt_cost > 8000 * info["input_cost_per_token"] +def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_rate( + local_model_cost_map: None, +) -> None: + without_published_rates: Final = [ + model + for model, info in litellm.model_cost.items() + if model.startswith("databricks/") + and info.get("input_cost_per_token") + and model not in PUBLISHED_DBU_PER_MILLION + ] + + assert len(without_published_rates) == 14 + for model in without_published_rates: + info = _model_info(model) + for field in CACHE_FIELDS: + assert info[field] == pytest.approx(info["input_cost_per_token"]), (model, field) + + @pytest.mark.parametrize("model", NEW_MODELS) def test_backup_price_map_matches_main(model: str) -> None: main_cost: Final = json.loads(MAIN_PRICES.read_text()) @@ -163,13 +221,7 @@ def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: No assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field -@pytest.mark.parametrize( - "model", - [ - "databricks/databricks-gemini-2-5-pro", - "databricks/databricks-gemini-2-5-flash", - ], -) +@pytest.mark.parametrize("model", OLDER_VINTAGE_MODELS) def test_entries_priced_at_an_older_vintage_keep_cache_rates_tied_to_their_own_input( local_model_cost_map: None, model: str, From 528d358c0540970c8bdb3802f6c480b0c3f24a9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:30:18 -0700 Subject: [PATCH 062/156] fix: leave a bounded insert, a bounded writable CTE and an uncalled routine alone A parenthesised VALUES list ended the search for an insert's row source only when no group followed it, so a RETURNING or an ON CONFLICT DO UPDATE carrying a subquery was read as the rows the insert copies. A writable CTE bounded by its own VALUES list was handed the query the statement ends with for the same reason: the WITH branch read the whole statement rather than the part holding the insert. A CREATE FUNCTION or CREATE PROCEDURE body was scanned as if it ran at boot, but defining a routine only stores it. The body is now read when the same migration names the routine somewhere else, so a migration that defines a backfill and then runs it is still caught, and one whose name needed quoting is read either way since quoting is blanked at the call sites too. main() had no test, so neither its exit codes nor the branch the CI gate reads were pinned; a mutant returning 0 on a violation passed the whole suite. Its four outcomes now have tests, along with both directions of each fix above. --- .../check_migrations_no_data_rewrites.py | 120 +++++++++-- .../test_check_migrations_no_data_rewrites.py | 199 +++++++++++++++++- 2 files changed, 299 insertions(+), 20 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 17bc90bd0d6..fa966a214b1 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -8,10 +8,12 @@ plus a doubled heap that plain autovacuum will not give back. What is banned is the row-rewriting DML behind that, not everything whose cost scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is -not binary coercible, and a volatile `DEFAULT` on a new column all read the whole -table and all pass. That is deliberate: a rule wide enough to reach them fires on -most ordinary migrations, and a marker everyone adds by reflex stops carrying -information. The outage this was written for was a backfill. +not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE ... AS +SELECT` or `SELECT ... INTO` filling a new table from an existing one, the rename +that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED VIEW` +all read the whole table and all pass. That is deliberate: a rule wide enough to +reach them fires on most ordinary migrations, and a marker everyone adds by reflex +stops carrying information. The outage this was written for was a backfill. Flagged, per statement, by its leading keyword: @@ -21,10 +23,15 @@ Flagged, per statement, by its leading keyword: INSERT only when its rows come from a query rather than a literal `VALUES` list. The query counts wherever it sits, since Postgres takes it parenthesised, and `TABLE t` is one as much as a `SELECT` is. An - insert bounded by a leading `VALUES` passes, scalar subqueries in that - list included, while a `VALUES` reached through a subquery or joined - to a query by a set operation bounds nothing - WITH a CTE-led statement containing any of the above + insert bounded by a `VALUES` list passes, written bare or in + parentheses, and so do the scalar subqueries in that list and the + `RETURNING` and `ON CONFLICT` clauses written after it, none of which + supply the rows. A `VALUES` reached through a subquery or joined to a + query by a set operation bounds nothing + WITH a CTE-led statement containing any of the above. An `INSERT` is read + against the part of the statement holding it, so a writable CTE + bounded by its own `VALUES` list is not handed the query the statement + ends with as the rows it copies Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a statement's leading keyword, so they pass. @@ -37,7 +44,12 @@ told not to run at boot, and a marker is a cheap answer if one ever does. Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise -hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the +hide. A `CREATE FUNCTION` or `CREATE PROCEDURE` body is the exception, because +defining a routine only stores it: that body is read when the same migration names +the routine somewhere else, which is what defining a backfill and then running it +looks like, and left alone when nothing calls it. A routine whose name needed +quoting is read either way, since quoting is blanked at the call sites too and a +call written there could never be found. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the same to Postgres whether it is spelled out or handed over as a string, and so is a literal parked in a variable some `EXECUTE` in the same body then runs by name, however it got there: an assignment with `:=`, the bare `=` PL/pgSQL takes as the @@ -110,6 +122,11 @@ LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTA WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?!:=])=(?![=>])") PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) +DEFINES_A_ROUTINE = re.compile( + r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE +) +QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)" +ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})") REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -152,6 +169,8 @@ NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) BIND_VALUES = re.compile(r"\bUSING\b", re.IGNORECASE) +WRITES_ROWS = re.compile(r"\bINSERT\b", re.IGNORECASE) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -370,13 +389,30 @@ def offending_keyword(statement: str) -> str | None: if nested is not None: return f"WITH ... {nested}" if contains(statement, "INSERT"): - source = row_source_keyword(statement) + source = insert_row_source(statement) if source is not None: return f"WITH ... INSERT ... {source}" return None +def insert_row_source(statement: str) -> str | None: + """Which keyword supplies the rows to an `INSERT` written somewhere inside a `WITH` + statement. Only the parts that hold that insert are read, because a writable CTE sits + beside the query the statement ends with and reading the whole thing hands the insert + the outer `SELECT` as its row source: `WITH c AS (INSERT ... VALUES (1) RETURNING "x") + SELECT * FROM c` adds one literal row and copies nothing. A CTE keeps its insert in a + parenthesised group, and the statement's own insert, if it is the one writing, runs from + the keyword to the end, found in the text outside every parenthesis so a group's insert + is not counted twice.""" + inserts = [group for group in parenthesised_groups(statement) if contains(group, "INSERT")] + written = WRITES_ROWS.search(strip_parens(statement)) + if written is not None: + inserts.append(statement[written.start() :]) + sources = (row_source_keyword(insert) for insert in inserts) + return next((source for source in sources if source is not None), None) + + def row_source_keyword(statement: str) -> str | None: """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list does. A query outside every parenthesis is the row source outright. Failing that, a @@ -387,9 +423,12 @@ def row_source_keyword(statement: str) -> str | None: a rewrite. Failing all three, the rows come from a parenthesised group, which Postgres accepts and which reading only the unparenthesised text would let through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table. Each group at that level is - read on its own terms and the first to name a row source is the answer, since the ones - around it are the column list, the conflict target and the rest of the clauses an insert - is allowed to carry, and any of those can be the last group written.""" + read on its own terms until one of them supplies the rows, since the ones before it are + the column list and the ones after it are the conflict target and the rest of the clauses + an insert is allowed to carry. A wrapped `VALUES` list is the row source as much as a + wrapped query is, so it ends the search rather than being skipped over: reading past it + reaches a `RETURNING (SELECT ...)` or a `DO UPDATE SET "a" = (SELECT ...)` written after + it and calls that scalar subquery the rows the insert copies.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: @@ -402,8 +441,13 @@ def row_source_keyword(statement: str) -> str | None: groups = list(parenthesised_groups(statement)) if not groups: return row_source_in(statement) - sources = (row_source_keyword(group) for group in groups) - return next((source for source in sources if source is not None), None) + for group in groups: + if contains(strip_parens(group), "VALUES"): + return None + source = row_source_keyword(group) + if source is not None: + return source + return None def set_operation_terms(statement: str, outer: str) -> Iterator[str]: @@ -594,10 +638,54 @@ def scan_region( continue yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) - for start, end in bodies: + for body in bodies: + if not runs_when_applied(masked, region, bodies, body): + continue + start, end = body yield from scan_region(document, region[start:end], migration, markers, offset + start) +def runs_when_applied( + masked: str, region: str, bodies: tuple[tuple[int, int], ...], body: tuple[int, int] +) -> bool: + """Whether a dollar-quoted body runs while the migration is being applied. A `DO` block runs + where it is written, and so does every other use of this quoting. A `CREATE FUNCTION` or a + `CREATE PROCEDURE` only stores its body, which runs when something calls the routine, so a + definition nothing calls rewrites no rows at boot and reporting it names a line that never + executes. Skipping every definition instead would let a migration define a backfill and then + run it unseen, which is the shape this check exists to catch, so the body is read whenever + the same migration names the routine anywhere outside the definition. The definition is + found in the masked text, where one written inside a comment has already been blanked, and + the name is read from the region at those same offsets, since masking blanks a quoted + identifier in place. A name that needed those quotes is blanked at its call sites too and + so can never be found there, which would read as uncalled however the migration runs it, + and the body is read rather than trusted.""" + start, end = body + opens = masked.rfind(";", 0, start) + 1 + defined = DEFINES_A_ROUTINE.search(masked, opens, start) + if defined is None: + return True + named = ROUTINE_NAME.match(region, defined.end(), start) + if named is None or named.group(1).startswith('"'): + return True + return contains(outside_definition(masked, region, bodies, opens, end), re.escape(named.group(1))) + + +def outside_definition( + masked: str, region: str, bodies: tuple[tuple[int, int], ...], opens: int, closes: int +) -> str: + """The migration's text with one routine definition blanked out and every dollar-quoted body + put back. Masking blanks the bodies alike, and a `DO` block is the ordinary way a migration + runs a routine it has just defined, so a call written inside one has to stay readable. The + definition is blanked after they are restored, which takes its own body with it, so a + routine that names itself recursively does not thereby count as called.""" + text = list(masked) + for start, end in bodies: + text[start:end] = region[start:end] + text[opens:closes] = blank(region[opens:closes]) + return "".join(text) + + def clauses(statement: str, start: int) -> Iterator[tuple[str, int]]: """The statements written inside one semicolon-delimited run, each with where it begins. A `FOR ... LOOP` header takes no semicolon of its own, so the first statement of the loop body diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index b78b34e52c7..6a55a7a7c8d 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -1,10 +1,9 @@ """Tests for tests/code_coverage_tests/check_migrations_no_data_rewrites.py. The checker reads migration.sql as SQL rather than as text, so the cases that matter -are the ones a grep would get wrong: `ON DELETE CASCADE` in a foreign key (60-odd -occurrences in the shipped migrations), an `UPDATE` inside a string literal or a -comment, and an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for -conditional DDL. +are the ones a grep would get wrong: the referential actions in a foreign key, of which +the shipped migrations carry 60, an `UPDATE` inside a string literal or a comment, and +an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for conditional DDL. """ import importlib.util @@ -218,6 +217,25 @@ class TestInsert: def test_a_table_named_in_the_insert_target_does_not_flag_it(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "audit table" ("id") VALUES (1);') == () + def test_a_returning_subquery_after_a_wrapped_values_list_is_not_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_update_after_a_wrapped_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES (1))' + ' ON CONFLICT ("id") DO UPDATE SET "id" = (SELECT max("id") FROM "Bar");' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_wrapped_values_list_of_several_rows_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1), (2)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_the_row_source_names_its_own_keyword_not_a_later_subquery(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (TABLE "Bar") RETURNING (SELECT count(*) FROM "Baz");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -248,6 +266,32 @@ class TestCommonTableExpressions: sql = 'WITH latest AS (SELECT max("id") AS "id" FROM "Bar")\nINSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT "id"::text FROM latest));' assert _keywords(tmp_path, sql) == () + def test_a_writable_cte_bounded_by_values_passes(self, tmp_path): + sql = 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id") SELECT * FROM added;' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_copying_a_query_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_bounded_writable_cte_does_not_hide_a_copying_one_beside_it(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id"),' + ' copied AS (INSERT INTO "Baz" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added, copied;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_writable_cte_wrapping_its_row_source_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + class TestDollarQuotedBlocks: def test_update_inside_do_block_is_flagged(self, tmp_path): @@ -326,6 +370,94 @@ class TestDollarQuotedBlocks: assert _scan(tmp_path, sql)[0].line == 7 +class TestStoredRoutines: + DEFINITION = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + PROCEDURE = ( + "CREATE OR REPLACE PROCEDURE sweep() AS $$\n" + "BEGIN\n" + ' DELETE FROM "Foo";\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + + def test_a_function_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION) == () + + def test_a_procedure_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE) == () + + def test_a_function_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION + "SELECT backfill();\n") == ("UPDATE",) + + def test_a_procedure_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE + "CALL sweep();\n") == ("DELETE",) + + def test_a_call_written_above_the_definition_still_counts(self, tmp_path): + assert _keywords(tmp_path, "SELECT backfill();\n" + self.DEFINITION) == ("UPDATE",) + + def test_a_call_from_inside_a_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_trigger_wiring_the_function_up_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TRIGGER t AFTER INSERT ON "Foo" EXECUTE FUNCTION backfill();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_definition_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION.replace("backfill()", "public.backfill()")) == () + + def test_a_schema_qualified_function_the_migration_calls_is_flagged(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", "public.backfill()") + "SELECT public.backfill();\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_the_name_written_only_in_a_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "-- backfill() is run by hand after the deploy\n" + assert _keywords(tmp_path, sql) == () + + def test_a_recursive_call_does_not_count_as_the_migration_calling_it(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill(n int) RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " PERFORM backfill(n - 1);\n" + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_routine_name_is_read_rather_than_trusted(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", '"back fill"()') + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_do_block_is_not_a_routine_definition(self, tmp_path): + sql = 'DO $$ BEGIN UPDATE "Foo" SET "a" = 1; END; $$;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_definition_written_after_another_statement_is_still_recognised(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;\n' + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_in_a_routine_the_migration_calls(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1; -- data-migration-ok: single config row\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + "SELECT backfill();\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_called_routine_reports_the_line_inside_its_body(self, tmp_path): + assert _scan(tmp_path, self.DEFINITION + "SELECT backfill();\n")[0].line == 3 + + class TestLoopBodies: def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): sql = ( @@ -1309,3 +1441,62 @@ class TestGrandfathering: class TestShippedMigrations: def test_the_repo_is_clean(self): assert checker.main() == 0 + + +CLEAN = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;' +DIRTY = 'UPDATE "Foo" SET "a" = 1;' +FIXTURE = "20260101000000_fixture" + + +def _tree(monkeypatch, tmp_path: Path, sql: str, grandfathered: frozenset = frozenset()) -> None: + """Stand a migrations directory holding one fixture migration in for the repo's own. The + root moves with it, since a rendered violation names the migration relative to the root and + the two are read off the same checkout everywhere but here.""" + directory = tmp_path / "migrations" / FIXTURE + directory.mkdir(parents=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "migrations") + monkeypatch.setattr(checker, "GRANDFATHERED", grandfathered) + + +class TestExitCode: + def test_a_clean_tree_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN) + assert checker.main() == 0 + + def test_a_violation_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY) + assert checker.main() == 1 + + def test_a_stale_grandfather_alone_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + assert checker.main() == 1 + + def test_a_grandfathered_violation_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY, frozenset({FIXTURE})) + assert checker.main() == 0 + + def test_a_missing_migrations_directory_is_an_error(self, tmp_path, monkeypatch): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + assert checker.main() == 2 + + def test_the_failure_names_the_migration_the_line_and_the_keyword( + self, tmp_path, monkeypatch, capsys + ): + _tree(monkeypatch, tmp_path, DIRTY) + checker.main() + printed = capsys.readouterr().out + assert f"migrations/{FIXTURE}/migration.sql:1" in printed + assert "UPDATE rewrites existing rows at boot" in printed + assert checker.GUIDANCE in printed + + def test_a_stale_grandfather_is_named(self, tmp_path, monkeypatch, capsys): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + checker.main() + assert f"{FIXTURE}: listed in GRANDFATHERED" in capsys.readouterr().out + + def test_a_missing_directory_is_reported_on_stderr(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + checker.main() + assert "migrations directory not found" in capsys.readouterr().err From a41ac5c1392171e59302926d88c70346fb1e1ea3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:35:19 -0700 Subject: [PATCH 063/156] fix(interactions): poll queued background creates and drop the poll's deployment identity `queued` is the Interactions API's not-started-yet status. It was in neither the pollable set nor the terminal one, so a create returning it got no poll task, counted as a response with nothing to charge for, and released its budget reservation: billed nowhere, alerting nobody. Poll it alongside `in_progress`, and pin the union of the pollable and terminal sets against the generated spec enum so a status Google adds later fails CI rather than shipping another unbilled path. A give-up on a status in neither set now names the status and logs at error, instead of the warning that reads as an interaction merely still running. Also drop `model_id` and `litellm_model_name` from the settled body next to the foreign `response_cost` already dropped there. All three come from the poll's own throwaway client call, and left in place the two identity fields overwrite the create's real deployment in the payload every logging integration reads. Rewrites the callback's per-status test to assert the observable outcome (reservation held vs released) across all eight statuses rather than comparing the gate to the function it delegates to, and pins the shipped 5-10-20-40-60 poll backoff and its timeout cutoff. --- .../interactions/background_cost_polling.py | 40 ++++++-- litellm/litellm_core_utils/litellm_logging.py | 8 +- .../test_background_cost_polling.py | 95 +++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 76 +++++++++++---- 4 files changed, 189 insertions(+), 30 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 279d2d6b32e..51325354e7d 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -32,7 +32,7 @@ the interaction is billed exactly once no matter who settles first. import asyncio from collections.abc import Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, TypeAlias +from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_logger from litellm.constants import ( @@ -47,9 +47,13 @@ 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", "requires_action"}) +_TERMINAL_STATUSES: Final = frozenset( + {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"} +) -_STATUSES_THAT_PRODUCED_OUTPUT = frozenset({"completed", "requires_action"}) +_POLLABLE_STATUSES: Final = frozenset({"in_progress", "queued"}) + +_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action"}) @dataclass(frozen=True, slots=True) @@ -113,6 +117,7 @@ async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = _fetch_interaction, ) -> None: + last_seen_status: str | None = None for interval in _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, @@ -130,6 +135,7 @@ async def poll_and_log_background_interaction_cost( e, ) continue + last_seen_status = response.status if response.status not in _TERMINAL_STATUSES: continue if not _claim_settlement(context.logging_obj): @@ -141,11 +147,21 @@ async def poll_and_log_background_interaction_cost( 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, - context.timeout_seconds, - ) + if last_seen_status is not None and last_seen_status not in _POLLABLE_STATUSES: + verbose_logger.error( + "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither " + "the pollable nor the terminal set, so this proxy never learned how to settle it and its usage " + "will not be tracked", + context.interaction_id, + context.timeout_seconds, + last_seen_status, + ) + else: + 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) @@ -193,8 +209,14 @@ def is_pollable_background_interaction(response: InteractionsAPIResponse) -> boo exactly these responses, on the promise that a poll task will settle them, so a response one site accepts and the other refuses strands its reservation on the spend counters with nothing left to reconcile it. + + ``queued`` belongs here alongside ``in_progress``. It is the API's + not-started-yet state, so it reaches a terminal status the same way and + needs polling for the same reason: nothing else in the proxy ever bills a + create that came back without usage, so a status missing from both this + set and ``_TERMINAL_STATUSES`` is billed nowhere and alerts nobody. """ - return response.status == "in_progress" and bool(response.id) + return response.status in _POLLABLE_STATUSES and bool(response.id) def missing_usage_is_expected(response: InteractionsAPIResponse) -> bool: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 853df91cb89..cdc12896678 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2212,10 +2212,16 @@ class Logging(LiteLLMLoggingBaseClass): cost breakdown at the zeros the usage-less create stamped and writing those zeros to the spend log. Dropping it makes this event price the settled body itself, against the deployment that served the create. + + The same throwaway call stamped the deployment identity that travels + with the price, so ``model_id`` and ``litellm_model_name`` go with it. + Left in place they overwrite the create's real deployment with the + poll's empty one in the payload every logging integration reads. """ settled_hidden_params: Final = getattr(result, "_hidden_params", None) if isinstance(settled_hidden_params, dict): - settled_hidden_params.pop("response_cost", None) + for poll_scoped_key in ("response_cost", "model_id", "litellm_model_name"): + settled_hidden_params.pop(poll_scoped_key, None) self._reset_success_emission_dedupe() await self.async_success_handler(result=result) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 759bdc71dfc..97f09de1b52 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -448,3 +448,98 @@ async def test_schedule_respects_kill_switch(monkeypatch): ) assert task is None + + +def test_every_status_the_api_can_return_is_either_pollable_or_terminal(): + """ + The proxy bills a usage-less create in exactly two ways: it polls the + interaction until it settles, or it recognises the status as terminal and + settles immediately. A status in neither set is billed by nobody, alerts + nobody, and releases its budget reservation, which is the zero-spend bug + this whole module exists to fix. + + Pinned against the generated spec enum rather than a hand-written list, so + a status Google adds later breaks this test instead of silently shipping + another unbilled path. + """ + from litellm.interactions.background_cost_polling import _POLLABLE_STATUSES, _TERMINAL_STATUSES + from litellm.types.interactions.generated import Status1 + + spec_statuses = {member.value for member in Status1} + handled = _POLLABLE_STATUSES | _TERMINAL_STATUSES + + assert spec_statuses - handled == set() + assert handled - spec_statuses == set() + + +@pytest.mark.asyncio +async def test_schedule_creates_poll_task_for_queued_create(): + """ + ``queued`` is the API's not-started-yet state. It carries no usage, so the + create cannot bill it, and it is not terminal, so nothing settles it: + without a poll task it is never charged at all. + """ + logging_obj = _logging_obj() + task = maybe_schedule_background_interaction_cost_polling( + response=_response("queued", 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 +async def test_poller_bills_an_interaction_that_started_out_queued(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("queued", with_usage=False), + _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) == 3 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +def test_poll_intervals_double_up_to_the_cap_and_stay_inside_the_timeout(): + """ + The degenerate cases are covered above; this pins the shape the proxy + actually ships, so an off-by-one in the doubling or in the remaining-budget + check cannot pass green. + """ + intervals = list(_poll_intervals(initial=5.0, maximum=60.0, timeout=3600.0)) + + assert intervals[:6] == [5.0, 10.0, 20.0, 40.0, 60.0, 60.0] + assert max(intervals) == 60.0 + assert sum(intervals) <= 3600.0 + assert sum(intervals) + 60.0 > 3600.0 + + +@pytest.mark.asyncio +async def test_giving_up_on_an_unrecognized_status_says_which_status_it_was(monkeypatch): + """ + A status outside both sets polls for the full timeout and then gives up. + The give-up line is the only trace it leaves, so it has to name the status + rather than reporting it as an interaction that was merely still running. + """ + import litellm.interactions.background_cost_polling as bg + + errors = [] + monkeypatch.setattr(bg.verbose_logger, "error", lambda *args, **kwargs: errors.append(args)) + + logging_obj = _logging_obj() + fetch, _ = _fetch_sequence(_response("halted_for_review", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), fetch_interaction=fetch + ) + + assert len(errors) == 1 + assert "halted_for_review" in errors[0] 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 d840611fc7b..ca517474a5c 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 @@ -785,11 +785,17 @@ def _in_progress_interaction_kwargs(reservation: dict) -> dict: @pytest.mark.asyncio -async def test_track_cost_callback_keeps_reservation_open_for_in_progress_background_interaction(): +@pytest.mark.parametrize("status", ["in_progress", "queued"]) +async def test_track_cost_callback_keeps_reservation_open_for_in_progress_background_interaction(status): """ 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. + + ``queued`` is in flight for the same reason ``in_progress`` is: it has not + reached a terminal status, so releasing its reservation here would drop the + estimate off the spend counters while the interaction is still going to run + and still going to cost money. """ from litellm.types.interactions import InteractionsAPIResponse @@ -798,7 +804,7 @@ async def test_track_cost_callback_keeps_reservation_open_for_in_progress_backgr in_progress_response = InteractionsAPIResponse( id="interactions/bg-abc", model="gemini-3-flash-preview", - status="in_progress", + status=status, ) with patch( @@ -982,29 +988,59 @@ async def test_track_cost_callback_releases_reservation_for_interaction_without_ mock_proxy_logging.failed_tracking_alert.assert_not_called() -@pytest.mark.parametrize( - "status", - ["in_progress", "completed", "failed", "cancelled", "incomplete", "requires_action"], -) -@pytest.mark.parametrize("interaction_id", ["interactions/bg-abc", ""]) -def test_callback_defers_exactly_the_interactions_the_scheduler_polls(status, interaction_id): +@pytest.mark.asyncio +async def test_callback_handles_every_status_the_interactions_api_can_return(): """ - Pins the invariant the two modules share: the callback may only hold a - budget reservation open for a response the scheduler will actually poll. - Any drift between the two gates leaks reservations onto live spend - counters, so assert they agree rather than restating either condition. + Whatever status a usage-less create comes back with, exactly one of two + things has to happen to its budget reservation: the callback holds it open + for a poll task that will settle it, or it releases it on the spot. A + status that falls through both leaves the pre-call estimate pinned to the + key, user, team and org spend counters forever, refusing traffic against + budget nobody spent. + + Driven off the generated spec enum so a status Google adds later fails here + instead of quietly leaking reservations in production. """ - from litellm.interactions.background_cost_polling import is_pollable_background_interaction - from litellm.proxy.hooks.proxy_track_cost_callback import _is_unbilled_in_progress_interaction from litellm.types.interactions import InteractionsAPIResponse + from litellm.types.interactions.generated import Status1 - response = InteractionsAPIResponse( - id=interaction_id, - model="gemini-3-flash-preview", - status=status, - ) + deferred = set() + released = set() - assert _is_unbilled_in_progress_interaction(response) is is_pollable_background_interaction(response) + for status in sorted(member.value for member in Status1): + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + 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=response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + (deferred if reservation["finalized"] is False else released).add(status) + + assert deferred == {"in_progress", "queued"} + assert released == { + "completed", + "requires_action", + "failed", + "cancelled", + "incomplete", + "budget_exceeded", + } @pytest.mark.asyncio From 9f1191e7a06e1bf2bb6e66a8ab26a60a2ff61a8e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:40:59 -0700 Subject: [PATCH 064/156] test: pin the gemini 2.5 promotional discount instead of calling it stale The two gemini 2.5 entries price a factor of 1.25 under the published DBU table because the published figures exclude a 20% promotion that runs to 2027-01-31. The previous constant name and test called them an older vintage awaiting a refresh, which would have led a future reader to scale them up and overcharge. Pin the discount and the cache relationship instead. --- .../databricks/test_databricks_cost_calculator.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 81517e23b53..d5abbc547e5 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -62,7 +62,8 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-gemini-2-5-pro": ("22.321", "178.571", "22.321", "2.232"), "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), } -OLDER_VINTAGE_MODELS: Final = ( +PROMOTIONAL_DISCOUNT: Final = 0.80 +PROMOTIONALLY_DISCOUNTED_MODELS: Final = ( "databricks/databricks-gemini-2-5-pro", "databricks/databricks-gemini-2-5-flash", ) @@ -135,7 +136,7 @@ def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, mod assert info[field] == _dollars_per_token(dbu_per_million), field -@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(OLDER_VINTAGE_MODELS))) +@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(PROMOTIONALLY_DISCOUNTED_MODELS))) def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] @@ -221,12 +222,17 @@ def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: No assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field -@pytest.mark.parametrize("model", OLDER_VINTAGE_MODELS) -def test_entries_priced_at_an_older_vintage_keep_cache_rates_tied_to_their_own_input( +@pytest.mark.parametrize("model", PROMOTIONALLY_DISCOUNTED_MODELS) +def test_promotionally_discounted_entries_price_below_the_published_table( local_model_cost_map: None, model: str, ) -> None: info: Final = _model_info(model) + input_dbu, output_dbu, _, _ = PUBLISHED_DBU_PER_MILLION[model] + assert info["input_cost_per_token"] == pytest.approx(_dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3) + assert info["output_cost_per_token"] == pytest.approx( + _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3 + ) assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) From 77cae927d5d9bc70ea19641a11b107a44bc444a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:54:43 -0700 Subject: [PATCH 065/156] fix(databricks): keep fable-5 declared text-only Databricks documents databricks-claude-fable-5 as accepting text only, where every sibling Claude endpoint accepts text and image. An earlier commit flipped the flag to true on the reasoning that fable-5 was the only entry of its generation declaring false, which had it backwards: it is the only one because the endpoint really does refuse images. Advertising vision here would surface the model in capability filters and hand the caller a provider-side rejection. --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1dc84f06aa7..8c8b909493a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14614,7 +14614,7 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true, + "supports_vision": false, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1dc84f06aa7..8c8b909493a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14614,7 +14614,7 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true, + "supports_vision": false, "thinking_always_on": true }, "databricks/databricks-claude-haiku-4-5": { From 15e8a35feb0bbdbe113c2389bf00c0738cb8bbc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:54:44 -0700 Subject: [PATCH 066/156] test: separate gemini entries storing the promo rate from those storing list The 20% promotion that runs to 2027-01-31 covers every gemini model, not just the 2.5 pair, so a constant naming two of them implied the other four were exempt. Six covered entries live in the registry: two store the discounted rate and four store list, which is a pre-existing overcharge this branch does not touch, since it only adds cache fields and derives them from each entry's own input rate. Name both groups for what they store, pin the expiry, and tighten the tolerance to 2e-4. --- .../test_databricks_cost_calculator.py | 40 +++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index d5abbc547e5..21f047b753c 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -63,10 +63,17 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), } PROMOTIONAL_DISCOUNT: Final = 0.80 -PROMOTIONALLY_DISCOUNTED_MODELS: Final = ( +PROMOTION_EXPIRES: Final = "2027-01-31" +ENTRIES_STORING_PROMOTIONAL_RATE: Final = ( "databricks/databricks-gemini-2-5-pro", "databricks/databricks-gemini-2-5-flash", ) +ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION: Final = ( + "databricks/databricks-gemini-3-1-pro", + "databricks/databricks-gemini-3-pro", + "databricks/databricks-gemini-3-flash", + "databricks/databricks-gemini-3-1-flash-lite", +) CACHE_FIELDS: Final = ("cache_creation_input_token_cost", "cache_read_input_token_cost") @@ -136,7 +143,7 @@ def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, mod assert info[field] == _dollars_per_token(dbu_per_million), field -@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(PROMOTIONALLY_DISCOUNTED_MODELS))) +@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(ENTRIES_STORING_PROMOTIONAL_RATE))) def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] @@ -222,17 +229,36 @@ def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: No assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field -@pytest.mark.parametrize("model", PROMOTIONALLY_DISCOUNTED_MODELS) -def test_promotionally_discounted_entries_price_below_the_published_table( +@pytest.mark.parametrize("model", ENTRIES_STORING_PROMOTIONAL_RATE) +def test_entries_storing_the_promotional_rate_price_below_the_published_table( local_model_cost_map: None, model: str, ) -> None: info: Final = _model_info(model) input_dbu, output_dbu, _, _ = PUBLISHED_DBU_PER_MILLION[model] + expiry_hint: Final = f"the gemini promotion expires {PROMOTION_EXPIRES}, after which the list rate applies" - assert info["input_cost_per_token"] == pytest.approx(_dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3) + assert info["input_cost_per_token"] == pytest.approx( + _dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint assert info["output_cost_per_token"] == pytest.approx( - _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3 - ) + _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) + + +@pytest.mark.parametrize("model", ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION) +def test_entries_storing_the_list_rate_bill_above_the_promotional_price( + local_model_cost_map: None, + model: str, +) -> None: + info: Final = _model_info(model) + input_dbu, _, _, _ = PUBLISHED_DBU_PER_MILLION[model] + list_rate: Final = _dollars_per_token(input_dbu) + + assert info["input_cost_per_token"] == pytest.approx(list_rate, rel=2e-4), ( + f"{model} moved off the list rate; if it now stores the discount that runs to " + f"{PROMOTION_EXPIRES}, move it into ENTRIES_STORING_PROMOTIONAL_RATE" + ) + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) From 0697188be40c0f5528b2ca18ec110fd95b55d9f8 Mon Sep 17 00:00:00 2001 From: milan Date: Sun, 23 Aug 2026 03:14:19 +0000 Subject: [PATCH 067/156] refactor(a2a): use direct typed access in protocol binding normalization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/a2a_protocol/card_resolver.py | 8 +++----- tests/test_litellm/a2a_protocol/test_main.py | 1 + 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 8070d88b761..47d6346b9a6 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -67,11 +67,9 @@ def normalize_agent_card_protocol_bindings(agent_card: "AgentCard") -> "AgentCar case-sensitively against its uppercase TransportProtocol constants and fails with "no compatible transports found." for spec-adjacent casings. """ - interfaces: Final = getattr(agent_card, "supported_interfaces", None) or () - for interface in interfaces: - binding: str = getattr(interface, "protocol_binding", "") or "" - canonical = _CANONICAL_PROTOCOL_BINDINGS.get(binding.lower()) - if canonical is not None and binding != canonical: + for interface in agent_card.supported_interfaces: + canonical: str | None = _CANONICAL_PROTOCOL_BINDINGS.get(interface.protocol_binding.lower()) + if canonical is not None and canonical != interface.protocol_binding: interface.protocol_binding = canonical return agent_card diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 29aa3aaabd7..aa988bc5b99 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -344,6 +344,7 @@ async def test_lowercase_protocol_binding_in_agent_card_still_gets_a_client(isol response = await _send_message(a2a_client, _send_request("lc")) assert type(response.root.result).__name__ == "Message" + assert a2a_client._litellm_agent_card.supported_interfaces[0].protocol_binding == "JSONRPC" @pytest.mark.asyncio From de0d8ceb25ea4cc3aa1df79b494ca95646a709f8 Mon Sep 17 00:00:00 2001 From: milan Date: Sun, 23 Aug 2026 03:20:53 +0000 Subject: [PATCH 068/156] refactor(a2a): return a normalized card copy instead of mutating the resolved card Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/a2a_protocol/card_resolver.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 47d6346b9a6..9b3e4ec13b9 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -67,11 +67,13 @@ def normalize_agent_card_protocol_bindings(agent_card: "AgentCard") -> "AgentCar case-sensitively against its uppercase TransportProtocol constants and fails with "no compatible transports found." for spec-adjacent casings. """ - for interface in agent_card.supported_interfaces: + normalized: Final = type(agent_card)() + normalized.CopyFrom(agent_card) + for interface in normalized.supported_interfaces: canonical: str | None = _CANONICAL_PROTOCOL_BINDINGS.get(interface.protocol_binding.lower()) - if canonical is not None and canonical != interface.protocol_binding: + if canonical is not None: interface.protocol_binding = canonical - return agent_card + return normalized def get_agent_card_url(agent_card: "AgentCard") -> str | None: From ee935cec230207a549f2939a6a9ca69023c609d6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 24 Aug 2026 09:46:35 -0700 Subject: [PATCH 069/156] refactor(proxy): trim the multi_items comment to the non-obvious clause --- litellm/proxy/common_utils/http_parsing_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 4cb55f6966e..96621b08ba1 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -275,9 +275,7 @@ async def get_form_data(request: Request) -> dict[str, Any]: """ form: Final = await request.form() parsed_form_data: Final[dict[str, Any]] = {} - # multi_items(), not dict(form): a dict drops every value but the last of a repeated key, - # which is the whole array this function exists to rebuild - for key, value in form.multi_items(): + for key, value in form.multi_items(): # not dict(form), which keeps only the last repeat if key.endswith("[]"): clean_key = key[:-2] parsed_form_data.setdefault(clean_key, []).append(value) From a72203eae4e98d216668b3a7216dda8a628431ca Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:12:10 -0700 Subject: [PATCH 070/156] fix(terraform): add soft_budget, tags, and soft_budget_alerting_emails to litellm_team (#37918) * fix(terraform): add soft_budget, tags, and soft_budget_alerting_emails to litellm_team The team resource rejected soft_budget and tags at plan time and had no way to express the list-valued metadata.soft_budget_alerting_emails the proxy reads for soft-budget alerts, even though /team/new and /team/update accept all three. Add the attributes, forward them in buildTeamData (alert emails merged under metadata, where the proxy stores them), and send the full metadata map whenever either half changes because /team/update replaces metadata wholesale. Read was decoding /team/info as if the team fields were top-level, but the proxy nests them under team_info, so every attribute silently fell back to prior state. Decode the envelope and split the proxy's metadata back into tags / soft_budget_alerting_emails / string metadata, dropping the server-managed team_member_budget_id. Verified with OpenTofu plan/apply against a live proxy: the attributes are accepted, land on the proxy, refresh into state, re-plan clean, propagate on update, and clear when removed from HCL. * fix(terraform): clear litellm_team.soft_budget in state when the proxy returns null Read only wrote soft_budget when the proxy returned a value, so a soft budget cleared outside Terraform stayed in state and never surfaced as drift. Set it from the response unconditionally so a null clears it. --- terraform/provider/CHANGELOG.md | 8 + terraform/provider/docs/resources/team.md | 15 +- terraform/provider/litellm/resource_team.go | 97 ++++++++- .../provider/litellm/resource_team_test.go | 184 ++++++++++++++++++ terraform/provider/litellm/types.go | 6 + 5 files changed, 300 insertions(+), 10 deletions(-) create mode 100644 terraform/provider/litellm/resource_team_test.go diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index ff2f3f817f9..3bd3c4d1d6c 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -14,6 +14,14 @@ longer signal it. ## [Unreleased] +### Added + +- **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it + +### Fixed + +- **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state + ### Changed - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 68535309f10..65ab4bf82d4 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -24,11 +24,18 @@ resource "litellm_team" "advanced_team" { # Budget and rate limiting max_budget = 1000.0 + soft_budget = 800.0 budget_duration = "1mo" tpm_limit = 500000 rpm_limit = 5000 blocked = false + # Who gets paged when spend crosses soft_budget + soft_budget_alerting_emails = ["finops@example.com"] + + # Tags for spend tracking and tag-based routing + tags = ["team:ai-research", "environment:production"] + # Team member permissions team_member_permissions = [ "create_key", @@ -91,7 +98,9 @@ The following arguments are supported: * `models` - (Optional) List of model names that this team can access. -* `metadata` - (Optional) A map of metadata key-value pairs associated with the team. +* `metadata` - (Optional) A map of string metadata key-value pairs associated with the team. `tags` and `soft_budget_alerting_emails` are stored by the proxy under metadata but are managed through their own attributes below, not this map. + +* `tags` - (Optional) List of tags applied to the team, used for [spend tracking](https://docs.litellm.ai/docs/proxy/enterprise#tracking-spend-for-custom-tags) and [tag-based routing](https://docs.litellm.ai/docs/proxy/tag_routing). * `blocked` - (Optional) Whether the team is blocked from making requests. Default is `false`. @@ -101,6 +110,10 @@ The following arguments are supported: * `max_budget` - (Optional) Maximum budget allocated to the team. +* `soft_budget` - (Optional) Spend threshold at which the proxy sends a soft budget alert without blocking requests. + +* `soft_budget_alerting_emails` - (Optional) List of email addresses notified when the team's spend crosses `soft_budget`. + * `budget_duration` - (Optional) Duration for the budget cycle. Valid values are: * `daily` * `weekly` diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 88e0dcd4811..2a167a1b5c4 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -53,6 +53,11 @@ func ResourceLiteLLMTeam() *schema.Resource { Type: schema.TypeFloat, Optional: true, }, + "soft_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Spend threshold that triggers a soft budget alert without blocking requests", + }, "budget_duration": { Type: schema.TypeString, Optional: true, @@ -72,6 +77,18 @@ func ResourceLiteLLMTeam() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, Description: "List of permissions granted to team members", }, + "tags": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Tags for spend tracking and tag-based routing", + }, + "soft_budget_alerting_emails": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Email addresses alerted when the team crosses soft_budget", + }, }, } } @@ -117,21 +134,20 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { return nil } - var teamResp TeamResponse - if err := json.NewDecoder(resp.Body).Decode(&teamResp); err != nil { + var infoResp TeamInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { return fmt.Errorf("error decoding team info response: %w", err) } + teamResp := infoResp.TeamInfo // Update the state with values from the response or fall back to the data passed in during creation d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string))) d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string))) - // Handle metadata separately as it's a map - if teamResp.Metadata != nil { - d.Set("metadata", teamResp.Metadata) - } else { - d.Set("metadata", d.Get("metadata")) - } + metadata, tags, alertEmails := splitTeamMetadata(teamResp.Metadata) + d.Set("metadata", metadata) + d.Set("tags", tags) + d.Set("soft_budget_alerting_emails", alertEmails) if teamResp.TPMLimit != nil { d.Set("tpm_limit", *teamResp.TPMLimit) @@ -142,6 +158,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { if teamResp.MaxBudget != nil { d.Set("max_budget", *teamResp.MaxBudget) } + d.Set("soft_budget", teamResp.SoftBudget) d.Set("budget_duration", GetStringValue(teamResp.BudgetDuration, d.Get("budget_duration").(string))) // Handle models separately as it's a list @@ -240,15 +257,77 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "team_alias": d.Get("team_alias").(string), } - for _, key := range []string{"organization_id", "metadata", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions"} { + for _, key := range []string{"organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions"} { if v, ok := d.GetOk(key); ok { teamData[key] = v } } + if v, ok := d.GetOk("soft_budget"); ok { + teamData["soft_budget"] = v + } else if d.HasChange("soft_budget") { + teamData["soft_budget"] = nil + } + + if v, ok := d.GetOk("tags"); ok || d.HasChange("tags") { + teamData["tags"] = v + } + + if metadata := buildTeamMetadata(d); metadata != nil { + teamData["metadata"] = metadata + } + return teamData } +// /team/update replaces metadata wholesale, so the full map must go out whenever either half changed. +func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} { + metadata := map[string]interface{}{} + for k, v := range d.Get("metadata").(map[string]interface{}) { + metadata[k] = v + } + if v, ok := d.GetOk("soft_budget_alerting_emails"); ok { + metadata["soft_budget_alerting_emails"] = v + } + if len(metadata) == 0 && !d.HasChange("metadata") && !d.HasChange("soft_budget_alerting_emails") { + return nil + } + return metadata +} + +func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) { + metadata := map[string]string{} + var tags, alertEmails []string + for k, v := range raw { + switch k { + case "tags": + tags = toStringSlice(v) + case "soft_budget_alerting_emails": + alertEmails = toStringSlice(v) + case "team_member_budget_id": + default: + if s, ok := v.(string); ok { + metadata[k] = s + } + } + } + return metadata, tags, alertEmails +} + +func toStringSlice(v interface{}) []string { + items, ok := v.([]interface{}) + if !ok { + return nil + } + out := make([]string, 0, len(items)) + for _, item := range items { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return out +} + func handleResponse(resp *http.Response, action string) error { if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go new file mode 100644 index 00000000000..1f74be4819d --- /dev/null +++ b/terraform/provider/litellm/resource_team_test.go @@ -0,0 +1,184 @@ +package litellm + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +func newTeamTestServer(t *testing.T, captured *map[string]interface{}, infoBody string) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case endpointTeamNew, endpointTeamUpdate: + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, captured) + w.Write([]byte(`{}`)) + case endpointTeamInfo: + w.Write([]byte(infoBody)) + case endpointTeamPermissionsList: + w.Write([]byte(`{"team_id":"team-1","team_member_permissions":[],"all_available_permissions":[]}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) +} + +const teamInfoWithSoftBudget = `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "insights", + "max_budget": 750.0, + "soft_budget": 600.0, + "models": ["claude-haiku-4-5"], + "metadata": { + "department": "customer-insights", + "tags": ["team:customer-insights", "environment:production"], + "soft_budget_alerting_emails": ["finops@example.com"], + "team_member_budget_id": "budget-1" + } + }, + "keys": [], + "team_memberships": [] +}` + +func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{ + "team_alias": "insights", + "max_budget": 750.0, + "soft_budget": 600.0, + "tags": []interface{}{"team:customer-insights", "environment:production"}, + "soft_budget_alerting_emails": []interface{}{"finops@example.com"}, + "metadata": map[string]interface{}{"department": "customer-insights"}, + }) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if got := captured["soft_budget"]; got != 600.0 { + t.Fatalf("payload soft_budget = %v, want 600", got) + } + wantTags := []interface{}{"team:customer-insights", "environment:production"} + if got := captured["tags"]; !reflect.DeepEqual(got, wantTags) { + t.Fatalf("payload tags = %v, want %v", got, wantTags) + } + wantMetadata := map[string]interface{}{ + "department": "customer-insights", + "soft_budget_alerting_emails": []interface{}{"finops@example.com"}, + } + if got := captured["metadata"]; !reflect.DeepEqual(got, wantMetadata) { + t.Fatalf("payload metadata = %v, want %v", got, wantMetadata) + } +} + +func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{}) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("team_alias"); got != "insights" { + t.Fatalf("team_alias = %v, want insights", got) + } + if got := d.Get("soft_budget"); got != 600.0 { + t.Fatalf("soft_budget = %v, want 600", got) + } + if got := d.Get("max_budget"); got != 750.0 { + t.Fatalf("max_budget = %v, want 750", got) + } + wantTags := []interface{}{"team:customer-insights", "environment:production"} + if got := d.Get("tags"); !reflect.DeepEqual(got, wantTags) { + t.Fatalf("tags = %v, want %v", got, wantTags) + } + wantEmails := []interface{}{"finops@example.com"} + if got := d.Get("soft_budget_alerting_emails"); !reflect.DeepEqual(got, wantEmails) { + t.Fatalf("soft_budget_alerting_emails = %v, want %v", got, wantEmails) + } + wantMetadata := map[string]interface{}{"department": "customer-insights"} + if got := d.Get("metadata"); !reflect.DeepEqual(got, wantMetadata) { + t.Fatalf("metadata = %v, want %v (server-managed team_member_budget_id dropped)", got, wantMetadata) + } +} + +func TestTeamUpdateClearsRemovedTagsAndSoftBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"insights"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "insights", + "soft_budget": 600.0, + "tags": []interface{}{"team:to-be-removed"}, + "soft_budget_alerting_emails": []interface{}{"ops@example.com"}, + "metadata": map[string]interface{}{"department": "eng"}, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "team_alias": "insights", + "metadata": map[string]interface{}{"department": "eng"}, + }) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if got, ok := captured["soft_budget"]; !ok || got != nil { + t.Fatalf("payload soft_budget = %v (present=%v), want explicit null", got, ok) + } + if got := captured["tags"]; !reflect.DeepEqual(got, []interface{}{}) { + t.Fatalf("payload tags = %v, want []", got) + } + if got := captured["metadata"]; !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) { + t.Fatalf("payload metadata = %v, want department only", got) + } +} + +func TestTeamReadClearsSoftBudgetWhenProxyReturnsNull(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"insights","soft_budget":null},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{ + "team_alias": "insights", + "soft_budget": 600.0, + }) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("soft_budget"); got != 0.0 { + t.Fatalf("soft_budget = %v, want cleared after the proxy returned null", got) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 069fe4b3e23..66d1f6a8ba9 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -33,6 +33,11 @@ type ModelRequest struct { Additional map[string]interface{} `json:"additional"` } +type TeamInfoResponse struct { + TeamID string `json:"team_id"` + TeamInfo TeamResponse `json:"team_info"` +} + // TeamResponse represents a response from the API containing team information. type TeamResponse struct { TeamID string `json:"team_id,omitempty"` @@ -42,6 +47,7 @@ type TeamResponse struct { TPMLimit *int `json:"tpm_limit,omitempty"` RPMLimit *int `json:"rpm_limit,omitempty"` MaxBudget *float64 `json:"max_budget,omitempty"` + SoftBudget *float64 `json:"soft_budget,omitempty"` BudgetDuration string `json:"budget_duration,omitempty"` Models []string `json:"models"` Blocked bool `json:"blocked,omitempty"` From 5f56be3294b753e8418defafe36e99f669263720 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:12:16 -0700 Subject: [PATCH 071/156] fix(ui): theme the created-key box so it follows dark mode (#37985) The virtual key shown after creating a key sits in a div with a hardcoded #f8f8f8 inline background, so in dark mode the box keeps the light background while the key text inherits the light foreground color, leaving the key nearly unreadable. Swap the inline styles for the bg-muted and text-foreground tokens, which resolve per theme. --- .../src/components/shared/CreatedKeyDisplay.test.tsx | 9 +++++++++ .../src/components/shared/CreatedKeyDisplay.tsx | 11 ++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx index 6c8cdae6deb..0e10d32f8d1 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx @@ -19,6 +19,15 @@ describe("CreatedKeyDisplay", () => { expect(screen.getByText("sk-test-123")).toBeInTheDocument(); }); + it("should theme the key box with tokens instead of a hardcoded light background", () => { + render(); + + const keyBox = screen.getByText("sk-test-123").parentElement as HTMLElement; + + expect(keyBox).not.toHaveAttribute("style"); + expect(keyBox).toHaveClass("bg-muted"); + }); + it("should display the security warning", () => { render(); expect(screen.getByText(/you will not be able to view it again/i)).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx index c14379378d3..5dbc17b13bf 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx @@ -29,15 +29,8 @@ const CreatedKeyDisplay: React.FC = ({ apiKey }) => {

Virtual Key:

-
- 
{apiKey}
+
+ 
{apiKey}
From 6db5a5d6605114c672e2ef3c0fff6a5f76f56721 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:12:46 -0700 Subject: [PATCH 072/156] fix(ui): restore the public model name tooltip layout in the add model flow (#37986) The tooltip popup is an inline-flex row, so the four sibling blocks passed as a fragment laid out side by side in four columns. Wrap them in a single flex-col container instead. The inline code samples also used bg-muted, which is defined against the page surface, not the inverted tooltip surface, so they rendered as near-white chips carrying near-white text. Tint them from the popup's own token instead. --- .../conditional_public_model_name.tsx | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/conditional_public_model_name.tsx b/ui/litellm-dashboard/src/components/add_model/conditional_public_model_name.tsx index b9c128a3d51..c58b177a6e1 100644 --- a/ui/litellm-dashboard/src/components/add_model/conditional_public_model_name.tsx +++ b/ui/litellm-dashboard/src/components/add_model/conditional_public_model_name.tsx @@ -34,6 +34,8 @@ const modelMappingsRule = { }, }; +const tooltipCodeClassName = "rounded-sm bg-background/20 px-1 py-0.5 font-mono text-xs"; + const ConditionalPublicModelName: React.FC = () => { const form = useFormContext(); @@ -123,22 +125,22 @@ const ConditionalPublicModelName: React.FC = () => { if (!showPublicModelName) return null; const publicNameTooltipContent = ( - <> -
The name you specify in your API calls to LiteLLM Proxy
-
+
+
The name you specify in your API calls to LiteLLM Proxy
+
Example: If you name your public model{" "} - example-name, and choose{" "} - openai/qwen-plus-latest as the LiteLLM model + example-name, and choose{" "} + openai/qwen-plus-latest as the LiteLLM model
-
+
Usage: You make an API call to the LiteLLM proxy with{" "} - model = "example-name" + model = "example-name"
-
- Result: LiteLLM sends{" "} - qwen-plus-latest to the provider +
+ Result: LiteLLM sends qwen-plus-latest to the + provider
- +
); const liteLLMModelTooltipContent =
The model name LiteLLM will send to the LLM API
; From 5b1c142c6e2aeee84360cad44192f382510c6c04 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:12:52 -0700 Subject: [PATCH 073/156] fix(ui): render team and org tpm/rpm limits of 0 as 0 instead of Unlimited (#37916) * fix(ui): render team and org tpm/rpm limits of 0 as 0 instead of Unlimited A tpm_limit or rpm_limit of 0 is a hard block on the backend (every request 429s) and only null means unlimited, but the team and organization views rendered both as "Unlimited" (and a team-member limit of 0 as "No Limit") because every display site used a falsy || fallback. The team member edit dialog also seeded its form with `tpm_limit || null`, so opening Edit Member on a member stored with 0 and clicking Save sent null to /team/member_update and silently turned the hard block into unlimited Every limit display site in TeamInfo, organization_view, the organizations list cell and the team members table now uses a nullish check, and both member form seeding paths keep 0 for max_budget_in_team, tpm_limit and rpm_limit. Regression tests cover each site and the existing memberFormValues test that asserted 0 -> null is flipped to assert 0 survives Resolves LIT-5760 * test(ui): assert a stored 0 member limit survives an untouched save The EditMembership integration test named the old 0 -> null collapse as the expected payload, so the related-tests CI job went red once the form kept 0. It now asserts 0 survives and only the empty budget_duration collapses to null. The TeamMemberTab fixture is built with a map instead of mutating the nested membership --- .../_components/OrganizationsTable.test.tsx | 14 +++++++ .../_components/OrganizationsTableColumns.tsx | 4 +- .../organization/organization_view.test.tsx | 36 ++++++++++++++++- .../organization/organization_view.tsx | 8 ++-- .../team/EditMembership.integration.test.tsx | 8 ++-- .../src/components/team/TeamInfo.test.tsx | 28 +++++++++++++ .../src/components/team/TeamInfo.tsx | 14 +++---- .../components/team/TeamMemberTab.test.tsx | 39 ++++++++++++++++++- .../src/components/team/TeamMemberTab.tsx | 10 ++--- .../components/team/memberFormValues.test.ts | 20 +++++++++- .../src/components/team/memberFormValues.ts | 6 +-- 11 files changed, 158 insertions(+), 29 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index a06c5c885e3..1ac33a27186 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -165,6 +165,20 @@ describe("OrganizationsTable", () => { expect(screen.getByText("RPM: Unlimited")).toBeInTheDocument(); }); + it("renders a tpm/rpm limit of 0 as 0, never as Unlimited", () => { + render( + , + ); + + expect(screen.getByText("TPM: 0")).toBeInTheDocument(); + expect(screen.getByText("RPM: 0")).toBeInTheDocument(); + expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument(); + expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument(); + }); + it("renders loading skeletons instead of rows while loading", () => { render( - TPM: {tpm_limit ? tpm_limit : "Unlimited"} - RPM: {rpm_limit ? rpm_limit : "Unlimited"} + TPM: {tpm_limit ?? "Unlimited"} + RPM: {rpm_limit ?? "Unlimited"}
); } diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index 0337fea4131..b8d8e3ba9c8 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { vi, test, expect, beforeEach } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -290,3 +290,37 @@ test("should keep unsaved settings edits when switching tabs and back", async () expect(screen.getByLabelText(/Organization Name/i)).toHaveValue("Renamed Org"); }); + +test("renders a tpm/rpm limit of 0 as 0 in the overview and settings tabs, never as Unlimited", async () => { + const zeroLimitOrg = { + ...mockOrg, + litellm_budget_table: { ...mockOrg.litellm_budget_table, tpm_limit: 0, rpm_limit: 0 }, + }; + mockUseOrganization.mockReturnValue({ data: zeroLimitOrg, isLoading: false } as unknown as ReturnType< + typeof useOrganization + >); + + const user = userEvent.setup(); + renderWithProviders( + {}} + accessToken="test-token" + is_org_admin={false} + is_proxy_admin={true} + userModels={[]} + editOrg={false} + />, + ); + + const overview = await screen.findByRole("tabpanel", { name: "Overview" }); + expect(within(overview).getByText("TPM: 0")).toBeInTheDocument(); + expect(within(overview).getByText("RPM: 0")).toBeInTheDocument(); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + const settings = await screen.findByRole("tabpanel", { name: "Settings" }); + expect(within(settings).getByText("TPM: 0")).toBeInTheDocument(); + expect(within(settings).getByText("RPM: 0")).toBeInTheDocument(); + expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument(); + expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.tsx index 2590f3bad52..7c250a5ef98 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.tsx @@ -206,8 +206,8 @@ const OrganizationInfoView: React.FC = ({

Rate Limits

-

TPM: {orgData.litellm_budget_table.tpm_limit || "Unlimited"}

-

RPM: {orgData.litellm_budget_table.rpm_limit || "Unlimited"}

+

TPM: {orgData.litellm_budget_table.tpm_limit ?? "Unlimited"}

+

RPM: {orgData.litellm_budget_table.rpm_limit ?? "Unlimited"}

{orgData.litellm_budget_table.max_parallel_requests && (

Max Parallel Requests: {orgData.litellm_budget_table.max_parallel_requests}

)} @@ -311,8 +311,8 @@ const OrganizationInfoView: React.FC = ({

Rate Limits

-
TPM: {orgData.litellm_budget_table.tpm_limit || "Unlimited"}
-
RPM: {orgData.litellm_budget_table.rpm_limit || "Unlimited"}
+
TPM: {orgData.litellm_budget_table.tpm_limit ?? "Unlimited"}
+
RPM: {orgData.litellm_budget_table.rpm_limit ?? "Unlimited"}

Budget

diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx index b1dd3a2efcb..a82f475512c 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx @@ -110,7 +110,7 @@ describe("EditMembership submit payload", () => { expect(submitted()).toStrictEqual(expected); }); - it("collapses falsy budget and limit values to null and a missing model list to an empty array", async () => { + it("keeps stored 0 budget and limits as 0 on an untouched save, collapsing only empty strings and a missing model list", async () => { renderEdit(teamMemberConfig, { user_id: "u1", user_email: "a@b.com", @@ -128,10 +128,10 @@ describe("EditMembership submit payload", () => { user_email: "a@b.com", user_id: "u1", role: "user", - max_budget_in_team: null, + max_budget_in_team: 0, budget_duration: null, - tpm_limit: null, - rpm_limit: null, + tpm_limit: 0, + rpm_limit: 0, allowed_models: [], }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a7e8e6788dd..c3cc362e29e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -319,6 +319,34 @@ describe("TeamInfoView", () => { expect(screen.getByText(/of \$1,000\.00/)).toBeInTheDocument(); }); + it("renders a tpm/rpm/budget limit of 0 as 0 in the overview and settings tabs, never as Unlimited or No Limit", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + tpm_limit: 0, + rpm_limit: 0, + team_member_budget_table: { max_budget: 0, budget_duration: null, tpm_limit: 0, rpm_limit: 0 }, + }), + ); + + renderWithProviders(); + + const overview = await screen.findByRole("tabpanel", { name: "Overview" }); + expect(within(overview).getByText("TPM: 0")).toBeInTheDocument(); + expect(within(overview).getByText("RPM: 0")).toBeInTheDocument(); + + await userEvent.setup({ delay: null }).click(screen.getByRole("tab", { name: "Settings" })); + const settings = await screen.findByRole("tabpanel", { name: "Settings" }); + expect(within(settings).getByText("TPM: 0")).toBeInTheDocument(); + expect(within(settings).getByText("RPM: 0")).toBeInTheDocument(); + expect(within(settings).getByText("TPM Limit: 0")).toBeInTheDocument(); + expect(within(settings).getByText("RPM Limit: 0")).toBeInTheDocument(); + expect(within(settings).getByText("Max Budget: 0")).toBeInTheDocument(); + expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument(); + expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument(); + expect(screen.queryByText("TPM Limit: No Limit")).not.toBeInTheDocument(); + expect(screen.queryByText("RPM Limit: No Limit")).not.toBeInTheDocument(); + }); + it("should display guardrails in overview when present", async () => { vi.mocked(networking.teamInfoCall).mockResolvedValue( createMockTeamData({ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 727445e4fcd..88a893fd3e4 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -971,8 +971,8 @@ const TeamInfoView: React.FC = ({

Rate Limits

-

TPM: {info.tpm_limit || "Unlimited"}

-

RPM: {info.rpm_limit || "Unlimited"}

+

TPM: {info.tpm_limit ?? "Unlimited"}

+

RPM: {info.rpm_limit ?? "Unlimited"}

{info.max_parallel_requests &&

Max Parallel Requests: {info.max_parallel_requests}

} {(() => { const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record; @@ -1760,8 +1760,8 @@ const TeamInfoView: React.FC = ({

Rate Limits

-
TPM: {info.tpm_limit || "Unlimited"}
-
RPM: {info.rpm_limit || "Unlimited"}
+
TPM: {info.tpm_limit ?? "Unlimited"}
+
RPM: {info.rpm_limit ?? "Unlimited"}
{(() => { const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record; const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record; @@ -1811,11 +1811,11 @@ const TeamInfoView: React.FC = ({

-
Max Budget: {info.team_member_budget_table?.max_budget || "No Limit"}
+
Max Budget: {info.team_member_budget_table?.max_budget ?? "No Limit"}
Budget Duration: {info.team_member_budget_table?.budget_duration || "No Limit"}
Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}
-
TPM Limit: {info.team_member_budget_table?.tpm_limit || "No Limit"}
-
RPM Limit: {info.team_member_budget_table?.rpm_limit || "No Limit"}
+
TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}
+
RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}

Router Settings

diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 04234cf5a5e..6711514bfe9 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -1,4 +1,4 @@ -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -322,6 +322,43 @@ describe("TeamMembersComponent", () => { expect(mockSetSelectedEditMember).toHaveBeenCalled(); }); + it("keeps a member's stored 0 limits as 0 in the table and in the edit payload, never unlimited", async () => { + const user = userEvent.setup(); + vi.mocked(isProxyAdminRole).mockReturnValue(true); + const baseTeamData = createMockTeamData(); + const teamData = { + ...baseTeamData, + team_memberships: baseTeamData.team_memberships.map((membership, index) => + index === 0 + ? { + ...membership, + litellm_budget_table: { ...membership.litellm_budget_table, max_budget: 0, tpm_limit: 0, rpm_limit: 0 }, + } + : membership, + ), + }; + + renderWithProviders( + , + ); + + const memberRow = screen.getByRole("row", { name: /user1@test\.com/ }); + expect(within(memberRow).getByText("0 RPM / 0 TPM")).toBeInTheDocument(); + expect(within(memberRow).queryByText("No Limits")).not.toBeInTheDocument(); + + await user.click(within(memberRow).getByTestId("edit-member")); + + const zeroLimitsMember = { user_id: "user1@test.com", max_budget_in_team: 0, tpm_limit: 0, rpm_limit: 0 }; + expect(mockSetSelectedEditMember).toHaveBeenCalledWith(expect.objectContaining(zeroLimitsMember)); + }); + it("should call setIsAddMemberModalVisible when Add Member button is clicked", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index a662a02f91e..4d7246eb077 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -71,8 +71,8 @@ export default function TeamMemberTab({ const rpmLimit = membership?.litellm_budget_table?.rpm_limit; const tpmLimit = membership?.litellm_budget_table?.tpm_limit; - const rpmText = rpmLimit ? `${formatNumber(rpmLimit)} RPM` : null; - const tpmText = tpmLimit ? `${formatNumber(tpmLimit)} TPM` : null; + const rpmText = rpmLimit != null ? `${formatNumber(rpmLimit)} RPM` : null; + const tpmText = tpmLimit != null ? `${formatNumber(tpmLimit)} TPM` : null; const limits = [rpmText, tpmText].filter(Boolean); return limits.length > 0 ? limits.join(" / ") : "No Limits"; @@ -191,9 +191,9 @@ export default function TeamMemberTab({ const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id); const enhancedMember = { ...record, - max_budget_in_team: membership?.litellm_budget_table?.max_budget || null, - tpm_limit: membership?.litellm_budget_table?.tpm_limit || null, - rpm_limit: membership?.litellm_budget_table?.rpm_limit || null, + max_budget_in_team: membership?.litellm_budget_table?.max_budget ?? null, + tpm_limit: membership?.litellm_budget_table?.tpm_limit ?? null, + rpm_limit: membership?.litellm_budget_table?.rpm_limit ?? null, budget_duration: membership?.litellm_budget_table?.budget_duration || null, allowed_models: membership?.litellm_budget_table?.allowed_models || [], }; diff --git a/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts b/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts index 398719c6197..7a1fcd85a80 100644 --- a/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts +++ b/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts @@ -82,7 +82,7 @@ describe("buildMemberFormValues", () => { }); }); - it("collapses falsy budgets and limits to null and a missing model list to an empty array", () => { + it("keeps a stored budget or limit of 0 as 0 because only null means unlimited", () => { expect( buildMemberFormValues( "edit", @@ -90,6 +90,19 @@ describe("buildMemberFormValues", () => { teamConfig, ), ).toStrictEqual({ + user_email: "a@b.com", + user_id: "u1", + role: "user", + max_budget_in_team: 0, + budget_duration: null, + tpm_limit: 0, + rpm_limit: 0, + allowed_models: [], + }); + }); + + it("collapses missing budgets and limits to null and a missing model list to an empty array", () => { + const unlimitedMember = { user_email: "a@b.com", user_id: "u1", role: "user", @@ -98,7 +111,10 @@ describe("buildMemberFormValues", () => { tpm_limit: null, rpm_limit: null, allowed_models: [], - }); + }; + expect( + buildMemberFormValues("edit", { user_email: "a@b.com", user_id: "u1", role: "user" }, teamConfig), + ).toStrictEqual(unlimitedMember); }); it("falls back to the configured default role when the member has none", () => { diff --git a/ui/litellm-dashboard/src/components/team/memberFormValues.ts b/ui/litellm-dashboard/src/components/team/memberFormValues.ts index 942b9b665fe..51b8fac71c1 100644 --- a/ui/litellm-dashboard/src/components/team/memberFormValues.ts +++ b/ui/litellm-dashboard/src/components/team/memberFormValues.ts @@ -43,9 +43,9 @@ export const buildMemberFormValues = ( const seeded: MemberFormValues = { ...initialData, role: (initialData.role as string) || config.defaultRole, - max_budget_in_team: initialData.max_budget_in_team || null, - tpm_limit: initialData.tpm_limit || null, - rpm_limit: initialData.rpm_limit || null, + max_budget_in_team: initialData.max_budget_in_team ?? null, + tpm_limit: initialData.tpm_limit ?? null, + rpm_limit: initialData.rpm_limit ?? null, budget_duration: initialData.budget_duration || null, allowed_models: initialData.allowed_models || [], }; From 7113685a76fbd77f35c031b9dd165427afe9e674 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:12:55 -0700 Subject: [PATCH 074/156] fix(ui): repoint the key detail URL to the rotated hash after regenerating (#37968) Regenerating a key from the key info page left the ?key= query param on the old hash, so dismissing the dialog or reloading landed on a key that no longer exists and the page rendered "Key not found". Two defects had to line up. POST /key/{key}/regenerate returns the rotated hash in token_id and leaves token null, but RegenerateKeyModal read response.token || response.key_id, neither of which the endpoint populates, so it always reported the old hash back to its parent. And KeyInfoView's onKeyDataUpdate prop had no caller anywhere in the tree: VirtualKeysTable owns the ?key= param and mounts the view but never passed it, so even a correct hash went nowhere. VirtualKeysTable now handles the update by pointing ?key= at the rotated hash and refetching. KeyInfoView holds that callback until the regenerate dialog is dismissed rather than firing it on the API response, because swapping the selected key mid-dialog unmounts the view and tears down the one-time plaintext key before the user can copy it. --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 28 +++++++++++++++++++ .../VirtualKeysPage/VirtualKeysTable.tsx | 13 ++++++++- .../organisms/RegenerateKeyModal.test.tsx | 17 +++++++++++ .../organisms/RegenerateKeyModal.tsx | 2 +- .../components/templates/key_info_view.tsx | 26 +++++++++++------ 5 files changed, 76 insertions(+), 10 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 9ef849dbeb2..643da71e8b1 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -8,6 +8,7 @@ import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; +import { regenerateKeyCall } from "../networking"; // Resolve debounced values synchronously so an applied filter lands in the useKeys query within the test tick. vi.mock("@tanstack/react-pacer/debouncer", async () => { @@ -25,6 +26,11 @@ vi.mock("@tanstack/react-pacer/debouncer", async () => { vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); +vi.mock("../networking", async (importOriginal) => ({ + ...(await importOriginal()), + regenerateKeyCall: vi.fn(), +})); + vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: vi.fn(() => ({ accessToken: "test-token", @@ -389,6 +395,28 @@ it("renders KeyInfoView when the URL has ?key= for a key on the current page, wi expect(screen.getByTestId("pagination-range")).toBeInTheDocument(); }); +it("repoints ?key= to the rotated hash once the regenerate dialog is dismissed", async () => { + const user = userEvent.setup(); + vi.mocked(regenerateKeyCall).mockResolvedValue({ + key: "sk-rotated-plaintext", + token: null, + token_id: "rotated-hash-456", + }); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: { key: mockKey.token }, onUrlUpdate }); + + await user.click(await screen.findByRole("button", { name: /regenerate key/i })); + await user.click(await screen.findByRole("button", { name: /^Regenerate$/ })); + expect(await screen.findAllByText("sk-rotated-plaintext")).not.toHaveLength(0); + expect(lastKeyParam(onUrlUpdate)).toBeUndefined(); + + await user.click(screen.getAllByRole("button", { name: "Close" })[0]); + + await waitFor(() => { + expect(lastKeyParam(onUrlUpdate)).toBe("rotated-hash-456"); + }); +}); + it("fetches the key by id when the URL has ?key= for a key not in the loaded page", async () => { mockUseKeyInfo.mockReturnValue( keyInfoResult({ ...mockKey, token: "other-key-hash", key_alias: "Fetched Key Alias" }), diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index b278b4f675d..fbdcf2d8d83 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -20,7 +20,7 @@ import { KeyRound } from "lucide-react"; import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; -import { Team } from "../key_team_helpers/key_list"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; import KeyInfoView from "../templates/key_info_view"; import { getKeyTableColumns, KEY_TABLE_HIDDEN_COLUMNS } from "./keyTableColumns"; @@ -139,6 +139,16 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { [organizations], ); + const handleSelectedKeyDataUpdate = useCallback( + (updated: Partial) => { + const rotatedToken = updated.token ?? updated.token_id; + if (!rotatedToken || rotatedToken === selectedKeyId) return; + void setSelectedKeyId(rotatedToken); + void refetch(); + }, + [refetch, selectedKeyId, setSelectedKeyId], + ); + const formatFilterValue = useCallback( (columnId: string, value: unknown): string => { const raw = String(value); @@ -165,6 +175,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { keyData={selectedKey} teams={allTeams} onDelete={refetch} + onKeyDataUpdate={handleSelectedKeyDataUpdate} />
); diff --git a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx index 180f9d1f265..262ce9669f2 100644 --- a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx @@ -427,4 +427,21 @@ describe("RegenerateKeyModal", () => { ); }); }); + + it("should report the rotated hash from token_id when the API leaves token null", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: null, + token_id: "rotated-hash-456", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(mockOnKeyUpdate).toHaveBeenCalledOnce(); + }); + expect(mockOnKeyUpdate.mock.calls[0][0].token).toBe("rotated-hash-456"); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx index d85e6cb5f91..a2658985f32 100644 --- a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx @@ -107,7 +107,7 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat // formatted preview, otherwise downstream expiry parsing breaks. const updatedKeyData: Partial = { ...response, - token: response.token || response.key_id || selectedToken.token, + token: response.token_id || response.token || selectedToken.token, key_name: response.key, max_budget: formValues.max_budget, tpm_limit: formValues.tpm_limit, diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 5cfa737397b..59950303c45 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -100,6 +100,9 @@ export default function KeyInfoView({ // Add local state to maintain key data and track regeneration const [currentKeyData, setCurrentKeyData] = useState(keyData); const [lastRegeneratedAt, setLastRegeneratedAt] = useState(null); + const [keyDataUpdateHeldUntilModalClose, setKeyDataUpdateHeldUntilModalClose] = useState | null>( + null, + ); const [isRecentlyRegenerated, setIsRecentlyRegenerated] = useState(false); const [policyGuardrails, setPolicyGuardrails] = useState>({}); const [loadingPolicies, setLoadingPolicies] = useState(false); @@ -352,6 +355,7 @@ export default function KeyInfoView({ }; const handleRegenerateKeyUpdate = (updatedKeyData: Partial) => { + const regeneratedAt = new Date(); // Update local state immediately with ALL the new data setCurrentKeyData((prevData) => { if (!prevData) return undefined; @@ -359,20 +363,26 @@ export default function KeyInfoView({ ...prevData, ...updatedKeyData, // This should include the new token (key-id) // Update the created_at to show when it was regenerated - created_at: new Date().toLocaleString(), + created_at: regeneratedAt.toLocaleString(), }; return newData; }); // Track regeneration timestamp - setLastRegeneratedAt(new Date()); + setLastRegeneratedAt(regeneratedAt); setIsRecentlyRegenerated(true); - if (onKeyDataUpdate) { - onKeyDataUpdate({ - ...updatedKeyData, - created_at: new Date().toLocaleString(), - }); + setKeyDataUpdateHeldUntilModalClose({ + ...updatedKeyData, + created_at: regeneratedAt.toLocaleString(), + }); + }; + + const handleRegenerateModalClose = () => { + setIsRegenerateModalOpen(false); + if (keyDataUpdateHeldUntilModalClose) { + setKeyDataUpdateHeldUntilModalClose(null); + onKeyDataUpdate?.(keyDataUpdateHeldUntilModalClose); } }; @@ -506,7 +516,7 @@ export default function KeyInfoView({ setIsRegenerateModalOpen(false)} + onClose={handleRegenerateModalClose} onKeyUpdate={handleRegenerateKeyUpdate} /> From 3fb1009f816fbe96e90c6cc14d4a2ea36463f5eb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 24 Aug 2026 10:13:24 -0700 Subject: [PATCH 075/156] fix(ui): make playground chat bubbles theme-aware (#37978) The playground message bubble painted its fill, border and avatar circle from inline hex values, so in dark mode both bubbles stayed near-white while the text inherited the dark foreground: the message body was unreadable. The MCP-events placeholder bubble in ChatUI carried the same three fills. They move onto the tokens the rest of the sweep already uses, so the assistant surface is bg-card over border-border and the user surface is the info tint at the same weight the other selected-state surfaces take. Light mode keeps the same colour family it had. The regression test asserts the token classes and that no inline style survives on either surface, which is the exact shape the bug took. --- .../chat_ui/ChatMessageBubble.test.tsx | 15 +++++++++++++++ .../components/chat_ui/ChatMessageBubble.tsx | 16 ++++++---------- .../playground/components/chat_ui/ChatUI.tsx | 14 ++------------ 3 files changed, 23 insertions(+), 22 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index c218db5b914..a83c11d1444 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -87,6 +87,21 @@ describe("ChatMessageBubble", () => { expect(screen.getByText("Hi there")).toBeInTheDocument(); }); + it.each([ + { role: "user" as const, bubble: ["bg-info/10", "border-info/20"], avatar: "bg-info/20" }, + { role: "assistant" as const, bubble: ["bg-card", "border-border"], avatar: "bg-muted" }, + ])("should paint the $role surface from theme tokens, not fixed colours", ({ role, bubble, avatar }) => { + render(); + + const header = screen.getByText(role).closest("div") as HTMLElement; + const surface = header.parentElement as HTMLElement; + + expect(surface).toHaveClass(...bubble); + expect(surface).not.toHaveAttribute("style"); + expect(header.firstElementChild).toHaveClass(avatar); + expect(header.firstElementChild).not.toHaveAttribute("style"); + }); + it("should show model badge for assistant messages when model is provided", () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index eb24463cfd3..8c54d9e89fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -46,20 +46,16 @@ function ChatMessageBubble({ return (
{/* Header: role icon + name + model badge */}
{isUser ? (