@osolmaz/pi-workflows 0.2.0 → 0.3.0

package/src/extension/recorder.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import {
   SESSION_BINDING_SCHEMA,
@@ -83,6 +84,12 @@ export class SessionRecorder {
   private readonly store: WorkflowRunStore;
   private readonly runDir: string;
   private readonly runId: string;
+  /**
+   * Capture segment for this recorder. Undefined writes the legacy flat
+   * stream; a second recorder attaching to the same bundle (resume or
+   * handoff) writes a segment instead so the first capture stays intact.
+   */
+  private segmentId: string | undefined;
   private cursor: string | null = null;
   private readonly recorded: string[] = [];
   private readonly unclaimedEntries: RecordedEntry[] = [];
@@ -128,22 +135,35 @@ export class SessionRecorder {
     this.bound = true;
     this.cursor = ctx.sessionManager.getLeafId();
     const sessionFile = ctx.sessionManager.getSessionFile();
-    await this.store.writeSessionBinding(this.runDir, {
-      schema: SESSION_BINDING_SCHEMA,
-      runId: this.runId,
-      piSessionId: ctx.sessionManager.getSessionId(),
-      ...(sessionFile !== undefined ? { piSessionFile: sessionFile } : {}),
-      cwd: ctx.cwd,
-      boundAt: new Date().toISOString(),
-    });
-    await this.store.writeSessionCapture(this.runDir, {
-      schema: SESSION_CAPTURE_SCHEMA,
-      eventSchema: SESSION_EVENT_SCHEMA,
-      status: "recording",
-      eventCount: 0,
-      entryCount: 0,
-      lastEventSeq: 0,
-    });
+    // A bundle that already has a binding belongs to an earlier capture
+    // attempt; this recorder writes its own segment under session/segments.
+    if (await this.store.hasSessionBinding(this.runDir)) {
+      this.segmentId = randomUUID();
+    }
+    await this.store.writeSessionBinding(
+      this.runDir,
+      {
+        schema: SESSION_BINDING_SCHEMA,
+        runId: this.runId,
+        piSessionId: ctx.sessionManager.getSessionId(),
+        ...(sessionFile !== undefined ? { piSessionFile: sessionFile } : {}),
+        cwd: ctx.cwd,
+        boundAt: new Date().toISOString(),
+      },
+      this.segmentId,
+    );
+    await this.store.writeSessionCapture(
+      this.runDir,
+      {
+        schema: SESSION_CAPTURE_SCHEMA,
+        eventSchema: SESSION_EVENT_SCHEMA,
+        status: "recording",
+        eventCount: 0,
+        entryCount: 0,
+        lastEventSeq: 0,
+      },
+      this.segmentId,
+    );
   }
 
   /** Fix the owner before the executor delivers an agent-step prompt. */
@@ -324,7 +344,7 @@ export class SessionRecorder {
       }
       const appended: RecordedEntry[] = [];
       for (const entry of branch.slice(startIndex)) {
-        await this.store.appendSessionEntry(this.runDir, entry);
+        await this.store.appendSessionEntry(this.runDir, entry, this.segmentId);
         const recorded = { id: entry.id, entry, claimed: false };
         appended.push(recorded);
         this.unclaimedEntries.push(recorded);
@@ -406,31 +426,39 @@ export class SessionRecorder {
       try {
         await this.flushAllEvents();
         await this.entryChain;
-        const counts = await this.store.sessionCounts(this.runDir);
-        await this.store.writeSessionCapture(this.runDir, {
-          schema: SESSION_CAPTURE_SCHEMA,
-          eventSchema: SESSION_EVENT_SCHEMA,
-          status: this.captureFailure ? "failed" : "complete",
-          ...counts,
-          ...(this.captureFailure ? { failure: this.captureFailure } : {}),
-        });
+        const counts = await this.store.sessionCounts(this.runDir, this.segmentId);
+        await this.store.writeSessionCapture(
+          this.runDir,
+          {
+            schema: SESSION_CAPTURE_SCHEMA,
+            eventSchema: SESSION_EVENT_SCHEMA,
+            status: this.captureFailure ? "failed" : "complete",
+            ...counts,
+            ...(this.captureFailure ? { failure: this.captureFailure } : {}),
+          },
+          this.segmentId,
+        );
       } catch (error) {
         // Capture is observational. A finalization failure must never reject
         // the workflow's terminal persistence hook.
         this.failCapture("capture_finalize_failed", failureMessage(error));
         try {
-          const counts = await this.store.sessionCounts(this.runDir);
-          await this.store.writeSessionCapture(this.runDir, {
-            schema: SESSION_CAPTURE_SCHEMA,
-            eventSchema: SESSION_EVENT_SCHEMA,
-            status: "failed",
-            ...counts,
-            failure: this.captureFailure ?? {
-              failedAt: new Date().toISOString(),
-              code: "capture_finalize_failed",
-              message: failureMessage(error),
+          const counts = await this.store.sessionCounts(this.runDir, this.segmentId);
+          await this.store.writeSessionCapture(
+            this.runDir,
+            {
+              schema: SESSION_CAPTURE_SCHEMA,
+              eventSchema: SESSION_EVENT_SCHEMA,
+              status: "failed",
+              ...counts,
+              failure: this.captureFailure ?? {
+                failedAt: new Date().toISOString(),
+                code: "capture_finalize_failed",
+                message: failureMessage(error),
+              },
             },
-          });
+            this.segmentId,
+          );
         } catch {
           // The viewer will report the missing/invalid capture file.
         }
@@ -576,6 +604,7 @@ export class SessionRecorder {
       .appendSessionEventBatch(
         this.runDir,
         batch.map((queued) => queued.record),
+        this.segmentId,
       )
       .catch((error: unknown) => {
         const message = failureMessage(error);

package/src/extension/workflow-tool.ts

@@ -0,0 +1,59 @@
+import { StringEnum } from "@earendil-works/pi-ai";
+import { Type, type TSchema } from "typebox";
+
+const noExtraProperties = { additionalProperties: false } as const;
+
+export const WorkflowToolParameters: TSchema = Type.Union([
+  Type.Object(
+    {
+      action: StringEnum(["list"] as const),
+      offset: Type.Optional(Type.Integer({ minimum: 0, description: "Workflow list offset" })),
+    },
+    noExtraProperties,
+  ),
+  Type.Object(
+    {
+      action: StringEnum(["start"] as const),
+      workflow: Type.String({ description: "Discovered workflow name or workflow file path" }),
+      input: Type.Optional(Type.Unknown({ description: "Structured workflow input" })),
+    },
+    noExtraProperties,
+  ),
+  Type.Object(
+    {
+      action: StringEnum(["status"] as const),
+      runId: Type.Optional(Type.String({ description: "Run id; omit for the active run" })),
+    },
+    noExtraProperties,
+  ),
+  Type.Object({ action: StringEnum(["pause"] as const) }, noExtraProperties),
+  Type.Object({ action: StringEnum(["resume"] as const) }, noExtraProperties),
+  Type.Object({ action: StringEnum(["cancel"] as const) }, noExtraProperties),
+  Type.Object(
+    {
+      action: StringEnum(["answer"] as const),
+      input: Type.Unknown({ description: "Checkpoint answer" }),
+      runId: Type.Optional(Type.String({ description: "Waiting run id" })),
+    },
+    noExtraProperties,
+  ),
+  Type.Object(
+    {
+      action: StringEnum(["submit"] as const),
+      step: Type.String({ description: "Step id from the workflow step contract" }),
+      attempt: Type.String({ description: "Attempt id from the workflow step contract" }),
+      output: Type.Unknown({ description: "Step output matching the expected shape" }),
+    },
+    noExtraProperties,
+  ),
+]);
+
+export type WorkflowToolInput =
+  | { action: "list"; offset?: number }
+  | { action: "start"; workflow: string; input?: unknown }
+  | { action: "status"; runId?: string }
+  | { action: "pause" }
+  | { action: "resume" }
+  | { action: "cancel" }
+  | { action: "answer"; input: unknown; runId?: string }
+  | { action: "submit"; step: string; attempt: string; output: unknown };

package/src/host/processes.ts

@@ -0,0 +1,119 @@
+import fs from "node:fs";
+import path from "node:path";
+
+/**
+ * Tracks headless child processes (spawned `pi --mode rpc` sessions) so a
+ * killed host never leaves orphans working. Children spawn in their own
+ * process group; the registry file lets a later host reap leftovers by
+ * killing the whole group. The file lives next to the project store and is
+ * only ever touched by one host at a time (the advisory lock ensures it).
+ */
+export class HostProcessRegistry {
+  private readonly filePath: string;
+  private readonly pids = new Set<number>();
+
+  constructor(storeDir: string) {
+    this.filePath = path.join(storeDir, "host.children.json");
+  }
+
+  register(pid: number): void {
+    this.pids.add(pid);
+    this.persist();
+  }
+
+  unregister(pid: number): void {
+    this.pids.delete(pid);
+    this.persist();
+  }
+
+  get size(): number {
+    return this.pids.size;
+  }
+
+  /** Kill every registered child's process group, escalating to SIGKILL. */
+  killAll(): void {
+    for (const pid of this.pids) {
+      killProcessGroup(pid, "SIGTERM");
+    }
+    for (const pid of this.pids) {
+      killProcessGroup(pid, "SIGKILL");
+    }
+    this.pids.clear();
+    this.persist();
+  }
+
+  /**
+   * Reap children recorded by a previous host that died without cleanup.
+   * Called once at host start, before the lock is taken.
+   */
+  reapOrphans(): number[] {
+    const recorded = this.readFile();
+    const reaped: number[] = [];
+    const stillAlive: number[] = [];
+    for (const pid of recorded) {
+      if (isAlive(pid)) {
+        killProcessGroup(pid, "SIGKILL");
+        reaped.push(pid);
+        if (isAlive(pid)) {
+          // A group kill that left the leader alive stays registered so a
+          // later start tries again instead of forgetting it.
+          stillAlive.push(pid);
+        }
+      }
+    }
+    this.pids.clear();
+    for (const pid of stillAlive) {
+      this.pids.add(pid);
+    }
+    this.persist();
+    return reaped;
+  }
+
+  private readFile(): number[] {
+    try {
+      const parsed = JSON.parse(fs.readFileSync(this.filePath, "utf8")) as unknown;
+      if (!Array.isArray(parsed)) {
+        return [];
+      }
+      return parsed.filter((pid): pid is number => Number.isSafeInteger(pid) && pid > 0);
+    } catch {
+      return [];
+    }
+  }
+
+  private persist(): void {
+    try {
+      fs.mkdirSync(path.dirname(this.filePath), { recursive: true, mode: 0o700 });
+      // Write exactly the live set. Unioning with the previous file would
+      // resurrect exited pids, and a later host reaping them could kill an
+      // unrelated process group that reused one.
+      fs.writeFileSync(this.filePath, `${JSON.stringify([...this.pids])}\n`, {
+        encoding: "utf8",
+        mode: 0o600,
+      });
+    } catch {
+      // Registry bookkeeping is best-effort; orphan reaping is the backstop.
+    }
+  }
+}
+
+function isAlive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function killProcessGroup(pid: number, signal: NodeJS.Signals): void {
+  try {
+    process.kill(-pid, signal);
+  } catch {
+    try {
+      process.kill(pid, signal);
+    } catch {
+      // Already gone.
+    }
+  }
+}

package/src/host/rpc-bridge.ts

@@ -0,0 +1,44 @@
+import { StringEnum } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+
+export const RPC_SUBMISSION_PREFIX = "PI_WORKFLOWS_STEP_SUBMISSION ";
+
+/**
+ * Loaded into headless `pi --mode rpc` children spawned by the standalone
+ * host. The child has no workflow engine, so this bridge only registers the
+ * `workflow` tool and reports every submission to the host over stderr; the
+ * host validates against the engine and re-prompts on rejection.
+ */
+export default function piWorkflowsRpcBridge(pi: ExtensionAPI) {
+  pi.registerTool({
+    name: "workflow",
+    label: "Workflow",
+    description: [
+      "Submit the output for the pending workflow step with action submit.",
+      "Only call this tool when a workflow step contract in the conversation asks you to.",
+      "Pass the exact step id from the contract and your result as the output.",
+    ].join(" "),
+    parameters: Type.Object(
+      {
+        action: StringEnum(["submit"] as const),
+        step: Type.String({ description: "Step id from the workflow step contract" }),
+        attempt: Type.String({ description: "Attempt id from the workflow step contract" }),
+        output: Type.Unknown({ description: "Step output matching the expected shape" }),
+      },
+      { additionalProperties: false },
+    ),
+    async execute(_toolCallId, params) {
+      process.stderr.write(`${RPC_SUBMISSION_PREFIX}${JSON.stringify(params)}\n`);
+      return {
+        content: [
+          {
+            type: "text",
+            text: "Submission recorded. Continue only when the workflow sends the next step.",
+          },
+        ],
+        details: {},
+      };
+    },
+  });
+}

package/src/host/rpc-executor.ts

@@ -0,0 +1,299 @@
+import { type ChildProcess, spawn } from "node:child_process";
+import fs from "node:fs";
+import { fileURLToPath } from "node:url";
+import type {
+  AgentStepExecutor,
+  AgentStepRequest,
+  AgentStepSubmission,
+} from "../workflows/types.js";
+import type { HostProcessRegistry } from "./processes.js";
+import { RPC_SUBMISSION_PREFIX } from "./rpc-bridge.js";
+
+// Production resolves the compiled .js; tests and dev checkouts load the .ts
+// source, which pi compiles itself when loading extensions.
+const BRIDGE_PATH = ["./rpc-bridge.js", "./rpc-bridge.ts"]
+  .map((candidate) => fileURLToPath(new URL(candidate, import.meta.url)))
+  .find((candidate) => fs.existsSync(candidate));
+const ABORT_GRACE_MS = 3_000;
+
+type StepSubmission = {
+  action: "submit";
+  step: string;
+  attempt: string;
+  output: unknown;
+};
+
+export type RpcStepExecutorOptions = {
+  cwd: string;
+  registry: HostProcessRegistry;
+  /** Absolute path to the pi binary; defaults to the installed `pi`. */
+  piBin?: string;
+  /** Extra environment for the child; the host's environment is inherited. */
+  env?: Record<string, string>;
+  /** Extra pi arguments, for example a model override. */
+  piArgs?: string[];
+};
+
+/**
+ * Runs agent steps in a headless `pi --mode rpc` child. One child serves one
+ * workflow run. The rpc-bridge extension inside the child registers the
+ * `workflow` tool and reports submissions over stderr; this executor
+ * validates them through `request.accept` and re-prompts on rejection, so
+ * the model sees the same tool contract as an in-session run.
+ */
+export class RpcStepExecutor implements AgentStepExecutor {
+  private readonly options: RpcStepExecutorOptions;
+  private child: ChildProcess | null = null;
+  private childExited: { code: number | null; signal: string | null } | null = null;
+  private stderrBuffer = "";
+  private submissions: StepSubmission[] = [];
+  private submissionWaiters: Array<() => void> = [];
+  private stdoutBuffer = "";
+
+  constructor(options: RpcStepExecutorOptions) {
+    this.options = options;
+  }
+
+  async runAgentStep(request: AgentStepRequest, signal: AbortSignal): Promise<AgentStepSubmission> {
+    this.ensureStarted();
+    let prompt = request.prompt;
+    for (;;) {
+      throwIfAborted(signal);
+      const submission = await this.promptForSubmission(prompt, request, signal);
+      const accepted = await request.accept(submission.output);
+      if (accepted.ok) {
+        return { output: accepted.value };
+      }
+      prompt = [
+        `The workflow rejected your submission for step ${JSON.stringify(request.contract.nodeId)}:`,
+        accepted.error,
+        "",
+        "Call the `workflow` tool again with a corrected output.",
+      ].join("\n");
+    }
+  }
+
+  /** Stop the child's whole process group, then unregister it. */
+  async close(): Promise<void> {
+    const child = this.child;
+    this.child = null;
+    if (child === null) {
+      return;
+    }
+    const pid = child.pid;
+    if (this.childExited === null) {
+      try {
+        child.stdin?.end();
+      } catch {
+        // The pipe may already be closed.
+      }
+      // The child spawns in its own group; tool calls can leave
+      // grandchildren that leader-only signals would orphan forever.
+      if (pid !== undefined) {
+        killGroup(pid, "SIGTERM");
+      } else {
+        child.kill("SIGTERM");
+      }
+      const exited = await waitForExit(child, ABORT_GRACE_MS);
+      if (!exited) {
+        if (pid !== undefined) {
+          killGroup(pid, "SIGKILL");
+        } else {
+          child.kill("SIGKILL");
+        }
+      }
+    }
+    if (pid !== undefined) {
+      this.options.registry.unregister(pid);
+    }
+  }
+
+  private currentExit(): { code: number | null; signal: string | null } | null {
+    return this.childExited;
+  }
+
+  private ensureStarted(): void {
+    if (this.child !== null) {
+      const exited = this.currentExit();
+      if (exited !== null) {
+        throw new Error(
+          `Headless pi session exited (code ${exited.code}, signal ${exited.signal})`,
+        );
+      }
+      return;
+    }
+    if (BRIDGE_PATH === undefined) {
+      throw new Error("The pi-workflows rpc-bridge extension is missing from this installation");
+    }
+    const piBin = this.options.piBin ?? "pi";
+    const child = spawn(
+      piBin,
+      [
+        "--mode",
+        "rpc",
+        "--no-session",
+        // Isolation is required: an installed pi-workflows extension would
+        // otherwise load beside the bridge, register a competing workflow
+        // tool, and start its own resume and controller workers.
+        "--no-extensions",
+        "-e",
+        BRIDGE_PATH,
+        ...(this.options.piArgs ?? []),
+      ],
+      {
+        cwd: this.options.cwd,
+        env: { ...process.env, ...this.options.env },
+        stdio: ["pipe", "pipe", "pipe"],
+        // Own process group: the host kills groups, never individual PIDs,
+        // so a killed host cannot leave an orphaned agent working.
+        detached: true,
+      },
+    );
+    this.child = child;
+    if (child.pid !== undefined) {
+      this.options.registry.register(child.pid);
+    }
+    child.on("exit", (code, signal) => {
+      this.childExited = { code, signal };
+      if (child.pid !== undefined) {
+        this.options.registry.unregister(child.pid);
+      }
+      this.wakeSubmissionWaiters();
+    });
+    // A failed spawn (missing binary, EACCES) emits error instead of exit;
+    // without a listener it would crash the host as an uncaught exception.
+    child.on("error", (error) => {
+      this.childExited = { code: null, signal: error.message };
+      this.wakeSubmissionWaiters();
+    });
+    child.stdin?.on("error", () => {
+      // Writes racing child exit surface asynchronously as EPIPE; the step
+      // fails through the exit path instead of an uncaught stream error.
+    });
+    child.stderr?.on("data", (chunk: Buffer) => {
+      this.stderrBuffer += chunk.toString("utf8");
+      const lines = this.stderrBuffer.split("\n");
+      this.stderrBuffer = lines.pop() ?? "";
+      for (const line of lines) {
+        if (!line.startsWith(RPC_SUBMISSION_PREFIX)) {
+          continue;
+        }
+        try {
+          const parsed = JSON.parse(line.slice(RPC_SUBMISSION_PREFIX.length)) as StepSubmission;
+          if (
+            parsed.action === "submit" &&
+            typeof parsed.step === "string" &&
+            typeof parsed.attempt === "string"
+          ) {
+            this.submissions.push(parsed);
+          }
+        } catch {
+          // A malformed marker line is ignored; the step times out instead.
+        }
+      }
+      this.wakeSubmissionWaiters();
+    });
+    // stdout carries the RPC event stream. The submission channel is stderr,
+    // so stdout is only drained to keep the pipe from blocking.
+    child.stdout?.on("data", (chunk: Buffer) => {
+      this.stdoutBuffer = (this.stdoutBuffer + chunk.toString("utf8")).slice(-64 * 1024);
+    });
+  }
+
+  private async promptForSubmission(
+    prompt: string,
+    request: AgentStepRequest,
+    signal: AbortSignal,
+  ): Promise<StepSubmission> {
+    const child = this.child;
+    if (child === null || this.childExited !== null) {
+      throw new Error("Headless pi session is not running");
+    }
+    const matching = this.takeMatchingSubmission(request);
+    if (matching !== undefined) {
+      return matching;
+    }
+    child.stdin?.write(`${JSON.stringify({ type: "prompt", message: prompt })}\n`);
+    for (;;) {
+      throwIfAborted(signal);
+      const exited = this.currentExit();
+      if (exited !== null) {
+        throw new Error(`Headless pi session exited mid-step (code ${exited.code})`);
+      }
+      const found = this.takeMatchingSubmission(request);
+      if (found !== undefined) {
+        return found;
+      }
+      await new Promise<void>((resolve, reject) => {
+        const waiter = () => {
+          signal.removeEventListener("abort", onAbort);
+          resolve();
+        };
+        const onAbort = () => {
+          this.sendAbortQuietly();
+          reject(signal.reason ?? new Error("Workflow step aborted"));
+        };
+        this.submissionWaiters.push(waiter);
+        signal.addEventListener("abort", onAbort, { once: true });
+      });
+    }
+  }
+
+  private takeMatchingSubmission(request: AgentStepRequest): StepSubmission | undefined {
+    const index = this.submissions.findIndex(
+      (candidate) =>
+        candidate.step === request.contract.nodeId &&
+        candidate.attempt === request.contract.attemptId,
+    );
+    return index === -1 ? undefined : (this.submissions.splice(index, 1)[0] as StepSubmission);
+  }
+
+  private sendAbortQuietly(): void {
+    try {
+      this.child?.stdin?.write(`${JSON.stringify({ type: "abort" })}\n`);
+    } catch {
+      // Best-effort.
+    }
+  }
+
+  private wakeSubmissionWaiters(): void {
+    const waiters = this.submissionWaiters;
+    this.submissionWaiters = [];
+    for (const waiter of waiters) {
+      waiter();
+    }
+  }
+}
+
+function killGroup(pid: number, signal: NodeJS.Signals): void {
+  try {
+    process.kill(-pid, signal);
+    return;
+  } catch {
+    // Group signals are unsupported here; the leader still needs it.
+  }
+  try {
+    process.kill(pid, signal);
+  } catch {
+    // Already gone.
+  }
+}
+
+function throwIfAborted(signal: AbortSignal): void {
+  if (signal.aborted) {
+    throw (signal.reason as unknown) ?? new Error("Workflow step aborted");
+  }
+}
+
+async function waitForExit(child: ChildProcess, timeoutMs: number): Promise<boolean> {
+  if (child.exitCode !== null || child.signalCode !== null) {
+    return true;
+  }
+  return await new Promise<boolean>((resolve) => {
+    const timer = setTimeout(() => resolve(false), timeoutMs);
+    child.once("exit", () => {
+      clearTimeout(timer);
+      resolve(true);
+    });
+  });
+}