Python: Fix A2AAgent to surface message content from in-progress TaskStatusUpdateEvents (#4798)

* Fix A2AAgent dropping message content from in-progress TaskStatusUpdateEvents (#4783) _updates_from_task() returned [] for working-state tasks when background=False, silently discarding all intermediate message content from task.status.message. Now extracts and yields message parts from in-progress status updates during streaming. Also fixed MockA2AClient.send_message to yield all queued responses (enabling multi-event streaming tests) and added text parameter to add_in_progress_task_response for tests that need status messages. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix: gate intermediate status updates behind emit_intermediate flag and add missing test coverage - Add emit_intermediate parameter to _updates_from_task and _map_a2a_stream - Thread stream flag from run() so only streaming callers see intermediate updates - Add IN_PROGRESS_TASK_STATES guard to emit_intermediate condition - Add role parameter to test helper add_in_progress_task_response - Add clarifying comment on MockA2AClient.send_message batch semantics - Add tests for user role mapping, background precedence, non-streaming behavior, terminal task with no artifacts, and empty parts edge case Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-03-25 19:08:47 -07:00
parent 0bdcaa5c07
commit c1435ac201
2 changed files with 202 additions and 10 deletions
@@ -91,9 +91,18 @@ class MockA2AClient:
        task_id: str,
        context_id: str = "test-context",
        state: TaskState = TaskState.working,
+        text: str | None = None,
+        role: A2ARole = A2ARole.agent,
    ) -> None:
        """Add a mock in-progress Task response (non-terminal)."""
-        status = TaskStatus(state=state, message=None)
+        message = None
+        if text is not None:
+            message = A2AMessage(
+                message_id=str(uuid4()),
+                role=role,
+                parts=[Part(root=TextPart(text=text))],
+            )
+        status = TaskStatus(state=state, message=message)
        task = Task(id=task_id, context_id=context_id, status=status)
        client_event = (task, None)
        self.responses.append(client_event)
@@ -102,9 +111,10 @@ class MockA2AClient:
        """Mock send_message method that yields responses."""
        self.call_count += 1

-        if self.responses:
-            response = self.responses.pop(0)
+        # All queued responses are delivered as a single streaming batch per call.
+        for response in self.responses:
            yield response
+        self.responses.clear()

    async def resubscribe(self, request: Any) -> AsyncIterator[Any]:
        """Mock resubscribe method that yields responses."""
@@ -1039,3 +1049,144 @@ async def test_run_with_continuation_token_does_not_require_messages(mock_a2a_cl


 # endregion
+
+# region Streaming with in-progress message content
+
+
+async def test_streaming_working_updates_yield_message_content(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that streaming working updates with status.message yield content."""
+    mock_a2a_client.add_in_progress_task_response("task-w", context_id="ctx-w", text="Processing step 1...")
+    mock_a2a_client.add_in_progress_task_response("task-w", context_id="ctx-w", text="Processing step 2...")
+    mock_a2a_client.add_task_response("task-w", [{"id": "art-w", "content": "Final result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 3
+    assert updates[0].contents[0].text == "Processing step 1..."
+    assert updates[1].contents[0].text == "Processing step 2..."
+    assert updates[2].contents[0].text == "Final result"
+
+
+async def test_streaming_single_working_update_with_message(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a single working update with message content is not dropped."""
+    mock_a2a_client.add_in_progress_task_response("task-s", context_id="ctx-s", text="Thinking...")
+    mock_a2a_client.add_task_response("task-s", [{"id": "art-s", "content": "Done"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "Thinking..."
+    assert updates[0].role == "assistant"
+    assert updates[1].contents[0].text == "Done"
+
+
+async def test_streaming_working_update_without_message_is_skipped(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that working updates without status.message are still silently skipped."""
+    mock_a2a_client.add_in_progress_task_response("task-n", context_id="ctx-n")
+    mock_a2a_client.add_task_response("task-n", [{"id": "art-n", "content": "Result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].contents[0].text == "Result"
+
+
+async def test_streaming_working_update_user_role_mapping(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that A2ARole.user in status message maps to role='user'."""
+    mock_a2a_client.add_in_progress_task_response("task-u", context_id="ctx-u", text="User echo", role=A2ARole.user)
+    mock_a2a_client.add_task_response("task-u", [{"id": "art-u", "content": "Done"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "User echo"
+    assert updates[0].role == "user"
+
+
+async def test_background_with_status_message_yields_continuation_token(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that background=True takes precedence over status message content."""
+    mock_a2a_client.add_in_progress_task_response("task-bg", context_id="ctx-bg", text="Should be ignored")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True, background=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].continuation_token is not None
+    assert updates[0].continuation_token["task_id"] == "task-bg"
+    assert updates[0].contents == []
+
+
+async def test_non_streaming_does_not_surface_intermediate_messages(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that run(stream=False) does not include intermediate status messages."""
+    mock_a2a_client.add_in_progress_task_response("task-ns", context_id="ctx-ns", text="Intermediate")
+    mock_a2a_client.add_task_response("task-ns", [{"id": "art-ns", "content": "Final"}])
+
+    response = await a2a_agent.run("Hello")
+
+    assert len(response.messages) == 1
+    assert response.messages[0].text == "Final"
+
+
+async def test_terminal_no_artifacts_after_working_with_content(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a terminal task with no artifacts after working-state messages does not re-emit the working content."""
+    mock_a2a_client.add_in_progress_task_response("task-t", context_id="ctx-t", text="Working on it...")
+    # Terminal task with no artifacts and no history
+    status = TaskStatus(state=TaskState.completed, message=None)
+    task = Task(id="task-t", context_id="ctx-t", status=status)
+    mock_a2a_client.responses.append((task, None))
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "Working on it..."
+    # Terminal task with no artifacts yields an empty-contents update
+    assert updates[1].contents == []
+
+
+async def test_streaming_working_update_with_empty_parts_is_skipped(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a working update with status.message but empty parts list is skipped."""
+    # Construct a message with an empty parts list (distinct from message=None)
+    message = A2AMessage(
+        message_id=str(uuid4()),
+        role=A2ARole.agent,
+        parts=[],
+    )
+    status = TaskStatus(state=TaskState.working, message=message)
+    task = Task(id="task-ep", context_id="ctx-ep", status=status)
+    mock_a2a_client.responses.append((task, None))
+    mock_a2a_client.add_task_response("task-ep", [{"id": "art-ep", "content": "Result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].contents[0].text == "Result"
+
+
+# endregion