From bb8a269079c975d5ca5422aa7d86c6d44d362a75 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 3 Dec 2025 18:25:10 -0800 Subject: [PATCH] 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 PiperOrigin-RevId: 839996774 --- src/google/adk/agents/remote_a2a_agent.py | 14 ++- .../unittests/agents/test_remote_a2a_agent.py | 108 +++++++++++++----- 2 files changed, 88 insertions(+), 34 deletions(-) diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 8d133060..46aefe45 100644 --- a/src/google/adk/agents/remote_a2a_agent.py +++ b/src/google/adk/agents/remote_a2a_agent.py @@ -343,18 +343,18 @@ class RemoteA2aAgent(BaseAgent): 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 @@ class RemoteA2aAgent(BaseAgent): 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 @@ class RemoteA2aAgent(BaseAgent): 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 @@ class RemoteA2aAgent(BaseAgent): # 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 @@ class RemoteA2aAgent(BaseAgent): parts=message_parts, role="user", context_id=context_id, + task_id=task_id, ) logger.debug(build_a2a_request_log(a2a_request)) diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py index e7865f39..e0ebecfe 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandlingFromFactory: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecutionFromFactory: 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 @@ class TestRemoteA2aAgentExecutionFromFactory: 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 @@ class TestRemoteA2aAgentExecutionFromFactory: 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 @@ class TestRemoteA2aAgentExecutionFromFactory: 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()