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

PiperOrigin-RevId: 840375992

Author: google-genai-bot

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], Optional[str]]:
+  ) -> tuple[list[A2APart], 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, task ID
+      List of A2A parts extracted from session events, context ID,
+      request metadata
     """
     message_parts: list[A2APart] = []
     context_id = None
-    task_id = None
 
     events_to_process = []
     for event in reversed(ctx.session.events):
@@ -364,7 +364,6 @@ 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)
 
@@ -385,7 +384,7 @@ def _construct_message_parts_from_session(
         else:
           logger.warning("Failed to convert part to A2A format: %s", part)
 
-    return message_parts, context_id, task_id
+    return message_parts, context_id
 
   async def _handle_a2a_response(
       self, a2a_response: A2AClientEvent | A2AMessage, ctx: InvocationContext
@@ -501,8 +500,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, task_id = (
-          self._construct_message_parts_from_session(ctx)
+      message_parts, context_id = self._construct_message_parts_from_session(
+          ctx
       )
 
       if not message_parts:
@@ -522,7 +521,6 @@ 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,7 +15,6 @@
 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
@@ -613,14 +612,13 @@ 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, task_id = (
-          self.agent._construct_message_parts_from_session(self.mock_context)
+      parts, context_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."""
@@ -648,54 +646,23 @@ def test_construct_message_parts_from_session_success_multiple_parts(self):
           mock_a2a_part2,
       ]
 
-      parts, context_id, task_id = (
-          self.agent._construct_message_parts_from_session(self.mock_context)
+      parts, context_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, task_id = (
-        self.agent._construct_message_parts_from_session(self.mock_context)
+    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
-
-  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):
@@ -819,14 +786,13 @@ def mock_converter(part):
 
       self.mock_genai_part_converter.side_effect = mock_converter
 
-      parts, context_id, task_id = (
-          self.agent._construct_message_parts_from_session(self.mock_context)
+      parts, context_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"
@@ -1143,14 +1109,24 @@ def test_construct_message_parts_from_session_success(self):
         mock_a2a_part = Mock()
         mock_convert_part.return_value = mock_a2a_part
 
-        parts, context_id, task_id = (
-            self.agent._construct_message_parts_from_session(self.mock_context)
+        parts, context_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_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
 
   @pytest.mark.asyncio
   async def test_handle_a2a_response_success_with_message(self):
@@ -1487,8 +1463,7 @@ async def test_run_async_impl_no_message_parts(self):
           mock_construct.return_value = (
               [],
               None,
-              None,
-          )  # Tuple with empty parts and no context/task ids
+          )  # Tuple with empty parts and no context_id
 
           events = []
           async for event in self.agent._run_async_impl(self.mock_context):
@@ -1518,8 +1493,7 @@ async def test_run_async_impl_successful_request(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-              "task-789",
-          )  # Tuple with parts and context/task ids
+          )  # Tuple with parts and context_id
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1571,13 +1545,6 @@ 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):
@@ -1598,8 +1565,7 @@ async def test_run_async_impl_a2a_client_error(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-              "task-789",
-          )  # Tuple with parts and context/task ids
+          )  # Tuple with parts and context_id
 
           # Mock A2A client that throws an exception
           mock_a2a_client = AsyncMock()
@@ -1666,8 +1632,7 @@ async def test_run_async_impl_with_meta_provider(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-              "task-789",
-          )  # Tuple with parts and context/task ids
+          )  # Tuple with parts and context_id
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1718,13 +1683,6 @@ 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:
@@ -1779,8 +1737,7 @@ async def test_run_async_impl_no_message_parts(self):
           mock_construct.return_value = (
               [],
               None,
-              None,
-          )  # Tuple with empty parts and no context/task ids
+          )  # Tuple with empty parts and no context_id
 
           events = []
           async for event in self.agent._run_async_impl(self.mock_context):
@@ -1810,8 +1767,7 @@ async def test_run_async_impl_successful_request(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-              "task-789",
-          )  # Tuple with parts and context/task ids
+          )  # Tuple with parts and context_id
 
           # Mock A2A client
           mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
@@ -1865,13 +1821,6 @@ 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):
@@ -1892,8 +1841,7 @@ async def test_run_async_impl_a2a_client_error(self):
           mock_construct.return_value = (
               [mock_a2a_part],
               "context-123",
-              "task-789",
-          )  # Tuple with parts and context/task ids
+          )  # Tuple with parts and context_id
 
           # Mock A2A client that throws an exception
           mock_a2a_client = AsyncMock()