fix: Update remote_a2a_agent to better handle streaming events and avoid duplicate responses

Currently, the A2A Task -> ADK event conversion is producing the same events on the last two update events (the last is a status update marking the task complete)

The change here based on A2AClientEvent(task, update):
 - if the update == None: handle the non-streaming task case and also streaming case for the initial task creation event
 - if the update = TaskStatusUpdateEvent AND a message is set: emit an event with that message
 - if a task status update AND no message is set: don't emit event (for example, the final status update)
 - if the update is ArtifactUpdateEvent and it's final artifact: emit the event

PiperOrigin-RevId: 812878869

tests/unittests/agents/test_remote_a2a_agent.py

@@ -39,9 +39,12 @@ try:
   from a2a.types import AgentCapabilities
   from a2a.types import AgentCard
   from a2a.types import AgentSkill
+  from a2a.types import Artifact
   from a2a.types import Message as A2AMessage
   from a2a.types import SendMessageSuccessResponse
   from a2a.types import Task as A2ATask
+  from a2a.types import TaskArtifactUpdateEvent
+  from a2a.types import TaskStatusUpdateEvent
   from google.adk.agents.invocation_context import InvocationContext
   from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX
   from google.adk.agents.remote_a2a_agent import AgentCardResolutionError
@@ -60,6 +63,9 @@ except ImportError as e:
     A2AMessage = DummyTypes()
     SendMessageSuccessResponse = DummyTypes()
     A2ATask = DummyTypes()
+    TaskStatusUpdateEvent = DummyTypes()
+    Artifact = DummyTypes()
+    TaskArtifactUpdateEvent = DummyTypes()
     InvocationContext = DummyTypes()
     RemoteA2aAgent = DummyTypes()
     AgentCardResolutionError = Exception
@@ -685,8 +691,8 @@ class TestRemoteA2aAgentMessageHandling:
       assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
 
   @pytest.mark.asyncio
-  async def test_handle_a2a_response_success_with_task(self):
-    """Test successful A2A response handling with task."""
+  async def test_handle_a2a_response_with_task_and_no_update(self):
+    """Test successful A2A response handling with task and no update."""
     mock_a2a_task = Mock(spec=A2ATask)
     mock_a2a_task.id = "task-123"
     mock_a2a_task.context_id = "context-123"
@@ -718,6 +724,116 @@ class TestRemoteA2aAgentMessageHandling:
       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."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+    mock_a2a_task.context_id = "context-123"
+
+    mock_a2a_message = Mock(spec=A2AMessage)
+    mock_update = Mock(spec=TaskStatusUpdateEvent)
+    mock_update.message = mock_a2a_message
+    mock_update.status = "COMPLETED"
+
+    # Create a proper Event mock that can handle custom_metadata
+    mock_event = Event(
+        author=self.agent.name,
+        invocation_id=self.mock_context.invocation_id,
+        branch=self.mock_context.branch,
+    )
+
+    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
+      )
+
+      assert result == mock_event
+      mock_convert.assert_called_once_with(
+          mock_a2a_message,
+          self.agent.name,
+          self.mock_context,
+      )
+      # 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_no_message(self):
+    """Test handling of a task status update with no message."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+
+    mock_update = Mock(spec=TaskStatusUpdateEvent)
+    mock_update.message = None
+    mock_update.status = "COMPLETED"
+
+    result = await self.agent._handle_a2a_response(
+        (mock_a2a_task, mock_update), self.mock_context
+    )
+
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_handle_a2a_response_with_artifact_update(self):
+    """Test successful A2A response handling with artifact update."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    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
+
+    # Create a proper Event mock that can handle custom_metadata
+    mock_event = Event(
+        author=self.agent.name,
+        invocation_id=self.mock_context.invocation_id,
+        branch=self.mock_context.branch,
+    )
+
+    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, mock_update), self.mock_context
+      )
+
+      assert result == mock_event
+      mock_convert.assert_called_once_with(
+          mock_a2a_task, self.agent.name, self.mock_context
+      )
+      # 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."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+
+    mock_update = Mock(spec=TaskArtifactUpdateEvent)
+    mock_update.artifact = Mock(spec=Artifact)
+    mock_update.append = True
+    mock_update.last_chunk = False
+
+    result = await self.agent._handle_a2a_response(
+        (mock_a2a_task, mock_update), self.mock_context
+    )
+
+    assert result is None
+
 
 class TestRemoteA2aAgentMessageHandlingFromFactory:
   """Test message handling functionality."""
