Repository navigation
fix(server): save a resumed turn's message with its first event - #1392
Draft
mykytanetipa wants to merge 1 commit into
Draft
mykytanetipa wants to merge 1 commit into
mykytanetipa wants to merge 1 commit into
Conversation
When a task waiting for input is resumed, the consumer added the user's message to the history with a separate save before applying the agent's first event. That save wrote the task with its previous state, and in cluster mode it also wrote a Task(INPUT_REQUIRED) row to the shared event log. A SubscribeToTask on another replica treated that row as the end of the stream and missed the rest of the turn. The message is now added in memory and persisted by the save of the event itself, so a resumed turn writes no stale state. If that save conflicts, the message is added again when the event is retried on the reloaded task.
🧪 Code Coverage (vs
|
| Base | PR | Delta | |
|---|---|---|---|
| src/a2a/server/agent_execution/active_task.py | 95.92% | 95.91% | 🔴 -0.01% |
| src/a2a/server/cluster/database_event_stream.py | 96.94% | 92.86% | 🔴 -4.08% |
| src/a2a/server/cluster/database_task_store.py | 96.55% | 99.31% | 🟢 +2.76% |
| src/a2a/server/events/event_queue_v2.py | 91.79% | 92.27% | 🟢 +0.48% |
| src/a2a/server/request_handlers/default_request_handler_v2.py | 87.68% | 88.54% | 🟢 +0.86% |
| src/a2a/utils/telemetry.py | 90.70% | 91.47% | 🟢 +0.78% |
| Total | 93.15% | 93.20% | 🟢 +0.05% |
Generated by coverage-comment.yml
4 tasks done
Iwaniukooo11
added a commit
that referenced
this pull request
Oct 30, 2026
# 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>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resuming a task that waits for input wrote the task once with its old state, which a subscriber on another replica took as the end of the stream. Found while working on #1324.
Example:
sequenceDiagram participant C as Client participant A as Replica A (runs agent) participant S as Shared store participant B as Replica B Note over S: T is INPUT_REQUIRED C->>A: SendMessage(T, follow-up) C->>B: SubscribeToTask(T) A->>S: save T + message, still INPUT_REQUIRED S-->>B: event log row Task(INPUT_REQUIRED) B-->>C: stream closed (interrupted state) A->>S: WORKING, artifact, COMPLETED Note over C: never sees the rest of the turnFix: the follow-up message is added to the task in memory and persisted by the save of the agent's first event, so a resumed turn no longer writes its previous state. If that save conflicts with another replica, the message is added again when the event is retried on the reloaded task, so it is not lost. In the example B follows the turn to completion, and the message is recorded once.
Tested: a remote subscriber follows a resumed turn to completion, and the message is recorded once, including after a conflict; full suite and Postgres/MySQL tests pass.