@osolmaz/pi-workflows 0.15.1 → 0.15.3

package/src/extension/index.ts

@@ -13,6 +13,8 @@ import { discoverWorkflows } from "../workflows/loader.js";
 import { createRunId, WorkflowRunStore } from "../workflows/store.js";
 import type { AgentStepContract, HumanDecisionResponse } from "../workflows/types.js";
 import { parseControllerArgs, type ParsedControllerArgs } from "./controller-command.js";
+import { SessionDeliveryCoordinator, type ClaimedSessionDelivery } from "./session-delivery.js";
+import { SessionWorkflowView } from "./session-view.js";
 import {
   recoverAssistantStep,
   registerWorkflowAgentStepMessageRenderer,
@@ -131,6 +133,8 @@ export default function piWorkflows(pi: ExtensionAPI): void {
   let pollTimer: ReturnType<typeof setInterval> | null = null;
   let presentationTail = Promise.resolve();
   let toolTail = Promise.resolve();
+  const sessionDelivery = new SessionDeliveryCoordinator();
+  const sessionView = new SessionWorkflowView();
 
   const presentInOrder = async (ctx: ExtensionContext): Promise<void> => {
     const prior = presentationTail;
@@ -140,12 +144,13 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     });
     await prior;
     try {
-      await presentPendingInteraction(pi, client, ctx);
-      await deliverPendingNotification(pi, client, ctx);
-      if (pendingInteractionForSession(ctx.sessionManager.getSessionId()) === undefined) {
-        await presentPendingTurn(pi, client, ctx);
-      }
+      await sessionDelivery.synchronize(ctx, [
+        () => claimPendingInteractionDelivery(pi, client, ctx),
+        () => claimPendingNotificationDelivery(pi, client, ctx),
+        () => claimPendingTurnDelivery(pi, client, ctx),
+      ]);
     } finally {
+      sessionView.refresh(ctx);
       release?.();
     }
   };
@@ -266,6 +271,16 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     },
   });
 
