fix: enforce task state-machine transitions and serialize cancels

[ez-lbz]

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. COMPLETED → SUBMITTED) 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=ReplicatedQueueManagerTest — 15 tests run, 0 failures across 5 consecutive runs (previously failed 3/3 locally with the same command).

[ehsavoie]

Trace the full call chain:

TaskManager.java:161-191 — process() handles A2AError by synthesizing a TASK_STATE_FAILED event and calling saveTaskEvent:
TaskStatusUpdateEvent failedEvent = ...status(new TaskStatus(TASK_STATE_FAILED)).build();
isFinal = saveTaskEvent(failedEvent, isReplicated, taskSnapshot); // line 191

TaskManager.java:101 — saveTaskEvent(TaskStatusUpdateEvent,...) calls the new guard:
validateStateTransition(currentState, newState, event.taskId());

If the task is already COMPLETED (or any other terminal state), validateStateTransition at line 266-276 throws A2AServerException. Before this PR, it would silently overwrite COMPLETED with FAILED.

MainEventBusProcessor.java:343-347 — updateTaskStore() catches this as the generic catch (Exception e) block (because A2AServerException is not InternalError, TaskSerializationException, or TaskPersistenceException):
} catch (Exception e) {
// Unexpected exception type - treat as permanent failure
throw new InternalError("TaskStore persistence failed: " + e.getMessage());
}

MainEventBusProcessor.java:230-234 — back in processEvent(), that InternalError is caught and set as the event to distribute to clients:
} catch (InternalError e) {
LOGGER.error("Failed to persist event for task {}, distributing error to clients", taskId, e);
eventToDistribute = e; // clients receive this InternalError instead of the original A2AError

Net effect: a client that sent a message to an already-completed task, whose agent then emits an A2AError, now receives a generic InternalError("TaskStore persistence failed: Task X is already in terminal state COMPLETED") instead of the original A2AError. The message is confusing
because the store didn't fail — the state machine rejected the transition.

The cleaner fix would be inside TaskManager.process() itself at line 161: before synthesizing the FAILED event, check if the task is already in a terminal state and skip the update (just return true):
} else if (event instanceof A2AError) {
// ... existing null checks ...
if (errorContextId != null) {
Task existing = getTask();
TaskState currentState = existing != null && existing.status() != null ? existing.status().state() : null;
if (currentState != null && currentState.isFinal()) {
// Task already terminal — no state update needed, A2AError still signals finality
return true;
}
// ... synthesize FAILED event as before

This would prevent the A2AServerException from propagating at all, and clients would still receive the original A2AError (since process() returns normally, eventToDistribute stays as the original event in processEvent()).

[ehsavoie]

Some helper method would help reduce the duplication code on creating those

[ehsavoie]

I think this should be guarded as well: if the task is already COMPLETED then it go wild

[ehsavoie]

This would set the Task to TASK_STATE_FAILED, but this could lead to a new Error because of the guard if the task is in any terminal state