Skip to content

Python: Sramsey/python continuation token fix #9002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Eduard van Valkenburg (eavanvalkenburg) merged 8 commits into from
Oct 9, 2026
Merged

Python: Sramsey/python continuation token fix #9002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
bf7f2e7
Python: add terminal state check for continuation token
sophia-ramsey Oct 3, 2026
9dd12d4
Python: update continuation token in background
sophia-ramsey Oct 3, 2026
df1a3e8
replace is_operation_terminal with _OperationState
sophia-ramsey Oct 8, 2026
c4312ca
Potential fix for pull request finding 'Unused import'
eavanvalkenburg Oct 8, 2026
809b2f3
Merge branch 'main' into sramsey/python-continuation-token-fix
eavanvalkenburg Oct 8, 2026
aea1a64
Restore Generic import for response classes
Copilot Oct 8, 2026
d545fcc
error handling and serialization fix
sophia-ramsey Oct 8, 2026
46d13cd
fix default event status
sophia-ramsey Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 149 additions & 76 deletions python/packages/a2a/agent_framework_a2a/_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
prepend_agent_framework_to_user_agent,
)
from agent_framework._telemetry import mark_feature_used
from agent_framework._types import AgentRunInputs
from agent_framework._types import AgentRunInputs, _set_operation_state # pyright: ignore[reportPrivateUsage]
from agent_framework.exceptions import AgentInvalidRequestException
from agent_framework.observability import AgentTelemetryLayer
from google.protobuf.json_format import MessageToDict
Expand All @@ -66,6 +66,15 @@ class A2AServiceSessionId(TypedDict):
task_state: TaskState | None


def _with_operation_state(
update: AgentResponseUpdate,
state: Literal["in_progress", "terminal"],
) -> AgentResponseUpdate:
"""Attach framework-internal resumable-operation state."""
_set_operation_state(update, state)
return update


