fix: Handling of A2ATaskStatusUpdateEvent when streaming in remote_a2a_agent

The proto has the Message object in the TaskStatus.

PiperOrigin-RevId: 813844289

Author: google-genai-bot

src/google/adk/agents/remote_a2a_agent.py

@@ -414,9 +414,14 @@ async def _handle_a2a_response(
           # response for a non-streaming task, which is the full task state.
           # We process this to get the initial message.
           event = convert_a2a_task_to_event(task, self.name, ctx)
-        elif isinstance(update, A2ATaskStatusUpdateEvent) and update.message:
+        elif (
+            isinstance(update, A2ATaskStatusUpdateEvent)
+            and update.status.message
+        ):
           # This is a streaming task status update with a message.
-          event = convert_a2a_message_to_event(update.message, self.name, ctx)
+          event = convert_a2a_message_to_event(
+              update.status.message, self.name, ctx
+          )
         elif isinstance(update, A2ATaskArtifactUpdateEvent) and (
             not update.append or update.last_chunk
         ):

tests/unittests/agents/test_remote_a2a_agent.py

@@ -44,6 +44,8 @@
   from a2a.types import SendMessageSuccessResponse
   from a2a.types import Task as A2ATask
   from a2a.types import TaskArtifactUpdateEvent
+  from a2a.types import TaskState
+  from a2a.types import TaskStatus
   from a2a.types import TaskStatusUpdateEvent
   from google.adk.agents.invocation_context import InvocationContext
   from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX
@@ -733,8 +735,9 @@ async def test_handle_a2a_response_with_task_status_update_with_message(self):
 
     mock_a2a_message = Mock(spec=A2AMessage)
     mock_update = Mock(spec=TaskStatusUpdateEvent)
-    mock_update.message = mock_a2a_message
-    mock_update.status = "COMPLETED"
+    mock_update.status = Mock(TaskStatus)
+    mock_update.status.state = TaskState.completed
+    mock_update.status.message = mock_a2a_message
 
     # Create a proper Event mock that can handle custom_metadata
     mock_event = Event(
@@ -770,8 +773,9 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self):
     mock_a2a_task.id = "task-123"
 
     mock_update = Mock(spec=TaskStatusUpdateEvent)
-    mock_update.message = None
-    mock_update.status = "COMPLETED"
+    mock_update.status = Mock(TaskStatus)
+    mock_update.status.state = TaskState.completed
+    mock_update.status.message = None
 
     result = await self.agent._handle_a2a_response(
         (mock_a2a_task, mock_update), self.mock_context
@@ -1021,8 +1025,9 @@ async def test_handle_a2a_response_with_task_status_update_with_message(self):
 
     mock_a2a_message = Mock(spec=A2AMessage)
     mock_update = Mock(spec=TaskStatusUpdateEvent)
-    mock_update.message = mock_a2a_message
-    mock_update.status = "COMPLETED"
+    mock_update.status = Mock(TaskStatus)
+    mock_update.status.state = TaskState.completed
+    mock_update.status.message = mock_a2a_message
 
     # Create a proper Event mock that can handle custom_metadata
     mock_event = Event(
@@ -1058,8 +1063,9 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self):
     mock_a2a_task.id = "task-123"
 
     mock_update = Mock(spec=TaskStatusUpdateEvent)
-    mock_update.message = None
-    mock_update.status = "COMPLETED"
+    mock_update.status = Mock(TaskStatus)
+    mock_update.status.state = TaskState.completed
+    mock_update.status.message = None
 
     result = await self.agent._handle_a2a_response(
         (mock_a2a_task, mock_update), self.mock_context