@osolmaz/pi-workflows 0.16.4 → 0.16.6

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

@@ -16,19 +16,27 @@ import {
 
 const WORKFLOW_MESSAGE_ID_FIELD = "workflowMessageId";
 
+type SettledTurn = { stopReason: WorkflowTurnStopReason; responseSessionEntryId: string | null };
+type BeforeTurnEnd = (message: WorkflowMessage, end: SettledTurn) => Promise<void>;
+type DeliveryCallbacks = {
+  beforeTurnEnd?: BeforeTurnEnd;
+  terminalDelivered?: (message: WorkflowMessage) => Promise<void>;
+};
+
 type PendingTurn = {
   workflowTurnId: string;
   workflowMessageId: string | null;
   runId: string | null;
   message: WorkflowMessage | null;
   startedReported: boolean;
-  end: { stopReason: WorkflowTurnStopReason; responseSessionEntryId: string | null } | null;
+  end: SettledTurn | null;
 };
 
 /** Adds every server-owned workflow message to one Pi session through one public API path. */
 export class WorkflowMessageCoordinator {
   private readonly queued = new Set<string>();
   private readonly closedTurnMessages = new Set<string>();
+  private readonly finalizedTerminals = new Set<string>();
   private synchronizing = false;
   private lastBranchEpoch: string | null = null;
   private view: WorkflowSessionView | null = null;
@@ -41,19 +49,14 @@ export class WorkflowMessageCoordinator {
     for (const messageId of this.closedTurnMessages) {
       if (!visibleIds.has(messageId)) this.closedTurnMessages.delete(messageId);
     }
+    for (const messageId of this.finalizedTerminals) {
+      if (!visibleIds.has(messageId)) this.finalizedTerminals.delete(messageId);
+    }
     for (const message of view.workflowMessages) {
       if (message.status === "sent" || message.status === "cancelled") {
         this.queued.delete(message.workflowMessageId);
       }
     }
-    if (
-      this.turn !== null &&
-      this.turn.workflowMessageId === null &&
-      view.openWorkflowMessageId !== null
-    ) {
-      const message = messageById(view, view.openWorkflowMessageId);
-      if (message !== undefined && messageStartsTurn(message)) this.bindTurn(message);
-    }
   }
 
   branchChanged(): void {
@@ -61,6 +64,8 @@ export class WorkflowMessageCoordinator {
   }
 
   startTurn(): void {
+    // Automatic Pi retries belong to the same unsettled workflow turn.
+    if (this.turn !== null) return;
     const awaited = this.awaitingTurnMessage;
     this.awaitingTurnMessage = null;
     this.turn = {
@@ -71,13 +76,12 @@ export class WorkflowMessageCoordinator {
       startedReported: false,
       end: null,
     };
-    const candidate = awaited ?? this.turnCandidate();
-    if (candidate !== undefined) this.bindTurn(candidate);
+    if (awaited !== null) this.bindTurn(awaited);
   }
 
   endTurn(stopReason: WorkflowTurnStopReason, responseSessionEntryId: string | null): void {
-    if (this.turn === null) this.startTurn();
-    if (this.turn !== null) this.turn.end = { stopReason, responseSessionEntryId };
+    if (this.turn !== null && this.turn.end === null)
+      this.turn.end = { stopReason, responseSessionEntryId };
   }
 
   activeTurnMessage(): WorkflowMessage | undefined {
@@ -89,13 +93,26 @@ export class WorkflowMessageCoordinator {
     pi: ExtensionAPI,
     client: WorkflowClient,
     ctx: Pick<ExtensionContext, "hasPendingMessages" | "isIdle" | "sessionManager">,
+    callbacks: DeliveryCallbacks = {},
   ): Promise<void> {
     if (this.synchronizing || this.view === null) return;
     this.synchronizing = true;
     try {
       const view = this.view;
       if (!view.coordinatorActive || view.coordinatorEpoch === null) return;
-      if (this.turn === null && !ctx.isIdle() && view.openWorkflowTurn !== null) {
+      const branchEntries = branchWorkflowEntries(ctx.sessionManager.getBranch());
+      if (this.turn?.workflowMessageId === null) {
+        const candidate = this.turnCandidate();
+        if (candidate !== undefined && branchEntries.has(candidate.workflowMessageId)) {
+          this.bindTurn(candidate);
+        }
+      }
+      if (
+        this.turn === null &&
+        !ctx.isIdle() &&
+        view.openWorkflowTurn !== null &&
+        branchEntries.has(view.openWorkflowTurn.workflowMessageId)
+      ) {
         this.turn = {
           workflowTurnId: view.openWorkflowTurn.workflowTurnId,
           workflowMessageId: view.openWorkflowTurn.workflowMessageId,
@@ -112,7 +129,19 @@ export class WorkflowMessageCoordinator {
       ) {
         await this.reportBranch(client, ctx, view);
       }
-      await this.flushTurn(client, view);
+      await this.flushTurn(client, view, callbacks.beforeTurnEnd);
+      if (this.turn !== null || !ctx.isIdle() || ctx.hasPendingMessages()) return;
+      for (const message of view.workflowMessages) {
+        if (
+          message.kind === "terminal" &&
+          message.status === "sent" &&
+          branchEntries.has(message.workflowMessageId) &&
+          !this.finalizedTerminals.has(message.workflowMessageId)
+        ) {
+          await callbacks.terminalDelivered?.(message);
+          this.finalizedTerminals.add(message.workflowMessageId);
+        }
+      }
       const messageId = view.nextWorkflowMessageId;
       if (messageId === null || this.queued.has(messageId)) return;
       const message = messageById(view, messageId);
@@ -140,7 +169,7 @@ export class WorkflowMessageCoordinator {
         this.queued.delete(messageId);
         return;
       }
-      if (message.content.triggerTurn && messageStartsTurn(message)) {
+      if (messageStartsTurn(message)) {
         this.awaitingTurnMessage = message;
       }
       try {
@@ -161,7 +190,7 @@ export class WorkflowMessageCoordinator {
         throw error;
       }
       await this.reportBranch(client, ctx, view);
-      await this.flushTurn(client, view);
+      await this.flushTurn(client, view, callbacks.beforeTurnEnd);
     } finally {
       this.synchronizing = false;
     }
@@ -170,6 +199,7 @@ export class WorkflowMessageCoordinator {
   clear(): void {
     this.queued.clear();
     this.closedTurnMessages.clear();
+    this.finalizedTerminals.clear();
     this.view = null;
     this.awaitingTurnMessage = null;
     this.turn = null;
@@ -204,7 +234,11 @@ export class WorkflowMessageCoordinator {
     this.turn.message = message;
   }
 
-  private async flushTurn(client: WorkflowClient, view: WorkflowSessionView): Promise<void> {
+  private async flushTurn(
+    client: WorkflowClient,
+    view: WorkflowSessionView,
+    beforeTurnEnd?: BeforeTurnEnd,
+  ): Promise<void> {
     const pending = this.turn;
     if (
       pending === null ||
@@ -219,43 +253,37 @@ export class WorkflowMessageCoordinator {
     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: "ended",
+      const receipt = await reportTurn(client, {
+        state: "started",
         workflowMessageId: pending.workflowMessageId,
         workflowTurnId: pending.workflowTurnId,
         runId: pending.runId,
         targetSessionId: view.sessionId,
         coordinatorEpoch: view.coordinatorEpoch,
-        stopReason: pending.end.stopReason,
-        responseSessionEntryId: pending.end.responseSessionEntryId,
       });
-    } catch (error) {
-      this.turn = null;
-      throw error;
+      if (
+        receipt.ownership !== "active" &&
+        !(receipt.ownership === "settled" && pending.end !== null)
+      ) {
+        this.turn = null;
+        return;
+      }
+      pending.message = message;
+      pending.startedReported = true;
     }
-    if (message?.kind === "terminal" || message?.kind === "followUp") {
+    if (pending.end === null) return;
+    if (message !== null) await beforeTurnEnd?.(message, pending.end);
+    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 === "followUp") {
       this.closedTurnMessages.add(pending.workflowMessageId);
     }
     for (const current of new Set([view, this.view])) {
@@ -264,7 +292,7 @@ export class WorkflowMessageCoordinator {
       }
       if (
         current?.openWorkflowMessageId === pending.workflowMessageId &&
-        (message?.kind === "terminal" || message?.kind === "followUp")
+        message?.kind === "followUp"
       ) {
         current.openWorkflowMessageId = null;
       }
@@ -351,7 +379,7 @@ function messageById(
 }
 
 function messageStartsTurn(message: WorkflowMessage): boolean {
-  return message.kind === "step" || message.kind === "terminal" || message.kind === "followUp";
+  return message.content.triggerTurn;
 }
 
 async function reportTurn(

package/src/resource-managers/index.ts

@@ -18,16 +18,8 @@ export {
 } from "./errors.js";
 export { ResourceManagerEffectService } from "./effects.js";
 export { canonicalJson, jsonFingerprint } from "./json.js";
-export { ResourceManagerRuntime, type ResourceManagerRuntimeOptions } from "./runtime.js";
 export { createResultHelpers, requeue, requeueAfter, settled } from "./results.js";
-export {
-  SqliteResourceManagerStore,
-  type RunEventRecord,
-  type WorkflowRunClaimOptions,
-  type WorkflowRunPreparationResult,
-  type WorkflowRunQueueRecord,
-  type WorkflowRunReservationOptions,
-} from "./sqlite.js";
+export { SqliteResourceManagerStore } from "./sqlite.js";
 export {
   type ResourceManagerStore,
   type EffectReservation,