fix(workflow): journal terminal transitions atomically

Author: Waishnav

src/workflow-cli.ts

@@ -457,12 +457,7 @@ export async function runWorkflowWorker(
       }
     }
 
-    store.completeRun(runId, { resultJson });
-    store.appendEvent({
-      runId,
-      type: "run_completed",
-      data: { callCount },
-    });
+    store.completeRun(runId, { resultJson, callCount });
   } catch (error) {
     if (store.isCancelRequested(runId) || abort.signal.aborted) {
       try {
@@ -476,11 +471,6 @@ export async function runWorkflowWorker(
     const errorKind = mapEngineErrorKind(error);
     try {
       store.failRun(runId, { error: message, errorKind });
-      store.appendEvent({
-        runId,
-        type: "run_failed",
-        data: { error: message, errorKind },
-      });
     } catch {
       // terminal race
     }

src/workflow-store.test.ts

@@ -1,4 +1,5 @@
 import assert from "node:assert/strict";
+import { spawnSync } from "node:child_process";
 import { mkdtempSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
@@ -56,12 +57,14 @@ try {
   const page1 = store.drainEvents(run.id, 0, 2);
   assert.equal(page1.events.length, 2);
   assert.equal(page1.nextSeq, 2);
+  assert.equal(page1.hasMore, true);
   assert.equal(page1.terminal, false);
 
   const page2 = store.drainEvents(run.id, 2, 10);
   assert.equal(page2.events.length, 1);
   assert.equal(page2.events[0]?.seq, 3);
   assert.equal(page2.nextSeq, 3);
+  assert.equal(page2.hasMore, false);
 
   store.beginAgentCall({
     runId: run.id,
@@ -120,7 +123,12 @@ try {
   assert.equal(terminal.errorKind, "cancelled");
   assert.equal(store.cancelRun(run.id).status, "cancelled");
 
+  const terminalPage1 = store.drainEvents(run.id, 0, 2);
+  assert.equal(terminalPage1.hasMore, true);
+  assert.equal(terminalPage1.terminal, false);
   const drainDone = store.drainEvents(run.id, 0, 100);
+  assert.equal(drainDone.events.at(-1)?.type, "run_cancelled");
+  assert.equal(drainDone.hasMore, false);
   assert.equal(drainDone.terminal, true);
 
   const run2 = store.createRun({
@@ -159,7 +167,7 @@ try {
     store.listRunsForWorkspace(join(root, "other-project"))[0]?.id,
     otherProjectRun.id,
   );
-  assert.deepEqual(store.listEvents(run.id, 2).map((event) => event.seq), [2, 3]);
+  assert.deepEqual(store.listEvents(run.id, 2).map((event) => event.seq), [3, 4]);
 
   // Reap: stale heartbeat + dead pid (force heartbeat via shared sqlite handle)
   const run3 = store.createRun({
@@ -169,7 +177,9 @@ try {
     scriptHash: "h3",
     workspaceRoot: join(root, "project"),
   });
-  store.claimRun(run3.id, 2_147_483_646);
+  const dead = spawnSync(process.execPath, ["-e", ""]);
+  assert.ok(dead.pid);
+  store.claimRun(run3.id, dead.pid);
   const db = openDatabase(root);
   try {
     db.sqlite
@@ -181,6 +191,7 @@ try {
   const reaped = store.reapStale(60_000);
   assert.ok(reaped.some((r) => r.id === run3.id && r.status === "failed"));
   assert.equal(store.getRun(run3.id)?.errorKind, "heartbeat");
+  assert.equal(store.listEvents(run3.id).at(-1)?.type, "run_failed");
 
   const run4 = store.createRun({
     name: "seq",

src/workflow-store.ts

@@ -87,6 +87,7 @@ export interface FailAgentCallInput {
 
 export interface CompleteRunInput {
   resultJson?: string;
+  callCount?: number;
 }
 
 export interface FailRunInput {
@@ -97,6 +98,7 @@ export interface FailRunInput {
 export interface DrainEventsResult {
   events: WorkflowEventRecord[];
   nextSeq: number;
+  hasMore: boolean;
   terminal: boolean;
   run: WorkflowRunRecord;
 }
@@ -409,11 +411,14 @@ export class WorkflowStore {
     const updated = Result.try({
       try: () => {
         const now = isoNow();
-        this.database.sqlite
+        const update = this.database.sqlite
           .prepare(
-            `update workflow_runs set cancel_requested = 'true', updated_at = ? where id = ?`,
+            `update workflow_runs
+             set cancel_requested = 'true', updated_at = ?
+             where id = ? and status in ('starting', 'running')`,
           )
           .run(now, id);
+        if (update.changes === 0) return this.getRun(id);
         return this.getRun(id);
       },
       catch: (cause) => new WorkflowStoreError("request_cancel", cause),
@@ -439,18 +444,31 @@ export class WorkflowStore {
     return this.transitionRunResult(id, "complete", () => {
       if (input.resultJson !== undefined) assertResultSize(input.resultJson);
       const now = isoNow();
-      return this.database.sqlite
-        .prepare(
-          `update workflow_runs set
-            status = 'completed',
-            result_json = ?,
-            completed_at = ?,
-            updated_at = ?,
-            error = null,
-            error_kind = null
-           where id = ? and status in ('starting', 'running')`,
-        )
-        .run(input.resultJson ?? null, now, now, id).changes;
+      const transaction = this.database.sqlite.transaction(() => {
+        const changes = this.database.sqlite
+          .prepare(
+            `update workflow_runs set
+              status = 'completed',
+              result_json = ?,
+              completed_at = ?,
+              updated_at = ?,
+              error = null,
+              error_kind = null
+             where id = ? and status in ('starting', 'running')`,
+          )
+          .run(input.resultJson ?? null, now, now, id).changes;
+        if (changes === 0) return 0;
+        this.insertEventRow(
+          {
+            runId: id,
+            type: "run_completed",
+            data: { callCount: input.callCount ?? 0 },
+          },
+          now,
+        );
+        return changes;
+      });
+      return transaction.immediate();
     });
   }
 
@@ -464,17 +482,31 @@ export class WorkflowStore {
   ): BetterResult<WorkflowRunRecord, WorkflowRunTransitionError> {
     return this.transitionRunResult(id, "fail", () => {
       const now = isoNow();
-      return this.database.sqlite
-        .prepare(
-          `update workflow_runs set
-            status = 'failed',
-            error = ?,
-            error_kind = ?,
-            completed_at = ?,
-            updated_at = ?
-           where id = ? and status in ('starting', 'running')`,
-        )
-        .run(input.error, input.errorKind ?? "internal", now, now, id).changes;
+      const errorKind = input.errorKind ?? "internal";
+      const transaction = this.database.sqlite.transaction(() => {
+        const changes = this.database.sqlite
+          .prepare(
+            `update workflow_runs set
+              status = 'failed',
+              error = ?,
+              error_kind = ?,
+              completed_at = ?,
+              updated_at = ?
+             where id = ? and status in ('starting', 'running')`,
+          )
+          .run(input.error, errorKind, now, now, id).changes;
+        if (changes === 0) return 0;
+        this.insertEventRow(
+          {
+            runId: id,
+            type: "run_failed",
+            data: { error: input.error, errorKind },
+          },
+          now,
+        );
+        return changes;
+      });
+      return transaction.immediate();
     });
   }
 
@@ -488,18 +520,31 @@ export class WorkflowStore {
   ): BetterResult<WorkflowRunRecord, WorkflowRunTransitionError> {
     return this.transitionRunResult(id, "cancel", () => {
       const now = isoNow();
-      return this.database.sqlite
-        .prepare(
-          `update workflow_runs set
-            status = 'cancelled',
-            error = ?,
-            error_kind = 'cancelled',
-            cancel_requested = 'true',
-            completed_at = ?,
-            updated_at = ?
-           where id = ? and status in ('starting', 'running')`,
-        )
-        .run(error, now, now, id).changes;
+      const transaction = this.database.sqlite.transaction(() => {
+        const changes = this.database.sqlite
+          .prepare(
+            `update workflow_runs set
+              status = 'cancelled',
+              error = ?,
+              error_kind = 'cancelled',
+              cancel_requested = 'true',
+              completed_at = ?,
+              updated_at = ?
+             where id = ? and status in ('starting', 'running')`,
+          )
+          .run(error, now, now, id).changes;
+        if (changes === 0) return 0;
+        this.insertEventRow(
+          {
+            runId: id,
+            type: "run_cancelled",
+            data: { reason: error },
+          },
+          now,
+        );
+        return changes;
+      });
+      return transaction.immediate();
     });
   }
 
@@ -540,47 +585,11 @@ export class WorkflowStore {
   }
 
   appendEvent(input: AppendWorkflowEventInput): WorkflowEventRecord {
-    const payload = parseWorkflowEventPayload(input.type, input.data);
-    const dataJson = truncateJson(payload, WORKFLOW_LIMITS.eventDataJsonBytes);
     const createdAt = isoNow();
-
-    const insert = this.database.sqlite.transaction(() => {
-      const next = this.database.sqlite
-        .prepare(
-          `select coalesce(max(seq), 0) + 1 as next_seq from workflow_events where run_id = ?`,
-        )
-        .get(input.runId) as { next_seq: number };
-      const seq = next.next_seq;
-      this.database.sqlite
-        .prepare(
-          `insert into workflow_events (run_id, seq, type, phase, label, data_json, created_at)
-           values (?, ?, ?, ?, ?, ?, ?)`,
-        )
-        .run(
-          input.runId,
-          seq,
-          input.type,
-          input.phase ?? null,
-          input.label ?? null,
-          dataJson,
-          createdAt,
-        );
-      this.database.sqlite
-        .prepare(`update workflow_runs set updated_at = ? where id = ?`)
-        .run(createdAt, input.runId);
-      return seq;
-    });
-
-    const seq = insert();
-    return {
-      runId: input.runId,
-      seq,
-      type: input.type,
-      phase: input.phase,
-      label: input.label,
-      dataJson,
-      createdAt,
-    };
+    const transaction = this.database.sqlite.transaction(() =>
+      this.insertEventRow(input, createdAt),
+    );
+    return transaction.immediate();
   }
 
   drainEvents(runId: string, sinceSeq = 0, limit: number = WORKFLOW_LIMITS.eventDrainDefault): DrainEventsResult {
@@ -593,13 +602,15 @@ export class WorkflowStore {
          order by seq asc
          limit ?`,
       )
-      .all(runId, sinceSeq, capped) as WorkflowEventRow[];
-    const events = rows.map(rowToEvent);
+      .all(runId, sinceSeq, capped + 1) as WorkflowEventRow[];
+    const hasMore = rows.length > capped;
+    const events = rows.slice(0, capped).map(rowToEvent);
     const nextSeq = events.length > 0 ? events[events.length - 1]!.seq : sinceSeq;
     return {
       events,
       nextSeq,
-      terminal: TERMINAL_STATUSES.has(run.status),
+      hasMore,
+      terminal: TERMINAL_STATUSES.has(run.status) && !hasMore,
       run,
     };
   }
@@ -755,16 +766,59 @@ export class WorkflowStore {
     const reaped: WorkflowRunRecord[] = [];
     for (const row of candidates) {
       if (row.pid !== null && isPidAlive(row.pid)) continue;
-      reaped.push(
-        this.failRun(row.id, {
-          error: "worker heartbeat lost",
-          errorKind: "heartbeat",
-        }),
-      );
+      const latest = this.getRun(row.id);
+      if (!latest || latest.status !== "running") continue;
+      const failed = this.failRun(row.id, {
+        error: "worker heartbeat lost",
+        errorKind: "heartbeat",
+      });
+      if (failed.status === "failed" && failed.errorKind === "heartbeat") {
+        reaped.push(failed);
+      }
     }
     return reaped;
   }
 
+  private insertEventRow(
+    input: AppendWorkflowEventInput,
+    createdAt: string,
+  ): WorkflowEventRecord {
+    const payload = parseWorkflowEventPayload(input.type, input.data);
+    const dataJson = truncateJson(payload, WORKFLOW_LIMITS.eventDataJsonBytes);
+    const next = this.database.sqlite
+      .prepare(
+        `select coalesce(max(seq), 0) + 1 as next_seq from workflow_events where run_id = ?`,
+      )
+      .get(input.runId) as { next_seq: number };
+    const seq = next.next_seq;
+    this.database.sqlite
+      .prepare(
+        `insert into workflow_events (run_id, seq, type, phase, label, data_json, created_at)
+         values (?, ?, ?, ?, ?, ?, ?)`,
+      )
+      .run(
+        input.runId,
+        seq,
+        input.type,
+        input.phase ?? null,
+        input.label ?? null,
+        dataJson,
+        createdAt,
+      );
+    this.database.sqlite
+      .prepare(`update workflow_runs set updated_at = ? where id = ?`)
+      .run(createdAt, input.runId);
+    return {
+      runId: input.runId,
+      seq,
+      type: input.type,
+      phase: input.phase,
+      label: input.label,
+      dataJson,
+      createdAt,
+    };
+  }
+
   close(): void {
     this.database.close();
   }
@@ -868,7 +922,8 @@ function isPidAlive(pid: number): boolean {
   try {
     process.kill(pid, 0);
     return true;
-  } catch {
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code === "EPERM") return true;
     return false;
   }
 }