Skip to content

Commit f0fb030

Browse files
marcinbelczewskimarcinbelczewski
committed
fix(server): keep follow-up message when the executor re-emits the current task
1 parent aa8e966 commit f0fb030

2 files changed

Lines changed: 197 additions & 17 deletions

File tree

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

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -334,14 +334,33 @@ async def _handle_task_event(
334334
async def _handle_initial_task(self, event: Task) -> None:
335335
existing_task = await self.active_task._task_manager.get_task()
336336
if existing_task:
337-
logger.error(
338-
'Task %s already exists. Ignoring task replacement.',
339-
self.active_task._task_id,
340-
)
337+
if self.message_to_save is None:
338+
logger.error(
339+
'Task %s already exists. Ignoring task replacement.',
340+
self.active_task._task_id,
341+
)
342+
else:
343+
await self._add_message_to_history(
344+
self.message_to_save, existing_task
345+
)
341346
else:
342347
await self.active_task._task_manager.save_task_event(event)
343348
self.message_to_save = None
344349

350+
async def _add_message_to_history(
351+
self, message_to_save: Message, task: Task
352+
) -> None:
353+
message_already_saved = any(
354+
message.message_id == message_to_save.message_id
355+
for message in task.history
356+
)
357+
if not message_already_saved:
358+
await self.active_task._task_manager.save_task_event(
359+
self.active_task._task_manager.update_with_message(
360+
message_to_save, task
361+
)
362+
)
363+
345364
async def _handle_task_modification_event(
346365
self,
347366
event: TaskStatusUpdateEvent
@@ -364,20 +383,9 @@ async def _handle_task_modification_event(
364383
)
365384

366385
if self.message_to_save is not None:
367-
message_already_saved = any(
368-
message.message_id == self.message_to_save.message_id
369-
for message in updated_task.history
386+
await self._add_message_to_history(
387+
self.message_to_save, updated_task
370388
)
371-
if not message_already_saved:
372-
updated_task = (
373-
self.active_task._task_manager.update_with_message(
374-
self.message_to_save,
375-
updated_task,
376-
)
377-
)
378-
await self.active_task._task_manager.save_task_event(
379-
updated_task
380-
)
381389
self.message_to_save = None
382390

383391
self.active_task._task_manager.context_id = event.context_id

‎tests/integration/test_scenarios.py‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1291,6 +1291,178 @@ def test_input_required_followup_across_per_rpc_event_loops():
12911291
assert result2.status.state == TaskState.TASK_STATE_COMPLETED
12921292

12931293

1294+
# Scenario: Streaming resumption that enqueues the current Task first
1295+
@pytest.mark.timeout(2.0)
1296+
@pytest.mark.asyncio
1297+
@pytest.mark.parametrize(
1298+
('task_enqueues', 'replacement_errors'),
1299+
[(1, 0), (2, 1)],
1300+
ids=['once', 'twice'],
1301+
)
1302+
async def test_scenario_streaming_resumption_enqueues_current_task(
1303+
caplog: pytest.LogCaptureFixture,
1304+
task_enqueues: int,
1305+
replacement_errors: int,
1306+
) -> None:
1307+
class CurrentTaskAgent(AgentExecutor):
1308+
async def execute(
1309+
self, context: RequestContext, event_queue: EventQueue
1310+
) -> None:
1311+
task = context.current_task or new_task_from_user_message(
1312+
context.message
1313+
)
1314+
for _ in range(task_enqueues if context.current_task else 1):
1315+
await event_queue.enqueue_event(task)
1316+
await event_queue.enqueue_event(
1317+
TaskStatusUpdateEvent(
1318+
task_id=task.id,
1319+
context_id=task.context_id,
1320+
status=TaskStatus(
1321+
state=TaskState.TASK_STATE_COMPLETED
1322+
if context.current_task
1323+
else TaskState.TASK_STATE_INPUT_REQUIRED
1324+
),
1325+
)
1326+
)
1327+
1328+
async def cancel(
1329+
self, context: RequestContext, event_queue: EventQueue
1330+
) -> None:
1331+
pass
1332+
1333+
handler = create_handler(CurrentTaskAgent(), use_legacy=False)
1334+
client = await create_client(
1335+
handler, agent_card=agent_card(), streaming=True
1336+
)
1337+
1338+
msg1 = Message(
1339+
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
1340+
)
1341+
it = client.send_message(
1342+
SendMessageRequest(
1343+
message=msg1,
1344+
configuration=SendMessageConfiguration(return_immediately=False),
1345+
)
1346+
)
1347+
events1 = [event async for event in it]
1348+
task_id = get_task_id(events1[0])
1349+
context_id = get_task_context_id(events1[0])
1350+
1351+
msg2 = Message(
1352+
task_id=task_id,
1353+
context_id=context_id,
1354+
message_id='msg-resume',
1355+
role=Role.ROLE_USER,
1356+
parts=[Part(text='here is input')],
1357+
)
1358+
it2 = client.send_message(
1359+
SendMessageRequest(
1360+
message=msg2,
1361+
configuration=SendMessageConfiguration(return_immediately=False),
1362+
)
1363+
)
1364+
events2 = [event async for event in it2]
1365+
1366+
assert [event.WhichOneof('payload') for event in events2] == [
1367+
*['task'] * task_enqueues,
1368+
'status_update',
1369+
]
1370+
assert [message.message_id for message in events2[0].task.history] == [
1371+
'msg-start',
1372+
'msg-resume',
1373+
]
1374+
validate_state(events2[-1], TaskState.TASK_STATE_COMPLETED)
1375+
1376+
final_task = await client.get_task(GetTaskRequest(id=task_id))
1377+
assert [message.message_id for message in final_task.history] == [
1378+
'msg-start',
1379+
'msg-resume',
1380+
]
1381+
error_logs = [
1382+
record.message
1383+
for record in caplog.records
1384+
if record.levelname == 'ERROR'
1385+
and 'Ignoring task replacement' in record.message
1386+
]
1387+
assert len(error_logs) == replacement_errors
1388+
1389+
1390+
# Scenario: Resending a follow-up message adds it to history once
1391+
@pytest.mark.timeout(2.0)
1392+
@pytest.mark.asyncio
1393+
@pytest.mark.parametrize(
1394+
'enqueue_current_task', [False, True], ids=['status_only', 'current_task']
1395+
)
1396+
async def test_scenario_resent_follow_up_message_saved_once(
1397+
enqueue_current_task: bool,
1398+
) -> None:
1399+
class InputRequiredAgent(AgentExecutor):
1400+
async def execute(
1401+
self, context: RequestContext, event_queue: EventQueue
1402+
) -> None:
1403+
task = context.current_task or new_task_from_user_message(
1404+
context.message
1405+
)
1406+
if enqueue_current_task or not context.current_task:
1407+
await event_queue.enqueue_event(task)
1408+
await event_queue.enqueue_event(
1409+
TaskStatusUpdateEvent(
1410+
task_id=task.id,
1411+
context_id=task.context_id,
1412+
status=TaskStatus(
1413+
state=TaskState.TASK_STATE_INPUT_REQUIRED
1414+
),
1415+
)
1416+
)
1417+
1418+
async def cancel(
1419+
self, context: RequestContext, event_queue: EventQueue
1420+
) -> None:
1421+
pass
1422+
1423+
handler = create_handler(InputRequiredAgent(), use_legacy=False)
1424+
client = await create_client(
1425+
handler, agent_card=agent_card(), streaming=True
1426+
)
1427+
1428+
msg1 = Message(
1429+
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
1430+
)
1431+
it = client.send_message(
1432+
SendMessageRequest(
1433+
message=msg1,
1434+
configuration=SendMessageConfiguration(return_immediately=False),
1435+
)
1436+
)
1437+
events1 = [event async for event in it]
1438+
task_id = get_task_id(events1[0])
1439+
context_id = get_task_context_id(events1[0])
1440+
1441+
msg2 = Message(
1442+
task_id=task_id,
1443+
context_id=context_id,
1444+
message_id='msg-resume',
1445+
role=Role.ROLE_USER,
1446+
parts=[Part(text='here is input')],
1447+
)
1448+
for _ in range(2):
1449+
it2 = client.send_message(
1450+
SendMessageRequest(
1451+
message=msg2,
1452+
configuration=SendMessageConfiguration(
1453+
return_immediately=False
1454+
),
1455+
)
1456+
)
1457+
_ = [event async for event in it2]
1458+
1459+
final_task = await client.get_task(GetTaskRequest(id=task_id))
1460+
assert [message.message_id for message in final_task.history] == [
1461+
'msg-start',
1462+
'msg-resume',
1463+
]
1464+
1465+
12941466
# Scenario: Auth required and side channel unblocking
12951467
# Migrated from: test_workflow_auth_required_side_channel in test_handler_comparison
12961468
@pytest.mark.timeout(2.0)

0 commit comments

Comments
 (0)