Skip to content

fix: enforce task state-machine transitions and serialize cancels - #1045

Open
ez-lbz wants to merge 2 commits into
a2aproject:mainfrom
ez-lbz:fix/state-machine-cancel-race
Open

fix: enforce task state-machine transitions and serialize cancels#1045
ez-lbz wants to merge 2 commits into
a2aproject:mainfrom
ez-lbz:fix/state-machine-cancel-race

Conversation

@ez-lbz

@ez-lbz ez-lbz commented Aug 10, 2026

Copy link
Copy Markdown

What changed

1. Enforce state-machine transitions in TaskManager

Problem: TaskManager.saveTaskEvent/process overwrote the persisted task status with whatever state the event carried, with no transition validation. A task in a terminal state (COMPLETED/FAILED/CANCELED/REJECTED) could be silently rewritten to a different state (e.g. COMPLETEDSUBMITTED) by a late, stale, or malformed event — including replicated events racing the local final event. Only Go partially blocks this today.

Fix (server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskManager.java):

  • Added validateStateTransition(currentState, newState, taskId) and call it in the status-update path (saveTaskEvent(TaskStatusUpdateEvent)) and the full-task path (saveTaskEvent(Task)).
  • Rule: a terminal state must not be overwritten by a different state (throws A2AServerException, which the event pipeline turns into an error to the client while preserving the persisted state). Re-arriving events carrying the same final state remain allowed, so replicated replays and idempotent retries keep working.
  • All transitions from non-terminal states stay permitted (SUBMITTED → WORKING → COMPLETED/FAILED/CANCELED, interrupted-state resume flows INPUT_REQUIRED/AUTH_REQUIRED → WORKING → COMPLETED), so normal AgentExecutor flows are unaffected.

Fix (server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java):

  • testRejectStatusUpdateOverwritingTerminalState, testRejectStatusUpdateToDifferentTerminalState, testRejectTaskEventOverwritingTerminalState — rejected transitions throw and the persisted terminal state is preserved.
  • testSameTerminalStateReplayAllowed — idempotent same-state replay still works.
  • testNormalStateFlowAllowed, testInterruptedStateResumeFlowAllowed — the standard flows keep working.

2. Serialize concurrent cancels per task

Problem: DefaultRequestHandler.onCancelTask performed a check-then-act sequence — read task → check isFinal() → invoke agentExecutor.cancel() — with no lock between the check and the act. Two concurrent cancels of the same task could both observe the pre-transition state and both "succeed", and a cancel could race a concurrent completion.

Fix (server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java):

  • Added a per-task lock registry (cancelLocks, a ConcurrentHashMap<String, Object> keyed by task ID) and moved the entire cancel body into synchronized (lock) via a doCancelTask helper. The second concurrent cancel now waits for the first to finish, observes the CANCELED terminal state, and fails with TaskNotCancelableError.
  • The state-machine validation additionally blocks a cancel from overwriting a task that a concurrent message/send completed first.

Fix (server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java):

  • testConcurrentCancelsAreSerialized — holds the first cancel inside agentExecutor.cancel(), asserts the second cancel blocks, then verifies the first succeeds with CANCELED and the second fails with TaskNotCancelableError.

Behavior change: (1) events attempting to change a terminal task's state are now rejected instead of silently overwriting the state; (2) concurrent cancels of the same task are serialized, so the second one gets TaskNotCancelableError instead of both succeeding.

2. Make the replicated queue manager parallel test deterministic

Problem: ReplicatedQueueManagerTest.testParallelReplicationBehavior was timing-dependent and failed intermittently in CI (observed counts 1, 2, 21 instead of the expected 25) and consistently locally (0 or 3). The replicated threads sent TASK_STATE_COMPLETED events; a COMPLETED event processed mid-stream finalizes the task and closes the queue, so overlapping normal enqueues no longer trigger replication and the final count assertion depends on thread interleaving.

Fix (extras/queue-manager-replicated/core/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java):

  • Replicated threads now send a non-terminal state (TASK_STATE_WORKING). The replication hook skips replicated events via isReplicated() regardless of state, so the test's intent (normal enqueues replicate, replicated events do not) is unchanged while the outcome is deterministic.

Testing

  • mvn -pl extras/queue-manager-replicated/core test -Dtest=ReplicatedQueueManagerTest15 tests run, 0 failures across 5 consecutive runs (previously failed 3/3 locally with the same command).

ez-lbz added 2 commits August 11, 2026 00:28
testParallelReplicationBehavior sent TASK_STATE_COMPLETED events from the
replicated threads. A COMPLETED event processed mid-stream finalizes the
task and closes the queue, so overlapping normal enqueues no longer
trigger replication and the final count assertion became timing-
dependent (observed 0/1/2/3/21 instead of 25, locally and in CI).

Use a non-terminal state for the replicated events; the replication hook
skips them via isReplicated() regardless of state, so the test's intent
(normal enqueues replicate, replicated events do not) is unchanged while
the outcome is now deterministic.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant