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 (BUG-44). 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;
@@ -464,6 +476,15 @@ 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 (BUG-44).
+        Object cancelLock = cancelLocks.computeIfAbsent(params.id(), k -> new Object());
+        synchronized (cancelLock) {
+            return doCancelTask(params, context);
+        }
+    }
+
+    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 (BUG-43).
+        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 (BUG-43). 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());
@@ -232,6 +246,36 @@ private Task ensureTask(String eventTaskId, String eventContextId) {
         return task;
     }
 
+    /**
+     * Validates a task state transition before it is persisted (BUG-43).
+     * <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;
@@ -1146,4 +1152,71 @@ public void onComplete() {
         assertEquals("1.0", pushConfigStore.getProtocolVersion(taskId, taskId),
             "Protocol version should be stored when push config is provided via onMessageSendStream");
     }
+
+    @Test
+    void testConcurrentCancelsAreSerialized() throws Exception {
+        // BUG-44 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();
+        }
+    }
 }

server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java

@@ -739,4 +739,123 @@ public void testUpdateWithMessage() throws A2AServerException {
         assertEquals("task message", ((TextPart) updated.history().get(1).parts().get(0)).text());
         assertEquals("update message", ((TextPart) updated.history().get(2).parts().get(0)).text());
     }
+
+    @Test
+    public void testRejectStatusUpdateOverwritingTerminalState() throws A2AServerException {
+        // Seed a COMPLETED task
+        Task completedTask = Task.builder()
+                .id("task-terminal")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                .build();
+        taskStore.save(completedTask, false);
+        TaskManager tm = new TaskManager("task-terminal", "ctx-1", taskStore, null);
+
+        // A status update to a different state after the terminal state must be rejected (BUG-43)
+        TaskStatusUpdateEvent workingEvent = TaskStatusUpdateEvent.builder()
+                .taskId("task-terminal")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_WORKING))
+                .build();
+        assertThrows(A2AServerException.class, () -> tm.saveTaskEvent(workingEvent, false));
+
+        // The persisted task must remain in its terminal state
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-terminal").status().state());
+    }
+
+    @Test
+    public void testRejectStatusUpdateToDifferentTerminalState() throws A2AServerException {
+        Task completedTask = Task.builder()
+                .id("task-terminal-2")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                .build();
+        taskStore.save(completedTask, false);
+        TaskManager tm = new TaskManager("task-terminal-2", "ctx-1", taskStore, null);
+
+        // COMPLETED must not be overwritten by FAILED either
+        TaskStatusUpdateEvent failedEvent = TaskStatusUpdateEvent.builder()
+                .taskId("task-terminal-2")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_FAILED))
+                .build();
+        assertThrows(A2AServerException.class, () -> tm.saveTaskEvent(failedEvent, false));
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-terminal-2").status().state());
+    }
+
+    @Test
+    public void testSameTerminalStateReplayAllowed() throws A2AServerException {
+        Task completedTask = Task.builder()
+                .id("task-replay")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                .build();
+        taskStore.save(completedTask, false);
+        TaskManager tm = new TaskManager("task-replay", "ctx-1", taskStore, null);
+
+        // Idempotent replay of the same final state must remain allowed (replication/replay)
+        TaskStatusUpdateEvent completedAgain = TaskStatusUpdateEvent.builder()
+                .taskId("task-replay")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                .build();
+        tm.saveTaskEvent(completedAgain, false);
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-replay").status().state());
+    }
+
+    @Test
+    public void testRejectTaskEventOverwritingTerminalState() throws A2AServerException {
+        Task completedTask = Task.builder()
+                .id("task-terminal-3")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED))
+                .build();
+        taskStore.save(completedTask, false);
+        TaskManager tm = new TaskManager("task-terminal-3", "ctx-1", taskStore, null);
+
+        // A full Task snapshot carrying a different (non-terminal) state must be rejected
+        Task submittedSnapshot = Task.builder()
+                .id("task-terminal-3")
+                .contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED))
+                .build();
+        assertThrows(A2AServerException.class, () -> tm.saveTaskEvent(submittedSnapshot, false));
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-terminal-3").status().state());
+    }
+
+    @Test
+    public void testNormalStateFlowAllowed() throws A2AServerException {
+        // SUBMITTED -> WORKING -> COMPLETED must keep working (BUG-43 must not break normal flows)
+        TaskManager tm = new TaskManager("task-flow", "ctx-1", taskStore, null);
+
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-flow").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED)).build(), false);
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-flow").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_WORKING)).build(), false);
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-flow").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)).build(), false);
+
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-flow").status().state());
+    }
+
+    @Test
+    public void testInterruptedStateResumeFlowAllowed() throws A2AServerException {
+        // INPUT_REQUIRED -> WORKING -> COMPLETED (resume flow) must keep working
+        TaskManager tm = new TaskManager("task-interrupted", "ctx-1", taskStore, null);
+
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-interrupted").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_INPUT_REQUIRED)).build(), false);
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-interrupted").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_WORKING)).build(), false);
+        tm.saveTaskEvent(TaskStatusUpdateEvent.builder()
+                .taskId("task-interrupted").contextId("ctx-1")
+                .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)).build(), false);
+
+        assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-interrupted").status().state());
+    }
 }