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
27 changes: 27 additions & 0 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -873,6 +873,7 @@ async def _run_async_impl(
# status/artifact updates are aggregated into a running task (matching the
# 0.3.x client behavior).
normalize_stream_item = _compat.make_stream_normalizer()
last_task = None
async with Aclosing(
_compat.send_message(
a2a_client,
Expand All @@ -888,6 +889,7 @@ async def _run_async_impl(
metadata = None
if isinstance(a2a_response, tuple):
task = a2a_response[0]
last_task = task
if task:
metadata = task.metadata
else:
Expand Down Expand Up @@ -926,6 +928,31 @@ async def _run_async_impl(

yield event

if last_task and last_task.status:
task_state = last_task.status.state
if task_state not in (
_compat.TS_COMPLETED,
_compat.TS_FAILED,
_compat.TS_CANCELED,
):
error_message = (
"A2A response stream ended before the task reached a terminal "
f"state (last state: {task_state})"
)
logger.error(error_message)
yield Event(
author=self.name,
error_message=error_message,
invocation_id=ctx.invocation_id,
branch=ctx.branch,
custom_metadata={
A2A_METADATA_PREFIX + "request": _compat.a2a_to_dict(
a2a_request
),
A2A_METADATA_PREFIX + "error": error_message,
},
)

except _compat.A2A_HTTP_ERRORS as e:
error_message = f"A2A request failed: {e}"
logger.error(error_message)
Expand Down
58 changes: 58 additions & 0 deletions tests/unittests/agents/test_remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -3023,6 +3023,64 @@ async def test_run_async_impl_closes_stream_when_abandoned(self):

mock_send_message.aclose.assert_awaited_once()

@pytest.mark.asyncio
async def test_run_async_impl_reports_non_terminal_stream_end(self):
"""A cleanly closed stream must not hide a still-running remote task."""
with patch.object(self.agent, "_ensure_resolved") as mock_ensure_resolved:
with patch.object(
self.agent, "_create_a2a_request_for_user_function_response"
) as mock_create_func:
mock_create_func.return_value = None

with patch.object(
self.agent, "_construct_message_parts_from_session"
) as mock_construct:
mock_a2a_part = _compat.make_text_part("test")
mock_construct.return_value = ([mock_a2a_part], "context-123")

mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
mock_send_message = AsyncMock()
update = _compat.make_task_status_update_event(
"task-123",
"context-123",
_compat.make_task_status(_compat.TS_WORKING),
final=False,
)
if _compat.IS_A2A_V1:
from a2a.types import StreamResponse

stream_response = StreamResponse()
stream_response.status_update.CopyFrom(update)
raw_update = stream_response
else:
raw_update = update
mock_send_message.__aiter__.return_value = [raw_update]
mock_a2a_client.send_message.return_value = mock_send_message
mock_ensure_resolved.return_value = mock_a2a_client

with patch.object(self.agent, "_handle_a2a_response", return_value=None):
with patch(
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
):
with patch(
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
):
with patch(
"google.adk.a2a._compat.a2a_to_dict",
return_value={"k": "v"},
):
events = [
event
async for event in self.agent._run_async_impl(
self.mock_context
)
]

assert len(events) == 1
assert "ended before the task reached a terminal state" in (
events[0].error_message
)

@pytest.mark.asyncio
async def test_run_async_impl_a2a_client_error(self):
"""Test _run_async_impl when A2A send_message fails."""
Expand Down