+  pi.registerShortcut("shift+up", {
+    description: "Scroll the workflow widget up",
+    handler: (ctx) => sessionView.scrollUp(ctx),
+  });
+
+  pi.registerShortcut("shift+down", {
+    description: "Scroll the workflow widget down",
+    handler: (ctx) => sessionView.scrollDown(ctx),
+  });
+
   pi.on("session_start", async (_event, ctx) => {
     sessionContext = ctx;
     try {
@@ -280,6 +295,44 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     pollTimer.unref?.();
   });
 
+  pi.on("agent_end", async (event, ctx) => {
+    const interrupted =
+      ctx.signal?.aborted === true ||
+      event.messages.some(
+        (message) =>
+          isRecord(message) && "stopReason" in message && message.stopReason === "aborted",
+      );
+    if (!interrupted) return;
+    const interaction = pendingInteractionForSession(ctx.sessionManager.getSessionId());
+    if (
+      interaction === undefined ||
+      interaction.kind === "decision" ||
+      workflowRunPaused(interaction.runId)
+    ) {
+      return;
+    }
+    const turnHasPrompt = event.messages.some(
+      (message) => interactionRequestId(message) === interaction.requestId,
+    );
+    if (!turnHasPrompt) return;
+    try {
+      const pauseId = `escape-pause-${interaction.runId}-${randomUUID()}`;
+      await requestAccepted(client, {
+        operation: "run.pause",
+        requestId: pauseId,
+        idempotencyKey: pauseId,
+        runId: interaction.runId,
+      });
+      sessionView.refresh(ctx);
+      ctx.ui.notify(
+        `Workflow ${interaction.runId} paused because its model turn was interrupted. Use /workflow resume to continue.`,
+        "info",
+      );
+    } catch (error) {
+      ctx.ui.notify(`Could not pause interrupted workflow: ${errorMessage(error)}`, "warning");
+    }
+  });
+
   pi.on("agent_settled", async (_event, ctx) => {
     try {
       await submitVisibleAssistantResponse(client, ctx);
@@ -289,8 +342,10 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     await presentInOrder(ctx).catch(() => undefined);
   });
 
-  pi.on("session_shutdown", async () => {
+  pi.on("session_shutdown", async (_event, ctx) => {
     sessionContext = null;
+    sessionDelivery.clear();
+    sessionView.clear(ctx);
     if (pollTimer !== null) clearInterval(pollTimer);
     pollTimer = null;
   });
@@ -514,94 +569,120 @@ async function executeControllerCommand(
   };
 }
 
-async function presentPendingInteraction(
+async function claimPendingInteractionDelivery(
   pi: ExtensionAPI,
   client: WorkflowHostClient,
   ctx: ExtensionContext,
-): Promise<void> {
+): Promise<ClaimedSessionDelivery | undefined> {
   const interaction = pendingInteractionForSession(ctx.sessionManager.getSessionId());
-  if (interaction === undefined) return;
-  const entries = ctx.sessionManager.getBranch();
-  const existing = entries.find((entry) => interactionRequestId(entry) === interaction.requestId);
-  if (existing !== undefined) {
-    const entryId = entryIdentifier(existing);
-    if (entryId !== undefined && interaction.presentationSessionEntryId !== entryId) {
+  if (interaction === undefined || interaction.presentationSessionEntryId !== null)
+    return undefined;
+
+  const findSessionEntryId = (entries: readonly unknown[]): string | undefined =>
+    entryIdentifier(entries.find((entry) => interactionRequestId(entry) === interaction.requestId));
+  const existingEntryId = findSessionEntryId(ctx.sessionManager.getBranch());
+  let presentationRevision = interaction.revision;
+  let claimExpiresAt = Number.POSITIVE_INFINITY;
+  if (existingEntryId === undefined) {
+    if (workflowRunPaused(interaction.runId)) return undefined;
+    if (
+      interaction.presentationClaimExpiresAt !== null &&
+      Date.parse(interaction.presentationClaimExpiresAt) > Date.now()
+    ) {
+      return undefined;
+    }
+    await client.ensureRunning();
+    const claim = await client.request({
+      operation: "interaction.update",
+      requestId: `claim-presentation-${interaction.requestId}-${interaction.revision}-${client.clientId}`,
+      idempotencyKey: `claim-presentation-${interaction.requestId}-${interaction.revision}-${client.clientId}`,
+      runId: interaction.runId,
+      expectedRevision: interaction.revision,
+      payload: { requestId: interaction.requestId, claimPresentation: true },
+    });
+    if (claim.outcome === "conflict") return undefined;
+    if (claim.outcome !== "accepted" && claim.outcome !== "adopted") {
+      throw new Error(claim.error ?? "Workflow host rejected interaction.update");
+    }
+    if (claim.revision === undefined) throw new Error("Presentation claim has no revision");
+    const receipt = isRecord(claim.receipt) ? claim.receipt : undefined;
+    presentationRevision = claim.revision;
+    claimExpiresAt = claimExpiry(receipt?.presentationClaimExpiresAt, "presentation claim");
+  }
+
+  const contract = interactionContract(interaction);
+  return {
+    deliveryId: `interaction:${interaction.requestId}`,
+    claimExpiresAt,
+    isStillDeliverable: () =>
+      existingEntryId === undefined &&
+      interactionPresentationClaimIsLive({
+        requestId: interaction.requestId,
+        runId: interaction.runId,
+        presenterId: client.clientId,
+        revision: presentationRevision,
+        claimExpiresAt,
+      }),
+    findSessionEntryId,
+    send: () => {
+      if (interaction.kind === "decision") {
+        pi.sendMessage(
+          {
+            customType: WORKFLOW_INTERACTION_MESSAGE_TYPE,
+            content: decisionPrompt(contract),
+            display: true,
+            details: {
+              requestId: interaction.requestId,
+              runId: interaction.runId,
+              kind: "decision",
+            },
+          },
+          { triggerTurn: false },
+        );
+        return;
+      }
+      const agent = agentContract(interaction);
+      if (agent === undefined) throw new Error("Stored workflow agent contract is invalid");
+      const details: WorkflowAgentStepMessageDetails & { requestId: string } = {
+        schema: WORKFLOW_AGENT_STEP_MESSAGE_SCHEMA,
+        kind: "step",
+        contract: agent,
+        requestId: interaction.requestId,
+        ...(isRecord(contract.presentation)
+          ? { presentation: contract.presentation as never }
+          : {}),
+      };
+      pi.sendMessage(
+        {
+          customType: WORKFLOW_AGENT_STEP_MESSAGE_TYPE,
+          content:
+            typeof contract.prompt === "string" ? contract.prompt : "Continue the workflow step.",
+          display: true,
+          details,
+        },
+        { triggerTurn: true },
+      );
+    },
+    settle: async (sessionEntryId) => {
       await requestAccepted(client, {
         operation: "interaction.update",
-        requestId: `present-${interaction.requestId}-${entryId}`,
-        idempotencyKey: `present-${interaction.requestId}-${entryId}`,
+        requestId: `present-${interaction.requestId}-${sessionEntryId}`,
+        idempotencyKey: `present-${interaction.requestId}-${sessionEntryId}`,
         runId: interaction.runId,
-        expectedRevision: interaction.revision,
-        payload: { requestId: interaction.requestId, sessionEntryId: entryId },
+        expectedRevision: presentationRevision,
+        payload: { requestId: interaction.requestId, sessionEntryId },
       });
-    }
-    return;
-  }
-  if (interaction.presentationSessionEntryId !== null) return;
-  const claim = await requestAccepted(client, {
-    operation: "interaction.update",
-    requestId: `claim-presentation-${interaction.requestId}-${interaction.revision}-${client.clientId}`,
-    idempotencyKey: `claim-presentation-${interaction.requestId}-${interaction.revision}-${client.clientId}`,
-    runId: interaction.runId,
-    expectedRevision: interaction.revision,
-    payload: { requestId: interaction.requestId, claimPresentation: true },
-  });
-  if (claim.revision === undefined) throw new Error("Presentation claim has no revision");
-  const presentationRevision = claim.revision;
-  const contract = interactionContract(interaction);
-  if (interaction.kind === "decision") {
-    pi.sendMessage(
-      {
-        customType: WORKFLOW_INTERACTION_MESSAGE_TYPE,
-        content: decisionPrompt(contract),
-        display: true,
-        details: { requestId: interaction.requestId, runId: interaction.runId, kind: "decision" },
-      },
-      { triggerTurn: false },
-    );
-  } else {
-    const agent = agentContract(interaction);
-    if (agent === undefined) throw new Error("Stored workflow agent contract is invalid");
-    const details: WorkflowAgentStepMessageDetails & { requestId: string } = {
-      schema: WORKFLOW_AGENT_STEP_MESSAGE_SCHEMA,
-      kind: "step",
-      contract: agent,
-      requestId: interaction.requestId,
-      ...(isRecord(contract.presentation) ? { presentation: contract.presentation as never } : {}),
-    };
-    pi.sendMessage(
-      {
-        customType: WORKFLOW_AGENT_STEP_MESSAGE_TYPE,
-        content:
-          typeof contract.prompt === "string" ? contract.prompt : "Continue the workflow step.",
-        display: true,
-        details,
-      },
-      { triggerTurn: true, deliverAs: "followUp" },
-    );
-  }
-  const inserted = ctx.sessionManager
-    .getBranch()
-    .find((entry) => interactionRequestId(entry) === interaction.requestId);
-  const entryId = entryIdentifier(inserted);
-  if (entryId === undefined) return;
-  await requestAccepted(client, {
-    operation: "interaction.update",
-    requestId: `present-${interaction.requestId}-${entryId}`,
-    idempotencyKey: `present-${interaction.requestId}-${entryId}`,
-    runId: interaction.runId,
-    expectedRevision: presentationRevision,
-    payload: { requestId: interaction.requestId, sessionEntryId: entryId },
-  });
+    },
+  };
 }
 
-async function deliverPendingNotification(
+async function claimPendingNotificationDelivery(
   pi: ExtensionAPI,
   client: WorkflowHostClient,
   ctx: ExtensionContext,
-): Promise<void> {
+): Promise<ClaimedSessionDelivery | undefined> {
   const sessionId = ctx.sessionManager.getSessionId();
-  if (!hasClaimableNotification(sessionId)) return;
+  if (!hasClaimableNotification(sessionId)) return undefined;
   const claimRequestId = randomUUID();
   const claimed = await requestAccepted(client, {
     operation: "notification.claim",
@@ -611,57 +692,67 @@ async function deliverPendingNotification(
   });
   const receipt = isRecord(claimed.receipt) ? claimed.receipt : undefined;
   const notification = isRecord(receipt?.notification) ? receipt.notification : undefined;
-  if (notification === undefined) return;
+  if (notification === undefined) return undefined;
   const claimId = requireText(receipt?.claimId, "notification claimId");
+  const claimExpiresAt = claimExpiry(receipt?.claimExpiresAt, "notification claim");
   const notificationId = requireText(notification.notificationId, "notificationId");
-  const branch = ctx.sessionManager.getBranch();
-  let entry = branch.find(
-    (candidate) =>
-      customMessageDetail(candidate, WORKFLOW_NOTIFICATION_MESSAGE_TYPE, "notificationId") ===
-      notificationId,
-  );
-  if (entry === undefined) {
-    pi.sendMessage(
-      {
-        customType: WORKFLOW_NOTIFICATION_MESSAGE_TYPE,
-        content: requireText(notification.content, "notification content"),
-        display: true,
-        details: {
-          notificationId,
-          runId: requireText(notification.runId, "notification runId"),
-          kind: notification.kind,
+  return {
+    deliveryId: `notification:${notificationId}`,
+    claimExpiresAt,
+    isStillDeliverable: () =>
+      hostDeliveryClaimIsLive(client, {
+        kind: "notification",
+        resourceId: notificationId,
+        targetSessionId: sessionId,
+        claimId,
+      }),
+    findSessionEntryId: (entries) =>
+      entryIdentifier(
+        entries.find(
+          (entry) =>
+            customMessageDetail(entry, WORKFLOW_NOTIFICATION_MESSAGE_TYPE, "notificationId") ===
+            notificationId,
+        ),
+      ),
+    send: () => {
+      pi.sendMessage(
+        {
+          customType: WORKFLOW_NOTIFICATION_MESSAGE_TYPE,
+          content: requireText(notification.content, "notification content"),
+          display: true,
+          details: {
+            notificationId,
+            runId: requireText(notification.runId, "notification runId"),
+            kind: notification.kind,
+          },
         },
-      },
-      { triggerTurn: false },
-    );
-    entry = ctx.sessionManager
-      .getBranch()
-      .find(
-        (candidate) =>
-          customMessageDetail(candidate, WORKFLOW_NOTIFICATION_MESSAGE_TYPE, "notificationId") ===
-          notificationId,
+        { triggerTurn: false },
       );
-  }
-  if (entry === undefined) return;
-  await requestAccepted(client, {
-    operation: "notification.deliver",
-    requestId: `notification-deliver-${notificationId}-${claimId}`,
-    idempotencyKey: `notification-deliver-${notificationId}-${claimId}`,
-    payload: {
-      notificationId,
-      targetSessionId: sessionId,
-      claimId,
     },
-  });
+    settle: async () => {
+      await requestAccepted(client, {
+        operation: "notification.deliver",
+        requestId: `notification-deliver-${notificationId}-${claimId}`,
+        idempotencyKey: `notification-deliver-${notificationId}-${claimId}`,
+        payload: {
+          notificationId,
+          targetSessionId: sessionId,
+          claimId,
+        },
+      });
+    },
+  };
 }
 
-async function presentPendingTurn(
+async function claimPendingTurnDelivery(
   pi: ExtensionAPI,
   client: WorkflowHostClient,
   ctx: ExtensionContext,
-): Promise<void> {
+): Promise<ClaimedSessionDelivery | undefined> {
   const sessionId = ctx.sessionManager.getSessionId();
-  if (!hasClaimableTurn(sessionId)) return;
+  if (pendingInteractionForSession(sessionId) !== undefined || !hasClaimableTurn(sessionId)) {
+    return undefined;
+  }
   const claimRequestId = randomUUID();
   const claimed = await requestAccepted(client, {
     operation: "turn.claim",
@@ -671,49 +762,55 @@ async function presentPendingTurn(
   });
   const receipt = isRecord(claimed.receipt) ? claimed.receipt : undefined;
   const turn = isRecord(receipt?.turn) ? receipt.turn : undefined;
-  if (turn === undefined) return;
+  if (turn === undefined) return undefined;
   const claimId = requireText(receipt?.claimId, "turn claimId");
+  const claimExpiresAt = claimExpiry(receipt?.claimExpiresAt, "turn claim");
   const intentId = requireText(turn.intentId, "turn intentId");
   const runId = requireText(turn.runId, "turn runId");
   const state = terminalRunState(runId);
-  if (state === undefined) return;
-  let entry = ctx.sessionManager
-    .getBranch()
-    .find(
-      (candidate) =>
-        customMessageDetail(candidate, WORKFLOW_PRESENTATION_MESSAGE_TYPE, "intentId") === intentId,
-    );
-  if (entry === undefined) {
-    pi.sendMessage(
-      {
-        customType: WORKFLOW_PRESENTATION_MESSAGE_TYPE,
-        content: presentationMessage(turn, state),
-        display: false,
-        details: { intentId, runId },
-      },
-      { triggerTurn: true, deliverAs: "followUp" },
-    );
-    entry = ctx.sessionManager
-      .getBranch()
-      .find(
-        (candidate) =>
-          customMessageDetail(candidate, WORKFLOW_PRESENTATION_MESSAGE_TYPE, "intentId") ===
-          intentId,
+  if (state === undefined) return undefined;
+  return {
+    deliveryId: `turn:${intentId}`,
+    claimExpiresAt,
+    isStillDeliverable: () =>
+      hostDeliveryClaimIsLive(client, {
+        kind: "turn",
+        resourceId: intentId,
+        targetSessionId: sessionId,
+        claimId,
+      }),
+    findSessionEntryId: (entries) =>
+      entryIdentifier(
+        entries.find(
+          (entry) =>
+            customMessageDetail(entry, WORKFLOW_PRESENTATION_MESSAGE_TYPE, "intentId") === intentId,
+        ),
+      ),
+    send: () => {
+      pi.sendMessage(
+        {
+          customType: WORKFLOW_PRESENTATION_MESSAGE_TYPE,
+          content: presentationMessage(turn, state),
+          display: false,
+          details: { intentId, runId },
+        },
+        { triggerTurn: true },
       );
-  }
-  const messageId = entryIdentifier(entry);
-  if (messageId === undefined) return;
-  await requestAccepted(client, {
-    operation: "turn.resolve",
-    requestId: `turn-resolve-${intentId}-${messageId}`,
-    idempotencyKey: `turn-resolve-${intentId}-${messageId}`,
-    payload: {
-      intentId,
-      targetSessionId: sessionId,
-      claimId,
-      messageId,
     },
-  });
+    settle: async (sessionEntryId) => {
+      await requestAccepted(client, {
+        operation: "turn.resolve",
+        requestId: `turn-resolve-${intentId}-${sessionEntryId}`,
+        idempotencyKey: `turn-resolve-${intentId}-${sessionEntryId}`,
+        payload: {
+          intentId,
+          targetSessionId: sessionId,
+          claimId,
+          messageId: sessionEntryId,
+        },
+      });
+    },
+  };
 }
 
 async function submitVisibleAssistantResponse(
@@ -721,7 +818,13 @@ async function submitVisibleAssistantResponse(
   ctx: ExtensionContext,
 ): Promise<void> {
   const interaction = pendingInteractionForSession(ctx.sessionManager.getSessionId());
-  if (interaction === undefined || interaction.kind !== "assistant") return;
+  if (
+    interaction === undefined ||
+    interaction.kind !== "assistant" ||
+    workflowRunPaused(interaction.runId)
+  ) {
+    return;
+  }
   const contract = agentContract(interaction);
   if (contract === undefined) return;
   const submission = recoverAssistantStep(ctx.sessionManager.getBranch(), contract);
@@ -872,6 +975,65 @@ function terminalRunState(runId: string): Record<string, unknown> | undefined {
   }
 }
 
+function workflowRunPaused(runId: string): boolean {
+  return terminalRunState(runId)?.paused === true;
+}
+
+function interactionPresentationClaimIsLive(options: {
+  requestId: string;
+  runId: string;
+  presenterId: string;
+  revision: number;
+  claimExpiresAt: number;
+}): boolean {
+  try {
+    const store = new HostStateStore(workflowStatePath(), { readOnly: true });
+    try {
+      const interaction = store.getInteraction(options.requestId);
+      return (
+        interaction?.runId === options.runId &&
+        interaction.status === "presenting" &&
+        interaction.presenterId === options.presenterId &&
+        interaction.revision === options.revision &&
+        interaction.presentationSessionEntryId === null &&
+        interaction.presentationClaimExpiresAt !== null &&
+        Date.parse(interaction.presentationClaimExpiresAt) === options.claimExpiresAt &&
+        !workflowRunPaused(options.runId)
+      );
+    } finally {
+      store.close();
+    }
+  } catch {
+    return false;
+  }
+}
+
+async function hostDeliveryClaimIsLive(
+  client: WorkflowHostClient,
+  options: {
+    kind: "notification" | "turn";
+    resourceId: string;
+    targetSessionId: string;
+    claimId: string;
+  },
+): Promise<boolean> {
+  try {
+    const validationId = randomUUID();
+    const response = await client.request({
+      operation: options.kind === "notification" ? "notification.claim" : "turn.claim",
+      requestId: `delivery-validate-${options.claimId}-${validationId}`,
+      idempotencyKey: validationId,
+      payload: { ...options, validateClaim: true },
+    });
+    const receipt = isRecord(response.receipt) ? response.receipt : undefined;
+    return (
+      (response.outcome === "accepted" || response.outcome === "adopted") && receipt?.live === true
+    );
+  } catch {
+    return false;
+  }
+}
+
 function presentationMessage(
   turn: Record<string, unknown>,
   state: Record<string, unknown>,
@@ -927,6 +1089,12 @@ function requireText(value: unknown, name: string): string {
   return value;
 }
 
+function claimExpiry(value: unknown, name: string): number {
+  const expiry = Date.parse(requireText(value, `${name} expiry`));
+  if (!Number.isFinite(expiry)) throw new Error(`${name} expiry must be a timestamp`);
+  return expiry;
+}
+
 function activeSessionRun(ctx: ExtensionContext): WorkflowRunQueueRecord | undefined {
   return sessionRun(ctx);
 }
@@ -1021,7 +1189,11 @@ function agentContract(interaction: InteractiveRequestRecord): AgentStepContract
 }
 
 function interactionRequestId(value: unknown): string | undefined {
-  if (!isRecord(value) || value.type !== "custom_message" || !isRecord(value.details)) {
+  if (
+    !isRecord(value) ||
+    (value.type !== "custom_message" && value.role !== "custom") ||
+    !isRecord(value.details)
+  ) {
     return undefined;
   }
   return typeof value.details.requestId === "string" ? value.details.requestId : undefined;