fix: enforce task state-machine transitions and serialize cancels - #1045
Open
ez-lbz wants to merge 2 commits into
Open
fix: enforce task state-machine transitions and serialize cancels#1045ez-lbz wants to merge 2 commits into
ez-lbz wants to merge 2 commits into
Conversation
testParallelReplicationBehavior sent TASK_STATE_COMPLETED events from the replicated threads. 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 became timing- dependent (observed 0/1/2/3/21 instead of 25, locally and in CI). Use a non-terminal state for the replicated events; the replication hook skips them via isReplicated() regardless of state, so the test's intent (normal enqueues replicate, replicated events do not) is unchanged while the outcome is now deterministic.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
1. Enforce state-machine transitions in
TaskManagerProblem:
TaskManager.saveTaskEvent/processoverwrote 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):
validateStateTransition(currentState, newState, taskId)and call it in the status-update path (saveTaskEvent(TaskStatusUpdateEvent)) and the full-task path (saveTaskEvent(Task)).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.AgentExecutorflows 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.onCancelTaskperformed a check-then-act sequence — read task → checkisFinal()→ invokeagentExecutor.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):
cancelLocks, aConcurrentHashMap<String, Object>keyed by task ID) and moved the entire cancel body intosynchronized (lock)via adoCancelTaskhelper. The second concurrent cancel now waits for the first to finish, observes theCANCELEDterminal state, and fails withTaskNotCancelableError.message/sendcompleted first.Fix (server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java):
testConcurrentCancelsAreSerialized— holds the first cancel insideagentExecutor.cancel(), asserts the second cancel blocks, then verifies the first succeeds withCANCELEDand the second fails withTaskNotCancelableError.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
TaskNotCancelableErrorinstead of both succeeding.2. Make the replicated queue manager parallel test deterministic
Problem:
ReplicatedQueueManagerTest.testParallelReplicationBehaviorwas 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 sentTASK_STATE_COMPLETEDevents; 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):
TASK_STATE_WORKING). The replication hook skips replicated events viaisReplicated()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).