Skip to content

Commit b594be3

Browse files
fix(server): keep follow-up message when the executor re-emits the current task
1 parent ddce6b8 commit b594be3

2 files changed

Lines changed: 84 additions & 5 deletions

File tree

‎src/a2a/server/agent_execution/active_task.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -274,13 +274,16 @@ async def _handle_task_event(
274274
async def _handle_initial_task(self, event: Task) -> None:
275275
existing_task = await self.active_task._task_manager.get_task()
276276
if existing_task:
277-
logger.error(
278-
'Task %s already exists. Ignoring task replacement.',
279-
self.active_task._task_id,
280-
)
277+
# A follow-up request may start by re-emitting the current task.
278+
# Keep message_to_save so the next update adds it to history.
279+
if self.message_to_save is None:
280+
logger.error(
281+
'Task %s already exists. Ignoring task replacement.',
282+
self.active_task._task_id,
283+
)
281284
else:
282285
await self.active_task._task_manager.save_task_event(event)
283-
self.message_to_save = None
286+
self.message_to_save = None
284287

285288
async def _handle_task_modification_event(
286289
self,

‎tests/integration/test_scenarios.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1190,6 +1190,82 @@ def test_input_required_followup_across_per_rpc_event_loops():
11901190
assert result2.status.state == TaskState.TASK_STATE_COMPLETED
11911191

11921192

1193+
# Scenario: Streaming resumption that enqueues the current Task first
1194+
@pytest.mark.timeout(2.0)
1195+
@pytest.mark.asyncio
1196+
async def test_scenario_streaming_resumption_enqueues_current_task(caplog):
1197+
class CurrentTaskAgent(AgentExecutor):
1198+
async def execute(
1199+
self, context: RequestContext, event_queue: EventQueue
1200+
):
1201+
task = context.current_task or new_task_from_user_message(
1202+
context.message
1203+
)
1204+
await event_queue.enqueue_event(task)
1205+
await event_queue.enqueue_event(
1206+
TaskStatusUpdateEvent(
1207+
task_id=task.id,
1208+
context_id=task.context_id,
1209+
status=TaskStatus(
1210+
state=TaskState.TASK_STATE_COMPLETED
1211+
if context.current_task
1212+
else TaskState.TASK_STATE_INPUT_REQUIRED
1213+
),
1214+
)
1215+
)
1216+
1217+
async def cancel(
1218+
self, context: RequestContext, event_queue: EventQueue
1219+
):
1220+
pass
1221+
1222+
handler = create_handler(CurrentTaskAgent(), use_legacy=False)
1223+
client = await create_client(
1224+
handler, agent_card=agent_card(), streaming=True
1225+
)
1226+
1227+
msg1 = Message(
1228+
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
1229+
)
1230+
it = client.send_message(
1231+
SendMessageRequest(
1232+
message=msg1,
1233+
configuration=SendMessageConfiguration(return_immediately=False),
1234+
)
1235+
)
1236+
events1 = [event async for event in it]
1237+
task_id = get_task_id(events1[0])
1238+
context_id = get_task_context_id(events1[0])
1239+
1240+
msg2 = Message(
1241+
task_id=task_id,
1242+
context_id=context_id,
1243+
message_id='msg-resume',
1244+
role=Role.ROLE_USER,
1245+
parts=[Part(text='here is input')],
1246+
)
1247+
it2 = client.send_message(
1248+
SendMessageRequest(
1249+
message=msg2,
1250+
configuration=SendMessageConfiguration(return_immediately=False),
1251+
)
1252+
)
1253+
events2 = [event async for event in it2]
1254+
1255+
assert events2[0].HasField('task')
1256+
validate_state(events2[-1], TaskState.TASK_STATE_COMPLETED)
1257+
1258+
final_task = await client.get_task(GetTaskRequest(id=task_id))
1259+
assert [message.message_id for message in final_task.history] == [
1260+
'msg-start',
1261+
'msg-resume',
1262+
]
1263+
assert not any(
1264+
'Ignoring task replacement' in record.message
1265+
for record in caplog.records
1266+
)
1267+
1268+
11931269
# Scenario: Auth required and side channel unblocking
11941270
# Migrated from: test_workflow_auth_required_side_channel in test_handler_comparison
11951271
@pytest.mark.timeout(2.0)

0 commit comments

Comments
 (0)