fix: enforce task state machine on task status writes

docs/architecture.md

@@ -393,13 +393,40 @@ heartbeat-sweep tick that handles offline-worker cleanup, controlled by
                └──────────┘ └──────────┘  └──────────────┘
 ```
 
-Transitions are validated — only the arrows above are permitted. Any other
-transition returns an error from `store.TransitionTask`. The auto-retry
-re-queue described below (`running` → `ready` on a transient failure) is a
-**separate, policy-driven store call** (`RequeueTaskForRetry`), not a
-`TransitionTask` arrow — the diagram above still reflects every state a task
-can be validated *into* via `TransitionTask`; auto-retry sends a task back to
-`ready` directly instead of landing it on `failed` at all.
+Transitions are validated by `store.ValidateTaskTransition`
+(`internal/store/statemachine.go`) and enforced by `UpdateTaskStatus` in both
+store implementations: the SQLite store reads the current status and writes the
+new one inside a single transaction, so the check cannot race a concurrent
+writer, and the in-memory fake does the same under its mutex. A transition
+outside the permitted set returns `store.ErrInvalidTransition` and leaves the
+row unchanged.
+
+Two rules keep enforcement safe given that task status arrives over JetStream
+(at-least-once delivery):
+
+- **Writing a task's current status is a no-op, not an error** — a redelivered
+  message must not fail.
+- **The consumer acks an invalid transition instead of Nak'ing it.** A message
+  describing a state the task has already left cannot become legal on
+  redelivery, so Nak'ing would loop forever. It is discarded with a warning,
+  the same treatment a malformed payload gets.
+
+Cancellation follows the same principle. `CancelTask` checks for a terminal
+status before writing, but the check and the write are separate operations, so
+a task can finish in between and the state machine then rejects the cancel.
+That is treated as the no-op it would have been had the check seen the newer
+value — canceling a completed task is not an error, regardless of which side
+of the race the caller landed on. Other store failures still propagate.
+
+Two arrows deserve note. `assigned` → `succeeded`/`failed` is permitted even
+though it appears to skip `running`: the worker publishes `running` first, but
+that publish is best-effort and gives up after `MaxRetries`, so it can be lost
+while the task still runs to completion. Rejecting the terminal message would
+strand finished work. Separately, the auto-retry re-queue described below
+(`running` → `ready` on a transient failure) is a **policy-driven store call**
+(`RequeueTaskForRetry`) with its own guarded SQL, as are the other bulk paths
+(`RetryTasks`, `TransitionStepPendingTasks`, `CancelJobTasks`, and the reclaim
+sweeps); none of them route through `UpdateTaskStatus`.
 
 ### Auto-retry on worker-reported failure
 

internal/api/tasks_test.go

@@ -53,7 +53,25 @@ func (f *fakeTaskCanceler) RetryTask(ctx context.Context, id string) error {
 		if status == "" {
 			status = store.TaskStatusReady
 		}
-		return f.retryStore.UpdateTaskStatus(ctx, id, status)
+		// Revive through RetryTasks, the same store call the real scheduler
+		// makes. UpdateTaskStatus would be wrong here: it enforces the task
+		// state machine, and failed → ready is not an arrow — production
+		// revives failed → pending inside RetryTasks and then promotes to ready
+		// via dependency resolution. Using UpdateTaskStatus made this double
+		// exercise a transition the real store rejects.
+		task, err := f.retryStore.GetTask(ctx, id)
+		if err != nil {
+			return err
+		}
+		if _, err := f.retryStore.RetryTasks(ctx, task.JobID, []string{id}, time.Now()); err != nil {
+			return err
+		}
+		if status != store.TaskStatusPending {
+			// RetryTasks lands on pending; walk the legal pending → ready arrow
+			// when the test wants the post-resolution status.
+			return f.retryStore.UpdateTaskStatus(ctx, id, status)
+		}
+		return nil
 	}
 	return nil
 }

internal/openjd/statemachine.go

@@ -9,72 +9,17 @@ import (
 	"github.com/uberware/sqi/internal/store"
 )
 
-// ErrInvalidTransition is returned when a requested status transition is not
-// permitted by the task or step state machine.
+// ErrInvalidTransition is returned when a requested step status transition is
+// not permitted by the step state machine.
 //
 // Use errors.Is to test:
 //
-//	err := ValidateTaskTransition(from, to)
+//	err := ValidateStepTransition(from, to)
 //	if errors.Is(err, ErrInvalidTransition) { ... }
-var ErrInvalidTransition = errors.New("openjd: invalid state transition")
-
-// ── Task state machine ────────────────────────────────────────────────────────
-//
-// Permitted task status transitions:
 //
-//	pending  → ready      dependency resolution: all dependency steps completed
-//	pending  → canceled   job canceled before step dependencies were satisfied
-//	ready    → assigned   scheduler assigns task to a worker
-//	ready    → canceled   task canceled while waiting for a worker
-//	assigned → running    worker confirms execution has started
-//	assigned → ready      reassignment: assigned worker disconnected or timed out
-//	assigned → canceled   task canceled after assignment but before confirmation
-//	running  → succeeded  worker reports clean exit (exit code 0)
-//	running  → failed     worker reports non-zero exit or fatal error
-//	running  → ready      reassignment: running worker became unreachable
-//	running  → canceled   task canceled while executing
-//
-// Terminal states (succeeded, failed, canceled) have no outgoing transitions.
-
-var validTaskTransitions = map[store.TaskStatus]map[store.TaskStatus]struct{}{
-	store.TaskStatusPending: {
-		store.TaskStatusReady:    {},
-		store.TaskStatusCanceled: {},
-	},
-	store.TaskStatusReady: {
-		store.TaskStatusAssigned: {},
-		store.TaskStatusCanceled: {},
-	},
-	store.TaskStatusAssigned: {
-		store.TaskStatusRunning:  {},
-		store.TaskStatusReady:    {},
-		store.TaskStatusCanceled: {},
-	},
-	store.TaskStatusRunning: {
-		store.TaskStatusSucceeded: {},
-		store.TaskStatusFailed:    {},
-		store.TaskStatusReady:     {},
-		store.TaskStatusCanceled:  {},
-	},
-	// Terminal states — no outgoing transitions.
-	store.TaskStatusSucceeded: {},
-	store.TaskStatusFailed:    {},
-	store.TaskStatusCanceled:  {},
-}
-
-// ValidateTaskTransition returns nil if transitioning a task from old to new
-// status is permitted by the state machine, or a descriptive error wrapping
-// [ErrInvalidTransition] otherwise.
-func ValidateTaskTransition(from, to store.TaskStatus) error {
-	targets, known := validTaskTransitions[from]
-	if !known {
-		return fmt.Errorf("%w: unknown task status %q", ErrInvalidTransition, from)
-	}
-	if _, ok := targets[to]; ok {
-		return nil
-	}
-	return fmt.Errorf("%w: task %q → %q not permitted", ErrInvalidTransition, from, to)
-}
+// The task state machine lives in package store, which enforces it on every
+// status write, and carries its own [store.ErrInvalidTransition].
+var ErrInvalidTransition = errors.New("openjd: invalid state transition")
 
 // ── Step state machine ────────────────────────────────────────────────────────
 //

internal/openjd/statemachine_test.go

@@ -14,68 +14,6 @@ import (
 	"github.com/uberware/sqi/internal/store"
 )
 
-// ── Task transitions ──────────────────────────────────────────────────────────
-
-func TestValidateTaskTransition(t *testing.T) {
-	legal := []struct {
-		from store.TaskStatus
-		to   store.TaskStatus
-	}{
-		{store.TaskStatusPending, store.TaskStatusReady},
-		{store.TaskStatusPending, store.TaskStatusCanceled},
-		{store.TaskStatusReady, store.TaskStatusAssigned},
-		{store.TaskStatusReady, store.TaskStatusCanceled},
-		{store.TaskStatusAssigned, store.TaskStatusRunning},
-		{store.TaskStatusAssigned, store.TaskStatusReady},
-		{store.TaskStatusAssigned, store.TaskStatusCanceled},
-		{store.TaskStatusRunning, store.TaskStatusSucceeded},
-		{store.TaskStatusRunning, store.TaskStatusFailed},
-		{store.TaskStatusRunning, store.TaskStatusReady},
-		{store.TaskStatusRunning, store.TaskStatusCanceled},
-	}
-	for _, tc := range legal {
-		if err := openjd.ValidateTaskTransition(tc.from, tc.to); err != nil {
-			t.Errorf("expected legal transition %q→%q, got error: %v", tc.from, tc.to, err)
-		}
-	}
-
-	illegal := []struct {
-		from store.TaskStatus
-		to   store.TaskStatus
-	}{
-		{store.TaskStatusPending, store.TaskStatusRunning},
-		{store.TaskStatusPending, store.TaskStatusSucceeded},
-		{store.TaskStatusReady, store.TaskStatusRunning},
-		{store.TaskStatusReady, store.TaskStatusSucceeded},
-		{store.TaskStatusSucceeded, store.TaskStatusRunning},
-		{store.TaskStatusSucceeded, store.TaskStatusFailed},
-		{store.TaskStatusFailed, store.TaskStatusRunning},
-		{store.TaskStatusFailed, store.TaskStatusSucceeded},
-		{store.TaskStatusCanceled, store.TaskStatusRunning},
-		{store.TaskStatusCanceled, store.TaskStatusSucceeded},
-	}
-	for _, tc := range illegal {
-		err := openjd.ValidateTaskTransition(tc.from, tc.to)
-		if err == nil {
-			t.Errorf("expected error for illegal transition %q→%q, got nil", tc.from, tc.to)
-			continue
-		}
-		if !errors.Is(err, openjd.ErrInvalidTransition) {
-			t.Errorf("transition %q→%q: error %v should wrap ErrInvalidTransition", tc.from, tc.to, err)
-		}
-	}
-}
-
-func TestValidateTaskTransition_UnknownStatus(t *testing.T) {
-	err := openjd.ValidateTaskTransition("bogus", store.TaskStatusReady)
-	if err == nil {
-		t.Fatal("expected error for unknown status, got nil")
-	}
-	if !errors.Is(err, openjd.ErrInvalidTransition) {
-		t.Errorf("expected ErrInvalidTransition, got %v", err)
-	}
-}
-
 // ── Step transitions ──────────────────────────────────────────────────────────
 
 func TestValidateStepTransition(t *testing.T) {

internal/scheduler/cancellation.go

@@ -114,7 +114,10 @@ func (s *Scheduler) CancelJob(ctx context.Context, jobID string) error {
 // task.cancel.<taskID> signal to the assigned worker.
 //
 // If the task is already in a terminal state (succeeded, failed, canceled),
-// CancelTask returns nil without modifying any state.
+// CancelTask returns nil without modifying any state. That holds whether the
+// terminal state was visible up front or the task reached it mid-cancel: the
+// guard below and the status write are separate operations, and losing that
+// race is reported the same way as never having had it.
 func (s *Scheduler) CancelTask(ctx context.Context, taskID string) error {
 	now := time.Now().UTC()
 
@@ -148,6 +151,24 @@ func (s *Scheduler) CancelTask(ctx context.Context, taskID string) error {
 	}
 
 	if err = s.store.UpdateTaskStatus(ctx, taskID, store.TaskStatusCanceled); err != nil {
+		// Losing a race to completion is a no-op, not a failure. The terminal
+		// guard above is a separate read, so the task can finish between that
+		// read and this write; the state machine then rejects it. Canceling an
+		// already-terminal task is documented to return nil, and which side of
+		// the race the caller landed on must not change that.
+		//
+		// Narrow by construction: every non-terminal status has a legal arrow
+		// to canceled, so ErrInvalidTransition on *this* write can only mean
+		// the task is already terminal. Any other store failure still
+		// propagates.
+		if errors.Is(err, store.ErrInvalidTransition) {
+			s.logger.DebugContext(
+				ctx, "scheduler: cancel task — reached terminal state first",
+				slog.String("task_id", taskID),
+				slog.Any("error", err),
+			)
+			return nil
+		}
 		return fmt.Errorf("scheduler: transition task %s to canceled: %w", taskID, err)
 	}
 

internal/scheduler/cancellation_race_test.go

@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+package scheduler
+
+// CancelTask's "already terminal" guard is a read followed by a separate write,
+// so a task can reach a terminal state in between: the guard sees `running` and
+// lets the cancel through, but by the time UpdateTaskStatus runs the row is
+// `succeeded` and the state machine rejects the write.
+//
+// The documented contract is that canceling an already-terminal task is a
+// silent no-op. Losing that race must therefore behave the same way it would
+// have if the guard had seen the newer value — return nil — not surface a 500
+// to the caller.
+//
+// staleReadStore reproduces the race deterministically. The underlying task is
+// already terminal; the wrapper's GetTask hands back the pre-completion status,
+// standing in for a guard that read a moment too early. No sleeps, no
+// goroutines, no flakiness.
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/uberware/sqi/internal/store"
+	"github.com/uberware/sqi/internal/store/fake"
+)
+
+type staleReadStore struct {
+	store.Store
+
+	taskID      string
+	staleStatus store.TaskStatus
+}
+
+func (s *staleReadStore) GetTask(ctx context.Context, id string) (store.Task, error) {
+	task, err := s.Store.GetTask(ctx, id)
+	if err != nil || id != s.taskID {
+		return task, err
+	}
+	task.Status = s.staleStatus // what the guard would have read pre-completion
+	return task, nil
+}
+
+func TestCancelTask_LosesRaceToCompletion_IsNoOp(t *testing.T) {
+	for _, terminal := range []store.TaskStatus{
+		store.TaskStatusSucceeded,
+		store.TaskStatusFailed,
+	} {
+		t.Run(string(terminal), func(t *testing.T) {
+			st := fake.New()
+			bus := &stubBus{}
+			job := seedCancelJob(t, st)
+			tk := seedTaskForJob(t, st, job, "w1", terminal)
+
+			// The guard reads "running"; the row is already terminal.
+			s := newTestScheduler(&staleReadStore{
+				Store:       st,
+				taskID:      tk.ID,
+				staleStatus: store.TaskStatusRunning,
+			}, bus)
+
+			if err := s.CancelTask(t.Context(), tk.ID); err != nil {
+				t.Fatalf("CancelTask losing the race to completion = %v, want nil (no-op)", err)
+			}
+
+			stored, err := st.GetTask(t.Context(), tk.ID)
+			if err != nil {
+				t.Fatalf("GetTask: %v", err)
+			}
+			if stored.Status != terminal {
+				t.Errorf("status = %q, want %q — a completed task must not be overwritten by a losing cancel",
+					stored.Status, terminal)
+			}
+		})
+	}
+}
+
+// TestCancelTask_RealErrorStillPropagates guards against the fix being written
+// as a blanket "swallow every error from UpdateTaskStatus".
+func TestCancelTask_RealErrorStillPropagates(t *testing.T) {
+	st := fake.New()
+	bus := &stubBus{}
+	job := seedCancelJob(t, st)
+	tk := seedTaskForJob(t, st, job, "w1", store.TaskStatusRunning)
+
+	s := newTestScheduler(&failingUpdateStore{Store: st, taskID: tk.ID}, bus)
+
+	err := s.CancelTask(t.Context(), tk.ID)
+	if err == nil {
+		t.Fatal("CancelTask = nil, want the underlying store error to propagate")
+	}
+	if errors.Is(err, store.ErrInvalidTransition) {
+		t.Errorf("error = %v, want the store failure, not ErrInvalidTransition", err)
+	}
+}
+
+var errStoreUnavailable = errors.New("store unavailable")
+
+type failingUpdateStore struct {
+	store.Store
+
+	taskID string
+}
+
+func (s *failingUpdateStore) UpdateTaskStatus(ctx context.Context, id string, status store.TaskStatus) error {
+	if id == s.taskID {
+		return errStoreUnavailable
+	}
+	return s.Store.UpdateTaskStatus(ctx, id, status)
+}