agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Eduard van Valkenburg
2026-02-10 21:37:43 +01:00
committed by GitHub
Unverified
parent 32ba81e990
commit 35097d8c75
12 changed files with 916 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/a2a/agent_framework_a2a/__init__.py
+2 -1
View File
@@ -2,7 +2,7 @@
import importlib.metadata
from ._agent import A2AAgent
from ._agent import A2AAgent, A2AContinuationToken
try:
__version__ = importlib.metadata.version(__name__)
@@ -11,5 +11,6 @@ except importlib.metadata.PackageNotFoundError:
__all__ = [
"A2AAgent",
"A2AContinuationToken",
"__version__",
]
python/packages/a2a/agent_framework_a2a/_agent.py
+126 -76
View File
@@ -20,6 +20,8 @@ from a2a.types import (
FileWithUri,
Message,
Task,
TaskIdParams,
TaskQueryParams,
TaskState,
TextPart,
TransportProtocol,
@@ -34,21 +36,39 @@ from agent_framework import (
BaseAgent,
ChatMessage,
Content,
ContinuationToken,
ResponseStream,
normalize_messages,
prepend_agent_framework_to_user_agent,
)
from agent_framework.observability import AgentTelemetryLayer
__all__ = ["A2AAgent"]
__all__ = ["A2AAgent", "A2AContinuationToken"]
URI_PATTERN = re.compile(r"^data:(?P<media_type>[^;]+);base64,(?P<base64_data>[A-Za-z0-9+/=]+)$")
class A2AContinuationToken(ContinuationToken):
"""Continuation token for A2A protocol long-running tasks."""
task_id: str
"""A2A protocol task ID."""
context_id: str
"""A2A protocol context ID."""
TERMINAL_TASK_STATES = [
TaskState.completed,
TaskState.failed,
TaskState.canceled,
TaskState.rejected,
]
IN_PROGRESS_TASK_STATES = [
TaskState.submitted,
TaskState.working,
TaskState.input_required,
TaskState.auth_required,
]
def _get_uri_data(uri: str) -> str:
@@ -193,6 +213,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@@ -203,6 +225,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
*,
stream: Literal[True],
thread: AgentThread | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
@@ -212,85 +236,62 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
*,
stream: bool = False,
thread: AgentThread | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a response from the agent.
This method returns the final result of the agent's execution
as a single AgentResponse object when stream=False. When stream=True,
it returns a ResponseStream that yields AgentResponseUpdate objects.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
continuation_token: Optional token to resume a long-running task
instead of starting a new one.
background: When True, in-progress task updates surface continuation
tokens so the caller can poll or resubscribe later. When False
(default), the agent internally waits for the task to complete.
kwargs: Additional keyword arguments.
Returns:
When stream=False: An Awaitable[AgentResponse].
When stream=True: A ResponseStream of AgentResponseUpdate items.
"""
if continuation_token is not None:
a2a_stream: AsyncIterable[Any] = self.client.resubscribe(TaskIdParams(id=continuation_token["task_id"]))
else:
normalized_messages = normalize_messages(messages)
a2a_message = self._prepare_message_for_a2a(normalized_messages[-1])
a2a_stream = self.client.send_message(a2a_message)
response = ResponseStream(
self._map_a2a_stream(a2a_stream, background=background),
finalizer=lambda updates: AgentResponse.from_updates(list(updates)),
)
if stream:
return self._run_stream_impl(messages=messages, thread=thread, **kwargs)
return self._run_impl(messages=messages, thread=thread, **kwargs)
return response
return response.get_final_response()
async def _run_impl(
async def _map_a2a_stream(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
a2a_stream: AsyncIterable[Any],
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse[Any]:
"""Non-streaming implementation of run."""
# Collect all updates and use framework to consolidate updates into response
updates: list[AgentResponseUpdate] = []
async for update in self._stream_updates(messages, thread=thread, **kwargs):
updates.append(update)
return AgentResponse.from_updates(updates)
def _run_stream_impl(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Streaming implementation of run."""
def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
return AgentResponse.from_updates(list(updates))
return ResponseStream(self._stream_updates(messages, thread=thread, **kwargs), finalizer=_finalize)
async def _stream_updates(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
background: bool = False,
) -> AsyncIterable[AgentResponseUpdate]:
"""Internal method to stream updates from the A2A agent.
"""Map raw A2A protocol items to AgentResponseUpdates.
Args:
messages: The message(s) to send to the agent.
a2a_stream: The raw A2A event stream.
Keyword Args:
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Yields:
AgentResponseUpdate items from the A2A agent.
background: When False, in-progress task updates are silently
consumed (the stream keeps iterating until a terminal state).
When True, they are yielded with a continuation token.
"""
normalized_messages = normalize_messages(messages)
a2a_message = self._prepare_message_for_a2a(normalized_messages[-1])
response_stream = self.client.send_message(a2a_message)
async for item in response_stream:
async for item in a2a_stream:
if isinstance(item, Message):
# Process A2A Message
contents = self._parse_contents_from_a2a(item.parts)
yield AgentResponseUpdate(
contents=contents,
@@ -300,33 +301,82 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
)
elif isinstance(item, tuple) and len(item) == 2: # ClientEvent = (Task, UpdateEvent)
task, _update_event = item
if isinstance(task, Task) and task.status.state in TERMINAL_TASK_STATES:
# Convert Task artifacts to ChatMessages and yield as separate updates
task_messages = self._parse_messages_from_task(task)
if task_messages:
for message in task_messages:
# Use the artifact's ID from raw_representation as message_id for unique identification
artifact_id = getattr(message.raw_representation, "artifact_id", None)
yield AgentResponseUpdate(
contents=message.contents,
role=message.role,
response_id=task.id,
message_id=artifact_id,
raw_representation=task,
)
else:
# Empty task
yield AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
raw_representation=task,
)
if isinstance(task, Task):
for update in self._updates_from_task(task, background=background):
yield update
else:
# Unknown response type
msg = f"Only Message and Task responses are supported from A2A agents. Received: {type(item)}"
raise NotImplementedError(msg)
# ------------------------------------------------------------------
# Task helpers
# ------------------------------------------------------------------
def _updates_from_task(self, task: Task, *, background: bool = False) -> list[AgentResponseUpdate]:
"""Convert an A2A Task into AgentResponseUpdate(s).
Terminal tasks produce updates from their artifacts/history.
In-progress tasks produce a continuation token update only when
``background=True``; otherwise they are silently skipped so the
caller keeps consuming the stream until completion.
"""
if task.status.state in TERMINAL_TASK_STATES:
task_messages = self._parse_messages_from_task(task)
if task_messages:
return [
AgentResponseUpdate(
contents=message.contents,
role=message.role,
response_id=task.id,
message_id=getattr(message.raw_representation, "artifact_id", None),
raw_representation=task,
)
for message in task_messages
]
return [AgentResponseUpdate(contents=[], role="assistant", response_id=task.id, raw_representation=task)]
if background and task.status.state in IN_PROGRESS_TASK_STATES:
token = self._build_continuation_token(task)
return [
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
continuation_token=token,
raw_representation=task,
)
]
return []
@staticmethod
def _build_continuation_token(task: Task) -> A2AContinuationToken | None:
"""Build an A2AContinuationToken from an A2A Task if it is still in progress."""
if task.status.state in IN_PROGRESS_TASK_STATES:
return A2AContinuationToken(task_id=task.id, context_id=task.context_id)
return None
async def poll_task(self, continuation_token: A2AContinuationToken) -> AgentResponse[Any]:
"""Poll for the current state of a long-running A2A task.
Unlike ``run(continuation_token=...)``, which resubscribes to the SSE
stream, this performs a single request to retrieve the task state.
Args:
continuation_token: A token previously obtained from a response's
``continuation_token`` field.
Returns:
An AgentResponse whose ``continuation_token`` is set when the task
is still in progress, or ``None`` when it has reached a terminal state.
"""
task_id = continuation_token["task_id"]
task = await self.client.get_task(TaskQueryParams(id=task_id))
updates = self._updates_from_task(task, background=True)
if updates:
return AgentResponse.from_updates(updates)
return AgentResponse(messages=[], response_id=task.id, raw_representation=task)
def _prepare_message_for_a2a(self, message: ChatMessage) -> A2AMessage:
"""Prepare a ChatMessage for the A2A protocol.