From f9c104faf73e2a002bb3092b50fb88f4eed78163 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Mon, 2 Mar 2026 10:28:46 -0800 Subject: [PATCH] fix: Preserve thought_signature in FunctionCall conversions between GenAI and A2A Close: https://github.com/google/adk-python/issues/4311 Co-authored-by: Xuan Yang PiperOrigin-RevId: 877465519 --- .../adk/a2a/converters/part_converter.py | 33 ++- .../a2a/converters/test_part_converter.py | 204 +++++++++++++++++- 2 files changed, 229 insertions(+), 8 deletions(-) diff --git a/src/google/adk/a2a/converters/part_converter.py b/src/google/adk/a2a/converters/part_converter.py index 7b501f75..ce65a3de 100644 --- a/src/google/adk/a2a/converters/part_converter.py +++ b/src/google/adk/a2a/converters/part_converter.py @@ -104,10 +104,25 @@ def convert_a2a_part_to_genai_part( part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)] == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL ): + # Restore thought_signature if present + thought_signature = None + thought_sig_key = _get_adk_metadata_key('thought_signature') + if thought_sig_key in part.metadata: + sig_value = part.metadata[thought_sig_key] + if isinstance(sig_value, bytes): + thought_signature = sig_value + elif isinstance(sig_value, str): + try: + thought_signature = base64.b64decode(sig_value) + except Exception: + logger.warning( + 'Failed to decode thought_signature: %s', sig_value + ) return genai_types.Part( function_call=genai_types.FunctionCall.model_validate( part.data, by_alias=True - ) + ), + thought_signature=thought_signature, ) if ( part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)] @@ -214,16 +229,22 @@ def convert_genai_part_to_a2a_part( # TODO once A2A defined how to service such information, migrate below # logic accordingly if part.function_call: + fc_metadata = { + _get_adk_metadata_key( + A2A_DATA_PART_METADATA_TYPE_KEY + ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL + } + # Preserve thought_signature if present + if part.thought_signature is not None: + fc_metadata[_get_adk_metadata_key('thought_signature')] = ( + base64.b64encode(part.thought_signature).decode('utf-8') + ) return a2a_types.Part( root=a2a_types.DataPart( data=part.function_call.model_dump( by_alias=True, exclude_none=True ), - metadata={ - _get_adk_metadata_key( - A2A_DATA_PART_METADATA_TYPE_KEY - ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL - }, + metadata=fc_metadata, ) ) diff --git a/tests/unittests/a2a/converters/test_part_converter.py b/tests/unittests/a2a/converters/test_part_converter.py index ec611fba..057b6c9e 100644 --- a/tests/unittests/a2a/converters/test_part_converter.py +++ b/tests/unittests/a2a/converters/test_part_converter.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 import json from unittest.mock import Mock from unittest.mock import patch @@ -74,7 +75,6 @@ class TestConvertA2aPartToGenaiPart: # Arrange test_bytes = b"test file content" # A2A FileWithBytes expects base64-encoded string - import base64 base64_encoded = base64.b64encode(test_bytes).decode("utf-8") a2a_part = a2a_types.Part( @@ -328,7 +328,6 @@ class TestConvertGenaiPartToA2aPart: assert isinstance(result.root, a2a_types.FilePart) assert isinstance(result.root.file, a2a_types.FileWithBytes) # A2A FileWithBytes now stores base64-encoded bytes to ensure round-trip compatibility - import base64 expected_base64 = base64.b64encode(test_bytes).decode("utf-8") assert result.root.file.bytes == expected_base64 @@ -841,3 +840,204 @@ class TestNewConstants: assert result.executable_code is not None assert result.executable_code.language == genai_types.Language.PYTHON assert result.executable_code.code == "print('Hello, World!')" + + +class TestThoughtSignaturePreservation: + """Tests for thought_signature preservation in function call conversions.""" + + def test_genai_function_call_with_thought_signature_to_a2a(self): + """Test that thought_signature is preserved when converting GenAI to A2A.""" + # Arrange + function_call = genai_types.FunctionCall( + id="fc_gemini3", + name="my_tool", + args={"document": "test content"}, + ) + genai_part = genai_types.Part( + function_call=function_call, + thought_signature=b"gemini3_signature_bytes", + ) + + # Act + result = convert_genai_part_to_a2a_part(genai_part) + + # Assert + assert result is not None + assert isinstance(result.root, a2a_types.DataPart) + assert ( + result.root.metadata[ + _get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY) + ] + == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL + ) + # thought_signature should be base64 encoded in metadata + thought_sig_key = _get_adk_metadata_key("thought_signature") + assert thought_sig_key in result.root.metadata + assert ( + base64.b64decode(result.root.metadata[thought_sig_key]) + == b"gemini3_signature_bytes" + ) + + def test_genai_function_call_without_thought_signature_to_a2a(self): + """Test function call without thought_signature doesn't add metadata key.""" + # Arrange + function_call = genai_types.FunctionCall( + id="fc_regular", + name="regular_tool", + args={}, + ) + genai_part = genai_types.Part(function_call=function_call) + + # Act + result = convert_genai_part_to_a2a_part(genai_part) + + # Assert + assert result is not None + assert isinstance(result.root, a2a_types.DataPart) + # thought_signature key should not be present + thought_sig_key = _get_adk_metadata_key("thought_signature") + assert thought_sig_key not in result.root.metadata + + def test_a2a_function_call_with_thought_signature_to_genai(self): + """Test that thought_signature is restored when converting A2A to GenAI.""" + # Arrange + a2a_part = a2a_types.Part( + root=a2a_types.DataPart( + data={ + "id": "fc_gemini3", + "name": "my_tool", + "args": {"document": "test content"}, + }, + metadata={ + _get_adk_metadata_key( + A2A_DATA_PART_METADATA_TYPE_KEY + ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL, + _get_adk_metadata_key("thought_signature"): ( + base64.b64encode(b"restored_signature").decode("utf-8") + ), + }, + ) + ) + + # Act + result = convert_a2a_part_to_genai_part(a2a_part) + + # Assert + assert result is not None + assert result.function_call is not None + assert result.function_call.name == "my_tool" + # thought_signature should be decoded back to bytes + assert result.thought_signature == b"restored_signature" + + def test_a2a_function_call_without_thought_signature_to_genai(self): + """Test function call without thought_signature returns None for it.""" + # Arrange + a2a_part = a2a_types.Part( + root=a2a_types.DataPart( + data={ + "id": "fc_regular", + "name": "regular_tool", + "args": {}, + }, + metadata={ + _get_adk_metadata_key( + A2A_DATA_PART_METADATA_TYPE_KEY + ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL, + }, + ) + ) + + # Act + result = convert_a2a_part_to_genai_part(a2a_part) + + # Assert + assert result is not None + assert result.function_call is not None + assert result.function_call.name == "regular_tool" + # thought_signature should be None + assert result.thought_signature is None + + def test_function_call_with_thought_signature_round_trip(self): + """Test thought_signature is preserved in GenAI -> A2A -> GenAI round trip.""" + # Arrange + original_signature = b"round_trip_signature_test" + function_call = genai_types.FunctionCall( + id="fc_round_trip", + name="round_trip_tool", + args={"key": "value"}, + ) + original_part = genai_types.Part( + function_call=function_call, + thought_signature=original_signature, + ) + + # Act - Convert GenAI -> A2A -> GenAI + a2a_part = convert_genai_part_to_a2a_part(original_part) + restored_part = convert_a2a_part_to_genai_part(a2a_part) + + # Assert + assert restored_part is not None + assert restored_part.function_call is not None + assert restored_part.function_call.name == "round_trip_tool" + assert restored_part.thought_signature == original_signature + + def test_a2a_function_call_with_bytes_thought_signature_to_genai(self): + """Test that bytes thought_signature is used directly without decoding.""" + # Arrange - metadata contains raw bytes (not base64 encoded) + a2a_part = a2a_types.Part( + root=a2a_types.DataPart( + data={ + "id": "fc_bytes", + "name": "bytes_tool", + "args": {}, + }, + metadata={ + _get_adk_metadata_key( + A2A_DATA_PART_METADATA_TYPE_KEY + ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL, + _get_adk_metadata_key( + "thought_signature" + ): b"raw_bytes_signature", + }, + ) + ) + + # Act + result = convert_a2a_part_to_genai_part(a2a_part) + + # Assert + assert result is not None + assert result.function_call is not None + # bytes should be used directly + assert result.thought_signature == b"raw_bytes_signature" + + def test_a2a_function_call_with_invalid_base64_thought_signature(self): + """Test that invalid base64 thought_signature logs warning and returns None.""" + # Arrange - metadata contains invalid base64 string + a2a_part = a2a_types.Part( + root=a2a_types.DataPart( + data={ + "id": "fc_invalid", + "name": "invalid_sig_tool", + "args": {}, + }, + metadata={ + _get_adk_metadata_key( + A2A_DATA_PART_METADATA_TYPE_KEY + ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL, + _get_adk_metadata_key( + "thought_signature" + ): "not_valid_base64!!!", + }, + ) + ) + + # Act + result = convert_a2a_part_to_genai_part(a2a_part) + + # Assert + assert result is not None + assert result.function_call is not None + assert result.function_call.name == "invalid_sig_tool" + # thought_signature should be None due to decode failure + assert result.thought_signature is None