fix(vertex_ai): validate Base64 thought signatures during Gemini session replay (#42201)

litellm/llms/vertex_ai/gemini/transformation.py

@@ -4,6 +4,8 @@ Transformation logic from OpenAI format to Gemini format.
 Why separate file? Make it easy to see how transformation works
 """
 
+import base64
+import binascii
 import json
 import os
 import re
@@ -12,8 +14,6 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast
 from urllib.parse import quote
 
 import httpx
-from pydantic import BaseModel
-
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.asyncify import asyncify
@@ -56,6 +56,7 @@ from litellm.types.llms.vertex_ai import (
     Tools,
 )
 from litellm.types.utils import GenericImageParsingChunk, LlmProviders
+from pydantic import BaseModel
 
 from ..common_utils import (
     _check_text_in_content,
@@ -179,7 +180,7 @@ def _apply_gemini_metadata(
     part: PartType,
     model: str | None,
     media_resolution_enum: dict[str, str] | None,
-    video_metadata: Mapping[str, object] | None,
+    video_metadata: dict[str, Any] | None,
 ) -> PartType:
     """
     Apply media_resolution and video_metadata parameters to a Gemini part.
@@ -637,6 +638,23 @@ def check_if_part_exists_in_parts(parts: list[PartType], part: PartType, exclude
     return False
 
 
+def _get_valid_base64_thought_signature(signature: object) -> str | None:
+    if not isinstance(signature, str):
+        return None
+    stripped: Final = signature.strip()
+    if not stripped:
+        return None
+
+    padded: Final = stripped + "=" * (-len(stripped) % 4)
+
+    try:
+        decoded: Final = base64.b64decode(padded.encode("utf-8"), altchars=b"-_", validate=True)
+    except (binascii.Error, ValueError):
+        return None
+
+    return padded if decoded else None
+
+
 def _collect_tool_call_thought_signatures(
     assistant_msg: ChatCompletionAssistantMessage,
 ) -> frozenset[str]:
@@ -682,8 +700,9 @@ def _collect_tool_call_thought_signatures(
             continue
         for key in ("thought_signature", "response_thought_signature"):
             invocation_signature = invocation.get(key)
-            if isinstance(invocation_signature, str) and invocation_signature:
-                signatures += (invocation_signature,)
+            valid_sig = _get_valid_base64_thought_signature(invocation_signature)
+            if valid_sig:
+                signatures += (valid_sig,)
 
     return frozenset(signatures)
 
@@ -889,19 +908,23 @@ def _gemini_convert_messages_with_history(
                         if block["type"] == "thinking":
                             block_thinking_str = block.get("thinking")
                             block_signature = block.get("signature")
-                            if block_thinking_str is not None and block_signature is not None:
+                            valid_block_sig = _get_valid_base64_thought_signature(block_signature)
+                            if block_thinking_str is not None:
+                                sig_kwargs: dict[str, Any] = (
+                                    {"thoughtSignature": valid_block_sig} if valid_block_sig is not None else {}
+                                )
                                 try:
                                     assistant_content.append(
                                         PartType(
-                                            thoughtSignature=block_signature,
+                                            **sig_kwargs,
                                             **json.loads(block_thinking_str),
                                         )
                                     )
                                 except Exception:
                                     assistant_content.append(
                                         PartType(
-                                            thoughtSignature=block_signature,
                                             text=block_thinking_str,
+                                            **sig_kwargs,
                                         )
                                     )
                 if _message_content is not None and isinstance(_message_content, list):
@@ -927,17 +950,18 @@ def _gemini_convert_messages_with_history(
                     # reasoning token count on gemini-3 and newer models
                     tool_call_signatures = _collect_tool_call_thought_signatures(assistant_msg)
 
-                    if (
-                        thought_signatures
-                        and isinstance(thought_signatures, list)
-                        and len(thought_signatures) > 0
-                        and thought_signatures[0] not in tool_call_signatures
-                    ):
+                    valid_text_signature = (
+                        _get_valid_base64_thought_signature(thought_signatures[0])
+                        if (thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0)
+                        else None
+                    )
+
+                    if valid_text_signature and valid_text_signature not in tool_call_signatures:
                         # Use the first signature for the text part (Gemini expects one signature per part)
                         assistant_content.append(
                             PartType(
                                 text=assistant_text,
-                                thoughtSignature=thought_signatures[0],
+                                thoughtSignature=valid_text_signature,
                             )
                         )
                     else: