Skip to content

fix(server): keep follow-up message when the executor re-emits the current task #1298

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

Open
marcinbelczewski wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(server): keep follow-up message when the executor re-emits the current task #1298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
42 changes: 25 additions & 17 deletions src/a2a/server/agent_execution/active_task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,14 +334,33 @@ async def _handle_task_event(
async def _handle_initial_task(self, event: Task) -> None:
existing_task = await self.active_task._task_manager.get_task()
if existing_task:
logger.error(
'Task %s already exists. Ignoring task replacement.',
self.active_task._task_id,
)
if self.message_to_save is None:

@ConnorMoss02 ConnorMoss02 Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ran this locally with an agent like the one in your test. Two things.

  1. The leading Task on the follow-up stream still has the old history: ['msg-start']. msg-resume only lands with the next status update. The final get_task is right, but the first streamed Task is what 3.1.2 is about.
  2. On a follow-up message_to_save is still set for every Task, so enqueueing current_task twice drops the second one with no ERROR.

Saving the message here fixes both. The first Task then has both messages, the double enqueue logs once, and tests/integration plus tests/server pass:

if existing_task:
    if self.message_to_save is None:
        logger.error(
            'Task %s already exists. Ignoring task replacement.',
            self.active_task._task_id,
        )
    else:
        await self.active_task._task_manager.save_task_event(
            self.active_task._task_manager.update_with_message(
                self.message_to_save, existing_task
            )
        )
else:
    await self.active_task._task_manager.save_task_event(event)
self.message_to_save = None

@marcinbelczewski marcinbelczewski Oct 31, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ConnorMoss02 thanks a lot for taking the time to actually run it and for the patch! You are right on both points and sorry for that. My test only checked that the stream starts with a Task and what ends up in history at the end, but never what is inside that first Task, which is the whole point of 3.1.2... And the claim that a second Task still logs the ERROR I based on test_scenario_enqueue_task_twice, which only covers a brand new task, not a follow-up. My bad.

I applied your approach with one small change: the follow-up message now goes through the same message_id check that _handle_task_modification_event already had (from #1106). I moved it to _add_message_to_history so both paths use it. Without it, if the executor enqueues current_task and the client resends the same follow-up message (same message_id), the message lands in history twice, while for an executor that only sends status updates it's saved once. Let me know if you'd rather keep it exactly as you proposed.

Tests:

  • test_scenario_streaming_resumption_enqueues_current_task is now parametrized with once/twice. It checks the history of the first streamed Task and that the ERROR is logged exactly once when the Task is enqueued twice.
  • New test_scenario_resent_follow_up_message_saved_once checks that both executor patterns save a resent message only once.

I checked that the streaming test fails on the previous version and the resend one fails with your patch as is. I also broke each part of the change on purpose to make sure some test catches it. Squashed into a single commit and force pushed.

Thanks again for catching this!

@marcinbelczewski marcinbelczewski Oct 31, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more thing I noticed after posting the above. The first streamed Task is right only if the executor enqueues context.current_task itself (as in the migration guide), because that's the same object TaskManager keeps. If it enqueues a copy, history is fine but the first streamed Task still has only ['msg-start']. Subscribers get the raw executor event here:

a2a-python/src/a2a/server/agent_execution/active_task.py

Lines 280 to 286 in aa8e966

handled_event = updated_task if isinstance(event, Task) else event
if updated_task is not None and handled_event is not None:
await self._update_task_state(updated_task, handled_event)
self.active_task._task_created.set()
await self._enqueue_to_subscribers(event, updated_task)

Looks like a leftover from #1053. Before it the consumer sent the stored task to subscribers, but maybe I'm missing something. Passing handled_event there fixes it and all tests pass, but it changes what's streamed for every Task event, so I'd rather not sneak it into this PR. Should I add it here or open a separate issue?

@marcinbelczewski marcinbelczewski Oct 31, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small update: the test now also has an agent message in the INPUT_REQUIRED status, to check the user 1 → agent → user 2 order.

logger.error(
'Task %s already exists. Ignoring task replacement.',
self.active_task._task_id,
)
else:
await self._add_message_to_history(
self.message_to_save, existing_task
)
else:
await self.active_task._task_manager.save_task_event(event)
self.message_to_save = None

async def _add_message_to_history(
self, message_to_save: Message, task: Task
) -> None:
message_already_saved = any(
message.message_id == message_to_save.message_id
for message in task.history
)
if not message_already_saved:
await self.active_task._task_manager.save_task_event(
self.active_task._task_manager.update_with_message(
message_to_save, task
)
)

async def _handle_task_modification_event(
self,
event: TaskStatusUpdateEvent
Expand All @@ -364,20 +383,9 @@ async def _handle_task_modification_event(
)

if self.message_to_save is not None:
message_already_saved = any(
message.message_id == self.message_to_save.message_id
for message in updated_task.history
await self._add_message_to_history(
self.message_to_save, updated_task
)
if not message_already_saved:
updated_task = (
self.active_task._task_manager.update_with_message(
self.message_to_save,
updated_task,
)
)
await self.active_task._task_manager.save_task_event(
updated_task
)
self.message_to_save = None

self.active_task._task_manager.context_id = event.context_id
Expand Down
182 changes: 182 additions & 0 deletions tests/integration/test_scenarios.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1291,6 +1291,188 @@ def test_input_required_followup_across_per_rpc_event_loops():
assert result2.status.state == TaskState.TASK_STATE_COMPLETED


# Scenario: Streaming resumption that enqueues the current Task first
@pytest.mark.timeout(2.0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
('task_enqueues', 'replacement_errors'),
[(1, 0), (2, 1)],
ids=['once', 'twice'],
)
async def test_scenario_streaming_resumption_enqueues_current_task(
caplog: pytest.LogCaptureFixture,
task_enqueues: int,
replacement_errors: int,
) -> None:
class CurrentTaskAgent(AgentExecutor):
async def execute(
self, context: RequestContext, event_queue: EventQueue
) -> None:
task = context.current_task or new_task_from_user_message(
context.message
)
for _ in range(task_enqueues if context.current_task else 1):
await event_queue.enqueue_event(task)
if context.current_task:
status = TaskStatus(state=TaskState.TASK_STATE_COMPLETED)
else:
status = TaskStatus(
state=TaskState.TASK_STATE_INPUT_REQUIRED,
message=Message(
message_id='agent-ask',
role=Role.ROLE_AGENT,
parts=[Part(text='which one?')],
),
)
await event_queue.enqueue_event(
TaskStatusUpdateEvent(
task_id=task.id,
context_id=task.context_id,
status=status,
)
)

async def cancel(
self, context: RequestContext, event_queue: EventQueue
) -> None:
pass

handler = create_handler(CurrentTaskAgent(), use_legacy=False)
client = await create_client(
handler, agent_card=agent_card(), streaming=True
)

msg1 = Message(
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
)
it = client.send_message(
SendMessageRequest(
message=msg1,
configuration=SendMessageConfiguration(return_immediately=False),
)
)
events1 = [event async for event in it]
task_id = get_task_id(events1[0])
context_id = get_task_context_id(events1[0])

msg2 = Message(
task_id=task_id,
context_id=context_id,
message_id='msg-resume',
role=Role.ROLE_USER,
parts=[Part(text='here is input')],
)
it2 = client.send_message(
SendMessageRequest(
message=msg2,
configuration=SendMessageConfiguration(return_immediately=False),
)
)
events2 = [event async for event in it2]

assert [event.WhichOneof('payload') for event in events2] == [
*['task'] * task_enqueues,
'status_update',
]
assert [message.message_id for message in events2[0].task.history] == [
'msg-start',
'agent-ask',
'msg-resume',
]
assert not events2[0].task.status.HasField('message')
validate_state(events2[-1], TaskState.TASK_STATE_COMPLETED)

final_task = await client.get_task(GetTaskRequest(id=task_id))
assert [message.message_id for message in final_task.history] == [
'msg-start',
'agent-ask',
'msg-resume',
]
error_logs = [
record.message
for record in caplog.records
if record.levelname == 'ERROR'
and 'Ignoring task replacement' in record.message
]
assert len(error_logs) == replacement_errors


# Scenario: Resending a follow-up message adds it to history once
@pytest.mark.timeout(2.0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
'enqueue_current_task', [False, True], ids=['status_only', 'current_task']
)
async def test_scenario_resent_follow_up_message_saved_once(
enqueue_current_task: bool,
) -> None:
class InputRequiredAgent(AgentExecutor):
async def execute(
self, context: RequestContext, event_queue: EventQueue
) -> None:
task = context.current_task or new_task_from_user_message(
context.message
)
if enqueue_current_task or not context.current_task:
await event_queue.enqueue_event(task)
await event_queue.enqueue_event(
TaskStatusUpdateEvent(
task_id=task.id,
context_id=task.context_id,
status=TaskStatus(
state=TaskState.TASK_STATE_INPUT_REQUIRED
),
)
)

async def cancel(
self, context: RequestContext, event_queue: EventQueue
) -> None:
pass

handler = create_handler(InputRequiredAgent(), use_legacy=False)
client = await create_client(
handler, agent_card=agent_card(), streaming=True
)

msg1 = Message(
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
)
it = client.send_message(
SendMessageRequest(
message=msg1,
configuration=SendMessageConfiguration(return_immediately=False),
)
)
events1 = [event async for event in it]
task_id = get_task_id(events1[0])
context_id = get_task_context_id(events1[0])

msg2 = Message(
task_id=task_id,
context_id=context_id,
message_id='msg-resume',
role=Role.ROLE_USER,
parts=[Part(text='here is input')],
)
for _ in range(2):
it2 = client.send_message(
SendMessageRequest(
message=msg2,
configuration=SendMessageConfiguration(
return_immediately=False
),
)
)
_ = [event async for event in it2]

final_task = await client.get_task(GetTaskRequest(id=task_id))
assert [message.message_id for message in final_task.history] == [
'msg-start',
'msg-resume',
]


# Scenario: Auth required and side channel unblocking
# Migrated from: test_workflow_auth_required_side_channel in test_handler_comparison
@pytest.mark.timeout(2.0)
Expand Down
Loading