From 3e6f11e1a05579a24ca33edfe7a40a8f98d33c47 Mon Sep 17 00:00:00 2001 From: saakshigupta2002 Date: Thu, 23 Jul 2026 22:37:26 +0930 Subject: [PATCH 1/2] fix(litellm): strip embedded thought_signature from tool call id When a Gemini thinking model exposes its thought_signature only embedded in the tool call id via the __thought__ separator, _message_to_generate_content_response extracted the signature onto part.thought_signature but assigned the full, unsplit id (including the __thought__ suffix) to part.function_call.id. That corrupted id is persisted to session history and, on later turns, replayed to the model as a literal tool_call_id, which OpenAI-compatible endpoints reject (422 'Tool Call ID required on tool calls'). Because the malformed id is now part of the event history, every subsequent turn on the thread fails identically, permanently breaking the conversation. Strip the embedded suffix so function_call.id is the clean, original id, while still preserving the signature separately on part.thought_signature. The extra_content and provider_specific_fields paths are unchanged. Adds regression and unit tests. Fixes #6454 --- src/google/adk/models/lite_llm.py | 34 +++++++++++++- tests/unittests/models/test_litellm.py | 62 ++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index def9aba4cf6..cc6340b9ff4 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -987,6 +987,33 @@ def _extract_thought_signature_from_tool_call( if len(parts) == 2: return _decode_thought_signature(parts[1]) + +def _strip_thought_signature_from_tool_call_id( + tool_call_id: Optional[str], +) -> Optional[str]: + """Strips an embedded thought_signature suffix from a tool call ID. + + Some Gemini thinking model paths embed the thought_signature in the tool + call ID via the ``__thought__`` separator (see + ``_extract_thought_signature_from_tool_call``). The signature is surfaced + separately on ``part.thought_signature``, so the ID assigned to + ``function_call.id`` must be the clean, original ID without the suffix. + Otherwise the corrupted ID is persisted to session history and later + replayed to the model as a ``tool_call_id``, which downstream + OpenAI-compatible endpoints reject. + + Args: + tool_call_id: The raw tool call ID, which may contain an embedded + thought_signature. + + Returns: + The tool call ID with any ``__thought__`` suffix removed, or the value + unchanged when no separator is present. + """ + if not tool_call_id: + return tool_call_id + return tool_call_id.split(_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + return None @@ -2182,7 +2209,12 @@ def _message_to_generate_content_response( name=tool_call.function.name, args=_parse_tool_call_arguments(tool_call.function.arguments), ) - part.function_call.id = tool_call.id + # Strip any embedded thought_signature suffix so the persisted + # function_call.id stays the clean, original tool call ID. The + # signature is preserved separately on part.thought_signature. + part.function_call.id = _strip_thought_signature_from_tool_call_id( + tool_call.id + ) if thought_signature: part.thought_signature = thought_signature parts.append(part) diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 3c59b01b8a0..7e5795c0d09 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -57,6 +57,7 @@ from google.adk.models.lite_llm import _safe_json_serialize from google.adk.models.lite_llm import _schema_to_dict from google.adk.models.lite_llm import _split_message_content_and_tool_calls +from google.adk.models.lite_llm import _strip_thought_signature_from_tool_call_id from google.adk.models.lite_llm import _THOUGHT_SIGNATURE_SEPARATOR from google.adk.models.lite_llm import _to_litellm_response_format from google.adk.models.lite_llm import _to_litellm_role @@ -2937,6 +2938,67 @@ def test_message_to_generate_content_response_no_thought_signature(): assert fc_part.thought_signature is None +def test_message_to_generate_content_response_strips_thought_signature_from_id(): + """function_call.id is cleaned when the signature is embedded in the ID. + + Regression test: when the thought_signature is only available embedded in + the tool call ID (the __thought__ fallback), the original code assigned the + entire unsplit ID to function_call.id, corrupting session history and + breaking every subsequent turn once the malformed ID is replayed as a + tool_call_id. + """ + sig_b64 = base64.b64encode(b"embedded_sig").decode("utf-8") + embedded_id = f"call_789{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}" + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=embedded_id, + function=Function( + name="test_function", + arguments='{"test_arg": "test_value"}', + ), + ) + ], + ) + + response = _message_to_generate_content_response(message) + fc_part = response.content.parts[0] + # The ID is stripped back to the clean, original tool call ID ... + assert fc_part.function_call.id == "call_789" + assert _THOUGHT_SIGNATURE_SEPARATOR not in fc_part.function_call.id + # ... and the signature is still preserved separately. + assert fc_part.thought_signature == b"embedded_sig" + + +def test_strip_thought_signature_from_tool_call_id_removes_suffix(): + """The embedded __thought__ suffix is removed from the ID.""" + sig_b64 = base64.b64encode(b"embedded_sig").decode("utf-8") + tool_call_id = f"call_789{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}" + assert _strip_thought_signature_from_tool_call_id(tool_call_id) == "call_789" + + +def test_strip_thought_signature_from_tool_call_id_leaves_plain_id(): + """A plain ID with no separator is returned unchanged.""" + assert _strip_thought_signature_from_tool_call_id("call_plain") == "call_plain" + + +@pytest.mark.parametrize("value", [None, ""]) +def test_strip_thought_signature_from_tool_call_id_handles_empty(value): + """None and empty IDs are returned unchanged.""" + assert _strip_thought_signature_from_tool_call_id(value) == value + + +def test_strip_thought_signature_from_tool_call_id_splits_once(): + """Only the first separator is used, mirroring signature extraction.""" + tool_call_id = ( + f"call_1{_THOUGHT_SIGNATURE_SEPARATOR}sig{_THOUGHT_SIGNATURE_SEPARATOR}x" + ) + assert _strip_thought_signature_from_tool_call_id(tool_call_id) == "call_1" + + @pytest.mark.asyncio async def test_content_to_message_param_preserves_thought_signature(): """thought_signature on Part is emitted on both tool call metadata paths.""" From 7c6904abe6c9922e7fe3cd8afd61a004469e525c Mon Sep 17 00:00:00 2001 From: saakshigupta2002 Date: Fri, 24 Jul 2026 21:48:38 +0930 Subject: [PATCH 2/2] fix(litellm): only strip tool call id suffix when it decodes as a signature Address review feedback: _strip_thought_signature_from_tool_call_id previously truncated any id containing the __thought__ separator, even when the suffix was not a valid base64 signature. Since the separator can legitimately appear inside an opaque tool call id, that mutated ids from which no signature was actually extracted (e.g. 'call_keep__thought__not-base64' became 'call_keep' while thought_signature stayed None), contradicting the Gemini requirement to echo back the exact function-call id. Now the suffix is decoded first (mirroring _extract_thought_signature_from_tool_call); the id is only stripped when the suffix is a genuine signature, otherwise it is returned unchanged. Adds malformed-suffix regression tests at both the helper and _message_to_generate_content_response levels. --- src/google/adk/models/lite_llm.py | 21 ++++++++++--- tests/unittests/models/test_litellm.py | 41 ++++++++++++++++++++++---- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index cc6340b9ff4..6beb4c2d815 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -1002,17 +1002,30 @@ def _strip_thought_signature_from_tool_call_id( replayed to the model as a ``tool_call_id``, which downstream OpenAI-compatible endpoints reject. + The suffix is only stripped when it actually decodes as a thought_signature, + mirroring ``_extract_thought_signature_from_tool_call``. If the text after the + separator does not decode, the separator is treated as part of the opaque + tool call ID and the ID is returned unchanged, so an ID is never mutated + unless a signature was genuinely extracted from it. This matches the Gemini + requirement to echo back the exact function-call ID. + Args: tool_call_id: The raw tool call ID, which may contain an embedded thought_signature. Returns: - The tool call ID with any ``__thought__`` suffix removed, or the value - unchanged when no separator is present. + The tool call ID with the ``__thought__`` suffix removed when that suffix + is a decodable signature; otherwise the value unchanged. """ - if not tool_call_id: + if not tool_call_id or _THOUGHT_SIGNATURE_SEPARATOR not in tool_call_id: + return tool_call_id + base_id, _, encoded_signature = tool_call_id.partition( + _THOUGHT_SIGNATURE_SEPARATOR + ) + if _decode_thought_signature(encoded_signature) is None: + # The suffix is not a valid signature, so it belongs to the opaque ID. return tool_call_id - return tool_call_id.split(_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + return base_id return None diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 7e5795c0d09..a1e47c8446c 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -2991,12 +2991,43 @@ def test_strip_thought_signature_from_tool_call_id_handles_empty(value): assert _strip_thought_signature_from_tool_call_id(value) == value -def test_strip_thought_signature_from_tool_call_id_splits_once(): - """Only the first separator is used, mirroring signature extraction.""" - tool_call_id = ( - f"call_1{_THOUGHT_SIGNATURE_SEPARATOR}sig{_THOUGHT_SIGNATURE_SEPARATOR}x" +def test_strip_thought_signature_from_tool_call_id_retains_undecodable_suffix(): + """An ID whose __thought__ suffix is not a valid signature is untouched. + + The separator can legitimately appear inside an opaque tool call ID. When the + text after it does not decode as a signature, no signature was extracted, so + the ID must be preserved exactly rather than truncated. + """ + tool_call_id = f"call_keep{_THOUGHT_SIGNATURE_SEPARATOR}not-base64" + assert ( + _strip_thought_signature_from_tool_call_id(tool_call_id) == tool_call_id + ) + + +def test_message_to_generate_content_response_retains_id_when_suffix_undecodable(): + """function_call.id is preserved when the embedded suffix is not a signature. + + Regression test for the review feedback: when the __thought__ suffix does not + decode, no signature is extracted (thought_signature stays None) and the + opaque ID must not be mutated before it is persisted. + """ + undecodable_id = f"call_keep{_THOUGHT_SIGNATURE_SEPARATOR}not-base64" + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=undecodable_id, + function=Function(name="test_function", arguments="{}"), + ) + ], ) - assert _strip_thought_signature_from_tool_call_id(tool_call_id) == "call_1" + + response = _message_to_generate_content_response(message) + fc_part = response.content.parts[0] + assert fc_part.function_call.id == undecodable_id + assert fc_part.thought_signature is None @pytest.mark.asyncio