From 82bd4f380bd8b4822191ea16e6140fe2613023ad Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Thu, 4 Dec 2025 13:25:19 -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 PiperOrigin-RevId: 840375992 --- src/google/adk/agents/remote_a2a_agent.py | 14 +-- .../unittests/agents/test_remote_a2a_agent.py | 108 +++++------------- 2 files changed, 34 insertions(+), 88 deletions(-) diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 46aefe45..8d133060 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], 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 @@ 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) @@ -385,7 +384,7 @@ class RemoteA2aAgent(BaseAgent): 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 @@ 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, 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 @@ 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 e0ebecfe..e7865f39 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandling: 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 @@ class TestRemoteA2aAgentMessageHandlingFromFactory: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ 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): @@ -1598,8 +1565,7 @@ class TestRemoteA2aAgentExecution: 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 @@ class TestRemoteA2aAgentExecution: 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 @@ 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: @@ -1779,8 +1737,7 @@ class TestRemoteA2aAgentExecutionFromFactory: 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 @@ class TestRemoteA2aAgentExecutionFromFactory: 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 @@ 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): @@ -1892,8 +1841,7 @@ class TestRemoteA2aAgentExecutionFromFactory: 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()