class A2AAgentSession(AgentSession):
"""Session for A2A-based agents.

Expand Down Expand Up @@ -680,18 +689,21 @@ async def _map_a2a_stream(
all_updates.append(update)
yield update
elif is_terminal or is_input_required:
if updates:
if updates and any(update.contents for update in updates):
# Terminal/input-required event with content — discard accumulated intermediates
pending_updates_by_task.pop(status_event.task_id, None)
for update in updates:
all_updates.append(update)
yield update
elif is_terminal:
# Terminal event with NO content — flush accumulated updates
# Terminal event with no content — flush accumulated updates before its metadata marker.
pending = pending_updates_by_task.pop(status_event.task_id, [])
for update in pending:
all_updates.append(update)
yield update
for update in updates:
all_updates.append(update)
yield update
else:
# Non-streaming intermediate: accumulate for later
if updates:
Expand All @@ -702,7 +714,7 @@ async def _map_a2a_stream(
last_task_id = artifact_event.task_id
if artifact_event.context_id:
last_context_id = artifact_event.context_id
updates = self._updates_from_task_update_event(artifact_event)
updates = self._updates_from_task_update_event(artifact_event, background=background)
# Always yield artifact updates — they carry actual response
# content (files, data). Track IDs so that a subsequent
# terminal Task doesn't duplicate the same artifacts.
Expand Down Expand Up @@ -814,22 +826,21 @@ def _updates_from_task(
if status.state == TaskState.TASK_STATE_INPUT_REQUIRED:
message = status.message if status.HasField("message") and status.message.parts else None
occurrence_id = input_request_occurrence_id or task.id
return [
AgentResponseUpdate(
contents=[
self._input_required_request(
task.id,
message,
fallback_occurrence_id=occurrence_id,
)
],
role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user",
response_id=task.id,
continuation_token=self._build_continuation_token(task) if background else None,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
)
]
update = AgentResponseUpdate(
contents=[
self._input_required_request(
task.id,
message,
fallback_occurrence_id=occurrence_id,
)
],
role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user",
response_id=task.id,
continuation_token=self._build_continuation_token(task) if background else None,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
)
return [_with_operation_state(update, "in_progress")] if background else [update]

if status.state in TERMINAL_TASK_STATES:
task_messages = self._parse_messages_from_task(task)
Expand All @@ -840,7 +851,7 @@ def _updates_from_task(
if getattr(message.raw_representation, "artifact_id", None) not in streamed_artifact_ids
]
if task_messages:
return [
updates = [
AgentResponseUpdate(
contents=message.contents,
role=message.role,
Expand All @@ -853,28 +864,49 @@ def _updates_from_task(
)
for message in task_messages
]
_set_operation_state(updates[-1], "terminal")
return updates
if task.artifacts:
return []
if not background:
return []
return [
_with_operation_state(
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
),
"terminal",
)
]
return [
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
_with_operation_state(
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
),
"terminal",
)
]

if background and 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,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
_with_operation_state(
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=task.id,
continuation_token=token,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
),
"in_progress",
)
]

Expand All @@ -888,7 +920,18 @@ def _updates_from_task(
contents = self._parse_contents_from_a2a(status.message.parts)
if contents:
return [
AgentResponseUpdate(
_with_operation_state(
AgentResponseUpdate(
contents=contents,
role="assistant" if status.message.role == A2ARole.ROLE_AGENT else "user",
response_id=task.id,
additional_properties={"a2a_metadata": task_metadata} if task_metadata else None,
raw_representation=task,
),
"in_progress",
)
if background
else AgentResponseUpdate(
contents=contents,
role="assistant" if status.message.role == A2ARole.ROLE_AGENT else "user",
response_id=task.id,
Expand All @@ -914,16 +957,15 @@ def _updates_from_task_update_event(
artifact_meta = MessageToDict(update_event.artifact.metadata) if update_event.artifact.metadata else {}
event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {}
merged_metadata = {**artifact_meta, **event_meta} or None
return [
AgentResponseUpdate(
contents=contents,
role="assistant",
response_id=update_event.task_id,
message_id=update_event.artifact.artifact_id,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
)
]
update = AgentResponseUpdate(
contents=contents,
role="assistant",
response_id=update_event.task_id,
message_id=update_event.artifact.artifact_id,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
)
return [_with_operation_state(update, "in_progress")] if background else [update]

if not isinstance(update_event, TaskStatusUpdateEvent):
return []
Expand All @@ -944,47 +986,78 @@ def _updates_from_task_update_event(
event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {}
merged_metadata = {**message_meta, **event_meta} or None
occurrence_id = input_request_occurrence_id or update_event.task_id
return [
AgentResponseUpdate(
contents=[
self._input_required_request(
update_event.task_id,
message,
fallback_occurrence_id=occurrence_id,
)
],
role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user",
response_id=update_event.task_id,
message_id=message.message_id if message is not None else None,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
)
]

if not update_event.status.HasField("message") or not update_event.status.message.parts:
return []
update = AgentResponseUpdate(
contents=[
self._input_required_request(
update_event.task_id,
message,
fallback_occurrence_id=occurrence_id,
)
],
role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user",
response_id=update_event.task_id,
message_id=message.message_id if message is not None else None,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
)
return [_with_operation_state(update, "in_progress")] if background else [update]

if state not in TERMINAL_TASK_STATES:
return []

event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {}
if not update_event.status.HasField("message") or not update_event.status.message.parts:
if not background:
return []
return [
_with_operation_state(
AgentResponseUpdate(
contents=[],
role="assistant",
response_id=update_event.task_id,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": event_meta} if event_meta else None,
raw_representation=update_event,
),
"terminal",
)
]

message = update_event.status.message
contents = self._parse_contents_from_a2a(message.parts)
if not contents:
return []
if not background:
return []
return [
_with_operation_state(
AgentResponseUpdate(
contents=[],
role="assistant" if message.role == A2ARole.ROLE_AGENT else "user",
response_id=update_event.task_id,
message_id=message.message_id,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": event_meta} if event_meta else None,
raw_representation=update_event,
),
"terminal",
)
]
msg_meta = MessageToDict(message.metadata) if message.metadata else {}
event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {}
merged_metadata = {**msg_meta, **event_meta} or None

return [
AgentResponseUpdate(
contents=contents,
role="assistant" if message.role == A2ARole.ROLE_AGENT else "user",
response_id=update_event.task_id,
message_id=message.message_id,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
_with_operation_state(
AgentResponseUpdate(
contents=contents,
role="assistant" if message.role == A2ARole.ROLE_AGENT else "user",
response_id=update_event.task_id,
message_id=message.message_id,
continuation_token=continuation_token,
additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None,
raw_representation=update_event,
),
"terminal",
)
]

Expand Down
Loading
Loading