fix: Preserve thought_signature in FunctionCall conversions between GenAI and A2A

Close: #4311

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 877465519

Author: xuanyang15

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,
         )
     )
 

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 @@ def test_convert_file_part_with_bytes(self):
     # 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 @@ def test_convert_inline_data_part(self):
     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 @@ def test_convert_a2a_data_part_with_executable_code_metadata(self):
     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