fix(vertex_ai): verify embedded thought signature checksum

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -66,6 +66,7 @@ BAD_MESSAGE_ERROR_STR: Final = "Invalid Message "
 # Separator used to embed Gemini thought signatures in tool call IDs
 # See: https://ai.google.dev/gemini-api/docs/thought-signatures
 THOUGHT_SIGNATURE_SEPARATOR: Final = "__thought__"
+THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR: Final = "__checksum__"
 
 # used to interweave user messages, to ensure user/assistant alternating
 DEFAULT_USER_CONTINUE_MESSAGE: Final = {
@@ -1134,7 +1135,7 @@ def _gemini_tool_call_invoke_helper(
         args=arguments_dict,
     )
     if tool_call_id:
-        clean_id: Final = tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0]
+        clean_id, _ = _decode_tool_call_id_with_signature(tool_call_id)
         if clean_id:
             function_call["id"] = clean_id
     return function_call
@@ -1150,28 +1151,21 @@ def _encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: st
 
     Returns:
         Tool call ID with embedded signature if present, otherwise original ID
-        Format: call_<uuid>__thought__<base64_signature>
+        Format: call_<uuid>__checksum__<sha256>__thought__<base64_signature>
 
     See: https://ai.google.dev/gemini-api/docs/thought-signatures
     """
     if thought_signature:
-        return f"{tool_call_id}{THOUGHT_SIGNATURE_SEPARATOR}{thought_signature}"
+        checksum: Final = hashlib.sha256(thought_signature.encode("utf-8")).hexdigest()
+        return (
+            f"{tool_call_id}{THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR}{checksum}"
+            f"{THOUGHT_SIGNATURE_SEPARATOR}{thought_signature}"
+        )
     return tool_call_id
 
 
 def _is_valid_thought_signature(signature: str) -> bool:
-    """
-    A Gemini ``thoughtSignature`` travels as standard-alphabet base64 bytes
-    (``A-Z a-z 0-9 + /`` with ``=`` padding). If a client normalizes the
-    ``tool_call_id`` that embeds the signature (e.g.
-    ``sanitized[:30] + "_" + sha256hex[:10]``), the round-tripped value can
-    still contain the ``__thought__`` marker but no longer base64-decode.
-    Forwarding such a value causes Vertex/AI Studio to reject the request with
-    ``Base64 decoding failed for "..."`` or ``Invalid thought signature``.
-
-    Use ``validate=True`` to match Vertex's strict decoder, and pad the tail
-    since Gemini often omits ``=`` on the wire.
-    """
+    """Check strict standard-alphabet base64 while allowing omitted padding."""
     if not signature:
         return False
     try:
@@ -1182,6 +1176,23 @@ def _is_valid_thought_signature(signature: str) -> bool:
     return True
 
 
+def _decode_tool_call_id_with_signature(tool_call_id: str) -> tuple[str, str | None]:
+    payload, signature_separator, signature = tool_call_id.partition(THOUGHT_SIGNATURE_SEPARATOR)
+    if not signature_separator:
+        return tool_call_id, None
+
+    base_id, checksum_separator, checksum = payload.rpartition(THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR)
+    if not checksum_separator:
+        return payload, None
+    if not _is_valid_thought_signature(signature):
+        return base_id, None
+
+    expected_checksum: Final = hashlib.sha256(signature.encode("utf-8")).hexdigest()
+    if checksum != expected_checksum:
+        return base_id, None
+    return base_id, signature
+
+
 def _get_thought_signature_from_tool(tool: dict) -> str | None:
     """Extract thought signature from tool call's provider_specific_fields.
 
@@ -1214,12 +1225,9 @@ def _get_thought_signature_from_tool(tool: dict) -> str | None:
                     return signature
     # Check if thought signature is embedded in tool call ID
     tool_call_id: Final = tool.get("id")
-    if tool_call_id and THOUGHT_SIGNATURE_SEPARATOR in tool_call_id:
-        parts: Final = tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)
-        if len(parts) == 2:
-            _, signature = parts
-            if _is_valid_thought_signature(signature):
-                return signature
+    if tool_call_id and isinstance(tool_call_id, str):
+        _, signature = _decode_tool_call_id_with_signature(tool_call_id)
+        return signature
     return None
 
 
