Skip to content

Commit 4e916d8

Browse files
fix(server): record FAILED when a task's terminal save fails (#1400)
# Description When the task store rejects the write of a task's terminal status (for example a dropped database connection during the COMPLETED save), `DefaultRequestHandler` never records FAILED: the task stays `TASK_STATE_WORKING` in the store, and clients polling `tasks/get` never see a terminal state. `TaskManager.save_task_event` applies a status or artifact event to the cached `_current_task` in place and only then calls the store. When `save()` raises, the cache already says COMPLETED. The consumer's failure handler, `_write_failed_status` in `active_task.py`, reads that cache through `task_manager.get_task()` and returns early for a terminal state, so the FAILED write it documents for this case never happens. Reproducer: an agent emits `Task`, `start_work()`, `complete()`; the store raises `ConnectionError` on the first COMPLETED write; then `tasks/get`. ``` # main (43d1eff) send raised: ConnectionError database connection lost tasks/get state: TASK_STATE_WORKING # this branch send raised: ConnectionError database connection lost tasks/get state: TASK_STATE_FAILED ``` ```mermaid flowchart TD E[status or artifact event] --> M{TaskManager.save_task_event} M -->|main| A[mutate cached _current_task in place] A --> S1[store.save raises] S1 --> C1[cache says COMPLETED, store says WORKING] C1 --> W1[_write_failed_status: cache is terminal, skip] W1 --> R1[task stays WORKING] M -->|this PR| B[apply event to a copy] B --> S2[store.save raises] S2 --> C2[cache still says WORKING, as persisted] C2 --> W2[_write_failed_status writes FAILED] W2 --> R2[task is FAILED] ``` Change, in `src/a2a/server/tasks/task_manager.py`: `save_task_event` copies the cached task and applies the event to the copy. `_save_task` already assigns `_current_task` only after `save()` returns, so a failed save leaves the cache as last persisted. The same change fixes `prev` for a `VersionedTaskStore`: it was the updated task itself rather than "the task as previously read" that `VersionedTaskStore.save` documents. Also: - The comment in `ActiveTask.cancel` that said ordinary status writes still update the shared object in place is updated; Task events and `update_with_message` still do. - Two existing tests in `test_default_request_handler.py` (legacy handler) relied on the in-place updates: - `test_on_message_send_non_blocking` compared `len(result.history)` with the stored task after polling. `result` was the cached object, so it kept growing after it was returned. It is now the task as returned, and the test checks that its history is a prefix of the stored history. - `test_on_message_send_error_does_not_hang` used an `AsyncMock(spec=TaskStore)` whose `get` returned an `AsyncMock` in place of a `Task`; it now returns `None` so the task is created and the first `save` raises as before. Behaviour notes: - `RequestContext.current_task` is no longer updated in place by later status and artifact events during a request; it is the task as read when the request started. Those in-place updates raced with the consumer anyway (`enqueue_event` returns before the event is applied). - With an unversioned store, a save that raised although it committed (the acknowledgement was lost) is now followed by a FAILED write over the committed state. Before this change the same store left the task WORKING when the write had not committed, which is the usual case. With `VersionedDatabaseTaskStore` the version check rejects the stale FAILED write. - Cost: one `Task.CopyFrom` per status or artifact event. The default path already copies the whole task per event (`_enqueue_to_subscribers`, `CopyingTaskStoreAdapter.save`, `DatabaseTaskStore` serialisation). Not in this PR: `update_with_message` still appends the user message to the cached task before its save (#1392 rewrites those lines in `active_task.py`). Tests: - `tests/server/tasks/test_task_manager.py`: `test_save_task_event_failed_save_keeps_cached_task` (status and artifact events) and `test_save_task_event_passes_previous_task_as_prev`. - `tests/server/request_handlers/test_default_request_handler_v2.py`: `test_failed_terminal_save_records_failed_task`, end to end through `on_message_send` with an `InMemoryTaskStore` that fails the first COMPLETED write. - All four fail on `main` (`assert 3 == 1`, artifact left in the cache, `assert 2 == 1`, `assert 2 == 4` for WORKING vs FAILED) and pass with the change. Checks, run on Windows 11 (CPU, Python 3.10.19); Linux CI will confirm: - `bash ./scripts/lint.sh`: ruff check, ruff format and `ty` pass (three existing unused `type: ignore` warnings in files this PR does not touch). - `uv run pytest tests/server/tasks/test_task_manager.py tests/server/request_handlers/test_default_request_handler.py tests/server/request_handlers/test_default_request_handler_v2.py tests/server/agent_execution`: 297 passed. - `uv run pytest --ignore=tests/integration/cross_version --ignore=tests/integration/test_scenarios.py --cov=src`: 2364 passed, 203 skipped, 4 xfailed; total coverage 93%. The 4 setup errors (`tests/integration/push_notifications`, `Can't pickle local object` under Windows spawn) happen on `main` here too. `tests/integration/cross_version` (uses `os.killpg`) and `tests/integration/test_scenarios.py` (hangs on Windows on `main`) were not run locally, so the Linux CI run is the check for those. - [x] Follow the [`CONTRIBUTING` Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md). - [x] Make your Pull Request title in the <https://www.conventionalcommits.org/> specification. - [x] Ensure the tests and linter pass - [x] Appropriate docs were updated (if necessary) Co-authored-by: Mateusz Iwaniuk <45974414+Iwaniukooo11@users.noreply.github.com>
1 parent ac48ce1 commit 4e916d8

