@osolmaz/pi-workflows 0.16.1 → 0.16.2

package/src/controllers/sqlite.ts

@@ -5,6 +5,7 @@ import { StateDatabase, workflowStatePath } from "../state/database.js";
 import { canonicalJson } from "../state/json.js";
 import { resourceIdFor, tokenHash } from "../state/mutation.js";
 import { initializeViewerRun, recordViewerDeltas } from "../state/viewer.js";
+import { WorkflowMessageStore } from "../state/workflow-messages.js";
 import {
   EffectRequestConflictError,
   ResourceConflictError,
@@ -270,6 +271,7 @@ type WorkflowRow = {
 export class SqliteControllerStore implements ControllerStore {
   readonly filePath: string;
   readonly state: StateDatabase;
+  private readonly workflowMessages: WorkflowMessageStore;
   private readonly ownsState: boolean;
   private readonly projectId: string | null;
   private closed = false;
@@ -293,6 +295,7 @@ export class SqliteControllerStore implements ControllerStore {
         checkLegacyState: filePath === workflowStatePath(),
       });
     this.filePath = this.state.filePath;
+    this.workflowMessages = new WorkflowMessageStore(this.state);
     if (options.global === true) {
       this.projectId = null;
     } else {
@@ -1556,7 +1559,7 @@ export class SqliteControllerStore implements ControllerStore {
       "(q.affinity_runner_id IS NULL OR q.affinity_runner_id = ? OR q.status IN ('parked', 'starting', 'running'))",
       "NOT (q.status = 'parked' AND (r.status = 'waiting' OR r.paused = 1))",
       "NOT EXISTS (SELECT 1 FROM interactive_requests i WHERE i.run_id = r.run_id AND i.status = 'pending')",
-      "(q.error_code IS NULL OR q.error_code <> 'workflowSourceChanged')",
+      "(q.error_code IS NULL OR q.error_code NOT IN ('workflowSourceChanged', 'workerNoProgress'))",
     ];
     const params: unknown[] = [now, now, options.runnerId, options.runnerId];
     if (this.projectId !== null) {
@@ -1669,6 +1672,78 @@ export class SqliteControllerStore implements ControllerStore {
     return this.releaseRunClaim(options.runId, options.claimToken, "parked", options.now);
   }
 
+  parkWorkflowRunForWorkerNoProgress(options: {
+    runId: string;
+    claimToken: string;
+    detail: string;
+    now?: string;
+  }): boolean {
+    const now = epoch(validTimestamp(options.now));
+    return this.state.transaction(() => {
+      if (
+        !this.verifyWorkflowRunClaim({
+          runId: options.runId,
+          claimToken: options.claimToken,
+          now: new Date(now).toISOString(),
+        })
+      ) {
+        return false;
+      }
+      const row = this.requireWorkflowRunRow(options.runId);
+      const lease = this.requireLease(row.resourceId);
+      const detail = options.detail.slice(0, 8_192);
+      const errorHash = this.state.putText(detail, now);
+      const run = this.state.connection
+        .prepare(
+          `UPDATE runs SET status_detail = ?, updated_at = ?, finished_at = NULL
+           WHERE run_id = ? AND status NOT IN ('completed', 'failed', 'timed_out', 'cancelled')`,
+        )
+        .run(detail, now, options.runId);
+      const queue = this.state.connection
+        .prepare(
+          `UPDATE run_queue
+           SET status = 'parked', error_code = 'workerNoProgress', error_hash = ?,
+               available_at = ?, updated_at = ?, finished_at = NULL
+           WHERE run_id = ? AND status IN ('queued', 'starting', 'running', 'parked')`,
+        )
+        .run(errorHash, now, now, options.runId);
+      if (run.changes !== 1 || queue.changes !== 1) {
+        throw new Error(`Workflow run ${options.runId} changed during worker recovery`);
+      }
+      const released = this.state.connection
+        .prepare(
+          `UPDATE leases
+           SET owner_type = NULL, owner_id = NULL, token_hash = NULL,
+               acquired_at = NULL, heartbeat_at = NULL, expires_at = NULL
+           WHERE resource_id = ? AND token_hash = ? AND generation = ? AND expires_at > ?`,
+        )
+        .run(row.resourceId, tokenHash(options.claimToken), lease.generation, now);
+      /* istanbul ignore if -- the exact live claim is stable in this transaction */
+      if (released.changes !== 1) {
+        throw new Error(`Workflow run ${options.runId} claim changed during worker recovery`);
+      }
+      const revision = this.resourceRevision(row.resourceId);
+      this.bumpResource(row.resourceId, revision, now);
+      this.insertEvent(
+        row.resourceId,
+        revision + 1,
+        "run.worker_no_progress",
+        lease.ownerType ?? "system",
+        lease.ownerId,
+        { status: "parked", code: "workerNoProgress" },
+        now,
+        lease.generation || undefined,
+      );
+      recordViewerDeltas(
+        this.state,
+        options.runId,
+        [{ targetType: "summary" }, { targetType: "replay" }],
+        now,
+      );
+      return true;
+    });
+  }
+
   pauseParkedWorkflowRun(options: { runId: string; now?: string }): boolean {
     const now = epoch(validTimestamp(options.now));
     return this.state.transaction(() => {
@@ -2015,6 +2090,7 @@ export class SqliteControllerStore implements ControllerStore {
       if (run.changes !== 1) {
         throw new Error(`Workflow run ${options.runId} has inconsistent durable state`);
       }
+      this.settleWorkflowRunMessages(options.runId, now);
       this.cancelWorkflowRunDependents(options.runId, controlId, errorHash, now);
 
       const released = this.state.connection
@@ -2924,6 +3000,7 @@ export class SqliteControllerStore implements ControllerStore {
            WHERE run_id = ?`,
         )
         .run(status, errorHash, now, now, runId);
+      this.settleWorkflowRunMessages(runId, now);
       if (status === "cancelled") {
         this.cancelWorkflowRunDependents(
           runId,
@@ -2953,6 +3030,11 @@ export class SqliteControllerStore implements ControllerStore {
     });
   }
 
+  private settleWorkflowRunMessages(runId: string, now: number): void {
+    this.workflowMessages.settleOpenTurnsForRun(runId, "lost", now);
+    this.workflowMessages.cancelPendingForRun(runId, now);
+  }
+
   private cancelWorkflowRunDependents(
     runId: string,
     actorId: string,

package/src/extension/index.ts

@@ -183,7 +183,9 @@ export default function piWorkflows(pi: ExtensionAPI): void {
   const workflowMessages = new WorkflowMessageCoordinator();
   const sessionView = new SessionWorkflowView();
   const sessionRecorders = new Map<string, SessionRecorder>();
+  let agentRunning = false;
   let activeRecorder: SessionRecorder | null = null;
+  let activeRecorderMessageId: string | null = null;
 
   const ensureRecorder = async (
     message: WorkflowMessage,
@@ -203,6 +205,33 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     return recorder;
   };
 
+  const activateRecorder = async (ctx: ExtensionContext): Promise<void> => {
+    if (!agentRunning) return;
+    const message = workflowMessages.activeTurnMessage();
+    if (message === undefined || activeRecorderMessageId === message.workflowMessageId) return;
+    const contract = agentContractForWorkflowMessage(message);
+    if (contract === undefined && message.kind !== "terminal" && message.kind !== "followUp") {
+      return;
+    }
+    const recorder = await ensureRecorder(message, ctx);
+    if (
+      !agentRunning ||
+      workflowMessages.activeTurnMessage()?.workflowMessageId !== message.workflowMessageId
+    ) {
+      return;
+    }
+    if (contract === undefined) {
+      recorder.beginWorkflowMessage(
+        message.workflowMessageId,
+        message.kind as "terminal" | "followUp",
+      );
+    } else {
+      recorder.beginAttempt(contract);
+    }
+    activeRecorder = recorder;
+    activeRecorderMessageId = message.workflowMessageId;
+  };
+
   const presentInOrder = async (ctx: ExtensionContext): Promise<void> => {
     const prior = presentationTail;
     let release: (() => void) | undefined;
@@ -212,6 +241,7 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     await prior;
     try {
       await workflowMessages.synchronize(pi, client, ctx);
+      await activateRecorder(ctx);
     } finally {
       sessionView.refresh(ctx);
       release?.();
@@ -651,27 +681,11 @@ export default function piWorkflows(pi: ExtensionAPI): void {
   });
 
   pi.on("agent_start", async (_event, ctx) => {
+    agentRunning = true;
+    activeRecorder = null;
+    activeRecorderMessageId = null;
     workflowMessages.startTurn();
     await presentInOrder(ctx).catch(() => undefined);
-    const message = workflowMessages.activeTurnMessage();
-    const contract = message === undefined ? undefined : agentContractForWorkflowMessage(message);
-    if (
-      message === undefined ||
-      (contract === undefined && message.kind !== "terminal" && message.kind !== "followUp")
-    ) {
-      activeRecorder = null;
-      return;
-    }
-    const recorder = await ensureRecorder(message, ctx);
-    if (contract === undefined) {
-      recorder.beginWorkflowMessage(
-        message.workflowMessageId,
-        message.kind as "terminal" | "followUp",
-      );
-    } else {
-      recorder.beginAttempt(contract);
-    }
-    activeRecorder = recorder;
   });
 
   pi.on("agent_end", async (event, ctx) => {
@@ -681,6 +695,8 @@ export default function piWorkflows(pi: ExtensionAPI): void {
       ctx.ui.notify(`Workflow response was rejected: ${errorMessage(error)}`, "error");
     }
     const finishedMessage = workflowMessages.activeTurnMessage();
+    agentRunning = false;
+    activeRecorderMessageId = null;
     workflowMessages.endTurn(
       workflowTurnStopReason(event.messages, ctx.signal?.aborted === true),
       responseEntryId(ctx.sessionManager.getBranch()),
@@ -747,7 +763,9 @@ export default function piWorkflows(pi: ExtensionAPI): void {
   pi.on("session_shutdown", async (_event, ctx) => {
     sessionGeneration += 1;
     sessionContext = null;
+    agentRunning = false;
     activeRecorder = null;
+    activeRecorderMessageId = null;
     await Promise.allSettled(
       [...sessionRecorders.values()].map(async (recorder) => recorder.stop()),
     );

package/src/extension/workflow-message-coordinator.ts

@@ -1,9 +1,18 @@
 import { randomUUID } from "node:crypto";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import type { WorkflowClient } from "../client/client.js";
-import type { WorkflowSessionView, WorkflowTurnReport } from "../client/view.js";
+import {
+  WORKFLOW_TURN_REPORT_RECEIPT_SCHEMA,
+  type WorkflowSessionView,
+  type WorkflowTurnReport,
+  type WorkflowTurnReportReceipt,
+} from "../client/view.js";
 import type { JsonValue } from "../state/json.js";
-import type { WorkflowMessage, WorkflowTurnStopReason } from "../state/workflow-messages.js";
+import {
+  WORKFLOW_TURN_SCHEMA,
+  type WorkflowMessage,
+  type WorkflowTurnStopReason,
+} from "../state/workflow-messages.js";
 
 const WORKFLOW_MESSAGE_ID_FIELD = "workflowMessageId";
 
@@ -11,6 +20,7 @@ type PendingTurn = {
   workflowTurnId: string;
   workflowMessageId: string | null;
   runId: string | null;
+  message: WorkflowMessage | null;
   startedReported: boolean;
   end: { stopReason: WorkflowTurnStopReason; responseSessionEntryId: string | null } | null;
 };
@@ -22,6 +32,7 @@ export class WorkflowMessageCoordinator {
   private synchronizing = false;
   private lastBranchEpoch: string | null = null;
   private view: WorkflowSessionView | null = null;
+  private awaitingTurnMessage: WorkflowMessage | null = null;
   private turn: PendingTurn | null = null;
 
   updateView(view: WorkflowSessionView): void {
@@ -50,15 +61,17 @@ export class WorkflowMessageCoordinator {
   }
 
   startTurn(): void {
-    if (this.turn !== null && this.turn.end === null) return;
+    const awaited = this.awaitingTurnMessage;
+    this.awaitingTurnMessage = null;
     this.turn = {
       workflowTurnId: `workflow-turn-${randomUUID()}`,
       workflowMessageId: null,
       runId: null,
+      message: null,
       startedReported: false,
       end: null,
     };
-    const candidate = this.turnCandidate();
+    const candidate = awaited ?? this.turnCandidate();
     if (candidate !== undefined) this.bindTurn(candidate);
   }
 
@@ -68,10 +81,8 @@ export class WorkflowMessageCoordinator {
   }
 
   activeTurnMessage(): WorkflowMessage | undefined {
-    if (this.view === null || this.turn?.workflowMessageId === null || this.turn === null) {
-      return undefined;
-    }
-    return messageById(this.view, this.turn.workflowMessageId);
+    if (this.turn === null || !this.turn.startedReported) return undefined;
+    return this.turn.message ?? undefined;
   }
 
   async synchronize(
@@ -89,7 +100,8 @@ export class WorkflowMessageCoordinator {
           workflowTurnId: view.openWorkflowTurn.workflowTurnId,
           workflowMessageId: view.openWorkflowTurn.workflowMessageId,
           runId: view.openWorkflowTurn.runId,
-          startedReported: false,
+          message: messageById(view, view.openWorkflowTurn.workflowMessageId) ?? null,
+          startedReported: true,
           end: null,
         };
       }
@@ -128,6 +140,9 @@ export class WorkflowMessageCoordinator {
         this.queued.delete(messageId);
         return;
       }
+      if (message.content.triggerTurn && messageStartsTurn(message)) {
+        this.awaitingTurnMessage = message;
+      }
       try {
         pi.sendMessage(
           {
@@ -140,6 +155,9 @@ export class WorkflowMessageCoordinator {
         );
       } catch (error) {
         this.queued.delete(messageId);
+        if (this.awaitingTurnMessage?.workflowMessageId === messageId) {
+          this.awaitingTurnMessage = null;
+        }
         throw error;
       }
       await this.reportBranch(client, ctx, view);
@@ -153,6 +171,7 @@ export class WorkflowMessageCoordinator {
     this.queued.clear();
     this.closedTurnMessages.clear();
     this.view = null;
+    this.awaitingTurnMessage = null;
     this.turn = null;
     this.lastBranchEpoch = null;
     this.synchronizing = false;
@@ -182,6 +201,7 @@ export class WorkflowMessageCoordinator {
     if (this.turn === null) return;
     this.turn.workflowMessageId = message.workflowMessageId;
     this.turn.runId = message.runId;
+    this.turn.message = message;
   }
 
   private async flushTurn(client: WorkflowClient, view: WorkflowSessionView): Promise<void> {
@@ -195,31 +215,47 @@ export class WorkflowMessageCoordinator {
     ) {
       return;
     }
-    const message = messageById(view, pending.workflowMessageId);
-    if (message?.status !== "sent") return;
+    let message = pending.message;
     if (!pending.startedReported) {
+      message = messageById(view, pending.workflowMessageId) ?? null;
+      if (message?.status !== "sent") return;
+      try {
+        const receipt = await reportTurn(client, {
+          state: "started",
+          workflowMessageId: pending.workflowMessageId,
+          workflowTurnId: pending.workflowTurnId,
+          runId: pending.runId,
+          targetSessionId: view.sessionId,
+          coordinatorEpoch: view.coordinatorEpoch,
+        });
+        if (receipt.ownership !== "active") {
+          this.turn = null;
+          return;
+        }
+        pending.message = message;
+        pending.startedReported = true;
+      } catch (error) {
+        this.turn = null;
+        throw error;
+      }
+    }
+    if (pending.end === null) return;
+    try {
       await reportTurn(client, {
-        state: "started",
+        state: "ended",
         workflowMessageId: pending.workflowMessageId,
         workflowTurnId: pending.workflowTurnId,
         runId: pending.runId,
         targetSessionId: view.sessionId,
         coordinatorEpoch: view.coordinatorEpoch,
+        stopReason: pending.end.stopReason,
+        responseSessionEntryId: pending.end.responseSessionEntryId,
       });
-      pending.startedReported = true;
+    } catch (error) {
+      this.turn = null;
+      throw error;
     }
-    if (pending.end === null) return;
-    await reportTurn(client, {
-      state: "ended",
-      workflowMessageId: pending.workflowMessageId,
-      workflowTurnId: pending.workflowTurnId,
-      runId: pending.runId,
-      targetSessionId: view.sessionId,
-      coordinatorEpoch: view.coordinatorEpoch,
-      stopReason: pending.end.stopReason,
-      responseSessionEntryId: pending.end.responseSessionEntryId,
-    });
-    if (message.kind === "terminal" || message.kind === "followUp") {
+    if (message?.kind === "terminal" || message?.kind === "followUp") {
       this.closedTurnMessages.add(pending.workflowMessageId);
     }
     for (const current of new Set([view, this.view])) {
@@ -228,7 +264,7 @@ export class WorkflowMessageCoordinator {
       }
       if (
         current?.openWorkflowMessageId === pending.workflowMessageId &&
-        (message.kind === "terminal" || message.kind === "followUp")
+        (message?.kind === "terminal" || message?.kind === "followUp")
       ) {
         current.openWorkflowMessageId = null;
       }
@@ -318,7 +354,10 @@ function messageStartsTurn(message: WorkflowMessage): boolean {
   return message.kind === "step" || message.kind === "terminal" || message.kind === "followUp";
 }
 
-async function reportTurn(client: WorkflowClient, report: WorkflowTurnReport): Promise<void> {
+async function reportTurn(
+  client: WorkflowClient,
+  report: WorkflowTurnReport,
+): Promise<WorkflowTurnReportReceipt> {
   const response = await client.request({
     operation: "workflowTurn.report",
     runId: report.runId,
@@ -327,6 +366,31 @@ async function reportTurn(client: WorkflowClient, report: WorkflowTurnReport): P
   if (response.outcome !== "accepted" && response.outcome !== "adopted") {
     throw new Error(response.error ?? "Workflow host rejected the model-turn report");
   }
+  const receipt = response.receipt;
+  if (
+    !isRecord(receipt) ||
+    receipt.schema !== WORKFLOW_TURN_REPORT_RECEIPT_SCHEMA ||
+    !["active", "settled", "absent"].includes(String(receipt.ownership))
+  ) {
+    throw new Error("Workflow host returned an invalid model-turn receipt");
+  }
+  const ownership = receipt.ownership as WorkflowTurnReportReceipt["ownership"];
+  if (ownership === "absent") {
+    if (receipt.turn !== null) {
+      throw new Error("Workflow host returned an invalid model-turn receipt");
+    }
+  } else if (
+    !isRecord(receipt.turn) ||
+    receipt.turn.schema !== WORKFLOW_TURN_SCHEMA ||
+    receipt.turn.workflowTurnId !== report.workflowTurnId ||
+    receipt.turn.workflowMessageId !== report.workflowMessageId ||
+    receipt.turn.runId !== report.runId ||
+    receipt.turn.targetSessionId !== report.targetSessionId ||
+    receipt.turn.state !== (ownership === "active" ? "started" : "ended")
+  ) {
+    throw new Error("Workflow host returned an invalid model-turn receipt");
+  }
+  return receipt as unknown as WorkflowTurnReportReceipt;
 }
 
 function isRecord(value: unknown): value is Record<string, unknown> {