fix: enforce task state-machine transitions and serialize cancels

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

@@ -370,7 +370,12 @@ public void onTaskFinalized(String tid) {
                         TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder()
                                 .taskId(taskId)  // Use same taskId as queue
                                 .contextId("test-context")
-                                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                                // Use a non-terminal state: a COMPLETED event processed mid-stream
+                                // finalizes the task and closes the queue, so overlapping normal
+                                // enqueues no longer trigger replication and the count assertion
+                                // becomes timing-dependent (flaky). Replicated events are skipped by
+                                // the replication hook via isReplicated() regardless of state.
+                                .status(new TaskStatus(TaskState.TASK_STATE_WORKING))
                                 .build();
                         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, event);
                         queueManager.onReplicatedEvent(replicatedEvent);

server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java

@@ -257,6 +257,18 @@ public class DefaultRequestHandler implements RequestHandler {
 
     private final ConcurrentMap<String, CompletableFuture<Void>> runningAgents = new ConcurrentHashMap<>();
 
+    /**
+     * Per-task lock registry serializing {@link #onCancelTask} check-then-act sequences.
+     * <p>
+     * The cancel path reads the task, checks it is not already in a terminal state, and
+     * then invokes the agent executor to enqueue the CANCELED event. Without a lock, two
+     * concurrent cancels (or a cancel racing a concurrent completion) could both observe
+     * the pre-transition state and both act on it. Entries are retained for the
+     * lifetime of the JVM; the registry is bounded by the number of distinct tasks that
+     * have ever been canceled, in the same way the in-memory task store is unbounded.
+     */
+    private final ConcurrentMap<String, Object> cancelLocks = new ConcurrentHashMap<>();
+
 
     private Executor executor;
     private Executor eventConsumerExecutor;
@@ -466,6 +478,24 @@ public ListTasksResult onListTasks(ListTasksParams params, ServerCallContext con
 
     @Override
     public Task onCancelTask(CancelTaskParams params, ServerCallContext context) throws A2AError {
+        // Serialize check-then-act per task so two concurrent cancels (or a cancel racing
+        // a concurrent completion) cannot both act on the pre-transition state.
+        Object cancelLock = cancelLocks.computeIfAbsent(params.id(), k -> new Object());
+        synchronized (cancelLock) {
+            try {
+                return doCancelTask(params, context);
+            } finally {
+                // After doCancelTask returns (successfully or via TaskNotCancelableError),
+                // the task is in a terminal state, so no future cancel can succeed — the
+                // state-machine guard rejects it even without the lock. Drop the entry to
+                // avoid unbounded growth of the map. The 2-arg remove keeps a concurrently
+                // created newer lock entry (for a future, already-terminal task) intact.
+                cancelLocks.remove(params.id(), cancelLock);
+            }
+        }
+    }
+
+    private Task doCancelTask(CancelTaskParams params, ServerCallContext context) throws A2AError {
         Task task = taskStore.get(params.id());
         if (task == null) {
             throw new TaskNotFoundError();

server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskManager.java

@@ -19,6 +19,7 @@
 import org.a2aproject.sdk.spec.Message;
 import org.a2aproject.sdk.spec.Task;
 import org.a2aproject.sdk.spec.TaskArtifactUpdateEvent;
+import org.a2aproject.sdk.spec.TaskState;
 import org.a2aproject.sdk.spec.TaskStatus;
 import org.a2aproject.sdk.spec.TaskStatusUpdateEvent;
 import org.jspecify.annotations.Nullable;
@@ -69,6 +70,13 @@ boolean saveTaskEvent(Task task, boolean isReplicated) throws A2AServerException
     boolean saveTaskEvent(Task task, boolean isReplicated, @Nullable AtomicReference<Task> taskSnapshot)
             throws A2AServerException {
         checkIdsAndUpdateIfNecessary(task.id(), task.contextId());
+        // Defensive state-machine check: a task that already reached a terminal state must
+        // not be overwritten by a task snapshot carrying a different state.
+        Task current = getTask();
+        if (current != null && current.status() != null && current.status().state() != null
+                && task.status() != null && task.status().state() != null) {
+            validateStateTransition(current.status().state(), task.status().state(), task.id());
+        }
         Task savedTask = saveTask(task, isReplicated);
         if (taskSnapshot != null) {
             taskSnapshot.set(savedTask);
@@ -85,6 +93,12 @@ boolean saveTaskEvent(TaskStatusUpdateEvent event, boolean isReplicated, @Nullab
         checkIdsAndUpdateIfNecessary(event.taskId(), event.contextId());
         Task task = ensureTask(event.taskId(), event.contextId());
 
+        // State-machine validation: reject transitions that would overwrite a terminal
+        // state with a different state. Re-arriving events carrying the same
+        // final state remain allowed (idempotent replays / replication).
+        TaskState currentState = task.status() != null ? task.status().state() : null;
+        TaskState newState = event.status() != null ? event.status().state() : null;
+        validateStateTransition(currentState, newState, event.taskId());
 
         Task.Builder builder = Task.builder(task)
                 .status(event.status());
@@ -174,7 +188,16 @@ public boolean process(Event event, boolean isReplicated, @Nullable AtomicRefere
                         .contextId(errorContextId)
                         .status(new TaskStatus(TASK_STATE_FAILED))
                         .build();
-                isFinal = saveTaskEvent(failedEvent, isReplicated, taskSnapshot);
+                try {
+                    isFinal = saveTaskEvent(failedEvent, isReplicated, taskSnapshot);
+                } catch (A2AServerException e) {
+                    // Task already in a terminal state: the state-machine guard in
+                    // validateStateTransition rejected the synthesized FAILED event
+                    // (terminal -> FAILED). No state update is needed — the A2AError
+                    // itself still signals finality to clients.
+                    LOGGER.debug("A2AError for task {} already in terminal state - skipping state update", taskId);
+                    isFinal = true;
+                }
             } else {
                 // Can't update status without contextId, but error is still terminal
                 LOGGER.debug("A2AError event for task {} without contextId - skipping state update", taskId);
@@ -232,6 +255,36 @@ private Task ensureTask(String eventTaskId, String eventContextId) {
         return task;
     }
 
+    /**
+     * Validates a task state transition before it is persisted.
+     * <p>
+     * A terminal (final) state must not be overwritten by a <em>different</em> state:
+     * once a task is {@code COMPLETED}/{@code FAILED}/{@code CANCELED}/{@code REJECTED}
+     * it stays in that state. Events re-arriving with the <em>same</em> final state are
+     * allowed, so replicated replays and idempotent retries keep working.
+     * <p>
+     * Transitions from any non-terminal state to any state are permitted (e.g.
+     * SUBMITTED → WORKING → COMPLETED/FAILED/CANCELED, interrupted-state resume flows),
+     * matching the transitions the A2A spec and the reference agents exercise.
+     *
+     * @param currentState the task's current state, or {@code null} if unknown
+     * @param newState     the state requested by the event, or {@code null} if unknown
+     * @param taskId       the task identifier, used in the error message
+     * @throws A2AServerException if the transition would overwrite a terminal state
+     */
+    private static void validateStateTransition(@Nullable TaskState currentState, @Nullable TaskState newState,
+            String taskId) throws A2AServerException {
+        if (currentState == null || newState == null) {
+            return;
+        }
+        if (currentState.isFinal() && currentState != newState) {
+            throw new A2AServerException(
+                    "Task " + taskId + " is already in terminal state " + currentState
+                            + " and cannot transition to " + newState,
+                    new InternalError("Task " + taskId + " is already in terminal state " + currentState));
+        }
+    }
+
     private Task createTask(String taskId, String contextId) {
         List<Message> history = initialMessage != null ? List.of(initialMessage) : Collections.emptyList();
         return Task.builder()

server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java

@@ -13,10 +13,14 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Flow;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.a2aproject.sdk.server.ServerCallContext;
@@ -36,6 +40,7 @@
 import org.a2aproject.sdk.server.tasks.PushNotificationSender;
 import org.a2aproject.sdk.server.tasks.TaskStore;
 import org.a2aproject.sdk.spec.A2AError;
+import org.a2aproject.sdk.spec.CancelTaskParams;
 import org.a2aproject.sdk.spec.Event;
 import org.a2aproject.sdk.spec.EventKind;
 import org.a2aproject.sdk.spec.InvalidParamsError;
@@ -46,6 +51,7 @@
 import org.a2aproject.sdk.spec.Task;
 import org.a2aproject.sdk.spec.TaskArtifactUpdateEvent;
 import org.a2aproject.sdk.spec.TaskNotFoundError;
+import org.a2aproject.sdk.spec.TaskNotCancelableError;
 import org.a2aproject.sdk.spec.TaskPushNotificationConfig;
 import org.a2aproject.sdk.spec.TaskState;
 import org.a2aproject.sdk.spec.TaskStatus;
@@ -1186,4 +1192,71 @@ private Task taskWithHistory(String id) {
                 .artifacts(List.of())
                 .build();
     }
+
+    @Test
+    void testConcurrentCancelsAreSerialized() throws Exception {
+        // Regression: two concurrent cancels of the same task must serialize on a
+        // per-task lock so the second one observes the CANCELED terminal state and fails
+        // with TaskNotCancelableError instead of both acting on the pre-transition state.
+        Task workingTask = Task.builder()
+                .id("task-cancel-lock")
+                .contextId("ctx-cancel")
+                .status(new TaskStatus(TaskState.TASK_STATE_WORKING))
+                .history(List.of())
+                .artifacts(List.of())
+                .build();
+        taskStore.save(workingTask, false);
+
+        CountDownLatch cancelEntered = new CountDownLatch(1);
+        CountDownLatch releaseCancel = new CountDownLatch(1);
+
+        agentExecutorCancel = (context, emitter) -> {
+            cancelEntered.countDown();
+            try {
+                releaseCancel.await(10, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            }
+            emitter.cancel();
+        };
+
+        ExecutorService cancelExec = Executors.newFixedThreadPool(2);
+        try {
+            CountDownLatch firstDone = new CountDownLatch(1);
+            Future<Task> first = cancelExec.submit(() -> {
+                try {
+                    Task result = requestHandler.onCancelTask(
+                            new CancelTaskParams("task-cancel-lock"), NULL_CONTEXT);
+                    firstDone.countDown();
+                    return result;
+                } catch (A2AError e) {
+                    firstDone.countDown();
+                    throw e;
+                }
+            });
+
+            // Wait until the first cancel is inside agentExecutor.cancel() (holding the per-task lock)
+            assertTrue(cancelEntered.await(5, TimeUnit.SECONDS),
+                    "First cancel should enter agentExecutor.cancel()");
+
+            // The second cancel must block on the per-task lock while the first is in progress
+            Future<Task> second = cancelExec.submit(() -> requestHandler.onCancelTask(
+                    new CancelTaskParams("task-cancel-lock"), NULL_CONTEXT));
+            assertThrows(TimeoutException.class, () -> second.get(300, TimeUnit.MILLISECONDS),
+                    "Second cancel should not complete while the first holds the per-task lock");
+
+            // Release the first cancel so it can enqueue CANCELED and finish
+            releaseCancel.countDown();
+            assertTrue(firstDone.await(10, TimeUnit.SECONDS), "First cancel should complete");
+            assertEquals(TaskState.TASK_STATE_CANCELED, first.get().status().state());
+
+            // The second cancel now observes the terminal state and is rejected
+            ExecutionException ex = assertThrows(ExecutionException.class, second::get);
+            assertInstanceOf(TaskNotCancelableError.class, ex.getCause(),
+                    "Second cancel should fail with TaskNotCancelableError");
+        } finally {
+            releaseCancel.countDown();
+            cancelExec.shutdownNow();
+        }
+    }
 }