Python: Add long-running agents and background responses support (#3808)

* Python: Add long-running agents and background responses support - Add ContinuationToken TypedDict to core types - Add continuation_token field to ChatResponse, ChatResponseUpdate, AgentResponse, and AgentResponseUpdate - Add background and continuation_token options to OpenAIResponsesOptions - Implement polling via responses.retrieve() and streaming resumption in RawOpenAIResponsesClient - Propagate continuation tokens through agent run() and map_chat_to_agent_update - Fix streaming telemetry 'Failed to detach context' error in both ChatTelemetryLayer and AgentTelemetryLayer by avoiding trace.use_span() context attachment for async-managed spans - Add 14 unit tests for continuation token types and background flows - Add background_responses sample showing polling and stream resumption Fixes #2478 * Python: Add A2A long-running task support via ContinuationToken - Make ContinuationToken provider-agnostic (total=False, optional task_id/context_id fields) - Add background param to A2AAgent.run() controlling token emission - Add poll_task() for single-request task state retrieval - Add resubscribe support via continuation_token param on run() - Extract _updates_from_task() and _map_a2a_stream() for cleaner code - Streamline run()/streaming by removing intermediate _stream_updates wrapper - Update A2A sample to show background=False (default) with link to background_responses sample - Remove stale BareAgent from __all__ - Add 12 new A2A continuation token tests * fix logic for overriding continuation token when done * refactored ContinuationToken setup
2026-06-16 21:04:09 +08:00 · 2026-02-10 21:37:43 +01:00
parent 32ba81e990
commit 35097d8c75
12 changed files with 916 additions and 127 deletions
@@ -29,6 +29,7 @@ from agent_framework import (
 from agent_framework.a2a import A2AAgent
 from pytest import fixture, raises

+from agent_framework_a2a import A2AContinuationToken
 from agent_framework_a2a._agent import _get_uri_data  # type: ignore


@@ -38,6 +39,8 @@ class MockA2AClient:
    def __init__(self) -> None:
        self.call_count: int = 0
        self.responses: list[Any] = []
+        self.resubscribe_responses: list[Any] = []
+        self.get_task_response: Task | None = None

    def add_message_response(self, message_id: str, text: str, role: str = "agent") -> None:
        """Add a mock Message response."""
@@ -80,6 +83,18 @@ class MockA2AClient:
        client_event = (task, update_event)
        self.responses.append(client_event)

+    def add_in_progress_task_response(
+        self,
+        task_id: str,
+        context_id: str = "test-context",
+        state: TaskState = TaskState.working,
+    ) -> None:
+        """Add a mock in-progress Task response (non-terminal)."""
+        status = TaskStatus(state=state, message=None)
+        task = Task(id=task_id, context_id=context_id, status=status)
+        client_event = (task, None)
+        self.responses.append(client_event)
+
    async def send_message(self, message: Any) -> AsyncIterator[Any]:
        """Mock send_message method that yields responses."""
        self.call_count += 1
@@ -88,6 +103,22 @@ class MockA2AClient:
            response = self.responses.pop(0)
            yield response

+    async def resubscribe(self, request: Any) -> AsyncIterator[Any]:
+        """Mock resubscribe method that yields responses."""
+        self.call_count += 1
+
+        for response in self.resubscribe_responses:
+            yield response
+        self.resubscribe_responses.clear()
+
+    async def get_task(self, request: Any) -> Task:
+        """Mock get_task method that returns a task."""
+        self.call_count += 1
+        if self.get_task_response is not None:
+            return self.get_task_response
+        msg = "No get_task response configured"
+        raise ValueError(msg)
+

@fixture
 def mock_a2a_client() -> MockA2AClient:
@@ -598,3 +629,158 @@ def test_a2a_agent_initialization_with_timeout_parameter() -> None:

        # Verify it's an httpx.Timeout object with our custom timeout applied to all components
        assert isinstance(timeout_arg, httpx.Timeout)
+
+
+# region Continuation Token Tests
+
+
+async def test_working_task_emits_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that a working (non-terminal) task yields an update with a continuation token when background=True."""
+    mock_a2a_client.add_in_progress_task_response("task-wip", context_id="ctx-1", state=TaskState.working)
+
+    response = await a2a_agent.run("Start long task", background=True)
+
+    assert isinstance(response, AgentResponse)
+    assert response.continuation_token is not None
+    assert response.continuation_token["task_id"] == "task-wip"
+    assert response.continuation_token["context_id"] == "ctx-1"
+
+
+async def test_submitted_task_emits_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that a submitted task yields a continuation token when background=True."""
+    mock_a2a_client.add_in_progress_task_response("task-sub", state=TaskState.submitted)
+
+    response = await a2a_agent.run("Submit task", background=True)
+
+    assert response.continuation_token is not None
+    assert response.continuation_token["task_id"] == "task-sub"
+
+
+async def test_input_required_task_emits_continuation_token(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that an input_required task yields a continuation token when background=True."""
+    mock_a2a_client.add_in_progress_task_response("task-input", state=TaskState.input_required)
+
+    response = await a2a_agent.run("Need input", background=True)
+
+    assert response.continuation_token is not None
+    assert response.continuation_token["task_id"] == "task-input"
+
+
+async def test_working_task_no_token_without_background(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that background=False (default) does not emit continuation tokens for in-progress tasks."""
+    mock_a2a_client.add_in_progress_task_response("task-fg", context_id="ctx-fg", state=TaskState.working)
+
+    response = await a2a_agent.run("Foreground task")
+
+    assert response.continuation_token is None
+
+
+async def test_completed_task_has_no_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that a completed task does not set a continuation token."""
+    mock_a2a_client.add_task_response("task-done", [{"id": "art-1", "content": "Result"}])
+
+    response = await a2a_agent.run("Quick task")
+
+    assert response.continuation_token is None
+    assert len(response.messages) == 1
+    assert response.messages[0].text == "Result"
+
+
+async def test_streaming_emits_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that streaming with background=True yields updates with continuation tokens."""
+    mock_a2a_client.add_in_progress_task_response("task-stream", context_id="ctx-s", state=TaskState.working)
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Stream task", 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-stream"
+    assert updates[0].continuation_token["context_id"] == "ctx-s"
+
+
+async def test_resume_via_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that run() with continuation_token uses resubscribe instead of send_message."""
+    # Set up the resubscribe response (completed task)
+    status = TaskStatus(state=TaskState.completed, message=None)
+    artifact = Artifact(
+        artifact_id="art-resume",
+        name="result",
+        parts=[Part(root=TextPart(text="Resumed result"))],
+    )
+    task = Task(id="task-resume", context_id="ctx-r", status=status, artifacts=[artifact])
+    mock_a2a_client.resubscribe_responses.append((task, None))
+
+    token = A2AContinuationToken(task_id="task-resume", context_id="ctx-r")
+    response = await a2a_agent.run(continuation_token=token)
+
+    assert isinstance(response, AgentResponse)
+    assert len(response.messages) == 1
+    assert response.messages[0].text == "Resumed result"
+    assert response.continuation_token is None
+
+
+async def test_resume_streaming_via_continuation_token(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that streaming run() with continuation_token and background=True uses resubscribe."""
+    # Still working
+    status_wip = TaskStatus(state=TaskState.working, message=None)
+    task_wip = Task(id="task-rs", context_id="ctx-rs", status=status_wip)
+    # Then completed
+    status_done = TaskStatus(state=TaskState.completed, message=None)
+    artifact = Artifact(
+        artifact_id="art-rs",
+        name="result",
+        parts=[Part(root=TextPart(text="Stream resumed"))],
+    )
+    task_done = Task(id="task-rs", context_id="ctx-rs", status=status_done, artifacts=[artifact])
+    mock_a2a_client.resubscribe_responses.extend([(task_wip, None), (task_done, None)])
+
+    token = A2AContinuationToken(task_id="task-rs", context_id="ctx-rs")
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run(stream=True, continuation_token=token, background=True):
+        updates.append(update)
+
+    # First update: in-progress with token, second: completed with content
+    assert len(updates) == 2
+    assert updates[0].continuation_token is not None
+    assert updates[0].continuation_token["task_id"] == "task-rs"
+    assert updates[1].continuation_token is None
+    assert updates[1].contents[0].text == "Stream resumed"
+
+
+async def test_poll_task_in_progress(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test poll_task returns continuation token when task is still in progress."""
+    status = TaskStatus(state=TaskState.working, message=None)
+    mock_a2a_client.get_task_response = Task(id="task-poll", context_id="ctx-p", status=status)
+
+    token = A2AContinuationToken(task_id="task-poll", context_id="ctx-p")
+    response = await a2a_agent.poll_task(token)
+
+    assert response.continuation_token is not None
+    assert response.continuation_token["task_id"] == "task-poll"
+
+
+async def test_poll_task_completed(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test poll_task returns result with no continuation token when task is complete."""
+    status = TaskStatus(state=TaskState.completed, message=None)
+    artifact = Artifact(
+        artifact_id="art-poll",
+        name="result",
+        parts=[Part(root=TextPart(text="Poll result"))],
+    )
+    mock_a2a_client.get_task_response = Task(
+        id="task-poll-done", context_id="ctx-pd", status=status, artifacts=[artifact]
+    )
+
+    token = A2AContinuationToken(task_id="task-poll-done", context_id="ctx-pd")
+    response = await a2a_agent.poll_task(token)
+
+    assert response.continuation_token is None
+    assert len(response.messages) == 1
+    assert response.messages[0].text == "Poll result"
+
+
+# endregion