@@ -865,8 +981,8 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
       assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
 
   @pytest.mark.asyncio
-  async def test_handle_a2a_response_success_with_task(self):
-    """Test successful A2A response handling with task."""
+  async def test_handle_a2a_response_with_task_and_no_update(self):
+    """Test successful A2A response handling with task and no update."""
     mock_a2a_task = Mock(spec=A2ATask)
     mock_a2a_task.id = "task-123"
     mock_a2a_task.context_id = "context-123"
@@ -896,6 +1012,116 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
       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."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+    mock_a2a_task.context_id = "context-123"
+
+    mock_a2a_message = Mock(spec=A2AMessage)
+    mock_update = Mock(spec=TaskStatusUpdateEvent)
+    mock_update.message = mock_a2a_message
+    mock_update.status = "COMPLETED"
+
+    # Create a proper Event mock that can handle custom_metadata
+    mock_event = Event(
+        author=self.agent.name,
+        invocation_id=self.mock_context.invocation_id,
+        branch=self.mock_context.branch,
+    )
+
+    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
+      )
+
+      assert result == mock_event
+      mock_convert.assert_called_once_with(
+          mock_a2a_message,
+          self.agent.name,
+          self.mock_context,
+      )
+      # 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_no_message(self):
+    """Test handling of a task status update with no message."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+
+    mock_update = Mock(spec=TaskStatusUpdateEvent)
+    mock_update.message = None
+    mock_update.status = "COMPLETED"
+
+    result = await self.agent._handle_a2a_response(
+        (mock_a2a_task, mock_update), self.mock_context
+    )
+
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_handle_a2a_response_with_artifact_update(self):
+    """Test successful A2A response handling with artifact update."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    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
+
+    # Create a proper Event mock that can handle custom_metadata
+    mock_event = Event(
+        author=self.agent.name,
+        invocation_id=self.mock_context.invocation_id,
+        branch=self.mock_context.branch,
+    )
+
+    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, mock_update), self.mock_context
+      )
+
+      assert result == mock_event
+      mock_convert.assert_called_once_with(
+          mock_a2a_task, self.agent.name, self.mock_context
+      )
+      # 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."""
+    mock_a2a_task = Mock(spec=A2ATask)
+    mock_a2a_task.id = "task-123"
+
+    mock_update = Mock(spec=TaskArtifactUpdateEvent)
+    mock_update.artifact = Mock(spec=Artifact)
+    mock_update.append = True
+    mock_update.last_chunk = False
+
+    result = await self.agent._handle_a2a_response(
+        (mock_a2a_task, mock_update), self.mock_context
+    )
+
+    assert result is None
+
 
 class TestRemoteA2aAgentExecution:
   """Test agent execution functionality."""
@@ -1019,6 +1245,7 @@ class TestRemoteA2aAgentExecution:
                   # Add model_dump to mock_response for metadata
                   mock_response.model_dump.return_value = {"test": "response"}
 
+                  # Execute
                   events = []
                   async for event in self.agent._run_async_impl(
                       self.mock_context
@@ -1211,6 +1438,7 @@ class TestRemoteA2aAgentExecutionFromFactory:
                       "test": "response"
                   }
 
+                  # Execute
                   events = []
                   async for event in self.agent._run_async_impl(
                       self.mock_context