@osolmaz/pi-workflows 0.13.3 → 0.14.0

package/src/builtins/autoplan.workflow.ts

@@ -96,9 +96,11 @@ function originalUserInstructions(outputs: Record<string, unknown>): string {
 }
 
 function originalInstructionsPrompt(outputs: Record<string, unknown>): string {
-  return ["Original user instructions (authoritative):", originalUserInstructions(outputs)].join(
-    "\n",
-  );
+  return [
+    "Continue in this Pi session. Do not delegate this workflow step or start another session.",
+    "Original user instructions (authoritative):",
+    originalUserInstructions(outputs),
+  ].join("\n");
 }
 
 function requireString(value: unknown, label: string): string {
@@ -286,7 +288,7 @@ function candidateSummary(
   ideal: AutoplanIdeal,
   id: string,
 ): { id: string; title: string; gist: string } {
-  if (id === "ideal") return { id, title: "Ideal end state", gist: ideal.ideal };
+  if (id === "ideal") return { id, title: "Holy grail", gist: ideal.ideal };
   const candidate = proposal.candidates.find((item) => item.id === id);
   if (candidate === undefined)
     throw new Error(`autoplan candidate ${JSON.stringify(id)} is missing`);
@@ -298,9 +300,9 @@ function summaryInput(
   blocked: boolean,
   input: AutoplanInput,
 ): PlainSummaryInput {
-  const proposal = outputs.propose as AutoplanProposal;
-  const selection = outputs.choose as AutoplanSelection;
-  const ideal = outputs.ideal as AutoplanIdeal;
+  const proposal = outputs.solutions as AutoplanProposal;
+  const selection = outputs.select as AutoplanSelection;
+  const ideal = outputs.holyGrail as AutoplanIdeal;
   const selected = candidateSummary(proposal, ideal, selection.selectedId);
   return {
     source: {
@@ -368,6 +370,7 @@ export const autoplanWorkflow = defineWorkflow({
       statusDetail: "capturing the user's instructions",
       prompt: () =>
         [
+          "Continue in this Pi session. Do not delegate this workflow step or start another session.",
           "Read the conversation that came before this workflow step.",
           "Return one text string named originalUserInstructions.",
           "Include everything that the user has instructed for the intended purpose in the given context.",
@@ -398,12 +401,13 @@ export const autoplanWorkflow = defineWorkflow({
       expectedOutput: `{ "problem": "concise statement", "success": ["criterion"], "inScope": ["change"], "outOfScope": ["change"], "constraints": ["constraint"], "controlBoundary": "what can change" }`,
       validate: (value) => requireRecord(value, "autoplan frame"),
     }),
-    propose: agent({
-      statusDetail: "devising candidate solutions",
+    solutions: agent({
+      statusDetail: "devising long-term solutions",
       prompt: ({ outputs, input }) =>
         [
           originalInstructionsPrompt(outputs),
           "Give two to four practical options that fit the allowed scope.",
+          "For each option, ask: Is this a Long term elegant and production ready solution?",
           "For each option return a stable lowercase id and a short title. Add a plain gist and full solution, explain the reason and trade-offs, and list the parts.",
           "Favor a few reusable parts with clear owners and use interfaces that already exist.",
           "Reject one-off machinery and infrastructure that the task does not need.",
@@ -416,42 +420,46 @@ export const autoplanWorkflow = defineWorkflow({
       validate: (value, { input }) =>
         parseProposal(value, (input as AutoplanInput).previousPlan !== undefined),
     }),
-    ideal: agent({
-      statusDetail: "describing the ideal end state",
+    holyGrail: agent({
+      statusDetail: "describing the Holy grail",
       prompt: ({ outputs, input }) =>
         [
           originalInstructionsPrompt(outputs),
-          "Describe the best possible end state separately from the practical options.",
-          "It may match one option or go beyond the current scope.",
+          "Describe the Holy grail separately from the practical options.",
+          "Ask: Is this the Holy grail for the problem?",
+          "The Holy grail may match one option or go beyond the current scope.",
+          "Explain what makes it more Long term elegant and production ready than the practical options.",
           "List each dependency we do not control. Do not assume that it can change.",
-          "State what this end state would improve beyond the practical options.",
+          "State what the Holy grail would improve beyond the practical options.",
           `Problem frame: ${JSON.stringify(outputs.frame)}`,
-          `Candidates: ${JSON.stringify(outputs.propose)}`,
+          `Candidates: ${JSON.stringify(outputs.solutions)}`,
           `New evidence: ${JSON.stringify((input as AutoplanInput).newEvidence ?? null)}`,
         ].join("\n"),
-      expectedOutput: `{ "ideal": "ideal end state", "outsideDependencies": ["dependency"], "additionalValue": ["benefit"] }`,
+      expectedOutput: `{ "ideal": "Holy grail end state", "outsideDependencies": ["dependency"], "additionalValue": ["benefit"] }`,
       validate: parseIdeal,
     }),
-    choose: agent({
-      statusDetail: "choosing the practical solution",
+    select: agent({
+      statusDetail: "selecting the solution",
       prompt: ({ outputs }) =>
         [
           originalInstructionsPrompt(outputs),
           "Select one option. Do not ask the user to choose.",
-          "Select the ideal only when it fits the allowed scope and is ready for production. Its value must justify the added complexity.",
-          "Otherwise select the best option we can build now that still moves toward the ideal.",
+          "Select the most Long term elegant and production ready option that is proportionate, in scope, and implementable through interfaces we control.",
+          "Select the Holy grail when it meets those conditions and its value justifies the added complexity.",
+          "Otherwise select the strongest practical option we can build now with a clear path toward the Holy grail.",
           "Work outside our control does not by itself make the plan blocked.",
           "Do not require changes to an upstream project or an unrelated repository. Do not require a new service or resource without approval.",
           "When two options solve the problem equally well, choose the simpler one.",
-          "Give one specific rejection reason for every option you do not select, including the ideal.",
+          "Give one specific rejection reason for every option you do not select, including the Holy grail.",
           "Return blocked only if no option inside the allowed scope can meet the success criteria.",
           `Frame: ${JSON.stringify(outputs.frame)}`,
-          `Candidates: ${JSON.stringify(outputs.propose)}`,
-          `Ideal candidate id: ideal`,
-          `Ideal: ${JSON.stringify(outputs.ideal)}`,
+          `Candidates: ${JSON.stringify(outputs.solutions)}`,
+          `Holy grail candidate id: ideal`,
+          `Holy grail: ${JSON.stringify(outputs.holyGrail)}`,
         ].join("\n"),
       expectedOutput: `{ "status": "ready" | "blocked", "selectedId": "candidate-id-or-ideal", "why": "reason", "relationshipToIdeal": "relationship", "rejected": [{ "id": "other-id", "reason": "why it lost" }], "compromises": ["compromise"], "blocker": "required only when blocked" }`,
-      validate: (value, { outputs }) => parseSelection(value, outputs.propose as AutoplanProposal),
+      validate: (value, { outputs }) =>
+        parseSelection(value, outputs.solutions as AutoplanProposal),
     }),
     plan: agent({
       timeoutMs: 30 * 60_000,
@@ -466,8 +474,8 @@ export const autoplanWorkflow = defineWorkflow({
           "Correct the earlier plan when the new evidence proves it wrong.",
           "Do not change files.",
           `Frame: ${JSON.stringify(outputs.frame)}`,
-          `Selection: ${JSON.stringify(outputs.choose)}`,
-          `Candidates: ${JSON.stringify(outputs.propose)}`,
+          `Selection: ${JSON.stringify(outputs.select)}`,
+          `Candidates: ${JSON.stringify(outputs.solutions)}`,
           `Previous plan: ${JSON.stringify((input as AutoplanInput).previousPlan ?? null)}`,
           `New evidence: ${JSON.stringify((input as AutoplanInput).newEvidence ?? null)}`,
         ].join("\n"),
@@ -476,13 +484,13 @@ export const autoplanWorkflow = defineWorkflow({
     }),
     blocked: compute({
       run: ({ outputs }) => {
-        const selection = outputs.choose as AutoplanSelection;
+        const selection = outputs.select as AutoplanSelection;
         return {
           status: "blocked",
           originalUserInstructions: originalUserInstructions(outputs),
           frame: outputs.frame,
-          proposal: outputs.propose as AutoplanProposal,
-          ideal: outputs.ideal as AutoplanIdeal,
+          proposal: outputs.solutions as AutoplanProposal,
+          ideal: outputs.holyGrail as AutoplanIdeal,
           selection,
           plainSummary: plainSummaryResult(outputs.blockedSummary, "autoplan blocked summary"),
           reason: requireString(selection.blocker, "autoplan blocker"),
@@ -499,9 +507,9 @@ export const autoplanWorkflow = defineWorkflow({
           status: "ready",
           originalUserInstructions: originalUserInstructions(outputs),
           frame: outputs.frame,
-          proposal: outputs.propose as AutoplanProposal,
-          ideal: outputs.ideal as AutoplanIdeal,
-          selection: outputs.choose as AutoplanSelection,
+          proposal: outputs.solutions as AutoplanProposal,
+          ideal: outputs.holyGrail as AutoplanIdeal,
+          selection: outputs.select as AutoplanSelection,
           plan: outputs.plan,
           plainSummary: plainSummaryResult(outputs.readySummary, "autoplan ready summary"),
           planDigest,
@@ -513,11 +521,11 @@ export const autoplanWorkflow = defineWorkflow({
   },
   edges: [
     { from: "captureIntent", to: "frame" },
-    { from: "frame", to: "propose" },
-    { from: "propose", to: "ideal" },
-    { from: "ideal", to: "choose" },
+    { from: "frame", to: "solutions" },
+    { from: "solutions", to: "holyGrail" },
+    { from: "holyGrail", to: "select" },
     {
-      from: "choose",
+      from: "select",
       switch: { on: "$.status", cases: { ready: "plan", blocked: "blockedSummary" } },
     },
     { from: "plan", to: "readySummary" },

package/src/builtins/catalog.ts

@@ -9,7 +9,7 @@ import sanityCheckWorkflow from "./sanity-check.workflow.js";
 
 export const builtinWorkflowCatalog = new BuiltinWorkflowCatalog([
   { id: "plain-summary", revision: "3", definition: plainSummaryWorkflow },
-  { id: "autoplan", revision: "5", definition: autoplanWorkflow },
+  { id: "autoplan", revision: "6", definition: autoplanWorkflow },
   { id: "autodoc", revision: "2", definition: autodocWorkflow },
   { id: "autoimplement", revision: "11", definition: autoimplementWorkflow },
   { id: "plan-approval", revision: "4", definition: planApprovalWorkflow },

package/src/controllers/index.ts

@@ -24,7 +24,10 @@ export {
   SqliteControllerStore,
   type RunEventRecord,
   type WorkflowNotificationRecord,
+  type WorkflowRunClaimOptions,
+  type WorkflowRunPreparationResult,
   type WorkflowRunQueueRecord,
+  type WorkflowRunReservationOptions,
 } from "./sqlite.js";
 export {
   type ControllerStore,

package/src/controllers/sqlite.ts

@@ -2,6 +2,7 @@ import { createHash, randomBytes, randomUUID } from "node:crypto";
 import fs from "node:fs";
 import path from "node:path";
 import { StateDatabase, workflowStatePath } from "../state/database.js";
+import { canonicalJson } from "../state/json.js";
 import { resourceIdFor, tokenHash } from "../state/mutation.js";
 import {
   EffectRequestConflictError,
@@ -39,6 +40,26 @@ export type WorkflowRunLaunchStatus =
   | "failed"
   | "cancelled";
 
+export type WorkflowRunReservationOptions = {
+  runId: string;
+  workflowName: string;
+  workflowSourceRef: string;
+  workflowSource: unknown;
+  definitionDigest: string;
+  definitionSnapshot: unknown;
+  input: unknown;
+  launchOptions?: unknown;
+  runnerId: string;
+  originSessionId: string;
+  parentRunId?: string;
+  now?: string;
+};
+
+export type WorkflowRunClaimOptions = WorkflowRunReservationOptions & {
+  claimToken: string;
+  leaseMs: number;
+};
+
 export type WorkflowRunQueueRecord = {
   runId: string;
   workflowName: string;
@@ -64,6 +85,16 @@ export type WorkflowRunQueueRecord = {
   finishedAt: string | null;
 };
 
+export type WorkflowRunPreparationResult =
+  | {
+      state: "claimed";
+      run: WorkflowRunQueueRecord & { claimToken: string };
+    }
+  | {
+      state: "adopted";
+      run: WorkflowRunQueueRecord;
+    };
+
 export type WorkflowNotificationRecord = {
   notificationId: string;
   runId: string;
@@ -84,7 +115,9 @@ export type WorkflowTurnIntentCause =
   | "failed"
   | "launchFailed"
   | "controllerInterrupted"
-  | "claimLost";
+  | "claimLost"
+  | "terminal"
+  | "cancelled";
 
 export type WorkflowTurnIntentResolution = "workflowPrompt" | "presentation" | "fallback";
 
@@ -996,139 +1029,147 @@ export class SqliteControllerStore implements ControllerStore {
     }));
   }
 
-  reserveWorkflowRun(options: {
-    runId: string;
-    workflowName: string;
-    workflowSourceRef: string;
-    workflowSource: unknown;
-    definitionDigest: string;
-    definitionSnapshot: unknown;
-    input: unknown;
-    launchOptions?: unknown;
-    runnerId: string;
-    originSessionId: string;
-    parentRunId?: string;
-    now?: string;
-  }): WorkflowRunQueueRecord {
+  reserveWorkflowRun(options: WorkflowRunReservationOptions): WorkflowRunQueueRecord {
     validateRunId(options.runId);
     const now = epoch(validTimestamp(options.now));
     const definitionDigest = digestBuffer(options.definitionDigest);
-    return this.state.transaction(() => {
-      if (this.getWorkflowRun(options.runId) !== undefined) {
-        throw new Error(`Workflow run already reserved: ${options.runId}`);
+    return this.state.transaction(() =>
+      this.reserveWorkflowRunInTransaction(options, now, definitionDigest),
+    );
+  }
+
+  private reserveWorkflowRunInTransaction(
+    options: WorkflowRunReservationOptions,
+    now: number,
+    definitionDigest: Buffer,
+  ): WorkflowRunQueueRecord {
+    if (this.getWorkflowRun(options.runId) !== undefined) {
+      throw new Error(`Workflow run already reserved: ${options.runId}`);
+    }
+    if (options.parentRunId !== undefined) {
+      const parent = this.requireWorkflowRunRow(options.parentRunId);
+      if (parent.originSessionId !== options.originSessionId) {
+        throw new Error("Continuation parent belongs to another Pi session");
       }
-      if (options.parentRunId !== undefined) {
-        const parent = this.requireWorkflowRunRow(options.parentRunId);
-        if (parent.originSessionId !== options.originSessionId) {
-          throw new Error("Continuation parent belongs to another Pi session");
+      if (parent.status === "parked") {
+        const parentLease = this.requireLease(parent.resourceId);
+        if (parentLease.ownerId !== null) {
+          throw new Error("Continuation parent still has an active owner");
         }
-        if (parent.status === "parked") {
-          const parentLease = this.requireLease(parent.resourceId);
-          if (parentLease.ownerId !== null) {
-            throw new Error("Continuation parent still has an active owner");
-          }
-          this.state.connection
-            .prepare(
-              `UPDATE run_queue
+        this.state.connection
+          .prepare(
+            `UPDATE run_queue
                SET status = 'done', updated_at = ?, finished_at = ?
                WHERE run_id = ? AND status = 'parked'`,
-            )
-            .run(now, now, parent.runId);
-          const parentRevision = this.resourceRevision(parent.resourceId);
-          this.bumpResource(parent.resourceId, parentRevision, now);
-          this.insertEvent(
-            parent.resourceId,
-            parentRevision + 1,
-            "run.queue_done_for_continuation",
-            "session",
-            options.originSessionId,
-            { continuationRunId: options.runId },
-            now,
-          );
-        } else if (parent.status === "done") {
-          throw new Error("Continuation parent already has a reserved continuation");
-        } else {
-          throw new Error(`Continuation parent queue is ${parent.status}`);
-        }
+          )
+          .run(now, now, parent.runId);
+        const parentRevision = this.resourceRevision(parent.resourceId);
+        this.bumpResource(parent.resourceId, parentRevision, now);
+        this.insertEvent(
+          parent.resourceId,
+          parentRevision + 1,
+          "run.queue_done_for_continuation",
+          "session",
+          options.originSessionId,
+          { continuationRunId: options.runId },
+          now,
+        );
+      } else if (parent.status === "done") {
+        throw new Error("Continuation parent already has a reserved continuation");
+      } else {
+        throw new Error(`Continuation parent queue is ${parent.status}`);
       }
-      const resourceId = resourceIdFor("run", options.runId);
-      const definitionHash = this.state.putJson(options.definitionSnapshot, now);
-      const queuedSource = queuedWorkflowSource(options.workflowSource);
-      const inputHash = this.state.putJson(options.input ?? null, now);
-      const launchHash = this.state.putJson(options.launchOptions ?? {}, now);
-      this.state.connection
-        .prepare(
-          `INSERT INTO workflow_definitions(
+    }
+    const resourceId = resourceIdFor("run", options.runId);
+    const definitionHash = this.state.putJson(options.definitionSnapshot, now);
+    const queuedSource = queuedWorkflowSource(options.workflowSource);
+    const inputHash = this.state.putJson(options.input ?? null, now);
+    const launchHash = this.state.putJson(options.launchOptions ?? {}, now);
+    this.state.connection
+      .prepare(
+        `INSERT INTO workflow_definitions(
              definition_digest, workflow_name, definition_hash, created_at
            ) VALUES (?, ?, ?, ?)
            ON CONFLICT(definition_digest) DO NOTHING`,
-        )
-        .run(definitionDigest, options.workflowName, definitionHash, now);
-      this.state.connection
-        .prepare(
-          `INSERT INTO resources(resource_id, resource_type, aggregate_key, revision, created_at, updated_at)
+      )
+      .run(definitionDigest, options.workflowName, definitionHash, now);
+    this.state.connection
+      .prepare(
+        `INSERT INTO resources(resource_id, resource_type, aggregate_key, revision, created_at, updated_at)
            VALUES (?, 'run', ?, 1, ?, ?)`,
-        )
-        .run(resourceId, options.runId, now, now);
-      this.state.connection
-        .prepare("INSERT INTO leases(resource_id, generation) VALUES (?, 0)")
-        .run(resourceId);
-      this.state.connection
-        .prepare(
-          `INSERT INTO runs(
+      )
+      .run(resourceId, options.runId, now, now);
+    this.state.connection
+      .prepare("INSERT INTO leases(resource_id, generation) VALUES (?, 0)")
+      .run(resourceId);
+    this.state.connection
+      .prepare(
+        `INSERT INTO runs(
              run_id, resource_id, project_id, parent_run_id, definition_digest,
              workflow_ref, launch_options_hash, status, paused,
              input_hash, created_at, updated_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?, ?, ?)`,
-        )
-        .run(
-          options.runId,
-          resourceId,
-          this.requireProjectId(),
-          options.parentRunId ?? null,
-          definitionDigest,
-          options.workflowSourceRef,
-          launchHash,
-          inputHash,
-          now,
-          now,
-        );
-      insertQueuedRunSources(this.state, options.runId, queuedSource);
-      this.state.connection
-        .prepare(
-          `INSERT INTO run_bindings(run_id, origin_session_id, execution_mode, created_at)
+      )
+      .run(
+        options.runId,
+        resourceId,
+        this.requireProjectId(),
+        options.parentRunId ?? null,
+        definitionDigest,
+        options.workflowSourceRef,
+        launchHash,
+        inputHash,
+        now,
+        now,
+      );
+    insertQueuedRunSources(this.state, options.runId, queuedSource);
+    this.state.connection
+      .prepare(
+        `INSERT INTO run_bindings(run_id, origin_session_id, execution_mode, created_at)
            VALUES (?, ?, 'interactive', ?)`,
-        )
-        .run(options.runId, options.originSessionId, now);
-      this.state.connection
-        .prepare(
-          `INSERT INTO run_queue(
+      )
+      .run(options.runId, options.originSessionId, now);
+    this.state.connection
+      .prepare(
+        `INSERT INTO run_queue(
              run_id, status, available_at, affinity_runner_id, origin_session_id,
              consecutive_errors, created_at, updated_at
            ) VALUES (?, 'queued', ?, ?, ?, 0, ?, ?)`,
-        )
-        .run(options.runId, now, options.runnerId, options.originSessionId, now, now);
-      this.insertEvent(resourceId, 1, "run.queued", "session", options.originSessionId, {}, now);
-      return this.requireWorkflowRun(options.runId);
+      )
+      .run(options.runId, now, options.runnerId, options.originSessionId, now, now);
+    this.insertEvent(resourceId, 1, "run.queued", "session", options.originSessionId, {}, now);
+    return this.requireWorkflowRun(options.runId);
+  }
+
+  prepareOrAdoptWorkflowRun(options: WorkflowRunClaimOptions): WorkflowRunPreparationResult {
+    validateRunId(options.runId);
+    const now = epoch(validTimestamp(options.now));
+    const definitionDigest = digestBuffer(options.definitionDigest);
+    return this.state.transaction(() => {
+      const existing = this.workflowRunRow(options.runId);
+      if (existing !== undefined) {
+        this.assertWorkflowRunPreparationCompatible(existing, options, definitionDigest);
+        return { state: "adopted", run: this.mapWorkflowRun(existing) };
+      }
+      this.reserveWorkflowRunInTransaction(options, now, definitionDigest);
+      const claimed = this.claimRunInTransaction(
+        options.runId,
+        options.runnerId,
+        options.claimToken,
+        options.leaseMs,
+        now,
+      );
+      if (claimed === undefined) {
+        throw new Error(`Workflow run could not be claimed: ${options.runId}`);
+      }
+      return {
+        state: "claimed",
+        run: claimed as WorkflowRunQueueRecord & { claimToken: string },
+      };
     });
   }
 
-  enqueueWorkflowRun(options: {
-    runId: string;
-    workflowName: string;
-    workflowSourceRef: string;
-    workflowSource: unknown;
-    definitionDigest: string;
-    definitionSnapshot: unknown;
-    input: unknown;
-    launchOptions?: unknown;
-    runnerId: string;
-    claimToken: string;
-    leaseMs: number;
-    originSessionId: string;
-    parentRunId?: string;
-    now?: string;
-  }): WorkflowRunQueueRecord {
+  enqueueWorkflowRun(options: WorkflowRunClaimOptions): WorkflowRunQueueRecord {
     if (this.getWorkflowRun(options.runId) === undefined) {
       this.reserveWorkflowRun({
         runId: options.runId,
@@ -2228,6 +2269,27 @@ export class SqliteControllerStore implements ControllerStore {
     return this.mapWorkflowRun(this.requireWorkflowRunRow(runId));
   }
 
+  private assertWorkflowRunPreparationCompatible(
+    row: RunRow,
+    options: WorkflowRunReservationOptions,
+    definitionDigest: Buffer,
+  ): void {
+    const compatible =
+      row.workflowName === options.workflowName &&
+      row.workflowRef === options.workflowSourceRef &&
+      row.definitionDigest.equals(definitionDigest) &&
+      canonicalJson(readQueuedRunSources(this.state, row)) ===
+        canonicalJson(queuedWorkflowSource(options.workflowSource)) &&
+      canonicalJson(this.state.readJson(row.inputHash)) === canonicalJson(options.input ?? null) &&
+      canonicalJson(this.state.readJson(row.launchOptionsHash)) ===
+        canonicalJson(options.launchOptions ?? {}) &&
+      row.originSessionId === options.originSessionId &&
+      row.parentRunId === (options.parentRunId ?? null);
+    if (!compatible) {
+      throw new Error(`Workflow run preparation conflicts: ${options.runId}`);
+    }
+  }
+
   /* istanbul ignore next -- pure projection covered by integration tests */
   private mapWorkflowRun(row: RunRow): WorkflowRunQueueRecord {
     return {
@@ -2267,56 +2329,65 @@ export class SqliteControllerStore implements ControllerStore {
     leaseMs: number,
     now: number,
   ): WorkflowRunQueueRecord | undefined {
-    return this.state.transaction(() => {
-      const row = this.workflowRunRow(runId);
-      if (row === undefined || ["done", "failed", "cancelled"].includes(row.status))
-        return undefined;
-      const lease = this.requireLease(row.resourceId);
-      if (
-        lease.ownerId !== null &&
-        lease.expiresAt !== null &&
-        lease.expiresAt > now &&
-        lease.ownerId !== runnerId
-      )
-        return undefined;
-      const generation = lease.generation + 1;
-      const expiresAt = now + leaseMs;
-      const result = this.state.connection
-        .prepare(
-          `UPDATE leases SET generation = ?, owner_type = ?, owner_id = ?, token_hash = ?,
+    return this.state.transaction(() =>
+      this.claimRunInTransaction(runId, runnerId, claimToken, leaseMs, now),
+    );
+  }
+
+  private claimRunInTransaction(
+    runId: string,
+    runnerId: string,
+    claimToken: string,
+    leaseMs: number,
+    now: number,
+  ): WorkflowRunQueueRecord | undefined {
+    const row = this.workflowRunRow(runId);
+    if (row === undefined || ["done", "failed", "cancelled"].includes(row.status)) return undefined;
+    const lease = this.requireLease(row.resourceId);
+    if (
+      lease.ownerId !== null &&
+      lease.expiresAt !== null &&
+      lease.expiresAt > now &&
+      lease.ownerId !== runnerId
+    )
+      return undefined;
+    const generation = lease.generation + 1;
+    const expiresAt = now + leaseMs;
+    const result = this.state.connection
+      .prepare(
+        `UPDATE leases SET generation = ?, owner_type = ?, owner_id = ?, token_hash = ?,
                   acquired_at = ?, heartbeat_at = ?, expires_at = ?
            WHERE resource_id = ? AND generation = ?`,
-        )
-        .run(
-          generation,
-          runnerId.startsWith("host-") ? "host" : "session",
-          runnerId,
-          tokenHash(claimToken),
-          now,
-          now,
-          expiresAt,
-          row.resourceId,
-          lease.generation,
-        );
-      /* istanbul ignore if -- impossible after exact schema and transaction checks */
-      if (result.changes !== 1) return undefined;
-      this.state.connection
-        .prepare("UPDATE run_queue SET status = 'starting', updated_at = ? WHERE run_id = ?")
-        .run(now, runId);
-      const revision = this.resourceRevision(row.resourceId);
-      this.bumpResource(row.resourceId, revision, now);
-      this.insertEvent(
-        row.resourceId,
-        revision + 1,
-        "lease.claimed",
+      )
+      .run(
+        generation,
         runnerId.startsWith("host-") ? "host" : "session",
         runnerId,
-        { expiresAt },
+        tokenHash(claimToken),
         now,
-        generation,
+        now,
+        expiresAt,
+        row.resourceId,
+        lease.generation,
       );
-      return { ...this.requireWorkflowRun(runId), claimToken };
-    });
+    /* istanbul ignore if -- impossible after exact schema and transaction checks */
+    if (result.changes !== 1) return undefined;
+    this.state.connection
+      .prepare("UPDATE run_queue SET status = 'starting', updated_at = ? WHERE run_id = ?")
+      .run(now, runId);
+    const revision = this.resourceRevision(row.resourceId);
+    this.bumpResource(row.resourceId, revision, now);
+    this.insertEvent(
+      row.resourceId,
+      revision + 1,
+      "lease.claimed",
+      runnerId.startsWith("host-") ? "host" : "session",
+      runnerId,
+      { expiresAt },
+      now,
+      generation,
+    );
+    return { ...this.requireWorkflowRun(runId), claimToken };
   }
 
   private updateClaimedRunStatus(