From a3bab876a981a6f2067c459e62bc2dcf502d7cd6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 20 May 2026 04:20:54 +0530 Subject: [PATCH 01/10] Day 0 support : Gemini 3.5 Flash (#28268) * Add day 0 support for gemini 3.5 flash * Fix pricing * Fix greptile review * Fix failing test * Fix tests * Fix: revert tool removing logic * fix greptile and test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit 3c3d131f017af22e09a1b32d033238a3933085dc) (cherry picked from commit cbf9ffec3050e1f97b1a1e78c7cccd14fa700527) --- .../prompt_templates/factory.py | 40 +++- .../llms/vertex_ai/gemini/transformation.py | 4 +- .../vertex_and_google_ai_studio_gemini.py | 128 ++++++++++++- ...odel_prices_and_context_window_backup.json | 178 ++++++++++++++++++ litellm/types/llms/vertex_ai.py | 21 ++- model_prices_and_context_window.json | 178 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 32 ++++ .../llms/vertex_ai/test_vertex.py | 40 ++-- 8 files changed, 576 insertions(+), 45 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d40ca4e3597..d39b7a5056e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1233,6 +1233,7 @@ def infer_protocol_value( def _gemini_tool_call_invoke_helper( function_call_params: ChatCompletionToolCallFunctionChunk, + tool_call_id: Optional[str] = None, ) -> Optional[VertexFunctionCall]: name = function_call_params.get("name", "") or "" arguments = function_call_params.get("arguments", "") @@ -1248,6 +1249,10 @@ def _gemini_tool_call_invoke_helper( name=name, args=arguments_dict, ) + if tool_call_id: + clean_id = tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + if clean_id: + function_call["id"] = clean_id return function_call @@ -1384,12 +1389,23 @@ def convert_to_gemini_tool_call_invoke( tool_calls = message.get("tool_calls", None) function_call = message.get("function_call", None) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + forward_tool_call_id = bool( + model and VertexGeminiConfig._is_gemini_3_or_newer(model) + ) + if tool_calls is not None: for idx, tool in enumerate(tool_calls): if "function" in tool: gemini_function_call: Optional[VertexFunctionCall] = ( _gemini_tool_call_invoke_helper( - function_call_params=tool["function"] + function_call_params=tool["function"], + tool_call_id=( + tool.get("id") if forward_tool_call_id else None + ), ) ) if gemini_function_call is not None: @@ -1429,10 +1445,6 @@ def convert_to_gemini_tool_call_invoke( thought_signature = provider_fields.get("thought_signature") # If no signature found and model is gemini-3, use dummy signature - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - if ( not thought_signature and model @@ -1462,6 +1474,7 @@ def convert_to_gemini_tool_call_invoke( def convert_to_gemini_tool_call_result( # noqa: PLR0915 message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], last_message_with_tool_calls: Optional[dict], + model: Optional[str] = None, ) -> Union[VertexPartType, List[VertexPartType]]: """ OpenAI message with a tool result looks like: @@ -1602,6 +1615,21 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 ): name = tool.get("function", {}).get("name", "") + # Echo the OpenAI tool_call_id on functionResponse (strip thought-signature suffix). + # Only Gemini 3+ accepts (and returns) an `id` on function_response parts; + # older Gemini models reject the field with a 400. + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + gemini_call_id: Optional[str] = None + if model and VertexGeminiConfig._is_gemini_3_or_newer(model): + raw_tool_call_id = message.get("tool_call_id") + if raw_tool_call_id and isinstance(raw_tool_call_id, str): + stripped_id = raw_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + if stripped_id: + gemini_call_id = stripped_id + if not name: raise Exception( "Missing corresponding tool call for tool response message. Received - message={}, last_message_with_tool_calls={}".format( @@ -1632,6 +1660,8 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 name=name, response=response_data, # type: ignore ) + if gemini_call_id: + _function_response["id"] = gemini_call_id # Create part with function_response, and optionally inline_data for images (Computer Use) _part: VertexPartType = {"function_response": _function_response} diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 9afa5dec465..44d31eb1b84 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -607,7 +607,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 and messages[msg_i]["role"] in tool_call_message_roles ): _part = convert_to_gemini_tool_call_result( - messages[msg_i], last_message_with_tool_calls # type: ignore + messages[msg_i], # type: ignore + last_message_with_tool_calls, # type: ignore + model=model, ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 6278de662f8..101828cbb57 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -280,6 +280,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): - gemini-3-pro-preview - gemini-3-flash - gemini-3-flash-preview (Gemini 3 Flash) + - gemini-3.1-pro-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview + - gemini-3.5-flash - Any future Gemini 3.x models """ # Check for Gemini 3 models @@ -300,6 +302,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): supported_params = [ "temperature", "top_p", + "top_k", "max_tokens", "max_completion_tokens", "stream", @@ -363,6 +366,66 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ return Tools(googleSearch={}) + @staticmethod + def _search_tool_keys() -> set: + return { + VertexToolName.GOOGLE_SEARCH.value, + VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value, + VertexToolName.ENTERPRISE_WEB_SEARCH.value, + VertexToolName.URL_CONTEXT.value, + "google_search", + "google_search_retrieval", + "enterprise_web_search", + "urlContext", + } + + @classmethod + def _drop_search_tools_mixed_with_functions(cls, optional_params: dict) -> None: + """ + Drop search tools from optional_params when mixed with function declarations + and include_server_side_tool_invocations is not enabled. + + Runs after map_openai_params merges tools and web_search_options so both + code paths (single _map_function call vs split tools + web_search_options) + get the same conflict resolution. + """ + if optional_params.get("include_server_side_tool_invocations"): + return + + tools = optional_params.get("tools") + if not isinstance(tools, list) or not tools: + return + + search_tool_keys = cls._search_tool_keys() + has_function_declarations = any( + isinstance(tool, dict) and tool.get("function_declarations") + for tool in tools + ) + if not has_function_declarations: + return + + has_search_tools = any( + isinstance(tool, dict) and any(key in tool for key in search_tool_keys) + for tool in tools + ) + if not has_search_tools: + return + + verbose_logger.warning( + "Vertex AI does not support mixing function declarations with " + "search tools (googleSearch, enterpriseWebSearch, urlContext, " + "googleSearchRetrieval) in the same request. Dropping search " + "tools and keeping function declarations. To use search tools, " + "send a request without function calling tools." + ) + optional_params["tools"] = [ + tool + for tool in tools + if not ( + isinstance(tool, dict) and any(key in tool for key in search_tool_keys) + ) + ] + def _map_service_tier_param(self, value: str, optional_params: dict) -> None: """ Map OpenAI service_tier (string) to Gemini serviceTier. @@ -884,9 +947,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): GeminiThinkingConfig with thinkingLevel and includeThoughts """ # Check if this is gemini-3-flash which supports MINIMAL thinking level - # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, etc. + # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, + # gemini-3.5-flash, and any future 3.x-flash variants. is_gemini3flash = model and ( - "gemini-3-flash" in model.lower() or "gemini-3.1-flash" in model.lower() + "flash" in model.lower() and "gemini-3" in model.lower() ) is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": @@ -982,8 +1046,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: is_gemini3flash = ( - "gemini-3-flash-preview" in model.lower() - or "gemini-3-flash" in model.lower() + "gemini-3" in model.lower() and "flash" in model.lower() ) params["thinkingLevel"] = ( "minimal" if is_gemini3flash else "low" @@ -1077,6 +1140,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model: str, drop_params: bool, ) -> Dict: + gemini_sampling_params_warned: bool = False for param, value in non_default_params.items(): if param == "temperature": if VertexGeminiConfig._is_gemini_3_or_newer(model): @@ -1086,9 +1150,41 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "can cause infinite loops, degraded reasoning performance, and failure on complex tasks. " "Strongly recommended to use temperature = 1.0 (default)." ) + if not gemini_sampling_params_warned: + verbose_logger.warning( + "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to " + f"function for Gemini 3+ ({model}) but are planned for removal in a " + "future release. Move sampling guidance into the `system` " + "instructions instead." + ) + gemini_sampling_params_warned = True optional_params["temperature"] = value elif param == "top_p": + if ( + VertexGeminiConfig._is_gemini_3_or_newer(model) + and not gemini_sampling_params_warned + ): + verbose_logger.warning( + "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to " + f"function for Gemini 3+ ({model}) but are planned for removal in a " + "future release. Move sampling guidance into the `system` " + "instructions instead." + ) + gemini_sampling_params_warned = True optional_params["top_p"] = value + elif param == "top_k": + if ( + VertexGeminiConfig._is_gemini_3_or_newer(model) + and not gemini_sampling_params_warned + ): + verbose_logger.warning( + "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to " + f"function for Gemini 3+ ({model}) but are planned for removal in a " + "future release. Move sampling guidance into the `system` " + "instructions instead." + ) + gemini_sampling_params_warned = True + optional_params["top_k"] = value elif ( param == "stream" and value is True ): # sending stream = False, can cause it to get passed unchecked and raise issues @@ -1139,11 +1235,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": - if value is False and not ( - drop_params or litellm.drop_params - ): # if drop params is True, then we should just ignore this - self.validate_parallel_tool_calls(value, non_default_params) - else: + tools_list = non_default_params.get( + "tools", non_default_params.get("functions") + ) + num_tools = len(tools_list) if isinstance(tools_list, list) else 0 + # Gemini does not support parallel_tool_calls=False with multiple + # tools. Drop the param instead of failing — Responses API clients + # often send parallel_tool_calls=false by default. + if not (value is False and num_tools > 1): optional_params["parallel_tool_calls"] = value elif param == "seed": optional_params["seed"] = value @@ -1216,6 +1315,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if "temperature" not in optional_params: optional_params["temperature"] = 1.0 + self._drop_search_tools_mixed_with_functions(optional_params) + return optional_params def get_mapped_special_auth_params(self) -> dict: @@ -1588,6 +1689,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): } # Extract thought signature if present thought_signature = part.get("thoughtSignature") + # Gemini 3.5+ returns a stable `id` per function call to enable + # strict response matching. Preserve it as the OpenAI + # tool_call_id so it can be echoed back unchanged. + gemini_call_id = part["functionCall"].get("id") if is_function_call is True: function_dict: Dict[str, Any] = dict(_function_chunk) @@ -1605,6 +1710,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "function": _function_chunk, "index": cumulative_tool_call_idx, } + # Gemini 3.5+ returns a stable native `id`; prefer it over + # the synthetic call_ so the same value can be echoed + # back on the matching `functionResponse`. + if gemini_call_id: + _tool_response_chunk["id"] = gemini_call_id # Embed thought signature in ID for OpenAI client compatibility if thought_signature: _tool_response_chunk["provider_specific_fields"] = { # type: ignore diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fb5bfa6cf4e..6d7c3eeb0b8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15611,6 +15611,64 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -16988,6 +17046,67 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "rpm": 2000, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -17173,6 +17292,65 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 87bf11a9026..b357b64156d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -14,13 +14,19 @@ from litellm.types.llms.openai import EmbeddingInput GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]] -class FunctionResponse(TypedDict): - name: str +class FunctionResponse(TypedDict, total=False): + # `id` correlates this response with the originating `functionCall` part. + # Required by Gemini 3.5+ for strict function-calling response matching. + id: str + name: Required[str] response: Optional[dict] -class FunctionCall(TypedDict): - name: str +class FunctionCall(TypedDict, total=False): + # `id` is returned by Gemini 3.5+ to correlate the corresponding + # `functionResponse`. Older Gemini models omit this field. + id: str + name: Required[str] args: Optional[dict] @@ -45,8 +51,11 @@ class PartType(TypedDict, total=False): media_resolution: Literal["low", "medium", "high"] -class HttpxFunctionCall(TypedDict): - name: str +class HttpxFunctionCall(TypedDict, total=False): + # `id` is returned by Gemini 3.5+ to correlate the corresponding + # `functionResponse`. Older Gemini models omit this field. + id: str + name: Required[str] args: dict diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 94f0f1e78d3..e7a03bb0984 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15645,6 +15645,64 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -17022,6 +17080,67 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "rpm": 2000, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -17207,6 +17326,65 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "output_cost_per_token_priority": 1.62e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "supports_service_tier": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 353d19b0198..59c27f56bb4 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2959,6 +2959,38 @@ def test_vertex_ai_gemini3_tool_combination_no_drop(): assert len(tools) == 3 +def test_vertex_ai_mixed_tools_and_web_search_options_drops_search(): + """ + When function tools and web_search_options are sent separately (Codex-style), + search tools are dropped unless include_server_side_tool_invocations is set. + """ + v = VertexGeminiConfig() + optional_params: dict = {} + non_default_params = { + "tools": [ + { + "type": "function", + "function": {"name": "exec_command", "description": "Run a command"}, + } + ], + "web_search_options": {}, + } + + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gemini-3.5-flash", + drop_params=True, + ) + + assert not result.get("include_server_side_tool_invocations") + tool_keys = set() + for tool in result.get("tools", []): + tool_keys.update(tool.keys()) + assert "function_declarations" in tool_keys + assert "googleSearch" not in tool_keys + + def test_vertex_ai_openai_web_search_tool_transformation(): """ Test that OpenAI-style web_search and web_search_preview tools are transformed to googleSearch. diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 2e9629f95de..5977fab68ef 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -1490,39 +1490,31 @@ def test_vertex_parallel_tool_calls_true(): assert "tools" in optional_params -def test_vertex_parallel_tool_calls_false_multiple_tools_error(): +def test_vertex_parallel_tool_calls_false_multiple_tools_dropped(): """ - Test that parallel_tool_calls = False with multiple tools raises UnsupportedParamsError - when drop_params is False. + parallel_tool_calls=False with multiple tools is dropped for Gemini + (unsupported upstream). Request should succeed without the param. """ tools = [ {"type": "function", "function": {"name": "get_weather"}}, {"type": "function", "function": {"name": "get_time"}}, ] - with pytest.raises(litellm.utils.UnsupportedParamsError) as excinfo: - get_optional_params( - model="gemini-1.5-pro", - custom_llm_provider="vertex_ai", - tools=tools, - parallel_tool_calls=False, - ) - assert ( - "`parallel_tool_calls=False` is not supported by Gemini when multiple tools are" - in str(excinfo.value) + optional_params = get_optional_params( + model="gemini-1.5-pro", + custom_llm_provider="vertex_ai", + tools=tools, + parallel_tool_calls=False, ) + assert "parallel_tool_calls" not in optional_params + assert "tools" in optional_params - # works when specified as "functions" - with pytest.raises(litellm.utils.UnsupportedParamsError) as excinfo: - get_optional_params( - model="gemini-1.5-pro", - custom_llm_provider="vertex_ai", - functions=tools, - parallel_tool_calls=False, - ) - assert ( - "`parallel_tool_calls=False` is not supported by Gemini when multiple tools are" - in str(excinfo.value) + optional_params = get_optional_params( + model="gemini-1.5-pro", + custom_llm_provider="vertex_ai", + functions=tools, + parallel_tool_calls=False, ) + assert "parallel_tool_calls" not in optional_params def test_vertex_parallel_tool_calls_false_single_tool(): From ecf1e6ee9c7145cfd05e0655807e8fd19166ce49 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Wed, 20 May 2026 20:57:08 +0300 Subject: [PATCH 02/10] fix(spend_counter): seed Redis counter via SET NX to prevent cross-pod double-seed (#27854) * fix(spend_counter): seed Redis counter via SET NX to prevent cross-pod double-seed Symptom ------- Customers on multi-pod deployments see team `spend` jump to ~2x (or N x the pod count) shortly after a Redis cache miss / TTL expiry, triggering spurious "Budget Crossed" alerts and blocked requests until the value is manually reset. Root cause ---------- `SpendCounterReseed.coalesced` warmed the primary spend counter by calling `redis.async_increment(key, value=db_spend, refresh_ttl=True)`, which lowers to Redis `INCRBYFLOAT`. That is additive, not idempotent. The per-counter `asyncio.Lock` only coalesces seeders inside one process. With N pods sharing one Redis, on a cold key (cold start, TTL expiry, manual delete) every pod independently passes its lock + Redis re-check, reads the same `db_spend`, and issues `INCRBYFLOAT db_spend`. Final value: N x db_spend. Fix --- Use `redis.async_set_cache(key, value=db_spend, nx=True)` for the seed. SET NX is atomic across pods: exactly one writer initializes the key; losers read the winner's value via `async_get_cache`. This is the same idiom already used by `coalesced_window` in the same file, so the two seed paths are now consistent. Per-request deltas continue to use `INCRBYFLOAT` (correct - additive behaviour is what we want for increments, not for initial seed). Verification ------------ Live two-process repro against the same Postgres + Redis (DB spend = 506): Unpatched: 4/4 runs -> Redis counter = ~1012 (~2 x db_spend) Patched: 12/12 runs -> Redis counter = ~506 Unit tests (`test_proxy_server.py`): - New `test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed` patches `_get_lock` to return a fresh lock per caller (otherwise the per-process lock masks the race), races two `coalesced` calls, and asserts final = 506 with exactly one of two SET NX attempts winning. - 4 existing tests updated for the new seed contract (SET NX for the seed, INCRBYFLOAT only for the per-request delta). - Full `spend_counter or reseed or budget` slice: 22 passed. Co-authored-by: Cursor * test(spend_counter): make SET NX mock atomic so loser branch is exercised Greptile flagged that `redis_set_cache` in test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed placed `await asyncio.sleep(0)` AFTER the NX membership check. Both concurrent tasks observed an empty `redis_store`, passed the guard, and both returned True - so the loser branch (else: read back winner's value) was never exercised. Fix the mock to model real atomic Redis SET NX: - Yield BEFORE the membership check so two concurrent callers interleave the way real SET NX does (first to resume runs check + write atomically and wins; second resumes after the key exists and loses). - Track set_cache return values; assert sorted([loser, winner]) so we know exactly one task wins and one loses. - Track async_get_cache calls that happen AFTER at least one SET NX has completed; assert at least one such read - that is the loser-path fallback (`current_value = float(cached)` when seeded is False). Verified by temporarily reverting the mock to the old order: the test now fails with `expected exactly one SET NX winner and one loser, got [True, True]`, exactly the failure mode Greptile described. No production code change. Co-authored-by: Cursor * test(spend_counter): mock async_set_cache to populate redis_store in concurrent read+write test `test_concurrent_read_and_write_paths_share_one_db_query` mocks `async_increment` to populate the in-memory `redis_store`, but did not mock `async_set_cache`. After the SET-NX seed change in `coalesced()`, the seed step writes via `async_set_cache(nx=True)` (default AsyncMock, no `redis_store` write), so the simulated Redis stays empty after the first reseed. The second `get_current_spend` then sees a clean Redis miss, re-enters the DB read path, and the test fails with `expected 1 DB query, got 2`. Fix: add a `redis_set_cache` side_effect that updates `redis_store` on `nx=True` (and rejects when the key already exists), matching the pattern used by the four sibling tests fixed in this branch's first commit. Pre-existing assertions are unchanged. Full `tests/test_litellm/proxy/test_proxy_server.py`: 158 passed. Co-authored-by: Cursor --------- Co-authored-by: Cursor (cherry picked from commit 0fb710400f80088fdcc1e382d3efa90a7ec895ea) (cherry picked from commit c621f58fffbb304960938cdeff009db3ded456f4) --- litellm/proxy/db/spend_counter_reseed.py | 27 ++- tests/test_litellm/proxy/test_proxy_server.py | 167 ++++++++++++++++-- 2 files changed, 176 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 19ec6699390..e7c5fa3f72c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -178,15 +178,28 @@ class SpendCounterReseed: if db_spend is None: return None # Warm even when 0 so subsequent reads hit cache, not DB. + # + # Seed via SET NX (cross-pod safe): only one pod initializes the + # Redis key with db_spend; concurrent seeders read the winner's + # value. INCRBYFLOAT-of-db_spend from N pods would multiply the + # counter (N x db_spend) and trigger spurious budget alerts. + current_value: float = float(db_spend) try: if spend_counter_cache.redis_cache is not None: - current_value = ( - await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=db_spend, - refresh_ttl=True, - ) + seeded = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=db_spend, + nx=True, ) + if seeded: + current_value = float(db_spend) + else: + cached = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + current_value = ( + float(cached) if cached is not None else float(db_spend) + ) spend_counter_cache.in_memory_cache.set_cache( key=counter_key, value=current_value, @@ -202,7 +215,7 @@ class SpendCounterReseed: ) if require_cache_warm: raise - return db_spend + return current_value @staticmethod async def window_from_spend_logs( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 73d53631622..f56ce716fcd 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5648,6 +5648,7 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing + fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis # Prisma returns spend=42.0 (authoritative) while the stale cached @@ -5684,16 +5685,131 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-9"} ) - # Two increments keyed on the counter: seed ($42) then request ($1.50). + # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. + # Only the per-request delta (1.5) goes through INCRBYFLOAT. + fake_redis.async_set_cache.assert_awaited_once_with( + key="spend:team:team-9", value=42.0, nx=True + ) writes = [(c["key"], c["value"]) for c in recorded_increments] - assert ("spend:team:team-9", 42.0) in writes - assert ("spend:team:team-9", 1.5) in writes + assert writes == [("spend:team:team-9", 1.5)] finally: ps.user_api_key_cache = orig_user ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma +@pytest.mark.asyncio +async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed(): + """Two pods both observing a missing Redis counter must not both + INCRBYFLOAT the full DB spend. SpendCounterReseed.coalesced uses SET NX + so the loser reads the winner's value; final Redis = db_spend, not + 2 * db_spend. + + The per-counter asyncio.Lock is per-process, so it does NOT coordinate + across pods. We simulate two pods by patching _get_lock to return a + fresh lock per call (each "pod" has its own lock registry in real life). + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-concurrent-seed" + redis_store: dict = {} + db_read_count = 0 + set_results: list = [] + get_after_set_count = 0 + set_completed_count = 0 + + async def redis_set_cache(key, value, nx=False, **_): + # Yield BEFORE the membership check so two concurrent callers + # interleave the way real atomic Redis SET NX does: the first + # to resume runs check + write atomically and wins; the second + # resumes after the key exists and loses. Yielding *after* the + # check would let both callers pass the empty-store check before + # either writes, so neither would ever lose. + await asyncio.sleep(0) + if nx and key in redis_store: + set_results.append(False) + return False + redis_store[key] = float(value) + set_results.append(True) + nonlocal set_completed_count + set_completed_count += 1 + return True + + async def redis_get_cache(key): + # Track reads that happen after at least one SET NX has completed + # — those are the loser-path fallback reads we want to verify. + if set_completed_count > 0: + nonlocal get_after_set_count + get_after_set_count += 1 + return redis_store.get(key) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def slow_find_unique(**_): + nonlocal db_read_count + db_read_count += 1 + # Both pods read DB before either's SET NX lands. + await asyncio.sleep(0) + row = MagicMock() + row.spend = 506.0 + return row + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=slow_find_unique + ) + + pod_a = DualCache() + pod_a.redis_cache = fake_redis + pod_b = DualCache() + pod_b.redis_cache = fake_redis + + # Each "pod" has its own per-process lock registry. Patch _get_lock to + # always return a fresh lock so the two coalesced calls do not serialize + # via one in-process lock (which is what would happen across pods). + async def fresh_lock(_counter_key): + return asyncio.Lock() + + with patch.object(SpendCounterReseed, "_get_lock", side_effect=fresh_lock): + results = await asyncio.gather( + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_a, + counter_key=counter_key, + ), + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_b, + counter_key=counter_key, + ), + ) + + assert all(r == 506.0 for r in results), results + assert redis_store[counter_key] == pytest.approx(506.0), redis_store + # Both pods read the DB and both attempted SET NX; exactly one wrote + # (winner) and one was rejected (loser). + assert db_read_count == 2 + assert fake_redis.async_set_cache.await_count == 2 + nx_writes = [ + call + for call in fake_redis.async_set_cache.await_args_list + if call.kwargs.get("nx") is True + ] + assert len(nx_writes) == 2 + assert sorted(set_results) == [False, True], ( + f"expected exactly one SET NX winner and one loser, got {set_results}" + ) + # Loser path executed: after the winner's SET NX returned True, the + # losing coalesced() call falls back to async_get_cache to read the + # winner's value rather than re-seeding. + assert get_after_set_count >= 1, ( + "loser branch (else: read back winner's value) was never exercised" + ) + + @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): """User and org counters reseed from their own DB tables. @@ -5817,9 +5933,16 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -5847,6 +5970,7 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-stale-local"} ) + # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. assert redis_store[counter_key] == pytest.approx(43.5) assert counter_cache.in_memory_cache.get_cache( key=counter_key @@ -6237,14 +6361,14 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): from litellm.proxy.proxy_server import get_current_spend counter_cache = DualCache() - recorded_warms: list = [] + recorded_seeds: list = [] - async def record_increment(key, value, ttl=None, **kwargs): - recorded_warms.append({"key": key, "value": value}) - return value + async def record_set_cache(key, value, nx=False, **kwargs): + recorded_seeds.append({"key": key, "value": value, "nx": nx}) + return True fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) fake_redis.async_get_cache = AsyncMock(return_value=None) counter_cache.redis_cache = fake_redis @@ -6269,9 +6393,9 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): f"expected DB reseed to return 362.0, got {spend} " f"(fallback would have returned 30.0 and caused bypass)" ) - # Counter warmed so subsequent reads are fast - assert ("spend:team_member:user-1:team-1", 362.0) in [ - (w["key"], w["value"]) for w in recorded_warms + # Counter warmed via SET NX so subsequent reads are fast. + assert ("spend:team_member:user-1:team-1", 362.0, True) in [ + (s["key"], s["value"], s["nx"]) for s in recorded_seeds ] assert counter_cache.in_memory_cache.get_cache( key="spend:team_member:user-1:team-1" @@ -6348,8 +6472,15 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -6456,9 +6587,16 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -6561,9 +6699,16 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_call_count = 0 From 957c1f8a1a29e4e16b17079faa8c26ffb9eddcf1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 21 May 2026 03:22:50 +0530 Subject: [PATCH 03/10] fix(vertex_ai): omit function_call id on Vertex Gemini 3.5+ tool turns (#28324) * fix(vertex_ai): omit function_call id on Vertex Gemini 3.5+ tool turns Vertex AI rejects `id` on function_call/function_response parts; only Google AI Studio accepts it for Gemini 3.5+ strict tool matching. Co-authored-by: Cursor * Update litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(vertex_ai): forward custom_llm_provider in context caching Pass custom_llm_provider through to _gemini_convert_messages_with_history in the context caching path so Gemini 3.5+ tool-call `id` forwarding behaves consistently between cached and non-cached completions on Google AI Studio. Co-authored-by: Claude --------- Co-authored-by: Cursor Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Claude (cherry picked from commit fecf212d7001cd409139249a69e6cbaa1f00afd6) (cherry picked from commit 75c72c51e28eaeb5de977bfd4f33e1230d862fc9) --- .../prompt_templates/factory.py | 15 +- litellm/llms/gemini/chat/transformation.py | 6 +- .../context_caching/transformation.py | 4 +- .../llms/vertex_ai/gemini/transformation.py | 6 +- .../vertex_and_google_ai_studio_gemini.py | 20 ++- litellm/types/llms/vertex_ai.py | 10 +- ...test_vertex_and_google_ai_studio_gemini.py | 148 +++++++++++++++++- 7 files changed, 188 insertions(+), 21 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d39b7a5056e..0663d4167d3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1344,6 +1344,7 @@ def _get_dummy_thought_signature() -> str: def convert_to_gemini_tool_call_invoke( message: ChatCompletionAssistantMessage, model: Optional[str] = None, + custom_llm_provider: Optional[str] = None, ) -> List[VertexPartType]: """ OpenAI tool invokes: @@ -1394,7 +1395,10 @@ def convert_to_gemini_tool_call_invoke( ) forward_tool_call_id = bool( - model and VertexGeminiConfig._is_gemini_3_or_newer(model) + model + and VertexGeminiConfig._forward_gemini_function_call_id( + model, custom_llm_provider + ) ) if tool_calls is not None: @@ -1475,6 +1479,7 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], last_message_with_tool_calls: Optional[dict], model: Optional[str] = None, + custom_llm_provider: Optional[str] = None, ) -> Union[VertexPartType, List[VertexPartType]]: """ OpenAI message with a tool result looks like: @@ -1616,14 +1621,16 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 name = tool.get("function", {}).get("name", "") # Echo the OpenAI tool_call_id on functionResponse (strip thought-signature suffix). - # Only Gemini 3+ accepts (and returns) an `id` on function_response parts; - # older Gemini models reject the field with a 400. + # Only Google AI Studio Gemini 3+ accepts `id` on function_response parts. + # Vertex AI and older Gemini models reject the field with HTTP 400. from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) gemini_call_id: Optional[str] = None - if model and VertexGeminiConfig._is_gemini_3_or_newer(model): + if model and VertexGeminiConfig._forward_gemini_function_call_id( + model, custom_llm_provider + ): raw_tool_call_id = message.get("tool_call_id") if raw_tool_call_id and isinstance(raw_tool_call_id, str): stripped_id = raw_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 72569e5c6cd..10ad373152e 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -151,4 +151,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): except Exception: # If conversion fails, leave as is and let the API handle it pass - return _gemini_convert_messages_with_history(messages=messages, model=model) + return _gemini_convert_messages_with_history( + messages=messages, + model=model, + custom_llm_provider="gemini", + ) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index 950edbeb478..3d532113ba0 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -174,7 +174,9 @@ def transform_openai_messages_to_gemini_context_caching( ) transformed_messages = _gemini_convert_messages_with_history( - messages=new_messages, model=model + messages=new_messages, + model=model, + custom_llm_provider=custom_llm_provider, ) model_name = "models/{}".format(model) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 44d31eb1b84..bec992b2bfb 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -311,6 +311,7 @@ def check_if_part_exists_in_parts( def _gemini_convert_messages_with_history( # noqa: PLR0915 messages: List[AllMessageValues], model: Optional[str] = None, + custom_llm_provider: Optional[str] = None, ) -> List[ContentType]: """ Converts given messages from OpenAI format to Gemini format @@ -548,7 +549,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 or assistant_msg.get("function_call") is not None ): # support assistant tool invoke conversion gemini_tool_call_parts = convert_to_gemini_tool_call_invoke( - assistant_msg, model=model + assistant_msg, + model=model, + custom_llm_provider=custom_llm_provider, ) ## check if gemini_tool_call already exists in assistant_content for gemini_tool_call_part in gemini_tool_call_parts: @@ -610,6 +613,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 messages[msg_i], # type: ignore last_message_with_tool_calls, # type: ignore model=model, + custom_llm_provider=custom_llm_provider, ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 101828cbb57..0822b364a60 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -289,6 +289,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return True return False + @staticmethod + def _forward_gemini_function_call_id( + model: str, custom_llm_provider: Optional[str] = None + ) -> bool: + """ + Whether to include `id` on function_call / function_response parts. + + Gemini 3+ on Google AI Studio accepts (and returns) `id` for strict + tool-call matching. Vertex AI rejects the field with HTTP 400. + """ + if custom_llm_provider != "gemini": + return False + return VertexGeminiConfig._is_gemini_3_or_newer(model) + def _supports_penalty_parameters(self, model: str) -> bool: # Gemini 3 models do not support penalty parameters if VertexGeminiConfig._is_gemini_3_or_newer(model): @@ -2645,7 +2659,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _transform_messages( self, messages: List[AllMessageValues], model: Optional[str] = None ) -> List[ContentType]: - return _gemini_convert_messages_with_history(messages=messages, model=model) + return _gemini_convert_messages_with_history( + messages=messages, + model=model, + custom_llm_provider="vertex_ai", + ) def get_error_class( self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index b357b64156d..a1d53978761 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -16,15 +16,15 @@ GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]] class FunctionResponse(TypedDict, total=False): # `id` correlates this response with the originating `functionCall` part. - # Required by Gemini 3.5+ for strict function-calling response matching. + # Supported on Google AI Studio Gemini 3.5+; Vertex AI rejects this field. id: str name: Required[str] response: Optional[dict] class FunctionCall(TypedDict, total=False): - # `id` is returned by Gemini 3.5+ to correlate the corresponding - # `functionResponse`. Older Gemini models omit this field. + # `id` correlates the corresponding `functionResponse` on Google AI Studio + # Gemini 3.5+. Vertex AI and older Gemini models omit/reject this field. id: str name: Required[str] args: Optional[dict] @@ -52,8 +52,8 @@ class PartType(TypedDict, total=False): class HttpxFunctionCall(TypedDict, total=False): - # `id` is returned by Gemini 3.5+ to correlate the corresponding - # `functionResponse`. Older Gemini models omit this field. + # `id` correlates the corresponding `functionResponse` on Google AI Studio + # Gemini 3.5+. Vertex AI and older Gemini models omit/reject this field. id: str name: Required[str] args: dict diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 59c27f56bb4..40ab4d1a893 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2097,6 +2097,125 @@ def test_is_gemini_3_or_newer(): assert VertexGeminiConfig._is_gemini_3_or_newer("") == False +def test_forward_gemini_function_call_id_vertex_vs_google_ai_studio(): + """Vertex AI rejects `id` on function_call/function_response; Google AI Studio accepts it on Gemini 3.5+.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + model = "gemini-3.5-flash" + assert ( + VertexGeminiConfig._forward_gemini_function_call_id(model, "vertex_ai") is False + ) + assert ( + VertexGeminiConfig._forward_gemini_function_call_id(model, "vertex_ai_beta") + is False + ) + assert VertexGeminiConfig._forward_gemini_function_call_id(model, "gemini") is True + assert VertexGeminiConfig._forward_gemini_function_call_id(model, None) is False + assert ( + VertexGeminiConfig._forward_gemini_function_call_id( + "gemini-2.5-flash", "gemini" + ) + is False + ) + + +def test_vertex_ai_gemini_35_tool_calls_omit_function_call_id(): + """Regression: Vertex must not send OpenAI tool_call id inside Gemini function_call parts.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Explore this directory"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_50e7e0fe0989464a89f188eda443", + "type": "function", + "function": { + "name": "read", + "arguments": '{"filePath": "/tmp"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_50e7e0fe0989464a89f188eda443", + "content": "ok", + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, + model="gemini-3.5-flash", + custom_llm_provider="vertex_ai", + ) + + for content in contents: + for part in content.get("parts", []): + fc = part.get("function_call") + if fc is not None: + assert "id" not in fc, f"Vertex payload must not include id: {fc}" + fr = part.get("function_response") + if fr is not None: + assert "id" not in fr, f"Vertex payload must not include id: {fr}" + + +def test_google_ai_studio_gemini_35_tool_calls_include_function_call_id(): + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + tool_call_id = "call_50e7e0fe0989464a89f188eda443" + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "read", + "arguments": '{"filePath": "/tmp"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": "ok", + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, + model="gemini-3.5-flash", + custom_llm_provider="gemini", + ) + + function_call_ids = [] + function_response_ids = [] + for content in contents: + for part in content.get("parts", []): + fc = part.get("function_call") + if fc is not None: + function_call_ids.append(fc.get("id")) + fr = part.get("function_response") + if fr is not None: + function_response_ids.append(fr.get("id")) + + assert function_call_ids == [tool_call_id] + assert function_response_ids == [tool_call_id] + + def test_reasoning_effort_maps_to_thinking_level_gemini_3(): """Test that reasoning_effort maps to thinking_level AND includeThoughts for Gemini 3+ models""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -3531,7 +3650,12 @@ def test_video_metadata_supported_for_all_gemini_models(): } ] - for model in ["gemini-1.5-pro", "gemini-2.5-flash", "gemini-2.5-pro", "gemini-3-pro-preview"]: + for model in [ + "gemini-1.5-pro", + "gemini-2.5-flash", + "gemini-2.5-pro", + "gemini-3-pro-preview", + ]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = None @@ -3541,19 +3665,25 @@ def test_video_metadata_supported_for_all_gemini_models(): break assert file_part is not None, f"{model}: file part should exist" - assert "video_metadata" in file_part, f"{model}: video_metadata should be present" + assert ( + "video_metadata" in file_part + ), f"{model}: video_metadata should be present" assert file_part["video_metadata"]["fps"] == 5, f"{model}: fps should be 5" # Per-part media_resolution is Gemini 3+ only; 2.x uses generation_config global for model in ["gemini-3-pro-preview"]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = next(p for p in contents[0]["parts"] if "file_data" in p) - assert "media_resolution" in file_part, f"{model}: media_resolution should be present" + assert ( + "media_resolution" in file_part + ), f"{model}: media_resolution should be present" for model in ["gemini-1.5-pro", "gemini-2.5-flash", "gemini-2.5-pro"]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = next(p for p in contents[0]["parts"] if "file_data" in p) - assert "media_resolution" not in file_part, f"{model}: per-part media_resolution should not be set" + assert ( + "media_resolution" not in file_part + ), f"{model}: per-part media_resolution should not be set" def test_chunk_parser_handles_prompt_feedback_block(): @@ -4186,8 +4316,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): # DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782 assert result.prompt_tokens_details is not None - assert result.prompt_tokens_details.text_tokens == 782, \ - "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + assert ( + result.prompt_tokens_details.text_tokens == 782 + ), "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" # Verify completion token details assert result.completion_tokens_details is not None @@ -4222,8 +4353,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_cached(): # DOCUMENT cached tokens map to cached_text_tokens, so: # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382 - assert result.prompt_tokens_details.text_tokens == 382, \ - "text_tokens should be (8 + 774) - 400 cached = 382" + assert ( + result.prompt_tokens_details.text_tokens == 382 + ), "text_tokens should be (8 + 774) - 400 cached = 382" assert result.prompt_tokens_details.cached_tokens == 400 From 1f0fdaa8ffb91951a8eda4f557dc12a439214df4 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 29 May 2026 22:23:24 -0700 Subject: [PATCH 04/10] [internal copy of #29089] fix: duplicate claude code traces (#29311) (cherry picked from commit fcdf0231d008b8e701695540292b0d624fbbe87c) --- litellm/litellm_core_utils/litellm_logging.py | 106 +++++++-- .../litellm_core_utils/streaming_handler.py | 20 +- litellm/proxy/common_request_processing.py | 49 +--- .../streaming_handler.py | 22 +- .../pass_through_endpoints/success_handler.py | 24 +- .../test_unit_test_streaming.py | 117 ++++++++++ .../test_proxy_reject_logging.py | 19 +- .../test_litellm_logging.py | 212 +++++++++++++++++- .../test_streaming_handler.py | 16 +- .../test_deferred_guardrail_logging.py | 103 ++++++++- 10 files changed, 566 insertions(+), 122 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 876f1b167db..652d45753c1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1613,6 +1613,90 @@ class Logging(LiteLLMLoggingBaseClass): ) -> Optional[float]: return self._response_cost_calculator(result=result, cache_hit=cache_hit) + @staticmethod + def _is_sync_litellm_request(litellm_params: dict) -> bool: + """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.).""" + return ( + litellm_params.get(CallTypes.acompletion.value, False) is not True + and litellm_params.get(CallTypes.aresponses.value, False) is not True + and litellm_params.get(CallTypes.aembedding.value, False) is not True + and litellm_params.get(CallTypes.aimage_generation.value, False) is not True + and litellm_params.get(CallTypes.atranscription.value, False) is not True + ) + + def _is_assembled_stream_success(self, result=None) -> bool: + """Final assembled stream export (not a per-chunk success call). + + Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the + final assembled response is any other non-``None`` value (typically a + ``ModelResponse``). Treating a chunk as the assembled response would + prematurely set the ``has_dispatched_final_stream_success`` dedup + guard and silently suppress the real final stream log. + """ + if self.stream is not True: + return False + if result is not None and not isinstance(result, ModelResponseStream): + return True + return ( + "async_complete_streaming_response" in self.model_call_details + or self.model_call_details.get("complete_streaming_response") is not None + ) + + async def dispatch_success_handlers( + self, + result=None, + start_time=None, + end_time=None, + cache_hit=None, + prefer_async_handlers: bool = False, + **kwargs, + ) -> None: + """Route success logging to async and/or sync handlers for this request. + + ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g. + ``async for`` on a stream from ``completion()``). Legacy string callbacks + still run via ``executor.submit(success_handler)`` when configured. + """ + from litellm.litellm_core_utils.thread_pool_executor import executor + + if self._is_assembled_stream_success(result): + if self.model_call_details.get("has_dispatched_final_stream_success"): + return + self.model_call_details["has_dispatched_final_stream_success"] = True + + litellm_params = self.model_call_details.get("litellm_params", {}) or {} + sync_sdk = self._is_sync_litellm_request(litellm_params) + passthrough = self.call_type == CallTypes.pass_through.value + if sync_sdk and not prefer_async_handlers and not passthrough: + self.success_handler( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + return + + await self.async_success_handler( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + + if not self._should_run_sync_callbacks_for_async_calls(): + return + + executor.submit( + self.success_handler, + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + def should_run_logging( self, event_type: Literal[ @@ -2032,13 +2116,7 @@ class Logging(LiteLLMLoggingBaseClass): standard_logging_object=kwargs.get("standard_logging_object", None), ) litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) + is_sync_request = self._is_sync_litellm_request(litellm_params) try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: Optional[ @@ -2494,9 +2572,11 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose( "Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit) ) - if not self.should_run_logging( + if not self._is_assembled_stream_success( + result + ) and not self.should_run_logging( event_type="async_success" - ): # prevent double logging + ): # prevent double logging (non-streaming) return ## CALCULATE COST FOR BATCH JOBS @@ -2946,13 +3026,7 @@ class Logging(LiteLLMLoggingBaseClass): ): # prevent double logging return litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) + is_sync_request = self._is_sync_litellm_request(litellm_params) try: start_time, end_time = self._failure_handler_helper_fn( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fa7faf3035d..29c0d0629e8 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1808,8 +1808,10 @@ class CustomStreamWrapper: processed_chunk, None, None, cache_hit ) ) - ## SYNC LOGGING - self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) + ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler + litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) + if self.logging_obj._is_sync_litellm_request(litellm_params): + self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): model_response = self.model_response_creator() @@ -2206,23 +2208,19 @@ class CustomStreamWrapper: cache_hit, ) else: + # prefer_async_handlers routes CustomLogger to async_success_handler + # when consumers use ``async for`` on sync-SDK streams. Legacy string + # callbacks still run via executor.submit inside dispatch_success_handlers. asyncio.create_task( - self.logging_obj.async_success_handler( + self.logging_obj.dispatch_success_handlers( complete_streaming_response, cache_hit=cache_hit, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) - executor.submit( - self.logging_obj.success_handler, - complete_streaming_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) - raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 038d2d81277..63fd16fe8a9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1266,7 +1266,7 @@ class ProxyBaseLLMRequestProcessing: # (ProxyLogging._fire_deferred_stream_logging) fires the # closure after the full streaming pipeline finishes. # The closure runs non-apply_guardrail hooks on the - # assembled response, then fires both logging handlers. + # assembled response, then fires success logging. # Only for CustomStreamWrapper — raw async generators from # passthrough routes bypass CSW and would orphan the closure. from litellm.litellm_core_utils.streaming_handler import ( @@ -1387,33 +1387,18 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr] try: asyncio.create_task( - logging_obj.async_success_handler( + logging_obj.dispatch_success_handlers( response, cache_hit=None, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) except Exception as e: verbose_proxy_logger.exception( "Error in orphaned streaming async logging: %s", e ) - try: - from litellm.litellm_core_utils.thread_pool_executor import ( - executor as _exc, - ) - - _exc.submit( - logging_obj.success_handler, - response, - cache_hit=None, - start_time=None, - end_time=None, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in orphaned streaming sync logging: %s", e - ) # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. @@ -1615,7 +1600,7 @@ class ProxyBaseLLMRequestProcessing: ) -> None: """ Run non-streaming post-call guardrail hooks on an assembled streaming - response, then fire both async and sync logging handlers. + response, then fire success logging via ``dispatch_success_handlers``. Called by ProxyLogging._fire_deferred_stream_logging after the full streaming pipeline (including unified_guardrail end-of-stream blocks) @@ -1631,8 +1616,6 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can call the production implementation directly rather than reimplementing the closure. """ - from litellm.litellm_core_utils.thread_pool_executor import executor - _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router @@ -1691,31 +1674,23 @@ class ProxyBaseLLMRequestProcessing: ) finally: try: + # Proxy streaming always runs in async context and proxy spend + # logging is async-only; force async dispatch so DB/spend + # callbacks fire regardless of the call-type heuristic in + # _is_sync_litellm_request (which only recognizes a subset of + # async markers stored in litellm_params). asyncio.create_task( - captured_logging_obj.async_success_handler( + captured_logging_obj.dispatch_success_handlers( _response, cache_hit=cache_hit, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) except Exception as e: verbose_proxy_logger.exception( - "Error in deferred streaming async logging: %s", - e, - ) - - try: - executor.submit( - captured_logging_obj.success_handler, - _response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in deferred streaming sync logging: %s", + "Error in deferred streaming success logging: %s", e, ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index cbfcd34c438..d69a66ae3f5 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -7,7 +7,6 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.proxy._types import PassThroughEndpointLoggingResultValues from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -173,25 +172,16 @@ class PassThroughStreamingHandler: standard_logging_response_object = StandardPassThroughResponseObject( response=f"cannot parse chunks to standard response object. Chunks={all_chunks}" ) - await litellm_logging_obj.async_success_handler( + # Always reached from an async context (anthropic_messages, + # google_genai, and proxy pass-through stream tasks). prefer_async_handlers + # keeps async-only loggers running even when call_type isn't pass_through + # and litellm_params lacks an async flag (e.g. aanthropic_messages). + await litellm_logging_obj.dispatch_success_handlers( result=standard_logging_response_object, start_time=start_time, end_time=end_time, cache_hit=False, - **kwargs, - ) - if ( - litellm_logging_obj._should_run_sync_callbacks_for_async_calls() - is False - ): - return - - executor.submit( - litellm_logging_obj.success_handler, - result=standard_logging_response_object, - end_time=end_time, - cache_hit=False, - start_time=start_time, + prefer_async_handlers=True, **kwargs, ) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 0bc0183aa7c..292871bae67 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) from litellm.types.utils import StandardPassThroughResponseObject -from litellm.utils import executor as thread_pool_executor from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -94,19 +93,15 @@ class PassThroughEndpointLogging: cache_hit: bool, **kwargs, ): - """Helper function to handle both sync and async logging operations""" - # Submit to thread pool for sync logging - thread_pool_executor.submit( - logging_obj.success_handler, - standard_logging_response_object, - start_time, - end_time, - cache_hit, - **kwargs, - ) - - # Handle async logging - await logging_obj.async_success_handler( + """Log pass-through success via the shared async dispatch path.""" + # Always reached from pass_through_async_success_handler, which runs in + # an async context. call_type is "pass_through_endpoint" here, so the + # passthrough guard in dispatch_success_handlers already forces the + # async handler to run; pass prefer_async_handlers explicitly to match + # the streaming sibling (_route_streaming_logging_to_handler) and keep + # async-only loggers (e.g. the proxy spend logger) firing regardless of + # how the call-type classification evolves. + await logging_obj.dispatch_success_handlers( result=( json.dumps(result) if isinstance(result, dict) @@ -115,6 +110,7 @@ class PassThroughEndpointLogging: start_time=start_time, end_time=end_time, cache_hit=False, + prefer_async_handlers=True, **kwargs, ) diff --git a/tests/pass_through_unit_tests/test_unit_test_streaming.py b/tests/pass_through_unit_tests/test_unit_test_streaming.py index 38b650121bd..63965320f2b 100644 --- a/tests/pass_through_unit_tests/test_unit_test_streaming.py +++ b/tests/pass_through_unit_tests/test_unit_test_streaming.py @@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route): ), "Collected chunks do not match raw chunks" +@pytest.mark.asyncio +async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough(): + """ + SDK pass-through streaming (anthropic_messages, google generate_content) must run + the async success handler so async-only loggers record the assembled stream. + + Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as + sync SDK requests because call_type is not ``pass_through_endpoint`` and + litellm_params carries no ``acompletion`` flag, so only the sync success_handler + ran and CustomLogger.async_log_success_event never fired. + """ + import time + + from litellm.types.utils import CallTypes + + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True} + + with ( + patch.object( + PassThroughStreamingHandler, + "_build_passthrough_logging_result", + return_value=({"id": "slp"}, {}), + ), + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + ): + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + raw_bytes=[], + end_time=datetime.now(), + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + + +@pytest.mark.asyncio +async def test_handle_logging_runs_async_handler_for_passthrough(): + """ + Non-streaming pass-through logging (_handle_logging) must always run the + async success handler so async-only loggers (e.g. the proxy spend logger) + record the request. + + _handle_logging is only ever reached from pass_through_async_success_handler + (an async context), so it forces async dispatch via prefer_async_handlers. + This pins that contract independent of the call-type classification: even a + call_type that _is_sync_litellm_request would classify as sync (here + "completion" with no async marker in litellm_params) must still reach + async_success_handler. Without prefer_async_handlers=True the sync-only + branch would return early and async_log_success_event would never fire. + """ + import time + + from litellm.types.utils import CallTypes + + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=CallTypes.completion.value, + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {} + + handler = PassThroughEndpointLogging() + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + ): + await handler._handle_logging( + logging_obj=logging_obj, + standard_logging_response_object={"id": "slp"}, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + + def test_convert_raw_bytes_to_str_lines(): """ Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/proxy_unit_tests/test_proxy_reject_logging.py index 51a92fa3b4b..e0b575f4a71 100644 --- a/tests/proxy_unit_tests/test_proxy_reject_logging.py +++ b/tests/proxy_unit_tests/test_proxy_reject_logging.py @@ -95,6 +95,21 @@ router = Router( ) +def _register_proxy_test_logger(callback_logger: testLogger) -> None: + """ + Register the test logger on global callback lists. + + ``function_setup`` dedupes by object identity; each parametrized case + constructs a new ``testLogger`` and must replace the global lists, not + only ``litellm.callbacks``. + """ + litellm.callbacks = [callback_logger] + litellm.success_callback = [callback_logger] + litellm.failure_callback = [callback_logger] + litellm._async_success_callback = [callback_logger] + litellm._async_failure_callback = [callback_logger] + + @pytest.mark.parametrize( "route, body", [ @@ -115,7 +130,7 @@ router = Router( "/v1/embeddings", { "input": "The food was delicious and the waiter...", - "model": "text-embedding-ada-002", + "model": "fake-model", "encoding_format": "float", }, ), @@ -133,7 +148,7 @@ async def test_chat_completion_request_with_redaction(route, body): setattr(proxy_server, "llm_router", router) _test_logger = testLogger() - litellm.callbacks = [_test_logger] + _register_proxy_test_logger(_test_logger) litellm.set_verbose = True # Prepare the query string 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 c6961477a58..8b10288522b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,6 +1,7 @@ import os import sys -from unittest.mock import MagicMock, patch +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -786,6 +787,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call dummy_logger.log_stream_event.assert_not_called() +def test_is_sync_litellm_request(): + assert LitellmLogging._is_sync_litellm_request({}) is True + assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream( + logging_obj, +): + """Second final-stream dispatch must not re-export (CSW + deferred guardrail paths).""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_async_callbacks = list(litellm._async_success_callback or []) + litellm._async_success_callback = [mock_callback] + + result = ModelResponse( + id="resp-dedupe", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + try: + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + + with ( + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object(mock_callback, "log_success_event") as mock_sync_log, + patch.object( + logging_obj, + "_success_handler_helper_fn", + return_value=(time.time(), time.time(), result), + ), + patch.object( + logging_obj, + "_get_assembled_streaming_response", + return_value=result, + ), + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=True, + ), + ): + await logging_obj.dispatch_success_handlers(result=result) + await logging_obj.dispatch_success_handlers(result=result) + + mock_async_log.assert_awaited_once() + mock_sync_log.assert_not_called() + finally: + litellm._async_success_callback = original_async_callbacks + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream( + logging_obj, +): + """Sync dispatch path must also dedupe when dispatch is called twice.""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_success_callbacks = list(litellm.success_callback or []) + litellm.success_callback = [mock_callback] + + result = ModelResponse( + id="resp-sync-dedupe", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + try: + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(mock_callback, "log_success_event") as mock_sync_log, + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object( + logging_obj, + "_success_handler_helper_fn", + return_value=(time.time(), time.time(), result), + ), + patch.object( + logging_obj, + "_get_assembled_streaming_response", + return_value=result, + ), + ): + await logging_obj.dispatch_success_handlers(result=result) + await logging_obj.dispatch_success_handlers(result=result) + + mock_sync_log.assert_called_once() + mock_async_log.assert_not_awaited() + finally: + litellm.success_callback = original_success_callbacks + + +@pytest.mark.asyncio +async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks( + logging_obj, +): + """``prefer_async_handlers`` must not skip executor.submit for string callbacks.""" + result = ModelResponse( + id="resp-prefer-async", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + ) + + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=True, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_success_handlers( + result=result, + prefer_async_handlers=True, + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + mock_submit.assert_called_once() + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through( + logging_obj, +): + """Pass-through must use async_success_handler (CustomLogger skips sync success_handler).""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import CallTypes + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_async_callbacks = list(litellm._async_success_callback or []) + litellm._async_success_callback = [mock_callback] + + logging_obj.call_type = CallTypes.pass_through.value + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {} + + try: + with ( + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object(mock_callback, "log_success_event") as mock_sync_log, + ): + await logging_obj.dispatch_success_handlers(result={"id": "pt-1"}) + + mock_async_log.assert_awaited_once() + mock_sync_log.assert_not_called() + finally: + litellm._async_success_callback = original_async_callbacks + + def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" import datetime @@ -1351,7 +1557,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path() Test that _generate_cold_storage_object_key uses s3_path from custom logger instance. """ from datetime import datetime, timezone - from unittest.mock import MagicMock, patch + from unittest.mock import AsyncMock, MagicMock, patch from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -1404,7 +1610,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path(): Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path. """ from datetime import datetime, timezone - from unittest.mock import MagicMock, patch + from unittest.mock import AsyncMock, MagicMock, patch from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 49d3c51e340..63e2cb7f35c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool): == final_usage_block ) - print(mock_log_success_event.call_args.kwargs.keys()) - def test_streaming_handler_with_stop_chunk( initialized_custom_stream_wrapper: CustomStreamWrapper, @@ -2036,23 +2034,19 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool chunks.append(chunk) # The prompt_filter chunk should be forwarded with choices=[] - assert len(chunks[0].choices) == 0, ( - f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" - ) + assert ( + len(chunks[0].choices) == 0 + ), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" # At least one chunk must have role='assistant' in its delta has_role = any( - len(c.choices) > 0 - and getattr(c.choices[0].delta, "role", None) == "assistant" + len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant" for c in chunks ) assert has_role, ( "No chunk contained role='assistant' in delta (issue #24221). " "Chunk deltas: " - + str([ - c.choices[0].delta if c.choices else "no choices" - for c in chunks - ]) + + str([c.choices[0].delta if c.choices else "no choices" for c in chunks]) ) diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e10258c0829..e9ff193e044 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -18,7 +18,7 @@ import asyncio import os import sys from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks # --------------------------------------------------------------------------- +def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn): + """Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch.""" + + async def dispatch_success_handlers( + result=None, start_time=None, end_time=None, cache_hit=None, **kwargs + ): + await async_success_fn( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + + mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers + mock_logging_obj.async_success_handler = async_success_fn + + class PostCallGuardrail(CustomGuardrail): """A post-call guardrail.""" @@ -454,7 +472,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) tracking_guardrail = TrackingGuardrail() tracking_logger = TrackingLogger() @@ -511,7 +529,7 @@ class TestDeferredStreamingClosure: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class ModifyingGuardrail(CustomGuardrail): def __init__(self): @@ -573,7 +591,7 @@ class TestDeferredStreamingClosure: nonlocal logging_called logging_called = True - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = BlockingGuardrail() @@ -621,7 +639,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = TransientErrorGuardrail() @@ -656,7 +674,7 @@ class TestDeferredStreamingClosure: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class TestGuardrail(CustomGuardrail): def __init__(self): @@ -739,7 +757,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = ApplyGuardrailType() @@ -792,7 +810,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = IteratorHookGuardrail() @@ -847,7 +865,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = InspectingGuardrail() @@ -914,7 +932,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail_a = TaggedGuardrail("guardrail-a") guardrail_b = TaggedGuardrail("guardrail-b") @@ -962,7 +980,7 @@ class TestDeferredStreamingClosure: nonlocal logging_called logging_called = True - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) def exploding_merge(data, llm_router): raise RuntimeError("Simulated init failure") @@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure: logging_called is True ), "Logging must fire even when guardrail initialization raises" + @pytest.mark.asyncio + async def test_deferred_logging_forces_async_for_sync_classified_call_type(self): + """ + Regression: proxy deferred streaming logging must reach the async success + handler (which runs the async-only DB/spend logger) even when the call + type is classified as a sync SDK request by _is_sync_litellm_request. + + Without prefer_async_handlers=True, an async proxy stream whose + litellm_params lacks a recognized async marker would enter the sync + branch of dispatch_success_handlers and silently skip spend tracking. + + Uses the real dispatch_success_handlers via the production + _run_deferred_stream_guardrails entrypoint. + """ + import time + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", # not pass_through_endpoint + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + # litellm_params with no recognized async marker -> classified sync. + logging_obj.model_call_details["litellm_params"] = {} + assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + patch("litellm.callbacks", [PostCallGuardrail()]), + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4o-mini", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + # --------------------------------------------------------------------------- # 7. _fire_deferred_stream_logging @@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class InfoWritingGuardrail(CustomGuardrail): def __init__(self): From c645952a2101423376f2697f0e4c2f4c1674a54e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 30 May 2026 14:04:22 -0700 Subject: [PATCH 05/10] refactor(proxy/auth): normalize Bearer prefix in safe-hash helper (#29343) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(proxy/auth): normalize Bearer prefix in safe-hash helper UserAPIKeyAuth._safe_hash_litellm_api_key now strips a leading "Bearer "/"bearer " prefix before its existing sk-/JWT classification, so the helper produces the same hashed output regardless of whether the caller stripped the Authorization header prefix or passed the header value through unchanged. * refactor(proxy/auth): make Bearer-prefix strip case-insensitive Per RFC 7235 the HTTP authorization scheme token is case-insensitive. Replace the two-prefix loop with a single case-insensitive check so the helper normalizes "Bearer ", "bearer ", "BEARER ", and any mixed-case variant before classifying the remainder as sk- or JWT. The contract test gains coverage of "BEARER " and "BeArEr ". * test(mcp): align auth-handler test expectations with safe-hash helper The two MCP auth tests asserted that UserAPIKeyAuth(api_key="Bearer ...") retained the raw header bytes on the api_key field. _safe_hash_litellm_api_key now normalizes that input — stripping the Bearer prefix and hashing the resulting sk- key — so the expectations move to the normalized form: the bare token in the parametrize case, and hash_token("sk-...") in the backward-compat assertion. This matches what the real auth flow produces (the builder strips Bearer and the DB stores the hashed token), so the mocks now line up with production rather than with the un-normalized validator output. (cherry picked from commit 87b0e47485796a4fd1da802c8eef3f9d923dde2f) --- litellm/proxy/_types.py | 13 ++++++++----- .../auth/test_user_api_key_auth_mcp.py | 6 ++++-- tests/test_litellm/proxy/test_proxy_types.py | 18 ++++++++++++++++++ 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4fa497698a..e1ccc3d181a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2713,13 +2713,16 @@ class UserAPIKeyAuth( 1. Regular API keys from LiteLLM DB 2. JWT tokens used for connecting to LiteLLM API """ - if api_key.startswith("sk-"): - return hash_token(api_key) + normalized = api_key + if normalized[:7].lower() == "bearer ": + normalized = normalized[7:] + if normalized.startswith("sk-"): + return hash_token(normalized) from litellm.proxy.auth.handle_jwt import JWTHandler - if JWTHandler.is_jwt(token=api_key): - return f"hashed-jwt-{hash_token(token=api_key)}" - return api_key + if JWTHandler.is_jwt(token=normalized): + return f"hashed-jwt-{hash_token(token=normalized)}" + return normalized @classmethod def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth": diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 88742c67a86..868d1ec7a09 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -213,7 +213,7 @@ class TestMCPRequestHandler: # Test case 2: Authorization header present (fallback) ( [(b"authorization", b"Bearer test-auth-token")], - "Bearer test-auth-token", + "test-auth-token", None, {}, ), @@ -674,7 +674,9 @@ class TestMCPOAuth2AuthFlow: ) = await MCPRequestHandler.process_mcp_request(scope) # Should succeed with the LiteLLM key from Authorization header - assert auth_result.api_key == "Bearer sk-litellm-valid-key" + from litellm.proxy.utils import hash_token + + assert auth_result.api_key == hash_token("sk-litellm-valid-key") mock_auth.assert_called_once() async def test_non_auth_http_exception_still_raises(self): diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index 0fa86798999..dbb952968ed 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -69,3 +69,21 @@ def test_internal_jobs_user_has_proxy_admin_role(): assert system_user.user_id == "system" assert system_user.team_id == "system" assert system_user.team_alias == "system" + + +def test_user_api_key_auth_hashes_authorization_header_form_of_key(): + from litellm.proxy._types import UserAPIKeyAuth + + raw_key = "sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789" + baseline = UserAPIKeyAuth(api_key=raw_key) + + for header_form in ( + f"Bearer {raw_key}", + f"bearer {raw_key}", + f"BEARER {raw_key}", + f"BeArEr {raw_key}", + ): + from_header = UserAPIKeyAuth(api_key=header_form) + assert from_header.api_key == baseline.api_key + assert from_header.token == baseline.token + assert not from_header.api_key.lower().startswith("bearer") From 98f0fd3b3efd9ac92d9dabce02cbaf75ba761e68 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 30 May 2026 17:48:16 -0700 Subject: [PATCH 06/10] fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter (#29358) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter ResetBudgetJob's batched update_data path shipped the full key/user/team model on each reset. Prisma rejects object_permission_id and budget_limits on the update input type, so any row carrying those fields detonated the entire batch -- spend never reset, budget_reset_at never advanced. After v1.84.0 started populating object_permission_id on UI-created keys, this fires routinely. _reset_budget_common also zeroed the cross-pod spend counter before the DB write, so failed resets left enforcement reading 0 from the counter while the DB still held the over-budget spend, admitting requests past the cap until the counter naturally re-saturated from new reservations. Switch the write to per-row narrow updates ({spend, budget_reset_at}) via db.batch_, and move the counter invalidation out of _reset_budget_common so it only fires after the DB write commits. On DB-write failure the counter is left untouched, enforcement continues to block, and the next scheduler tick can retry without leaving a bypass window. Fixes #27730. * fix(reset_budget): address Greptile review on #29358 - Strengthen the bypass-half regression test: replace the for-loop over call_args_list (vacuously true when empty) with assert_not_called(), so the test would actually flag a re-introduction of counter-zeroing via any code path. - Add the same explanatory docstring on _write_user_reset_updates and _write_team_reset_updates that _write_key_reset_updates already has, so all three helpers point future maintainers at #27730. * test(reset_budget): update test_proxy_budget_reset for new batch-write path Same shape as the previous test_reset_budget_job.py update: keys/users/teams now write through prisma.db.batch_()..update, not update_data, so the tests need a batcher mock and updated assertions. Adds: - _wire_batcher_for_test helper that returns a list which accumulates per-row batch updates captured from prisma_client.db.batch_(). - _attrify helper that wraps dict fixtures so getattr(item, "token") works alongside the dict item-access the fake_reset_* mocks rely on. The new narrow-write helpers use getattr to pull out the row's id, and would silently skip plain dicts otherwise. - Updates 3 partial_failure tests to assert against the batch-call list (rows by id, payload contains only {spend, budget_reset_at}) instead of update_data.assert_awaited_once + data_list inspection. - Updates test_reset_budget_continues_other_categories_on_failure: only budget + enduser still flow through update_data; key/user/team go through the batch path now. - Wires the batcher mock into 3 service_logger_*_success tests so commit() is actually awaitable and the success hook fires. These tests were silently passing locally only because the editable install in .venv pointed at the main repo, not the worktree — running pytest with PYTHONPATH overridden to the worktree (matching CI) reproduces the failures. (cherry picked from commit a06ec43b36006f73dee8aa8f6a3c5cf54300316e) --- .../proxy/common_utils/reset_budget_job.py | 129 ++++++----- .../test_proxy_budget_reset.py | 201 ++++++++++++----- .../common_utils/test_reset_budget_job.py | 211 ++++++++++++++++-- 3 files changed, 409 insertions(+), 132 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 52bbeaf2ad3..40c8caa49e5 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -414,6 +414,72 @@ class ResetBudgetJob: ) return [LiteLLM_EndUserTable(**row.dict()) for row in rows] + async def _write_key_reset_updates( + self, updated_keys: List[LiteLLM_VerificationToken] + ) -> None: + """ + Write per-row {spend, budget_reset_at} updates for keys. + + Avoids the batched full-model update path, which trips + prisma.errors.DataError on any row carrying object_permission_id or + budget_limits (see #27730). Both fields are rejected by Prisma's + update input type for LiteLLM_VerificationToken, and the failure + aborts the entire batch — silently leaving spend over the cap and + budget_reset_at unchanged forever. + """ + batcher = self.prisma_client.db.batch_() + for k in updated_keys: + token = getattr(k, "token", None) + if token is None: + continue + batcher.litellm_verificationtoken.update( + where={"token": token}, + data={"spend": 0, "budget_reset_at": k.budget_reset_at}, + ) + await batcher.commit() + + async def _write_user_reset_updates( + self, updated_users: List[LiteLLM_UserTable] + ) -> None: + """ + Write per-row {spend, budget_reset_at} updates for users. + + Mirrors _write_key_reset_updates — avoids the full-model update path + that trips Prisma's DataError on rows carrying unrecognised fields + (see #27730). + """ + batcher = self.prisma_client.db.batch_() + for u in updated_users: + user_id = getattr(u, "user_id", None) + if user_id is None: + continue + batcher.litellm_usertable.update( + where={"user_id": user_id}, + data={"spend": 0, "budget_reset_at": u.budget_reset_at}, + ) + await batcher.commit() + + async def _write_team_reset_updates( + self, updated_teams: List[LiteLLM_TeamTable] + ) -> None: + """ + Write per-row {spend, budget_reset_at} updates for teams. + + Mirrors _write_key_reset_updates — avoids the full-model update path + that trips Prisma's DataError on rows carrying unrecognised fields + (see #27730). + """ + batcher = self.prisma_client.db.batch_() + for t in updated_teams: + team_id = getattr(t, "team_id", None) + if team_id is None: + continue + batcher.litellm_teamtable.update( + where={"team_id": team_id}, + data={"spend": 0, "budget_reset_at": t.budget_reset_at}, + ) + await batcher.commit() + async def reset_budget_for_litellm_keys(self): """ Resets the budget for all the litellm keys @@ -455,11 +521,7 @@ class ResetBudgetJob: ) if updated_keys: - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_keys, - table_name="key", - ) + await self._write_key_reset_updates(updated_keys=updated_keys) for k in updated_keys: token = getattr(k, "token", None) if token: @@ -544,11 +606,7 @@ class ResetBudgetJob: "Updated users %s", json.dumps(updated_users, indent=4, default=str) ) if updated_users: - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_users, - table_name="user", - ) + await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: user_id = getattr(u, "user_id", None) if user_id: @@ -641,11 +699,7 @@ class ResetBudgetJob: "Updated teams %s", json.dumps(updated_teams, indent=4, default=str) ) if updated_teams: - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_teams, - table_name="team", - ) + await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: team_id = getattr(t, "team_id", None) if team_id: @@ -816,49 +870,16 @@ class ResetBudgetJob: """ In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration - Common logic for resetting budget for a team, user, or key + Common logic for resetting budget for a team, user, or key. + + Spend-counter invalidation happens in the caller, AFTER the DB write + commits. Zeroing the counter here would open a bypass window when the + DB write fails: get_current_spend reads 0 from Redis while the DB + still holds the pre-reset value, admitting requests past the cap. """ try: item.spend = 0.0 - - # Reset the cross-pod spend counter. - # Reset Redis directly (not via DualCache) so a Redis failure - # doesn't silently leave a stale counter that get_current_spend - # would read as authoritative, permanently blocking the user. - from litellm.proxy.proxy_server import spend_counter_cache - - counter_key = None - if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr] - counter_key = f"spend:key:{item.token}" # type: ignore[union-attr] - elif ( - item_type == "team" - and hasattr(item, "team_id") - and item.team_id is not None # type: ignore[union-attr] - ): - counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr] - - if counter_key is not None: - # Always reset in-memory (local fallback) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset spend counter in Redis for %s key=%s: %s. " - "Budget may be over-enforced until counter expires.", - item_type, - counter_key, - redis_err, - ) - if hasattr(item, "budget_duration") and item.budget_duration is not None: - # Get standardized reset time based on budget duration from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_time, ) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 6240bedd3e6..5c96eb619bf 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -22,6 +22,60 @@ from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob # In a real-world scenario, these would be instances of LiteLLM_VerificationToken, LiteLLM_UserTable, etc. +def _attrify(d: dict): + """ + Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`, + etc.) works alongside the existing item-access the fake_reset_* helpers + rely on. The reset job's narrow-write helpers use `getattr(item, "token", + None)` (et al), which returns None for plain dicts — that would silently + skip the row. + """ + class _AttrDict(dict): + def __getattr__(self, k): + try: + return self[k] + except KeyError: + raise AttributeError(k) + + def __setattr__(self, k, v): + self[k] = v + + return _AttrDict(d) + + +def _wire_batcher_for_test(prisma_client): + """ + Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is + awaitable and whose per-table .update() calls get captured. The reset job + writes key/user/team resets via prisma.db.batch_().
.update — not via + prisma_client.update_data — so tests must let that batch path complete. + + Returns the list that will accumulate {table, where, data} dicts from + each captured update call. + """ + batch_calls = [] + + def make_batcher(): + class _Table: + def __init__(self, table_name): + self._table_name = table_name + + def update(self, where=None, data=None): + batch_calls.append( + {"table": self._table_name, "where": where, "data": data} + ) + + batcher = MagicMock() + batcher.litellm_verificationtoken = _Table("key") + batcher.litellm_usertable = _Table("user") + batcher.litellm_teamtable = _Table("team") + batcher.commit = AsyncMock(return_value=None) + return batcher + + prisma_client.db.batch_ = MagicMock(side_effect=make_batcher) + return batch_calls + + @pytest.mark.asyncio async def test_reset_budget_keys_partial_failure(): """ @@ -45,6 +99,9 @@ async def test_reset_budget_keys_partial_failure(): return_value=[key1, key2, key3, key4, key5, key6] ) prisma_client.update_data = AsyncMock() + # Reset job writes key resets via prisma.db.batch_().
.update — not + # via update_data — so wire that path. + batch_calls = _wire_batcher_for_test(prisma_client) # Using a dummy logging object with async hooks mocked out. proxy_logging_obj = MagicMock() @@ -56,6 +113,15 @@ async def test_reset_budget_keys_partial_failure(): now = datetime.utcnow() + # token is needed because the new write path uses where={"token": ...} + # and _AttrDict makes getattr work alongside item access used by fake_reset_key. + for k in [key1, key2, key3, key4, key5, key6]: + k.setdefault("token", k["id"]) + key1, key2, key3, key4, key5, key6 = ( + _attrify(k) for k in [key1, key2, key3, key4, key5, key6] + ) + prisma_client.get_data = AsyncMock(return_value=[key1, key2, key3, key4, key5, key6]) + async def fake_reset_key(key, current_time): if key["id"] == "key1": # Simulate a failure on key1 (for example, this might be due to an invariant check) @@ -80,17 +146,17 @@ async def test_reset_budget_keys_partial_failure(): # Assert that the helper was called for 6 keys assert mock_reset_key.call_count == 6 - # Assert that update_data was called once with a list containing all 6 keys - prisma_client.update_data.assert_awaited_once() - update_call = prisma_client.update_data.call_args - assert update_call.kwargs.get("table_name") == "key" - updated_keys = update_call.kwargs.get("data_list", []) - assert len(updated_keys) == 5 - assert updated_keys[0]["id"] == "key2" - assert updated_keys[1]["id"] == "key3" - assert updated_keys[2]["id"] == "key4" - assert updated_keys[3]["id"] == "key5" - assert updated_keys[4]["id"] == "key6" + # Assert that the new narrow write path got 5 batched updates (key1 failed). + # update_data must NOT have been called for keys. + prisma_client.update_data.assert_not_awaited() + key_writes = [c for c in batch_calls if c["table"] == "key"] + assert len(key_writes) == 5 + written_ids = [c["where"]["token"] for c in key_writes] + assert written_ids == ["key2", "key3", "key4", "key5", "key6"] + # And every write must carry only {spend, budget_reset_at} — never the full row. + for c in key_writes: + assert set(c["data"].keys()) == {"spend", "budget_reset_at"} + assert c["data"]["spend"] == 0 # Verify that the failure logging hook was scheduled (due to the failure for key1) failure_hook_calls = ( @@ -125,6 +191,7 @@ async def test_reset_budget_users_partial_failure(): return_value=[user1, user2, user3, user4, user5, user6] ) prisma_client.update_data = AsyncMock() + batch_calls = _wire_batcher_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -133,6 +200,15 @@ async def test_reset_budget_users_partial_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) + # user_id required for the new write path's where clause; _AttrDict so + # getattr(u, 'user_id') works alongside the dict access fake_reset_user uses. + for u in [user1, user2, user3, user4, user5, user6]: + u.setdefault("user_id", u["id"]) + user1, user2, user3, user4, user5, user6 = ( + _attrify(u) for u in [user1, user2, user3, user4, user5, user6] + ) + prisma_client.get_data = AsyncMock(return_value=[user1, user2, user3, user4, user5, user6]) + async def fake_reset_user(user, current_time): if user["id"] == "user1": raise Exception("Simulated failure for user1") @@ -150,16 +226,14 @@ async def test_reset_budget_users_partial_failure(): await asyncio.sleep(0.1) assert mock_reset_user.call_count == 6 - prisma_client.update_data.assert_awaited_once() - update_call = prisma_client.update_data.call_args - assert update_call.kwargs.get("table_name") == "user" - updated_users = update_call.kwargs.get("data_list", []) - assert len(updated_users) == 5 - assert updated_users[0]["id"] == "user2" - assert updated_users[1]["id"] == "user3" - assert updated_users[2]["id"] == "user4" - assert updated_users[3]["id"] == "user5" - assert updated_users[4]["id"] == "user6" + prisma_client.update_data.assert_not_awaited() + user_writes = [c for c in batch_calls if c["table"] == "user"] + assert len(user_writes) == 5 + written_ids = [c["where"]["user_id"] for c in user_writes] + assert written_ids == ["user2", "user3", "user4", "user5", "user6"] + for c in user_writes: + assert set(c["data"].keys()) == {"spend", "budget_reset_at"} + assert c["data"]["spend"] == 0 failure_hook_calls = ( proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list @@ -308,6 +382,7 @@ async def test_reset_budget_teams_partial_failure(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=[team1, team2]) prisma_client.update_data = AsyncMock() + batch_calls = _wire_batcher_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -316,6 +391,12 @@ async def test_reset_budget_teams_partial_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) + # team_id required for the new write path's where clause; _AttrDict for getattr. + for t in [team1, team2]: + t.setdefault("team_id", t["id"]) + team1, team2 = _attrify(team1), _attrify(team2) + prisma_client.get_data = AsyncMock(return_value=[team1, team2]) + async def fake_reset_team(team, current_time): if team["id"] == "team1": raise Exception("Simulated failure for team1") @@ -333,12 +414,12 @@ async def test_reset_budget_teams_partial_failure(): await asyncio.sleep(0.1) assert mock_reset_team.call_count == 2 - prisma_client.update_data.assert_awaited_once() - update_call = prisma_client.update_data.call_args - assert update_call.kwargs.get("table_name") == "team" - updated_teams = update_call.kwargs.get("data_list", []) - assert len(updated_teams) == 1 - assert updated_teams[0]["id"] == "team2" + prisma_client.update_data.assert_not_awaited() + team_writes = [c for c in batch_calls if c["table"] == "team"] + assert len(team_writes) == 1 + assert team_writes[0]["where"] == {"team_id": "team2"} + assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"} + assert team_writes[0]["data"]["spend"] == 0 failure_hook_calls = ( proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list @@ -402,6 +483,18 @@ async def test_reset_budget_continues_other_categories_on_failure(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + batch_calls = _wire_batcher_for_test(prisma_client) + # ID fields required by the new write path's where clauses; _AttrDict + # lets getattr() see them alongside the item-access fake_reset_* helpers use. + for k in [key1, key2]: + k.setdefault("token", k["id"]) + for u in [user1, user2]: + u.setdefault("user_id", u["id"]) + for t in [team1, team2]: + t.setdefault("team_id", t["id"]) + key1, key2 = _attrify(key1), _attrify(key2) + user1, user2 = _attrify(user1), _attrify(user2) + team1, team2 = _attrify(team1), _attrify(team2) # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( return_value={"count": 0} @@ -488,32 +581,29 @@ async def test_reset_budget_continues_other_categories_on_failure(): "team_membership", } - # Verify that update_data was called three times (one per category, enduser update includes two) - assert prisma_client.update_data.await_count == 5 + # After the fix, keys/users/teams write via prisma.db.batch_().
.update, + # so only budget + enduser still go through update_data. calls = prisma_client.update_data.await_args_list - - # Check keys update: both keys succeed. - keys_call = calls[0] - assert keys_call.kwargs.get("table_name") == "key" - assert len(keys_call.kwargs.get("data_list", [])) == 2 - - # Check users update: only user2 succeeded. - users_call = calls[1] - assert users_call.kwargs.get("table_name") == "user" - users_updated = users_call.kwargs.get("data_list", []) - assert len(users_updated) == 1 - assert users_updated[0]["id"] == "user2" - - # Check teams update: both teams succeed. - teams_call = calls[2] - assert teams_call.kwargs.get("table_name") == "team" - assert len(teams_call.kwargs.get("data_list", [])) == 2 + update_data_tables = [c.kwargs.get("table_name") for c in calls] + assert sorted(update_data_tables) == ["budget", "enduser"] # Check enduser update: enduser succeed. - enduser_call = calls[4] - assert enduser_call.kwargs.get("table_name") == "enduser" + enduser_call = next(c for c in calls if c.kwargs.get("table_name") == "enduser") assert len(enduser_call.kwargs.get("data_list", [])) == 1 + # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams. + key_writes = [c for c in batch_calls if c["table"] == "key"] + user_writes = [c for c in batch_calls if c["table"] == "user"] + team_writes = [c for c in batch_calls if c["table"] == "team"] + assert len(key_writes) == 2 + assert len(user_writes) == 1 + assert user_writes[0]["where"] == {"user_id": "user2"} + assert len(team_writes) == 2 + # Every batched write must carry only the two reset fields, never the full row. + for c in key_writes + user_writes + team_writes: + assert set(c["data"].keys()) == {"spend", "budget_reset_at"} + assert c["data"]["spend"] == 0 + # --------------------------------------------------------------------------- # Additional tests for service logger behavior (keys, users, teams, endusers) @@ -527,12 +617,13 @@ async def test_service_logger_keys_success(): logger success hook is called with the correct event metadata and no exception is logged. """ keys = [ - {"id": "key1", "spend": 10.0, "budget_duration": 60}, - {"id": "key2", "spend": 15.0, "budget_duration": 60}, + {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"}, + {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"}, ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=keys) prisma_client.update_data = AsyncMock() + _wire_batcher_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -644,12 +735,13 @@ async def test_service_logger_users_success(): the correct metadata and no exception is logged. """ users = [ - {"id": "user1", "spend": 20.0, "budget_duration": 120}, - {"id": "user2", "spend": 25.0, "budget_duration": 120}, + {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"}, + {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"}, ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=users) prisma_client.update_data = AsyncMock() + _wire_batcher_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -756,12 +848,13 @@ async def test_service_logger_teams_success(): the proper metadata and nothing is logged as an exception. """ teams = [ - {"id": "team1", "spend": 30.0, "budget_duration": 180}, - {"id": "team2", "spend": 35.0, "budget_duration": 180}, + {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"}, + {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"}, ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=teams) prisma_client.update_data = AsyncMock() + _wire_batcher_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 8a47c78db05..0b683745369 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -92,6 +92,37 @@ class MockLiteLLMEndUserTable: return self._find_many_results +class MockBatcher: + """Captures per-row update calls and exposes them after commit(). + + Mirrors prisma's `db.batch_()` ergonomics enough that the reset job's + narrow-write helpers (`_write_key_reset_updates` et al) can run against + the mock and the test can assert on what would have been written. + """ + + def __init__(self): + self.calls: List[Dict[str, Any]] = [] + self.committed: bool = False + + class _Table: + def __init__(_self, table_name: str, outer: "MockBatcher"): + _self._table_name = table_name + _self._outer = outer + + def update(_self, where, data): + _self._outer.calls.append( + {"table": _self._table_name, "where": where, "data": data} + ) + + self.litellm_verificationtoken = _Table("key", self) + self.litellm_usertable = _Table("user", self) + self.litellm_teamtable = _Table("team", self) + + async def commit(self): + self.committed = True + return self.calls + + class MockDB: def __init__(self): self.litellm_teammembership = MockLiteLLMTeamMembership() @@ -99,6 +130,19 @@ class MockDB: self.litellm_endusertable = MockLiteLLMEndUserTable() self.litellm_organizationtable = MockLiteLLMOrganizationTable() self.litellm_tagtable = MockLiteLLMTagTable() + self.batch_calls: List[Dict[str, Any]] = [] + + def batch_(self): + batcher = MockBatcher() + # Aggregate calls across all batches so tests can assert on cumulative writes. + original_commit = batcher.commit + + async def _record_and_commit(): + self.batch_calls.extend(batcher.calls) + return await original_commit() + + batcher.commit = _record_and_commit # type: ignore[assignment] + return batcher class MockPrismaClient: @@ -205,6 +249,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): "budget_duration": "30d", "budget_reset_at": now, "id": "test-key-1", + "token": "tok-key-1", }, ) @@ -213,11 +258,16 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): # Run the test asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - # Verify results - assert len(mock_prisma_client.updated_data["key"]) == 1 - updated_key = mock_prisma_client.updated_data["key"][0] - assert updated_key.spend == 0.0 - assert updated_key.budget_reset_at > now + # The reset writes only {spend, budget_reset_at} per row via batch_(). + # Full-row writes would re-detonate the Prisma DataError on rows carrying + # object_permission_id / budget_limits (see #27730). + key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] + assert len(key_writes) == 1 + write = key_writes[0] + assert write["where"] == {"token": "tok-key-1"} + assert write["data"]["spend"] == 0 + assert write["data"]["budget_reset_at"] > now + assert set(write["data"].keys()) == {"spend", "budget_reset_at"} def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): @@ -231,6 +281,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): "budget_duration": "7d", "budget_reset_at": now, "id": "test-user-1", + "user_id": "uid-1", }, ) @@ -239,11 +290,13 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): # Run the test asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - # Verify results - assert len(mock_prisma_client.updated_data["user"]) == 1 - updated_user = mock_prisma_client.updated_data["user"][0] - assert updated_user.spend == 0.0 - assert updated_user.budget_reset_at > now + user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"] + assert len(user_writes) == 1 + write = user_writes[0] + assert write["where"] == {"user_id": "uid-1"} + assert write["data"]["spend"] == 0 + assert write["data"]["budget_reset_at"] > now + assert set(write["data"].keys()) == {"spend", "budget_reset_at"} def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): @@ -257,6 +310,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): "budget_duration": "1mo", "budget_reset_at": now, "id": "test-team-1", + "team_id": "tid-1", }, ) @@ -265,11 +319,13 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): # Run the test asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) - # Verify results - assert len(mock_prisma_client.updated_data["team"]) == 1 - updated_team = mock_prisma_client.updated_data["team"][0] - assert updated_team.spend == 0.0 - assert updated_team.budget_reset_at > now + team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"] + assert len(team_writes) == 1 + write = team_writes[0] + assert write["where"] == {"team_id": "tid-1"} + assert write["data"]["spend"] == 0 + assert write["data"]["budget_reset_at"] > now + assert set(write["data"].keys()) == {"spend", "budget_reset_at"} def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): @@ -324,6 +380,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "budget_duration": "30d", "budget_reset_at": now, "id": "test-key-1", + "token": "tok-all-1", }, ) @@ -335,6 +392,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "budget_duration": "7d", "budget_reset_at": now, "id": "test-user-1", + "user_id": "uid-all-1", }, ) @@ -346,6 +404,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "budget_duration": "1mo", "budget_reset_at": now, "id": "test-team-1", + "team_id": "tid-all-1", }, ) @@ -379,17 +438,22 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Run the test asyncio.run(reset_budget_job.reset_budget()) - # Verify results - assert len(mock_prisma_client.updated_data["key"]) == 1 - assert len(mock_prisma_client.updated_data["user"]) == 1 - assert len(mock_prisma_client.updated_data["team"]) == 1 + # key/user/team rows are written via batch_().
.update — verify each + # one fired exactly once with the narrow {spend, budget_reset_at} payload. + for table_name, where in [ + ("key", {"token": "tok-all-1"}), + ("user", {"user_id": "uid-all-1"}), + ("team", {"team_id": "tid-all-1"}), + ]: + writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == table_name] + assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}" + assert writes[0]["where"] == where + assert writes[0]["data"]["spend"] == 0 + assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"} + + # Enduser + budget rows still go through update_data (not narrowed; different path). assert len(mock_prisma_client.updated_data["enduser"]) == 1 assert len(mock_prisma_client.updated_data["budget"]) == 1 - - # Check that all spends were reset to 0 - assert mock_prisma_client.updated_data["key"][0].spend == 0.0 - assert mock_prisma_client.updated_data["user"][0].spend == 0.0 - assert mock_prisma_client.updated_data["team"][0].spend == 0.0 assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0 @@ -1399,6 +1463,105 @@ def test_reset_budget_for_teams_invalidates_redis_counter( ) +def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): + """ + Regression for #27730 (the bypass-half). + + If the DB write inside the reset job raises (e.g. Prisma DataError on a + row carrying object_permission_id or budget_limits), the Redis spend + counter MUST NOT be zeroed — that would let get_current_spend admit + requests past the cap while the DB row still holds the over-budget + spend. + + Pre-fix: _reset_budget_common pre-zeroed the counter before the DB + write attempt, opening the bypass window. + Post-fix: counter invalidation lives in the caller, AFTER the DB write + commits. If the write raises, the post-write invalidation never runs. + """ + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + prisma_client = MagicMock() + + matching_key = type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": now - timedelta(seconds=1), + "token": "sk-failing", + }, + ) + + # get_data returns one key needing reset; the batched DB write then explodes. + async def fake_get_data(table_name, query_type, **kwargs): + if table_name == "key": + return [matching_key] + return [] + + prisma_client.get_data = fake_get_data + + batcher = MagicMock() + batcher.litellm_verificationtoken.update = MagicMock() + + async def failing_commit(): + raise RuntimeError("simulated Prisma DataError on update") + + batcher.commit = failing_commit + prisma_client.db.batch_ = MagicMock(return_value=batcher) + + job = ResetBudgetJob( + proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client + ) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + # CRITICAL: counter invalidation must NOT have been called at all — + # the DB write raised before the post-write invalidation loop. Using + # assert_not_called() instead of iterating call_args_list, because the + # latter is vacuously true when the list is empty (would pass even if + # the bypass were re-introduced via a different code path). + counter_cache.in_memory_cache.set_cache.assert_not_called() + + +def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, mock_prisma_client): + """ + Regression for #27730 (the trigger-half). + + The reset job must write only {spend, budget_reset_at} per row — never + the full key object. Sending the full object via the old update_data + batcher path made Prisma reject any row carrying object_permission_id + or budget_limits (both became non-NULL on UI-created keys after v1.84.0). + """ + now = datetime.now(timezone.utc) + key_with_problematic_fields = type( + "LiteLLM_VerificationToken", + (), + { + "spend": 50.0, + "budget_duration": "30d", + "budget_reset_at": now, + "token": "sk-problematic", + "object_permission_id": "perm-abc", # would be rejected on update + "budget_limits": [{"max_budget": 5}], # would be rejected on update + "metadata": {"some": "thing"}, + }, + ) + mock_prisma_client.data["key"] = [key_with_problematic_fields] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] + assert len(key_writes) == 1 + payload_keys = set(key_writes[0]["data"].keys()) + assert payload_keys == {"spend", "budget_reset_at"}, ( + f"reset payload must not include any field besides spend / budget_reset_at, " + f"got: {payload_keys}. Any extra field (object_permission_id, budget_limits, etc.) " + f"trips Prisma DataError and detonates the whole batch." + ) + + def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch): """Resetting keys via budget tier must clear each linked key's counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) From 0e1ff6ac232d7f6e6fdd0d5f27f9d16b4ade4761 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 1 Jun 2026 14:04:42 -0700 Subject: [PATCH 07/10] fix: stop use_chat_completions_api flag from leaking into provider request body (#29447) * fix: stop use_chat_completions_api flag from leaking into provider request body use_chat_completions_api is a LiteLLM control flag that forces the /responses -> /chat/completions bridge. It was missing from all_litellm_params, so get_non_default_completion_params treated it as a model-specific param and forwarded it to the upstream provider. A model-level "use_chat_completions_api: true" in the proxy config therefore reached the chat-completions path and was rejected by strict providers (OpenAI/Anthropic) with HTTP 400 for an unknown body field. Register it as a known internal param so it is stripped on every path (completion, the responses bridge that calls litellm.completion, and filter_out_litellm_params). Adds a regression test driving litellm.completion() with a mocked OpenAI client that asserts the flag never reaches the request body. * test: clarify extra_body assertion in use_chat_completions_api leak test Replace the misleading 'not in ... or {}' precedence idiom with an explicit parenthesized guard that also handles extra_body being None. (cherry picked from commit acbbfe9cae4f04d85be8df89ac18bffc7b2098ee) --- litellm/types/utils.py | 1 + .../test_use_chat_completions_api_no_leak.py | 74 +++++++++++++++++++ 2 files changed, 75 insertions(+) create mode 100644 tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 400edcac889..db598d85e55 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3148,6 +3148,7 @@ all_litellm_params = ( "allowed_openai_params", "litellm_session_id", "use_litellm_proxy", + "use_chat_completions_api", "prompt_label", "shared_session", "search_tool_name", diff --git a/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py new file mode 100644 index 00000000000..9a266fca81f --- /dev/null +++ b/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py @@ -0,0 +1,74 @@ +""" +Regression test for issue #28146. + +`use_chat_completions_api` is a LiteLLM-internal control flag (it forces the +/responses -> /chat/completions bridge). When set as a model-level param in the +proxy config, it must never be forwarded to the upstream provider's request +body. OpenAI/Anthropic reject unknown body params with HTTP 400. +""" + +import os +import sys +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.types.utils import all_litellm_params +from litellm.utils import get_non_default_completion_params + + +def test_use_chat_completions_api_is_a_known_litellm_param(): + assert "use_chat_completions_api" in all_litellm_params + + +def test_use_chat_completions_api_not_forwarded_as_provider_param(): + forwarded = get_non_default_completion_params( + {"use_chat_completions_api": True, "temperature": 0.5} + ) + assert "use_chat_completions_api" not in forwarded + + +def test_completion_does_not_leak_flag_into_provider_request_body(): + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + } + + mock_raw_response = MagicMock() + mock_raw_response.headers = {} + mock_raw_response.parse.return_value = mock_response + + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_raw_response + ) + + litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + use_chat_completions_api=True, + api_key="sk-test", + client=mock_client, + ) + + create_kwargs = ( + mock_client.chat.completions.with_raw_response.create.call_args.kwargs + ) + assert "use_chat_completions_api" not in create_kwargs + assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {}) From b7f514eeedb26d1d055303442bb4ceb19024a64f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 1 Jun 2026 17:31:05 -0700 Subject: [PATCH 08/10] fix(passthrough): extract _build_passthrough_logging_result helper The #29311 cherry-pick onto stable/1.85.x carried the test test_route_streaming_logging_runs_async_handler_for_sdk_passthrough, which patches PassThroughStreamingHandler._build_passthrough_logging_result to verify the SDK-passthrough dispatch contract. The manual conflict resolution kept the per-endpoint if/elif/elif chain inline in _route_streaming_logging_to_handler (matching v1.84.4's resolution), so the patched attribute did not exist and the test errored at collection with AttributeError. Extract the chain into the static _build_passthrough_logging_result helper as #29089 originally designed it. _route_streaming_logging_to_handler now resolves (standard_logging_response_object, kwargs) through the helper and dispatches via dispatch_success_handlers; the helper itself is synchronous and CPU-bound, suitable for the unit test's patch target. Verified locally: tests/pass_through_unit_tests/test_unit_test_streaming.py passes (5/5) and tests/test_litellm/litellm_core_utils/test_litellm_logging.py passes (83/83). v1.84.4 ships with the same broken test; this strictly improves on that resolution. (cherry picked from commit 8824745c4ee5ade0c6ad083fc7d07010dada5abd) --- .../streaming_handler.py | 155 +++++++++++------- 1 file changed, 97 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index d69a66ae3f5..7e7f0b42b4d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import List, Optional +from typing import List, Optional, Tuple import httpx @@ -114,64 +114,20 @@ class PassThroughStreamingHandler: - OpenAI """ try: - all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( - raw_bytes + ( + standard_logging_response_object, + kwargs, + ) = PassThroughStreamingHandler._build_passthrough_logging_result( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=end_time, + model=model, ) - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None - kwargs: dict = {} - if endpoint_type == EndpointType.ANTHROPIC: - anthropic_passthrough_logging_handler_result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - ) - standard_logging_response_object = ( - anthropic_passthrough_logging_handler_result["result"] - ) - kwargs = anthropic_passthrough_logging_handler_result["kwargs"] - elif endpoint_type == EndpointType.VERTEX_AI: - vertex_passthrough_logging_handler_result = VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - model=model, - ) - standard_logging_response_object = ( - vertex_passthrough_logging_handler_result["result"] - ) - kwargs = vertex_passthrough_logging_handler_result["kwargs"] - elif endpoint_type == EndpointType.OPENAI: - openai_passthrough_logging_handler_result = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - ) - standard_logging_response_object = ( - openai_passthrough_logging_handler_result["result"] - ) - kwargs = openai_passthrough_logging_handler_result["kwargs"] - - if standard_logging_response_object is None: - standard_logging_response_object = StandardPassThroughResponseObject( - response=f"cannot parse chunks to standard response object. Chunks={all_chunks}" - ) # Always reached from an async context (anthropic_messages, # google_genai, and proxy pass-through stream tasks). prefer_async_handlers # keeps async-only loggers running even when call_type isn't pass_through @@ -189,6 +145,89 @@ class PassThroughStreamingHandler: f"Error in _route_streaming_logging_to_handler: {str(e)}" ) + @staticmethod + def _build_passthrough_logging_result( + litellm_logging_obj: LiteLLMLoggingObj, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: dict, + endpoint_type: EndpointType, + start_time: datetime, + raw_bytes: List[bytes], + end_time: datetime, + model: Optional[str], + ) -> Tuple[PassThroughEndpointLoggingResultValues, dict]: + """ + Synchronous, CPU-bound reconstruction of the standard logging payload + from collected raw SSE bytes. Extracted from + _route_streaming_logging_to_handler so the per-endpoint dispatch can + be unit-tested in isolation. Still invoked synchronously on the event + loop; an off-loop dispatch is a future change, not part of this PR. + """ + all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( + raw_bytes + ) + standard_logging_response_object: Optional[ + PassThroughEndpointLoggingResultValues + ] = None + kwargs: dict = {} + if endpoint_type == EndpointType.ANTHROPIC: + anthropic_passthrough_logging_handler_result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + ) + standard_logging_response_object = ( + anthropic_passthrough_logging_handler_result["result"] + ) + kwargs = anthropic_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.VERTEX_AI: + vertex_passthrough_logging_handler_result = ( + VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + model=model, + ) + ) + standard_logging_response_object = ( + vertex_passthrough_logging_handler_result["result"] + ) + kwargs = vertex_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.OPENAI: + openai_passthrough_logging_handler_result = ( + OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + ) + ) + standard_logging_response_object = ( + openai_passthrough_logging_handler_result["result"] + ) + kwargs = openai_passthrough_logging_handler_result["kwargs"] + + if standard_logging_response_object is None: + standard_logging_response_object = StandardPassThroughResponseObject( + response=f"cannot parse chunks to standard response object. Chunks={all_chunks}" + ) + return standard_logging_response_object, kwargs + @staticmethod def _extract_model_for_cost_injection( request_body: Optional[dict], From 018351bd4a36c8d2b444653aab443eeeb0733ccc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 2 Jun 2026 16:40:12 -0700 Subject: [PATCH 09/10] =?UTF-8?q?bump:=20version=201.86.2=20=E2=86=92=201.?= =?UTF-8?q?86.3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9e5a01c9107..b13d0d1a930 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.86.2" +version = "1.86.3" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -251,7 +251,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.86.2" +version = "1.86.3" version_files = [ "pyproject.toml:^version", ] From 09186bf6bbf765f555e81b0c06c56445fc2baef5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 2 Jun 2026 16:40:58 -0700 Subject: [PATCH 10/10] chore: update uv.lock for 1.86.3 --- uv.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uv.lock b/uv.lock index cd7a717e312..52f673245e1 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-05-24T07:24:03.668568Z" +exclude-newer = "2026-05-30T23:40:12.713734Z" exclude-newer-span = "P3D" [manifest] @@ -3189,7 +3189,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.86.2" +version = "1.86.3" source = { editable = "." } dependencies = [ { name = "aiohttp" },