fix: Extract and propagate task_id in RemoteA2aAgent

The RemoteA2aAgent now extracts a "task_id" from the custom metadata of the last agent event in the session, alongside the existing "context_id". This task_id is then included in the A2AMessage sent to the remote A2A service.

Close #3765

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 839996774

Author: GWeale

src/google/adk/agents/remote_a2a_agent.py

@@ -343,18 +343,18 @@ def _create_a2a_request_for_user_function_response(
 
   def _construct_message_parts_from_session(
       self, ctx: InvocationContext
-  ) -> tuple[list[A2APart], Optional[str]]:
+  ) -> tuple[list[A2APart], Optional[str], Optional[str]]:
     """Construct A2A message parts from session events.
 
     Args:
       ctx: The invocation context
 
     Returns:
-      List of A2A parts extracted from session events, context ID,
-      request metadata
+      List of A2A parts extracted from session events, context ID, task ID
     """
     message_parts: list[A2APart] = []
     context_id = None
+    task_id = None
 
     events_to_process = []
     for event in reversed(ctx.session.events):
@@ -364,6 +364,7 @@ def _construct_message_parts_from_session(
         if event.custom_metadata:
           metadata = event.custom_metadata
           context_id = metadata.get(A2A_METADATA_PREFIX + "context_id")
+          task_id = metadata.get(A2A_METADATA_PREFIX + "task_id")
         break
       events_to_process.append(event)
 
@@ -384,7 +385,7 @@ def _construct_message_parts_from_session(
         else:
           logger.warning("Failed to convert part to A2A format: %s", part)
 
-    return message_parts, context_id
+    return message_parts, context_id, task_id
 
   async def _handle_a2a_response(
       self, a2a_response: A2AClientEvent | A2AMessage, ctx: InvocationContext
@@ -500,8 +501,8 @@ async def _run_async_impl(
     # Create A2A request for function response or regular message
     a2a_request = self._create_a2a_request_for_user_function_response(ctx)
     if not a2a_request:
-      message_parts, context_id = self._construct_message_parts_from_session(
-          ctx
+      message_parts, context_id, task_id = (
+          self._construct_message_parts_from_session(ctx)
       )
 
       if not message_parts:
@@ -521,6 +522,7 @@ async def _run_async_impl(
           parts=message_parts,
           role="user",
           context_id=context_id,
+          task_id=task_id,
       )
 
     logger.debug(build_a2a_request_log(a2a_request))

tests/unittests/agents/test_remote_a2a_agent.py

@@ -15,6 +15,7 @@
 import json
 from pathlib import Path
 import tempfile
+from unittest import mock
 from unittest.mock import AsyncMock
 from unittest.mock import create_autospec
 from unittest.mock import Mock
@@ -612,13 +613,14 @@ def test_construct_message_parts_from_session_success(self):
       mock_a2a_part = Mock()
       self.mock_genai_part_converter.return_value = mock_a2a_part
 
-      parts, context_id = self.agent._construct_message_parts_from_session(
-          self.mock_context
+      parts, context_id, task_id = (
+          self.agent._construct_message_parts_from_session(self.mock_context)
       )
 
       assert len(parts) == 1
       assert parts[0] == mock_a2a_part
       assert context_id is None
+      assert task_id is None
 
   def test_construct_message_parts_from_session_success_multiple_parts(self):
     """Test successful message parts construction from session."""
@@ -646,23 +648,54 @@ def test_construct_message_parts_from_session_success_multiple_parts(self):
           mock_a2a_part2,
       ]
 
-      parts, context_id = self.agent._construct_message_parts_from_session(
-          self.mock_context
+      parts, context_id, task_id = (
+          self.agent._construct_message_parts_from_session(self.mock_context)
       )
 
       assert parts == [mock_a2a_part1, mock_a2a_part2]
       assert context_id is None
+      assert task_id is None
 
   def test_construct_message_parts_from_session_empty_events(self):
     """Test message parts construction with empty events."""
     self.mock_session.events = []
 
-    parts, context_id = self.agent._construct_message_parts_from_session(
-        self.mock_context
+    parts, context_id, task_id = (
+        self.agent._construct_message_parts_from_session(self.mock_context)
     )
 
     assert parts == []
     assert context_id is None
+    assert task_id is None
+
+  def test_construct_message_parts_from_session_reads_ids_from_metadata(self):
+    """Metadata from last agent event is reused for context and task IDs."""
+    mock_part = Mock()
+    mock_part.text = "User message"
+    mock_content = Mock()
+    mock_content.parts = [mock_part]
+    user_event = Mock()
+    user_event.content = mock_content
+    user_event.author = "user"
+
+    agent_event = Mock()
+    agent_event.author = self.agent.name
+    agent_event.custom_metadata = {
+        A2A_METADATA_PREFIX + "context_id": "context-xyz",
+        A2A_METADATA_PREFIX + "task_id": "task-abc",
+    }
+
+    # Agent reply is before the latest user message (chronological order).
+    self.mock_session.events = [agent_event, user_event]
+    self.mock_genai_part_converter.return_value = Mock()
+
+    parts, context_id, task_id = (
+        self.agent._construct_message_parts_from_session(self.mock_context)
+    )
+
+    assert len(parts) == 1  # the latest user message
+    assert context_id == "context-xyz"
+    assert task_id == "task-abc"
 
   @pytest.mark.asyncio
   async def test_handle_a2a_response_success_with_message(self):
@@ -786,13 +819,14 @@ def mock_converter(part):
 
       self.mock_genai_part_converter.side_effect = mock_converter
 
-      parts, context_id = self.agent._construct_message_parts_from_session(
-          self.mock_context
+      parts, context_id, task_id = (
+          self.agent._construct_message_parts_from_session(self.mock_context)
       )
 
       # Verify the parts are in correct order
       assert len(parts) == 3  # 1 user part + 2 other agent parts
       assert context_id is None
+      assert task_id is None
 
       # Verify order: user part, then "For context:", then agent message
       assert converted_parts[0].original_text == "User question"
@@ -1109,24 +1143,14 @@ def test_construct_message_parts_from_session_success(self):
         mock_a2a_part = Mock()
         mock_convert_part.return_value = mock_a2a_part
 
-        parts, context_id = self.agent._construct_message_parts_from_session(
-            self.mock_context
+        parts, context_id, task_id = (
+            self.agent._construct_message_parts_from_session(self.mock_context)
         )
 
         assert len(parts) == 1
         assert parts[0] == mock_a2a_part
         assert context_id is None
-
-  def test_construct_message_parts_from_session_empty_events(self):
-    """Test message parts construction with empty events."""
-    self.mock_session.events = []
-
-    parts, context_id = self.agent._construct_message_parts_from_session(
-        self.mock_context
-    )
-
-    assert parts == []
-    assert context_id is None
+        assert task_id is None
 
   @pytest.mark.asyncio
   async def test_handle_a2a_response_success_with_message(self):
@@ -1463,7 +1487,8 @@ async def test_run_async_impl_no_message_parts(self):
           mock_construct.return_value = (
               [],
               None,
-          )  # Tuple with empty parts and no context_id
+              None,
+          )  # Tuple with empty parts and no context/task ids
 
           events = []
           async for event in self.agent._run_async_impl(self.mock_context):
@@ -1493,7 +1518,8 @@ async def test_run_async_impl_successful_request(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-          )  # Tuple with parts and context_id
+              "task-789",
+          )  # Tuple with parts and context/task ids
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1545,6 +1571,13 @@ async def test_run_async_impl_successful_request(self):
                       A2A_METADATA_PREFIX + "request"
                       in mock_event.custom_metadata
                   )
+                  mock_message_class.assert_called_once_with(
+                      message_id=mock.ANY,
+                      parts=[mock_a2a_part],
+                      role="user",
+                      context_id="context-123",
+                      task_id="task-789",
+                  )
 
   @pytest.mark.asyncio
   async def test_run_async_impl_a2a_client_error(self):
@@ -1565,7 +1598,8 @@ async def test_run_async_impl_a2a_client_error(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-          )  # Tuple with parts and context_id
+              "task-789",
+          )  # Tuple with parts and context/task ids
 
           # Mock A2A client that throws an exception
           mock_a2a_client = AsyncMock()
@@ -1632,7 +1666,8 @@ async def test_run_async_impl_with_meta_provider(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-          )  # Tuple with parts and context_id
+              "task-789",
+          )  # Tuple with parts and context/task ids
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1683,6 +1718,13 @@ async def test_run_async_impl_with_meta_provider(self):
                       request=mock_message,
                       request_metadata=request_metadata,
                   )
+                  mock_message_class.assert_called_once_with(
+                      message_id=mock.ANY,
+                      parts=[mock_a2a_part],
+                      role="user",
+                      context_id="context-123",
+                      task_id="task-789",
+                  )
 
 
 class TestRemoteA2aAgentExecutionFromFactory:
@@ -1737,7 +1779,8 @@ async def test_run_async_impl_no_message_parts(self):
           mock_construct.return_value = (
               [],
               None,
-          )  # Tuple with empty parts and no context_id
+              None,
+          )  # Tuple with empty parts and no context/task ids
 
           events = []
           async for event in self.agent._run_async_impl(self.mock_context):
@@ -1767,7 +1810,8 @@ async def test_run_async_impl_successful_request(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-          )  # Tuple with parts and context_id
+              "task-789",
+          )  # Tuple with parts and context/task ids
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1821,6 +1865,13 @@ async def test_run_async_impl_successful_request(self):
                       A2A_METADATA_PREFIX + "request"
                       in mock_event.custom_metadata
                   )
+                  mock_message_class.assert_called_once_with(
+                      message_id=mock.ANY,
+                      parts=[mock_a2a_part],
+                      role="user",
+                      context_id="context-123",
+                      task_id="task-789",
+                  )
 
   @pytest.mark.asyncio
   async def test_run_async_impl_a2a_client_error(self):
@@ -1841,7 +1892,8 @@ async def test_run_async_impl_a2a_client_error(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-          )  # Tuple with parts and context_id
+              "task-789",
+          )  # Tuple with parts and context/task ids
 
           # Mock A2A client that throws an exception
           mock_a2a_client = AsyncMock()