refactor(control-plane): introduce revision-based CAS for operation state updates

Add revision tracking and compare-and-set (CAS) semantics to operation state
stores, enabling atomic state transitions and preventing lost updates. Extend
the OperationStateStore interface with getOpStateRecord and compareAndSetOpState
methods. Update orchestrator and worker runtime to utilize CAS for marking
operations as failed only if the state is unchanged. Enhance in-memory,
JetStream, and test control plane implementations to support revision logic.

This change improves concurrency safety and correctness of operation state
management across distributed components.

compute/core/src/control-plane/index.ts

@@ -1,5 +1,4 @@
 export * from './types';
 export * from './state-machine';
 export * from './orchestrator';
-export * from './in-memory';
 export * from './sse';

compute/core/src/control-plane/orchestrator.ts

@@ -263,6 +263,36 @@ export class OperationOrchestrator {
     return next;
   }
 
+  async markFailedIfUnchanged(input: {
+    current: WorkerOperationState;
+    expectedRevision: number;
+    error: WorkerJobErrorShape | string;
+    updatedAt?: number;
+    timing?: WorkerJobTiming;
+  }): Promise<WorkerOperationState | null> {
+    const now = input.updatedAt ?? this.clock.now();
+    const shape = typeof input.error === 'string' ? { message: input.error } : input.error;
+
+    const next: WorkerOperationState = {
+      ...input.current,
+      status: 'failed',
+      startedAt: input.current.startedAt ?? now,
+      updatedAt: now,
+      error: shape,
+      ...(input.timing ? { timing: input.timing } : {}),
+    };
+
+    const updated = await this.stateStore.compareAndSetOpState({
+      opId: input.current.opId,
+      expectedRevision: input.expectedRevision,
+      newState: next,
+    });
+    if (!updated) return null;
+
+    await this.eventStream.append(next.opId, next);
+    return next;
+  }
+
   async explainReuseDecision(input: {
     current: WorkerOperationState;
     requestKind: WorkerOperationKind;

compute/core/src/control-plane/types.ts

@@ -26,9 +26,20 @@ export interface OperationIndexEntry {
   opId: string;
 }
 
+export interface OperationStateRecord<Result = unknown> {
+  state: OperationState<Result>;
+  revision: number;
+}
+
 export interface OperationStateStore<Result = unknown> {
   getOpState(opId: string): Promise<OperationState<Result> | null>;
+  getOpStateRecord(opId: string): Promise<OperationStateRecord<Result> | null>;
   putOpState(state: OperationState<Result>): Promise<void>;
+  compareAndSetOpState(input: {
+    opId: string;
+    expectedRevision: number;
+    newState: OperationState<Result>;
+  }): Promise<boolean>;
   getOpIndex(opKey: string): Promise<OperationIndexEntry | null>;
   compareAndSetOpIndex(input: {
     opKey: string;

compute/core/tests/control-plane/orchestrator.test.ts

@@ -1,11 +1,11 @@
 import { describe, expect, test } from 'vitest';
 import type { WorkerOperationRequest } from '../../src/api-contracts';
+import { OperationOrchestrator } from '../../src/control-plane';
 import {
   InMemoryOperationEventStream,
   InMemoryOperationQueue,
   InMemoryOperationStateStore,
-  OperationOrchestrator,
-} from '../../src/control-plane';
+} from '../helpers/in-memory-control-plane';
 
 function buildRequest(opKey: string): WorkerOperationRequest {
   return {
@@ -65,7 +65,9 @@ describe('operation orchestrator', () => {
     let firstAttempt = true;
     const conflictStore = {
       getOpState: store.getOpState.bind(store),
+      getOpStateRecord: store.getOpStateRecord.bind(store),
       putOpState: store.putOpState.bind(store),
+      compareAndSetOpState: store.compareAndSetOpState.bind(store),
       getOpIndex: store.getOpIndex.bind(store),
       compareAndSetOpIndex: async (input: { opKey: string; newOpId: string; expectedOpId: string | null }) => {
         if (firstAttempt && input.expectedOpId === null) {
@@ -92,4 +94,45 @@ describe('operation orchestrator', () => {
     expect(created.opId).toMatch(/^op-/);
     expect(await store.getOpIndex('cas-key')).toEqual({ opId: created.opId });
   });
+
+  test('markFailedIfUnchanged only writes once for the expected revision', async () => {
+    const queue = new InMemoryOperationQueue();
+    const stateStore = new InMemoryOperationStateStore();
+    const eventStream = new InMemoryOperationEventStream();
+    const orchestrator = new OperationOrchestrator({
+      queue,
+      stateStore,
+      eventStream,
+      config: { opStaleMs: 2_000, maxCasRetries: 5 },
+    });
+
+    const created = await orchestrator.enqueueOrReuse(buildRequest('stale-op'));
+    await orchestrator.markRunning({ opId: created.opId, updatedAt: 2_000 });
+
+    const record = await stateStore.getOpStateRecord(created.opId);
+    expect(record).not.toBeNull();
+
+    const first = await orchestrator.markFailedIfUnchanged({
+      current: record!.state,
+      expectedRevision: record!.revision,
+      error: { code: 'WORKER_ORPHANED_OP', message: 'stale op' },
+      updatedAt: 3_000,
+    });
+    const second = await orchestrator.markFailedIfUnchanged({
+      current: record!.state,
+      expectedRevision: record!.revision,
+      error: { code: 'WORKER_ORPHANED_OP', message: 'stale op' },
+      updatedAt: 3_000,
+    });
+
+    expect(first).toMatchObject({
+      opId: created.opId,
+      status: 'failed',
+      error: { code: 'WORKER_ORPHANED_OP' },
+    });
+    expect(second).toBeNull();
+
+    const events = await eventStream.listSince(created.opId, 0);
+    expect(events.filter((event) => event.snapshot.status === 'failed')).toHaveLength(1);
+  });
 });

compute/core/tests/helpers/in-memory-control-plane.ts

@@ -1,4 +1,5 @@
 import { EventEmitter } from 'node:events';
+import type { WorkerOperationKind } from '../../src/api-contracts';
 import type {
   OperationEvent,
   OperationEventStream,
@@ -7,8 +8,7 @@ import type {
   OperationState,
   OperationStateStore,
   QueuedOperation,
-} from './types';
-import type { WorkerOperationKind } from '../api-contracts';
+} from '../../src/control-plane/types';
 
 function topicFor(opId: string): string {
   return `op.${opId}`;
@@ -49,14 +49,39 @@ export class InMemoryOperationQueue implements OperationQueue {
 
 export class InMemoryOperationStateStore implements OperationStateStore {
   private readonly stateByOpId = new Map<string, OperationState>();
+  private readonly revisionByOpId = new Map<string, number>();
   private readonly opIndexByKey = new Map<string, string>();
 
   async getOpState(opId: string): Promise<OperationState | null> {
     return this.stateByOpId.get(opId) ?? null;
   }
 
+  async getOpStateRecord(opId: string): Promise<{ state: OperationState; revision: number } | null> {
+    const state = this.stateByOpId.get(opId);
+    if (!state) return null;
+    return {
+      state,
+      revision: this.revisionByOpId.get(opId) ?? 0,
+    };
+  }
+
   async putOpState(state: OperationState): Promise<void> {
     this.stateByOpId.set(state.opId, state);
+    this.revisionByOpId.set(state.opId, (this.revisionByOpId.get(state.opId) ?? 0) + 1);
+  }
+
+  async compareAndSetOpState(input: {
+    opId: string;
+    expectedRevision: number;
+    newState: OperationState;
+  }): Promise<boolean> {
+    const currentState = this.stateByOpId.get(input.opId);
+    if (!currentState) return false;
+    const currentRevision = this.revisionByOpId.get(input.opId) ?? 0;
+    if (currentRevision !== input.expectedRevision) return false;
+    this.stateByOpId.set(input.opId, input.newState);
+    this.revisionByOpId.set(input.opId, currentRevision + 1);
+    return true;
   }
 
   async getOpIndex(opKey: string): Promise<OperationIndexEntry | null> {

compute/worker/src/control-plane/jetstream.ts

@@ -81,10 +81,18 @@ export class JetStreamOperationStateStore<Result = unknown> implements Operation
   }
 
   async getOpState(opId: string): Promise<OperationState<Result> | null> {
+    const record = await this.getOpStateRecord(opId);
+    return record?.state ?? null;
+  }
+
+  async getOpStateRecord(opId: string): Promise<{ state: OperationState<Result>; revision: number } | null> {
     const kv = await this.getKv();
     const entry = await kv.get(opStateKvKey(opId));
     if (!isPut(entry)) return null;
-    return this.opStateCodec.decode(entry.value);
+    return {
+      state: this.opStateCodec.decode(entry.value),
+      revision: entry.revision,
+    };
   }
 
   async putOpState(state: OperationState<Result>): Promise<void> {
@@ -92,6 +100,25 @@ export class JetStreamOperationStateStore<Result = unknown> implements Operation
     await kv.put(opStateKvKey(state.opId), this.opStateCodec.encode(state));
   }
 
+  async compareAndSetOpState(input: {
+    opId: string;
+    expectedRevision: number;
+    newState: OperationState<Result>;
+  }): Promise<boolean> {
+    const kv = await this.getKv();
+    try {
+      await kv.update(
+        opStateKvKey(input.opId),
+        this.opStateCodec.encode(input.newState),
+        input.expectedRevision,
+      );
+      return true;
+    } catch (error) {
+      if (isCasConflictError(error)) return false;
+      throw error;
+    }
+  }
+
   async listOpStates(): Promise<OperationState<Result>[]> {
     const kv = await this.getKv();
     const keys = await kv.keys('op_state.*');

compute/worker/src/runtime.ts

@@ -120,6 +120,7 @@ interface OperationEventStreamLike {
 
 interface OperationStateStoreLike {
   getOpState(opId: string): Promise<StreamedOperationState | null>;
+  getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>;
   listOpStates?(): Promise<StreamedOperationState[]>;
 }
 
@@ -149,6 +150,13 @@ interface OrchestratorLike {
     updatedAt?: number;
     timing?: WorkerJobTiming;
   }): Promise<StreamedOperationState>;
+  markFailedIfUnchanged?(input: {
+    current: StreamedOperationState;
+    expectedRevision: number;
+    error: { message: string; code?: string } | string;
+    updatedAt?: number;
+    timing?: WorkerJobTiming;
+  }): Promise<StreamedOperationState | null>;
 }
 
 export interface ComputeWorkerRouteDeps {
@@ -183,6 +191,20 @@ function isInflightStatus(status: WorkerJobState): boolean {
   return status === 'queued' || status === 'running';
 }
 
+function getOrphanRecoveryThresholdMs(input: {
+  state: StreamedOperationState;
+  whisperTimeoutMs: number;
+  pdfTimeoutMs: number;
+  opStaleMs: number;
+}): number | null {
+  if (!isInflightStatus(input.state.status)) return null;
+  if (input.state.status === 'running') {
+    return input.state.kind === 'whisper_align' ? input.whisperTimeoutMs : input.pdfTimeoutMs;
+  }
+  if (input.state.kind !== 'pdf_layout') return null;
+  return input.opStaleMs;
+}
+
 function readIntEnv(name: string, fallback: number): number {
   const raw = process.env[name]?.trim();
   if (!raw) return fallback;
@@ -744,6 +766,8 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti
 
   const ensureOrphanedOpRecovery = async (): Promise<void> => {
     if (typeof operationStateStore.listOpStates !== 'function') return;
+    if (typeof operationStateStore.getOpStateRecord !== 'function') return;
+    if (typeof orchestrator.markFailedIfUnchanged !== 'function') return;
     if (orphanRecoveryDoneForGeneration === sessionGeneration) return;
     if (orphanRecoveryPromise) {
       await orphanRecoveryPromise;
@@ -753,35 +777,51 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti
     orphanRecoveryPromise = (async () => {
       const now = Date.now();
       const states = await operationStateStore.listOpStates!();
-      const staleStates = states.filter((state) => {
-        if (!isInflightStatus(state.status)) return false;
-        const ageMs = now - state.updatedAt;
-        if (state.status === 'running') {
-          const runningTimeoutMs = state.kind === 'whisper_align' ? whisperTimeoutMs : pdfTimeoutMs;
-          return ageMs > runningTimeoutMs;
-        }
-        if (state.kind !== 'pdf_layout') return false;
-        return ageMs > opStaleMs;
-      });
+      const candidateStates = states.filter((state) => (
+        getOrphanRecoveryThresholdMs({
+          state,
+          whisperTimeoutMs,
+          pdfTimeoutMs,
+          opStaleMs,
+        }) !== null
+      ));
+      const recoveredStates: Array<Pick<StreamedOperationState, 'opId' | 'kind' | 'status'>> = [];
 
-      for (const state of staleStates) {
-        const staleAfterMs = state.status === 'running'
-          ? (state.kind === 'whisper_align' ? whisperTimeoutMs : pdfTimeoutMs)
-          : opStaleMs;
-        await orchestrator.markFailed({
-          opId: state.opId,
+      for (const candidate of candidateStates) {
+        const record = await operationStateStore.getOpStateRecord!(candidate.opId);
+        if (!record) continue;
+        const staleAfterMs = getOrphanRecoveryThresholdMs({
+          state: record.state,
+          whisperTimeoutMs,
+          pdfTimeoutMs,
+          opStaleMs,
+        });
+        if (staleAfterMs === null) continue;
+        const ageMs = now - record.state.updatedAt;
+        if (ageMs <= staleAfterMs) continue;
+
+        const recovered = await orchestrator.markFailedIfUnchanged!({
+          current: record.state,
+          expectedRevision: record.revision,
           error: {
             code: 'WORKER_ORPHANED_OP',
             message: `Worker stopped before completion; stale operation recovered on startup after ${staleAfterMs}ms`,
           },
           updatedAt: now,
         });
+        if (!recovered) continue;
+
+        recoveredStates.push({
+          opId: recovered.opId,
+          kind: recovered.kind,
+          status: record.state.status,
+        });
       }
 
-      if (staleStates.length > 0) {
+      if (recoveredStates.length > 0) {
         app.log.warn({
-          recoveredCount: staleStates.length,
-          ops: staleStates.map((state) => ({
+          recoveredCount: recoveredStates.length,
+          ops: recoveredStates.map((state) => ({
             opId: state.opId,
             kind: state.kind,
             status: state.status,

compute/worker/tests/fixtures/fake-control-plane.ts

@@ -13,6 +13,7 @@ type ComputeEvent = WorkerOperationEvent<ComputeResult>;
 
 export class FakeControlPlane {
   private readonly stateByOpId = new Map<string, ComputeState>();
+  private readonly revisionByOpId = new Map<string, number>();
   private readonly opIdByOpKey = new Map<string, string>();
   private readonly eventsByOpId = new Map<string, ComputeEvent[]>();
   private nextOpId = 1;
@@ -44,9 +45,18 @@ export class FakeControlPlane {
         updatedAt: input.updatedAt,
         timing: input.timing,
       }),
+      markFailedIfUnchanged: async (input) => this.compareAndSetFailed(input),
     },
     operationStateStore: {
       getOpState: async (opId) => this.stateByOpId.get(opId) ?? null,
+      getOpStateRecord: async (opId) => {
+        const state = this.stateByOpId.get(opId);
+        if (!state) return null;
+        return {
+          state,
+          revision: this.revisionByOpId.get(opId) ?? 0,
+        };
+      },
       listOpStates: async () => Array.from(this.stateByOpId.values()),
     },
     operationEventStream: {
@@ -63,6 +73,7 @@ export class FakeControlPlane {
 
   seedState(state: ComputeState): void {
     this.stateByOpId.set(state.opId, state);
+    this.revisionByOpId.set(state.opId, (this.revisionByOpId.get(state.opId) ?? 0) + 1);
     this.opIdByOpKey.set(state.opKey, state.opId);
   }
 
@@ -96,11 +107,30 @@ export class FakeControlPlane {
     };
 
     this.stateByOpId.set(opId, state);
+    this.revisionByOpId.set(opId, 1);
     this.opIdByOpKey.set(request.opKey, opId);
     this.seedEvent(opId, { eventId: 1, snapshot: state });
     return state;
   }
 
+  private async compareAndSetFailed(input: {
+    current: ComputeState;
+    expectedRevision: number;
+    error: { message: string; code?: string } | string;
+    updatedAt?: number;
+    timing?: ComputeState['timing'];
+  }): Promise<ComputeState | null> {
+    const currentRevision = this.revisionByOpId.get(input.current.opId) ?? 0;
+    if (currentRevision !== input.expectedRevision) return null;
+    return this.updateState(input.current.opId, {
+      ...input.current,
+      status: 'failed',
+      error: typeof input.error === 'string' ? { message: input.error } : input.error,
+      updatedAt: input.updatedAt,
+      timing: input.timing,
+    });
+  }
+
   private async updateState(
     opId: string,
     patch: Partial<ComputeState>,
@@ -115,6 +145,7 @@ export class FakeControlPlane {
       updatedAt: patch.updatedAt ?? Date.now(),
     };
     this.stateByOpId.set(opId, next);
+    this.revisionByOpId.set(opId, (this.revisionByOpId.get(opId) ?? 0) + 1);
     const currentEvents = this.eventsByOpId.get(opId) ?? [];
     const nextEventId = (currentEvents.at(-1)?.eventId ?? 0) + 1;
     currentEvents.push({ eventId: nextEventId, snapshot: next });