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/273] 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/273] 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/273] 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/273] 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/273] 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/273] 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/273] 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/273] 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 e4c2ad4627b71d280603f721234ec8990f3aa6bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:23:34 -0700 Subject: [PATCH 009/273] fix(anthropic): buffer streamed responses carrying server-fulfilled tools so retrieval tool calls never reach the client --- .../compression_interception/handler.py | 4 +- litellm/integrations/custom_logger.py | 4 +- .../messages/agentic_streaming_iterator.py | 59 ++++++ litellm/llms/custom_httpx/llm_http_handler.py | 23 +++ .../guardrail_hooks/headroom/headroom.py | 1 + .../test_agentic_streaming_iterator.py | 178 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 71 +++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 + 8 files changed, 345 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 7ea60053e6f..76720682101 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Final, cast +from typing import Any, ClassVar, Final, cast from litellm._logging import verbose_logger from litellm.compression import compress @@ -72,6 +72,8 @@ class CompressionInterceptionLogger(CustomLogger): 4. Build typed rerun plan with tool_result blocks from the compressed cache. """ + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({LITELLM_CONTENT_RETRIEVE_TOOL_NAME}) + def __init__( self, enabled: bool = True, diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a0c78674ac8..60af4063f84 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -3,7 +3,7 @@ import re import traceback from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional from pydantic import BaseModel @@ -60,6 +60,8 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() + def __init__( self, turn_off_message_logging: bool = False, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 5c4fa4700c0..0699595821e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -6,14 +6,27 @@ yields every chunk to the caller (preserving real streaming), collects all bytes, and on stream exhaustion rebuilds the full Anthropic response to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. + +In hold-back mode (``hold_back=True``), chunks are buffered instead of +yielded live, with SSE ping events emitted while the upstream message is +in flight. On exhaustion the hooks run first: if a follow-up response +replaces the message, only the follow-up is yielded and the buffered +message is dropped; otherwise the buffer is replayed verbatim. This is +required for server-fulfilled tools (e.g. ``headroom_retrieve``), whose +tool_use blocks must never reach a client that cannot execute them. """ +import asyncio +import contextlib import json from collections.abc import AsyncIterator from typing import Any, Final, cast from litellm._logging import verbose_logger +PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' +HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 + # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- @@ -156,6 +169,8 @@ class AgenticAnthropicStreamingIterator: logging_obj: Any, custom_llm_provider: str, kwargs: dict, + hold_back: bool = False, + ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS, ): self._inner = completion_stream.__aiter__() self._http_handler = http_handler @@ -166,16 +181,23 @@ class AgenticAnthropicStreamingIterator: self._logging_obj = logging_obj self._custom_llm_provider = custom_llm_provider self._kwargs = kwargs + self._hold_back = hold_back + self._ping_interval_seconds = ping_interval_seconds self._collected_bytes: list[bytes] = [] self._stream_exhausted = False self._hook_processing_done = False self._follow_up_iterator: AsyncIterator | None = None + self._drain_task: asyncio.Task | None = None + self._replay_index = 0 def __aiter__(self): return self async def __anext__(self) -> bytes: + if self._hold_back: + return await self._anext_held_back() + # Phase 1: yield from upstream, collect bytes if not self._stream_exhausted: try: @@ -194,11 +216,48 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + async def _drain_upstream(self) -> None: + try: + while True: + self._collected_bytes.append(await self._inner.__anext__()) + except StopAsyncIteration: + return + + async def _anext_held_back(self) -> bytes: + if self._drain_task is None: + self._drain_task = asyncio.create_task(self._drain_upstream()) + return PING_SSE_BYTES + + while not self._stream_exhausted: + try: + await asyncio.wait_for(asyncio.shield(self._drain_task), timeout=self._ping_interval_seconds) + except asyncio.TimeoutError: + return PING_SSE_BYTES + self._stream_exhausted = True + await self._process_agentic_hooks() + + if self._follow_up_iterator is not None: + return await self._follow_up_iterator.__anext__() + + if self._replay_index < len(self._collected_bytes): + chunk: Final = self._collected_bytes[self._replay_index] + self._replay_index += 1 + return chunk + + raise StopAsyncIteration + async def aclose(self) -> None: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, ) + if self._drain_task is not None and self._drain_task.done(): + if not self._drain_task.cancelled(): + self._drain_task.exception() + elif self._drain_task is not None: + self._drain_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._drain_task await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a58397c9184..913cedcfa55 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2189,6 +2189,10 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, + hold_back=self._should_hold_back_stream( + logging_obj=logging_obj, + tools=anthropic_messages_optional_request_params.get("tools"), + ), ) return AnthropicMessagesStreamingResponse( completion_stream=initial_response, @@ -5033,6 +5037,25 @@ class BaseLLMHTTPHandler: return True return False + @staticmethod + def _should_hold_back_stream(logging_obj: LiteLLMLoggingObj, tools: object) -> bool: + """ + True when the request carries a tool that a registered callback fulfills + server-side (e.g. ``headroom_retrieve``). The model's tool_use for such a + tool must never reach the client, which cannot execute it: the agentic + loop replaces the whole message with a follow-up response, so the stream + is buffered (with ping keepalives) instead of forwarded live. + """ + if not isinstance(tools, list) or not tools: + return False + from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name + + return any( + has_tool_with_name(tools, name) + for cb in _custom_logger_callbacks(logging_obj) + for name in getattr(cb, "server_fulfilled_tool_names", frozenset()) + ) + @staticmethod def _check_agentic_loop_safety( tool_calls: object, diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 8bfd5cca58a..84c6b220e62 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -339,6 +339,7 @@ def _build_responses_followup_items( class HeadroomGuardrail(CustomGuardrail): records_own_guardrail_information: ClassVar[bool] = True + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({HEADROOM_RETRIEVE_TOOL_NAME}) @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index b9bda07336f..a6b071bbab5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -2,6 +2,7 @@ Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers. """ +import asyncio import json import os import sys @@ -13,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + PING_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, _handle_content_block_start, @@ -230,6 +232,51 @@ class MockAsyncStream: return chunk +class MockSlowAsyncStream(MockAsyncStream): + """Async iterator that sleeps before every chunk.""" + + def __init__(self, chunks: List[bytes], delay_seconds: float): + super().__init__(chunks) + self._delay_seconds = delay_seconds + + async def __anext__(self) -> bytes: + await asyncio.sleep(self._delay_seconds) + return await super().__anext__() + + +class MockFailingAsyncStream(MockAsyncStream): + """Async iterator that raises after yielding its chunks.""" + + def __init__(self, chunks: List[bytes], error: Exception): + super().__init__(chunks) + self._error = error + + async def __anext__(self) -> bytes: + if self._idx >= len(self._chunks): + raise self._error + return await super().__anext__() + + +def _build_hold_back_iterator( + stream: MockAsyncStream, + mock_handler: MagicMock, + ping_interval_seconds: float = 15.0, +) -> AgenticAnthropicStreamingIterator: + return AgenticAnthropicStreamingIterator( + completion_stream=stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + ping_interval_seconds=ping_interval_seconds, + ) + + # --------------------------------------------------------------------------- # Tests for _parse_sse_events # --------------------------------------------------------------------------- @@ -790,3 +837,134 @@ class TestAgenticStreamingIteratorErrorHandling: call_kwargs = mock_handler._call_agentic_completion_hooks.call_args assert call_kwargs.kwargs["stream"] is True + + +class TestAgenticStreamingIteratorHoldBack: + @pytest.mark.asyncio + async def test_should_not_leak_intercepted_message_when_follow_up_fires(self): + """The buffered tool_use message must be dropped: only pings and follow-up bytes reach the client.""" + phase1_chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=MockAsyncStream(phase2_chunks)) + + iterator = _build_hold_back_iterator(MockAsyncStream(phase1_chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + non_ping = [c for c in collected if c != PING_SSE_BYTES] + assert non_ping == phase2_chunks + assert b"litellm_content_retrieve" not in b"".join(collected) + assert collected[0] == PING_SSE_BYTES + + @pytest.mark.asyncio + async def test_should_replay_buffer_verbatim_when_no_hook_fires(self): + """Without interception the buffered message is replayed byte-identical after the pings.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + mock_handler._call_agentic_completion_hooks.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_emit_pings_while_upstream_is_slow(self): + """Pings keep the client connection alive while the upstream message is buffered.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockSlowAsyncStream(chunks, delay_seconds=0.05), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected.count(PING_SSE_BYTES) >= 2 + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + + @pytest.mark.asyncio + async def test_should_propagate_upstream_error_instead_of_partial_message(self): + """An upstream failure surfaces as an error; the client never receives a truncated message.""" + chunks = _build_simple_text_stream()[:2] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockFailingAsyncStream(chunks, RuntimeError("upstream died")), + mock_handler, + ) + + collected = [] + with pytest.raises(RuntimeError, match="upstream died"): + async for chunk in iterator: + collected.append(chunk) + + assert all(c == PING_SSE_BYTES for c in collected) + mock_handler._call_agentic_completion_hooks.assert_not_awaited() + + @pytest.mark.asyncio + async def test_should_replay_buffer_when_hook_processing_errors(self): + """A hook crash degrades to replaying the original message rather than dropping it.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded")) + + mock_logging = MagicMock() + mock_logging.litellm_call_id = "test_call_holdback" + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=MockAsyncStream(chunks), + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging, + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + + @pytest.mark.asyncio + async def test_aclose_cancels_drain_task(self): + """Closing the iterator mid-buffer must cancel the background drain task.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockSlowAsyncStream(chunks, delay_seconds=5.0), + mock_handler, + ) + + first = await iterator.__anext__() + assert first == PING_SSE_BYTES + assert iterator._drain_task is not None + + await iterator.aclose() + assert iterator._drain_task.cancelled() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..6c2727e7da2 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2071,3 +2071,74 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques retry_authorization = posts[1]["headers"]["Authorization"] assert retry_authorization.startswith("AWS4-HMAC-SHA256") assert retry_authorization != first_attempt_headers["Authorization"] + + +class TestShouldHoldBackStream: + """_should_hold_back_stream gates the buffered (non-leaking) streaming mode + for server-fulfilled tools like headroom_retrieve.""" + + @staticmethod + def _logging_obj_with(callbacks): + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = callbacks + return logging_obj + + def test_should_hold_back_when_callback_owns_tool_in_request(self): + from litellm.integrations.custom_logger import CustomLogger + + class RetrievalCallback(CustomLogger): + server_fulfilled_tool_names = frozenset({"headroom_retrieve"}) + + tools = [ + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "headroom_retrieve", "input_schema": {"type": "object"}}, + ] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) + is True + ) + + def test_should_stream_live_when_tool_absent_from_request(self): + from litellm.integrations.custom_logger import CustomLogger + + class RetrievalCallback(CustomLogger): + server_fulfilled_tool_names = frozenset({"headroom_retrieve"}) + + tools = [{"name": "Bash", "input_schema": {"type": "object"}}] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) + is False + ) + + def test_should_stream_live_when_no_callback_declares_tool_names(self): + from litellm.integrations.custom_logger import CustomLogger + + tools = [{"name": "headroom_retrieve", "input_schema": {"type": "object"}}] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([CustomLogger()]), tools=tools + ) + is False + ) + + def test_should_stream_live_without_tools(self): + assert BaseLLMHTTPHandler._should_hold_back_stream(logging_obj=self._logging_obj_with([]), tools=None) is False + + def test_interception_callbacks_declare_their_retrieval_tools(self): + from litellm.integrations.compression_interception.handler import ( + LITELLM_CONTENT_RETRIEVE_TOOL_NAME, + CompressionInterceptionLogger, + ) + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( + HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + ) + + assert HeadroomGuardrail.server_fulfilled_tool_names == frozenset({HEADROOM_RETRIEVE_TOOL_NAME}) + assert CompressionInterceptionLogger.server_fulfilled_tool_names == frozenset( + {LITELLM_CONTENT_RETRIEVE_TOOL_NAME} + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f1660e77ad9..8e950874a10 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21391,6 +21391,13 @@ export interface components { * @description What the routed traffic actually cost */ spend: number; + /** + * Tier Turns + * @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns + */ + tier_turns?: { + [key: string]: number; + }; /** Turns */ turns: number; }; From f994068a7338b4bb54fa7a53d76fb009dd3e9f6a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 07:30:24 +0000 Subject: [PATCH 010/273] fix(anthropic): keep pinging during agentic hooks and fail instead of replaying server-fulfilled tool_use Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 87 ++++++++++++--- litellm/llms/custom_httpx/llm_http_handler.py | 29 ++--- .../test_agentic_streaming_iterator.py | 105 +++++++++++++++--- .../custom_httpx/test_llm_http_handler.py | 28 ++--- 4 files changed, 190 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 0699595821e..4cf348dda9e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -9,11 +9,13 @@ follow-up response is chained as Phase 2 of the same iterator. In hold-back mode (``hold_back=True``), chunks are buffered instead of yielded live, with SSE ping events emitted while the upstream message is -in flight. On exhaustion the hooks run first: if a follow-up response -replaces the message, only the follow-up is yielded and the buffered -message is dropped; otherwise the buffer is replayed verbatim. This is -required for server-fulfilled tools (e.g. ``headroom_retrieve``), whose -tool_use blocks must never reach a client that cannot execute them. +in flight and while the agentic hooks run. On exhaustion the hooks run +first: if a follow-up response replaces the message, only the follow-up +is yielded and the buffered message is dropped; otherwise the buffer is +replayed verbatim, unless it holds a tool_use for a server-fulfilled tool +(e.g. ``headroom_retrieve``), in which case an SSE ``error`` event is +emitted because such a block must never reach a client that cannot +execute it. """ import asyncio @@ -26,6 +28,11 @@ from litellm._logging import verbose_logger PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 +SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "api_error", "message": ' + b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' +) # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) @@ -170,6 +177,7 @@ class AgenticAnthropicStreamingIterator: custom_llm_provider: str, kwargs: dict, hold_back: bool = False, + server_fulfilled_tool_names: frozenset[str] = frozenset(), ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS, ): self._inner = completion_stream.__aiter__() @@ -182,6 +190,7 @@ class AgenticAnthropicStreamingIterator: self._custom_llm_provider = custom_llm_provider self._kwargs = kwargs self._hold_back = hold_back + self._server_fulfilled_tool_names = server_fulfilled_tool_names self._ping_interval_seconds = ping_interval_seconds self._collected_bytes: list[bytes] = [] @@ -189,7 +198,9 @@ class AgenticAnthropicStreamingIterator: self._hook_processing_done = False self._follow_up_iterator: AsyncIterator | None = None self._drain_task: asyncio.Task | None = None + self._hook_task: asyncio.Task | None = None self._replay_index = 0 + self._error_emitted = False def __aiter__(self): return self @@ -223,22 +234,42 @@ class AgenticAnthropicStreamingIterator: except StopAsyncIteration: return + async def _completed_within_ping_interval(self, task: asyncio.Task) -> bool: + try: + await asyncio.wait_for(asyncio.shield(task), timeout=self._ping_interval_seconds) + except asyncio.TimeoutError: + return False + return True + async def _anext_held_back(self) -> bytes: if self._drain_task is None: self._drain_task = asyncio.create_task(self._drain_upstream()) return PING_SSE_BYTES - while not self._stream_exhausted: - try: - await asyncio.wait_for(asyncio.shield(self._drain_task), timeout=self._ping_interval_seconds) - except asyncio.TimeoutError: + if not self._stream_exhausted: + if not await self._completed_within_ping_interval(self._drain_task): return PING_SSE_BYTES self._stream_exhausted = True - await self._process_agentic_hooks() + + if self._hook_task is None: + self._hook_task = asyncio.create_task(self._process_agentic_hooks()) + if not await self._completed_within_ping_interval(self._hook_task): + return PING_SSE_BYTES if self._follow_up_iterator is not None: return await self._follow_up_iterator.__anext__() + if self._buffer_holds_server_fulfilled_tool_use(): + if self._error_emitted: + raise StopAsyncIteration + self._error_emitted = True + verbose_logger.error( + "AgenticStreamingIterator: hooks did not replace a message containing a server-fulfilled " + "tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client", + self._model, + ) + return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES + if self._replay_index < len(self._collected_bytes): chunk: Final = self._collected_bytes[self._replay_index] self._replay_index += 1 @@ -246,18 +277,40 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + def _buffer_holds_server_fulfilled_tool_use(self) -> bool: + if not self._server_fulfilled_tool_names: + return False + started_blocks: Final = ( + data.get("content_block") + for event_type, data in _parse_sse_events(b"".join(self._collected_bytes)) + if event_type == "content_block_start" + ) + return any( + isinstance(block, dict) + and block.get("type") == "tool_use" + and block.get("name") in self._server_fulfilled_tool_names + for block in started_blocks + ) + + @staticmethod + async def _settle_task(task: asyncio.Task | None) -> None: + if task is None: + return + if task.done(): + if not task.cancelled(): + task.exception() + return + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + async def aclose(self) -> None: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, ) - if self._drain_task is not None and self._drain_task.done(): - if not self._drain_task.cancelled(): - self._drain_task.exception() - elif self._drain_task is not None: - self._drain_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await self._drain_task + await self._settle_task(self._drain_task) + await self._settle_task(self._hook_task) await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 913cedcfa55..e9b88e45219 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2179,6 +2179,10 @@ class BaseLLMHTTPHandler: AgenticAnthropicStreamingIterator, ) + held_back_tool_names: Final = self._server_fulfilled_tools_in_request( + logging_obj=logging_obj, + tools=anthropic_messages_optional_request_params.get("tools"), + ) initial_response = AgenticAnthropicStreamingIterator( completion_stream=completion_stream, http_handler=self, @@ -2189,10 +2193,8 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, - hold_back=self._should_hold_back_stream( - logging_obj=logging_obj, - tools=anthropic_messages_optional_request_params.get("tools"), - ), + hold_back=bool(held_back_tool_names), + server_fulfilled_tool_names=held_back_tool_names, ) return AnthropicMessagesStreamingResponse( completion_stream=initial_response, @@ -5038,22 +5040,23 @@ class BaseLLMHTTPHandler: return False @staticmethod - def _should_hold_back_stream(logging_obj: LiteLLMLoggingObj, tools: object) -> bool: + def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]: """ - True when the request carries a tool that a registered callback fulfills - server-side (e.g. ``headroom_retrieve``). The model's tool_use for such a - tool must never reach the client, which cannot execute it: the agentic - loop replaces the whole message with a follow-up response, so the stream - is buffered (with ping keepalives) instead of forwarded live. + The request's tools that a registered callback fulfills server-side (e.g. + ``headroom_retrieve``). The model's tool_use for such a tool must never + reach the client, which cannot execute it: the agentic loop replaces the + whole message with a follow-up response, so a stream carrying any of + these is buffered (with ping keepalives) instead of forwarded live. """ if not isinstance(tools, list) or not tools: - return False + return frozenset() from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name - return any( - has_tool_with_name(tools, name) + return frozenset( + name for cb in _custom_logger_callbacks(logging_obj) for name in getattr(cb, "server_fulfilled_tool_names", frozenset()) + if has_tool_with_name(tools, name) ) @staticmethod diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index a6b071bbab5..d59f76232cd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -15,6 +15,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( PING_SSE_BYTES, + SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, _handle_content_block_start, @@ -261,6 +262,7 @@ def _build_hold_back_iterator( stream: MockAsyncStream, mock_handler: MagicMock, ping_interval_seconds: float = 15.0, + server_fulfilled_tool_names: frozenset = frozenset({"litellm_content_retrieve"}), ) -> AgenticAnthropicStreamingIterator: return AgenticAnthropicStreamingIterator( completion_stream=stream, @@ -273,6 +275,7 @@ def _build_hold_back_iterator( custom_llm_provider="anthropic", kwargs={}, hold_back=True, + server_fulfilled_tool_names=server_fulfilled_tool_names, ping_interval_seconds=ping_interval_seconds, ) @@ -920,27 +923,76 @@ class TestAgenticStreamingIteratorHoldBack: mock_handler._call_agentic_completion_hooks.assert_not_awaited() @pytest.mark.asyncio - async def test_should_replay_buffer_when_hook_processing_errors(self): - """A hook crash degrades to replaying the original message rather than dropping it.""" + async def test_should_emit_pings_while_hooks_are_slow(self): + """Retrieval and follow-up generation can outlast a client's idle timeout, so hooks get keepalives too.""" + chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk"] + + async def slow_hooks(**_kwargs): + await asyncio.sleep(0.12) + return MockAsyncStream(phase2_chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=slow_hooks) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected.count(PING_SSE_BYTES) >= 4 + assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_error_instead_of_replaying_server_fulfilled_tool_use_when_hook_crashes(self): + """A hook crash must not replay the buffered retrieval tool_use: that is the unknown-tool bug.""" chunks = _build_tool_use_stream() mock_handler = MagicMock() mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded")) - mock_logging = MagicMock() - mock_logging.litellm_call_id = "test_call_holdback" + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) - iterator = AgenticAnthropicStreamingIterator( - completion_stream=MockAsyncStream(chunks), - http_handler=mock_handler, - model="claude-sonnet-4-20250514", - messages=[], - anthropic_messages_provider_config=MagicMock(), - anthropic_messages_optional_request_params={}, - logging_obj=mock_logging, - custom_llm_provider="anthropic", - kwargs={}, - hold_back=True, + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert b"litellm_content_retrieve" not in b"".join(collected) + + @pytest.mark.asyncio + async def test_should_error_instead_of_replaying_when_no_hook_fires_on_tool_use(self): + """Hooks returning None on a retrieval tool_use is still a leak, so the turn fails loudly.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + + @pytest.mark.asyncio + async def test_should_replay_client_owned_tool_use_verbatim(self): + """Only server-fulfilled tools are withheld: a client's own tool_use still reaches it byte-identical.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + server_fulfilled_tool_names=frozenset({"headroom_retrieve"}), ) collected = [] @@ -968,3 +1020,26 @@ class TestAgenticStreamingIteratorHoldBack: await iterator.aclose() assert iterator._drain_task.cancelled() + + @pytest.mark.asyncio + async def test_aclose_cancels_in_flight_hook_task(self): + """Closing while hooks are running must not leave the retrieval follow-up task orphaned.""" + chunks = _build_tool_use_stream() + + async def never_finishing_hooks(**_kwargs): + await asyncio.sleep(5.0) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=never_finishing_hooks) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + while iterator._hook_task is None: + await iterator.__anext__() + + await iterator.aclose() + assert iterator._hook_task.cancelled() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 6c2727e7da2..798fb4c92e2 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2073,9 +2073,9 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques assert retry_authorization != first_attempt_headers["Authorization"] -class TestShouldHoldBackStream: - """_should_hold_back_stream gates the buffered (non-leaking) streaming mode - for server-fulfilled tools like headroom_retrieve.""" +class TestServerFulfilledToolsInRequest: + """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming + mode for server-fulfilled tools like headroom_retrieve.""" @staticmethod def _logging_obj_with(callbacks): @@ -2093,12 +2093,9 @@ class TestShouldHoldBackStream: {"name": "Bash", "input_schema": {"type": "object"}}, {"name": "headroom_retrieve", "input_schema": {"type": "object"}}, ] - assert ( - BaseLLMHTTPHandler._should_hold_back_stream( - logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools - ) - is True - ) + assert BaseLLMHTTPHandler._server_fulfilled_tools_in_request( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) == frozenset({"headroom_retrieve"}) def test_should_stream_live_when_tool_absent_from_request(self): from litellm.integrations.custom_logger import CustomLogger @@ -2108,10 +2105,10 @@ class TestShouldHoldBackStream: tools = [{"name": "Bash", "input_schema": {"type": "object"}}] assert ( - BaseLLMHTTPHandler._should_hold_back_stream( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request( logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools ) - is False + == frozenset() ) def test_should_stream_live_when_no_callback_declares_tool_names(self): @@ -2119,14 +2116,17 @@ class TestShouldHoldBackStream: tools = [{"name": "headroom_retrieve", "input_schema": {"type": "object"}}] assert ( - BaseLLMHTTPHandler._should_hold_back_stream( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request( logging_obj=self._logging_obj_with([CustomLogger()]), tools=tools ) - is False + == frozenset() ) def test_should_stream_live_without_tools(self): - assert BaseLLMHTTPHandler._should_hold_back_stream(logging_obj=self._logging_obj_with([]), tools=None) is False + assert ( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request(logging_obj=self._logging_obj_with([]), tools=None) + == frozenset() + ) def test_interception_callbacks_declare_their_retrieval_tools(self): from litellm.integrations.compression_interception.handler import ( From 398e3d214cc97be0531428f3fb506a7cc42e2683 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 19:28:40 +0000 Subject: [PATCH 011/273] refactor(anthropic): trim hold-back commentary and drop dead rebuilt-content expression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 17 ++++------------- litellm/llms/custom_httpx/llm_http_handler.py | 8 +------- 2 files changed, 5 insertions(+), 20 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 4cf348dda9e..a88b148e92c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -7,14 +7,10 @@ all bytes, and on stream exhaustion rebuilds the full Anthropic response to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. -In hold-back mode (``hold_back=True``), chunks are buffered instead of -yielded live, with SSE ping events emitted while the upstream message is -in flight and while the agentic hooks run. On exhaustion the hooks run -first: if a follow-up response replaces the message, only the follow-up -is yielded and the buffered message is dropped; otherwise the buffer is -replayed verbatim, unless it holds a tool_use for a server-fulfilled tool -(e.g. ``headroom_retrieve``), in which case an SSE ``error`` event is -emitted because such a block must never reach a client that cannot +In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded +live, keepalive pings run until the hooks finish, and then either the follow-up +replaces the message or the buffer replays, except that a buffered tool_use for +a server-fulfilled tool fails the turn rather than reaching a client that cannot execute it. """ @@ -329,11 +325,6 @@ class AgenticAnthropicStreamingIterator: verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes") return - [ - (f"{b.get('type')}({b.get('name', '')})" if b.get("type") == "tool_use" else b.get("type")) - for b in rebuilt.get("content", []) - ] - result: Final = await self._http_handler._call_agentic_completion_hooks( response=rebuilt, model=self._model, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e9b88e45219..193a38a5404 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5041,13 +5041,7 @@ class BaseLLMHTTPHandler: @staticmethod def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]: - """ - The request's tools that a registered callback fulfills server-side (e.g. - ``headroom_retrieve``). The model's tool_use for such a tool must never - reach the client, which cannot execute it: the agentic loop replaces the - whole message with a follow-up response, so a stream carrying any of - these is buffered (with ping keepalives) instead of forwarded live. - """ + """The request's tools that a registered callback fulfills server-side (e.g. ``headroom_retrieve``).""" if not isinstance(tools, list) or not tools: return frozenset() from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name From cbefb1ce5ffbdc90c9e6691206752b40ec9ef6e0 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 19:41:42 +0000 Subject: [PATCH 012/273] fix(anthropic): keep pinging while the held-back follow-up stream is in flight Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 27 ++++++++- .../test_agentic_streaming_iterator.py | 59 +++++++++++++++++++ 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index a88b148e92c..3aaba0b139c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -8,8 +8,8 @@ to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded -live, keepalive pings run until the hooks finish, and then either the follow-up -replaces the message or the buffer replays, except that a buffered tool_use for +live, keepalive pings run whenever no other byte is ready, and then either the +follow-up replaces the message or the buffer replays, except that a tool_use for a server-fulfilled tool fails the turn rather than reaching a client that cannot execute it. """ @@ -30,6 +30,14 @@ SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' ) + +async def _anext_or_none(iterator: AsyncIterator) -> bytes | None: + try: + return await iterator.__anext__() + except StopAsyncIteration: + return None + + # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- @@ -195,6 +203,7 @@ class AgenticAnthropicStreamingIterator: self._follow_up_iterator: AsyncIterator | None = None self._drain_task: asyncio.Task | None = None self._hook_task: asyncio.Task | None = None + self._follow_up_chunk_task: asyncio.Task | None = None self._replay_index = 0 self._error_emitted = False @@ -253,7 +262,7 @@ class AgenticAnthropicStreamingIterator: return PING_SSE_BYTES if self._follow_up_iterator is not None: - return await self._follow_up_iterator.__anext__() + return await self._next_follow_up_chunk(self._follow_up_iterator) if self._buffer_holds_server_fulfilled_tool_use(): if self._error_emitted: @@ -273,6 +282,17 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + async def _next_follow_up_chunk(self, follow_up_iterator: AsyncIterator) -> bytes: + if self._follow_up_chunk_task is None: + self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator)) + if not await self._completed_within_ping_interval(self._follow_up_chunk_task): + return PING_SSE_BYTES + chunk: Final = self._follow_up_chunk_task.result() + self._follow_up_chunk_task = None + if chunk is None: + raise StopAsyncIteration + return chunk + def _buffer_holds_server_fulfilled_tool_use(self) -> bool: if not self._server_fulfilled_tool_names: return False @@ -307,6 +327,7 @@ class AgenticAnthropicStreamingIterator: await self._settle_task(self._drain_task) await self._settle_task(self._hook_task) + await self._settle_task(self._follow_up_chunk_task) await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index d59f76232cd..d4aebf099d1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -1001,6 +1001,65 @@ class TestAgenticStreamingIteratorHoldBack: assert [c for c in collected if c != PING_SSE_BYTES] == chunks + @pytest.mark.asyncio + async def test_should_emit_pings_while_the_follow_up_stream_is_slow(self): + """The corrected answer can be slow to generate, so the follow-up stream gets keepalives too.""" + chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream(phase2_chunks, delay_seconds=0.06) + ) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + first_follow_up_index = collected.index(phase2_chunks[0]) + assert collected[first_follow_up_index + 1] == PING_SSE_BYTES + assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_propagate_follow_up_stream_error(self): + """A failing follow-up stream surfaces its error instead of hanging on pings forever.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockFailingAsyncStream([b"follow-up-chunk"], RuntimeError("follow-up died")) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + with pytest.raises(RuntimeError, match="follow-up died"): + async for _ in iterator: + pass + + @pytest.mark.asyncio + async def test_aclose_cancels_in_flight_follow_up_chunk_task(self): + """Closing while a follow-up chunk is pending must not orphan that task.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream([b"follow-up-chunk"], delay_seconds=5.0) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + while iterator._follow_up_chunk_task is None: + await iterator.__anext__() + + await iterator.aclose() + assert iterator._follow_up_chunk_task.cancelled() + @pytest.mark.asyncio async def test_aclose_cancels_drain_task(self): """Closing the iterator mid-buffer must cancel the background drain task.""" From bb0bb48da8c3a5fe3557812e79a31c71c608e006 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 20:05:45 +0000 Subject: [PATCH 013/273] fix(proxy): do not let held-back keepalive pings block the budget reservation refund on client disconnect Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ .../messages/agentic_streaming_iterator.py | 10 +++---- litellm/proxy/common_request_processing.py | 6 ++-- litellm/proxy/common_utils/sse_keepalive.py | 4 ++- .../test_agentic_streaming_iterator.py | 30 +++++++++---------- .../proxy/test_budget_reservation.py | 28 +++++++++++++++++ 6 files changed, 57 insertions(+), 23 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..9db9fb36b65 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -434,6 +434,8 @@ CONNECTION_ERROR_PATTERNS: Final[list[str]] = [ ] STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " +STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' +STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float( os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 3aaba0b139c..d6f4e51a09a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -21,8 +21,8 @@ from collections.abc import AsyncIterator from typing import Any, Final, cast from litellm._logging import verbose_logger +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES -PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b"event: error\n" @@ -249,17 +249,17 @@ class AgenticAnthropicStreamingIterator: async def _anext_held_back(self) -> bytes: if self._drain_task is None: self._drain_task = asyncio.create_task(self._drain_upstream()) - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES if not self._stream_exhausted: if not await self._completed_within_ping_interval(self._drain_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES self._stream_exhausted = True if self._hook_task is None: self._hook_task = asyncio.create_task(self._process_agentic_hooks()) if not await self._completed_within_ping_interval(self._hook_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES if self._follow_up_iterator is not None: return await self._next_follow_up_chunk(self._follow_up_iterator) @@ -286,7 +286,7 @@ class AgenticAnthropicStreamingIterator: if self._follow_up_chunk_task is None: self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator)) if not await self._completed_within_ping_interval(self._follow_up_chunk_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES chunk: Final = self._follow_up_chunk_task.result() self._follow_up_chunk_task = None if chunk is None: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..2607ff411a9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -27,6 +27,7 @@ from litellm.constants import ( MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, + STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -2953,8 +2954,9 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. - delivered_chunk = True + # False and refunds. A keepalive ping carries no provider output, + # so it must not suppress that refund. + delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 6700700ff7c..6e0ea4db431 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -6,7 +6,9 @@ from typing import Final import anyio -ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n' +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_CHUNK + +ANTHROPIC_PING_SSE_CHUNK: Final = STREAM_SSE_KEEPALIVE_PING_CHUNK def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index d4aebf099d1..b0467430533 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -13,8 +13,8 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( - PING_SSE_BYTES, SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, @@ -858,10 +858,10 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - non_ping = [c for c in collected if c != PING_SSE_BYTES] + non_ping = [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] assert non_ping == phase2_chunks assert b"litellm_content_retrieve" not in b"".join(collected) - assert collected[0] == PING_SSE_BYTES + assert collected[0] == STREAM_SSE_KEEPALIVE_PING_BYTES @pytest.mark.asyncio async def test_should_replay_buffer_verbatim_when_no_hook_fires(self): @@ -877,7 +877,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks mock_handler._call_agentic_completion_hooks.assert_awaited_once() @pytest.mark.asyncio @@ -898,8 +898,8 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert collected.count(PING_SSE_BYTES) >= 2 - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert collected.count(STREAM_SSE_KEEPALIVE_PING_BYTES) >= 2 + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks @pytest.mark.asyncio async def test_should_propagate_upstream_error_instead_of_partial_message(self): @@ -919,7 +919,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert all(c == PING_SSE_BYTES for c in collected) + assert all(c == STREAM_SSE_KEEPALIVE_PING_BYTES for c in collected) mock_handler._call_agentic_completion_hooks.assert_not_awaited() @pytest.mark.asyncio @@ -945,8 +945,8 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert collected.count(PING_SSE_BYTES) >= 4 - assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + assert collected.count(STREAM_SSE_KEEPALIVE_PING_BYTES) >= 4 + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == phase2_chunks @pytest.mark.asyncio async def test_should_error_instead_of_replaying_server_fulfilled_tool_use_when_hook_crashes(self): @@ -962,7 +962,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] assert b"litellm_content_retrieve" not in b"".join(collected) @pytest.mark.asyncio @@ -979,7 +979,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] @pytest.mark.asyncio async def test_should_replay_client_owned_tool_use_verbatim(self): @@ -999,7 +999,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks @pytest.mark.asyncio async def test_should_emit_pings_while_the_follow_up_stream_is_slow(self): @@ -1023,8 +1023,8 @@ class TestAgenticStreamingIteratorHoldBack: collected.append(chunk) first_follow_up_index = collected.index(phase2_chunks[0]) - assert collected[first_follow_up_index + 1] == PING_SSE_BYTES - assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + assert collected[first_follow_up_index + 1] == STREAM_SSE_KEEPALIVE_PING_BYTES + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == phase2_chunks @pytest.mark.asyncio async def test_should_propagate_follow_up_stream_error(self): @@ -1074,7 +1074,7 @@ class TestAgenticStreamingIteratorHoldBack: ) first = await iterator.__anext__() - assert first == PING_SSE_BYTES + assert first == STREAM_SSE_KEEPALIVE_PING_BYTES assert iterator._drain_task is not None await iterator.aclose() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 34adb4d2091..e9a0e80752d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -7,6 +7,7 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, @@ -2453,6 +2454,33 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() +@pytest.mark.asyncio +async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token, reservation = await _reserve_for_stream( + counter_cache, key_cache, proxy_logging_obj, "key-cancel-after-ping" + ) + + async def cancel_after_ping(user_api_key_dict, response, request_data): + yield STREAM_SSE_KEEPALIVE_PING_BYTES + raise asyncio.CancelledError() + + generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_ping) + received = [] + with pytest.raises(asyncio.CancelledError): + async for chunk in generator: + received.append(chunk) + + assert received == [STREAM_SSE_KEEPALIVE_PING_BYTES] + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-cancel-after-ping" + ) == pytest.approx(0.5) + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_swallows_release_errors(): # If the release itself fails (e.g. Redis unavailable) it must not escape From 2d1ee3aab2fe6a37c80085009f789416fe191d4e Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 20:39:55 +0000 Subject: [PATCH 014/273] fix(proxy): keep the reservation when a disconnect happens while provider output is held back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 5 ++ litellm/proxy/common_request_processing.py | 9 ++- .../proxy/test_budget_reservation.py | 63 +++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index d6f4e51a09a..3d3d3a12b17 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -207,6 +207,11 @@ class AgenticAnthropicStreamingIterator: self._replay_index = 0 self._error_emitted = False + @property + def has_buffered_provider_output(self) -> bool: + """Whether provider output was received but withheld from the client behind keepalive pings.""" + return self._hold_back and bool(self._collected_bytes) + def __aiter__(self): return self diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2607ff411a9..3aeb7729c81 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -40,6 +40,9 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, +) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe @@ -95,6 +98,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } +def _withheld_provider_output(response: object) -> bool: + return isinstance(response, AgenticAnthropicStreamingIterator) and response.has_buffered_provider_output + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -2970,7 +2977,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk: + if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index e9a0e80752d..c3210dd7f4d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -8,6 +8,9 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, +) from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, @@ -2376,6 +2379,11 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token return valid_token, reservation +async def _never_ending_stream(): + yield b'event: message_start\ndata: {"type": "message_start"}\n\n' + await asyncio.sleep(30) + + def _drive_streaming_cancel(valid_token, iterator_hook): streaming_logging_obj = MagicMock() streaming_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook @@ -2481,6 +2489,61 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_streaming_cancel_while_holding_back_provider_output_keeps_reservation( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token, reservation = await _reserve_for_stream( + counter_cache, key_cache, proxy_logging_obj, "key-cancel-held-back" + ) + + held_back = AgenticAnthropicStreamingIterator( + completion_stream=_never_ending_stream(), + http_handler=MagicMock(), + model="claude-haiku-4-5", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + server_fulfilled_tool_names=frozenset({"headroom_retrieve"}), + ping_interval_seconds=0.01, + ) + + async def ping_then_cancel(user_api_key_dict, response, request_data): + yield await response.__anext__() + while not response.has_buffered_provider_output: + yield await response.__anext__() + raise asyncio.CancelledError() + + streaming_logging_obj = MagicMock() + streaming_logging_obj.async_post_call_streaming_iterator_hook = ping_then_cancel + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() + generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=held_back, + user_api_key_dict=valid_token, + request_data=_request_body(), + proxy_logging_obj=streaming_logging_obj, + serialize_chunk=lambda chunk: chunk, + serialize_error=lambda exc: str(exc), + ) + + received = [] + with pytest.raises(asyncio.CancelledError): + async for chunk in generator: + received.append(chunk) + + assert received and received == [STREAM_SSE_KEEPALIVE_PING_BYTES] * len(received) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-cancel-held-back" + ) == pytest.approx(2.0) + assert reservation.get("finalized") is not True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_swallows_release_errors(): # If the release itself fails (e.g. Redis unavailable) it must not escape 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 015/273] 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 016/273] 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 017/273] 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 018/273] 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 019/273] 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 020/273] 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 021/273] 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 b8680e6baed05863712a4c57a10b128ecd95475a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:22:31 +0000 Subject: [PATCH 022/273] fix(ui): render tag-based guardrail mode instead of crashing guardrails page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/guardrailTableColumns.tsx | 13 +++++--- .../_components/guardrail_info.test.tsx | 30 +++++++++++++++++++ .../guardrails/_components/guardrail_info.tsx | 5 ++-- .../guardrail_info_helpers.test.tsx | 29 ++++++++++++++++++ .../_components/guardrail_info_helpers.tsx | 13 ++++++++ .../_components/guardrail_table.test.tsx | 12 ++++++++ .../src/components/guardrails/types.ts | 7 ++++- 7 files changed, 102 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx index ec3d05a6907..53f1b1a03d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -15,7 +15,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; @@ -117,9 +117,14 @@ export const getGuardrailTableColumns = ({ header: "Mode", size: 130, enableSorting: false, - cell: ({ row }) => ( - {row.original.litellm_params.mode} - ), + cell: ({ row }) => { + const mode = formatGuardrailMode(row.original.litellm_params.mode); + return ( + + {mode || "-"} + + ); + }, }, { id: "default_on", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index b6ee130d50a..3f6317ed366 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -82,6 +82,36 @@ describe("Guardrail Info", () => { expect(getByText("Settings")).toBeInTheDocument(); }); + it("should render a tag-based mode object rather than crashing the detail view", async () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "database", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { findAllByText } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + }); + it("should render the provider logo from the bundled guardrail logo map", async () => { vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ guardrail_id: "123", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index e80ddac932f..5e476a8accd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -35,6 +35,7 @@ import { import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; import { + formatGuardrailMode, getGuardrailLogoAndName, guardrail_provider_map, skipSystemMessageToChoice, @@ -559,7 +560,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{guardrailData.litellm_params?.mode || "-"}

