Skip to content

Commit bd0ecc6

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

2 files changed

Lines changed: 207 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: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1291,6 +1291,188 @@ 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+
if context.current_task:
1317+
status = TaskStatus(state=TaskState.TASK_STATE_COMPLETED)
1318+
else:
1319+
status = TaskStatus(
1320+
state=TaskState.TASK_STATE_INPUT_REQUIRED,
1321+
message=Message(
1322+
message_id='agent-ask',
1323+
role=Role.ROLE_AGENT,
1324+
parts=[Part(text='which one?')],
1325+
),
1326+
)
1327+
await event_queue.enqueue_event(
1328+
TaskStatusUpdateEvent(
1329+
task_id=task.id,
1330+
context_id=task.context_id,
1331+
status=status,
1332+
)
1333+
)
1334+
1335+
async def cancel(
1336+
self, context: RequestContext, event_queue: EventQueue
1337+
) -> None:
1338+
pass
1339+
1340+
handler = create_handler(CurrentTaskAgent(), use_legacy=False)
1341+
client = await create_client(
1342+
handler, agent_card=agent_card(), streaming=True
1343+
)
1344+
1345+
msg1 = Message(
1346+
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
1347+
)
1348+
it = client.send_message(
1349+
SendMessageRequest(
1350+
message=msg1,
1351+
configuration=SendMessageConfiguration(return_immediately=False),
1352+
)
1353+
)
1354+
events1 = [event async for event in it]
1355+
task_id = get_task_id(events1[0])
1356+
context_id = get_task_context_id(events1[0])
1357+
1358+
msg2 = Message(
1359+
task_id=task_id,
1360+
context_id=context_id,
1361+
message_id='msg-resume',
1362+
role=Role.ROLE_USER,
1363+
parts=[Part(text='here is input')],
1364+
)
1365+
it2 = client.send_message(
1366+
SendMessageRequest(
1367+
message=msg2,
1368+
configuration=SendMessageConfiguration(return_immediately=False),
1369+
)
1370+
)
1371+
events2 = [event async for event in it2]
1372+
1373+
assert [event.WhichOneof('payload') for event in events2] == [
1374+
*['task'] * task_enqueues,
1375+
'status_update',
1376+
]
1377+
assert [message.message_id for message in events2[0].task.history] == [
1378+
'msg-start',
1379+
'agent-ask',
1380+
'msg-resume',
1381+
]
1382+
assert not events2[0].task.status.HasField('message')
1383+
validate_state(events2[-1], TaskState.TASK_STATE_COMPLETED)
1384+
1385+
final_task = await client.get_task(GetTaskRequest(id=task_id))
1386+
assert [message.message_id for message in final_task.history] == [
1387+
'msg-start',
1388+
'agent-ask',
1389+
'msg-resume',
1390+
]
1391+
error_logs = [
1392+
record.message
1393+
for record in caplog.records
1394+
if record.levelname == 'ERROR'
1395+
and 'Ignoring task replacement' in record.message
1396+
]
1397+
assert len(error_logs) == replacement_errors
1398+
1399+
1400+
# Scenario: Resending a follow-up message adds it to history once
1401+
@pytest.mark.timeout(2.0)
1402+
@pytest.mark.asyncio
1403+
@pytest.mark.parametrize(
1404+
'enqueue_current_task', [False, True], ids=['status_only', 'current_task']
1405+
)
1406+
async def test_scenario_resent_follow_up_message_saved_once(
1407+
enqueue_current_task: bool,
1408+
) -> None:
1409+
class InputRequiredAgent(AgentExecutor):
1410+
async def execute(
1411+
self, context: RequestContext, event_queue: EventQueue
1412+
) -> None:
1413+
task = context.current_task or new_task_from_user_message(
1414+
context.message
1415+
)
1416+
if enqueue_current_task or not context.current_task:
1417+
await event_queue.enqueue_event(task)
1418+
await event_queue.enqueue_event(
1419+
TaskStatusUpdateEvent(
1420+
task_id=task.id,
1421+
context_id=task.context_id,
1422+
status=TaskStatus(
1423+
state=TaskState.TASK_STATE_INPUT_REQUIRED
1424+
),
1425+
)
1426+
)
1427+
1428+
async def cancel(
1429+
self, context: RequestContext, event_queue: EventQueue
1430+
) -> None:
1431+
pass
1432+
1433+
handler = create_handler(InputRequiredAgent(), use_legacy=False)
1434+
client = await create_client(
1435+
handler, agent_card=agent_card(), streaming=True
1436+
)
1437+
1438+
msg1 = Message(
1439+
message_id='msg-start', role=Role.ROLE_USER, parts=[Part(text='start')]
1440+
)
1441+
it = client.send_message(
1442+
SendMessageRequest(
1443+
message=msg1,
1444+
configuration=SendMessageConfiguration(return_immediately=False),
1445+
)
1446+
)
1447+
events1 = [event async for event in it]
1448+
task_id = get_task_id(events1[0])
1449+
context_id = get_task_context_id(events1[0])
1450+
1451+
msg2 = Message(
1452+
task_id=task_id,
1453+
context_id=context_id,
1454+
message_id='msg-resume',
1455+
role=Role.ROLE_USER,
1456+
parts=[Part(text='here is input')],
1457+
)
1458+
for _ in range(2):
1459+
it2 = client.send_message(
1460+
SendMessageRequest(
1461+
message=msg2,
1462+
configuration=SendMessageConfiguration(
1463+
return_immediately=False
1464+
),
1465+
)
1466+
)
1467+
_ = [event async for event in it2]
1468+
1469+
final_task = await client.get_task(GetTaskRequest(id=task_id))
1470+
assert [message.message_id for message in final_task.history] == [
1471+
'msg-start',
1472+
'msg-resume',
1473+
]
1474+
1475+
12941476
# Scenario: Auth required and side channel unblocking
12951477
# Migrated from: test_workflow_auth_required_side_channel in test_handler_comparison
12961478
@pytest.mark.timeout(2.0)

0 commit comments

Comments
 (0)