Python: handle streamed A2A update events (#4919)

* Python: handle streamed A2A update events

* Python: preserve terminal A2A artifacts during streaming

* Python: harden streamed A2A update event handling

* Python: simplify streamed A2A update guard

---------

Co-authored-by: sztoplover-bit <253473756+sztoplover-bit@users.noreply.github.com>
Co-authored-by: Giles Odigwe <79032838+giles17@users.noreply.github.com>
This commit is contained in:
perry
2026-04-02 23:20:47 +08:00
committed by GitHub
Unverified
parent 524c0216e4
commit 5f06b68535
2 changed files with 266 additions and 2 deletions
@@ -365,6 +365,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
)
all_updates: list[AgentResponseUpdate] = []
streamed_artifact_ids_by_task: dict[str, set[str]] = {}
async for item in a2a_stream:
if isinstance(item, A2AMessage):
# Process A2A Message
@@ -378,12 +379,21 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
all_updates.append(update)
yield update
elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], Task):
task, _update_event = item
for update in self._updates_from_task(
task, update_event = item
updates = self._updates_from_task(
task,
update_event=update_event,
background=background,
emit_intermediate=emit_intermediate,
streamed_artifact_ids=streamed_artifact_ids_by_task.get(task.id),
)
if isinstance(update_event, TaskArtifactUpdateEvent) and any(
update.raw_representation is update_event for update in updates
):
streamed_artifact_ids_by_task.setdefault(task.id, set()).add(update_event.artifact.artifact_id)
if task.status.state in TERMINAL_TASK_STATES:
streamed_artifact_ids_by_task.pop(task.id, None)
for update in updates:
all_updates.append(update)
yield update
else:
@@ -403,8 +413,10 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
self,
task: Task,
*,
update_event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None,
background: bool = False,
emit_intermediate: bool = False,
streamed_artifact_ids: set[str] | None = None,
) -> list[AgentResponseUpdate]:
"""Convert an A2A Task into AgentResponseUpdate(s).
@@ -418,8 +430,21 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
"""
status = task.status
if (
emit_intermediate
and update_event is not None
and (event_updates := self._updates_from_task_update_event(update_event))
):
return event_updates
if status.state in TERMINAL_TASK_STATES:
task_messages = self._parse_messages_from_task(task)
if task.artifacts is not None and streamed_artifact_ids:
task_messages = [
message
for message in task_messages
if getattr(message.raw_representation, "artifact_id", None) not in streamed_artifact_ids
]
if task_messages:
return [
AgentResponseUpdate(
@@ -431,6 +456,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
)
for message in task_messages
]
if task.artifacts is not None:
return []
return [AgentResponseUpdate(contents=[], role="assistant", response_id=task.id, raw_representation=task)]
if background and status.state in IN_PROGRESS_TASK_STATES:
@@ -467,6 +494,44 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
return []
def _updates_from_task_update_event(
self, update_event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent
) -> list[AgentResponseUpdate]:
"""Convert A2A task update events into streaming AgentResponseUpdates."""
if isinstance(update_event, TaskArtifactUpdateEvent):
contents = self._parse_contents_from_a2a(update_event.artifact.parts)
if not contents:
return []
return [
AgentResponseUpdate(
contents=contents,
role="assistant",
response_id=update_event.task_id,
message_id=update_event.artifact.artifact_id,
raw_representation=update_event,
)
]
if not isinstance(update_event, TaskStatusUpdateEvent):
return []
message = update_event.status.message
if message is None or not message.parts:
return []
contents = self._parse_contents_from_a2a(message.parts)
if not contents:
return []
return [
AgentResponseUpdate(
contents=contents,
role="assistant" if message.role == A2ARole.agent else "user",
response_id=update_event.task_id,
raw_representation=update_event,
)
]
@staticmethod
def _build_continuation_token(task: Task) -> A2AContinuationToken | None:
"""Build an A2AContinuationToken from an A2A Task if it is still in progress."""