+

{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} @@ -852,7 +853,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-
{guardrailData.litellm_params?.mode || "-"}
+
{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

Default On

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx index ec910673b8f..c5e07fe9624 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx @@ -14,6 +14,7 @@ import { choiceToSkipSystemForCreate, skipToolMessageToChoice, choiceToSkipToolForCreate, + formatGuardrailMode, } from "./guardrail_info_helpers"; describe("guardrail_info_helpers", () => { @@ -210,6 +211,34 @@ describe("guardrail_info_helpers", () => { }); }); + describe("formatGuardrailMode", () => { + it("renders a single mode and a list of modes", () => { + expect(formatGuardrailMode("pre_call")).toBe("pre_call"); + expect(formatGuardrailMode(["pre_call", "post_call"])).toBe("pre_call, post_call"); + }); + + it("flattens a tag-based mode object into deduped modes instead of returning it verbatim", () => { + const mode = { + tags: { "Service-Type: internal-service": "post_call", "Service-Type: batch": ["during_call", "post_call"] }, + default: ["pre_call", "post_call"], + }; + + expect(formatGuardrailMode(mode)).toBe("pre_call, post_call, during_call (tag-based)"); + }); + + it("handles a tag-based mode with no default and with no tags", () => { + expect(formatGuardrailMode({ tags: { "team: a": "post_call" } })).toBe("post_call (tag-based)"); + expect(formatGuardrailMode({ default: "pre_call" })).toBe("pre_call (tag-based)"); + }); + + it("returns an empty string for missing or unusable modes", () => { + expect(formatGuardrailMode(undefined)).toBe(""); + expect(formatGuardrailMode(null)).toBe(""); + expect(formatGuardrailMode({})).toBe(""); + expect(formatGuardrailMode({ tags: {}, default: null })).toBe(""); + }); + }); + describe("skipSystemMessageToChoice / choiceToSkipSystemForCreate", () => { it("maps API values to form choices and back for create", () => { expect(skipSystemMessageToChoice(undefined)).toBe("inherit"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 12aaba0d696..c12529e6326 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,6 +110,19 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; +// Turns a guardrail mode into a renderable string. A mode is a single mode, a list of modes, or a +// tag-based `{ tags, default }` object, which React refuses to render as a child +export const formatGuardrailMode = (raw: unknown): string => { + const flat: string[] = toModeArray(raw); + if (flat.length > 0) return flat.join(", "); + if (raw === null || typeof raw !== "object") return ""; + + const { tags, default: fallback } = raw as { tags?: Record; default?: unknown }; + const tagged: string[] = tags && typeof tags === "object" ? Object.values(tags).flatMap(toModeArray) : []; + const modes: string[] = Array.from(new Set([...toModeArray(fallback), ...tagged])); + return modes.length > 0 ? `${modes.join(", ")} (tag-based)` : ""; +}; + // Resolves the supported modes for the selected provider, falling back to the global list export const getSupportedModesForProvider = ( settings: { supported_modes?: string[]; supported_modes_by_provider?: Record } | null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index ee619dc7468..561a89a191a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -46,6 +46,18 @@ describe("GuardrailTable", () => { expect(screen.getByText("m")).toBeInTheDocument(); }); + it("renders a tag-based mode object instead of crashing the table", () => { + const guardrail = makeGuardrail({ + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + }); + render(); + expect(screen.getByText("pre_call, post_call (tag-based)")).toBeInTheDocument(); + }); + it("deletes a DB guardrail through the actions menu", async () => { const user = userEvent.setup(); const onDeleteClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/guardrails/types.ts b/ui/litellm-dashboard/src/components/guardrails/types.ts index e8ed27d9e45..0f5ce1c883d 100644 --- a/ui/litellm-dashboard/src/components/guardrails/types.ts +++ b/ui/litellm-dashboard/src/components/guardrails/types.ts @@ -18,12 +18,17 @@ export interface PiiConfigurationProps { entityCategories?: PiiEntityCategory[]; } +export type GuardrailMode = + | string + | string[] + | { tags?: Record; default?: string | string[] | null }; + export interface Guardrail { guardrail_id: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; pii_entities_config?: { [key: string]: string }; [key: string]: any; From 881aa2080871052a2173f7b3352df39fc0e61e03 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:31:42 +0000 Subject: [PATCH 023/273] fix(ui): format tag-based guardrail mode in delete modal, playground, and policy picker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/GuardrailTestPlayground.tsx | 8 ++++++-- .../guardrails/_components/GuardrailsPanel.tsx | 4 ++-- .../(dashboard)/guardrails/_components/guardrail_info.tsx | 4 +++- .../policies/_components/guardrail_selection_modal.tsx | 5 ++++- 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx index fd8ed22867b..c64b5d7cb5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx @@ -6,13 +6,15 @@ import { toast } from "@/lib/toast"; import { Card, CardContent } from "@/components/ui/card"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { GuardrailMode } from "@/components/guardrails/types"; +import { formatGuardrailMode } from "./guardrail_info_helpers"; interface GuardrailItem { guardrail_id?: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; }; guardrail_info: Record | null; @@ -171,7 +173,9 @@ const GuardrailTestPlayground: React.FC = ({
Mode: - {guardrail.litellm_params.mode} + + {formatGuardrailMode(guardrail.litellm_params.mode)} +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index b4c29bd9c40..7e59abf8e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -18,7 +18,7 @@ import GuardrailTestPlayground from "./GuardrailTestPlayground"; import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -211,7 +211,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { label: "Name", value: guardrailToDelete?.guardrail_name }, { label: "ID", value: guardrailToDelete?.guardrail_id, code: true }, { label: "Provider", value: providerDisplayName }, - { label: "Mode", value: guardrailToDelete?.litellm_params.mode }, + { label: "Mode", value: formatGuardrailMode(guardrailToDelete?.litellm_params.mode) }, { label: "Default On", value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 5e476a8accd..d4a1885146d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -560,7 +560,9 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

+

+ {formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"} +

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx index 0b439462c1a..f87155db719 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx @@ -12,6 +12,7 @@ import { } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; import { CheckCircle2, Info } from "lucide-react"; +import { formatGuardrailMode } from "@/app/(dashboard)/guardrails/_components/guardrail_info_helpers"; interface GuardrailInfo { guardrail_name: string; @@ -163,7 +164,9 @@ const GuardrailSelectionModal: React.FC = ({ {/* Show guardrail type and mode */}
{guardrail.definition?.litellm_params?.guardrail || "unknown"} - {guardrail.definition?.litellm_params?.mode || "unknown"} + + {formatGuardrailMode(guardrail.definition?.litellm_params?.mode) || "unknown"} + {guardrail.definition?.litellm_params?.patterns && ( {guardrail.definition.litellm_params.patterns.length} pattern(s) From f80cb0d9f8e37539b39bf6412ef7f673c2074e58 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:33:06 +0000 Subject: [PATCH 024/273] refactor(ui): drop redundant comment above guardrail mode formatter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/guardrail_info_helpers.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index c12529e6326..83038b8e0e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,8 +110,6 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; -// Turns a guardrail mode into a renderable string. A mode is a single mode, a list of modes, or a -// tag-based `{ tags, default }` object, which React refuses to render as a child export const formatGuardrailMode = (raw: unknown): string => { const flat: string[] = toModeArray(raw); if (flat.length > 0) return flat.join(", "); From 9e86cfa7e994edd3ac77456a7b0edb974e8012ff Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 01:56:02 +0000 Subject: [PATCH 025/273] fix(auth): support wildcard prefixes in jwt team_allowed_routes team_allowed_routes and admin_allowed_routes only matched exact strings or named route groups, so a whole prefix of pass-through endpoints had to be listed route by route in config. Match trailing-wildcard patterns with the same helper the key-level allowed_routes check uses, so "/prefix/*" covers endpoints registered later. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 2 +- litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/auth_utils.py | 2 +- litellm/proxy/auth/route_checks.py | 10 +-- litellm/proxy/policy_engine/policy_matcher.py | 4 +- .../policy_engine/policy_resolve_endpoints.py | 8 +- .../proxy/auth/test_auth_checks.py | 79 +++++++++++++++++++ .../policies/_components/scope_validation.ts | 2 +- 8 files changed, 96 insertions(+), 16 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 776aecbd883..46a06f77861 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1823 + "limit": 1817 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 12d6b44a648..9d8eedaa7dc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1128,7 +1128,8 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. + - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/tempus/*"). """ from starlette.routing import compile_path @@ -1138,7 +1139,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: regex, _, _ = compile_path(template) if regex.match(user_route): return True - elif allowed_route == user_route: + elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): return True return False diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ce662ee0374..1e6d8137d53 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -608,7 +608,7 @@ def route_in_additonal_public_routes(current_route: str): # Check wildcard patterns for route_pattern in routes_defined: - if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): + if RouteChecks.route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): return True return False diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cea21ca088b..4dba2497bb9 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -181,7 +181,7 @@ class RouteChecks: # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True if denied_auth_enforced_pass_through_route: @@ -329,7 +329,7 @@ class RouteChecks: route_allowed = True break - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): route_allowed = True break @@ -397,7 +397,7 @@ class RouteChecks: return True # Check for wildcard patterns like "/containers/*" if RouteChecks._is_wildcard_pattern(pattern=openai_route): - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): return True # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" @@ -517,7 +517,7 @@ class RouteChecks: return pattern.endswith("*") @staticmethod - def _route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: """ Check if route matches the wildcard pattern @@ -594,7 +594,7 @@ class RouteChecks: # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* ######################################################### if any( - RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route) + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes if RouteChecks._is_wildcard_pattern(pattern=allowed_route) ): diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index f66dc4e7bbe..001e4115374 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -30,7 +30,7 @@ class PolicyMatcher: """ Check if a value matches any of the given patterns. - Uses the existing RouteChecks._route_matches_wildcard_pattern helper. + Uses the existing RouteChecks.route_matches_wildcard_pattern helper. Args: value: The value to check (e.g., team alias, key alias, model) @@ -45,7 +45,7 @@ class PolicyMatcher: for pattern in patterns: # Use existing wildcard pattern matching helper - if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern): + if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern): return True return False diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 346586c1e5a..70b98933d0f 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -100,7 +100,7 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: key_alias = key.key_alias or "" key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None)) if key_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in key_tags for pat in tag_patterns ): @@ -123,7 +123,7 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple: team_alias = team.team_alias or "" team_tags = _get_tags_from_metadata(team.metadata) if team_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in team_tags for pat in tag_patterns ): @@ -152,7 +152,7 @@ async def _find_affected_by_team_patterns( for team in all_teams: team_alias = team.team_alias or "" if team_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns + RouteChecks.route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns ): if team_alias not in existing_teams: new_teams.append(team_alias) @@ -190,7 +190,7 @@ async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list for key in keys: key_alias = key.key_alias or "" if key_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns + RouteChecks.route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns ): if key_alias not in existing_keys: affected.append(key_alias) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 6b40fa1b324..0b174cda9d5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6897,3 +6897,82 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is router = _router_with_a_group_priced_through_model_info() assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False + +@pytest.mark.parametrize( + "user_route, expected", + [ + ("/tempus/v1/chat/completions", True), + ("/tempus/newly-registered-model/predict", True), + ("/tempus-other/v1/chat/completions", False), + ("/anthropic/v1/messages", False), + ], +) +def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_routes(user_route, expected): + """A `/prefix/*` entry in `team_allowed_routes` must cover every route under that prefix, so + passthrough endpoints registered after the proxy config was written are reachable without an + exact-route config change.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=user_route, + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/tempus/*"]), + ) + is expected + ) + + +def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["/tempus/model-a"]) + + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-a", litellm_proxy_roles=roles) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-b", litellm_proxy_roles=roles) + is False + ) + + +def test_admin_allowed_routes_wildcard_prefix_is_honored(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/tempus/*"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/tempus/anything", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/other/anything", litellm_proxy_roles=roles + ) + is False + ) + + +def test_team_allowed_routes_named_route_group_still_resolves(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/v1/chat/completions", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/key/generate", litellm_proxy_roles=roles) + is False + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts index 7c49117088c..53a76dc5a1c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts @@ -1,4 +1,4 @@ -// Mirrors request-time matching (RouteChecks._route_matches_wildcard_pattern): only a +// Mirrors request-time matching (RouteChecks.route_matches_wildcard_pattern): only a // trailing "*" is a wildcard (prefix match). Anything else - including a "?" or a // non-trailing "*" - is compared by exact equality when a request is matched, so it is // treated as a concrete alias that must exist. From 07416344cc8865c1867c51dd733582e04236aeef Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 02:11:18 +0000 Subject: [PATCH 026/273] test(auth): use a generic route prefix in wildcard route tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- .../proxy/auth/test_auth_checks.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9d8eedaa7dc..bf7a6a8f6c3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1129,7 +1129,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name - (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/tempus/*"). + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). """ from starlette.routing import compile_path diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0b174cda9d5..7fa16508054 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6901,9 +6901,9 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is @pytest.mark.parametrize( "user_route, expected", [ - ("/tempus/v1/chat/completions", True), - ("/tempus/newly-registered-model/predict", True), - ("/tempus-other/v1/chat/completions", False), + ("/internal-models/v1/chat/completions", True), + ("/internal-models/newly-registered-model/predict", True), + ("/internal-models-other/v1/chat/completions", False), ("/anthropic/v1/messages", False), ], ) @@ -6918,7 +6918,7 @@ def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_ro allowed_routes_check( user_role=LitellmUserRoles.TEAM, user_route=user_route, - litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/tempus/*"]), + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/*"]), ) is expected ) @@ -6928,14 +6928,14 @@ def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.auth_checks import allowed_routes_check - roles = LiteLLM_JWTAuth(team_allowed_routes=["/tempus/model-a"]) + roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-a", litellm_proxy_roles=roles) + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) is True ) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-b", litellm_proxy_roles=roles) + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) is False ) @@ -6944,11 +6944,11 @@ def test_admin_allowed_routes_wildcard_prefix_is_honored(): from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.auth_checks import allowed_routes_check - roles = LiteLLM_JWTAuth(admin_allowed_routes=["/tempus/*"]) + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/internal-models/*"]) assert ( allowed_routes_check( - user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/tempus/anything", litellm_proxy_roles=roles + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/internal-models/anything", litellm_proxy_roles=roles ) is True ) From cafc8c1455a7691b4cf2082bc809abeb3cc45af4 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 03:01:26 +0000 Subject: [PATCH 027/273] fix(proxy): store the actual selected model in spend logs for Azure Model Router Co-authored-by: Filippo Mattia Menghi Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/spend_tracking_utils.py | 4 +- .../test_spend_tracking_utils.py | 42 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 0b56f0d8246..822f03873b3 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -411,7 +411,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9710dc44e99..b1a45fb84a3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3241,3 +3241,45 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["model_group"] == "" assert payload["api_base"] == "" assert payload["custom_llm_provider"] == "" + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> dict[str, Any]: + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": slp_model, + "metadata": {}, + "model_map_information": StandardLoggingModelInformation( + model_map_key="azure_ai/model_router", model_map_value=None + ), + }, + ) + return { + "model": "azure_ai/model_router/model-router", + "litellm_params": {"metadata": {"user_api_key": "sk-test-key"}}, + "standard_logging_object": standard_logging_payload, + } + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_uses_standard_logging_payload_model(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/gpt-5-mini" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model=None), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/model_router/model-router" From 57b367c78e6f691839a4c6dccf8ffe57bfb25478 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 03:25:06 +0000 Subject: [PATCH 028/273] refactor(tests): type the model router spend log kwargs helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/spend_tracking/test_spend_tracking_utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index b1a45fb84a3..9c97b2683b2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -6,6 +6,8 @@ import sys from datetime import timezone from typing import Any, Final, cast +from typing_extensions import ReadOnly, TypedDict + import pytest from fastapi.testclient import TestClient @@ -3243,7 +3245,13 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["custom_llm_provider"] == "" -def _model_router_spend_log_kwargs(slp_model: str | None) -> dict[str, Any]: +class _ModelRouterSpendLogKwargs(TypedDict): + model: ReadOnly[str] + litellm_params: ReadOnly[dict[str, dict[str, str]]] + standard_logging_object: ReadOnly[StandardLoggingPayload] + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs: standard_logging_payload: Final = cast( StandardLoggingPayload, { From 4e37425a782b7bf09d18e42ef9f072f6108155e4 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 21 Aug 2026 15:26:42 -0700 Subject: [PATCH 029/273] 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 030/273] 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 031/273] 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 032/273] 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 033/273] 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 034/273] 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 035/273] 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 036/273] 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 037/273] 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 038/273] 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 039/273] 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 040/273] 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 041/273] 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 042/273] 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 043/273] 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 044/273] 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 045/273] 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 046/273] 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 047/273] 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 048/273] 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 70e4273ba1f50fd921c05198fdfc868291dc2d57 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:12:17 -0700 Subject: [PATCH 049/273] fix(responses): make previous_response_id resolve on the bridged path Streaming /v1/responses over the completion bridge minted a fresh resp_{uuid4} for every response, while spend tracking stored the inner chat completion id as request_id. The session lookup queries on request_id, so a follow-up sent with that response id matched no rows and the prior conversation was silently dropped. The iterator now pulls the first upstream chunk before emitting response.created, so created, in_progress and completed all carry the same encoded chat completion id. Two more ways the same history went missing: - The session lookup only read spend logs already written to the DB, so a follow-up sent inside the batch writer's window found nothing. It now also reads the rows still queued in memory. - Input was only accepted as a string or a single dict, so the list shape the Responses API actually sends dropped every user turn from the reconstructed history. --- litellm/proxy/utils.py | 10 + .../session_handler.py | 67 ++++- .../streaming_iterator.py | 82 +++++- .../test_session_handler.py | 235 ++++++++++++++++++ .../test_streaming_iterator_response_id.py | 130 ++++++++++ 5 files changed, 508 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c616d9e8723..9978fa04f40 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6005,6 +6005,16 @@ async def enqueue_spend_logs( ) +async def peek_spend_logs(prisma_client: PrismaClient) -> tuple[SpendLogsPayload, ...]: + """Snapshot the spend logs still waiting for the next flush, leaving the queue intact. + + Reads that need a just-finished request use this, since the batch writer only + reaches the DB every ``PROXY_BATCH_WRITE_AT`` seconds. + """ + async with prisma_client._spend_log_transactions_lock: + return tuple(prisma_client.spend_log_transactions) + + async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index dcff26c5b0c..935c78bc9a1 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -104,10 +105,10 @@ class ResponsesSessionHandler: if proxy_server_request_dict: _response_input_param: Final = proxy_server_request_dict.get("input", None) _messages = proxy_server_request_dict.get("messages", None) - if isinstance(_response_input_param, str): + if isinstance(_response_input_param, (str, list)): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast(ResponseInputParam, _response_input_param) + response_input_param = cast(ResponseInputParam, [_response_input_param]) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -131,8 +132,8 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = spend_log.get("response", "{}") - if isinstance(_response_output, dict) and _response_output and _response_output != {}: + _response_output: Final = ResponsesSessionHandler._get_response_dict_from_spend_log(spend_log) + if _response_output: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -140,6 +141,23 @@ class ResponsesSessionHandler: chat_completion_message_history.append(getattr(choice, "message")) return chat_completion_message_history + @staticmethod + def _get_response_dict_from_spend_log(spend_log: SpendLogsPayload) -> Mapping[str, Any] | None: + """ + Spend logs read from the DB hold `response` as a dict, ones still queued in memory + hold it as a JSON string. + """ + _response_output: Final = spend_log.get("response") + if isinstance(_response_output, dict): + return _response_output or None + if isinstance(_response_output, str): + try: + parsed: Final = json.loads(_response_output) + except json.JSONDecodeError: + return None + return parsed if isinstance(parsed, dict) and parsed else None + return None + @staticmethod async def get_proxy_server_request_from_spend_log( spend_log: SpendLogsPayload, @@ -256,11 +274,12 @@ class ResponsesSessionHandler: SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id """ from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.utils import peek_spend_logs verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) decoded_response_id: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) - previous_response_id = decoded_response_id.get("response_id", previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) if prisma_client is None: return [] @@ -276,12 +295,46 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - spend_logs: Final = await prisma_client.db.query_raw(query, previous_response_id) + written_spend_logs: Final = await prisma_client.db.query_raw(query, response_id) + queued_spend_logs: Final = await peek_spend_logs(prisma_client) + spend_logs: Final = list( + ResponsesSessionHandler._merge_queued_spend_logs( + response_id=response_id, + written_spend_logs=written_spend_logs, + queued_spend_logs=queued_spend_logs, + ) + ) verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", - previous_response_id, + response_id, json.dumps(spend_logs, indent=4, default=str), ) return spend_logs + + @staticmethod + def _merge_queued_spend_logs( + response_id: str, + written_spend_logs: Sequence[SpendLogsPayload], + queued_spend_logs: Sequence[SpendLogsPayload], + ) -> tuple[SpendLogsPayload, ...]: + """ + Append the session's spend logs that the batch writer has not flushed to the DB yet. + + Without this a follow-up sent inside the ``PROXY_BATCH_WRITE_AT`` window sees an + empty session and silently drops the conversation. The queue is FIFO, so anything + still on it is newer than every row already written. + """ + session_ids: Final = frozenset( + session_id + for spend_log in (*written_spend_logs, *queued_spend_logs) + if spend_log.get("request_id") == response_id and (session_id := spend_log.get("session_id")) + ) | frozenset(session_id for spend_log in written_spend_logs if (session_id := spend_log.get("session_id"))) + written_request_ids: Final = frozenset(spend_log.get("request_id") for spend_log in written_spend_logs) + unflushed: Final = tuple( + spend_log + for spend_log in queued_spend_logs + if spend_log.get("session_id") in session_ids and spend_log.get("request_id") not in written_request_ids + ) + return (*written_spend_logs, *unflushed) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index aa5708088b7..a8092edc625 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -86,6 +86,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.final_text: str = "" self._cached_item_id: str | None = None self._cached_response_id: str | None = None + self._buffered_chunk: ModelResponseStream | None = None + self._upstream_exhausted: bool = False + self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} @@ -330,6 +333,59 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) self._pending_tool_events.append(item_done_event) + def _adopt_response_id_from_chunk(self, chunk: ModelResponseStream) -> None: + if self._cached_response_id is not None: + return + chunk_id: Final = getattr(chunk, "id", None) + if chunk_id and isinstance(chunk_id, str): + self._cached_response_id = chunk_id + + async def _aprime_response_id(self) -> None: + """ + Pull the first upstream chunk before `response.created` is emitted so every event + carries the chat completion id that spend tracking stores as `request_id`. + """ + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = await self.litellm_custom_stream_wrapper.__anext__() + except StopAsyncIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _prime_response_id(self) -> None: + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = self.litellm_custom_stream_wrapper.__next__() + except StopIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _take_buffered_chunk(self) -> ModelResponseStream | None: + buffered: Final = self._buffered_chunk + self._buffered_chunk = None + return buffered + + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: + return ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response, + custom_llm_provider=self.custom_llm_provider, + litellm_metadata=self.litellm_metadata, + ) + def _default_response_created_event_data(self) -> dict: # Use cached response ID if available, otherwise generate a new one if self._cached_response_id is None: @@ -388,7 +444,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, - response=ResponsesAPIResponse(**response_created_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -399,7 +455,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, - response=ResponsesAPIResponse(**response_in_progress_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -811,6 +867,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self.finished is True: raise StopAsyncIteration + await self._aprime_response_id() result = self.return_default_initial_events() if result: return result @@ -822,7 +879,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._pending_tool_events.pop(0) try: - chunk = await self.litellm_custom_stream_wrapper.__anext__() + chunk = self._take_buffered_chunk() + if chunk is None: + if self._upstream_exhausted: + raise StopAsyncIteration + chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) self._ensure_output_item_for_chunk(chunk) @@ -912,6 +973,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): while True: if self.finished is True: raise StopIteration + self._prime_response_id() result = self.return_default_initial_events() if result: return result @@ -922,7 +984,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._pending_tool_events: return self._pending_tool_events.pop(0) try: - chunk = self.litellm_custom_stream_wrapper.__next__() + buffered_chunk = self._take_buffered_chunk() + if buffered_chunk is not None: + chunk = buffered_chunk + elif self._upstream_exhausted: + raise StopIteration + else: + chunk = self.litellm_custom_stream_wrapper.__next__() self._ensure_output_item_for_chunk(chunk) # Accumulate provider_specific_fields from chunk and delta for src in ( @@ -1082,11 +1150,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_response.id = self._cached_response_id # Encode the response ID to match non-streaming behavior - encoded_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=responses_api_response, - custom_llm_provider=self.custom_llm_provider, - litellm_metadata=self.litellm_metadata, - ) + encoded_response: Final = self._with_encoded_response_id(responses_api_response) return ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 19f240fa3d4..926e9e0af2a 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -1,3 +1,4 @@ +import asyncio import json from unittest.mock import AsyncMock, patch @@ -10,6 +11,7 @@ from litellm.responses.litellm_completion_transformation import session_handler from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) +from litellm.responses.utils import ResponsesAPIRequestUtils @pytest.mark.asyncio @@ -430,3 +432,236 @@ async def test_get_chat_completion_message_history_empty_response_dict(): # Verify the session was still created correctly assert result["litellm_session_id"] == "test-session" + + +def _chat_completion_response(request_id: str, content: str) -> dict: + return { + "id": request_id, + "object": "chat.completion", + "created": 1748575031, + "model": "claude-haiku-4-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + } + + +class _FakePrismaDB: + def __init__(self, rows): + self._rows = rows + self.calls = [] + + async def query_raw(self, query, *args): + self.calls.append(args) + return list(self._rows) + + +class _FakePrismaClient: + def __init__(self, written_rows, queued_rows): + self.db = _FakePrismaDB(written_rows) + self.spend_log_transactions = list(queued_rows) + self._spend_log_transactions_lock = asyncio.Lock() + + +@pytest.mark.asyncio +async def test_message_history_reconstructs_list_shaped_input(): + """ + The Responses API sends `input` as a list of items, which is what lands in the stored + proxy_server_request. The user turns have to survive session reconstruction. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + mock_spend_logs = [ + { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701", + "proxy_server_request": { + "input": [ + { + "role": "user", + "content": "Remember this: my favorite color is chartreuse.", + } + ], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, "OK"), + } + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs: + mock_get_spend_logs.return_value = mock_spend_logs + + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + assert result["litellm_session_id"] == "a96757c4-c6dc-4c76-b37e-e7dfa526b701" + + +@pytest.mark.asyncio +async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_writer(): + """ + A follow-up sent right after the previous turn arrives before the batch writer has + flushed that turn's spend log, so the row is only in memory. The history has to + include it anyway. + """ + request_id = "chatcmpl-6c1f5f6c-6a2b-4c62-8d1f-0d9d4ce0a1b2" + queued_spend_log = { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3", + "proxy_server_request": json.dumps( + { + "input": [ + { + "role": "user", + "content": "Remember this: my favorite color is chartreuse.", + } + ], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response(request_id, "OK")), + } + fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[queued_spend_log]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + assert result["litellm_session_id"] == "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" + assert fake_prisma_client.spend_log_transactions == [queued_spend_log] + + +@pytest.mark.asyncio +async def test_message_history_merges_written_and_queued_turns_in_order(): + """ + Turn 1 already flushed to the DB, turn 2 still queued: the follow-up sees the whole + conversation, in order, with no row counted twice. + """ + session_id = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1" + first_request_id = "chatcmpl-1111" + second_request_id = "chatcmpl-2222" + written_spend_log = { + "request_id": first_request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": { + "input": [{"role": "user", "content": "My favorite color is chartreuse."}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(first_request_id, "Got it."), + } + queued_spend_log = { + "request_id": second_request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": json.dumps( + { + "input": [{"role": "user", "content": "And my favorite city is Lisbon."}], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response(second_request_id, "Noted.")), + } + fake_prisma_client = _FakePrismaClient( + written_rows=[written_spend_log], + queued_rows=[written_spend_log, queued_spend_log], + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + second_request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "My favorite color is chartreuse."), + ("assistant", "Got it."), + ("user", "And my favorite city is Lisbon."), + ("assistant", "Noted."), + ] + assert result["litellm_session_id"] == session_id + + +@pytest.mark.asyncio +async def test_message_history_ignores_queued_spend_logs_from_other_sessions(): + request_id = "chatcmpl-3333" + written_spend_log = { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "session-a", + "proxy_server_request": { + "input": [{"role": "user", "content": "Hello from session a."}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, "Hi."), + } + other_session_spend_log = { + "request_id": "chatcmpl-4444", + "call_type": "aresponses", + "session_id": "session-b", + "proxy_server_request": json.dumps( + { + "input": [{"role": "user", "content": "Hello from session b."}], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response("chatcmpl-4444", "Hi there.")), + } + fake_prisma_client = _FakePrismaClient( + written_rows=[written_spend_log], + queued_rows=[other_session_spend_log], + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Hello from session a."), + ("assistant", "Hi."), + ] + + +@pytest.mark.asyncio +async def test_message_history_looks_up_the_decoded_chat_completion_id(): + """ + A `previous_response_id` handed back by the proxy is base64 encoded; spend logs store + the bare chat completion id, so that is what the lookup has to query on. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="anthropic", + model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", + response_id=request_id, + ) + fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + encoded_response_id + ) + + assert fake_prisma_client.db.calls == [(request_id,)] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py new file mode 100644 index 00000000000..97f35900e9d --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py @@ -0,0 +1,130 @@ +from unittest.mock import MagicMock + +import pytest + +from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + +CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" +RESPONSE_ID_EVENT_TYPES = frozenset( + {"response.created", "response.in_progress", "response.completed"} +) + + +def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +class _FakeStreamWrapper: + def __init__(self, chunks): + self._chunks = list(chunks) + self.logging_obj = MagicMock() + + def __iter__(self): + return self + + def __next__(self): + if not self._chunks: + raise StopIteration + return self._chunks.pop(0) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + +def _response_ids(events) -> list[str]: + return [ + event.response.id + for event in events + if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + +@pytest.mark.asyncio +async def test_streaming_events_share_the_chat_completion_response_id(): + """ + Every event of a bridged stream has to carry the same id, and that id has to decode + to the chat completion id spend tracking stores as `request_id`. Otherwise a + follow-up `previous_response_id` matches no session and the conversation is dropped. + """ + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) + assert decoded["response_id"] == CHAT_COMPLETION_ID + assert decoded["custom_llm_provider"] == "anthropic" + + +def test_sync_streaming_events_share_the_chat_completion_response_id(): + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = list(iterator) + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + assert ( + ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] + == CHAT_COMPLETION_ID + ) + + +@pytest.mark.asyncio +async def test_streaming_emits_every_chunk_after_priming_the_response_id(): + iterator = _build_iterator( + [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] + ) + + events = [event async for event in iterator] + + deltas = "".join( + event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert deltas == "Hello!" + + +@pytest.mark.asyncio +async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): + iterator = _build_iterator([]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert response_ids + assert len(set(response_ids)) == 1 + assert response_ids[0].startswith("resp_") 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 050/273] 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 051/273] 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 052/273] 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 053/273] 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 f89a3693baabf3dba081ba213032ffe7acd39b65 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:46:24 -0700 Subject: [PATCH 054/273] fix(responses): resolve previous_response_id for a just-written turn The session lookup reads spend logs straight out of the database, so a follow-up sent seconds after the turn it chains off found nothing while the row was still queued in the worker that served it, and the conversation was dropped without an error. Responses calls now ask the spend-log writer to flush on its next pass instead of waiting out its poll interval, and the lookup gives a just-finished turn a short second chance. Replaying a session also accepted `input` only as a string or a single dict, so the standard list shape dropped every user turn and left the model with assistant messages alone. --- litellm/constants.py | 2 + litellm/proxy/db/db_spend_update_writer.py | 8 +- litellm/proxy/utils.py | 26 ++- .../session_handler.py | 86 +++------ .../proxy/db/test_db_spend_update_writer.py | 28 +++ .../prisma_and_spend/test_spend_functions.py | 50 ++++- .../test_session_handler.py | 176 +++++++----------- 7 files changed, 192 insertions(+), 184 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c33e5a53b76..aaaddd063e7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1542,6 +1542,8 @@ SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BA SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000"))) SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) +RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3"))) +RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 283194bad7c..0c8c9a853ec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -65,6 +65,7 @@ from litellm.proxy.spend_tracking.savings import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable +from litellm.types.utils import CallTypes if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -73,6 +74,9 @@ else: ProxyLogging = Any +RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -820,9 +824,11 @@ class DBSpendUpdateWriter: ) ) if prisma_client is not None and spend_logs_url is not None or prisma_client is not None: - from litellm.proxy.utils import enqueue_spend_logs + from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) + if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: + request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9978fa04f40..86d954c0913 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3341,6 +3341,7 @@ class _StaleReadEngine: class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() + spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event() spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None tool_usage_transactions: list["ToolUsageTransaction"] = [] @@ -6005,14 +6006,24 @@ async def enqueue_spend_logs( ) -async def peek_spend_logs(prisma_client: PrismaClient) -> tuple[SpendLogsPayload, ...]: - """Snapshot the spend logs still waiting for the next flush, leaving the queue intact. +def request_spend_log_flush() -> None: + """Wake the queue monitor now rather than leaving the rows for its next poll. - Reads that need a just-finished request use this, since the batch writer only - reaches the DB every ``PROXY_BATCH_WRITE_AT`` seconds. + The Responses API hands the client an id it can chain from straight away, and that + lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. + Repeated requests coalesce into the monitor's next pass, so the batching holds. """ - async with prisma_client._spend_log_transactions_lock: - return tuple(prisma_client.spend_log_transactions) + PrismaClient.spend_log_flush_requested.set() + + +async def _wait_for_spend_log_flush_request(interval: float) -> bool: + """Wait out ``interval``, returning early and True when a flush was requested.""" + try: + await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval) + except asyncio.TimeoutError: + return False + PrismaClient.spend_log_flush_requested.clear() + return True async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: @@ -6460,7 +6471,8 @@ async def _monitor_spend_logs_queue( # Exponential backoff when no logs to process current_interval = min(current_interval * backoff_multiplier, max_backoff) - await asyncio.sleep(current_interval) + if await _wait_for_spend_log_flush_request(current_interval): + current_interval = base_interval except Exception as e: spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e) # Continue monitoring even if there's an error, with exponential backoff diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 935c78bc9a1..15267533957 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,5 +1,5 @@ +import asyncio import json -from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -132,8 +132,8 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = ResponsesSessionHandler._get_response_dict_from_spend_log(spend_log) - if _response_output: + _response_output: Final = spend_log.get("response", "{}") + if isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -141,23 +141,6 @@ class ResponsesSessionHandler: chat_completion_message_history.append(getattr(choice, "message")) return chat_completion_message_history - @staticmethod - def _get_response_dict_from_spend_log(spend_log: SpendLogsPayload) -> Mapping[str, Any] | None: - """ - Spend logs read from the DB hold `response` as a dict, ones still queued in memory - hold it as a JSON string. - """ - _response_output: Final = spend_log.get("response") - if isinstance(_response_output, dict): - return _response_output or None - if isinstance(_response_output, str): - try: - parsed: Final = json.loads(_response_output) - except json.JSONDecodeError: - return None - return parsed if isinstance(parsed, dict) and parsed else None - return None - @staticmethod async def get_proxy_server_request_from_spend_log( spend_log: SpendLogsPayload, @@ -272,9 +255,16 @@ class ResponsesSessionHandler: SQL query SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + + A just-finished turn gets a short second chance: the worker that served it may + still be writing its spend log when the follow-up arrives, and an empty result + drops the whole conversation instead of erroring. """ + from litellm.constants import ( + RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS, + RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL, + ) from litellm.proxy.proxy_server import prisma_client - from litellm.proxy.utils import peek_spend_logs verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) @@ -295,46 +285,16 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - written_spend_logs: Final = await prisma_client.db.query_raw(query, response_id) - queued_spend_logs: Final = await peek_spend_logs(prisma_client) - spend_logs: Final = list( - ResponsesSessionHandler._merge_queued_spend_logs( - response_id=response_id, - written_spend_logs=written_spend_logs, - queued_spend_logs=queued_spend_logs, - ) - ) + for attempt in range(RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS): + if attempt: + await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) + if spend_logs := await prisma_client.db.query_raw(query, response_id): + verbose_proxy_logger.debug( + "Found the following spend logs for previous response id %s: %s", + response_id, + json.dumps(spend_logs, indent=4, default=str), + ) + return spend_logs - verbose_proxy_logger.debug( - "Found the following spend logs for previous response id %s: %s", - response_id, - json.dumps(spend_logs, indent=4, default=str), - ) - - return spend_logs - - @staticmethod - def _merge_queued_spend_logs( - response_id: str, - written_spend_logs: Sequence[SpendLogsPayload], - queued_spend_logs: Sequence[SpendLogsPayload], - ) -> tuple[SpendLogsPayload, ...]: - """ - Append the session's spend logs that the batch writer has not flushed to the DB yet. - - Without this a follow-up sent inside the ``PROXY_BATCH_WRITE_AT`` window sees an - empty session and silently drops the conversation. The queue is FIFO, so anything - still on it is newer than every row already written. - """ - session_ids: Final = frozenset( - session_id - for spend_log in (*written_spend_logs, *queued_spend_logs) - if spend_log.get("request_id") == response_id and (session_id := spend_log.get("session_id")) - ) | frozenset(session_id for spend_log in written_spend_logs if (session_id := spend_log.get("session_id"))) - written_request_ids: Final = frozenset(spend_log.get("request_id") for spend_log in written_spend_logs) - unflushed: Final = tuple( - spend_log - for spend_log in queued_spend_logs - if spend_log.get("session_id") in session_ids and spend_log.get("request_id") not in written_request_ids - ) - return (*written_spend_logs, *unflushed) + verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) + return [] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 76a80ac2651..ca1827aa38e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2712,3 +2712,31 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey assert mock_prisma_client.db.tx.call_count == 2 proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, expects_flush", + [("aresponses", True), ("responses", True), ("acompletion", False)], +) +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( + call_type: str, expects_flush: bool +): + """ + A `previous_response_id` chained straight off the previous turn reads the DB, so a + Responses row cannot sit in this worker's queue until the monitor's next poll. + """ + from litellm.proxy.utils import PrismaClient + + db_writer = DBSpendUpdateWriter() + prisma = _tool_usage_prisma() + PrismaClient.spend_log_flush_requested.clear() + + await db_writer._insert_spend_log_to_db( + payload={"request_id": "req-1", "call_type": call_type}, + prisma_client=prisma, + ) + + assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] + assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush + PrismaClient.spend_log_flush_requested.clear() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 54d59e690f9..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -11,7 +11,8 @@ Symbols pinned here: from __future__ import annotations import asyncio -from typing import Any, Dict, List +from contextlib import suppress +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock import pytest @@ -526,6 +527,53 @@ async def test_monitor_spend_logs_queue_swallows_errors_and_backs_off( assert sleep_count["n"] == 3 +@pytest.mark.asyncio +async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A requested flush wakes the monitor mid-poll, so a Responses row reaches the DB + before the client can chain a `previous_response_id` off it. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import PrismaClient, request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + PrismaClient.spend_log_flush_requested.clear() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [] + + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.sleep(0.05) + assert not flushed.is_set() + + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) + request_spend_log_flush() + + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + PrismaClient.spend_log_flush_requested.clear() + + def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 926e9e0af2a..4fc288a47cb 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -1,4 +1,3 @@ -import asyncio import json from unittest.mock import AsyncMock, patch @@ -451,20 +450,38 @@ def _chat_completion_response(request_id: str, content: str) -> dict: class _FakePrismaDB: - def __init__(self, rows): - self._rows = rows + def __init__(self, results): + self._results = list(results) self.calls = [] async def query_raw(self, query, *args): self.calls.append(args) - return list(self._rows) + if not self._results: + return [] + return list(self._results.pop(0)) class _FakePrismaClient: - def __init__(self, written_rows, queued_rows): - self.db = _FakePrismaDB(written_rows) - self.spend_log_transactions = list(queued_rows) - self._spend_log_transactions_lock = asyncio.Lock() + def __init__(self, results): + self.db = _FakePrismaDB(results) + + +def _spend_log(request_id: str, session_id: str, prompt: str, answer: str) -> dict: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": { + "input": [{"role": "user", "content": prompt}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, answer), + } + + +@pytest.fixture +def instant_session_lookup_retries(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.constants, "RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.0) @pytest.mark.asyncio @@ -475,21 +492,12 @@ async def test_message_history_reconstructs_list_shaped_input(): """ request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" mock_spend_logs = [ - { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701", - "proxy_server_request": { - "input": [ - { - "role": "user", - "content": "Remember this: my favorite color is chartreuse.", - } - ], - "model": "claude-bridge", - }, - "response": _chat_completion_response(request_id, "OK"), - } + _spend_log( + request_id, + "a96757c4-c6dc-4c76-b37e-e7dfa526b701", + "Remember this: my favorite color is chartreuse.", + "OK", + ) ] with patch.object( @@ -512,31 +520,22 @@ async def test_message_history_reconstructs_list_shaped_input(): @pytest.mark.asyncio -async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_writer(): +async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flushed_yet( + instant_session_lookup_retries: None, +): """ - A follow-up sent right after the previous turn arrives before the batch writer has - flushed that turn's spend log, so the row is only in memory. The history has to - include it anyway. + A follow-up sent right after the previous turn can beat that turn's spend log to the + DB. The lookup has to try again instead of handing back an empty conversation. """ request_id = "chatcmpl-6c1f5f6c-6a2b-4c62-8d1f-0d9d4ce0a1b2" - queued_spend_log = { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3", - "proxy_server_request": json.dumps( - { - "input": [ - { - "role": "user", - "content": "Remember this: my favorite color is chartreuse.", - } - ], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response(request_id, "OK")), - } - fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[queued_spend_log]) + session_id = "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" + spend_log = _spend_log( + request_id, + session_id, + "Remember this: my favorite color is chartreuse.", + "OK", + ) + fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -548,44 +547,22 @@ async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_wr ("user", "Remember this: my favorite color is chartreuse."), ("assistant", "OK"), ] - assert result["litellm_session_id"] == "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" - assert fake_prisma_client.spend_log_transactions == [queued_spend_log] + assert result["litellm_session_id"] == session_id + assert fake_prisma_client.db.calls == [(request_id,), (request_id,)] @pytest.mark.asyncio -async def test_message_history_merges_written_and_queued_turns_in_order(): - """ - Turn 1 already flushed to the DB, turn 2 still queued: the follow-up sees the whole - conversation, in order, with no row counted twice. - """ +async def test_message_history_reconstructs_every_turn_of_the_session_in_order(): session_id = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1" first_request_id = "chatcmpl-1111" second_request_id = "chatcmpl-2222" - written_spend_log = { - "request_id": first_request_id, - "call_type": "aresponses", - "session_id": session_id, - "proxy_server_request": { - "input": [{"role": "user", "content": "My favorite color is chartreuse."}], - "model": "claude-bridge", - }, - "response": _chat_completion_response(first_request_id, "Got it."), - } - queued_spend_log = { - "request_id": second_request_id, - "call_type": "aresponses", - "session_id": session_id, - "proxy_server_request": json.dumps( - { - "input": [{"role": "user", "content": "And my favorite city is Lisbon."}], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response(second_request_id, "Noted.")), - } fake_prisma_client = _FakePrismaClient( - written_rows=[written_spend_log], - queued_rows=[written_spend_log, queued_spend_log], + results=[ + [ + _spend_log(first_request_id, session_id, "My favorite color is chartreuse.", "Got it."), + _spend_log(second_request_id, session_id, "And my favorite city is Lisbon.", "Noted."), + ] + ] ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): @@ -604,45 +581,18 @@ async def test_message_history_merges_written_and_queued_turns_in_order(): @pytest.mark.asyncio -async def test_message_history_ignores_queued_spend_logs_from_other_sessions(): - request_id = "chatcmpl-3333" - written_spend_log = { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "session-a", - "proxy_server_request": { - "input": [{"role": "user", "content": "Hello from session a."}], - "model": "claude-bridge", - }, - "response": _chat_completion_response(request_id, "Hi."), - } - other_session_spend_log = { - "request_id": "chatcmpl-4444", - "call_type": "aresponses", - "session_id": "session-b", - "proxy_server_request": json.dumps( - { - "input": [{"role": "user", "content": "Hello from session b."}], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response("chatcmpl-4444", "Hi there.")), - } - fake_prisma_client = _FakePrismaClient( - written_rows=[written_spend_log], - queued_rows=[other_session_spend_log], - ) +async def test_session_lookup_stops_retrying_once_the_budget_is_spent( + instant_session_lookup_retries: None, +): + fake_prisma_client = _FakePrismaClient(results=[]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" ) - messages = result["messages"] - assert [(message.get("role"), message.get("content")) for message in messages] == [ - ("user", "Hello from session a."), - ("assistant", "Hi."), - ] + assert spend_logs == [] + assert len(fake_prisma_client.db.calls) == litellm.constants.RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS @pytest.mark.asyncio @@ -657,7 +607,9 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[]) + fake_prisma_client = _FakePrismaClient( + results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] + ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( 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 055/273] 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 056/273] 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 9d22acab110fb0407059481cc32755b0d3e095b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:22:10 -0700 Subject: [PATCH 057/273] fix(responses): skip the session lookup retry when spend logs are off --- .../session_handler.py | 8 ++++--- .../test_session_handler.py | 21 +++++++++++++++++++ ...ponse_id.py => test_streaming_iterator.py} | 0 3 files changed, 26 insertions(+), 3 deletions(-) rename tests/test_litellm/responses/litellm_completion_transformation/{test_streaming_iterator_response_id.py => test_streaming_iterator.py} (100%) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 15267533957..59ff492a79f 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -258,13 +258,14 @@ class ResponsesSessionHandler: A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result - drops the whole conversation instead of erroring. + drops the whole conversation instead of erroring. Deployments that write no spend + logs at all have nothing to wait for, so they keep the single original query. """ from litellm.constants import ( RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS, RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL, ) - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import disable_spend_logs, prisma_client verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) @@ -285,7 +286,8 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - for attempt in range(RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS): + max_attempts: Final = 1 if disable_spend_logs else RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS + for attempt in range(max_attempts): if attempt: await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) if spend_logs := await prisma_client.db.query_raw(query, response_id): diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 4fc288a47cb..df477f6d01e 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -617,3 +617,24 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): ) assert fake_prisma_client.db.calls == [(request_id,)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( + instant_session_lookup_retries: None, +): + """ + A deployment that writes no spend logs has nothing to wait for, so the miss path keeps + the single query it always had. + """ + fake_prisma_client = _FakePrismaClient(results=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( + "litellm.proxy.proxy_server.disable_spend_logs", True + ): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" + ) + + assert spend_logs == [] + assert fake_prisma_client.db.calls == [("chatcmpl-does-not-exist",)] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py 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 058/273] 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 d2b5034fea9e90e2258b6d91e191371606717f69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:33:23 -0700 Subject: [PATCH 059/273] test(responses): fold the bridged streaming regressions into the mapped test file --- .../test_streaming_iterator.py | 130 ----------------- ...test_streaming_iterator_transformation.py} | 132 +++++++++++++++++- 2 files changed, 129 insertions(+), 133 deletions(-) delete mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py rename tests/test_litellm/responses/litellm_completion_transformation/{test_tool_call_streaming_transformation.py => test_streaming_iterator_transformation.py} (76%) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py deleted file mode 100644 index 97f35900e9d..00000000000 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py +++ /dev/null @@ -1,130 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, -) -from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - -CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" -RESPONSE_ID_EVENT_TYPES = frozenset( - {"response.created", "response.in_progress", "response.completed"} -) - - -def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: - return ModelResponseStream( - id=CHAT_COMPLETION_ID, - created=1748575031, - model="claude-haiku-4-5", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - index=0, - delta=Delta(role="assistant", content=content), - finish_reason=finish_reason, - ) - ], - ) - - -class _FakeStreamWrapper: - def __init__(self, chunks): - self._chunks = list(chunks) - self.logging_obj = MagicMock() - - def __iter__(self): - return self - - def __next__(self): - if not self._chunks: - raise StopIteration - return self._chunks.pop(0) - - def __aiter__(self): - return self - - async def __anext__(self): - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - -def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: - return LiteLLMCompletionStreamingIterator( - model="claude-haiku-4-5", - litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), - request_input="What is the weather in San Francisco?", - responses_api_request={}, - custom_llm_provider="anthropic", - litellm_metadata={}, - ) - - -def _response_ids(events) -> list[str]: - return [ - event.response.id - for event in events - if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] - - -@pytest.mark.asyncio -async def test_streaming_events_share_the_chat_completion_response_id(): - """ - Every event of a bridged stream has to carry the same id, and that id has to decode - to the chat completion id spend tracking stores as `request_id`. Otherwise a - follow-up `previous_response_id` matches no session and the conversation is dropped. - """ - iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) - - events = [event async for event in iterator] - - response_ids = _response_ids(events) - assert len(response_ids) == 3 - assert len(set(response_ids)) == 1 - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) - assert decoded["response_id"] == CHAT_COMPLETION_ID - assert decoded["custom_llm_provider"] == "anthropic" - - -def test_sync_streaming_events_share_the_chat_completion_response_id(): - iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) - - events = list(iterator) - - response_ids = _response_ids(events) - assert len(response_ids) == 3 - assert len(set(response_ids)) == 1 - assert ( - ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] - == CHAT_COMPLETION_ID - ) - - -@pytest.mark.asyncio -async def test_streaming_emits_every_chunk_after_priming_the_response_id(): - iterator = _build_iterator( - [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] - ) - - events = [event async for event in iterator] - - deltas = "".join( - event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) - assert deltas == "Hello!" - - -@pytest.mark.asyncio -async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): - iterator = _build_iterator([]) - - events = [event async for event in iterator] - - response_ids = _response_ids(events) - assert response_ids - assert len(set(response_ids)) == 1 - assert response_ids[0].startswith("resp_") diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py similarity index 76% rename from tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index fa6f42609ca..823f656ddc5 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1,18 +1,23 @@ """ -Tests for streaming tool-calls in Responses API transformation. +Tests for the Responses API streaming bridge in +litellm/responses/litellm_completion_transformation/streaming_iterator.py. Ensures that when the underlying chat-completions stream includes tool_calls deltas, LiteLLM emits Responses API streaming events (output_item.added + function_call_arguments.*). Also ensures that tool calls that only appear in the final built response still get emitted -before response.completed. +before response.completed, and that every event of a bridged stream carries the response id +spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock + +import pytest from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import ( Delta, @@ -21,6 +26,68 @@ from litellm.types.utils import ( StreamingChoices, ) +CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" +RESPONSE_ID_EVENT_TYPES = frozenset( + {"response.created", "response.in_progress", "response.completed"} +) + + +def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +class _FakeStreamWrapper: + def __init__(self, chunks): + self._chunks = list(chunks) + self.logging_obj = MagicMock() + + def __iter__(self): + return self + + def __next__(self): + if not self._chunks: + raise StopIteration + return self._chunks.pop(0) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + +def _response_ids(events) -> list[str]: + return [ + event.response.id + for event in events + if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + def test_tool_call_delta_is_emitted_as_responses_events(): iterator = LiteLLMCompletionStreamingIterator( @@ -397,3 +464,62 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): assert arguments_by_call_id["call_b"] == '{"b":' assert arguments_by_call_id["call_a"] != '{"a":1}' assert arguments_by_call_id["call_b"] != '{"b":1}' + + +@pytest.mark.asyncio +async def test_streaming_events_share_the_chat_completion_response_id(): + """ + Every event of a bridged stream has to carry the same id, and that id has to decode + to the chat completion id spend tracking stores as `request_id`. Otherwise a + follow-up `previous_response_id` matches no session and the conversation is dropped. + """ + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) + assert decoded["response_id"] == CHAT_COMPLETION_ID + assert decoded["custom_llm_provider"] == "anthropic" + + +def test_sync_streaming_events_share_the_chat_completion_response_id(): + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = list(iterator) + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + assert ( + ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] + == CHAT_COMPLETION_ID + ) + + +@pytest.mark.asyncio +async def test_streaming_emits_every_chunk_after_priming_the_response_id(): + iterator = _build_iterator( + [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] + ) + + events = [event async for event in iterator] + + deltas = "".join( + event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert deltas == "Hello!" + + +@pytest.mark.asyncio +async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): + iterator = _build_iterator([]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert response_ids + assert len(set(response_ids)) == 1 + assert response_ids[0].startswith("resp_") 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 060/273] 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 061/273] 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 062/273] 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 063/273] 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 064/273] 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 065/273] 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 066/273] 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 23e64c8b3d89b781642703e4543aa45e56002881 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:25:59 -0700 Subject: [PATCH 067/273] chore(responses): keep the session lookup inside the type-discipline budget --- .../litellm_completion_transformation/session_handler.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 2008006e6cf..1566bb1bdd7 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -108,7 +108,10 @@ class ResponsesSessionHandler: if isinstance(_response_input_param, (str, list)): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast(ResponseInputParam, [_response_input_param]) + response_input_param = cast( + ResponseInputParam, + [_response_input_param], # mutable-ok: a lone input item still has to arrive as a list + ) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -301,4 +304,4 @@ class ResponsesSessionHandler: return spend_logs verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) - return [] + return [] # mutable-ok: an empty result the caller only reads 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 068/273] 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 069/273] 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 070/273] 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 071/273] 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 072/273] 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 073/273] 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 074/273] 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 075/273] 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 20e92d1e68c10c6e856b2618aa58341947abd587 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 22 Aug 2026 23:10:36 +0000 Subject: [PATCH 076/273] fix(anthropic/bedrock): request summarized adaptive thinking for reasoning_effort and use provider thinking token counts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 9 +- .../bedrock/chat/converse_transformation.py | 24 ++++- litellm/llms/bedrock/chat/invoke_handler.py | 5 ++ litellm/types/llms/anthropic.py | 1 + .../test_reasoning_effort_translation.py | 2 +- .../test_anthropic_reasoning_effort.py | 12 +++ ...azure_anthropic_messages_transformation.py | 2 +- .../chat/test_converse_transformation.py | 90 +++++++++++++++++++ .../llms/bedrock/chat/test_invoke_handler.py | 23 +++++ .../test_anthropic_claude3_transformation.py | 4 +- ...artner_models_anthropic_messages_config.py | 2 +- 11 files changed, 165 insertions(+), 9 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ef278c8f723..27caa9efc44 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1184,8 +1184,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if reasoning_effort is None or reasoning_effort == "none": return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): + # without display, Anthropic defaults adaptive thinking to + # display="omitted" and returns a blank thinking block return AnthropicThinkingParam( type="adaptive", + display="summarized", ) elif reasoning_effort == "low": return AnthropicThinkingParam( @@ -2113,7 +2116,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: details: Final = usage_object.get("output_tokens_details") if not isinstance(details, Mapping): return None @@ -2145,7 +2148,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reported_thinking_tokens: Final = ( iteration_thinking_tokens if iteration_thinking_tokens is not None - else self._thinking_tokens_from_usage(usage_object) + else self.thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) @@ -2168,7 +2171,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b437e25d24b..52366da8c35 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params + if "thinking" in additional_request_params: + data["additionalModelResponseFieldPaths"] = ["/usage/output_tokens_details"] if system_content_blocks: data["system"] = system_content_blocks @@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None: + """Converse omits thinking tokens from its usage block; they only arrive under + ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested.""" + if not isinstance(additional_fields, Mapping): + return None + usage: Final = additional_fields.get("usage") + if not isinstance(usage, Mapping): + return None + return AnthropicConfig.thinking_tokens_from_usage(usage) + @staticmethod def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool: """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" @@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig): usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, thinking_ran: bool = False, + provider_reasoning_tokens: int | None = None, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens: Final = ( + estimated_reasoning_tokens: Final = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens: Final = ( + min(max(0, provider_reasoning_tokens), output_tokens) + if provider_reasoning_tokens is not None + else estimated_reasoning_tokens + ) completion_tokens_details: Final = ( CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, @@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig): completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), thinking_ran=reasoningContentBlocks is not None, + provider_reasoning_tokens=self.thinking_tokens_from_additional_fields( + completion_response.get("additionalModelResponseFields") + ), ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index ce89c6c23e2..3937b36aca0 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -331,6 +331,7 @@ class AWSEventStreamDecoder: self.json_mode = json_mode self._current_tool_name: str | None = None self._thinking_ran = False + self._provider_reasoning_tokens: int | None = None def check_empty_tool_call_args(self) -> bool: """ @@ -559,10 +560,14 @@ class AWSEventStreamDecoder: tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) + self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields( + chunk_data.get("additionalModelResponseFields") + ) elif "usage" in chunk_data: usage = converse_config.transform_usage( chunk_data.get("usage", {}), thinking_ran=self._thinking_ran, + provider_reasoning_tokens=self._provider_reasoning_tokens, ) if thinking_blocks: self._thinking_ran = True diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cc6eccbf3e0..d3b0f334163 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int + display: ReadOnly[Literal["summarized", "omitted"]] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index f393a7b50b1..48a96d011d5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -44,7 +44,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py index ef74249ca8e..288817dff07 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +import pytest + from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -35,6 +37,16 @@ class TestMapReasoningEffort: ) assert result["type"] == "adaptive" + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_mapping_requests_summarized_display(self, effort): + """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic + return a blank thinking block, so reasoning_effort callers always got + ``reasoning_content: ""``.""" + result = AnthropicConfig._map_reasoning_effort( + reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic" + ) + assert result["display"] == "summarized" + def test_other_model_low_returns_enabled_with_budget(self): result = AnthropicConfig._map_reasoning_effort( reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 53a432427d3..326edde743d 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 604f3414775..b648b6322f7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_reasoning_effort_requests_summarized_display_converse(): + """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must + request the summarized display, otherwise the provider returns a blank thinking + block and reasoning_content is always empty.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-7", + drop_params=False, + ) + + assert optional_params["thinking"]["type"] == "adaptive" + assert optional_params["thinking"]["display"] == "summarized" + + +def test_thinking_request_adds_output_tokens_details_response_path(): + """Regression LIT-5714: the Converse usage block has no thinking-token field, so + thinking requests must ask for ``/usage/output_tokens_details`` via + ``additionalModelResponseFieldPaths``.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["additionalModelResponseFieldPaths"] == ["/usage/output_tokens_details"] + + +def test_request_without_thinking_omits_response_field_paths(): + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={"maxTokens": 256}, + litellm_params={}, + headers={}, + ) + + assert "additionalModelResponseFieldPaths" not in result + + +def test_transform_usage_prefers_provider_reasoning_tokens(): + """Regression LIT-5714: provider-reported thinking tokens must win over the + token_counter estimate derived from visible reasoning text.""" + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + provider_reasoning_tokens=1033, + ) + + assert usage.completion_tokens_details.reasoning_tokens == 1033 + assert usage.completion_tokens_details.text_tokens == 3002 - 1033 + + +def test_transform_usage_falls_back_to_estimate_without_provider_tokens(): + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + ) + + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert usage.completion_tokens_details.reasoning_tokens < 300 + + +def test_thinking_tokens_parsed_from_additional_model_response_fields(): + parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields( + {"usage": {"output_tokens_details": {"thinking_tokens": 92}}} + ) + assert parsed == 92 + assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None + assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None + + @pytest.mark.parametrize( "model,effort,expected_effort", [ 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 e2892a6ccee..2c5ff118c85 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -206,6 +206,29 @@ def test_bedrock_converse_streaming_consistent_id(): ), "All chunk IDs must match the one captured from the messageStart event" +def test_converse_streaming_usage_uses_provider_thinking_tokens(): + """Regression LIT-5714: the messageStop event carries provider thinking tokens + under ``additionalModelResponseFields``; the usage chunk must report them instead + of a token_counter estimate.""" + chunks = [ + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "thinking about it"}}, + }, + { + "stopReason": "end_turn", + "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}}, + }, + {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}}, + ] + + decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7") + parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks] + + usage = parsed[-1].usage + assert usage.completion_tokens_details.reasoning_tokens == 1033 + + @pytest.mark.asyncio async def test_make_call_does_not_rechunk_stream_by_default(): """Re-chunking the event stream into fixed 1024-byte blocks holds small diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3c28302bf9..1e09afd6919 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ba2f20e2337..f19e169dc9e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( 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 077/273] 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 078/273] 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 079/273] 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 418e8ca5e8db8bd8a0e916d6579aec3279cd2395 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 00:00:21 +0000 Subject: [PATCH 080/273] fix(bedrock): build response field paths as an immutable sequence to satisfy the type discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/chat/converse_transformation.py | 2 +- litellm/types/llms/bedrock.py | 3 ++- .../llms/bedrock/chat/test_converse_transformation.py | 2 +- type-discipline-budget.json | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 52366da8c35..767677cbcbf 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1618,7 +1618,7 @@ class AmazonConverseConfig(BaseConfig): if additional_request_params: data["additionalModelRequestFields"] = additional_request_params if "thinking" in additional_request_params: - data["additionalModelResponseFieldPaths"] = ["/usage/output_tokens_details"] + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..6ae2e31fe60 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal @@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False): class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict - additionalModelResponseFieldPaths: list[str] + additionalModelResponseFieldPaths: Sequence[str] inferenceConfig: InferenceConfig system: list[SystemContentBlock] toolConfig: ToolConfigBlock diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index b648b6322f7..4d2c077b548 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -401,7 +401,7 @@ def test_thinking_request_adds_output_tokens_details_response_path(): headers={}, ) - assert result["additionalModelResponseFieldPaths"] == ["/usage/output_tokens_details"] + assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) def test_request_without_thinking_omits_response_field_paths(): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..05098546325 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { "limit": 26873 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 081/273] 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 082/273] 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 083/273] 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 084/273] 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 085/273] 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 086/273] 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 d447be15b965a052e49a2d3c0b9f538d85c9d865 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 22 Aug 2026 19:13:48 -0700 Subject: [PATCH 087/273] feat(newrelic): per-team New Relic trace routing via team callbacks (#37603) --- litellm/integrations/callback_configs.json | 14 +- litellm/integrations/otel/emitter.py | 6 + litellm/integrations/otel/logger.py | 36 +++- litellm/integrations/otel/model/config.py | 10 + litellm/integrations/otel/model/metadata.py | 6 + .../integrations/otel/plumbing/providers.py | 2 + litellm/integrations/otel/plumbing/routing.py | 29 ++- litellm/integrations/otel/presets/__init__.py | 70 +++++-- litellm/integrations/otel/presets/newrelic.py | 104 +++++++++++ .../initialize_dynamic_callback_params.py | 32 +++- litellm/litellm_core_utils/litellm_logging.py | 7 + .../callback_config_validation.py | 77 ++++++++ litellm/proxy/litellm_pre_call_utils.py | 6 + .../key_management_endpoints.py | 15 ++ .../team_callback_endpoints.py | 13 ++ litellm/types/utils.py | 5 + .../integrations/otel/test_otel_v2_dynamic.py | 164 +++++++++++----- .../integrations/otel/test_otel_v2_logger.py | 175 ++++++++++++++++++ .../integrations/otel/test_otel_v2_presets.py | 52 ++++++ ...test_initialize_dynamic_callback_params.py | 46 +++++ .../test_litellm_logging.py | 78 ++++++++ .../test_callback_management_endpoints.py | 145 +++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 89 +++++++++ .../src/components/callback_info_helpers.tsx | 12 ++ 24 files changed, 1111 insertions(+), 82 deletions(-) create mode 100644 litellm/integrations/otel/presets/newrelic.py create mode 100644 litellm/proxy/common_utils/callback_config_validation.py diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 590c848767a..6d2bcea8bae 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -294,12 +294,18 @@ "id": "newrelic", "displayName": "New Relic", "logo": "newrelic.png", - "supports_key_team_logging": false, + "supports_key_team_logging": true, "dynamic_params": { - "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": { + "newrelic_api_key": { + "type": "password", + "ui_name": "New Relic Ingest License Key", + "description": "Per-team ingest (license) key. Team traces export to this key's New Relic account over OTLP.", + "required": false + }, + "newrelic_region": { "type": "text", - "ui_name": "Record AI Content (default: true)", - "description": "Whether to record AI message content. Set to false to disable.", + "ui_name": "New Relic Region (us or eu)", + "description": "Data center region for this team's account. Defaults to us.", "required": false } }, diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index 0850d867c7b..244e58eddf3 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -156,6 +156,12 @@ class SpanEmitter: links=list(links) if links else None, ) + def mark_emitted(self, dedup_key: str | None, role: SpanRole) -> None: + """Register a span emitted outside :meth:`emit` (the boundary-opened + LLM-call span closed via :meth:`finish_span`) so a later :meth:`emit` + for the same ``(dedup_key, role)`` deduplicates against it.""" + self._seen(dedup_key, role) + def _seen(self, dedup_key: str | None, role: SpanRole) -> bool: """Return True once a ``(dedup_key, role)`` pair has been emitted. diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 53b9829023c..4359b222d06 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -484,10 +484,15 @@ class OpenTelemetryV2(CustomLogger): # ``pop`` is the dedup: this method runs from both the success and failure # paths, and whichever fires first removes the carrier and closes the span. carrier: Final = self._open_llm_calls.pop(call_id, None) if call_id else None - if carrier is None: + # A missing carrier does not always mean nothing happened: a team/key-scoped + # logger is a success/failure callback only, so ``pre_call`` never reaches it + # and no carrier exists. The payload plus the request-level provider-handoff + # stamp (``upstream_started``) is the affirmative signal of a real call; a + # gate rejection carries ``is_no_upstream_call`` and gets no span. + if carrier is None and (call.is_no_upstream_call or not call.upstream_started or call.payload is None): return None try: - return self._finish_carrier(carrier, call, end_time) + return self._finish_carrier(carrier, call, start_time, end_time) finally: # After the span has ended, so a release-triggered provider shutdown # force-flushes it out rather than racing its enqueue. @@ -497,8 +502,11 @@ class OpenTelemetryV2(CustomLogger): """Remember an in-flight LLM call, evicting the oldest if over budget. A call that opens but never closes (a stream that only fires stream - events) would linger otherwise; the evicted span is simply dropped - (never exported). + events) would linger otherwise. Eviction only drops the boundary carrier, + not the call: if that call later closes as a real completed call, it still + emits through the deferred branch in ``_close_llm_call`` (the same path a + team/key-scoped logger uses, since it never opens a carrier), deduplicated + by call id. Only a call that is evicted and never closes goes unexported. """ self._open_llm_calls[call_id] = carrier if len(self._open_llm_calls) > _OPEN_CALLS_MAX: @@ -512,15 +520,20 @@ class OpenTelemetryV2(CustomLogger): def _finish_carrier( self, - carrier: _LLMCallSpan, + carrier: "_LLMCallSpan | None", call: LLMCallEvent, + start_time: datetime | float | None, end_time: datetime | float | None, ) -> Span | None: payload: Final = call.payload + call_id: Final = call.call_id if payload is None: - if carrier.span is not None: + if carrier is not None and carrier.span is not None: # Opened at the boundary but the payload never materialized — end - # it (named provisionally) so it isn't leaked as an open span. + # it (named provisionally) so it isn't leaked as an open span, and + # register the dedup marker so a later payload-carrying close for + # the same call id cannot re-emit through the deferred branch. + self._emitter.mark_emitted(call_id, SpanRole.LLM_CALL) carrier.span.end(end_time=to_ns(end_time)) return None data: Final = LLMCallSpanData.from_standard_logging_payload( @@ -529,10 +542,13 @@ class OpenTelemetryV2(CustomLogger): time_to_first_chunk_seconds=call.time_to_first_chunk_seconds, ) end_time_ns: Final = to_ns(end_time) - if carrier.span is not None: + if carrier is not None and carrier.span is not None: # Born at the boundary: stamp attributes from the typed payload, set # status, and end it. Its parent (the server span) was captured at - # creation from real ambient context. + # creation from real ambient context. Register the dedup marker so a + # second close for the same call id (success then failure on one + # logging object) cannot re-emit through the deferred branch. + self._emitter.mark_emitted(call_id, SpanRole.LLM_CALL) self._emitter.finish_span(SpanRole.LLM_CALL, carrier.span, data, end_time_ns=end_time_ns) return carrier.span # Deferred: ``pre_call`` saw no recordable parent, so create the span now. @@ -549,7 +565,7 @@ class OpenTelemetryV2(CustomLogger): SpanRole.LLM_CALL, data, parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx), - start_time_ns=carrier.start_time_ns, + start_time_ns=(carrier.start_time_ns if carrier is not None else to_ns(start_time)), end_time_ns=end_time_ns, tracer=route.tracer, links=_request_trace_links(parent_ctx) if route.detached else None, diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 53178b48991..9e3064c2bff 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -39,6 +39,7 @@ class ExporterOwner(str, Enum): WEAVE_OTEL = "weave_otel" LEVO = "levo" AGENTOPS = "agentops" + NEWRELIC = "newrelic" class _OTelV2Flag(BaseSettings): @@ -97,6 +98,15 @@ class ExporterSpec(BaseModel): "auto (Simple for console/in_memory, Batch otherwise)." ), ) + requires_headers: bool = Field( + default=False, + description=( + "Skip this exporter when no headers are resolved. For destinations " + "that reject unauthenticated exports (e.g. New Relic), a spec kept " + "only as the per-request credential-stamping target would otherwise " + "export keyless traffic and produce a 4xx for every span batch." + ), + ) class OpenTelemetryV2Config(BaseSettings): diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index de6366e7dbd..062b2ca20b4 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -203,6 +203,11 @@ class LLMCallEvent: # True for synthetic proxy-gate logs (auth / rate-limit rejections): they fire # the ``pre_call`` hook but never made an upstream call, so they get no span. is_no_upstream_call: bool + # True once the request handed off to a provider (``pre_call`` stamped + # ``api_call_start_time``). The affirmative signal that an LLM call was + # actually attempted — router pre-call rejections, SDK failures before the + # provider handoff, and standalone guardrail runs all lack it. + upstream_started: bool # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The # span is renamed from the typed payload at close (``finish_span``); this only # needs to be reasonable for a span that never gets closed (a leak). @@ -221,6 +226,7 @@ class LLMCallEvent: dynamic_params=kwargs.get("standard_callback_dynamic_params"), auth_metadata=auth_metadata(payload, kwargs), is_no_upstream_call=bool(kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)), + upstream_started=kwargs.get("api_call_start_time") is not None, provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), ) diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 80a8d01c061..81b00f788ee 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -436,6 +436,8 @@ def build_tracer_provider( # ``config._normalize`` guarantees at least one spec (it folds the top-level # ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty). for spec in config.exporters: + if spec.requires_headers and not spec.headers: + continue exp = _exporter_from_spec(spec) provider.add_span_processor( _processor_for( diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index f231df9e914..6a04dbb9bc8 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -28,6 +28,7 @@ from litellm.integrations.otel.plumbing.providers import ( get_tracer, ) from litellm.integrations.otel.presets import ( + dynamic_otlp_endpoint, dynamic_otlp_headers, project_routing_headers, ) @@ -129,7 +130,9 @@ class TenantTracerCache: # thread-pool workers concurrently with the event loop, so cache # updates, span counts, and retirement must be atomic. self._lock: Final = threading.Lock() - self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems], TracerProvider] = OrderedDict() + self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems, str | None], TracerProvider] = ( + OrderedDict() # mutable-ok: bounded LRU; eviction needs in-place ordered mutation + ) self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state # Oldest-first so an overflow of draining providers sheds the stalest. self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers @@ -182,12 +185,16 @@ class TenantTracerCache: project_headers: Final = self._project_headers(auth_metadata) if not credential_headers and not project_headers: return TenantRoute(tracer=default, detached=False) + # A fixed per-integration region endpoint (New Relic us/eu), never a + # caller-supplied host; ``None`` keeps the preset's own endpoint. + endpoint: Final = dynamic_otlp_endpoint(self._callback_name, dynamic_params) cache_key: Final = ( tuple(sorted(credential_headers.items())), tuple(sorted(project_headers.items())), + endpoint, ) with self._lock: - provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers) + provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers, endpoint) self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1 evicted: Final = self._evicted_on_overflow_locked() if evicted is not None: @@ -200,15 +207,16 @@ class TenantTracerCache: def _cached_provider_locked( self, - cache_key: tuple[_HeaderItems, _HeaderItems], + cache_key: tuple[_HeaderItems, _HeaderItems, str | None], credential_headers: Mapping[str, str], project_headers: Mapping[str, str], + endpoint: str | None, ) -> TracerProvider: cached: Final = self._providers.get(cache_key) if cached is not None: self._providers.move_to_end(cache_key) return cached - built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers)) + built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers, endpoint)) self._providers[cache_key] = built return built @@ -257,6 +265,7 @@ class TenantTracerCache: self, credential_headers: Mapping[str, str], project_headers: Mapping[str, str], + endpoint: str | None = None, ) -> OpenTelemetryV2Config: """Clone the config, rewriting headers on the callback's own exporter. @@ -272,7 +281,8 @@ class TenantTracerCache: ``Authorization``), which must survive routing to a project. """ exporters: Final = [ - self._routed_exporter(spec, credential_headers, project_headers) for spec in self._config.exporters + self._routed_exporter(spec, credential_headers, project_headers, endpoint) + for spec in self._config.exporters ] return self._config.model_copy(update={"exporters": exporters}) @@ -281,6 +291,7 @@ class TenantTracerCache: spec: ExporterSpec, credential_headers: Mapping[str, str], project_headers: Mapping[str, str], + endpoint: str | None = None, ) -> ExporterSpec: kind: Final = spec.kind.lower() if spec.owner != self._callback_name or kind in _NON_OTLP_KINDS: @@ -291,4 +302,10 @@ class TenantTracerCache: if project_headers and kind not in _GRPC_KINDS else base ) - return spec if routed == spec.headers else spec.model_copy(update={"headers": routed}) + update: Final = { # mutable-ok: model_copy(update=...) requires a plain dict + field: value + for field, value in (("headers", routed), ("endpoint", endpoint)) + if (field == "headers" and routed != spec.headers) + or (field == "endpoint" and endpoint is not None and endpoint != spec.endpoint) + } + return spec if not update else spec.model_copy(update=update) diff --git a/litellm/integrations/otel/presets/__init__.py b/litellm/integrations/otel/presets/__init__.py index 35b0584c697..a0cd5b3fd98 100644 --- a/litellm/integrations/otel/presets/__init__.py +++ b/litellm/integrations/otel/presets/__init__.py @@ -21,6 +21,11 @@ from litellm.integrations.otel.presets.langfuse import ( ) from litellm.integrations.otel.presets.langtrace import langtrace_preset from litellm.integrations.otel.presets.levo import levo_preset +from litellm.integrations.otel.presets.newrelic import ( + newrelic_dynamic_endpoint, + newrelic_dynamic_headers, + newrelic_preset, +) from litellm.integrations.otel.presets.phoenix import ( phoenix_preset, phoenix_project_headers, @@ -30,25 +35,45 @@ from litellm.types.utils import StandardCallbackDynamicParams #: Callback name → preset. The ``Preset`` annotation makes mypy verify every #: registered value matches the preset interface. -PRESET_BY_CALLBACK: Final[dict[str, Preset]] = { - "agentops": agentops_preset, - "arize": arize_preset, - "arize_phoenix": phoenix_preset, - "langfuse_otel": langfuse_preset, - "langtrace": langtrace_preset, - "levo": levo_preset, - "weave_otel": weave_preset, -} +PRESET_BY_CALLBACK: Final[Mapping[str, Preset]] = MappingProxyType( + { + "agentops": agentops_preset, + "arize": arize_preset, + "arize_phoenix": phoenix_preset, + "langfuse_otel": langfuse_preset, + "langtrace": langtrace_preset, + "levo": levo_preset, + "newrelic": newrelic_preset, + "weave_otel": weave_preset, + } +) #: Callback name → per-request OTLP header builder (team/key multi-tenant #: routing). Only integrations that support dynamic credentials appear here — #: Arize-Phoenix/Langtrace/Levo/AgentOps don't, so they use the logger's #: default tracer. -DYNAMIC_HEADERS_BY_CALLBACK: Final[dict[str, Callable[[StandardCallbackDynamicParams], dict[str, str]]]] = { - "arize": arize_dynamic_headers, - "langfuse_otel": langfuse_dynamic_headers, - "weave_otel": weave_dynamic_headers, -} +DYNAMIC_HEADERS_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], dict[str, str]]]] = ( + MappingProxyType( + { + "arize": arize_dynamic_headers, + "langfuse_otel": langfuse_dynamic_headers, + "newrelic": newrelic_dynamic_headers, + "weave_otel": weave_dynamic_headers, + } + ) +) + +#: Callback name → per-request OTLP endpoint resolver. Only integrations whose +#: destination host varies per tenant (from a fixed region table, never a +#: caller-supplied URL) appear here; for everyone else the preset's endpoint is +#: authoritative. +DYNAMIC_ENDPOINT_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], str | None]]] = ( + MappingProxyType( + { + "newrelic": newrelic_dynamic_endpoint, + } + ) +) #: Callback name → per-request *routing* header builder, sourced from the key/team @@ -98,17 +123,34 @@ def project_routing_headers( return builder(auth_metadata) +def dynamic_otlp_endpoint( + callback_name: str | None, + dynamic_params: StandardCallbackDynamicParams | None, +) -> str | None: + """Per-request OTLP endpoint for ``callback_name``, or ``None`` if N/A. + + ``None`` means "keep the preset's own endpoint". + """ + resolver: Final = DYNAMIC_ENDPOINT_BY_CALLBACK.get(callback_name or "") + if resolver is None or not dynamic_params: + return None + return resolver(dynamic_params) + + __all__ = [ + "DYNAMIC_ENDPOINT_BY_CALLBACK", "DYNAMIC_HEADERS_BY_CALLBACK", "PRESET_BY_CALLBACK", "PROJECT_HEADERS_BY_CALLBACK", "Preset", "agentops_preset", "arize_preset", + "dynamic_otlp_endpoint", "dynamic_otlp_headers", "langfuse_preset", "langtrace_preset", "levo_preset", + "newrelic_preset", "phoenix_preset", "project_routing_headers", "weave_preset", diff --git a/litellm/integrations/otel/presets/newrelic.py b/litellm/integrations/otel/presets/newrelic.py new file mode 100644 index 00000000000..4660a707355 --- /dev/null +++ b/litellm/integrations/otel/presets/newrelic.py @@ -0,0 +1,104 @@ +"""New Relic preset — OTLP/HTTP exporter to New Relic + GenAI vocabulary.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_logger +from litellm.integrations.otel.model.config import ( + ExporterOwner, + ExporterSpec, + OpenTelemetryV2Config, +) +from litellm.integrations.otel.presets.utils import ensure_mappers +from litellm.types.utils import StandardCallbackDynamicParams + +#: Region -> OTLP base endpoint. A fixed table by design: team config picks a +#: region enum rather than a free-form endpoint, so callback vars can never +#: redirect telemetry to an arbitrary host. +NEWRELIC_OTLP_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType( + { + "us": "https://otlp.nr-data.net", + "eu": "https://otlp.eu01.nr-data.net", + } +) + +_DEFAULT_REGION: Final = "us" + + +class _NewRelicSettings(BaseSettings): + model_config = SettingsConfigDict(case_sensitive=False, extra="ignore") + + # The same env vars the agent-based integration documents; the key is the + # operator-level fallback for traffic without team credentials, the region + # picks that fallback's data center, and the record-content flag keeps its + # documented meaning when the OTel path replaces the agent. + license_key: str | None = Field(default=None, validation_alias="NEW_RELIC_LICENSE_KEY") + region: str | None = Field(default=None, validation_alias="NEW_RELIC_REGION") + record_content: bool | None = Field(default=None, validation_alias="NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED") + + +def newrelic_preset( + *, + config_overrides: OpenTelemetryV2Config | None = None, +) -> OpenTelemetryV2Config: + settings: Final = _NewRelicSettings() + base: Final = config_overrides or OpenTelemetryV2Config() + endpoint: Final = NEWRELIC_OTLP_ENDPOINT_BY_REGION.get( + (settings.region or _DEFAULT_REGION).lower(), NEWRELIC_OTLP_ENDPOINT_BY_REGION[_DEFAULT_REGION] + ) + return base.model_copy( + update={ + "exporters": [ + *base.exporters, + ExporterSpec( + kind="otlp_http", + endpoint=endpoint, + headers=(f"api-key={settings.license_key}" if settings.license_key else None), + owner=ExporterOwner.NEWRELIC, + requires_headers=True, + ), + ], + # New Relic ingests the OTLP GenAI semantic conventions natively. + "mapper_names": ensure_mappers(base.mapper_names, "genai"), + **( + {"capture_message_content": ("span_only" if settings.record_content else "no_content")} + if settings.record_content is not None + else {} + ), + } + ) + + +def newrelic_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str, str]: + """Per-request New Relic OTLP headers from team/key dynamic params.""" + api_key: Final = params.get("newrelic_api_key") + return {header: value for header, value in (("api-key", api_key),) if value} + + +def newrelic_dynamic_endpoint(params: StandardCallbackDynamicParams) -> str: + """Per-request OTLP endpoint for the team's ``newrelic_region``. + + Always the team's own region endpoint, defaulting to US when the team left + the region unset. It never falls through to the preset's endpoint, which + follows the operator's ``NEW_RELIC_REGION`` env; a team that saved only its + ingest key must not inherit the operator's region and have its US-account + spans rejected by an EU-configured default (or vice versa). An unknown + region likewise resolves to the documented US default rather than a guess. + """ + region: Final = params.get("newrelic_region") + default_endpoint: Final = NEWRELIC_OTLP_ENDPOINT_BY_REGION[_DEFAULT_REGION] + if not region: + return default_endpoint + endpoint: Final = NEWRELIC_OTLP_ENDPOINT_BY_REGION.get(region.lower()) + if endpoint is None: + verbose_logger.warning( + "New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.", + region, + ", ".join(sorted(NEWRELIC_OTLP_ENDPOINT_BY_REGION)), + ) + return default_endpoint + return endpoint diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 284989ab20f..3b42ca4eaaf 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -46,7 +46,7 @@ def validate_no_callback_env_reference(param: str, value: object, *, source: str # Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict -_supported_callback_params: Final = [ +_supported_callback_params: Final[tuple[str, ...]] = ( "langfuse_public_key", "langfuse_secret", "langfuse_secret_key", @@ -72,8 +72,10 @@ _supported_callback_params: Final = [ "dd_site", "dd_agent_host", "dd_agent_port", + "newrelic_api_key", + "newrelic_region", "turn_off_message_logging", -] +) _request_blocked_callback_params: Final = frozenset( { @@ -83,6 +85,20 @@ _request_blocked_callback_params: Final = frozenset( "dd_site", "dd_agent_host", "dd_agent_port", + "newrelic_api_key", + "newrelic_region", + } +) + +# Request-blocked params that must still reach ``standard_callback_dynamic_params`` +# when the proxy itself stamped them from admin-configured team/key callback +# settings (the trusted-vars channel). The OTel per-tenant tracer routing reads +# ``standard_callback_dynamic_params``, so without this overlay a blocked param +# could never drive routing at all. +_trusted_overlay_callback_params: Final = frozenset( + { + "newrelic_api_key", + "newrelic_region", } ) @@ -121,7 +137,9 @@ def initialize_standard_callback_dynamic_params( if param in kwargs: _param_value = kwargs.get(param) validate_no_callback_env_reference(param, _param_value, source="request body") - standard_callback_dynamic_params[param] = _param_value + standard_callback_dynamic_params[param] = ( # pyright: ignore[reportGeneralTypeIssues] # several supported params predate their StandardCallbackDynamicParams fields + _param_value + ) for slot_label, metadata in iter_client_callback_metadata_dicts(kwargs): for param in _supported_callback_params: @@ -130,6 +148,12 @@ def initialize_standard_callback_dynamic_params( if param not in standard_callback_dynamic_params and param in metadata: _param_value = metadata.get(param) validate_no_callback_env_reference(param, _param_value, source=slot_label) - standard_callback_dynamic_params[param] = _param_value + standard_callback_dynamic_params[param] = ( # pyright: ignore[reportGeneralTypeIssues] # several supported params predate their StandardCallbackDynamicParams fields + _param_value + ) + + for param, trusted_value in get_trusted_callback_params(kwargs): + if param in _trusted_overlay_callback_params: + standard_callback_dynamic_params[param] = trusted_value return standard_callback_dynamic_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c14dd6c3d8b..3ad4c187b6d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4503,6 +4503,9 @@ def _init_custom_logger_compatible_class( _in_memory_loggers.append(gitlab_logger) return gitlab_logger elif logging_integration == "newrelic": + _v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers) + if _v2 is not None: + return _v2 for callback in _in_memory_loggers: if isinstance(callback, NewRelicLogger): return callback @@ -4789,7 +4792,11 @@ def get_custom_logger_compatible_class( if isinstance(callback, SMTPEmailLogger): return callback elif logging_integration == "newrelic": + from litellm.integrations.otel.logger import OpenTelemetryV2 + for callback in _in_memory_loggers: + if isinstance(callback, OpenTelemetryV2) and callback.callback_name == "newrelic": + return callback if isinstance(callback, NewRelicLogger): return callback return None diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py new file mode 100644 index 00000000000..680cc226d18 --- /dev/null +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -0,0 +1,77 @@ +"""Save-time validation of team/key logging configs the runtime cannot honor. + +Team callbacks arrive as a single ``AddTeamCallback``, key callbacks arrive as a +``logging`` list inside the key metadata, so both shapes funnel into the same +per-integration checks here. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +_NEWRELIC_CALLBACK: Final = "newrelic" +_NEWRELIC_VAR_PREFIX: Final = "newrelic_" + + +def callback_config_error(callback_name: str | None, callback_vars: Mapping[str, str] | None) -> str | None: + if callback_name != _NEWRELIC_CALLBACK or not callback_vars: + return None + return _newrelic_config_error(callback_vars) + + +def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None: + """Validate every ``logging`` entry of a team/key metadata payload.""" + if not metadata: + return None + entries: Final = metadata.get("logging") + if not isinstance(entries, Sequence) or isinstance(entries, (str, bytes)): + return None + return next( + (error for error in (_logging_entry_error(entry) for entry in entries) if error is not None), + None, + ) + + +def _logging_entry_error(entry: object) -> str | None: + if not isinstance(entry, Mapping): + return None + callback_name: Final = entry.get("callback_name") + callback_vars: Final = entry.get("callback_vars") + if not isinstance(callback_name, str) or not isinstance(callback_vars, Mapping): + return None + return callback_config_error( + callback_name, + MappingProxyType({str(key): str(value) for key, value in callback_vars.items()}), + ) + + +def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None: + """Per-team New Relic routing runs on the OTel v2 path only. + + Accepting the config with the flag off would silently ship the team's traffic + through the operator's env-configured agent instead of the team's account. A + region outside the fixed table, or a region without a key, would likewise be + accepted and then silently ignored or misrouted at request time. + """ + if not any(key.startswith(_NEWRELIC_VAR_PREFIX) for key in callback_vars): + return None + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.integrations.otel.presets.newrelic import NEWRELIC_OTLP_ENDPOINT_BY_REGION + + if not is_otel_v2_enabled(): + return "Per-team New Relic routing requires the proxy to run with LITELLM_OTEL_V2=true." + + region: Final = callback_vars.get("newrelic_region") + if region is not None and region.lower() not in NEWRELIC_OTLP_ENDPOINT_BY_REGION: + return ( + f"Unknown newrelic_region {region!r}. " + f"Supported regions: {', '.join(sorted(NEWRELIC_OTLP_ENDPOINT_BY_REGION))}." + ) + + # ``callback_vars`` values are str()-coerced upstream, so a JSON ``null`` key + # arrives as the literal ``"None"``; treat that and the empty string as absent. + api_key: Final = callback_vars.get("newrelic_api_key") + if region is not None and (not api_key or api_key == "None"): + return "newrelic_region requires newrelic_api_key; the region rides the team's own key." + return None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4794da05a3e..cb5002e431b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -754,6 +754,12 @@ def convert_key_logging_metadata_to_callback( team_callback_settings_obj.callbacks.append(data.callback_name) for var, value in data.callback_vars.items(): + # New Relic routing reads these from the trusted-vars overlay with no + # callback-name check, so scope them to the newrelic entry: a team that + # put newrelic_* under a different callback never asked for New Relic and + # must not export to it. + if var.startswith("newrelic_") and data.callback_name != "newrelic": + continue if team_callback_settings_obj.callback_vars is None: team_callback_settings_obj.callback_vars = {} team_callback_settings_obj.callback_vars[var] = str(value) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 34a91dc59da..54f567b7aa2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_utils import ( enforce_output_token_estimates_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, @@ -553,6 +554,17 @@ def key_generation_check( return _personal_key_generation_check(user_api_key_dict=user_api_key_dict, data=data) +def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -> None: + """Key-level logging writes go through key metadata, not /team/callback. + + Without this the same New Relic config the team endpoint rejects would be + accepted here and then silently ignored or misrouted at request time. + """ + error: Final = logging_metadata_config_error(metadata) + if error is not None: + raise HTTPException(status_code=400, detail={"error": error}) # mutable-ok: FastAPI detail contract + + def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest | UpdateKeyRequest, @@ -891,6 +903,7 @@ async def _common_key_generation_helper( ) validate_budget_duration(data.budget_duration) + raise_on_invalid_key_logging_config(data.metadata) if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -1992,6 +2005,8 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ """ Check LiteLLM_ManagementEndpoint_MetadataFields (proxy/_types.py) for fields that are allowed to be updated """ + raise_on_invalid_key_logging_config(non_default_values.get("metadata")) + if "metadata" not in non_default_values: # allow user to set metadata to none non_default_values["metadata"] = existing_metadata.copy() diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 0e8c4e1825d..14a2a8a98a5 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -28,6 +28,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.callback_config_validation import callback_config_error from litellm.proxy.common_utils.callback_utils import ( _CALLBACK_VAR_ENCRYPTED_PREFIX, decrypt_callback_vars, @@ -51,6 +52,16 @@ router: Final = APIRouter() _CALLBACK_VARS_REDACTED: Final = "***REDACTED***" +def _callback_config_error(message: str) -> HTTPException: + return HTTPException(status_code=400, detail={"error": message}) # mutable-ok: FastAPI detail contract + + +def _validate_team_callback(data: "AddTeamCallback") -> None: + error: Final = callback_config_error(data.callback_name, data.callback_vars) + if error is not None: + raise _callback_config_error(error) + + def _redact_callback_secrets(metadata: Any) -> Any: """Strip secret values out of a team-metadata snapshot before audit logging. @@ -304,6 +315,8 @@ async def add_team_callbacks( user_api_key_dict=user_api_key_dict, ) + _validate_team_callback(data) + # store team callback settings in metadata team_metadata = _existing_team.metadata team_callback_settings: list[dict] = team_metadata.get("logging") # will be dict of type AddTeamCallback diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 94526de0757..67eae2b4f21 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3300,6 +3300,11 @@ class StandardCallbackDynamicParams(TypedDict, total=False): dd_agent_host: str | None dd_agent_port: str | None + # New Relic dynamic params (proxy-stamped team/key callback vars only; + # request-supplied values are blocked) + newrelic_api_key: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns into the dict + newrelic_region: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns into the dict + # Logging settings turn_off_message_logging: bool | None # when true will not log messages litellm_disabled_callbacks: list[str] | None diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py index a9a78dcb2f1..633be9f105f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py @@ -7,6 +7,7 @@ from opentelemetry.trace import NoOpTracer from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config from litellm.integrations.otel.presets import ( + dynamic_otlp_endpoint, dynamic_otlp_headers, project_routing_headers, ) @@ -23,31 +24,23 @@ def _cache(callback_name, exporters=None): def test_arize_dynamic_headers(): - headers = dynamic_otlp_headers( - "arize", {"arize_space_id": "S", "arize_api_key": "K"} - ) + headers = dynamic_otlp_headers("arize", {"arize_space_id": "S", "arize_api_key": "K"}) assert headers == {"arize-space-id": "S", "api_key": "K"} def test_arize_space_key_overrides_space_id(): - headers = dynamic_otlp_headers( - "arize", {"arize_space_id": "S", "arize_space_key": "SK"} - ) + headers = dynamic_otlp_headers("arize", {"arize_space_id": "S", "arize_space_key": "SK"}) assert headers == {"arize-space-id": "SK"} def test_langfuse_dynamic_headers_need_both_keys(): assert dynamic_otlp_headers("langfuse_otel", {"langfuse_public_key": "pk"}) is None - headers = dynamic_otlp_headers( - "langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"} - ) + headers = dynamic_otlp_headers("langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}) assert headers is not None and "Authorization" in headers def test_langfuse_dynamic_headers_carry_v4_ingestion_version(): - headers = dynamic_otlp_headers( - "langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"} - ) + headers = dynamic_otlp_headers("langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}) expected_auth = "Basic " + base64.b64encode(b"pk:sk").decode() assert headers == { "Authorization": expected_auth, @@ -56,9 +49,7 @@ def test_langfuse_dynamic_headers_carry_v4_ingestion_version(): def test_weave_dynamic_headers(): - headers = dynamic_otlp_headers( - "weave_otel", {"wandb_api_key": "w", "weave_project_id": "p"} - ) + headers = dynamic_otlp_headers("weave_otel", {"wandb_api_key": "w", "weave_project_id": "p"}) assert headers is not None assert "Authorization" in headers and headers["project_id"] == "p" @@ -100,9 +91,7 @@ def test_provider_cache_is_bounded_and_evicts_lru(monkeypatch): monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 2) shut_down = [] - monkeypatch.setattr( - routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) - ) + monkeypatch.setattr(routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)) cache = _cache("arize") default = NoOpTracer() @@ -179,9 +168,7 @@ def test_dynamic_headers_do_not_leak_to_other_owners_exporter(): ), ], ) - new_cfg = cache._routed_config( - {"arize-space-id": "TEAMX", "api_key": "TEAMX_KEY"}, {} - ) + new_cfg = cache._routed_config({"arize-space-id": "TEAMX", "api_key": "TEAMX_KEY"}, {}) by_owner = {e.owner: e.headers for e in new_cfg.exporters} assert by_owner["arize"] == "arize-space-id=TEAMX,api_key=TEAMX_KEY" assert by_owner[None] == "x=base-collector" @@ -206,16 +193,14 @@ def _phoenix_cache(kind="otlp_http"): def test_phoenix_project_headers_precedence_and_blanks(): - assert project_routing_headers( - "arize_phoenix", {"phoenix_project_name": "team-proj"} - ) == {"x-project-name": "team-proj"} + assert project_routing_headers("arize_phoenix", {"phoenix_project_name": "team-proj"}) == { + "x-project-name": "team-proj" + } assert project_routing_headers( "arize_phoenix", {"phoenix_project_name_override": "override", "phoenix_project_name": "base"}, ) == {"x-project-name": "override"} - assert ( - project_routing_headers("arize_phoenix", {"phoenix_project_name": " "}) == {} - ) + assert project_routing_headers("arize_phoenix", {"phoenix_project_name": " "}) == {} assert project_routing_headers("arize_phoenix", None) == {} # Only Phoenix participates in project routing. assert project_routing_headers("arize", {"phoenix_project_name": "p"}) == {} @@ -286,10 +271,7 @@ def test_client_dynamic_params_cannot_choose_phoenix_project(): cache = _phoenix_cache() default = NoOpTracer() assert cache.route_for(default, {"phoenix_project_name": "attacker"}).tracer is default - assert ( - cache.route_for(default, {"phoenix_project_name_override": "attacker"}).tracer - is default - ) + assert cache.route_for(default, {"phoenix_project_name_override": "attacker"}).tracer is default assert cache._providers == {} @@ -321,9 +303,7 @@ def test_eviction_defers_shutdown_while_a_span_is_open(monkeypatch): monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) shut_down = [] - monkeypatch.setattr( - routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) - ) + monkeypatch.setattr(routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)) cache = _cache("arize") default = NoOpTracer() @@ -347,18 +327,13 @@ def test_retired_providers_are_capped(monkeypatch): monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) monkeypatch.setattr(routing_mod, "_MAX_RETIRED_PROVIDERS", 2) shut_down = [] - monkeypatch.setattr( - routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) - ) + monkeypatch.setattr(routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)) cache = _cache("arize") default = NoOpTracer() # Every route stays held (no release), so each one evicts and retires its # predecessor instead of shutting it down. - routes = [ - cache.route_for(default, {"arize_space_id": str(i), "arize_api_key": "K"}) - for i in range(5) - ] + routes = [cache.route_for(default, {"arize_space_id": str(i), "arize_api_key": "K"}) for i in range(5)] assert len(cache._providers) == 1 assert len(cache._retired) == 2 # capped, not one retiree per open call @@ -374,11 +349,112 @@ def test_release_without_eviction_keeps_provider_alive(monkeypatch): from litellm.integrations.otel.plumbing import routing as routing_mod shut_down = [] - monkeypatch.setattr( - routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) - ) + monkeypatch.setattr(routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)) cache = _cache("arize") route = cache.route_for(NoOpTracer(), {"arize_space_id": "A", "arize_api_key": "K"}) cache.release(route.provider) assert shut_down == [] # still cached, never retired cache.release(None) # default-route release is a no-op + + +# --- New Relic: per-team api-key header + fixed-table region endpoint --- # + + +def test_newrelic_dynamic_headers(): + assert dynamic_otlp_headers("newrelic", {"newrelic_api_key": "NRAL-KEY"}) == {"api-key": "NRAL-KEY"} + assert dynamic_otlp_headers("newrelic", {"newrelic_region": "eu"}) is None + + +def test_newrelic_dynamic_endpoint_resolves_from_fixed_table(): + from litellm.integrations.otel.presets import dynamic_otlp_endpoint + + assert dynamic_otlp_endpoint("newrelic", {"newrelic_region": "eu"}) == "https://otlp.eu01.nr-data.net" + assert dynamic_otlp_endpoint("newrelic", {"newrelic_region": "US"}) == "https://otlp.nr-data.net" + # A key-only team (no region) resolves to the fixed US default deterministically, + # never the operator's NEW_RELIC_REGION-configured preset endpoint. + assert dynamic_otlp_endpoint("newrelic", {"newrelic_api_key": "k"}) == "https://otlp.nr-data.net" + # An unknown region also resolves to the documented US default, not a guess. + assert dynamic_otlp_endpoint("newrelic", {"newrelic_region": "mars"}) == "https://otlp.nr-data.net" + # Callbacks without an endpoint resolver keep their preset endpoint. + assert dynamic_otlp_endpoint("arize", {"newrelic_region": "eu"}) is None + + +def test_newrelic_endpoint_stamped_onto_owned_exporter_only(): + cache = _cache( + "newrelic", + exporters=[ + ExporterSpec( + kind="otlp_http", + endpoint="http://self-hosted-collector:4318", + headers="x=base-collector", + owner=None, + ), + ExporterSpec( + kind="otlp_http", + endpoint="https://otlp.nr-data.net", + owner="newrelic", + requires_headers=True, + ), + ], + ) + new_cfg = cache._routed_config({"api-key": "TEAM-EU-KEY"}, {}, "https://otlp.eu01.nr-data.net") + by_owner = {e.owner: e for e in new_cfg.exporters} + assert by_owner["newrelic"].endpoint == "https://otlp.eu01.nr-data.net" + assert by_owner["newrelic"].headers == "api-key=TEAM-EU-KEY" + assert by_owner[None].endpoint == "http://self-hosted-collector:4318" + assert by_owner[None].headers == "x=base-collector" + + +def test_newrelic_provider_cached_per_key_and_region(): + cache = _cache( + "newrelic", + exporters=[ExporterSpec(kind="in_memory"), ExporterSpec(kind="otlp_http", owner="newrelic")], + ) + default = NoOpTracer() + cache.route_for(default, {"newrelic_api_key": "K1", "newrelic_region": "us"}) + cache.route_for(default, {"newrelic_api_key": "K1", "newrelic_region": "us"}) + assert len(cache._providers) == 1 + # Same key, different region → distinct provider (distinct endpoint). + cache.route_for(default, {"newrelic_api_key": "K1", "newrelic_region": "eu"}) + assert len(cache._providers) == 2 + cache.route_for(default, {"newrelic_api_key": "K2", "newrelic_region": "eu"}) + assert len(cache._providers) == 3 + + +def test_requires_headers_spec_skipped_without_headers(): + from litellm.integrations.otel.plumbing.providers import build_tracer_provider + + cfg = OpenTelemetryV2Config( + exporters=[ExporterSpec(kind="otlp_http", endpoint="https://otlp.nr-data.net", requires_headers=True)] + ) + provider = build_tracer_provider(cfg) + processors = provider._active_span_processor._span_processors + # Only the baggage processor: the keyless spec must not export (New Relic + # rejects unauthenticated posts with a 4xx per span batch). + assert [type(p).__name__ for p in processors] == ["LiteLLMBaggageSpanProcessor"] + + keyed = OpenTelemetryV2Config( + exporters=[ + ExporterSpec( + kind="otlp_http", endpoint="https://otlp.nr-data.net", headers="api-key=k", requires_headers=True + ) + ] + ) + keyed_provider = build_tracer_provider(keyed) + assert len(keyed_provider._active_span_processor._span_processors) == 2 + + +def test_newrelic_key_only_team_routes_to_us_not_operator_region(monkeypatch): + """A team that saves a key but no region must export to the fixed US default, + independent of the operator's NEW_RELIC_REGION, so its spans are never + silently dropped by an operator-configured region its key does not match.""" + monkeypatch.setenv("NEW_RELIC_REGION", "eu") + cache = _cache( + "newrelic", + exporters=[ExporterSpec(kind="otlp_http", endpoint="https://otlp.eu01.nr-data.net", owner="newrelic")], + ) + new_cfg = cache._routed_config( + {"api-key": "US-KEY"}, {}, dynamic_otlp_endpoint("newrelic", {"newrelic_api_key": "US-KEY"}) + ) + owned = next(e for e in new_cfg.exporters if e.owner == "newrelic") + assert owned.endpoint == "https://otlp.nr-data.net" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index e5d5b62b856..704a1d3a7bb 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -341,6 +341,35 @@ def test_idempotent_on_repeat_callback(): assert len(exporter.get_finished_spans()) == 1 +def test_evicted_carrier_completed_call_emits_one_deferred_span(): + """Eviction over the concurrency budget drops only the boundary carrier, not + the call. When an evicted call later closes as a real completed call + (``upstream_started``, payload present) it still emits exactly one span + through the deferred branch, and a second close for the same id dedups. Only + an evicted call that never closes goes unexported.""" + logger, exporter = _logger() + kwargs = {**_kwargs(), "api_call_start_time": datetime(2026, 5, 26, 12, 0, 0, tzinfo=timezone.utc)} + logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs) + assert "call_1" in logger._open_llm_calls + + # Evict exactly as ``_store_open_call`` does over budget: drop the oldest + # carrier and release its routed provider. + _, evicted = logger._open_llm_calls.popitem(last=False) + logger._release_carrier(evicted) + assert not logger._open_llm_calls + assert exporter.get_finished_spans() == () # the evicted boundary span is never exported + + asyncio.run(logger.async_log_success_event(kwargs, None, None, None)) + spans = exporter.get_finished_spans() + assert len(spans) == 1, "the evicted call's real close re-emits one deferred span, not zero" + assert spans[0].name == "chat gpt-4o" + assert spans[0].attributes[LiteLLM.CALL_ID] == "call_1" + + # Success-then-failure on one logging object: the deferred branch dedups by id. + asyncio.run(logger.async_log_failure_event(kwargs, None, None, None)) + assert len(exporter.get_finished_spans()) == 1, "second close for the same id must not duplicate" + + # --------------------------------------------------------------------------- # # MCP tool-call spans # --------------------------------------------------------------------------- # @@ -2522,3 +2551,149 @@ def test_deferred_pre_call_does_not_churn_tenant_cache(monkeypatch): server.end() headers_b = next(h for h in captured if "proj-b" in h) assert [s.name for s in captured[headers_b].get_finished_spans()] == ["chat gpt-4o"] + + +# --- New Relic team-scoped deferred emit + dedup (no pre_call carrier) --- # + + +def test_no_span_when_request_never_reached_upstream(): + """A request rejected before the upstream call — at the auth/budget gate, or + blocked by a pre-call guardrail — carries the ``no upstream call`` marker + (stamped in ``proxy/utils.py`` before its handlers fire), so the failure log + produces no phantom CLIENT span even though a payload exists.""" + from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL + + logger, exporter = _logger() + payload = _payload( + status="failure", + error_information={"error_class": "ProxyException", "error_code": "401"}, + ) + kwargs = _kwargs(payload=payload) + kwargs[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True + # No log_pre_api_call: the call never started. + asyncio.run(logger.async_log_failure_event(kwargs, None, None, None)) + assert exporter.get_finished_spans() == () # no phantom LLM span + + +def test_success_without_pre_call_emits_deferred_span(): + """A team/key-scoped logger is registered as a success callback only, so + ``pre_call`` never reaches it and no carrier exists. A completed call (it has + its payload, no ``no upstream call`` marker) must still get its span — the + deferred branch — or team-scoped destinations receive nothing at all.""" + logger, exporter = _logger() + # No log_pre_api_call: this logger never receives the input hook. The + # request-level provider-handoff stamp is present (pre_call ran globally). + asyncio.run( + logger.async_log_success_event({**_kwargs(), "api_call_start_time": 100.0}, None, 100.0, 101.5) + ) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes.get("gen_ai.operation.name") + # Start time comes from the callback's start_time, not a bogus zero. + assert spans[0].start_time == 100_000_000_000 + assert spans[0].end_time == 101_500_000_000 + + +def test_no_carrier_and_no_payload_is_noop(): + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event({"litellm_params": {}}, None, None, None) + ) + assert exporter.get_finished_spans() == () + + +def test_second_close_for_same_call_does_not_duplicate_span(): + """Success and failure can both fire on one logging object for the same call + id. The first close pops the carrier and finishes the boundary span; the + second must dedup against it, not fabricate a duplicate through the + deferred branch.""" + logger, exporter = _logger() + kwargs = {**_kwargs(), "api_call_start_time": 100.0} + # Boundary open: the span is born at pre_call under a live server span and + # closed via finish_span, which never passes through emit()'s dedup. + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + _emit_llm(logger, kwargs, ambient=server) + server.end() + llm_before = [s for s in exporter.get_finished_spans() if s.name.startswith("chat")] + assert len(llm_before) == 1 + # Second close: carrier already popped, payload still present, handoff stamped. + asyncio.run(logger.async_log_failure_event(kwargs, None, None, None)) + llm_after = [s for s in exporter.get_finished_spans() if s.name.startswith("chat")] + assert len(llm_after) == 1 + + +def test_failure_without_pre_call_emits_deferred_error_span(): + """A team-scoped logger registered as a failure callback only still gets an + ERROR span for a real provider failure (payload present, no marker).""" + from opentelemetry.trace import StatusCode + + logger, exporter = _logger() + payload = _payload( + status="failure", + error_information={"error_class": "RateLimitError", "error_code": "429"}, + ) + asyncio.run( + logger.async_log_failure_event( + {**_kwargs(payload=payload), "api_call_start_time": 100.0}, None, None, None + ) + ) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + + +def test_boundary_open_with_no_payload_ends_provisional_span(): + """Opened at pre_call but the payload never materialized: the boundary span + is ended provisionally (no payload attributes) rather than leaked open.""" + logger, exporter = _logger() + kwargs = _kwargs() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + with trace.use_span(server, end_on_exit=False): + logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs) + asyncio.run( + logger.async_log_success_event( + {**kwargs, "standard_logging_object": None, "litellm_call_id": "call_1"}, None, None, None + ) + ) + server.end() + llm_spans = [s for s in exporter.get_finished_spans() if s.name.startswith("chat")] + assert len(llm_spans) == 1 + assert "gen_ai.usage.input_tokens" not in llm_spans[0].attributes + + +def test_failure_before_provider_handoff_emits_nothing(): + """A failure event whose request never handed off to a provider (router + pre-call rejection, SDK error before the call, standalone guardrail run) + has a payload but no ``api_call_start_time``; without a carrier it must not + fabricate an LLM-call span.""" + logger, exporter = _logger() + payload = _payload( + status="failure", + error_information={"error_class": "RateLimitError", "error_code": "429"}, + ) + asyncio.run(logger.async_log_failure_event(_kwargs(payload=payload), None, None, None)) + assert exporter.get_finished_spans() == () + + +def test_provisional_close_then_payload_close_does_not_duplicate(): + """Streaming shape: the success close arrives with no assembled payload (the + boundary span is ended provisionally), then the failure close arrives with a + payload for the same call id. Exactly one exported span.""" + logger, exporter = _logger() + kwargs = {**_kwargs(), "api_call_start_time": 100.0} + payload = kwargs["standard_logging_object"] + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + with trace.use_span(server, end_on_exit=False): + logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs) + asyncio.run( + logger.async_log_success_event( + {**kwargs, "standard_logging_object": None, "litellm_call_id": payload["litellm_call_id"]}, + None, + None, + None, + ) + ) + asyncio.run(logger.async_log_failure_event(kwargs, None, None, None)) + server.end() + llm_spans = [s for s in exporter.get_finished_spans() if s.name.startswith("chat")] + assert len(llm_spans) == 1 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_presets.py b/tests/test_litellm/integrations/otel/test_otel_v2_presets.py index add2caf9a48..58cfc1ceb3f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_presets.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_presets.py @@ -160,3 +160,55 @@ def test_agentops_endpoint_points_at_live_host(): # so a typo or stale domain can never ship again. assert _AGENTOPS_ENDPOINT == "https://otlp.agentops.ai/v1/traces" assert "agentops.cloud" not in _AGENTOPS_ENDPOINT + + +def test_newrelic_preset_reads_env_license_key(monkeypatch): + monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "env-license-key") + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.presets.newrelic import newrelic_preset + + cfg = newrelic_preset() + spec = next(e for e in cfg.exporters if e.owner == ExporterOwner.NEWRELIC) + assert spec.kind == "otlp_http" + assert spec.endpoint == "https://otlp.nr-data.net" + assert spec.headers == "api-key=env-license-key" + assert spec.requires_headers is True + assert "genai" in cfg.mapper_names + + +def test_newrelic_preset_without_key_still_contributes_owned_spec(monkeypatch): + # The owned spec is the stamping target for per-team credentials, so it must + # exist even with no operator env key; requires_headers keeps the keyless + # copy from ever exporting. + monkeypatch.delenv("NEW_RELIC_LICENSE_KEY", raising=False) + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.presets.newrelic import newrelic_preset + + cfg = newrelic_preset() + spec = next(e for e in cfg.exporters if e.owner == ExporterOwner.NEWRELIC) + assert spec.headers is None + assert spec.requires_headers is True + + +def test_newrelic_preset_operator_region_and_content_knob(monkeypatch): + monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "env-license-key") + monkeypatch.setenv("NEW_RELIC_REGION", "EU") + monkeypatch.setenv("NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", "true") + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.presets.newrelic import newrelic_preset + + cfg = newrelic_preset() + spec = next(e for e in cfg.exporters if e.owner == ExporterOwner.NEWRELIC) + assert spec.endpoint == "https://otlp.eu01.nr-data.net" + assert cfg.capture_span_content is True + + monkeypatch.setenv("NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", "false") + assert newrelic_preset().capture_span_content is False + + +def test_newrelic_preset_unset_content_knob_keeps_default(monkeypatch): + monkeypatch.delenv("NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", raising=False) + monkeypatch.delenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", raising=False) + from litellm.integrations.otel.presets.newrelic import newrelic_preset + + assert newrelic_preset().capture_span_content is False diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index f9ddc47cc7c..9b2bd5e2585 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -187,3 +187,49 @@ def test_empty_kwargs_returns_empty_params(): params = initialize_standard_callback_dynamic_params({}) assert dict(params) == {} + + +def test_newrelic_callback_params_are_not_extracted_from_request_kwargs(): + kwargs = { + "newrelic_api_key": "caller-key", + "metadata": {"newrelic_api_key": "caller-key-2", "newrelic_region": "eu"}, + "litellm_params": {"metadata": {"newrelic_region": "eu"}}, + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("newrelic_api_key") is None + assert params.get("newrelic_region") is None + + +def test_newrelic_trusted_vars_overlay_reaches_standard_params(): + from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD + + kwargs = { + # A caller-supplied copy must lose to the proxy-stamped trusted value. + "newrelic_api_key": "caller-key", + TRUSTED_CALLBACK_VARS_FIELD: { + "newrelic_api_key": "team-key", + "newrelic_region": "eu", + # Non-overlay trusted vars must not be copied by the overlay. + "langfuse_public_key": "pk-team", + }, + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("newrelic_api_key") == "team-key" + assert params.get("newrelic_region") == "eu" + assert params.get("langfuse_public_key") is None + + +def test_trusted_vars_overlay_uses_shared_parser_semantics(): + # The overlay rides get_trusted_callback_params, the same parser the + # datadog handler consumes, so values are str()-coerced identically. + from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD + + params = initialize_standard_callback_dynamic_params( + {TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": 12345}} + ) + + assert params.get("newrelic_api_key") == "12345" 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 873da28fc34..1c706be51fa 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5200,3 +5200,81 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa ) for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook) +def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch): + """With LITELLM_OTEL_V2 on, the "newrelic" callback builds the OTel v2 + logger (per-team credential routing); with the flag off (default) it keeps + the legacy agent-based logger, so existing deployments are untouched.""" + from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.litellm_core_utils import litellm_logging as logging_module + + logging_module._in_memory_loggers.clear() + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + try: + v2_logger = logging_module._init_custom_logger_compatible_class( + logging_integration="newrelic", + internal_usage_cache=None, + llm_router=None, + custom_logger_init_args={}, + ) + assert isinstance(v2_logger, OpenTelemetryV2) + assert v2_logger.callback_name == "newrelic" + # Same name resolves to the same instance, not a second logger. + again = logging_module._init_custom_logger_compatible_class( + logging_integration="newrelic", + internal_usage_cache=None, + llm_router=None, + custom_logger_init_args={}, + ) + assert again is v2_logger + finally: + logging_module._in_memory_loggers.clear() + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + + +def test_newrelic_dispatch_keeps_legacy_agent_when_flag_off(monkeypatch): + from litellm.integrations.newrelic import NewRelicLogger + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.litellm_core_utils import litellm_logging as logging_module + + logging_module._in_memory_loggers.clear() + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + legacy = logging_module._init_custom_logger_compatible_class( + logging_integration="newrelic", + internal_usage_cache=None, + llm_router=None, + custom_logger_init_args={}, + ) + assert isinstance(legacy, NewRelicLogger) + finally: + logging_module._in_memory_loggers.clear() + is_otel_v2_enabled.cache_clear() + + +def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch): + """Under LITELLM_OTEL_V2 the "newrelic" instance is an OpenTelemetryV2; the + cached-lookup must find it or hook resolution (post-call failure/success + hooks) silently skips the callback.""" + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.litellm_core_utils import litellm_logging as logging_module + + logging_module._in_memory_loggers.clear() + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + try: + created = logging_module._init_custom_logger_compatible_class( + logging_integration="newrelic", + internal_usage_cache=None, + llm_router=None, + custom_logger_init_args={}, + ) + found = logging_module.get_custom_logger_compatible_class("newrelic") + assert found is created + finally: + logging_module._in_memory_loggers.clear() + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() diff --git a/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py index b2a242bf8f2..272a8ffa972 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py @@ -264,3 +264,148 @@ class TestCallbackManagementEndpoints: assert galileo_config["displayName"] == "Galileo" assert "GALILEO_API_KEY" in galileo_config["dynamic_params"] assert "GALILEO_PROJECT_ID" in galileo_config["dynamic_params"] + + +class TestNewRelicCallbackConfig: + def test_newrelic_entry_supports_team_logging_with_dynamic_params(self): + client = TestClient(app) + response = client.get("/callbacks/configs", headers={"Authorization": "Bearer sk-1234"}) + assert response.status_code == 200 + newrelic = next( + (config for config in response.json() if config.get("id") == "newrelic"), + None, + ) + assert newrelic is not None + assert newrelic["supports_key_team_logging"] is True + params = newrelic["dynamic_params"] + assert params["newrelic_api_key"]["type"] == "password" + assert "newrelic_region" in params + # The operator-only agent env flag must not appear as a team-configurable + # field: it is not a StandardCallbackDynamicParams key and would be rejected. + assert "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED" not in params + + +class TestNewRelicTeamCallbackValidation: + def _data(self, callback_vars): + from litellm.proxy._types import AddTeamCallback + + return AddTeamCallback(callback_name="newrelic", callback_type="success", callback_vars=callback_vars) + + def test_rejects_when_otel_v2_off(self, monkeypatch): + from fastapi import HTTPException + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy.management_endpoints.team_callback_endpoints import _validate_team_callback + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + with pytest.raises(HTTPException) as exc: + _validate_team_callback(self._data({"newrelic_api_key": "k"})) + assert "LITELLM_OTEL_V2" in str(exc.value.detail) + finally: + is_otel_v2_enabled.cache_clear() + + def test_rejects_unknown_region_and_region_without_key(self, monkeypatch): + from fastapi import HTTPException + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy.management_endpoints.team_callback_endpoints import _validate_team_callback + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + try: + with pytest.raises(HTTPException) as exc: + _validate_team_callback(self._data({"newrelic_api_key": "k", "newrelic_region": "mars"})) + assert "Unknown newrelic_region" in str(exc.value.detail) + with pytest.raises(HTTPException) as exc: + _validate_team_callback(self._data({"newrelic_region": "eu"})) + assert "requires newrelic_api_key" in str(exc.value.detail) + _validate_team_callback(self._data({"newrelic_api_key": "k", "newrelic_region": "EU"})) + # A JSON-null key is str()-coerced to "None" upstream; it must not + # slip past the region-requires-key guard. + with pytest.raises(HTTPException) as exc: + _validate_team_callback(self._data({"newrelic_api_key": None, "newrelic_region": "eu"})) + assert "requires newrelic_api_key" in str(exc.value.detail) + finally: + is_otel_v2_enabled.cache_clear() + + def test_ignores_other_callbacks_and_bare_newrelic(self, monkeypatch): + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy._types import AddTeamCallback + from litellm.proxy.management_endpoints.team_callback_endpoints import _validate_team_callback + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + assert _validate_team_callback(self._data({})) is None + assert ( + _validate_team_callback( + AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, + ) + ) + is None + ) + finally: + is_otel_v2_enabled.cache_clear() + + +class TestNewRelicKeyLoggingValidation: + """Key-level logging is written through key metadata, not /team/callback.""" + + def _metadata(self, callback_vars): + return {"logging": [{"callback_name": "newrelic", "callback_type": "success", "callback_vars": callback_vars}]} + + def test_rejects_same_configs_as_team_endpoint(self, monkeypatch): + from fastapi import HTTPException + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_invalid_key_logging_config, + ) + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + with pytest.raises(HTTPException) as exc: + raise_on_invalid_key_logging_config(self._metadata({"newrelic_api_key": "k"})) + assert "LITELLM_OTEL_V2" in str(exc.value.detail) + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + with pytest.raises(HTTPException) as exc: + raise_on_invalid_key_logging_config( + self._metadata({"newrelic_api_key": "k", "newrelic_region": "mars"}) + ) + assert "Unknown newrelic_region" in str(exc.value.detail) + with pytest.raises(HTTPException) as exc: + raise_on_invalid_key_logging_config(self._metadata({"newrelic_region": "eu"})) + assert "requires newrelic_api_key" in str(exc.value.detail) + raise_on_invalid_key_logging_config(self._metadata({"newrelic_api_key": "k", "newrelic_region": "EU"})) + finally: + is_otel_v2_enabled.cache_clear() + + def test_ignores_metadata_without_newrelic_logging(self, monkeypatch): + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_invalid_key_logging_config, + ) + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + assert raise_on_invalid_key_logging_config(None) is None + assert raise_on_invalid_key_logging_config({"logging": "not-a-list"}) is None + assert raise_on_invalid_key_logging_config({"tags": ["a"]}) is None + assert raise_on_invalid_key_logging_config(self._metadata({})) is None + assert ( + raise_on_invalid_key_logging_config( + {"logging": [{"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}]} + ) + is None + ) + finally: + is_otel_v2_enabled.cache_clear() diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 81a97a70efa..50ef6f29ec2 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7329,3 +7329,92 @@ def test_vertex_sends_exactly_one_authorization_header(): vertex_request_headers.update(forwarded) assert _authorization_values(vertex_request_headers) == [GOOGLE_ACCESS_TOKEN] +@pytest.mark.asyncio +async def test_newrelic_team_callback_vars_reach_trusted_field(): + """A key with a newrelic team callback stamps its vars into the proxy-owned + trusted field, and a caller-supplied newrelic_api_key in the body is + stripped rather than merged.""" + key_with_newrelic_callback = UserAPIKeyAuth( + api_key="hashed-key", + metadata={ + "logging": [ + { + "callback_name": "newrelic", + "callback_type": "success", + "callback_vars": {"newrelic_api_key": "team-nr-key", "newrelic_region": "eu"}, + } + ] + }, + ) + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + "newrelic_api_key": "attacker-key", + } + + updated = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=key_with_newrelic_callback, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated[TRUSTED_CALLBACK_VARS_FIELD] == { + "newrelic_api_key": "team-nr-key", + "newrelic_region": "eu", + } + assert updated["success_callback"] == ["newrelic"] + + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, + ) + + params = initialize_standard_callback_dynamic_params(updated) + assert params.get("newrelic_api_key") == "team-nr-key" + assert params.get("newrelic_region") == "eu" + + from litellm.integrations.otel.presets import dynamic_otlp_endpoint, dynamic_otlp_headers + + assert dynamic_otlp_headers("newrelic", params) == {"api-key": "team-nr-key"} + assert dynamic_otlp_endpoint("newrelic", params) == "https://otlp.eu01.nr-data.net" + + from litellm.utils import get_non_default_completion_params + + forwarded = get_non_default_completion_params(updated) + assert not any(param.startswith("newrelic_") for param in forwarded) + assert TRUSTED_CALLBACK_VARS_FIELD not in forwarded + + +def test_newrelic_vars_scoped_to_newrelic_callback_entry(): + """New Relic routing reads these vars from the trusted overlay with no + callback-name check, so a team that puts newrelic_* under a different + callback's vars must not have them enter the shared bag (and so never + exports to New Relic). Vars under a real newrelic entry are kept.""" + from litellm.proxy._types import AddTeamCallback + from litellm.proxy.litellm_pre_call_utils import convert_key_logging_metadata_to_callback + + smuggled = convert_key_logging_metadata_to_callback( + AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "newrelic_api_key": "SMUGGLED", + "newrelic_region": "eu", + }, + ), + None, + ) + assert smuggled.callback_vars == {"langfuse_public_key": "pk"} + + legit = convert_key_logging_metadata_to_callback( + AddTeamCallback( + callback_name="newrelic", + callback_type="success", + callback_vars={"newrelic_api_key": "REAL", "newrelic_region": "us"}, + ), + None, + ) + assert legit.callback_vars == {"newrelic_api_key": "REAL", "newrelic_region": "us"} diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index 7aa121dcca5..3906aa744f7 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -6,6 +6,7 @@ import galileoLogo from "../../public/assets/logos/galileo.ico"; import lagoLogo from "../../public/assets/logos/lago.svg"; import langfuseLogo from "../../public/assets/logos/langfuse.png"; import langsmithLogo from "../../public/assets/logos/langsmith.png"; +import newrelicLogo from "../../public/assets/logos/newrelic.png"; import openmeterLogo from "../../public/assets/logos/openmeter.png"; import otelLogo from "../../public/assets/logos/otel.png"; @@ -77,6 +78,17 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [ }, description: "Datadog Logging Integration", }, + { + id: "newrelic", + displayName: "New Relic", + logo: newrelicLogo.src, + supports_key_team_logging: true, + dynamic_params: { + newrelic_api_key: "password", + newrelic_region: "text", + }, + description: "New Relic Logging Integration", + }, { id: "lago", displayName: "Lago", From 0697188be40c0f5528b2ca18ec110fd95b55d9f8 Mon Sep 17 00:00:00 2001 From: milan Date: Sun, 23 Aug 2026 03:14:19 +0000 Subject: [PATCH 088/273] 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 089/273] 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 01595d2fbf232f49c45177c4c0f2108041966f28 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:44:32 -0700 Subject: [PATCH 090/273] perf(ci): cache uv dependencies in the lint job (#37783) The lint job installs its dependencies from scratch on every run. That step measures 2.8 minutes of a job whose p50 is 9.5, and lint is the slowest required check on 9 of the last 10 merged staging PRs, so it sets the critical path for the whole PR. _test-unit-base.yml already caches ~/.cache/uv and .venv keyed on uv.lock. This mirrors that block. The key carries its own `lint` namespace rather than sharing the unit tier's: the two jobs sync different group sets (proxy-dev + e2e-dev here, ci + proxy-dev + four extras there), so a shared .venv entry would be pruned and rebuilt on alternating runs. --- .github/workflows/test-linting.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index ccb58f5cc9c..9f1283da19e 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -67,6 +67,17 @@ jobs: with: version: "0.10.9" + - name: Cache uv dependencies + if: steps.changes.outputs.decision != 'skip' + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cache/uv + .venv + key: ${{ runner.os }}-uv-lint-${{ hashFiles('uv.lock') }} + restore-keys: | + ${{ runner.os }}-uv-lint- + - name: Clean Python cache if: steps.changes.outputs.decision != 'skip' run: | From 6e23288b47ecd89b658e6c5188f090dcf886a227 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:44:58 -0700 Subject: [PATCH 091/273] perf(ci): fan the budget checkers out across cores (#37784) * perf(ci): fan the budget checkers out across cores check_type_discipline.py and check_test_quality.py each walk a few thousand files and parse every one, single-threaded. In the lint job those two steps measure 2.3 and 1.6 minutes, second and third behind dependency install, and lint is the slowest required check on 9 of the last 10 merged staging PRs. check_file is already pure per-file work, so the walk fans out over a process pool with no change to what either rule reports. Callers sort, which is what keeps output order stable when results land out of order. Runs below PARALLEL_MIN_PATHS stay serial rather than pay for process startup, and the worker count is capped so a large runner does not oversubscribe. Measured locally over the same trees, output byte-identical both times: type-discipline 17.8s -> 3.0s over litellm/ (78,768 report lines), test-quality 14.4s -> 2.3s over tests/ (6,321 report lines), per-rule counts unchanged. * test(ci): type the fan-out helpers and skip the comparison on one core --- scripts/check_test_quality.py | 27 +++++++- scripts/check_type_discipline.py | 27 +++++++- tests/test_litellm/test_check_test_quality.py | 62 +++++++++++++++++++ .../test_check_type_discipline.py | 62 +++++++++++++++++++ 4 files changed, 176 insertions(+), 2 deletions(-) diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index e3ffbac9808..6964aed56e4 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -102,11 +102,13 @@ from __future__ import annotations import ast import io +import os import re import sys import tokenize from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass +from multiprocessing import Pool from pathlib import Path from types import MappingProxyType from typing import Final, NamedTuple @@ -692,13 +694,36 @@ def collect_paths(raw: Iterable[str]) -> Iterator[Path]: yield candidate +PARALLEL_MIN_PATHS = 200 +MAX_WORKERS = 8 + + +def _worker_count(path_count: int) -> int: + """1 when the run is too small to repay process startup, else one worker per + core up to MAX_WORKERS.""" + if path_count < PARALLEL_MIN_PATHS: + return 1 + return max(1, min(os.cpu_count() or 1, MAX_WORKERS)) + + +def scan_paths(paths: Sequence[Path]) -> tuple[Violation, ...]: + """check_file over every path. Pure per-file work, so it fans out across + processes; callers sort, which is what keeps output order stable.""" + workers = _worker_count(len(paths)) + if workers == 1: + return tuple(v for path in paths for v in check_file(path)) + with Pool(workers) as pool: + return tuple(v for found in pool.imap_unordered(check_file, paths, chunksize=32) for v in found) + + def main(argv: Sequence[str]) -> int: paths: Final = tuple(a for a in argv if not a.startswith("-")) if not paths: print("usage: check_test_quality.py ...", file=sys.stderr) return 2 - violations: Final = sorted(v for path in collect_paths(paths) for v in check_file(path)) + targets: Final = tuple(collect_paths(paths)) + violations: Final = sorted(scan_paths(targets)) for violation in violations: print(violation.render()) diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index 0706c8a7bd8..a2ab4760c4f 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -114,10 +114,12 @@ from __future__ import annotations import ast import io +import os import re import sys import tokenize from dataclasses import dataclass +from multiprocessing import Pool from pathlib import Path from collections.abc import Iterable, Iterator, Mapping, Sequence from typing import NamedTuple @@ -1070,13 +1072,36 @@ def collect_paths(raw: Iterable[str]) -> Iterator[Path]: yield p +PARALLEL_MIN_PATHS = 200 +MAX_WORKERS = 8 + + +def _worker_count(path_count: int) -> int: + """1 when the run is too small to repay process startup, else one worker per + core up to MAX_WORKERS.""" + if path_count < PARALLEL_MIN_PATHS: + return 1 + return max(1, min(os.cpu_count() or 1, MAX_WORKERS)) + + +def scan_paths(paths: Sequence[Path]) -> tuple[Violation, ...]: + """check_file over every path. Pure per-file work, so it fans out across + processes; callers sort, which is what keeps output order stable.""" + workers = _worker_count(len(paths)) + if workers == 1: + return tuple(v for path in paths for v in check_file(path)) + with Pool(workers) as pool: + return tuple(v for found in pool.imap_unordered(check_file, paths, chunksize=32) for v in found) + + def main(argv: Sequence[str]) -> int: paths = tuple(a for a in argv if not a.startswith("-")) if not paths: print("usage: check_type_discipline.py ...", file=sys.stderr) return 2 - violations = sorted(v for path in collect_paths(paths) for v in check_file(path)) + targets = tuple(collect_paths(paths)) + violations = sorted(scan_paths(targets)) for v in violations: print(v.render()) diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 4fea5761cc8..7d59e5a5dba 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -8,9 +8,13 @@ in the test body. """ import importlib.util +import os +import subprocess import sys from pathlib import Path +import pytest + _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "check_test_quality.py" _spec = importlib.util.spec_from_file_location("check_test_quality", _MODULE_PATH) @@ -548,3 +552,61 @@ def test_the_read_may_sit_a_statement_above_the_store(tmp_path): def test_a_loop_storing_under_a_key_that_is_not_the_loop_variable_is_not_an_inventory(tmp_path): source = _HELPER_DICT_CONFTEST.replace("state[attr] =", 'state["fixed"] =') assert [v.code for v in checker.check_file(_written(tmp_path, source))] == [] + + +_FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1 +_SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare" + + +def _corpus(tmp_path: Path, count: int) -> tuple[Path, ...]: + for index in range(count): + (tmp_path / f"test_gen_{index}.py").write_text( + f"def test_flagged_{index}():\n compute()\n\n\ndef test_clean_{index}():\n assert compute() == {index}\n", + encoding="utf-8", + ) + return tuple(sorted(tmp_path.rglob("*.py"))) + + +def _run_checker(target: Path) -> list[str]: + completed = subprocess.run( + [sys.executable, str(_MODULE_PATH), str(target)], + capture_output=True, text=True, timeout=300, + ) + return completed.stdout.splitlines() + + +def test_worker_count_stays_serial_below_the_threshold(): + assert checker._worker_count(checker.PARALLEL_MIN_PATHS - 1) == 1 + + +def test_worker_count_fans_out_at_the_threshold(): + assert checker._worker_count(checker.PARALLEL_MIN_PATHS) == max( + 1, min(os.cpu_count() or 1, checker.MAX_WORKERS) + ) + + +def test_worker_count_never_exceeds_the_cap(): + assert checker._worker_count(100_000) <= checker.MAX_WORKERS + + +def test_scan_paths_below_the_threshold_returns_every_violation(tmp_path): + paths = _corpus(tmp_path, 3) + assert checker._worker_count(len(paths)) == 1 + assert [v.code for v in checker.scan_paths(paths)] == ["TQ001"] * 3 + + +@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY) +def test_a_fanned_out_run_reports_exactly_what_a_serial_run_reports(tmp_path): + paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5) + serial = [v.render() for v in sorted(v for path in paths for v in checker.check_file(path))] + assert serial, "corpus must produce violations or the comparison proves nothing" + assert _run_checker(tmp_path) == serial + + +@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY) +def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): + paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5) + reported = _run_checker(tmp_path) + assert len(reported) == len(paths) + assert len({line.split(":")[0] for line in reported}) == len(paths) + assert all(" TQ001 " in line for line in reported) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 84dd547ad80..2d49332e687 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -8,10 +8,14 @@ a test fail. The comment-scanner cases are the regression for the readline path: import importlib.util import json +import os import re +import subprocess import sys from pathlib import Path +import pytest + _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "check_type_discipline.py" _spec = importlib.util.spec_from_file_location("check_type_discipline", _MODULE_PATH) @@ -695,3 +699,61 @@ def test_budget_covers_exactly_the_checker_rules(): for spec in budget.values(): assert isinstance(spec["limit"], int) assert spec["limit"] >= 0 + + +_FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1 +_SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare" + + +def _corpus(tmp_path: Path, count: int) -> tuple[Path, ...]: + for index in range(count): + (tmp_path / f"mod_{index}.py").write_text( + f"def build_{index}(items: list[int]) -> None:\n return None\n", + encoding="utf-8", + ) + return tuple(sorted(tmp_path.rglob("*.py"))) + + +def _run_checker(target: Path) -> list[str]: + completed = subprocess.run( + [sys.executable, str(_MODULE_PATH), str(target)], + capture_output=True, text=True, timeout=300, + ) + return completed.stdout.splitlines() + + +def test_worker_count_stays_serial_below_the_threshold(): + assert checker._worker_count(checker.PARALLEL_MIN_PATHS - 1) == 1 + + +def test_worker_count_fans_out_at_the_threshold(): + assert checker._worker_count(checker.PARALLEL_MIN_PATHS) == max( + 1, min(os.cpu_count() or 1, checker.MAX_WORKERS) + ) + + +def test_worker_count_never_exceeds_the_cap(): + assert checker._worker_count(100_000) <= checker.MAX_WORKERS + + +def test_scan_paths_below_the_threshold_returns_every_violation(tmp_path): + paths = _corpus(tmp_path, 3) + assert checker._worker_count(len(paths)) == 1 + found = checker.scan_paths(paths) + assert found and len({v.path for v in found}) == 3 + + +@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY) +def test_a_fanned_out_run_reports_exactly_what_a_serial_run_reports(tmp_path): + paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5) + serial = [v.render() for v in sorted(v for path in paths for v in checker.check_file(path))] + assert serial, "corpus must produce violations or the comparison proves nothing" + assert _run_checker(tmp_path) == serial + + +@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY) +def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): + paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5) + reported = _run_checker(tmp_path) + assert reported + assert len({line.split(":")[0] for line in reported}) == len(paths) From 346c69386026c245df93a3ef284e4c525df69640 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:45:30 -0700 Subject: [PATCH 092/273] ci: port the Postgres suites off CircleCI onto service containers (#37785) * ci: port the Postgres suites off CircleCI onto service containers proxy_behavior_tests, proxy_security_tests and schema_migration_check were near-identical CircleCI jobs: a Postgres sidecar, a schema seed, and one pytest tree each. They ran nowhere else, and CircleCI holds none of the branch ruleset's required checks, so the signal they produced gated nothing. test-postgres.yml runs the same three trees on a Postgres service container as one matrix, keeping each suite's own seeding rather than normalising it: the behavior and security trees keep `prisma db push`, and the migration tree keeps an empty database, which is what it needs to apply every committed migration itself. Their CircleCI definitions and workflow entries go with them, taking the config from 47 jobs to 44. assert_ci_coverage.py stays green: dropping the new workflow fails the census on exactly these trees, so the coverage moved rather than went missing. auth_ui_unit_tests is deliberately left behind. Ported, two of its tests fail because prepare_metadata_fields refuses enterprise-only keys without LITELLM_LICENSE, which exists as a CircleCI project variable and has no GitHub Actions secret. Creating that secret is a human action, so the job stays on CircleCI until it exists rather than shipping a red shard or quietly deselecting the two tests. * chore(ci): drop the narrative header from test-postgres.yml --- .circleci/config.yml | 126 ------------------------ .github/workflows/test-postgres.yml | 145 ++++++++++++++++++++++++++++ 2 files changed, 145 insertions(+), 126 deletions(-) create mode 100644 .github/workflows/test-postgres.yml diff --git a/.circleci/config.yml b/.circleci/config.yml index 5e77729df29..2586485e79c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -651,126 +651,6 @@ jobs: - auth_ui_unit_tests_coverage.xml - auth_ui_unit_tests_coverage - proxy_behavior_tests: - docker: - - *python312_image - - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 - environment: - POSTGRES_USER: postgres - POSTGRES_PASSWORD: postgres - POSTGRES_DB: litellm_test - working_directory: ~/project - environment: - DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" - steps: - - checkout - - skip_if_unrelated_changes - - setup_google_dns - - install_uv - - install_rust - - run: - name: Install Dependencies - command: | - uv sync --frozen --all-groups --all-extras --python 3.12 - - wait_for_service: - url: tcp://localhost:5432 - timeout: "60" - - run: - name: Seed DB schema via prisma db push - command: | - uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss - - run: - name: Generate Prisma Client - command: uv run --no-sync python -m prisma generate - - run: - name: Run proxy management behavior tests - command: | - mkdir -p test-results - uv run --no-sync python -m pytest tests/proxy_behavior \ - -v --junitxml=test-results/junit.xml --durations=10 - no_output_timeout: 15m - - store_test_results: - path: test-results - - proxy_security_tests: - docker: - - *python312_image - - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 - environment: - POSTGRES_USER: postgres - POSTGRES_PASSWORD: postgres - POSTGRES_DB: litellm_test - working_directory: ~/project - environment: - DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" - steps: - - checkout - - skip_if_unrelated_changes - - setup_google_dns - - install_uv - - install_rust - - run: - name: Install Dependencies - command: | - uv sync --frozen --all-groups --all-extras --python 3.12 - - wait_for_service: - url: tcp://localhost:5432 - timeout: "60" - - run: - name: Seed DB schema via prisma db push - command: | - uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss - - run: - name: Generate Prisma Client - command: uv run --no-sync python -m prisma generate - - run: - name: Run proxy security tests - command: | - mkdir -p test-results - uv run --no-sync python -m pytest tests/proxy_security_tests \ - -v --junitxml=test-results/junit.xml --durations=10 - no_output_timeout: 15m - - store_test_results: - path: test-results - - schema_migration_check: - docker: - - *python312_image - - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 - environment: - POSTGRES_USER: postgres - POSTGRES_PASSWORD: postgres - POSTGRES_DB: litellm_test - working_directory: ~/project - environment: - # An empty database; the test applies every committed migration itself. - DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" - steps: - - checkout - - skip_if_unrelated_changes - - setup_google_dns - - install_uv - - install_rust - - run: - name: Install Dependencies - command: | - uv sync --frozen --all-groups --all-extras --python 3.12 - - wait_for_service: - url: tcp://localhost:5432 - timeout: "60" - - run: - name: Generate Prisma Client - command: uv run --no-sync python -m prisma generate - - run: - name: Check schema.prisma is in sync with committed migrations - command: | - mkdir -p test-results - uv run --no-sync python -m pytest tests/proxy_migration_tests \ - -v --junitxml=test-results/junit.xml --durations=10 - no_output_timeout: 15m - - store_test_results: - path: test-results - litellm_router_testing: # Runs all tests with the "router" keyword docker: - *python312_image @@ -3105,12 +2985,6 @@ workflows: filters: *main_branches - auth_ui_unit_tests: filters: *main_branches - - proxy_behavior_tests: - filters: *main_branches - - proxy_security_tests: - filters: *main_branches - - schema_migration_check: - filters: *main_branches - build_docker_database_image: filters: *main_branches - e2e_ui_testing: diff --git a/.github/workflows/test-postgres.yml b/.github/workflows/test-postgres.yml new file mode 100644 index 00000000000..96c514dff7c --- /dev/null +++ b/.github/workflows/test-postgres.yml @@ -0,0 +1,145 @@ +name: "Postgres Tests" + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + push: + branches: + - main + - litellm_internal_staging + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + postgres: + name: ${{ matrix.shard }} + runs-on: ubuntu-latest + timeout-minutes: ${{ matrix.job-timeout-minutes }} + permissions: + contents: read + + services: + postgres: + image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + env: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 10 + + strategy: + fail-fast: false + matrix: + include: + - shard: proxy-behavior + test-path: "tests/proxy_behavior" + seed: db-push + workers: 0 + timeout-minutes: 25 + job-timeout-minutes: 50 + + - shard: proxy-security + test-path: "tests/proxy_security_tests" + seed: db-push + workers: 0 + timeout-minutes: 15 + job-timeout-minutes: 40 + + - shard: schema-migration + test-path: "tests/proxy_migration_tests" + seed: none + workers: 0 + timeout-minutes: 20 + job-timeout-minutes: 45 + + env: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + timeout-minutes: 3 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + timeout-minutes: 2 + uses: ./.github/actions/detect-changes + + - name: Set up Python + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache uv dependencies + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cache/uv + .venv + key: ${{ runner.os }}-uv-postgres-${{ hashFiles('uv.lock') }} + restore-keys: | + ${{ runner.os }}-uv-postgres- + + - name: Install dependencies + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 12 + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras + + - name: Cache Prisma binaries + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Seed database schema + if: steps.changes.outputs.decision != 'skip' && matrix.seed != 'none' + timeout-minutes: 10 + run: | + uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss + + - name: Run tests + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: ${{ matrix.timeout-minutes }} + env: + TEST_PATH: ${{ matrix.test-path }} + WORKERS: ${{ matrix.workers }} + run: | + if [ "${WORKERS}" = "0" ]; then + uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10 + else + uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10 -n "${WORKERS}" + fi From a734afca322959557694734e179b095f2e6f1039 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:54:30 -0700 Subject: [PATCH 093/273] feat(ci): gate patching of SDK internals in tests as TQ008 (#37787) * feat(ci): gate patching of SDK internals in tests as TQ008 TQ002 catches the narrowest symptom of the suite's dominant mocking idiom, patch X then assert only that X was called. The idiom itself is wider: tests reach for litellm's own functions instead of faking the wire, so they pin how the code is wired rather than what it does, and a test that patches internals but makes weak real assertions trips nothing today. TQ008 counts patch targets rooted at `litellm`, both the dotted string form and the attribute chain handed to patch.object, and ratchets like every other rule. Mocking anything outside the SDK is untouched: respx, httpx transports and third-party clients do not trip it, which is the point, since those are the patterns this is meant to move the suite toward. Seeded at 9,643, in line with the ~9.4k patch sites an independent grep found in the mirror. The burn-down horizon is long; the value here is stopping the flow rather than clearing the stock. Five existing rule tests patched `litellm.completion` incidentally and now report TQ008 alongside what they were pinning. Their expected values are updated to the accurate pair rather than loosened, so they keep failing on a regression in either rule. * test: add TQ008 to the shipped-budget rule canary * fix(ci): resolve imported SDK names in TQ008 patch.object(handler.OpenAIChatCompletion, ...) after a from-import reaches the same internal as the dotted string form, but the rule only saw the bare local name and let it through. Import bindings are now resolved to the path they stand for, so the aliased, renamed and from-imported forms all read alike and the reported target is the real one. That is 1,496 patches the ratchet could not see, so the TQ008 limit moves from 9,643 to 11,139. Third-party names and locals with no SDK import behind them stay unflagged. --- scripts/check_test_quality.py | 66 +++++++++ test-quality-budget.json | 3 + tests/test_litellm/test_check_test_quality.py | 137 +++++++++++++++++- tests/test_litellm/test_test_quality_gate.py | 2 +- 4 files changed, 202 insertions(+), 6 deletions(-) diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 6964aed56e4..41342acd23a 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -48,6 +48,10 @@ TQ006 A `pytest.skip` reached only when a credential-shaped environment variab deliberate branch. The gate follows one local or module-level binding, which is the `key = os.getenv(...)` then `if not key: pytest.skip(...)` shape most of these use. +TQ008 A `patch(...)` whose target is a `litellm.` internal. Patching the SDK's own + functions pins the test to the current wiring instead of the behaviour, and it + is the idiom the suite reaches for instead of faking the HTTP boundary. Mocking + a third-party client, a transport, or anything outside `litellm.` is untouched. TQ007 A module global that a conftest saves before every test and restores after it. The save/restore list is a hand-maintained inventory of the leaks the suite already knows about, so it is allowed to shrink and never to grow: a new entry @@ -469,6 +473,67 @@ def iter_global_mutation_violations(path: Path, tree: ast.Module) -> Iterator[Vi ) +def _is_sdk_internal(dotted: str) -> bool: + return dotted == SDK_MODULE or dotted.startswith(f"{SDK_MODULE}.") + + +def _sdk_import_bindings(tree: ast.Module) -> Iterator[tuple[str, str]]: + """(local name, dotted path) for every import that binds something under `litellm`.""" + for node in ast.walk(tree): + if isinstance(node, ast.Import): + yield from ( + (alias.asname, alias.name) if alias.asname else (root, root) + for alias in node.names + if _is_sdk_internal(alias.name) + for root in (alias.name.partition(".")[0],) + ) + elif isinstance(node, ast.ImportFrom) and node.module and _is_sdk_internal(node.module): + yield from ((alias.asname or alias.name, f"{node.module}.{alias.name}") for alias in node.names) + + +def _sdk_aliases(tree: ast.Module) -> Mapping[str, str]: + """Local names bound to something under `litellm`, mapped to the path they stand for. + + `from litellm.llms.openai.chat import handler` then `patch.object(handler.X, ...)` + reaches the same internal as the dotted string form and has to read the same way. + """ + return MappingProxyType({name: dotted for name, dotted in _sdk_import_bindings(tree)}) + + +def _resolved(dotted: str, aliases: Mapping[str, str]) -> str: + root, _, rest = dotted.partition(".") + base: Final = aliases.get(root, root) + return f"{base}.{rest}" if rest else base + + +def _patch_targets(call: ast.Call, aliases: Mapping[str, str]) -> Iterator[str]: + """What a patch installer is replacing: the dotted string it names, or the + attribute chain handed to `patch.object` / `patch.dict`, resolved through the + module's imports so a locally bound SDK object reads as its full path.""" + for first in call.args[:1]: + if isinstance(first, ast.Constant) and isinstance(first.value, str): + yield first.value + elif dotted := _dotted_name(first): + yield _resolved(dotted, aliases) + + +def iter_internal_patch_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + aliases: Final = _sdk_aliases(tree) + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and _is_patch_installer(_dotted_name(node.func))): + continue + for target in _patch_targets(node, aliases): + if _is_sdk_internal(target): + yield Violation( + path, + node.lineno, + "TQ008", + f"patches `{target}`, an SDK internal, so the test is pinned to how the code is " + "wired rather than what it does; fake the HTTP boundary (respx / MockTransport) " + f"or inject the collaborator (suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + + def _environ_keys(node: ast.AST) -> Iterator[str]: for inner in ast.walk(node): if isinstance(inner, ast.Call) and _dotted_name(inner.func) in ENVIRON_READERS: @@ -680,6 +745,7 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_global_mutation_violations(path, tree), *iter_credential_skip_violations(path, tree), *iter_conftest_inventory_violations(path, tree), + *iter_internal_patch_violations(path, tree), ) if violation.line not in skip ) diff --git a/test-quality-budget.json b/test-quality-budget.json index 0dea4e8fe93..4a7bc7edff2 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -19,5 +19,8 @@ }, "TQ007": { "limit": 117 + }, + "TQ008": { + "limit": 11139 } } diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 7d59e5a5dba..a75b1e43fb7 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -187,7 +187,7 @@ def test_mock_echo_is_flagged(tmp_path): " run()\n" " mock_completion.assert_called_once()\n" ) - assert _codes(tmp_path, source) == ["TQ002"] + assert _codes(tmp_path, source) == ["TQ002", "TQ008"] def test_call_args_inspection_is_mock_echo(tmp_path): @@ -200,7 +200,7 @@ def test_call_args_inspection_is_mock_echo(tmp_path): " run()\n" " assert mock_completion.call_args[1]['model'] == 'gpt-4o'\n" ) - assert _codes(tmp_path, source) == ["TQ002"] + assert _codes(tmp_path, source) == ["TQ002", "TQ008"] def test_patch_decorator_counts_as_installing_a_patch(tmp_path): @@ -213,7 +213,7 @@ def test_patch_decorator_counts_as_installing_a_patch(tmp_path): " run()\n" " mock_completion.assert_called_once()\n" ) - assert _codes(tmp_path, source) == ["TQ002"] + assert _codes(tmp_path, source) == ["TQ002", "TQ008"] def test_patching_but_asserting_the_output_is_not_mock_echo(tmp_path): @@ -227,7 +227,7 @@ def test_patching_but_asserting_the_output_is_not_mock_echo(tmp_path): " mock_completion.assert_called_once()\n" " assert result.choices[0].message.content == 'pong'\n" ) - assert _codes(tmp_path, source) == [] + assert _codes(tmp_path, source) == ["TQ008"] def test_asserting_without_patching_is_not_mock_echo(tmp_path): @@ -244,7 +244,7 @@ def test_a_test_with_no_assertions_is_tq001_not_tq002(tmp_path): " with patch('litellm.completion'):\n" " run()\n" ) - assert _codes(tmp_path, source) == ["TQ001"] + assert _codes(tmp_path, source) == ["TQ001", "TQ008"] def test_sys_path_insert_is_flagged(tmp_path): @@ -554,6 +554,133 @@ def test_a_loop_storing_under_a_key_that_is_not_the_loop_variable_is_not_an_inve assert [v.code for v in checker.check_file(_written(tmp_path, source))] == [] +def test_patching_an_sdk_function_by_string_is_flagged(tmp_path): + source = 'from unittest.mock import patch\n\n\n@patch("litellm.completion")\ndef test_x(m):\n assert m\n' + assert "TQ008" in _codes(tmp_path, source) + + +def test_patching_a_deep_sdk_path_is_flagged(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch("litellm.llms.openai.chat.handler.OpenAIChatCompletion.completion"):\n' + " assert True\n" + ) + assert "TQ008" in _codes(tmp_path, source) + + +def test_patch_object_rooted_at_the_sdk_is_flagged(tmp_path): + source = ( + "import litellm\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(litellm, "api_key", "x"):\n' + " assert True\n" + ) + assert "TQ008" in _codes(tmp_path, source) + + +def test_patch_object_on_a_from_imported_sdk_module_is_flagged(tmp_path): + source = ( + "from litellm.llms.openai.chat import handler\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(handler.OpenAIChatCompletion, "completion"):\n' + " assert True\n" + ) + assert "TQ008" in _codes(tmp_path, source) + + +def test_patch_object_on_an_aliased_sdk_module_is_flagged(tmp_path): + source = ( + "import litellm.llms.openai.chat.handler as oai\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(oai.OpenAIChatCompletion, "completion"):\n' + " assert True\n" + ) + assert "TQ008" in _codes(tmp_path, source) + + +def test_patch_object_on_a_renamed_sdk_symbol_is_flagged(tmp_path): + source = ( + "from litellm.utils import get_llm_provider as glp\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(glp, "__wrapped__"):\n' + " assert True\n" + ) + assert "TQ008" in _codes(tmp_path, source) + + +def test_the_reported_target_is_the_resolved_sdk_path(tmp_path): + source = ( + "from litellm.llms.openai.chat import handler\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(handler.OpenAIChatCompletion, "completion"):\n' + " assert True\n" + ) + reported = [v.message for v in checker.check_file(_written(tmp_path, source)) if v.code == "TQ008"] + assert reported + assert "litellm.llms.openai.chat.handler.OpenAIChatCompletion" in reported[0] + + +def test_patch_object_on_a_from_imported_third_party_is_not_flagged(tmp_path): + source = ( + "from openai import OpenAI\nfrom unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch.object(OpenAI, "chat"):\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + +def test_a_local_name_with_no_sdk_import_behind_it_is_not_flagged(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x(handler):\n" + ' with patch.object(handler, "completion"):\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + +def test_mocking_a_third_party_client_is_not_flagged(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch("openai.OpenAI.chat"):\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + +def test_mocking_the_http_transport_is_not_flagged(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch("httpx.AsyncClient.send"):\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + +def test_a_name_merely_starting_with_litellm_is_not_the_sdk(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch("litellm_enterprise.thing.go"):\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + +def test_an_sdk_patch_can_be_suppressed(tmp_path): + source = ( + "from unittest.mock import patch\n\n\n" + "def test_x():\n" + ' with patch("litellm.completion"): # test-quality-ok: pinning the router seam\n' + " assert True\n" + ) + assert "TQ008" not in _codes(tmp_path, source) + + _FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1 _SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare" diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 3bf4b89ac4e..8cce6bc735a 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -144,5 +144,5 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): import json budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) - assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007"} + assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"} assert all(spec["limit"] >= 0 for spec in budget.values()) From 6c30b4331d2d9b5571efcb0ff94fb731651d500d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:55:01 -0700 Subject: [PATCH 094/273] ci: measure enterprise/ coverage (#37788) codecov.yaml has carried an `Enterprise` component scoped to `enterprise/**` since it was written, and it has never received a line of data. Every one of the 19 coverage invocations across the unit base, the MCP workflow and the CircleCI config passes `--cov=./litellm` and nothing else, so 11,203 lines of paid-customer code sat outside the measured universe while the reported number described only the rest. litellm-enterprise is a uv workspace member and a direct dependency, so every job that syncs already has it installed and importable; only the measurement was missing. Measured on tests/test_litellm/enterprise, the shard that exercises this code: 0 enterprise files in the report before, 142 after, at `enterprise/...` paths that match the component's existing glob. That shard alone puts enterprise at 30.8%, which nudged its own total from 24.09% to 24.20% rather than down. The aggregate direction across every shard is not knowable until they all report, and a drop there is the instrument working, not a regression. --- .circleci/config.yml | 32 +++++++++++++-------------- .github/workflows/_test-unit-base.yml | 4 ++-- .github/workflows/test-mcp.yml | 2 +- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 2586485e79c..a261b7c78fc 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -430,7 +430,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv \ - --cov=./litellm \ + --cov=./litellm --cov=./enterprise/litellm_enterprise \ --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=20 \ @@ -504,7 +504,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv \ - --cov=./litellm \ + --cov=./litellm --cov=./enterprise/litellm_enterprise \ --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=20 \ @@ -631,7 +631,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -v -x \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 2" @@ -738,7 +738,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -v -x \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 4" @@ -865,7 +865,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=20 \ -n 4 \ @@ -910,7 +910,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5" no_output_timeout: 15m @@ -954,7 +954,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 2 \ @@ -1000,7 +1000,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ --retries 3 --retry-delay 5" @@ -1091,7 +1091,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 4" @@ -1135,7 +1135,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 4" @@ -1213,7 +1213,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 2" @@ -1257,7 +1257,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 2" @@ -1302,7 +1302,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 \ -n 4" @@ -1381,7 +1381,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ -n 4 \ --junitxml=test-results/junit.xml \ --durations=5 \ @@ -1426,7 +1426,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5" no_output_timeout: 15m @@ -1479,7 +1479,7 @@ jobs: --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ -vv -x -s \ - --cov=./litellm --cov-report=xml \ + --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \ --junitxml=test-results/junit.xml \ --durations=5 -n 2 \ --reruns 2 --reruns-delay 1" diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index b7d185bd0b9..c4045a08ffb 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -149,7 +149,7 @@ jobs: --reruns "${RERUNS}" \ --reruns-delay 1 \ --durations=20 \ - --cov=./litellm \ + --cov=./litellm --cov=./enterprise/litellm_enterprise \ --cov-report=xml:coverage.xml \ --cov-config=pyproject.toml else @@ -161,7 +161,7 @@ jobs: --reruns-delay 1 \ --dist="${DIST}" \ --durations=20 \ - --cov=./litellm \ + --cov=./litellm --cov=./enterprise/litellm_enterprise \ --cov-report=xml:coverage.xml \ --cov-config=pyproject.toml fi diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 6ea814dc2de..93ffcbe0586 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -60,4 +60,4 @@ jobs: - name: Run MCP tests if: steps.changes.outputs.decision != 'skip' run: | - uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov-report=xml --durations=5 + uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5 From b31484ed196d310878dfa3b1b90571d77d4f7f76 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:56:43 -0700 Subject: [PATCH 095/273] ci: run the keyless caching tests that ran in no job (#37790) The allowlist recorded eight files in tests/local_testing, 118 tests, that every job globbing that directory then deselects: local_testing_part1 and part2 carry `-k "... and not caching and not cache"`, and the other three keep one unrelated keyword each. They counted as covered while running nowhere. Five of the eight need nothing. Measured with no provider credentials and no Redis: test_cache_preset_key, test_caching_handler, test_prompt_caching, test_responses_stream_cache_keys and test_unit_test_caching pass, 45 tests together, and they now run as a caching-local shard. The other three stay allowlisted with what they actually need recorded rather than a question: test_caching wants Redis and a provider key for 37 of its 65, disk-cache wants OPENAI_API_KEY for 2 of 4, gcs-cache wants GCS credentials for all 4. Taking them off the allowlist exposed a gap in the slice guard itself: it reasoned only about CircleCI `-k` expressions, so a file every slice drops read as unrun even when a workflow names it outright. It now credits workflow test-paths the way the census already does, and only workflows, so a tree only CircleCI globs is still reported. --- .github/ci-coverage-allowlist.yml | 25 ++++++++----------- .github/scripts/assert_ci_coverage.py | 16 ++++++++++++ .github/workflows/test-unit.yml | 13 ++++++++++ tests/test_litellm/test_assert_ci_coverage.py | 24 ++++++++++++++++++ 4 files changed, 64 insertions(+), 14 deletions(-) diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 918589f84d1..b6ec08bd295 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -5,24 +5,21 @@ description: >- test_paths: - reason: >- - The caching suite in tests/local_testing, which runs nowhere. Every job that globs that - directory either deselects it (local_testing_part1 and part2 carry `-k "... and not caching - and not cache"`) or keeps only another keyword (langfuse, router, assistants), and no job - names these files the way redis_caching_unit_tests names test_dual_cache.py. Measured - 2026-08-20 by collecting the directory under each job's own selector: 118 tests across - these eight files are selected by none of them. Listed so the gap is a decision rather - than an accident, and so the --slices guard has a baseline to ratchet down from. Revisit - when tests/local_testing is ported off CircleCI, where the keyless part of this suite - belongs in a real job + What is left of the caching suite in tests/local_testing that runs nowhere. Every job that + globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and + not caching and not cache"`) or keeps only another keyword (langfuse, router, assistants), + and no job names these files the way redis_caching_unit_tests names test_dual_cache.py. + The gap was eight files and 118 tests when measured 2026-08-20; the five keyless ones now + run in the caching-local shard, leaving these three. Measured 2026-08-21 with no provider + credentials and no Redis: test_caching.py needs both (37 of 65 fail without them), + test_disk_cache_unit_tests.py needs OPENAI_API_KEY for 2 of its 4, and + test_gcs_cache_unit_tests.py needs GCS credentials for all 4. They want the keyless/live + split that porting tests/local_testing off CircleCI will force, not a job that is red by + construction paths: - - tests/local_testing/test_cache_preset_key.py - tests/local_testing/test_caching.py - - tests/local_testing/test_caching_handler.py - tests/local_testing/test_disk_cache_unit_tests.py - tests/local_testing/test_gcs_cache_unit_tests.py - - tests/local_testing/test_prompt_caching.py - - tests/local_testing/test_responses_stream_cache_keys.py - - tests/local_testing/test_unit_test_caching.py - reason: >- The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than from a pull request; it needs a live gateway and provider credentials no PR job holds diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index c8572d9f6ef..411852acb98 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -312,8 +312,23 @@ def _matchable_names(relative_path: str) -> frozenset[str]: ) +def _workflow_named_tokens() -> frozenset[str]: + """Test tokens a GitHub Actions job names directly. + + A CircleCI `-k` that deselects a file no longer means the file runs nowhere once a + workflow names it, so the slice check has to credit those the same way the census does. + """ + return _invoked_test_tokens( + scalar + for path in _config_files() + if path != CIRCLECI_CONFIG + for scalar in _scalars(yaml.safe_load(path.read_text(encoding="utf-8")), path.name) + ) + + def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]: slices: Final = _slices() + named_by_workflow: Final = _workflow_named_tokens() globbed: Final = tuple( path for path in _test_files() @@ -326,6 +341,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]: ) for path in globbed if not allowlist.covers_test(path) + and not any(_token_covers(token, path) for token in named_by_workflow) and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices) ) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 71eb0958bec..5a75b4af7db 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -200,6 +200,19 @@ jobs: timeout-minutes: 20 job-timeout-minutes: 60 + - shard: caching-local + artifact-name: caching-local + test-path: >- + tests/local_testing/test_cache_preset_key.py + tests/local_testing/test_caching_handler.py + tests/local_testing/test_prompt_caching.py + tests/local_testing/test_responses_stream_cache_keys.py + tests/local_testing/test_unit_test_caching.py + workers: 2 + reruns: 2 + timeout-minutes: 20 + job-timeout-minutes: 55 + - shard: responses-caching-types artifact-name: responses-caching-types test-path: >- diff --git a/tests/test_litellm/test_assert_ci_coverage.py b/tests/test_litellm/test_assert_ci_coverage.py index 4c8f4cc2984..d948c1a4155 100644 --- a/tests/test_litellm/test_assert_ci_coverage.py +++ b/tests/test_litellm/test_assert_ci_coverage.py @@ -280,3 +280,27 @@ def test_a_dockerfile_directory_entry_is_stale_because_only_an_exact_path_exempt dockerfiles=("docker/Dockerfile.database",), ) assert [f.subject for f in findings] == ["docker"] + + +def test_a_workflow_that_names_a_file_clears_it_from_the_slice_check(): + named = coverage._workflow_named_tokens() + assert named, "the workflows must name some test paths or the check proves nothing" + assert any( + coverage._token_covers(token, "tests/local_testing/test_caching_handler.py") + for token in named + ) + + +def test_the_slice_check_credits_only_workflows_never_the_circleci_config(): + named = coverage._workflow_named_tokens() + circleci_only = "tests/proxy_admin_ui_tests" + assert not any(coverage._token_covers(token, f"{circleci_only}/test_key_management.py") for token in named), ( + "a tree only CircleCI globs must not be credited to a workflow" + ) + + +def test_a_file_no_workflow_names_is_still_reported_when_every_slice_drops_it(): + named = coverage._workflow_named_tokens() + assert not any( + coverage._token_covers(token, "tests/local_testing/test_caching.py") for token in named + ), "test_caching.py is allowlisted, not run; crediting it would hide a real gap" From ae0e8a20dbf9200a5ec324d8381a1dc3c952fa46 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:57:16 -0700 Subject: [PATCH 096/273] fix(ci): run the migration DDL guard, and stop it reading comments as SQL (#37791) * fix(ci): make the migration DDL guard run, and stop it reading comments as SQL TestMigrationSQLIdempotency requires guarded DDL across litellm-proxy-extras and has never run in any job, so the convention eroded quietly. Four of its assertions fail today, and it was allowlisted rather than wired up because fixing the migrations is not an option: Prisma checksums an applied migration, so editing one breaks `migrate deploy` for every existing install. Two things were wrong with the guard itself. It scanned raw lines, so Prisma's own `-- CREATE INDEX CONCURRENTLY ...` explanations counted as the statements they describe, which is two of the reported migrations. And it had no way to say "these predate the rule", so the only options were editing immutable files or leaving the whole file unrun. Comments are now stripped before matching, on the drop-column rule too, and the migrations that already violate are named once in _PRE_GUARD_MIGRATIONS. The rules bind everything after them, so a new migration with bare CREATE TABLE, ADD COLUMN, CREATE INDEX or an unguarded ADD CONSTRAINT now fails a check instead of landing unnoticed. That set is 14 migrations, not the 13 previously recorded, measured after comment-stripping. It can only shrink: a test fails if an entry names no migration on disk, and another fails if an entry no longer violates anything. The file now runs as a proxy-extras shard and comes off the coverage allowlist. * fix(ci): strip block comments in the migration guard too Prisma opens a destructive migration with a /* Warnings: You are about to drop the column ... */ header. Nothing in the tree trips a rule on that text today, but it is prose about a statement rather than the statement, and the line-comment fix left the class open. Bodies are blanked rather than removed so the reported line number still points at the real statement. --- .github/ci-coverage-allowlist.yml | 11 - .github/workflows/test-unit.yml | 8 + .../test_litellm_proxy_extras_utils.py | 172 ++++++++++++-- whitelisted_bedrock_models.txt | 219 ++++++++++++++++++ 4 files changed, 379 insertions(+), 31 deletions(-) create mode 100644 whitelisted_bedrock_models.txt diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index b6ec08bd295..f7a785b3b80 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -89,17 +89,6 @@ test_paths: - tests/integration/sandbox/test_e2b_sandbox.py - tests/integration/test_oci_integration.py - tests/integration/test_oci_proxy_integration.py - - reason: >- - A unit test for the proxy-extras package that no job invokes, while the package's other - tests live under tests/proxy_migration_tests. Measured 2026-08-20: 24 of its 28 tests pass - and the 4 in TestMigrationSQLIdempotency fail, because 13 migrations from 2026-03 onward use - bare CREATE TABLE, ADD COLUMN, CREATE INDEX and ADD CONSTRAINT rather than the guarded forms - this file requires. It also matches those keywords inside SQL comments, so two further - migrations are reported that are in fact fine. Wiring it up means deciding what to do about - the 13 first, and they cannot simply be edited: Prisma checksums an applied migration, so a - changed one breaks migrate deploy for existing installs - paths: - - tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py dockerfiles: - reason: >- diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 5a75b4af7db..368680d9ec2 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -213,6 +213,14 @@ jobs: timeout-minutes: 20 job-timeout-minutes: 55 + - shard: proxy-extras + artifact-name: proxy-extras + test-path: "tests/litellm-proxy-extras" + workers: 2 + reruns: 2 + timeout-minutes: 20 + job-timeout-minutes: 55 + - shard: responses-caching-types artifact-name: responses-caching-types test-path: >- diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 6f4979b9b84..09f3e0ba34f 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -161,6 +161,58 @@ def _get_all_migrations(): return results +_LINE_COMMENT = re.compile(r"--.*$") +_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL) + +_PRE_GUARD_MIGRATIONS = frozenset({ + "20260331000000_add_prompt_environment_and_created_by", + "20260418000000_add_adaptive_router_tables", + "20260429161855_workflow_runs_tables", + "20260605182307_add_timeout_to_mcp_server_table", + "20260626120000_add_mcp_tool_search_enabled", + "20260629000000_add_max_concurrent_requests_to_mcp_server_table", + "20260710000000_add_dcr_bridge_to_mcp_server_table", + "20260713230852_add_key_type_to_litellm_verification_token", + "20260811172448_add_shadow_eval", + "20260813180408_add_shadow_eval_direction", + "20260814000000_add_proxy_worker_heartbeat", + "20260817143646_add_daily_guardrail_usage_units", + "20260818224500_add_shadow_eval_stopped_by", + "20260819000000_shadow_eval_max_budget", +}) + + +def _blanked_block_comments(sql): + """`sql` with every `/* ... */` body blanked out, newlines kept so lines still count. + + Prisma opens a destructive migration with a `/* Warnings: You are about to drop the + column ... */` header, which is prose about the statement rather than the statement. + """ + return _BLOCK_COMMENT.sub(lambda m: re.sub(r"[^\n]", " ", m.group(0)), sql) + + +def _statements(sql): + """(line_number, sql) for each line, with comments removed. + + Prisma writes its own explanations as `-- CREATE INDEX CONCURRENTLY ...`, which a + raw-line scan reads as the statement it is describing. + """ + return [ + (number, _LINE_COMMENT.sub("", line)) + for number, line in enumerate(_blanked_block_comments(sql).splitlines(), 1) + ] + + +def _guarded_migrations(all_migrations): + """Migrations the DDL rules apply to. Prisma checksums an applied migration, so the + ones that predate these rules cannot be edited without breaking `migrate deploy` + for existing installs; they are named once, and the rules bind everything after. + """ + return [ + (name, sql) for name, sql in all_migrations if name not in _PRE_GUARD_MIGRATIONS + ] + + class TestMigrationSQLIdempotency: """Ensure all migration SQL files use idempotent DDL (IF [NOT] EXISTS). @@ -181,8 +233,8 @@ class TestMigrationSQLIdempotency: def test_create_table_uses_if_not_exists(self, all_migrations): """CREATE TABLE statements must use IF NOT EXISTS""" violations = [] - for migration_name, sql in all_migrations: - for line_num, line in enumerate(sql.splitlines(), 1): + for migration_name, sql in _guarded_migrations(all_migrations): + for line_num, line in _statements(sql): if re.search( r"CREATE\s+TABLE\s+", line, re.IGNORECASE ) and not re.search( @@ -198,8 +250,8 @@ class TestMigrationSQLIdempotency: def test_add_column_uses_if_not_exists(self, all_migrations): """ADD COLUMN statements must use IF NOT EXISTS""" violations = [] - for migration_name, sql in all_migrations: - for line_num, line in enumerate(sql.splitlines(), 1): + for migration_name, sql in _guarded_migrations(all_migrations): + for line_num, line in _statements(sql): if re.search(r"ADD\s+COLUMN\s+", line, re.IGNORECASE) and not re.search( r"ADD\s+COLUMN\s+IF\s+NOT\s+EXISTS", line, re.IGNORECASE ): @@ -212,8 +264,8 @@ class TestMigrationSQLIdempotency: def test_drop_column_uses_if_exists(self, all_migrations): """DROP COLUMN statements must use IF EXISTS""" violations = [] - for migration_name, sql in all_migrations: - for line_num, line in enumerate(sql.splitlines(), 1): + for migration_name, sql in _guarded_migrations(all_migrations): + for line_num, line in _statements(sql): if re.search( r"DROP\s+COLUMN\s+", line, re.IGNORECASE ) and not re.search( @@ -239,7 +291,7 @@ class TestMigrationSQLIdempotency: for migration_name, sql in all_migrations: if migration_name in self._DROP_COLUMN_ALLOWLIST: continue - for line_num, line in enumerate(sql.splitlines(), 1): + for line_num, line in _statements(sql): if re.search(r"DROP\s+COLUMN", line, re.IGNORECASE): violations.append(f" {migration_name}:{line_num}: {line.strip()}") assert ( @@ -251,8 +303,8 @@ class TestMigrationSQLIdempotency: def test_drop_index_uses_if_exists(self, all_migrations): """DROP INDEX statements must use IF EXISTS""" violations = [] - for migration_name, sql in all_migrations: - for line_num, line in enumerate(sql.splitlines(), 1): + for migration_name, sql in _guarded_migrations(all_migrations): + for line_num, line in _statements(sql): if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search( r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE ): @@ -266,8 +318,8 @@ class TestMigrationSQLIdempotency: def test_create_index_uses_if_not_exists(self, all_migrations): """CREATE INDEX statements must use IF NOT EXISTS""" violations = [] - for migration_name, sql in all_migrations: - for line_num, line in enumerate(sql.splitlines(), 1): + for migration_name, sql in _guarded_migrations(all_migrations): + for line_num, line in _statements(sql): if re.search( r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+", line, re.IGNORECASE ) and not re.search( @@ -284,10 +336,9 @@ class TestMigrationSQLIdempotency: def test_rename_column_is_guarded(self, all_migrations): """RENAME COLUMN must be inside a DO $$ IF EXISTS block""" violations = [] - for migration_name, sql in all_migrations: - lines = sql.splitlines() + for migration_name, sql in _guarded_migrations(all_migrations): in_do_block = False - for line_num, line in enumerate(lines, 1): + for line_num, line in _statements(sql): if re.search(r"DO\s+\$\$", line, re.IGNORECASE): in_do_block = True if re.search(r"END\s+\$\$", line, re.IGNORECASE): @@ -305,10 +356,9 @@ class TestMigrationSQLIdempotency: def test_add_constraint_is_guarded(self, all_migrations): """ADD CONSTRAINT must be inside a DO $$ IF NOT EXISTS block""" violations = [] - for migration_name, sql in all_migrations: - lines = sql.splitlines() + for migration_name, sql in _guarded_migrations(all_migrations): in_do_block = False - for line_num, line in enumerate(lines, 1): + for line_num, line in _statements(sql): if re.search(r"DO\s+\$\$", line, re.IGNORECASE): in_do_block = True if re.search(r"END\s+\$\$", line, re.IGNORECASE): @@ -326,10 +376,9 @@ class TestMigrationSQLIdempotency: def test_drop_constraint_is_guarded(self, all_migrations): """DROP CONSTRAINT must be inside a DO $$ IF EXISTS block""" violations = [] - for migration_name, sql in all_migrations: - lines = sql.splitlines() + for migration_name, sql in _guarded_migrations(all_migrations): in_do_block = False - for line_num, line in enumerate(lines, 1): + for line_num, line in _statements(sql): if re.search(r"DO\s+\$\$", line, re.IGNORECASE): in_do_block = True if re.search(r"END\s+\$\$", line, re.IGNORECASE): @@ -343,3 +392,86 @@ class TestMigrationSQLIdempotency: "DROP CONSTRAINT without DO $$ IF EXISTS guard found in migrations:\n" + "\n".join(violations) ) + + +class TestMigrationGuardScope: + """The guard must ignore SQL comments, exempt only the named pre-guard migrations, + and still fail on a new migration that uses bare DDL.""" + + _NEW = "20990101000000_a_new_migration" + + def _run_rules(self, migrations): + suite = TestMigrationSQLIdempotency() + failures = [] + for name in ( + "test_create_table_uses_if_not_exists", + "test_add_column_uses_if_not_exists", + "test_create_index_uses_if_not_exists", + "test_add_constraint_is_guarded", + ): + try: + getattr(suite, name)(migrations) + except AssertionError: + failures.append(name) + return failures + + def test_a_comment_describing_ddl_is_not_the_ddl(self): + sql = '-- CREATE TABLE "Foo" (id TEXT);\n-- ADD COLUMN "bar" TEXT;\n' + assert self._run_rules([(self._NEW, sql)]) == [] + + def test_a_prisma_warning_block_is_not_the_ddl_it_describes(self): + sql = ( + "/*\n" + " Warnings:\n" + "\n" + " - You are about to CREATE TABLE \"Foo\" and ADD COLUMN \"bar\".\n" + "\n" + "*/\n" + 'CREATE TABLE IF NOT EXISTS "Foo" (id TEXT);\n' + ) + assert self._run_rules([(self._NEW, sql)]) == [] + + def test_a_block_comment_does_not_shift_the_reported_line(self): + sql = "/* filler\nfiller */\n" + 'CREATE TABLE "Foo" (id TEXT);\n' + suite = TestMigrationSQLIdempotency() + with pytest.raises(AssertionError) as failure: + suite.test_create_table_uses_if_not_exists([(self._NEW, sql)]) + assert f"{self._NEW}:3:" in str(failure.value) + + def test_a_new_migration_with_bare_create_table_fails(self): + assert "test_create_table_uses_if_not_exists" in self._run_rules( + [(self._NEW, 'CREATE TABLE "Foo" (id TEXT);\n')] + ) + + def test_a_new_migration_with_bare_add_column_fails(self): + assert "test_add_column_uses_if_not_exists" in self._run_rules( + [(self._NEW, 'ALTER TABLE "Foo" ADD COLUMN "bar" TEXT;\n')] + ) + + def test_the_guarded_forms_pass(self): + sql = ( + 'CREATE TABLE IF NOT EXISTS "Foo" (id TEXT);\n' + 'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "bar" TEXT;\n' + 'CREATE INDEX IF NOT EXISTS "Foo_bar_idx" ON "Foo"("bar");\n' + ) + assert self._run_rules([(self._NEW, sql)]) == [] + + def test_a_pre_guard_migration_is_exempt_but_a_new_one_is_not(self): + bare = 'CREATE TABLE "Foo" (id TEXT);\n' + exempt = sorted(_PRE_GUARD_MIGRATIONS)[0] + assert self._run_rules([(exempt, bare)]) == [] + assert self._run_rules([(self._NEW, bare)]) != [] + + def test_every_pre_guard_migration_still_exists_on_disk(self): + present = {name for name, _ in _get_all_migrations()} + missing = _PRE_GUARD_MIGRATIONS - present + assert not missing, f"pre-guard entries naming no migration: {sorted(missing)}" + + def test_no_pre_guard_entry_is_already_clean(self): + by_name = dict(_get_all_migrations()) + redundant = [ + name + for name in sorted(_PRE_GUARD_MIGRATIONS) + if not self._run_rules([(TestMigrationGuardScope._NEW, by_name[name])]) + ] + assert not redundant, f"these no longer violate and should be removed: {redundant}" diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt new file mode 100644 index 00000000000..7e20081988d --- /dev/null +++ b/whitelisted_bedrock_models.txt @@ -0,0 +1,219 @@ +ai21.j2-mid-v1 +ai21.j2-ultra-v1 +ai21.jamba-1-5-large-v1:0 +ai21.jamba-1-5-mini-v1:0 +ai21.jamba-instruct-v1:0 +twelvelabs.pegasus-1-2-v1:0 +us.twelvelabs.pegasus-1-2-v1:0 +eu.twelvelabs.pegasus-1-2-v1:0 +amazon.titan-text-express-v1 +amazon.titan-text-lite-v1 +amazon.titan-text-premier-v1:0 +anthropic.claude-3-5-haiku-20241022-v1:0 +anthropic.claude-3-5-sonnet-20240620-v1:0 +anthropic.claude-3-5-sonnet-20241022-v2:0 +anthropic.claude-3-7-sonnet-20240620-v1:0 +anthropic.claude-3-haiku-20240307-v1:0 +anthropic.claude-3-opus-20240229-v1:0 +anthropic.claude-3-sonnet-20240229-v1:0 +anthropic.claude-instant-v1 +anthropic.claude-mythos-preview +anthropic.claude-v1 +anthropic.claude-v2:1 +apac.anthropic.claude-3-5-sonnet-20240620-v1:0 +apac.anthropic.claude-3-5-sonnet-20241022-v2:0 +apac.anthropic.claude-3-haiku-20240307-v1:0 +apac.anthropic.claude-3-sonnet-20240229-v1:0 +bedrock/*/1-month-commitment/cohere.command-light-text-v14 +bedrock/*/1-month-commitment/cohere.command-text-v14 +bedrock/*/6-month-commitment/cohere.command-light-text-v14 +bedrock/*/6-month-commitment/cohere.command-text-v14 +bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-instant-v1 +bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-v1 +bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-v2:1 +bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-instant-v1 +bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-v1 +bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-v2:1 +bedrock/ap-northeast-1/anthropic.claude-instant-v1 +bedrock/ap-northeast-1/anthropic.claude-v1 +bedrock/ap-northeast-1/anthropic.claude-v2:1 +bedrock/ap-northeast-1/deepseek.v3.2 +bedrock/ap-northeast-1/minimax.minimax-m2.1 +bedrock/ap-northeast-1/minimax.minimax-m2.5 +bedrock/ap-northeast-1/moonshotai.kimi-k2-thinking +bedrock/ap-northeast-1/moonshotai.kimi-k2.5 +bedrock/ap-northeast-1/qwen.qwen3-coder-next +bedrock/moonshotai.kimi-k2-thinking +bedrock/moonshotai.kimi-k2.5 +bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0 +bedrock/ap-south-1/meta.llama3-8b-instruct-v1:0 +bedrock/ap-south-1/deepseek.v3.2 +bedrock/ap-south-1/minimax.minimax-m2.1 +bedrock/ap-south-1/minimax.minimax-m2.5 +bedrock/ap-south-1/moonshotai.kimi-k2-thinking +bedrock/ap-south-1/moonshotai.kimi-k2.5 +bedrock/ap-south-1/qwen.qwen3-coder-next +bedrock/ap-southeast-2/minimax.minimax-m2.5 +bedrock/ap-southeast-3/deepseek.v3.2 +bedrock/ap-southeast-3/minimax.minimax-m2.1 +bedrock/ap-southeast-3/minimax.minimax-m2.5 +bedrock/ap-southeast-3/moonshotai.kimi-k2.5 +bedrock/ap-southeast-3/qwen.qwen3-coder-next +bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0 +bedrock/ca-central-1/meta.llama3-8b-instruct-v1:0 +bedrock/eu-north-1/deepseek.v3.2 +bedrock/eu-north-1/minimax.minimax-m2.1 +bedrock/eu-north-1/minimax.minimax-m2.5 +bedrock/eu-north-1/moonshotai.kimi-k2.5 +bedrock/eu-central-1/1-month-commitment/anthropic.claude-instant-v1 +bedrock/eu-central-1/1-month-commitment/anthropic.claude-v1 +bedrock/eu-central-1/1-month-commitment/anthropic.claude-v2:1 +bedrock/eu-central-1/6-month-commitment/anthropic.claude-instant-v1 +bedrock/eu-central-1/6-month-commitment/anthropic.claude-v1 +bedrock/eu-central-1/6-month-commitment/anthropic.claude-v2:1 +bedrock/eu-central-1/anthropic.claude-instant-v1 +bedrock/eu-central-1/anthropic.claude-v1 +bedrock/eu-central-1/anthropic.claude-v2:1 +bedrock/eu-central-1/minimax.minimax-m2.1 +bedrock/eu-central-1/minimax.minimax-m2.5 +bedrock/eu-central-1/qwen.qwen3-coder-next +bedrock/eu-west-1/meta.llama3-70b-instruct-v1:0 +bedrock/eu-west-1/meta.llama3-8b-instruct-v1:0 +bedrock/eu-west-1/minimax.minimax-m2.1 +bedrock/eu-west-1/minimax.minimax-m2.5 +bedrock/eu-west-1/qwen.qwen3-coder-next +bedrock/eu-west-2/meta.llama3-70b-instruct-v1:0 +bedrock/eu-west-2/meta.llama3-8b-instruct-v1:0 +bedrock/eu-west-2/minimax.minimax-m2.1 +bedrock/eu-west-2/minimax.minimax-m2.5 +bedrock/eu-west-2/qwen.qwen3-coder-next +bedrock/eu-west-3/mistral.mistral-7b-instruct-v0:2 +bedrock/eu-west-3/mistral.mistral-large-2402-v1:0 +bedrock/eu-west-3/mistral.mixtral-8x7b-instruct-v0:1 +bedrock/eu-south-1/minimax.minimax-m2.1 +bedrock/eu-south-1/minimax.minimax-m2.5 +bedrock/eu-south-1/qwen.qwen3-coder-next +bedrock/invoke/anthropic.claude-3-5-sonnet-20240620-v1:0 +bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0 +bedrock/sa-east-1/meta.llama3-8b-instruct-v1:0 +bedrock/sa-east-1/deepseek.v3.2 +bedrock/sa-east-1/minimax.minimax-m2.1 +bedrock/sa-east-1/minimax.minimax-m2.5 +bedrock/sa-east-1/moonshotai.kimi-k2-thinking +bedrock/sa-east-1/moonshotai.kimi-k2.5 +bedrock/sa-east-1/qwen.qwen3-coder-next +bedrock/us-east-1/1-month-commitment/anthropic.claude-instant-v1 +bedrock/us-east-1/1-month-commitment/anthropic.claude-v1 +bedrock/us-east-1/1-month-commitment/anthropic.claude-v2:1 +bedrock/us-east-1/6-month-commitment/anthropic.claude-instant-v1 +bedrock/us-east-1/6-month-commitment/anthropic.claude-v1 +bedrock/us-east-1/6-month-commitment/anthropic.claude-v2:1 +bedrock/us-east-1/anthropic.claude-instant-v1 +bedrock/us-east-1/anthropic.claude-v1 +bedrock/us-east-1/anthropic.claude-v2:1 +bedrock/us-east-1/meta.llama3-70b-instruct-v1:0 +bedrock/us-east-1/meta.llama3-8b-instruct-v1:0 +bedrock/us-east-1/mistral.mistral-7b-instruct-v0:2 +bedrock/us-east-1/mistral.mistral-large-2402-v1:0 +bedrock/us-east-1/mistral.mixtral-8x7b-instruct-v0:1 +bedrock/us-east-1/deepseek.v3.2 +bedrock/us-east-1/minimax.minimax-m2.1 +bedrock/us-east-1/minimax.minimax-m2.5 +bedrock/us-east-1/moonshotai.kimi-k2-thinking +bedrock/us-east-1/moonshotai.kimi-k2.5 +bedrock/us-east-1/qwen.qwen3-coder-next +bedrock/us-east-2/deepseek.v3.2 +bedrock/us-east-2/minimax.minimax-m2.1 +bedrock/us-east-2/minimax.minimax-m2.5 +bedrock/us-east-2/moonshotai.kimi-k2-thinking +bedrock/us-east-2/moonshotai.kimi-k2.5 +bedrock/us-east-2/qwen.qwen3-coder-next +bedrock/us-gov-east-1/amazon.nova-pro-v1:0 +bedrock/us-gov-east-1/amazon.titan-text-express-v1 +bedrock/us-gov-east-1/amazon.titan-text-lite-v1 +bedrock/us-gov-east-1/amazon.titan-text-premier-v1:0 +bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0 +bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0 +bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0 +bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0 +bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0 +bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0 +bedrock/us-gov-west-1/amazon.nova-pro-v1:0 +bedrock/us-gov-west-1/amazon.titan-text-express-v1 +bedrock/us-gov-west-1/amazon.titan-text-lite-v1 +bedrock/us-gov-west-1/amazon.titan-text-premier-v1:0 +bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0 +bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0 +bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0 +bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0 +bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0 +bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0 +bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0 +bedrock/us-west-1/meta.llama3-70b-instruct-v1:0 +bedrock/us-west-1/meta.llama3-8b-instruct-v1:0 +bedrock/us-west-2/1-month-commitment/anthropic.claude-instant-v1 +bedrock/us-west-2/1-month-commitment/anthropic.claude-v1 +bedrock/us-west-2/1-month-commitment/anthropic.claude-v2:1 +bedrock/us-west-2/6-month-commitment/anthropic.claude-instant-v1 +bedrock/us-west-2/6-month-commitment/anthropic.claude-v1 +bedrock/us-west-2/6-month-commitment/anthropic.claude-v2:1 +bedrock/us-west-2/anthropic.claude-instant-v1 +bedrock/us-west-2/anthropic.claude-v1 +bedrock/us-west-2/anthropic.claude-v2:1 +bedrock/us-west-2/mistral.mistral-7b-instruct-v0:2 +bedrock/us-west-2/mistral.mistral-large-2402-v1:0 +bedrock/us-west-2/mistral.mixtral-8x7b-instruct-v0:1 +bedrock/us-west-2/deepseek.v3.2 +bedrock/us-west-2/minimax.minimax-m2.1 +bedrock/us-west-2/minimax.minimax-m2.5 +bedrock/us-west-2/moonshotai.kimi-k2-thinking +bedrock/us-west-2/moonshotai.kimi-k2.5 +bedrock/us-west-2/qwen.qwen3-coder-next +bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0 +claude-sonnet-4-5-20250929-v1:0 +cohere.command-light-text-v14 +cohere.command-r-plus-v1:0 +cohere.command-r-v1:0 +cohere.command-text-v14 +eu.anthropic.claude-3-5-haiku-20241022-v1:0 +eu.anthropic.claude-3-5-sonnet-20240620-v1:0 +eu.anthropic.claude-3-5-sonnet-20241022-v2:0 +eu.anthropic.claude-3-7-sonnet-20250219-v1:0 +eu.anthropic.claude-3-haiku-20240307-v1:0 +eu.anthropic.claude-3-opus-20240229-v1:0 +eu.anthropic.claude-3-sonnet-20240229-v1:0 +eu.meta.llama3-2-1b-instruct-v1:0 +eu.meta.llama3-2-3b-instruct-v1:0 +meta.llama2-13b-chat-v1 +meta.llama2-70b-chat-v1 +meta.llama3-1-405b-instruct-v1:0 +meta.llama3-1-70b-instruct-v1:0 +meta.llama3-1-8b-instruct-v1:0 +meta.llama3-2-11b-instruct-v1:0 +meta.llama3-2-1b-instruct-v1:0 +meta.llama3-2-3b-instruct-v1:0 +meta.llama3-2-90b-instruct-v1:0 +meta.llama3-70b-instruct-v1:0 +meta.llama3-8b-instruct-v1:0 +mistral.mistral-7b-instruct-v0:2 +mistral.mistral-large-2402-v1:0 +mistral.mistral-large-2407-v1:0 +mistral.mistral-small-2402-v1:0 +mistral.mixtral-8x7b-instruct-v0:1 +us.anthropic.claude-3-5-haiku-20241022-v1:0 +us.anthropic.claude-3-5-sonnet-20240620-v1:0 +us.anthropic.claude-3-5-sonnet-20241022-v2:0 +us.anthropic.claude-3-haiku-20240307-v1:0 +us.anthropic.claude-3-opus-20240229-v1:0 +us.anthropic.claude-3-sonnet-20240229-v1:0 +us.meta.llama3-1-405b-instruct-v1:0 +us.meta.llama3-1-70b-instruct-v1:0 +us.meta.llama3-1-8b-instruct-v1:0 +us.meta.llama3-2-11b-instruct-v1:0 +us.meta.llama3-2-1b-instruct-v1:0 +us.meta.llama3-2-3b-instruct-v1:0 +us.meta.llama3-2-90b-instruct-v1:0 +bedrock/us-east-1/zai.glm-5 +bedrock/us-west-2/zai.glm-5 +bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0 +bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0 From 68489f62ff5db52005bfd4084b8c4f46dafebe53 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:57:50 -0700 Subject: [PATCH 097/273] ci: run the enterprise package suite in GitHub Actions (#37798) tests/enterprise is 13 files and 244 tests that only CircleCI runs, and CircleCI gates nothing: it triggers on PR labeled events, none of its jobs are required, and red runs get merged past. So the suite that covers the enterprise package's guardrails, auth and management endpoints has had no say in whether a change lands. Measured on 2026-08-21 with every credential stripped from the environment: 240 passed, 4 skipped, nothing failed. It needs no provider key, so it can be a required shard rather than a scheduled lane, unlike the other CircleCI suites in this group, which each carry a live-API minority. The CircleCI job is removed in the same commit so the suite runs once, not twice. --- .circleci/config.yml | 37 --------------------------------- .github/workflows/test-unit.yml | 8 +++++++ 2 files changed, 8 insertions(+), 37 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index a261b7c78fc..4615a6a5a7e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1154,40 +1154,6 @@ jobs: paths: - search_coverage.xml - search_coverage - litellm_mapped_enterprise_tests: - docker: - - *python312_image - working_directory: ~/project - resource_class: large - - steps: - - checkout - - skip_if_unrelated_changes - - setup_google_dns - - install_uv - - install_rust - - run: - name: Install Dependencies - command: | - uv sync --frozen --all-groups --all-extras --python 3.12 - - setup_litellm_enterprise_pip - - run: - name: Run enterprise tests - command: | - uv run --no-sync python -m prisma generate - mkdir -p test-results - TEST_FILES=$(circleci tests glob "tests/enterprise/**/test_*.py") - echo "$TEST_FILES" | circleci tests run \ - --verbose \ - --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ - -v -x \ - --junitxml=test-results/junit-enterprise.xml \ - --durations=10 \ - -n 4" - no_output_timeout: 15m - # Store test results - - store_test_results: - path: test-results batches_testing: docker: - *python312_image @@ -3041,8 +3007,6 @@ workflows: filters: *main_branches - search_testing: filters: *main_branches - - litellm_mapped_enterprise_tests: - filters: *main_branches - batches_testing: filters: *main_branches - litellm_utils_testing: @@ -3065,7 +3029,6 @@ workflows: - guardrails_testing - ocr_testing - search_testing - - litellm_mapped_enterprise_tests - batches_testing - litellm_utils_testing - pass_through_unit_testing diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 368680d9ec2..1370a05181a 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -221,6 +221,14 @@ jobs: timeout-minutes: 20 job-timeout-minutes: 55 + - shard: enterprise-package + artifact-name: enterprise-package + test-path: "tests/enterprise" + workers: 4 + reruns: 2 + timeout-minutes: 20 + job-timeout-minutes: 55 + - shard: responses-caching-types artifact-name: responses-caching-types test-path: >- From d1f37788492b98b19974d981fff0498a5e090f64 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:58:26 -0700 Subject: [PATCH 098/273] perf(ci): give the two longest unit shards the runner's spare cores (#37804) proxy-endpoints and proxy-infra are the unit tier's critical path at 358s and 325s of pytest, measured on staging 2026-08-21, and both run two xdist workers on a four-vCPU runner. proxy-server already runs four. This is the cheaper half of splitting them: no second job, so no second setup to pay for. --- .github/workflows/test-unit.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1370a05181a..a7c67f2b35d 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -168,7 +168,7 @@ jobs: tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/config_resolvers tests/test_litellm/proxy/utils - workers: 2 + workers: 4 reruns: 2 timeout-minutes: 20 job-timeout-minutes: 60 @@ -195,7 +195,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py - workers: 2 + workers: 4 reruns: 2 timeout-minutes: 20 job-timeout-minutes: 60 From f005afa1460385a218be8ef1fdfa49998bf93523 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 22:59:02 -0700 Subject: [PATCH 099/273] test(exception-mapping): pin the status and error-shape table every provider maps to (#37807) `exception_type` decides the class and status a caller sees for every provider failure, across 190 raise sites, and the tests for it were written one incident at a time. Nothing said what a plain 401 from any given provider should be, so mutating a raise site went unnoticed: swapping the class at each of the 190 in turn, the mapped test file caught 15. Adds two tables asserted end to end through `exception_type`: 25 providers by the 9 upstream statuses, and the three error shapes the router branches on (a full context window, a content policy block, a timeout). The same 190 mutants now fail 97 of them. The tables record today's behavior, uneven where it is uneven. cloudflare, ollama and vllm map no status at all, so every failure reaches the caller as a 500. A full context window is recognised by 15 of the 25, and a content policy block by 11, which bounds where `context_window_fallbacks` and the content policy retry policy can fire. --- .../test_exception_mapping_utils.py | 285 ++++++++++++++++++ 1 file changed, 285 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index cc0a52247a4..599ad016827 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1,5 +1,6 @@ import httpx +import openai import pytest import litellm @@ -759,6 +760,290 @@ def test_azure_404_with_invalid_request_error_type_maps_to_not_found(): assert "Response with id 'resp_abc' not found." in excinfo.value.message +class _UpstreamHTTPError(Exception): + def __init__(self, status_code: int) -> None: + super().__init__("upstream failure") + self.message = "upstream failure" + self.status_code = status_code + self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions") + self.response = httpx.Response( + status_code=status_code, request=self.request, text="upstream failure" + ) + + +UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503) + +OPENAI_SHAPED = { + 400: (litellm.BadRequestError, 400), + 401: (litellm.AuthenticationError, 401), + 403: (litellm.APIError, 403), + 404: (litellm.NotFoundError, 404), + 408: (litellm.Timeout, 408), + 422: (litellm.BadRequestError, 422), + 429: (litellm.RateLimitError, 429), + 500: (litellm.InternalServerError, 500), + 503: (litellm.ServiceUnavailableError, 503), +} + +UPSTREAM_STATUS_DISCARDED = (litellm.APIConnectionError, 500) + +PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS = ("cloudflare", "ollama", "vllm") + +DEVIATIONS_FROM_THE_OPENAI_SHAPE = { + "anthropic": {403: UPSTREAM_STATUS_DISCARDED, 422: UPSTREAM_STATUS_DISCARDED}, + "azure": {500: (litellm.APIError, 500)}, + "bedrock": { + 403: UPSTREAM_STATUS_DISCARDED, + 500: (litellm.ServiceUnavailableError, 503), + }, + "cohere": { + 401: UPSTREAM_STATUS_DISCARDED, + 403: UPSTREAM_STATUS_DISCARDED, + 404: UPSTREAM_STATUS_DISCARDED, + 422: UPSTREAM_STATUS_DISCARDED, + 429: UPSTREAM_STATUS_DISCARDED, + 503: UPSTREAM_STATUS_DISCARDED, + }, + "databricks": { + 403: (litellm.AuthenticationError, 401), + 422: (litellm.BadRequestError, 400), + }, + "gemini": { + 403: (litellm.PermissionDeniedError, 403), + 422: UPSTREAM_STATUS_DISCARDED, + }, + "huggingface": { + 404: (litellm.APIError, 404), + 422: (litellm.APIError, 422), + 500: (litellm.APIError, 500), + }, + "nlp_cloud": { + 403: (litellm.AuthenticationError, 403), + 404: (litellm.APIError, 404), + 408: (litellm.APIError, 408), + 500: (litellm.APIError, 500), + 503: (litellm.APIError, 503), + }, + "openrouter": {500: (litellm.APIError, 500)}, + "replicate": { + 403: (litellm.APIError, 500), + 404: (litellm.APIError, 500), + 422: (litellm.UnprocessableEntityError, 422), + 500: (litellm.ServiceUnavailableError, 503), + 503: (litellm.APIError, 500), + }, + "sagemaker": { + 403: UPSTREAM_STATUS_DISCARDED, + 500: (litellm.ServiceUnavailableError, 503), + }, + "vertex_ai": { + 403: (litellm.PermissionDeniedError, 403), + 422: UPSTREAM_STATUS_DISCARDED, + }, + **{ + provider: dict.fromkeys(UPSTREAM_STATUS_CODES, UPSTREAM_STATUS_DISCARDED) + for provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS + }, +} + +PROVIDERS_WITH_A_HANDLER = ( + "ai21", + "anthropic", + "azure", + "azure_ai", + "bedrock", + "cloudflare", + "cohere", + "databricks", + "deepseek", + "fireworks_ai", + "gemini", + "groq", + "huggingface", + "mistral", + "nlp_cloud", + "ollama", + "openai", + "openrouter", + "perplexity", + "replicate", + "sagemaker", + "together_ai", + "vertex_ai", + "vllm", + "xai", +) + + +def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]: + return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get( + status_code, OPENAI_SHAPED[status_code] + ) + + +@pytest.fixture +def quiet_exception_mapping(monkeypatch): + monkeypatch.setattr(litellm, "suppress_debug_info", True) + + +@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_an_upstream_status_maps_to_one_exception_per_provider( + provider, status_code, quiet_exception_mapping +): + expected_class, expected_status = _expected_for(provider, status_code) + + with pytest.raises(openai.APIError) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamHTTPError(status_code=status_code), + custom_llm_provider=provider, + ) + + assert type(raised.value) is expected_class + assert raised.value.status_code == expected_status + + +@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from( + provider, status_code, quiet_exception_mapping +): + with pytest.raises(openai.APIError) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamHTTPError(status_code=status_code), + custom_llm_provider=provider, + ) + + assert raised.value.llm_provider == provider + assert raised.value.model == "test-model" + + +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_an_already_mapped_litellm_exception_passes_through_untouched( + provider, quiet_exception_mapping +): + already_mapped = litellm.RateLimitError( + message="already mapped", llm_provider=provider, model="test-model" + ) + + returned = exception_type( + model="test-model", + original_exception=already_mapped, + custom_llm_provider=provider, + ) + + assert returned is already_mapped + + +CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens." +CONTENT_POLICY_MESSAGE = ( + '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' +) +TIMEOUT_MESSAGE = "Request timed out." + +PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = ( + "ai21", + "anthropic", + "azure", + "azure_ai", + "databricks", + "deepseek", + "fireworks_ai", + "gemini", + "groq", + "mistral", + "openai", + "perplexity", + "together_ai", + "vertex_ai", + "xai", +) + +PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK = ( + "ai21", + "azure", + "azure_ai", + "deepseek", + "fireworks_ai", + "groq", + "mistral", + "openai", + "perplexity", + "together_ai", + "xai", +) + + +class _UpstreamErrorWithMessage(_UpstreamHTTPError): + def __init__(self, message: str, status_code: int) -> None: + super().__init__(status_code=status_code) + self.args = (message,) + self.message = message + self.response = httpx.Response( + status_code=status_code, request=self.request, text=message + ) + + +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( + provider, quiet_exception_mapping +): + if provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS: + expected_class, expected_status = UPSTREAM_STATUS_DISCARDED + elif provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW: + expected_class, expected_status = litellm.ContextWindowExceededError, 400 + else: + expected_class, expected_status = litellm.BadRequestError, 400 + + with pytest.raises(openai.APIError) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamErrorWithMessage(CONTEXT_WINDOW_MESSAGE, 400), + custom_llm_provider=provider, + ) + + assert type(raised.value) is expected_class + assert raised.value.status_code == expected_status + + +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( + provider, quiet_exception_mapping +): + if provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS: + expected_class, expected_status = UPSTREAM_STATUS_DISCARDED + elif provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK: + expected_class, expected_status = litellm.ContentPolicyViolationError, 400 + else: + expected_class, expected_status = litellm.BadRequestError, 400 + + with pytest.raises(openai.APIError) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamErrorWithMessage(CONTENT_POLICY_MESSAGE, 400), + custom_llm_provider=provider, + ) + + assert type(raised.value) is expected_class + assert raised.value.status_code == expected_status + + +@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) +def test_a_timed_out_request_is_a_timeout_for_every_provider( + provider, quiet_exception_mapping +): + with pytest.raises(litellm.Timeout) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamErrorWithMessage(TIMEOUT_MESSAGE, 408), + custom_llm_provider=provider, + ) + + assert raised.value.status_code == 408 + + def test_bedrock_mantle_400_maps_to_bad_request(): from litellm.llms.base_llm.chat.transformation import BaseLLMException From ee935cec230207a549f2939a6a9ca69023c609d6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 24 Aug 2026 09:46:35 -0700 Subject: [PATCH 100/273] 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 101/273] 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 102/273] 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 103/273] 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 104/273] 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 105/273] 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 106/273] 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 ? (