Skip to content
Open
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
26 changes: 15 additions & 11 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -543,18 +543,22 @@ async def _handle_a2a_response(
for part in event.content.parts:
part.thought = True
_add_mock_function_call(event, update.status.state)
elif isinstance(update, A2ATaskArtifactUpdateEvent) and (
not update.append or update.last_chunk
):
elif isinstance(update, A2ATaskArtifactUpdateEvent):
# This is a streaming task artifact update.
# We only handle full artifact updates and ignore partial updates.
# Note: Depends on the server implementation, there is no clear
# definition of what a partial update is currently. We use the two
# signals:
# 1. append: True for partial updates, False for full updates.
# 2. last_chunk: True for full updates, False for partial updates.
event = convert_a2a_task_to_event(
task, self.name, ctx, self._a2a_part_converter
# Convert only the parts carried by this update. Converting the
# accumulated task here would re-emit earlier chunks of the same
# artifact, duplicating already-streamed content.
if not update.artifact.parts:
return None
event = convert_a2a_message_to_event(
_compat.make_message(
message_id="",
role="agent",
parts=update.artifact.parts,
),
self.name,
ctx,
self._a2a_part_converter,
)
if not event:
return None
Expand Down
205 changes: 148 additions & 57 deletions tests/unittests/agents/test_remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,39 @@ def _make_stream_message(message: A2AMessage):
return message


def _make_artifact_chunk(text: str, *, append: bool, last_chunk: bool):
"""Build one streamed chunk of an artifact, version-agnostically."""
return TaskArtifactUpdateEvent(
task_id="task-123",
context_id="context-123",
append=append,
last_chunk=last_chunk,
artifact=_compat.make_artifact(
artifact_id="artifact-1",
parts=[_compat.make_text_part(text)],
),
)


def _make_accumulated_task(part_texts):
"""Build the running Task the stream normalizer yields alongside an update.

The task carries the artifact parts accumulated across all chunks received
so far, mirroring the 0.3.x ClientTaskManager / 1.x stream normalizer.
"""
return _compat.make_task(
id="task-123",
status=_compat.make_task_status(_compat.TS_WORKING),
context_id="context-123",
artifacts=[
_compat.make_artifact(
artifact_id="artifact-1",
parts=[_compat.make_text_part(text) for text in part_texts],
)
],
)


# Helper function to create a proper AgentCard for testing
def create_test_agent_card(
name: str = "test-agent",
Expand Down Expand Up @@ -1326,11 +1359,7 @@ async def test_handle_a2a_response_with_artifact_update(self):
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"

mock_artifact = Mock(spec=Artifact)
mock_update = Mock(spec=TaskArtifactUpdateEvent)
mock_update.artifact = mock_artifact
mock_update.append = False
mock_update.last_chunk = True
update = _make_artifact_chunk("chunk", append=False, last_chunk=True)

# Create a proper Event mock that can handle custom_metadata
mock_event = Event(
Expand All @@ -1339,45 +1368,53 @@ async def test_handle_a2a_response_with_artifact_update(self):
branch=self.mock_context.branch,
)

with patch.object(
remote_a2a_agent,
"convert_a2a_task_to_event",
autospec=True,
with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

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

assert result == mock_event
mock_convert.assert_called_once_with(
mock_a2a_task,
self.agent.name,
self.mock_context,
self.agent._a2a_part_converter,
)
mock_convert.assert_called_once()
# Only the parts carried by this update are converted, not the
# accumulated task.
converted_message = mock_convert.call_args[0][0]
assert list(converted_message.parts) == list(update.artifact.parts)
# 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_partial_artifact_update(self):
"""Test that partial artifact updates are ignored."""
async def test_handle_a2a_response_with_appended_artifact_chunk(self):
"""An appended (middle) artifact chunk emits only its own parts."""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"

mock_update = Mock(spec=TaskArtifactUpdateEvent)
mock_update.artifact = Mock(spec=Artifact)
mock_update.append = True
mock_update.last_chunk = False
update = _make_artifact_chunk("middle", append=True, last_chunk=False)

result = await self.agent._handle_a2a_response(
(mock_a2a_task, mock_update), self.mock_context
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
)

assert result is None
with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

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

assert result == mock_event
converted_message = mock_convert.call_args[0][0]
assert list(converted_message.parts) == list(update.artifact.parts)

@pytest.mark.asyncio
async def test_handle_a2a_response_with_real_empty_status_message(self):
Expand Down Expand Up @@ -1407,6 +1444,58 @@ async def test_handle_a2a_response_with_real_empty_status_message(self):
assert result is None


class TestRemoteA2aAgentStreamingArtifactChunks:
"""Regression tests for chunked artifact streams (#6343)."""

def setup_method(self):
"""Setup test fixtures."""
self.agent = RemoteA2aAgent(
name="test_agent",
agent_card=create_test_agent_card(),
)
self.mock_context = Mock(spec=InvocationContext)
self.mock_context.invocation_id = "invocation-123"
self.mock_context.branch = "main"

@pytest.mark.asyncio
async def test_chunked_artifact_stream_emits_each_part_exactly_once(self):
"""A two-chunk artifact stream renders its parts without duplication."""
chunk1 = _make_artifact_chunk("Hello, ", append=False, last_chunk=False)
chunk2 = _make_artifact_chunk("world!", append=True, last_chunk=True)
# (task, update) pairs as the client stream yields them: the task carries
# the artifact parts accumulated so far.
stream = [
(_make_accumulated_task(["Hello, "]), chunk1),
(_make_accumulated_task(["Hello, ", "world!"]), chunk2),
]

rendered = []
for pair in stream:
event = await self.agent._handle_a2a_response(pair, self.mock_context)
if event and event.content and event.content.parts:
rendered.extend(part.text for part in event.content.parts if part.text)

assert "".join(rendered) == "Hello, world!"

@pytest.mark.asyncio
async def test_artifact_update_without_parts_is_ignored(self):
"""An artifact update carrying no parts must not emit a spurious event."""
update = TaskArtifactUpdateEvent(
task_id="task-123",
context_id="context-123",
append=False,
last_chunk=True,
artifact=_compat.make_artifact(artifact_id="artifact-1", parts=[]),
)
task = _make_accumulated_task(["already streamed"])

result = await self.agent._handle_a2a_response(
(task, update), self.mock_context
)

assert result is None


class TestRemoteA2aAgentMessageHandlingFromFactory:
"""Test message handling functionality."""

Expand Down Expand Up @@ -1770,11 +1859,7 @@ async def test_handle_a2a_response_with_artifact_update(self):
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"

mock_artifact = Mock(spec=Artifact)
mock_update = Mock(spec=TaskArtifactUpdateEvent)
mock_update.artifact = mock_artifact
mock_update.append = False
mock_update.last_chunk = True
update = _make_artifact_chunk("chunk", append=False, last_chunk=True)

# Create a proper Event mock that can handle custom_metadata
mock_event = Event(
Expand All @@ -1783,45 +1868,53 @@ async def test_handle_a2a_response_with_artifact_update(self):
branch=self.mock_context.branch,
)

with patch.object(
remote_a2a_agent,
"convert_a2a_task_to_event",
autospec=True,
with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

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

assert result == mock_event
mock_convert.assert_called_once_with(
mock_a2a_task,
self.agent.name,
self.mock_context,
self.agent._a2a_part_converter,
)
mock_convert.assert_called_once()
# Only the parts carried by this update are converted, not the
# accumulated task.
converted_message = mock_convert.call_args[0][0]
assert list(converted_message.parts) == list(update.artifact.parts)
# 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_partial_artifact_update(self):
"""Test that partial artifact updates are ignored."""
async def test_handle_a2a_response_with_appended_artifact_chunk(self):
"""An appended (middle) artifact chunk emits only its own parts."""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"

mock_update = Mock(spec=TaskArtifactUpdateEvent)
mock_update.artifact = Mock(spec=Artifact)
mock_update.append = True
mock_update.last_chunk = False
update = _make_artifact_chunk("middle", append=True, last_chunk=False)

result = await self.agent._handle_a2a_response(
(mock_a2a_task, mock_update), self.mock_context
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
)

assert result is None
with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

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

assert result == mock_event
converted_message = mock_convert.call_args[0][0]
assert list(converted_message.parts) == list(update.artifact.parts)


class TestRemoteA2aAgentMessageHandlingV2:
Expand Down Expand Up @@ -2250,23 +2343,21 @@ async def test_legacy_message_converter_returns_none_status_update(self):
assert result is None

@pytest.mark.asyncio
async def test_legacy_task_converter_returns_none_artifact_update(self):
"""Legacy handler must not crash when task converter returns None for artifact update."""
async def test_legacy_message_converter_returns_none_artifact_update(self):
"""Legacy handler must not crash when message converter returns None for artifact update."""
mock_task = Mock(spec=A2ATask)
mock_task.id = "task-123"
mock_task.context_id = None

mock_update = Mock(spec=TaskArtifactUpdateEvent)
mock_update.append = False
mock_update.last_chunk = True
update = _make_artifact_chunk("chunk", append=False, last_chunk=True)

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

result = await self.legacy_agent._handle_a2a_response(
(mock_task, mock_update), self.mock_context
(mock_task, update), self.mock_context
)

assert result is None
Expand Down
Loading