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: Return final task result in task artifact instead of status message

Browse Source 
According to a2a protocol task artifact is a different concept from adk artifact.
if a task is completed the final result should be in task artifact.

PiperOrigin-RevId: 782154265
This commit is contained in:
Xiang (Sean) Zhou
2025-07-11 16:47:28 -07:00
committed by Copybara-Service
parent a57d629bb9
commit a8fcc1b8ab
4 changed files with 352 additions and 441 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/a2a/converters/event_converter.py
+5 -85
View File
@@ -24,13 +24,11 @@ from typing import Optional
import uuid
from a2a.server.events import Event as A2AEvent
from a2a.types import Artifact
from a2a.types import DataPart
from a2a.types import Message
from a2a.types import Part as A2APart
from a2a.types import Role
from a2a.types import Task
from a2a.types import TaskArtifactUpdateEvent
from a2a.types import TaskState
from a2a.types import TaskStatus
from a2a.types import TaskStatusUpdateEvent
@@ -145,81 +143,6 @@ def _create_artifact_id(
return ARTIFACT_ID_SEPARATOR.join(components)
def _convert_artifact_to_a2a_events(
event: Event,
invocation_context: InvocationContext,
filename: str,
version: int,
task_id: Optional[str] = None,
context_id: Optional[str] = None,
) -> TaskArtifactUpdateEvent:
"""Converts a new artifact version to an A2A TaskArtifactUpdateEvent.
Args:
event: The ADK event containing the artifact information.
invocation_context: The invocation context.
filename: The name of the artifact file.
version: The version number of the artifact.
task_id: Optional task ID to use for generated events. If not provided, new UUIDs will be generated.
Returns:
A TaskArtifactUpdateEvent representing the artifact update.
Raises:
ValueError: If required parameters are invalid.
RuntimeError: If artifact loading fails.
"""
if not filename:
raise ValueError("Filename cannot be empty")
if version < 0:
raise ValueError("Version must be non-negative")
try:
artifact_part = invocation_context.artifact_service.load_artifact(
app_name=invocation_context.app_name,
user_id=invocation_context.user_id,
session_id=invocation_context.session.id,
filename=filename,
version=version,
)
converted_part = convert_genai_part_to_a2a_part(part=artifact_part)
if not converted_part:
raise RuntimeError(f"Failed to convert artifact part for {filename}")
artifact_id = _create_artifact_id(
invocation_context.app_name,
invocation_context.user_id,
invocation_context.session.id,
filename,
version,
)
return TaskArtifactUpdateEvent(
taskId=task_id,
append=False,
contextId=context_id,
lastChunk=True,
artifact=Artifact(
artifactId=artifact_id,
name=filename,
metadata={
"filename": filename,
"version": version,
},
parts=[converted_part],
),
)
except Exception as e:
logger.error(
"Failed to convert artifact for %s, version %s: %s",
filename,
version,
e,
)
raise RuntimeError(f"Artifact conversion failed: {e}") from e
def _process_long_running_tool(a2a_part: A2APart, event: Event) -> None:
"""Processes long-running tool metadata for an A2A part.
@@ -268,7 +191,11 @@ def convert_a2a_task_to_event(
try:
# Extract message from task status or history
message = None
if a2a_task.status and a2a_task.status.message:
if a2a_task.artifacts:
message = Message(
messageId="", role=Role.agent, parts=a2a_task.artifacts[-1].parts
)
elif a2a_task.status and a2a_task.status.message:
message = a2a_task.status.message
elif a2a_task.history:
message = a2a_task.history[-1]
@@ -573,13 +500,6 @@ def convert_event_to_a2a_events(
a2a_events = []
try:
# Handle artifact deltas
if event.actions.artifact_delta:
for filename, version in event.actions.artifact_delta.items():
artifact_event = _convert_artifact_to_a2a_events(
event, invocation_context, filename, version, task_id, context_id
)
a2a_events.append(artifact_event)
# Handle error scenarios
if event.error_code:
src/google/adk/a2a/executor/a2a_agent_executor.py
+33 -4
View File
@@ -28,8 +28,10 @@ try:
from a2a.server.agent_execution import AgentExecutor
from a2a.server.agent_execution.context import RequestContext
from a2a.server.events.event_queue import EventQueue
from a2a.types import Artifact
from a2a.types import Message
from a2a.types import Role
from a2a.types import TaskArtifactUpdateEvent
from a2a.types import TaskState
from a2a.types import TaskStatus
from a2a.types import TaskStatusUpdateEvent
@@ -218,15 +220,42 @@ class A2aAgentExecutor(AgentExecutor):
await event_queue.enqueue_event(a2a_event)
# publish the task result event - this is final
if (
task_result_aggregator.task_state == TaskState.working
and task_result_aggregator.task_status_message is not None
and task_result_aggregator.task_status_message.parts
):
# if task is still working properly, publish the artifact update event as
# the final result according to a2a protocol.
await event_queue.enqueue_event(
TaskArtifactUpdateEvent(
taskId=context.task_id,
lastChunk=True,
contextId=context.context_id,
artifact=Artifact(
artifactId=str(uuid.uuid4()),
parts=task_result_aggregator.task_status_message.parts,
),
)
)
# public the final status update event
await event_queue.enqueue_event(
TaskStatusUpdateEvent(
taskId=context.task_id,
status=TaskStatus(
state=(
task_result_aggregator.task_state
if task_result_aggregator.task_state != TaskState.working
else TaskState.completed
state=TaskState.completed,
timestamp=datetime.now(timezone.utc).isoformat(),
),
contextId=context.context_id,
final=True,
)
)
else:
await event_queue.enqueue_event(
TaskStatusUpdateEvent(
taskId=context.task_id,
status=TaskStatus(
state=task_result_aggregator.task_state,
timestamp=datetime.now(timezone.utc).isoformat(),
message=task_result_aggregator.task_status_message,
),
tests/unittests/a2a/converters/test_event_converter.py
+74 -312
View File
File diff suppressed because it is too large Load Diff
tests/unittests/a2a/executor/test_a2a_agent_executor.py
+210 -10
View File
@@ -153,9 +153,10 @@ class TestA2aAgentExecutor:
0
]
assert final_event.final == True
# The TaskResultAggregator is created with default state (working), so final state should be completed
# The TaskResultAggregator is created with default state (working), and since no messages
# are processed, it will publish a status event with the current state
assert hasattr(final_event.status, "message")
assert final_event.status.state == TaskState.completed
assert final_event.status.state == TaskState.working
@pytest.mark.asyncio
async def test_execute_no_message_error(self):
@@ -224,9 +225,10 @@ class TestA2aAgentExecutor:
0
]
assert final_event.final == True
# The TaskResultAggregator is created with default state (working), so final state should be completed
# The TaskResultAggregator is created with default state (working), and since no messages
# are processed, it will publish a status event with the current state
assert hasattr(final_event.status, "message")
assert final_event.status.state == TaskState.completed
assert final_event.status.state == TaskState.working
@pytest.mark.asyncio
async def test_prepare_session_new_session(self):
@@ -456,9 +458,10 @@ class TestA2aAgentExecutor:
0
]
assert final_event.final == True
# The TaskResultAggregator is created with default state (working), so final state should be completed
# The TaskResultAggregator is created with default state (working), and since no messages
# are processed, it will publish a status event with the current state
assert hasattr(final_event.status, "message")
assert final_event.status.state == TaskState.completed
assert final_event.status.state == TaskState.working
@pytest.mark.asyncio
async def test_execute_with_async_callable_runner(self):
@@ -522,9 +525,10 @@ class TestA2aAgentExecutor:
0
]
assert final_event.final == True
# The TaskResultAggregator is created with default state (working), so final state should be completed
# The TaskResultAggregator is created with default state (working), and since no messages
# are processed, it will publish a status event with the current state
assert hasattr(final_event.status, "message")
assert final_event.status.state == TaskState.completed
assert final_event.status.state == TaskState.working
@pytest.mark.asyncio
async def test_handle_request_integration(self):
@@ -608,8 +612,8 @@ class TestA2aAgentExecutor:
assert (
final_event.status.message == mock_aggregator.task_status_message
)
# When aggregator state is working, final event should be completed
assert final_event.status.state == TaskState.completed
# When aggregator state is working but no message, final event should be working
assert final_event.status.state == TaskState.working
@pytest.mark.asyncio
async def test_cancel_with_task_id(self):
@@ -745,6 +749,7 @@ class TestA2aAgentExecutor:
assert len(final_events) >= 1
final_event = final_events[-1] # Get the last final event
assert final_event.status.message == test_message
# When aggregator state is completed (not working), final event should be completed
assert final_event.status.state == TaskState.completed
@pytest.mark.asyncio
@@ -827,3 +832,198 @@ class TestA2aAgentExecutor:
assert final_event.status.message == test_message
# When aggregator state is failed (not working), final event should keep failed state
assert final_event.status.state == TaskState.failed
@pytest.mark.asyncio
async def test_handle_request_with_working_state_publishes_artifact_and_completed(
self,
):
"""Test that when aggregator state is working, it publishes artifact update and completed status."""
# Setup context with task_id
self.mock_context.task_id = "test-task-id"
self.mock_context.context_id = "test-context-id"
# Create a test message to be returned by the aggregator
from a2a.types import Message
from a2a.types import Part
from a2a.types import Role
from a2a.types import TextPart
test_message = Mock(spec=Message)
test_message.messageId = "test-message-id"
test_message.role = Role.agent
test_message.parts = [Part(root=TextPart(text="test content"))]
# Setup detailed mocks
with patch(
"google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args"
) as mock_convert:
mock_convert.return_value = {
"user_id": "test-user",
"session_id": "test-session",
"new_message": Mock(),
"run_config": Mock(),
}
# Mock session service
mock_session = Mock()
mock_session.id = "test-session"
self.mock_runner.session_service.get_session = AsyncMock(
return_value=mock_session
)
# Mock invocation context
mock_invocation_context = Mock()
self.mock_runner._new_invocation_context.return_value = (
mock_invocation_context
)
# Mock agent run with multiple events using proper async generator
mock_events = [Mock(spec=Event), Mock(spec=Event)]
# Configure run_async to return the async generator when awaited
async def mock_run_async(**kwargs):
async for item in self._create_async_generator(mock_events):
yield item
self.mock_runner.run_async = mock_run_async
with patch(
"google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events"
) as mock_convert_events:
mock_convert_events.return_value = [Mock()]
with patch(
"google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator"
) as mock_aggregator_class:
mock_aggregator = Mock()
# Test with working state - should publish artifact update and completed status
mock_aggregator.task_state = TaskState.working
mock_aggregator.task_status_message = test_message
mock_aggregator_class.return_value = mock_aggregator
# Execute
await self.executor._handle_request(
self.mock_context, self.mock_event_queue
)
# Verify artifact update event was published
artifact_events = [
call[0][0]
for call in self.mock_event_queue.enqueue_event.call_args_list
if hasattr(call[0][0], "artifact")
and call[0][0].lastChunk == True
]
assert len(artifact_events) == 1
artifact_event = artifact_events[0]
assert artifact_event.taskId == "test-task-id"
assert artifact_event.contextId == "test-context-id"
# Check that artifact parts correspond to message parts
assert len(artifact_event.artifact.parts) == len(test_message.parts)
assert artifact_event.artifact.parts == test_message.parts
# Verify final status event was published with completed state
final_events = [
call[0][0]
for call in self.mock_event_queue.enqueue_event.call_args_list
if hasattr(call[0][0], "final") and call[0][0].final == True
]
assert len(final_events) >= 1
final_event = final_events[-1] # Get the last final event
assert final_event.status.state == TaskState.completed
assert final_event.taskId == "test-task-id"
assert final_event.contextId == "test-context-id"
@pytest.mark.asyncio
async def test_handle_request_with_non_working_state_publishes_status_only(
self,
):
"""Test that when aggregator state is not working, it publishes only the status event."""
# Setup context with task_id
self.mock_context.task_id = "test-task-id"
self.mock_context.context_id = "test-context-id"
# Create a test message to be returned by the aggregator
from a2a.types import Message
from a2a.types import Part
from a2a.types import Role
from a2a.types import TextPart
test_message = Mock(spec=Message)
test_message.messageId = "test-message-id"
test_message.role = Role.agent
test_message.parts = [Part(root=TextPart(text="test content"))]
# Setup detailed mocks
with patch(
"google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args"
) as mock_convert:
mock_convert.return_value = {
"user_id": "test-user",
"session_id": "test-session",
"new_message": Mock(),
"run_config": Mock(),
}
# Mock session service
mock_session = Mock()
mock_session.id = "test-session"
self.mock_runner.session_service.get_session = AsyncMock(
return_value=mock_session
)
# Mock invocation context
mock_invocation_context = Mock()
self.mock_runner._new_invocation_context.return_value = (
mock_invocation_context
)
# Mock agent run with multiple events using proper async generator
mock_events = [Mock(spec=Event), Mock(spec=Event)]
# Configure run_async to return the async generator when awaited
async def mock_run_async(**kwargs):
async for item in self._create_async_generator(mock_events):
yield item
self.mock_runner.run_async = mock_run_async
with patch(
"google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events"
) as mock_convert_events:
mock_convert_events.return_value = [Mock()]
with patch(
"google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator"
) as mock_aggregator_class:
mock_aggregator = Mock()
# Test with auth_required state - should publish only status event
mock_aggregator.task_state = TaskState.auth_required
mock_aggregator.task_status_message = test_message
mock_aggregator_class.return_value = mock_aggregator
# Execute
await self.executor._handle_request(
self.mock_context, self.mock_event_queue
)
# Verify no artifact update event was published
artifact_events = [
call[0][0]
for call in self.mock_event_queue.enqueue_event.call_args_list
if hasattr(call[0][0], "artifact")
and call[0][0].lastChunk == True
]
assert len(artifact_events) == 0
# Verify final status event was published with the actual state and message
final_events = [
call[0][0]
for call in self.mock_event_queue.enqueue_event.call_args_list
if hasattr(call[0][0], "final") and call[0][0].final == True
]
assert len(final_events) >= 1
final_event = final_events[-1] # Get the last final event
assert final_event.status.state == TaskState.auth_required
assert final_event.status.message == test_message
assert final_event.taskId == "test-task-id"
assert final_event.contextId == "test-context-id"