encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Extract and propagate task_id in RemoteA2aAgent

Browse Source 
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
This commit is contained in:
George Weale
2025-12-03 18:25:10 -08:00
committed by Copybara-Service
parent 507424acb9
commit bb8a269079
2 changed files with 88 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/agents/remote_a2a_agent.py
+8 -6
View File
@@ -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))
tests/unittests/agents/test_remote_a2a_agent.py
+80 -28
View File
@@ -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()