@osolmaz/pi-workflows 0.1.0 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@osolmaz/pi-workflows",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "description": "Workflow engine, JSON control-flow tool, and live terminal viewer for the pi coding agent",
   "keywords": [
     "pi-package"
@@ -42,6 +42,7 @@
   "scripts": {
     "build": "tsc -p tsconfig.build.json",
     "check": "npm run format:check && npm run lint && npm run typecheck && npm run build && npm run test:coverage",
+    "fixtures": "node scripts/export-layout-fixtures.mjs",
     "format": "oxfmt --write",
     "format:check": "oxfmt --check",
     "lint": "oxlint --deny-warnings src test examples",

package/src/extension/executor.ts

@@ -2,6 +2,7 @@ import type {
   AgentStepExecutor,
   AgentStepRequest,
   AgentStepSubmission,
+  ConversationRange,
 } from "../workflows/types.js";
 
 export type SubmissionResult =
@@ -14,11 +15,24 @@ export type PromptDelivery = {
   streaming: boolean;
 };
 
+/**
+ * Bracketing hooks for conversation linkage: a mark is taken when the step
+ * prompt is first delivered, and the recorded entry range since that mark is
+ * attached to the accepted submission.
+ */
+export type ConversationHooks = {
+  beginAttempt?: (contract: AgentStepRequest["contract"]) => void;
+  mark: () => number;
+  rangeSince: (mark: number) => ConversationRange | undefined;
+};
+
 export type ConversationStepExecutorOptions = {
   /** Deliver a prompt into the pi conversation. */
   sendPrompt: (delivery: PromptDelivery) => void;
   /** Reminders sent when the agent settles without submitting. Default 2. */
   maxNudges?: number;
+  /** Conversation linkage hooks, wired to the session recorder. */
+  conversation?: ConversationHooks;
 };
 
 type PendingStep = {
@@ -26,6 +40,8 @@ type PendingStep = {
   resolve: (submission: AgentStepSubmission) => void;
   reject: (error: unknown) => void;
   nudgesSent: number;
+  /** Conversation mark taken when the prompt was first delivered. */
+  mark: number | null;
   cleanup: () => void;
   /** Resolves when this step stops being the pending step. */
   cleared: Promise<void>;
@@ -44,6 +60,7 @@ const DEFAULT_MAX_NUDGES = 2;
 export class ConversationStepExecutor implements AgentStepExecutor {
   private readonly sendPrompt: (delivery: PromptDelivery) => void;
   private readonly maxNudges: number;
+  private readonly conversation: ConversationHooks | undefined;
   private pending: PendingStep | null = null;
   private streaming = false;
   private heldByUser = false;
@@ -51,6 +68,7 @@ export class ConversationStepExecutor implements AgentStepExecutor {
   constructor(options: ConversationStepExecutorOptions) {
     this.sendPrompt = options.sendPrompt;
     this.maxNudges = options.maxNudges ?? DEFAULT_MAX_NUDGES;
+    this.conversation = options.conversation;
   }
 
   /** Track agent streaming state (wire to agent_start / agent_settled). */
@@ -88,6 +106,7 @@ export class ConversationStepExecutor implements AgentStepExecutor {
       return;
     }
     try {
+      this.conversation?.beginAttempt?.(pending.request.contract);
       this.sendPrompt({ prompt: pending.request.prompt, streaming: this.streaming });
     } catch (error) {
       this.clearPending();
@@ -110,11 +129,13 @@ export class ConversationStepExecutor implements AgentStepExecutor {
       const cleared = new Promise<void>((resolveCleared) => {
         markCleared = resolveCleared;
       });
+      this.conversation?.beginAttempt?.(request.contract);
       this.pending = {
         request,
         resolve,
         reject,
         nudgesSent: 0,
+        mark: this.conversation?.mark() ?? null,
         cleanup: () => signal.removeEventListener("abort", onAbort),
         cleared,
         markCleared,
@@ -185,7 +206,12 @@ export class ConversationStepExecutor implements AgentStepExecutor {
       };
     }
     this.clearPending();
-    pending.resolve({ output: result.value });
+    const conversation =
+      pending.mark !== null ? this.conversation?.rangeSince(pending.mark) : undefined;
+    pending.resolve({
+      output: result.value,
+      ...(conversation !== undefined ? { conversation } : {}),
+    });
     return {
       accepted: true,
       message: [
@@ -224,6 +250,7 @@ export class ConversationStepExecutor implements AgentStepExecutor {
     pending.nudgesSent += 1;
     const { nodeId, attemptId } = pending.request.contract;
     try {
+      this.conversation?.beginAttempt?.(pending.request.contract);
       this.sendPrompt({
         prompt: [
           `Reminder: workflow step ${JSON.stringify(nodeId)} is still awaiting your output.`,

package/src/extension/index.ts

@@ -7,7 +7,7 @@ import { Type } from "typebox";
 import { WorkflowEngine } from "../workflows/engine.js";
 import { errorMessage } from "../workflows/errors.js";
 import { discoverWorkflows, loadWorkflowFile, resolveWorkflowRef } from "../workflows/loader.js";
-import { createDefinitionSnapshot } from "../workflows/store.js";
+import { WorkflowRunStore, createDefinitionSnapshot } from "../workflows/store.js";
 import type {
   WorkflowDefinition,
   WorkflowDefinitionSnapshot,
@@ -15,6 +15,7 @@ import type {
   WorkflowRunState,
 } from "../workflows/types.js";
 import { ConversationStepExecutor } from "./executor.js";
+import { SessionRecorder } from "./recorder.js";
 import { buildWidgetView } from "./widget.js";
 
 const WIDGET_KEY = "pi-workflows";
@@ -37,6 +38,7 @@ type ActiveRun = {
   workflowName: string;
   engine: WorkflowEngine;
   executor: ConversationStepExecutor;
+  recorder: SessionRecorder | null;
   snapshot: WorkflowDefinitionSnapshot;
   presentationPrompt: WorkflowDefinition["presentationPrompt"];
   generation: number;
@@ -294,6 +296,9 @@ export default function piWorkflows(pi: ExtensionAPI) {
     if (activeRun === run) {
       activeRun = null;
     }
+    // Normally already stopped via onRunFinishing; this covers observers of
+    // runs that ended without reaching that hook.
+    void run.recorder?.stop();
     stopWidgetTicker();
     const { state } = result;
     updateWidget(ctx, state, run.snapshot);
@@ -335,9 +340,36 @@ export default function piWorkflows(pi: ExtensionAPI) {
       sendPrompt: ({ prompt, streaming }) => {
         pi.sendUserMessage(prompt, streaming ? { deliverAs: "steer" } : undefined);
       },
+      conversation: {
+        beginAttempt: (contract) => run.recorder?.beginAttempt(contract),
+        mark: () => run.recorder?.mark() ?? 0,
+        rangeSince: (mark) => run.recorder?.rangeSince(mark),
+      },
     });
+    // The store is shared between the engine and the session recorder so the
+    // trace sequence stays single-writer (see docs/run-bundles.md).
+    const store = new WorkflowRunStore();
     const engine = new WorkflowEngine({
       executor,
+      store,
+      // Awaited by the engine after run_started is persisted, so the session
+      // binding and its trace event always precede node and terminal events.
+      onRunStarted: async (runDir, state) => {
+        const recorder = new SessionRecorder(store, runDir, state.runId);
+        try {
+          await recorder.bind(ctx);
+          run.recorder = recorder;
+        } catch {
+          // Binding is best-effort: a session without UI access or an
+          // ephemeral context must not fail the run.
+        }
+      },
+      // Awaited by the engine before the terminal snapshot. If completion was
+      // submitted from the workflow tool, capture stays open through Pi's
+      // final tool, message, and turn hooks before it drains.
+      onRunFinishing: async () => {
+        await run.recorder?.finish();
+      },
       onEvent: (_event, state: WorkflowRunState) => {
         if (run.runId === null) {
           run.runId = state.runId;
@@ -351,6 +383,7 @@ export default function piWorkflows(pi: ExtensionAPI) {
       workflowName: workflow.name,
       engine,
       executor,
+      recorder: null,
       snapshot,
       presentationPrompt: workflow.presentationPrompt,
       generation,
@@ -368,6 +401,7 @@ export default function piWorkflows(pi: ExtensionAPI) {
         if (activeRun === run) {
           activeRun = null;
         }
+        void run.recorder?.stop();
         stopWidgetTicker();
         clearWidget(ctx);
         notify(ctx, `Workflow ${workflow.name} crashed: ${errorMessage(error)}`, "error");
@@ -488,12 +522,21 @@ export default function piWorkflows(pi: ExtensionAPI) {
       attempt: Type.String({ description: "The attempt id from the workflow step contract" }),
       output: Type.Unknown({ description: "The step output, matching the expected output shape" }),
     }),
-    async execute(_toolCallId, params) {
+    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
       if (!activeRun) {
         throw new Error(
           "No workflow is running. Do not call the workflow tool outside a workflow.",
         );
       }
+      // Flush the conversation into the bundle before accepting, so the
+      // attempt's recorded range includes the assistant message that carries
+      // this submission. Pi guarantees ctx.sessionManager is synchronized
+      // through the current assistant tool-calling message before tool_call
+      // dispatch, which precedes tool execution (docs/extensions.md, "Tool
+      // Events"). The tool result of this call itself lands after the range
+      // by design: it is the submission receipt, not the submission.
+      await activeRun.recorder?.record(ctx).catch(() => undefined);
+      await activeRun.recorder?.synchronize(ctx).catch(() => undefined);
       const result = await activeRun.executor.submit(params.step, params.attempt, params.output);
       if (!result.accepted) {
         throw new Error(result.message);
@@ -552,19 +595,56 @@ export default function piWorkflows(pi: ExtensionAPI) {
     );
   });
 
-  pi.on("agent_settled", () => {
-    if (!activeRun) {
+  pi.on("turn_start", (event) => {
+    activeRun?.recorder?.handleTurnStart(event);
+  });
+
+  pi.on("turn_end", async (event, ctx) => {
+    await activeRun?.recorder?.handleTurnEnd(event, ctx).catch(() => undefined);
+  });
+
+  pi.on("message_start", async (event, ctx) => {
+    await activeRun?.recorder?.handleMessageStart(event, ctx).catch(() => undefined);
+  });
+
+  pi.on("message_update", (event) => {
+    activeRun?.recorder?.handleMessageUpdate(event);
+  });
+
+  pi.on("message_end", (event) => {
+    activeRun?.recorder?.handleMessageEnd(event);
+  });
+
+  pi.on("tool_execution_start", (event) => {
+    activeRun?.recorder?.handleToolStart(event);
+  });
+
+  pi.on("tool_execution_update", (event) => {
+    activeRun?.recorder?.handleToolUpdate(event);
+  });
+
+  pi.on("tool_execution_end", (event) => {
+    activeRun?.recorder?.handleToolEnd(event);
+  });
+
+  pi.on("agent_settled", async (_event, ctx) => {
+    const run = activeRun;
+    if (!run) {
       presentationPending = null;
       return;
     }
-    activeRun.executor.setStreaming(false);
-    activeRun.executor.handleAgentSettled();
+    await run.recorder?.synchronize(ctx).catch(() => undefined);
+    run.recorder?.settleAttempt();
+    run.executor.setStreaming(false);
+    run.executor.handleAgentSettled();
   });
 
   pi.on("session_shutdown", () => {
     sessionClosed = true;
     supersedePresentation();
-    activeRun?.engine.cancel();
+    const run = activeRun;
+    run?.engine.cancel();
+    void run?.recorder?.stop();
     activeRun = null;
     presentationPending = null;
     clearWidgetTimer();