From b80f3054301bcf203205be7e95350c5e5e25f9bd Mon Sep 17 00:00:00 2001 From: Baojiang Lee <101562714+libaojiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:19:25 +0800 Subject: [PATCH] fix(vertex_ai): reject client-mangled thought_signature embedded in tool_call_id A Gemini `thoughtSignature` is standard-alphabet base64 bytes, but LiteLLM appends it to `tool_calls[].id` after a `__thought__` marker. Any client that normalizes ids (Claude Code, OpenClaw, editor integrations that shorten + hash the id) mangles the base64 and replays it verbatim on the next turn. `_get_thought_signature_from_tool` used to extract that mangled tail as-is, so `convert_to_gemini_tool_call_invoke` forwarded it to Vertex, which rejects the request with `Base64 decoding failed for "..."` or `Invalid thought signature`. Once a corrupted value entered the history, every following turn failed the same way and the session was dead. Validate the extracted tail against Vertex's strict base64 decoder (padding tolerant) and return None when it fails. Callers already handle the "no signature" case correctly: on Gemini 3+ the first function call falls back to the documented skip-validator dummy signature; on Gemini 2.x the signature is simply omitted. Either way the hard 400 becomes a graceful degradation, and clients that keep the signature in `provider_specific_fields` (LiteLLM SDK) are unaffected because that path is checked first. Fixes #37849 --- .../prompt_templates/factory.py | 34 ++++- .../test_thought_signature_in_tool_call_id.py | 142 ++++++++++++++++-- 2 files changed, 162 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 795fb36961e..536ec175c16 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1,4 +1,5 @@ import base64 +import binascii import copy import hashlib import json @@ -1158,6 +1159,29 @@ def _encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: st 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. + """ + if not signature: + return False + try: + padding = "=" * (-len(signature) % 4) + base64.b64decode(signature + padding, validate=True) + return True + except (binascii.Error, ValueError): + return False + + def _get_thought_signature_from_tool(tool: dict) -> str | None: """Extract thought signature from tool call's provider_specific_fields. @@ -1188,13 +1212,19 @@ def _get_thought_signature_from_tool(tool: dict) -> str | None: signature = function.provider_specific_fields.get("thought_signature") if signature: return signature - # Check if thought signature is embedded in tool call ID + # Check if thought signature is embedded in tool call ID. + # Any client that normalizes the id shortens/hashes it, so the segment + # after ``__thought__`` may no longer base64-decode. Reject those here so + # the caller can decide between "no signature" and the dummy fallback, + # instead of forwarding a corrupted value that would 400 on Vertex. See + # issue #37849. 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 - return signature + if _is_valid_thought_signature(signature): + return signature return None diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py index 208cba519f3..de053a11f9e 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py @@ -9,13 +9,17 @@ Note: Embedding signatures in tool call IDs is a beta feature that requires enable_preview_features=True to be enabled. """ +import base64 + 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, + _is_valid_thought_signature, convert_to_gemini_tool_call_invoke, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -94,10 +98,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 @@ -234,9 +235,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 +246,10 @@ 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" + # Real Gemini signatures are standard-alphabet base64; using a valid + # placeholder keeps the parallel-tool-call path lit up under the + # base64-validated extraction added for issue #37849. + signature1 = base64.b64encode(b"signature_for_first_call").decode("ascii") # Only first call has signature (Gemini behavior for parallel calls) gemini_parts = [ @@ -271,10 +273,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"] @@ -285,3 +284,122 @@ def test_parallel_tool_calls_with_signatures(enable_preview_features): assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"] sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"}) assert sig2 is None + + +REAL_SIGNATURE = ( + "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdT" + "tfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed" + "0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7C" + "vykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" +) + + +# Signature values a client-side id sanitizer produces after stripping the raw +# base64 down to id-shaped characters. Standard-alphabet base64 has no ``_`` or +# trailing hex fragment, so Vertex rejects these with either "Base64 decoding +# failed" or "Invalid thought signature". See issue #37849. +@pytest.mark.parametrize( + "mangled_signature", + [ + # from the issue reproducer, /v1/responses turn 2 + "AY89a1/_57b05e78dc", + # from the issue reproducer, /v1/chat/completions turn 2 + "AY89a1/_ee781c9832", + # short SHA-suffix-only sanitization + "AY89a1_S3YvIpUCcBTFSgDfesRLDnA_775ff49bcd", + ], +) +def test_is_valid_thought_signature_rejects_client_normalized_values(mangled_signature): + assert _is_valid_thought_signature(REAL_SIGNATURE) is True + assert _is_valid_thought_signature(mangled_signature) is False + + +def test_is_valid_thought_signature_tolerates_missing_padding(): + # Gemini omits ``=`` on the wire; the strict decoder must still accept the + # value once we re-pad it. + encoded = base64.b64encode(b"hello").decode("ascii").rstrip("=") + assert "=" not in encoded + assert _is_valid_thought_signature(encoded) is True + + +def test_get_thought_signature_drops_client_mangled_id_suffix(): + """When the segment after ``__thought__`` isn't valid base64, the extractor + must return ``None`` so the caller can either fall back to the dummy + 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" + tool = {"id": mangled_id, "type": "function"} + + assert _get_thought_signature_from_tool(tool) is None + + +def test_get_thought_signature_still_prefers_provider_fields_even_when_id_mangled(): + """A valid signature in ``provider_specific_fields`` must win over the + mangled tail so we don't downgrade a good signal.""" + mangled_id = f"call_abc{THOUGHT_SIGNATURE_SEPARATOR}not_base64!!" + tool = { + "id": mangled_id, + "type": "function", + "function": {"name": "get_temperature", "arguments": '{"location": "Paris"}'}, + "provider_specific_fields": {"thought_signature": REAL_SIGNATURE}, + } + + assert _get_thought_signature_from_tool(tool) == REAL_SIGNATURE + + +def test_convert_to_gemini_uses_dummy_signature_when_client_mangles_id_on_gemini_3(): + """On Gemini 3+ a mangled id must degrade to the documented skip-validator + dummy signature rather than forwarding the corrupted bytes.""" + mangled_id = f"call_2156408{THOUGHT_SIGNATURE_SEPARATOR}AY89a1/_57b05e78dc" + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": mangled_id, + "type": "function", + "function": { + "name": "get_temperature", + "arguments": '{"location": "Paris"}', + }, + } + ], + } + + parts = convert_to_gemini_tool_call_invoke( + assistant_message, + model="vertex_ai/gemini-3.1-pro-preview", + ) + + assert len(parts) == 1 + assert parts[0].get("thoughtSignature") == _get_dummy_thought_signature() + + +def test_convert_to_gemini_drops_signature_when_client_mangles_id_on_gemini_2(): + """On older Gemini models the placeholder fallback doesn't apply, so a + mangled id must simply drop the signature (rather than forward a value + Vertex will 400 on).""" + mangled_id = f"call_2158562{THOUGHT_SIGNATURE_SEPARATOR}AY89a1/_ee781c9832" + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": mangled_id, + "type": "function", + "function": { + "name": "get_temperature", + "arguments": '{"location": "Paris"}', + }, + } + ], + } + + parts = convert_to_gemini_tool_call_invoke( + assistant_message, + model="vertex_ai/gemini-2.5-pro", + ) + + assert len(parts) == 1 + assert "thoughtSignature" not in parts[0]