Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -412,7 +412,15 @@ async def _handle_a2a_response(
)
# for streaming task, we update the event with the task status.
# We update the event as Thought updates.
if task and task.status and task.status.state == TaskState.submitted:
if (
task
and task.status
and task.status.state
in (
TaskState.submitted,
TaskState.working,
)
):
event.content.parts[0].thought = True
elif (
isinstance(update, A2ATaskStatusUpdateEvent)
Expand All @@ -423,10 +431,10 @@ async def _handle_a2a_response(
event = convert_a2a_message_to_event(
update.status.message, self.name, ctx, self._a2a_part_converter
)
if event.content and update.status.state in [
if event.content and update.status.state in (
TaskState.submitted,
TaskState.working,
]:
):
for part in event.content.parts:
part.thought = True
elif isinstance(update, A2ATaskArtifactUpdateEvent) and (
Expand Down
42 changes: 42 additions & 0 deletions tests/unittests/agents/test_remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -845,6 +845,48 @@ async def test_handle_a2a_response_with_task_submitted_and_no_update(self):
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata

@pytest.mark.asyncio
async def test_handle_a2a_response_with_task_working_and_no_update(self):
"""Test successful A2A response handling with streaming task and no update."""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
mock_a2a_task.status.state = TaskState.working

# Create a proper Event mock that can handle custom_metadata
mock_a2a_part = Mock(spec=TextPart)
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
)

with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_task_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

result = await self.agent._handle_a2a_response(
(mock_a2a_task, None), self.mock_context
)

assert result == mock_event
mock_convert.assert_called_once_with(
mock_a2a_task,
self.agent.name,
self.mock_context,
self.mock_a2a_part_converter,
)
# Check the parts are updated as Thought
assert result.content.parts[0].thought is True
assert result.content.parts[0].thought_signature is None
# Check that metadata was added
assert result.custom_metadata is not None
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata

@pytest.mark.asyncio
async def test_handle_a2a_response_with_task_status_update_with_message(self):
"""Test handling of a task status update with a message."""
Expand Down
Loading