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

PiperOrigin-RevId: 840375992
This commit is contained in:
Google Team Member
2025-12-04 13:25:19 -08:00
committed by Copybara-Service
parent c557b0a1f2
commit 82bd4f380b
2 changed files with 34 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/agents/remote_a2a_agent.py
+6 -8
View File
@@ -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))
tests/unittests/agents/test_remote_a2a_agent.py
+28 -80
View File
@@ -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()