From 91789da29662e11f81622bd8d2abd0c4876bca20 Mon Sep 17 00:00:00 2001 From: Baojiang Lee <101562714+libaojiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:19:26 +0800 Subject: [PATCH] fix(vertex_ai): verify embedded thought signature checksum --- .../prompt_templates/factory.py | 52 +++++++++++-------- .../test_thought_signature_in_tool_call_id.py | 29 +++++++---- 2 files changed, 49 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 5f3568c64c2..bddd0f7613e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/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___thought__ + Format: call___checksum____thought__ 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 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 489564c3e0c..964f5e0e673 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 @@ -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