Merge 553c3339b0 into c2c2a623c0

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -1195,6 +1195,12 @@ def _get_thought_signature_from_tool(tool: dict) -> str | None:
         parts: Final = tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)
         if len(parts) == 2:
             _, signature = parts
+            if not signature:
+                return None
+            try:
+                base64.b64decode(signature, validate=True)
+            except ValueError:
+                return None
             return signature
     return None
 

tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py

@@ -11,10 +11,10 @@ enable_preview_features=True to be enabled.
 
 import pytest
 
-import litellm
 from litellm.litellm_core_utils.prompt_templates.factory import (
     THOUGHT_SIGNATURE_SEPARATOR,
     _encode_tool_call_id_with_signature,
+    _get_dummy_thought_signature,
     _get_thought_signature_from_tool,
     convert_to_gemini_tool_call_invoke,
 )
@@ -94,10 +94,7 @@ def test_tool_call_id_includes_signature_in_response(enable_preview_features):
     tool_call_id = tools[0]["id"]
 
     # Verify signature is always in provider_specific_fields
-    assert (
-        tools[0].get("provider_specific_fields", {}).get("thought_signature")
-        == test_signature
-    )
+    assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == test_signature
 
     # When preview features enabled, signature should be embedded in ID
     assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id
@@ -181,6 +178,29 @@ def test_convert_to_gemini_with_embedded_signature():
     assert gemini_parts[0]["thoughtSignature"] == test_signature
 
 
+def test_convert_to_gemini_rejects_malformed_embedded_signature():
+    malformed_signature = "AY89a1/_57b05e78dc"
+    assistant_message = {
+        "role": "assistant",
+        "content": None,
+        "tool_calls": [
+            {
+                "id": f"call_abc123{THOUGHT_SIGNATURE_SEPARATOR}{malformed_signature}",
+                "type": "function",
+                "function": {
+                    "name": "get_current_temperature",
+                    "arguments": '{"location": "Paris"}',
+                },
+            }
+        ],
+    }
+
+    gemini_parts = convert_to_gemini_tool_call_invoke(assistant_message, model="gemini-3-pro-preview")
+
+    assert gemini_parts[0]["thoughtSignature"] == _get_dummy_thought_signature()
+    assert gemini_parts[0]["thoughtSignature"] != malformed_signature
+
+
 @pytest.mark.parametrize("enable_preview_features", [True, False])
 def test_openai_client_e2e_flow(enable_preview_features):
     """
@@ -234,9 +254,7 @@ def test_openai_client_e2e_flow(enable_preview_features):
         ],
     }
     # Step 4: LiteLLM converts back to Gemini format, extracting signature
-    gemini_parts_converted = convert_to_gemini_tool_call_invoke(
-        openai_assistant_message
-    )
+    gemini_parts_converted = convert_to_gemini_tool_call_invoke(openai_assistant_message)
 
     # Verify signature is preserved through the round trip
     assert len(gemini_parts_converted) == 1
@@ -247,7 +265,7 @@ def test_openai_client_e2e_flow(enable_preview_features):
 @pytest.mark.parametrize("enable_preview_features", [True, False])
 def test_parallel_tool_calls_with_signatures(enable_preview_features):
     """Test that parallel tool calls preserve signatures correctly"""
-    signature1 = "signature_for_first_call"
+    signature1 = "c2lnbmF0dXJlX2Zvcl9maXJzdF9jYWxs"
     # Only first call has signature (Gemini behavior for parallel calls)
 
     gemini_parts = [
@@ -271,10 +289,7 @@ def test_parallel_tool_calls_with_signatures(enable_preview_features):
     assert len(tools) == 2
 
     # First tool call should have signature in provider_specific_fields
-    assert (
-        tools[0].get("provider_specific_fields", {}).get("thought_signature")
-        == signature1
-    )
+    assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == signature1
 
     # When preview features enabled, first tool call has signature in ID
     assert THOUGHT_SIGNATURE_SEPARATOR in tools[0]["id"]