diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index def9aba4cf..6beb4c2d81 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -987,6 +987,46 @@ 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. + + 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 the ``__thought__`` suffix removed when that suffix + is a decodable signature; otherwise the value unchanged. + """ + 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 base_id + return None @@ -2182,7 +2222,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 3c59b01b8a..a1e47c8446 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,98 @@ 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_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="{}"), + ) + ], + ) + + 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 async def test_content_to_message_param_preserves_thought_signature(): """thought_signature on Part is emitted on both tool call metadata paths."""