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

[ez-lbz]

Thanks for the detailed call-chain trace — that's a real regression from the state-machine guard. Fixed as you suggested: in TaskManager.process(), the A2AError branch now checks whether the task is already terminal before synthesizing the TASK_STATE_FAILED event; if so, it skips the state update and lets the A2AError signal finality (returns true). The exception no longer propagates, so clients receive the original A2AError instead of the misleading "TaskStore persistence failed" InternalError.

Added regression tests for both paths (terminal task + non-terminal task). TaskManagerTest 36 passed.

[ehsavoie]

1. Unbounded cancelLocks map (Medium)

The problem: cancelLocks at DefaultRequestHandler.java:270 grows monotonically — computeIfAbsent adds an entry per canceled task, but there's no corresponding remove anywhere. By contrast, runningAgents (line 258) is explicitly cleaned up at lines 744 and 908.

Why it matters practically: The Javadoc (line 267) argues it's "bounded by the number of distinct tasks that have ever been canceled, in the same way the in-memory task store is unbounded." That's true today with InMemoryTaskStore, but with the JPA-backed task-store-database-jpa
extra, the task store is not unbounded in memory — tasks live in the DB. In that configuration, cancelLocks becomes the only in-memory structure that grows without bound.

The fix is straightforward. After doCancelTask returns (line 485) — whether successfully or via TaskNotCancelableError — the task is guaranteed to be in a terminal state. No future cancel can succeed (the state-machine guard in TaskManager.validateStateTransition will reject it
even without the lock). So the lock entry is no longer needed:

synchronized (cancelLock) {
try {
return doCancelTask(params, context);
} finally {
cancelLocks.remove(params.id());
}
}

The finally ensures cleanup on both the success path and when doCancelTask throws TaskNotCancelableError (task was already terminal). A brief window exists where a third concurrent cancel could create a new lock object after the remove, but that's safe — it will just immediately
see the terminal state and throw.

---

2. Duplicate terminal-state guard logic (Medium)

The three locations:

1. TaskManager.java:78 — saveTaskEvent(Task, ...) calls validateStateTransition
2. TaskManager.java:101 — saveTaskEvent(TaskStatusUpdateEvent, ...) calls validateStateTransition
3. TaskManager.java:191-197 — process(A2AError, ...) has its own inline check:
   Task existingTask = getTask();
   TaskState currentState = existingTask != null && existingTask.status() != null
   ? existingTask.status().state() : null;
   if (currentState != null && currentState.isFinal()) { ... }

Why they diverge: Locations 1 and 2 reject the transition by throwing. Location 3 silently skips it and sets isFinal = true. The skip is intentional — an A2AError arriving for a terminal task should not throw, because the error itself still needs to propagate to clients. But the
terminal-state detection logic (the null-check chain + isFinal() call) is duplicated.

The risk: If the definition of "terminal" ever changes (e.g., a new TASK_STATE_ARCHIVED that is final but allows certain transitions), you'd need to update both validateStateTransition and the inline check. They're in the same file, but the inline check at line 191 doesn't call
validateStateTransition or even reference it, so the connection is invisible.

A possible consolidation: Extract a isInTerminalState() query method (or just use getTask() + isFinal() consistently) and restructure process() to attempt the saveTaskEvent and catch the A2AServerException:

if (errorContextId != null) {
TaskStatusUpdateEvent failedEvent = TaskStatusUpdateEvent.builder()
.taskId(taskId).contextId(errorContextId)
.status(new TaskStatus(TASK_STATE_FAILED)).build();
try {
isFinal = saveTaskEvent(failedEvent, isReplicated, taskSnapshot);
} catch (A2AServerException e) {
// Task already terminal — the state-machine guard rejected FAILED.
// The A2AError still signals finality to clients.
LOGGER.debug("A2AError for task {} already in terminal state — skipping state update", taskId);
isFinal = true;
}
}

This eliminates the duplicate check entirely — validateStateTransition inside saveTaskEvent is the single authority. The downside is using exception flow for a non-exceptional path, but this is a rare edge case (error arriving after completion), not a hot path.