@osolmaz/pi-workflows 0.16.4 → 0.16.6

package/src/extension/index.ts

@@ -67,7 +67,7 @@ export type ParsedWorkflowArgs =
   | { kind: "resume" }
   | { kind: "status"; runId?: string }
   | { kind: "clear"; runId?: string }
-  | { kind: "restart"; runId?: string }
+  | { kind: "restart"; runId?: string; expectedRevision?: number }
   | {
       kind: "change-settings";
       patch: unknown;
@@ -77,7 +77,7 @@ export type ParsedWorkflowArgs =
     }
   | { kind: "queue-follow-up"; prompt: string; runId?: string }
   | { kind: "remove-follow-up"; followUpId: string; runId?: string }
-  | { kind: "answer"; input: unknown; runId?: string | undefined }
+  | { kind: "answer"; requestId: string; input: unknown }
   | { kind: "run"; ref: string; input: unknown };
 
 /** Parse `/workflow` arguments. Exported for tests. */
@@ -131,35 +131,22 @@ export function parseWorkflowArgs(args: string): ParsedWorkflowArgs {
     return { kind: "status", runId };
   }
   if (trimmed === "answer" || trimmed.startsWith("answer ")) {
-    let rest = trimmed === "answer" ? "" : trimmed.slice("answer".length).trim();
-    if (rest.length === 0) {
+    const rest = trimmed === "answer" ? "" : trimmed.slice("answer".length).trim();
+    const firstSpace = rest.search(/\s/);
+    if (firstSpace <= 0 || rest.slice(firstSpace).trim().length === 0) {
       throw new Error(
-        'answer requires a JSON value or text, e.g. /workflow answer {"approved":true}',
+        'answer requires an exact request id and a response: /workflow answer REQUEST_ID {"choice":"approve"}',
       );
     }
-    let runId: string | undefined;
-    const firstSpace = rest.search(/\s/);
-    if (firstSpace > 0) {
-      const candidate = rest.slice(0, firstSpace);
-      const remainder = rest.slice(firstSpace).trim();
-      if (
-        validRunId(candidate) &&
-        remainder.length > 0 &&
-        (remainder.startsWith("{") || remainder.startsWith("["))
-      ) {
-        runId = candidate;
-        rest = remainder;
-      }
-    }
+    const requestId = rest.slice(0, firstSpace);
+    if (!/^[A-Za-z0-9][A-Za-z0-9._-]*$/.test(requestId))
+      throw new Error("answer requires a valid request id");
+    const text = rest.slice(firstSpace).trim();
     try {
-      return {
-        kind: "answer",
-        input: JSON.parse(rest) as unknown,
-        ...(runId === undefined ? {} : { runId }),
-      };
+      return { kind: "answer", requestId, input: JSON.parse(text) as unknown };
     } catch {
-      if (runId !== undefined) throw new Error("answer with a run id requires a JSON value");
-      return { kind: "answer", input: { answer: rest } };
+      if (/^[{["]/.test(text)) throw new Error("answer contains invalid JSON");
+      return { kind: "answer", requestId, input: { answer: text } };
     }
   }
   const spaceIndex = trimmed.search(/\s/);
@@ -190,6 +177,7 @@ export default function piWorkflows(pi: ExtensionAPI): void {
   const sessionView = new SessionWorkflowView();
   const sessionRecorders = new Map<string, SessionRecorder>();
   let agentRunning = false;
+  let lastStopReason: ReturnType<typeof workflowTurnStopReason> = "completed";
   let activeRecorder: SessionRecorder | null = null;
   let activeRecorderMessageId: string | null = null;
 
@@ -216,7 +204,7 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     const message = workflowMessages.activeTurnMessage();
     if (message === undefined || activeRecorderMessageId === message.workflowMessageId) return;
     const contract = agentContractForWorkflowMessage(message);
-    if (contract === undefined && message.kind !== "terminal" && message.kind !== "followUp") {
+    if (contract === undefined && message.kind !== "followUp") {
       return;
     }
     const recorder = await ensureRecorder(message, ctx);
@@ -227,10 +215,7 @@ export default function piWorkflows(pi: ExtensionAPI): void {
       return;
     }
     if (contract === undefined) {
-      recorder.beginWorkflowMessage(
-        message.workflowMessageId,
-        message.kind as "terminal" | "followUp",
-      );
+      recorder.beginFollowUp(message.workflowMessageId);
     } else {
       recorder.beginAttempt(contract);
     }
@@ -246,7 +231,28 @@ export default function piWorkflows(pi: ExtensionAPI): void {
     });
     await prior;
     try {
-      await workflowMessages.synchronize(pi, client, ctx);
+      const finishRecording = async (message: WorkflowMessage): Promise<void> => {
+        const recorder = sessionRecorders.get(message.runId);
+        if (recorder === undefined) return;
+        await recorder.record(ctx).catch((error) => {
+          ctx.ui.notify(
+            `Workflow conversation recording failed: ${errorMessage(error)}`,
+            "warning",
+          );
+        });
+        await recorder.finish();
+        sessionRecorders.delete(message.runId);
+        if (activeRecorder === recorder) activeRecorder = null;
+      };
+      await workflowMessages.synchronize(pi, client, ctx, {
+        beforeTurnEnd: async (message, end) => {
+          if (end.stopReason === "completed") {
+            await submitVisibleAssistantResponse(client, ctx, message, end.responseSessionEntryId);
+          }
+          if (message.kind === "followUp") await finishRecording(message);
+        },
+        terminalDelivered: finishRecording,
+      });
       await activateRecorder(ctx);
     } finally {
       sessionView.refresh(ctx);
@@ -495,7 +501,7 @@ export default function piWorkflows(pi: ExtensionAPI): void {
       "List, start, restart, inspect, change settings, queue or remove follow-ups, pause, resume, cancel, answer ordinary checkpoints, update, or complete hosted workflow runs.",
       "Protected human decisions cannot be answered with this model-facing tool.",
       "When the user asks to continue or resume the active workflow, call workflow resume immediately.",
-      "Use update or submit only when a workflow step contract asks for it, and pass the exact step and attempt ids.",
+      "Use update or submit only when a workflow step contract asks for it, and pass its exact requestId.",
       "Do not start repeated work without the user's request.",
     ].join(" "),
     parameters: WorkflowToolParameters,
@@ -503,27 +509,16 @@ export default function piWorkflows(pi: ExtensionAPI): void {
       return await runToolInOrder(async () => {
         const params = parseWorkflowToolInput(rawParams);
         if (params.action === "update" || params.action === "submit") {
-          const interaction = pendingInteractionForSession(ctx.sessionManager.getSessionId());
-          if (interaction === undefined) throw new Error("No workflow step is waiting for output");
-          const contract = agentContract(interaction);
-          if (contract === undefined)
-            throw new Error("The pending interaction is not an agent step");
-          if (contract.nodeId !== params.step || contract.attemptId !== params.attempt) {
-            throw new Error("Workflow step or attempt does not match the durable request");
-          }
           let response: ClientResponse;
           try {
             response = await requestAccepted(client, {
               operation: params.action === "update" ? "interaction.update" : "interaction.submit",
               requestId: `${params.action}-${toolCallId}-${randomUUID()}`,
               idempotencyKey: toolCallId,
-              runId: interaction.runId,
-              expectedRevision: interaction.revision,
               payload: jsonValue({
-                requestId: interaction.requestId,
+                ...sessionCommandPayload(ctx),
+                requestId: params.requestId,
                 submissionId: toolCallId,
-                step: params.step,
-                attempt: params.attempt,
                 value:
                   params.action === "update"
                     ? { update: params.update }
@@ -692,40 +687,17 @@ export default function piWorkflows(pi: ExtensionAPI): void {
 
   pi.on("agent_start", async (_event, ctx) => {
     agentRunning = true;
+    lastStopReason = "completed";
     activeRecorder = null;
     activeRecorderMessageId = null;
     workflowMessages.startTurn();
     await presentInOrder(ctx).catch(() => undefined);
   });
 
-  pi.on("agent_end", async (event, ctx) => {
-    try {
-      await submitVisibleAssistantResponse(client, ctx);
-    } catch (error) {
-      ctx.ui.notify(`Workflow response was rejected: ${errorMessage(error)}`, "error");
-    }
-    const finishedMessage = workflowMessages.activeTurnMessage();
+  pi.on("agent_end", (event, ctx) => {
     agentRunning = false;
-    activeRecorderMessageId = null;
-    workflowMessages.endTurn(
-      workflowTurnStopReason(event.messages, ctx.signal?.aborted === true),
-      responseEntryId(ctx.sessionManager.getBranch()),
-    );
-    if (
-      activeRecorder !== null &&
-      finishedMessage !== undefined &&
-      (finishedMessage.kind === "terminal" || finishedMessage.kind === "followUp")
-    ) {
-      const finishedRecorder = activeRecorder;
-      activeRecorder = null;
-      await finishedRecorder.finish();
-      if (sessionRecorders.get(finishedMessage.runId) === finishedRecorder) {
-        sessionRecorders.delete(finishedMessage.runId);
-      }
-    }
-    await presentInOrder(ctx).catch((error) => {
-      ctx.ui.notify(`Could not record workflow model activity: ${errorMessage(error)}`, "warning");
-    });
+    lastStopReason = workflowTurnStopReason(event.messages, ctx.signal?.aborted === true);
+    // Pi can retry after agent_end. Only agent_settled completes the owned turn.
   });
 
   pi.on("turn_start", (event) => {
@@ -762,12 +734,11 @@ export default function piWorkflows(pi: ExtensionAPI): void {
 
   pi.on("agent_settled", async (_event, ctx) => {
     activeRecorder?.settleAttempt();
-    try {
-      await submitVisibleAssistantResponse(client, ctx);
-    } catch (error) {
-      ctx.ui.notify(`Workflow response was rejected: ${errorMessage(error)}`, "error");
-    }
-    await presentInOrder(ctx).catch(() => undefined);
+    workflowMessages.endTurn(lastStopReason, responseEntryId(ctx.sessionManager.getBranch()));
+    activeRecorderMessageId = null;
+    await presentInOrder(ctx).catch((error) => {
+      ctx.ui.notify(`Could not record workflow model activity: ${errorMessage(error)}`, "warning");
+    });
   });
 
   pi.on("session_shutdown", async (_event, ctx) => {
@@ -920,17 +891,16 @@ async function executeCommand(
       if (runId === undefined)
         throw new Error("No workflow terminal result is available to restart");
       const session = sessionCommandPayload(ctx);
-      const marker = terminalMessageForSession(ctx, runId);
+      const expectedRevision =
+        command.expectedRevision ?? (await client.getRun(runId))?.runRevision;
+      if (expectedRevision === undefined) throw new Error("Workflow run revision is unavailable");
       const response = await requestAccepted(client, {
         operation: "run.restart",
         requestId: `restart-${idempotencyKey}`,
         idempotencyKey,
         runId,
-        payload: {
-          ...session,
-          workflowMessageId: marker.workflowMessageId,
-          ...(marker.workflowTurnId === undefined ? {} : { workflowTurnId: marker.workflowTurnId }),
-        },
+        expectedRevision,
+        payload: session,
       });
       return {
         message: `Restarted workflow ${runId}.`,
@@ -1033,11 +1003,30 @@ async function executeCommand(
       };
     }
     case "answer": {
-      const interaction = pendingDecision(ctx.sessionManager.getSessionId(), command.runId);
-      if (interaction !== undefined) {
-        if (authority !== "human") {
-          throw new Error("Protected human decisions cannot be answered by the workflow tool");
-        }
+      if (authority === "model") {
+        const response = await requestAccepted(client, {
+          operation: "checkpoint.answer",
+          requestId: `checkpoint-${idempotencyKey}`,
+          idempotencyKey,
+          payload: {
+            ...sessionCommandPayload(ctx),
+            requestId: command.requestId,
+            submissionId: idempotencyKey,
+            input: jsonValue(command.input),
+          },
+        });
+        return {
+          message: `Answered checkpoint ${command.requestId}.`,
+          details: { action: "answer", response: response.receipt ?? null },
+        };
+      }
+      const interaction = sessionSnapshots
+        .get(ctx.sessionManager.getSessionId())
+        ?.pendingInteractions.map(parseInteractiveRequest)
+        .find((request) => request?.requestId === command.requestId);
+      if (interaction === undefined)
+        throw new Error("No matching checkpoint request is waiting in this session");
+      if (interaction.kind === "decision") {
         const response = await requestAccepted(client, {
           operation: "decision.answer",
           requestId: `answer-${idempotencyKey}`,
@@ -1045,6 +1034,7 @@ async function executeCommand(
           runId: interaction.runId,
           expectedRevision: interaction.revision,
           payload: {
+            ...sessionCommandPayload(ctx),
             requestId: interaction.requestId,
             submissionId: idempotencyKey,
             response: decisionResponse(command.input),
@@ -1059,33 +1049,26 @@ async function executeCommand(
           },
         };
       }
-      const parent = sessionRun(ctx, command.runId);
-      if (parent === undefined) throw new Error("No checkpoint is waiting in this session");
-      const continuationRunId = createRunId(parent.workflowName);
+      if (interaction.kind !== "checkpoint")
+        throw new Error("Only an ordinary checkpoint accepts an answer");
       const response = await requestAccepted(client, {
         operation: "checkpoint.answer",
         requestId: `checkpoint-${idempotencyKey}`,
         idempotencyKey,
-        runId: parent.runId,
+        runId: interaction.runId,
+        expectedRevision: interaction.revision,
         payload: {
-          continuationRunId,
+          ...sessionCommandPayload(ctx),
+          requestId: interaction.requestId,
+          submissionId: idempotencyKey,
           input: jsonValue(command.input),
         },
       });
-      const receipt = isRecord(response.receipt) ? response.receipt : undefined;
-      const actualRunId =
-        receipt !== undefined && typeof receipt.runId === "string"
-          ? receipt.runId
-          : continuationRunId;
       return {
-        message:
-          receipt?.alreadyAnswered === true
-            ? `Checkpoint ${parent.runId} was already answered; continuation ${actualRunId} exists.`
-            : `Answered checkpoint ${parent.runId}; continuation ${actualRunId} started.`,
+        message: `Answered checkpoint ${interaction.requestId}; run ${interaction.runId} can continue.`,
         details: {
           action: "answer",
-          parentRunId: parent.runId,
-          runId: actualRunId,
+          runId: interaction.runId,
           response: response.receipt ?? null,
         },
       };
@@ -1162,54 +1145,33 @@ async function executeResourceManagerCommand(
 async function submitVisibleAssistantResponse(
   client: WorkflowClient,
   ctx: ExtensionContext,
+  message: WorkflowMessage,
+  settledResponseEntryId: string | null,
 ): Promise<void> {
-  const interaction = pendingInteractionForSession(ctx.sessionManager.getSessionId());
-  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);
+  const contract = agentContractForWorkflowMessage(message);
+  if (contract?.completion !== "assistant" || workflowRunPaused(message.runId)) return;
+  if (settledResponseEntryId === null) return;
+  const branch = ctx.sessionManager.getBranch();
+  const responseIndex = branch.findIndex((entry) => entry.id === settledResponseEntryId);
+  if (responseIndex < 0) return;
+  const submission = recoverAssistantStep(branch.slice(0, responseIndex + 1), contract);
   if (submission === undefined) return;
   const responseId = submission.conversation?.lastEntryId ?? submission.assistantMessage?.entryId;
-  if (responseId === undefined) return;
-  await requestAccepted(client, {
-    operation: "interaction.submit",
-    requestId: `assistant-${interaction.requestId}-${responseId}`,
-    idempotencyKey: `assistant-${interaction.requestId}-${responseId}`,
-    runId: interaction.runId,
-    expectedRevision: interaction.revision,
+  if (responseId !== settledResponseEntryId) return;
+  const response = await client.request({
+    operation: "interaction.assistant",
+    requestId: `assistant-${message.sourceId}-${responseId}`,
+    idempotencyKey: `assistant-${message.sourceId}-${responseId}`,
     payload: {
-      requestId: interaction.requestId,
+      ...sessionCommandPayload(ctx),
+      requestId: message.sourceId,
       submissionId: `assistant-${responseId}`,
-      step: contract.nodeId,
-      attempt: contract.attemptId,
       value: submission as unknown as JsonValue,
     },
   });
-}
-
-function pendingInteractionForSession(sessionId: string): ClientInteractiveRequest | undefined {
-  return sessionSnapshots
-    .get(sessionId)
-    ?.pendingInteractions.map(parseInteractiveRequest)
-    .find((request): request is ClientInteractiveRequest => request !== undefined);
-}
-
-function pendingDecision(sessionId: string, runId?: string): ClientInteractiveRequest | undefined {
-  return sessionSnapshots
-    .get(sessionId)
-    ?.pendingInteractions.map(parseInteractiveRequest)
-    .find(
-      (request): request is ClientInteractiveRequest =>
-        request !== undefined &&
-        request.kind === "decision" &&
-        (runId === undefined || request.runId === runId),
-    );
+  if (response.outcome !== "accepted" && response.outcome !== "adopted") {
+    ctx.ui.notify(`Workflow response was rejected: ${response.error ?? response.outcome}`, "error");
+  }
 }
 
 function workflowRunPaused(runId: string): boolean {
@@ -1240,25 +1202,6 @@ function sessionCommandPayload(ctx: ExtensionContext): {
   return { targetSessionId, coordinatorEpoch: session.coordinatorEpoch };
 }
 
-function terminalMessageForSession(
-  ctx: ExtensionContext,
-  runId: string,
-): { workflowMessageId: string; workflowTurnId?: string } {
-  const session = sessionSnapshots.get(ctx.sessionManager.getSessionId());
-  const message = session?.workflowMessages
-    .filter((item) => item.kind === "terminal" && item.runId === runId && item.status === "sent")
-    .sort((left, right) => right.order - left.order)[0];
-  if (message === undefined)
-    throw new Error(`Workflow run ${runId} has no current terminal result`);
-  const turn = session?.openWorkflowTurn;
-  return {
-    workflowMessageId: message.workflowMessageId,
-    ...(turn?.workflowMessageId === message.workflowMessageId
-      ? { workflowTurnId: turn.workflowTurnId }
-      : {}),
-  };
-}
-
 function sessionRun(ctx: ExtensionContext, runId?: string): WorkflowRunQueueView | undefined {
   const session = sessionSnapshots.get(ctx.sessionManager.getSessionId());
   if (session?.run === null || session?.run === undefined || !isRecord(session.run.queue)) {
@@ -1288,7 +1231,7 @@ function toolInputToCommand(params: ReturnType<typeof parseWorkflowToolInput>):
     case "start":
       return { kind: "run", ref: params.workflow, input: params.input ?? {} };
     case "restart":
-      return { kind: "restart", runId: params.runId };
+      return { kind: "restart", runId: params.runId, expectedRevision: params.expectedRevision };
     case "change-settings":
       return {
         kind: "change-settings",
@@ -1324,8 +1267,8 @@ function toolInputToCommand(params: ReturnType<typeof parseWorkflowToolInput>):
     case "answer":
       return {
         kind: "answer",
-        input: params.input ?? {},
-        ...(params.runId === undefined ? {} : { runId: params.runId }),
+        requestId: params.requestId,
+        input: params.input,
       };
     case "update":
     case "submit":
@@ -1348,31 +1291,13 @@ async function requestAccepted(
   return response;
 }
 
-function interactionContract(interaction: ClientInteractiveRequest): Record<string, unknown> {
-  if (!isRecord(interaction.contract)) throw new Error("Stored interaction contract is invalid");
-  return interaction.contract;
-}
-
-function agentContract(interaction: ClientInteractiveRequest): AgentStepContract | undefined {
-  const value = interactionContract(interaction).contract;
-  if (
-    !isRecord(value) ||
-    typeof value.runId !== "string" ||
-    typeof value.workflowName !== "string" ||
-    typeof value.nodeId !== "string" ||
-    typeof value.attemptId !== "string" ||
-    (value.completion !== "submit" && value.completion !== "assistant")
-  ) {
-    return undefined;
-  }
-  return value as unknown as AgentStepContract;
-}
-
 function agentContractForWorkflowMessage(message: WorkflowMessage): AgentStepContract | undefined {
   if (message.kind !== "step" || !isRecord(message.content.details)) return undefined;
   const value = message.content.details.contract;
   if (
     !isRecord(value) ||
+    value.requestId !== message.sourceId ||
+    value.runId !== message.runId ||
     typeof value.runId !== "string" ||
     typeof value.workflowName !== "string" ||
     typeof value.nodeId !== "string" ||
@@ -1484,7 +1409,10 @@ function parseInteractiveRequest(value: unknown): ClientInteractiveRequest | und
     typeof value.runId !== "string" ||
     typeof value.targetSessionId !== "string" ||
     typeof value.revision !== "number" ||
-    (value.kind !== "agent" && value.kind !== "assistant" && value.kind !== "decision")
+    (value.kind !== "agent" &&
+      value.kind !== "assistant" &&
+      value.kind !== "checkpoint" &&
+      value.kind !== "decision")
   ) {
     return undefined;
   }

package/src/extension/recorder.ts

@@ -190,8 +190,8 @@ export class SessionRecorder {
   }
 
   /** Record a workflow-owned turn that is not tied to one engine node attempt. */
-  beginWorkflowMessage(workflowMessageId: string, kind: "terminal" | "followUp"): void {
-    this.currentAttempt = { nodeId: kind, attemptId: workflowMessageId };
+  beginFollowUp(workflowMessageId: string): void {
+    this.currentAttempt = { nodeId: "followUp", attemptId: workflowMessageId };
   }
 
   /** Release ownership of the turn that Pi has fully settled. */

package/src/extension/step-message.ts

@@ -19,7 +19,7 @@ import {
 export const WORKFLOW_AGENT_STEP_MESSAGE_TYPE = WORKFLOW_STEP_MESSAGE_TYPE;
 export const WORKFLOW_AGENT_STEP_MESSAGE_SCHEMA = "pi-workflows.agent-step-message.v1";
 
-type PromptDeliveryReason = "initial" | "reminder" | "resumed";
+type PromptDeliveryReason = "initial" | "resumed";
 
 export type WorkflowAgentStepMessageDetails = {
   schema: typeof WORKFLOW_AGENT_STEP_MESSAGE_SCHEMA;
@@ -152,6 +152,7 @@ export function recoverAssistantStep(
 
 function sameAttempt(left: AgentStepContract, right: AgentStepContract): boolean {
   return (
+    left.requestId === right.requestId &&
     left.runId === right.runId &&
     left.workflowName === right.workflowName &&
     left.nodeId === right.nodeId &&
@@ -225,9 +226,7 @@ function parseDetails(value: unknown): WorkflowAgentStepMessageDetails | undefin
   if (candidate.schema !== WORKFLOW_AGENT_STEP_MESSAGE_SCHEMA) return undefined;
   if (
     candidate.kind !== "step" ||
-    (candidate.reason !== "initial" &&
-      candidate.reason !== "reminder" &&
-      candidate.reason !== "resumed") ||
+    (candidate.reason !== "initial" && candidate.reason !== "resumed") ||
     typeof candidate.requestId !== "string" ||
     typeof candidate.workflowMessageId !== "string"
   ) {
@@ -237,6 +236,7 @@ function parseDetails(value: unknown): WorkflowAgentStepMessageDetails | undefin
   if (
     contract === null ||
     typeof contract !== "object" ||
+    typeof contract.requestId !== "string" ||
     typeof contract.runId !== "string" ||
     typeof contract.workflowName !== "string" ||
     typeof contract.nodeId !== "string" ||