@@ -1478,7 +1486,7 @@ def convert_to_gemini_tool_call_result(
     if forward_function_call_id:
         raw_tool_call_id: Final = message.get("tool_call_id")
         if raw_tool_call_id and isinstance(raw_tool_call_id, str):
-            stripped_id: Final = raw_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0]
+            stripped_id, _ = _decode_tool_call_id_with_signature(raw_tool_call_id)
             if stripped_id:
                 gemini_call_id = stripped_id
 

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

@@ -15,7 +15,9 @@ import pytest
 
 import litellm
 from litellm.litellm_core_utils.prompt_templates.factory import (
+    THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR,
     THOUGHT_SIGNATURE_SEPARATOR,
+    _decode_tool_call_id_with_signature,
     _encode_tool_call_id_with_signature,
     _get_dummy_thought_signature,
     _get_thought_signature_from_tool,
@@ -36,6 +38,7 @@ def test_encode_decode_tool_call_id_with_signature():
     # Test encoding
     encoded_id = _encode_tool_call_id_with_signature(base_id, test_signature)
     assert THOUGHT_SIGNATURE_SEPARATOR in encoded_id
+    assert THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR in encoded_id
     assert encoded_id.startswith(base_id)
 
     # Test decoding using factory function with realistic tool call structure
@@ -52,7 +55,7 @@ def test_encode_decode_tool_call_id_with_signature():
     assert extracted_signature == test_signature
 
     # Verify base ID is preserved
-    decoded_base_id = encoded_id.split(THOUGHT_SIGNATURE_SEPARATOR)[0]
+    decoded_base_id, _ = _decode_tool_call_id_with_signature(encoded_id)
     assert decoded_base_id == base_id
 
 
@@ -310,16 +313,22 @@ def test_is_valid_thought_signature_tolerates_missing_padding():
     assert _is_valid_thought_signature(encoded) is True
 
 
-def test_is_valid_thought_signature_scope_is_syntactic():
-    """The extractor validates the wire encoding, not semantic integrity: a
-    client normalizer that keeps standard base64 characters produces a value
-    that still decodes even though Vertex will reject it as an unauthenticated
-    signature. Callers that need semantic verification have to compare against
-    the provider-issued signature; the ID-embedded fallback can only catch
-    encoding damage. Documented so future readers understand issue #37849's
-    guarantee is bounded."""
+def test_get_thought_signature_rejects_decodable_tampering():
+    encoded_id = _encode_tool_call_id_with_signature("call_abc123", REAL_SIGNATURE)
+    checksum = encoded_id.split(THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR, 1)[1].split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0]
     decodable_but_wrong = base64.b64encode(b"not-the-real-signature").decode("ascii")
+    tampered_id = (
+        f"call_abc123{THOUGHT_SIGNATURE_CHECKSUM_SEPARATOR}{checksum}{THOUGHT_SIGNATURE_SEPARATOR}{decodable_but_wrong}"
+    )
+
     assert _is_valid_thought_signature(decodable_but_wrong) is True
+    assert _get_thought_signature_from_tool({"id": tampered_id, "type": "function"}) is None
+
+
+def test_get_thought_signature_rejects_unsigned_embedded_signature():
+    unsigned_id = f"call_abc123{THOUGHT_SIGNATURE_SEPARATOR}{REAL_SIGNATURE}"
+
+    assert _get_thought_signature_from_tool({"id": unsigned_id, "type": "function"}) is None
 
 
 def test_get_thought_signature_drops_client_mangled_id_suffix():
@@ -328,7 +337,7 @@ def test_get_thought_signature_drops_client_mangled_id_suffix():
     skip-validator signature (Gemini 3+) or drop the signature entirely
     (Gemini 2.x), instead of forwarding a corrupted value that Vertex would
     reject. Regression test for issue #37849."""
-    mangled_id = f"call_2156408{THOUGHT_SIGNATURE_SEPARATOR}AY89a1/_57b05e78dc"
+    mangled_id = _encode_tool_call_id_with_signature("call_2156408", "AY89a1/_57b05e78dc")
     tool = {"id": mangled_id, "type": "function"}
 
     assert _get_thought_signature_from_tool(tool) is None