diff --git a/extras/queue-manager-replicated/core/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java b/extras/queue-manager-replicated/core/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java index e524a717d..daced98d2 100644 --- a/extras/queue-manager-replicated/core/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java +++ b/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); diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java index d41b1e9c0..4b63c6716 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java @@ -257,6 +257,18 @@ public class DefaultRequestHandler implements RequestHandler { private final ConcurrentMap> runningAgents = new ConcurrentHashMap<>(); + /** + * Per-task lock registry serializing {@link #onCancelTask} check-then-act sequences. + *

+ * 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 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(); diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskManager.java b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskManager.java index e7a8b1568..775975727 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/TaskManager.java +++ b/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 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. + *

+ * A terminal (final) state must not be overwritten by a different state: + * once a task is {@code COMPLETED}/{@code FAILED}/{@code CANCELED}/{@code REJECTED} + * it stays in that state. Events re-arriving with the same final state are + * allowed, so replicated replays and idempotent retries keep working. + *

+ * 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 history = initialMessage != null ? List.of(initialMessage) : Collections.emptyList(); return Task.builder() diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java index cf768231e..48205fe7b 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/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 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 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(); + } + } } diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java index 75e04f2cd..e36ad6cc9 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/tasks/TaskManagerTest.java @@ -15,6 +15,7 @@ import java.util.Map; import org.a2aproject.sdk.spec.A2AServerException; +import org.a2aproject.sdk.spec.A2AError; import org.a2aproject.sdk.spec.Artifact; import org.a2aproject.sdk.spec.Message; import org.a2aproject.sdk.spec.Task; @@ -739,4 +740,159 @@ 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 + 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 + 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()); + } + + @Test + public void testA2AErrorOnTerminalTaskSkipsStateUpdate() throws A2AServerException { + // Seed a COMPLETED task + Task completedTask = Task.builder() + .id("task-a2aerror-terminal") + .contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + taskStore.save(completedTask, false); + TaskManager tm = new TaskManager("task-a2aerror-terminal", "ctx-1", taskStore, null); + + // An A2AError for an already-terminal task must not attempt the FAILED + // transition (which the state machine rejects); process() returns true + // (final) without throwing and the terminal state is preserved. + A2AError error = new A2AError(-32603, "agent failed", null); + assertTrue(tm.process(error, false)); + + assertEquals(TaskState.TASK_STATE_COMPLETED, taskStore.get("task-a2aerror-terminal").status().state()); + } + + @Test + public void testA2AErrorOnNonTerminalTaskTransitionsToFailed() throws A2AServerException { + Task workingTask = Task.builder() + .id("task-a2aerror-working") + .contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .build(); + taskStore.save(workingTask, false); + TaskManager tm = new TaskManager("task-a2aerror-working", "ctx-1", taskStore, null); + + A2AError error = new A2AError(-32603, "agent failed", null); + assertTrue(tm.process(error, false)); + + assertEquals(TaskState.TASK_STATE_FAILED, taskStore.get("task-a2aerror-working").status().state()); + } }