5 files changed

Lines changed: 152 additions & 10 deletions

File tree

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -964,10 +964,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
964964
# Write a copy rather than mutating the task in place: get_task()
965965
# returns the shared _task_manager._current_task, which may already
966966
# have been yielded to a subscriber. Copying keeps THIS terminal
967-
# write off the yielded reference. It is not a file-wide guarantee:
968-
# ordinary status writes (save_task_event -> ensure_task -> CopyFrom)
969-
# still update the shared object in place. Making every write
970-
# copy-on-yield is the general property tracked in #1175 and #1191.
967+
# write off the yielded reference. Status and artifact events are
968+
# applied to a copy too (TaskManager.save_task_event), but a Task
969+
# event or update_with_message still writes the object it is given.
970+
# Making every write copy-on-yield is tracked in #1175 and #1191.
971971
updated = Task()
972972
updated.CopyFrom(task)
973973
updated.status.state = TaskState.TASK_STATE_CANCELED

‎src/a2a/server/tasks/task_manager.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,12 @@ async def save_task_event(
267267
self._pending_event = None
268268
return event
269269

270-
task: Task = await self.ensure_task(event)
270+
# Apply the event to a copy; _save_task caches it only once the store
271+
# accepts it. If the save raises, the cached task still matches the
272+
# store, so the failure path does not take an unsaved terminal state
273+
# for a saved one, and `prev` is the task as last read.
274+
task = Task()
275+
task.CopyFrom(await self.ensure_task(event))
271276

272277
if isinstance(event, TaskStatusUpdateEvent):
273278
logger.debug(

‎tests/server/request_handlers/test_default_request_handler.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1098,11 +1098,10 @@ async def test_on_message_send_non_blocking(agent_card):
10981098

10991099
assert task is not None
11001100
assert task.status.state == TaskState.TASK_STATE_COMPLETED
1101-
assert (
1102-
result.history
1103-
and task.history
1104-
and len(result.history) == len(task.history)
1105-
)
1101+
# The result is the task as it was returned, with its whole history; the
1102+
# agent's later status messages are added to the stored task only.
1103+
assert result.history
1104+
assert list(task.history[: len(result.history)]) == list(result.history)
11061105

11071106

11081107
@pytest.mark.asyncio
@@ -2941,6 +2940,7 @@ async def test_on_message_send_error_does_not_hang(agent_card):
29412940
"""Test that if the consumer raises an exception during blocking wait, the producer is cancelled and no deadlock occurs."""
29422941
agent = HelloWorldAgentExecutor()
29432942
task_store = AsyncMock(spec=TaskStore)
2943+
task_store.get.return_value = None
29442944
task_store.save.side_effect = RuntimeError('This is an Error!')
29452945

29462946
request_handler = DefaultRequestHandler(

‎tests/server/request_handlers/test_default_request_handler_v2.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,78 @@ async def test_status_validation_failure_returns_failed_task() -> None:
190190
await handler.aclose()
191191

192192

193+
class _CompletingAgent(AgentExecutor):
194+
async def execute(
195+
self, context: RequestContext, event_queue: EventQueue
196+
) -> None:
197+
assert context.message is not None
198+
await event_queue.enqueue_event(
199+
new_task_from_user_message(context.message)
200+
)
201+
updater = TaskUpdater(
202+
event_queue,
203+
str(context.task_id or ''),
204+
str(context.context_id or ''),
205+
)
206+
await updater.start_work()
207+
await updater.complete()
208+
209+
async def cancel(
210+
self, context: RequestContext, event_queue: EventQueue
211+
) -> None:
212+
pass
213+
214+
215+
class _FailCompletedSaveTaskStore(InMemoryTaskStore):
216+
"""Rejects the first COMPLETED write, as a lost database connection would."""
217+
218+
def __init__(self) -> None:
219+
super().__init__()
220+
self.failed_save = False
221+
222+
async def save(self, task: Task, context: ServerCallContext) -> None:
223+
if (
224+
task.status.state == TaskState.TASK_STATE_COMPLETED
225+
and not self.failed_save
226+
):
227+
self.failed_save = True
228+
raise ConnectionError('database connection lost')
229+
await super().save(task, context)
230+
231+
232+
@pytest.mark.asyncio
233+
async def test_failed_terminal_save_records_failed_task() -> None:
234+
store = _FailCompletedSaveTaskStore()
235+
handler = DefaultRequestHandlerV2(
236+
agent_executor=_CompletingAgent(),
237+
task_store=store,
238+
agent_card=create_default_agent_card(),
239+
)
240+
context = create_server_call_context()
241+
try:
242+
with pytest.raises(ConnectionError, match='database connection lost'):
243+
await asyncio.wait_for(
244+
handler.on_message_send(
245+
SendMessageRequest(
246+
message=Message(
247+
role=Role.ROLE_USER,
248+
message_id='failed-terminal-save',
249+
parts=[Part(text='test failed save')],
250+
)
251+
),
252+
context,
253+
),
254+
timeout=5,
255+
)
256+
assert store.failed_save
257+
listed = await store.list(ListTasksRequest(), context)
258+
assert len(listed.tasks) == 1
259+
# The COMPLETED write never landed, so the task must not stay WORKING.
260+
assert listed.tasks[0].status.state == TaskState.TASK_STATE_FAILED
261+
finally:
262+
await handler.aclose()
263+
264+
193265
def test_init_default_dependencies():
194266
"""Test that default dependencies are created if not provided."""
195267
agent_executor = MockAgentExecutor()

‎tests/server/tasks/test_task_manager.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import pytest
55

66
from a2a.auth.user import User
7+
from a2a.server.cluster import VersionedTaskStore
78
from a2a.server.context import ServerCallContext
89
from a2a.server.tasks import TaskManager
910
from a2a.server.tasks.task_manager import append_artifact_to_task
@@ -196,6 +197,70 @@ async def test_save_task_event_metadata_update(
196197
assert updated_task.metadata == new_metadata
197198

198199

200+
@pytest.mark.asyncio
201+
@pytest.mark.parametrize(
202+
'event',
203+
[
204+
TaskStatusUpdateEvent(
205+
task_id=MINIMAL_TASK_ID,
206+
context_id=MINIMAL_CONTEXT_ID,
207+
status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
208+
),
209+
TaskArtifactUpdateEvent(
210+
task_id=MINIMAL_TASK_ID,
211+
context_id=MINIMAL_CONTEXT_ID,
212+
artifact=Artifact(
213+
artifact_id='artifact-id', parts=[Part(text='chunk')]
214+
),
215+
),
216+
],
217+
ids=['status', 'artifact'],
218+
)
219+
async def test_save_task_event_failed_save_keeps_cached_task(
220+
task_manager: TaskManager,
221+
mock_task_store: AsyncMock,
222+
event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
223+
) -> None:
224+
"""A rejected write leaves the cached task as it was last persisted."""
225+
mock_task_store.get.return_value = create_minimal_task()
226+
await task_manager.get_task()
227+
mock_task_store.save.side_effect = ConnectionError('connection lost')
228+
229+
with pytest.raises(ConnectionError):
230+
await task_manager.save_task_event(event)
231+
232+
task = await task_manager.get_task()
233+
assert task is not None
234+
assert task.status.state == TaskState.TASK_STATE_SUBMITTED
235+
assert not task.artifacts
236+
237+
238+
@pytest.mark.asyncio
239+
async def test_save_task_event_passes_previous_task_as_prev() -> None:
240+
"""A versioned store receives the task as it was before the event."""
241+
store = AsyncMock(spec=VersionedTaskStore)
242+
task_manager = TaskManager(
243+
task_id=MINIMAL_TASK_ID,
244+
context_id=MINIMAL_CONTEXT_ID,
245+
task_store=store,
246+
initial_message=None,
247+
context=TEST_CONTEXT,
248+
)
249+
await task_manager.save_task_event(create_minimal_task())
250+
251+
await task_manager.save_task_event(
252+
TaskStatusUpdateEvent(
253+
task_id=MINIMAL_TASK_ID,
254+
context_id=MINIMAL_CONTEXT_ID,
255+
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
256+
)
257+
)
258+
259+
call = store.save.call_args
260+
assert call.args[0].status.state == TaskState.TASK_STATE_WORKING
261+
assert call.kwargs['prev'].status.state == TaskState.TASK_STATE_SUBMITTED
262+
263+
199264
@pytest.mark.asyncio
200265
async def test_ensure_task_existing(
201266
task_manager: TaskManager, mock_task_store: AsyncMock

0 commit comments

Comments
 (0)