From 1e3905622b9806ec506fe545c0c2c49eb2b33e4c Mon Sep 17 00:00:00 2001
From: daleselaji-dev <265319989+daleselaji-dev@users.noreply.github.com>
Date: Wed, 5 Aug 2026 10:41:05 +0800
Subject: [PATCH] fix(a2a): report truncated remote task streams

---
 src/google/adk/agents/remote_a2a_agent.py     | 27 +++++++++
 .../unittests/agents/test_remote_a2a_agent.py | 58 +++++++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py
index 5435954706..4de5206386 100644
--- a/src/google/adk/agents/remote_a2a_agent.py
+++ b/src/google/adk/agents/remote_a2a_agent.py
@@ -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,
@@ -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:
@@ -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)
diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py
index fe39a29c26..3f0b94f891 100644
--- a/tests/unittests/agents/test_remote_a2a_agent.py
+++ b/tests/unittests/agents/test_remote_a2a_agent.py
@@ -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."""