pi-workflow-engine 0.2.0 → 0.4.0

package/.pi/extensions/pi-workflow-engine/src/concurrency.ts

@@ -0,0 +1,171 @@
+import { performance } from "node:perf_hooks";
+import { abortReason, throwIfAborted } from "./cancellation.ts";
+
+/**
+ * A counting semaphore: bounds how many async tasks run at once.
+ *
+ * This is the single global concurrency cap for a workflow run. Every `agent()`
+ * call acquires a slot before spawning a session and releases it when done, so the
+ * cap holds no matter how `parallel`/`pipeline` are nested — exactly like the
+ * built-in Workflow tool, where "concurrent agent() calls are capped and excess queue".
+ */
+export class Semaphore {
+  private active = 0;
+  private readonly waiters: Array<() => void> = [];
+
+  constructor(private readonly max: number) {}
+
+  async run<T>(fn: () => Promise<T>, options: { onQueueWaitMs?: (durationMs: number) => void; signal?: AbortSignal } = {}): Promise<T> {
+    throwIfAborted(options.signal);
+    const queuedAt = performance.now();
+    if (this.active >= this.max) {
+      await this.waitForSlot(options.signal);
+    }
+    options.onQueueWaitMs?.(performance.now() - queuedAt);
+    throwIfAborted(options.signal);
+    this.active++;
+    try {
+      return await fn();
+    } finally {
+      this.active--;
+      this.waiters.shift()?.();
+    }
+  }
+
+  private async waitForSlot(signal: AbortSignal | undefined): Promise<void> {
+    await new Promise<void>((resolve, reject) => {
+      let settled = false;
+      const cleanup = () => {
+        if (signal) signal.removeEventListener("abort", onAbort);
+      };
+      const waiter = () => {
+        if (settled) return;
+        settled = true;
+        cleanup();
+        resolve();
+      };
+      const onAbort = () => {
+        if (settled) return;
+        settled = true;
+        const index = this.waiters.indexOf(waiter);
+        if (index !== -1) this.waiters.splice(index, 1);
+        cleanup();
+        reject(abortReason(signal));
+      };
+      if (signal?.aborted) {
+        reject(abortReason(signal));
+        return;
+      }
+      this.waiters.push(waiter);
+      signal?.addEventListener("abort", onAbort, { once: true });
+    });
+  }
+}
+
+export interface ParallelOptions {
+  readonly signal?: AbortSignal;
+  readonly abortController?: AbortController;
+  readonly limit?: number;
+}
+
+/** Run every thunk concurrently and wait for all results (a barrier). */
+export async function parallel<T>(thunks: Array<() => Promise<T>>, options: ParallelOptions = {}): Promise<T[]> {
+  throwIfAborted(options.signal);
+  const results = new Array<T>(thunks.length);
+  let remaining = thunks.length;
+  if (remaining === 0) return results;
+
+  const limit = normalizeLimit(options.limit, thunks.length);
+  return await new Promise<T[]>((resolve, reject) => {
+    let rejected = false;
+    let active = 0;
+    let next = 0;
+    const cleanup = () => options.signal?.removeEventListener("abort", onAbort);
+    const rejectOnce = (error: unknown) => {
+      if (rejected) return;
+      rejected = true;
+      options.abortController?.abort(error);
+      cleanup();
+      reject(error);
+    };
+    const onAbort = () => rejectOnce(abortReason(options.signal));
+    const launchMore = () => {
+      while (!rejected && active < limit && next < thunks.length) {
+        const index = next++;
+        active++;
+        void Promise.resolve()
+          .then(() => {
+            throwIfAborted(options.signal);
+            return thunks[index]();
+          })
+          .then(
+            (value) => {
+              if (rejected) return;
+              active--;
+              results[index] = value;
+              remaining--;
+              if (remaining === 0) {
+                cleanup();
+                resolve(results);
+                return;
+              }
+              launchMore();
+            },
+            (error: unknown) => rejectOnce(error),
+          );
+      }
+    };
+
+    options.signal?.addEventListener("abort", onAbort, { once: true });
+    launchMore();
+  });
+}
+
+function normalizeLimit(limit: number | undefined, itemCount: number): number {
+  if (limit === undefined || !Number.isFinite(limit)) return Math.max(1, itemCount);
+  return Math.max(1, Math.min(itemCount, Math.trunc(limit)));
+}
+
+/**
+ * Run each item through all stages independently. There is NO barrier between
+ * stages: item A can be in stage 3 while item B is still in stage 1. Actual agent
+ * concurrency is bounded by the run's Semaphore (inside each `agent()`), not here —
+ * so launching every item at once is safe.
+ *
+ * Each stage receives (previousResult, originalItem, index).
+ */
+export async function pipeline<Item, A>(
+  items: readonly Item[],
+  stage1: (prev: Item, item: Item, index: number) => Promise<A>,
+): Promise<A[]>;
+export async function pipeline<Item, A, B>(
+  items: readonly Item[],
+  stage1: (prev: Item, item: Item, index: number) => Promise<A>,
+  stage2: (prev: A, item: Item, index: number) => Promise<B>,
+): Promise<B[]>;
+export async function pipeline<Item, A, B, C>(
+  items: readonly Item[],
+  stage1: (prev: Item, item: Item, index: number) => Promise<A>,
+  stage2: (prev: A, item: Item, index: number) => Promise<B>,
+  stage3: (prev: B, item: Item, index: number) => Promise<C>,
+): Promise<C[]>;
+export async function pipeline(
+  items: readonly unknown[],
+  ...stages: Array<(prev: unknown, item: unknown, index: number) => Promise<unknown>>
+): Promise<unknown[]>;
+export async function pipeline(
+  items: readonly unknown[],
+  ...stages: Array<(prev: unknown, item: unknown, index: number) => Promise<unknown>>
+): Promise<unknown[]> {
+  return Promise.all(
+    items.map(async (item, index) => {
+      let acc: unknown = item;
+      for (const stage of stages) {
+        acc = await stage(acc, item, index);
+      }
+      return acc;
+    }),
+  );
+}
+
+export type Pipeline = typeof pipeline;

package/.pi/extensions/pi-workflow-engine/src/diff-capture.ts

@@ -0,0 +1,152 @@
+import { spawn } from "node:child_process";
+import { performance } from "node:perf_hooks";
+import type { PerfSink } from "./perf.ts";
+
+export interface AllowedDiffCommand {
+  readonly file: "git" | "gh";
+  readonly args: readonly string[];
+}
+
+export interface DiffCaptureResult {
+  readonly ok: boolean;
+  readonly stdout: string;
+  readonly durationMs: number;
+  readonly bytes: number;
+  readonly error?: string;
+}
+
+export interface DiffCaptureOptions {
+  readonly cwd: string;
+  readonly signal?: AbortSignal;
+  readonly timeoutMs: number;
+  readonly maxBufferBytes: number;
+  readonly perf?: PerfSink;
+  readonly env?: NodeJS.ProcessEnv;
+  readonly killGraceMs?: number;
+}
+
+const SAFE_REF_OR_PATH = /^[A-Za-z0-9_./:@~+=,\-]+$/;
+const SAFE_GH_FLAG = /^--[A-Za-z0-9-]+(=[A-Za-z0-9_./:@~+=,\-]+)?$/;
+const SAFE_GIT_DIFF_FLAGS = new Set(["--cached", "--staged", "--no-color", "--color=never", "--no-ext-diff"]);
+
+export function parseAllowedDiffCommand(command: string): AllowedDiffCommand | { error: string } {
+  const tokens = command.trim().split(/\s+/).filter(Boolean);
+  if (tokens.length < 2) return { error: "empty or incomplete diff command" };
+  if (tokens[0] === "git" && tokens[1] === "diff") return parseGitDiff(tokens);
+  if (tokens[0] === "gh" && tokens[1] === "pr" && tokens[2] === "diff" && /^\d+$/.test(tokens[3] ?? "")) {
+    const flags = tokens.slice(4);
+    const unsafe = flags.find((token) => !SAFE_GH_FLAG.test(token));
+    if (unsafe) return { error: `unsupported gh pr diff token: ${unsafe}` };
+    return { file: "gh", args: ["pr", "diff", tokens[3], ...flags] };
+  }
+  return { error: "diff command is not in the git/gh allowlist" };
+}
+
+function parseGitDiff(tokens: readonly string[]): AllowedDiffCommand | { error: string } {
+  const args = tokens.slice(2);
+  const safeArgs = ["diff", "--no-ext-diff"];
+  let pathMode = false;
+
+  for (const token of args) {
+    if (!SAFE_REF_OR_PATH.test(token)) return { error: `unsupported git diff token: ${token}` };
+    if (token === "--") {
+      pathMode = true;
+      safeArgs.push(token);
+      continue;
+    }
+    if (!pathMode && token.startsWith("-")) {
+      if (!isAllowedGitDiffFlag(token)) return { error: `unsupported git diff option: ${token}` };
+      if (token !== "--no-ext-diff") safeArgs.push(token);
+      continue;
+    }
+    safeArgs.push(token);
+  }
+
+  return { file: "git", args: safeArgs };
+}
+
+function isAllowedGitDiffFlag(token: string): boolean {
+  return SAFE_GIT_DIFF_FLAGS.has(token) || /^-U\d+$/.test(token) || /^--unified=\d+$/.test(token) || /^--inter-hunk-context=\d+$/.test(token);
+}
+
+export async function captureDiff(command: string, options: DiffCaptureOptions): Promise<DiffCaptureResult> {
+  const parsed = parseAllowedDiffCommand(command);
+  if ("error" in parsed) {
+    return { ok: false, stdout: "", durationMs: 0, bytes: 0, error: parsed.error };
+  }
+
+  const start = performance.now();
+  let stdout = "";
+  let stderr = "";
+  let bytes = 0;
+  let error: string | undefined;
+  const child = spawn(parsed.file, [...parsed.args], {
+    cwd: options.cwd,
+    env: diffCaptureEnv(options.env),
+    stdio: ["ignore", "pipe", "pipe"],
+  });
+
+  return await new Promise<DiffCaptureResult>((resolve) => {
+    let settled = false;
+    let forceKillTimer: ReturnType<typeof setTimeout> | undefined;
+    const killGraceMs = Math.max(1, options.killGraceMs ?? 100);
+    const finish = (ok: boolean, finishError?: string) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      const durationMs = performance.now() - start;
+      options.perf?.observe("diff.capture_ms", durationMs);
+      options.perf?.observe("diff.bytes", bytes);
+      resolve({ ok, stdout, durationMs, bytes, error: finishError });
+    };
+    const kill = (message: string) => {
+      error = error ?? message;
+      child.kill("SIGTERM");
+      forceKillTimer ??= setTimeout(() => {
+        child.kill("SIGKILL");
+        finish(false, error);
+      }, killGraceMs);
+    };
+    const onAbort = () => kill("diff capture aborted");
+    const timeout = setTimeout(() => kill(`diff capture timed out after ${options.timeoutMs}ms`), options.timeoutMs);
+    const cleanup = () => {
+      clearTimeout(timeout);
+      if (forceKillTimer) clearTimeout(forceKillTimer);
+      options.signal?.removeEventListener("abort", onAbort);
+    };
+
+    if (options.signal?.aborted) {
+      kill("diff capture aborted");
+    } else {
+      options.signal?.addEventListener("abort", onAbort, { once: true });
+    }
+
+    child.stdout?.on("data", (chunk: Buffer) => {
+      bytes += chunk.length;
+      if (bytes > options.maxBufferBytes) {
+        kill(`diff capture exceeded ${options.maxBufferBytes} bytes`);
+        return;
+      }
+      stdout += chunk.toString("utf8");
+    });
+    child.stderr?.on("data", (chunk: Buffer) => {
+      stderr += chunk.toString("utf8");
+    });
+    child.on("error", (spawnError) => finish(false, spawnError.message));
+    child.on("close", (code, signal) => {
+      if (error) {
+        finish(false, error);
+        return;
+      }
+      if (code === 0) {
+        finish(true);
+        return;
+      }
+      finish(false, stderr.trim() || `diff command exited with code ${code ?? `signal ${signal ?? "unknown"}`}`);
+    });
+  });
+}
+
+function diffCaptureEnv(env: NodeJS.ProcessEnv | undefined): NodeJS.ProcessEnv {
+  return { ...(env ?? process.env), GIT_EXTERNAL_DIFF: "", GIT_DIFF_OPTS: "" };
+}

package/.pi/extensions/pi-workflow-engine/src/discovery.ts

@@ -0,0 +1,114 @@
+import { readdir } from "node:fs/promises";
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { pathToFileURL } from "node:url";
+import type { WorkflowMeta, WorkflowModule } from "./types.ts";
+import type { PerfSink } from "./perf.ts";
+import { BUILTIN_WORKFLOW_FILES, BUILTIN_WORKFLOWS } from "./workflows.ts";
+
+type WorkflowModuleCandidate = {
+  meta?: { name?: unknown; description?: unknown; phases?: unknown };
+  default?: unknown;
+};
+
+export interface DiscoverWorkflowsOptions {
+  readonly refresh?: boolean;
+  readonly perf?: PerfSink;
+  readonly userWorkflowDir?: string;
+}
+
+const discoveryCache = new Map<string, Map<string, WorkflowModule>>();
+
+function parseWorkflowModule(value: unknown): { module: WorkflowModule } | { reason: string } {
+  if (!value || typeof value !== "object") return { reason: "module export is not an object" };
+  const candidate = value as WorkflowModuleCandidate;
+  const meta = candidate.meta;
+  if (!meta || typeof meta !== "object") return { reason: "missing meta export" };
+  if (typeof meta.name !== "string") return { reason: "meta.name must be a string" };
+
+  const description = meta.description;
+  if (description !== undefined && typeof description !== "string") return { reason: "meta.description must be a string when provided" };
+
+  const phases = meta.phases;
+  if (phases !== undefined && !isWorkflowPhases(phases)) return { reason: "meta.phases must be an array of { title: string }" };
+  if (!isWorkflowRun(candidate.default)) return { reason: "default export must be a function" };
+
+  const workflowMeta: WorkflowMeta = { name: meta.name, description: description ?? "" };
+  if (phases !== undefined) workflowMeta.phases = phases;
+  return { module: { meta: workflowMeta, default: candidate.default } };
+}
+
+function isWorkflowRun(value: unknown): value is WorkflowModule["default"] {
+  return typeof value === "function";
+}
+
+function isWorkflowPhases(value: unknown): value is Array<{ title: string }> {
+  return Array.isArray(value) && value.every((phase) => typeof phase === "object" && phase !== null && typeof (phase as { title?: unknown }).title === "string");
+}
+
+/** Best-effort dynamic load of every `*.ts` workflow in a directory. */
+async function loadDir(dir: string, excludeFiles: ReadonlySet<string> = new Set()): Promise<WorkflowModule[]> {
+  let entries: string[];
+  try {
+    entries = await readdir(dir);
+  } catch {
+    return [];
+  }
+
+  const modules: WorkflowModule[] = [];
+  for (const name of entries) {
+    if (excludeFiles.has(name)) continue;
+    if (!name.endsWith(".ts") || name.startsWith("_") || name.startsWith(".")) continue;
+    try {
+      const loaded: unknown = await import(pathToFileURL(join(dir, name)).href);
+      const parsed = parseWorkflowModule(loaded);
+      if ("module" in parsed) modules.push(parsed.module);
+      else console.error(`[workflow-engine] skipped ${name}: ${parsed.reason}`);
+    } catch (error) {
+      // Drop-in loading depends on the runtime resolving TS + the bundled typebox.
+      // Failures here are non-fatal: the static registry still works.
+      console.error(`[workflow-engine] skipped ${name}: ${error instanceof Error ? error.message : String(error)}`);
+    }
+  }
+  return modules;
+}
+
+/**
+ * All available workflows by name. Static registry wins on name collisions, so the
+ * bundled example is always the verified one even if a same-named file is dropped in.
+ */
+export async function discoverWorkflows(repoDir: string, options: DiscoverWorkflowsOptions = {}): Promise<Map<string, WorkflowModule>> {
+  const userWorkflowDir = options.userWorkflowDir ?? join(homedir(), ".pi", "agent", "workflows");
+  const cacheKey = `${repoDir}\0${userWorkflowDir}`;
+  const cached = discoveryCache.get(cacheKey);
+  if (cached && !options.refresh) {
+    options.perf?.counter("discovery.cache_hit");
+    return new Map(cached);
+  }
+
+  const byName = await timed(options.perf, "discovery.total_ms", async () => {
+    const next = new Map<string, WorkflowModule>();
+    for (const mod of BUILTIN_WORKFLOWS) next.set(mod.meta.name, mod);
+
+    const repoWorkflowDir = join(repoDir, "workflows");
+    const [repoDynamic, userDynamic] = await Promise.all([
+      timed(options.perf, "discovery.repo_dir_ms", () => loadDir(repoWorkflowDir, BUILTIN_WORKFLOW_FILES)),
+      timed(options.perf, "discovery.user_dir_ms", () => loadDir(userWorkflowDir)),
+    ]);
+
+    for (const mod of repoDynamic) {
+      if (!next.has(mod.meta.name)) next.set(mod.meta.name, mod);
+    }
+    for (const mod of userDynamic) {
+      if (!next.has(mod.meta.name)) next.set(mod.meta.name, mod);
+    }
+    return next;
+  });
+
+  discoveryCache.set(cacheKey, byName);
+  return new Map(byName);
+}
+
+async function timed<T>(perf: PerfSink | undefined, name: string, fn: () => Promise<T>): Promise<T> {
+  return perf ? await perf.time(name, fn) : await fn();
+}

package/{src → .pi/extensions/pi-workflow-engine/src}/engine.ts

@@ -1,13 +1,15 @@
-import { cpus } from "node:os";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { parallel, pipeline, Semaphore } from "./concurrency.ts";
+import { linkAbortSignal } from "./cancellation.ts";
 import { runAgent, type RunContext } from "./agent-runner.ts";
 import { ProgressTracker } from "./progress.ts";
+import { createPerfRecorder, type PerfSnapshot } from "./perf.ts";
+import { defaultConcurrency, resolveWorkflowRunOptions } from "./options.ts";
 import type { AgentOptions, WorkflowApi, WorkflowModule, WorkflowRunOptions } from "./types.ts";
 import { WorkflowInspector } from "./ui/workflow-inspector.ts";
 
 /** Default global cap on concurrent agents per run. */
-const DEFAULT_CONCURRENCY = Math.min(8, Math.max(2, cpus().length));
+const DEFAULT_CONCURRENCY = defaultConcurrency();
 
 /**
  * Run a workflow module: build the per-run primitives (binding agent/parallel/pipeline
@@ -19,8 +21,9 @@ export async function runWorkflow(
   args: string,
   options: WorkflowRunOptions = {},
 ): Promise<unknown> {
+  const resolvedOptions = resolveWorkflowRunOptions(options);
   const progress = new ProgressTracker(ctx, mod.meta.name);
-  if (options.inspect && ctx.hasUI) {
+  if (resolvedOptions.inspect && ctx.hasUI) {
     void ctx.ui
       .custom<void>(
         (tui, theme, _keybindings, done) => new WorkflowInspector(() => progress.snapshot(), tui, theme, () => done(undefined)),
@@ -29,32 +32,57 @@ export async function runWorkflow(
       .catch((error: unknown) => progress.log(`inspector failed: ${error instanceof Error ? error.message : String(error)}`));
   }
 
+  const perf = resolvedOptions.perfRecorder ?? createPerfRecorder(resolvedOptions.perf);
+  const runAbortController = new AbortController();
+  const unlinkContextAbortSignal = linkAbortSignal(ctx.signal, runAbortController);
+  const unlinkOptionAbortSignal = linkAbortSignal(resolvedOptions.signal, runAbortController);
   const rc: RunContext = {
     cwd: ctx.cwd,
     hostModel: ctx.model,
     modelRegistry: ctx.modelRegistry,
-    semaphore: new Semaphore(DEFAULT_CONCURRENCY),
+    semaphore: new Semaphore(resolvedOptions.concurrency ?? DEFAULT_CONCURRENCY),
     progress,
-    signal: ctx.signal,
+    signal: runAbortController.signal,
+    perf,
   };
 
   const agent = ((prompt: string, opts?: AgentOptions) => runAgent(rc, prompt, opts)) as WorkflowApi["agent"];
 
   const api: WorkflowApi = {
     agent,
-    parallel,
+    parallel: (thunks) =>
+      parallel(thunks, {
+        signal: runAbortController.signal,
+        abortController: runAbortController,
+        limit: resolvedOptions.parallelSubmissionLimit ?? resolvedOptions.concurrency * 2,
+      }),
     pipeline,
     phase: (title) => progress.phase(title),
     log: (message) => progress.log(message),
     progress: (event) => progress.event(event),
     args,
     cwd: ctx.cwd,
-    signal: ctx.signal,
+    signal: runAbortController.signal,
   };
 
   try {
-    return await mod.default(api);
+    return await perf.time("workflow.total_ms", () => mod.default(api));
   } finally {
+    const snapshot = perf.snapshot();
+    if (resolvedOptions.perf) {
+      resolvedOptions.onPerfSnapshot?.(snapshot);
+      progress.log(formatPerfSummary(snapshot));
+    }
+    unlinkContextAbortSignal();
+    unlinkOptionAbortSignal();
     progress.done();
   }
 }
+
+function formatPerfSummary(snapshot: PerfSnapshot): string {
+  const parts = snapshot.aggregates
+    .filter((aggregate) => aggregate.count > 0)
+    .slice(0, 5)
+    .map((aggregate) => `${aggregate.name} ${Math.round(aggregate.total)}ms`);
+  return parts.length > 0 ? `perf: ${parts.join(", ")}` : "perf: no samples";
+}

package/.pi/extensions/pi-workflow-engine/src/options.ts

@@ -0,0 +1,46 @@
+import { cpus } from "node:os";
+import type { WorkflowRunOptions } from "./types.ts";
+
+export interface ResolvedWorkflowRunOptions extends WorkflowRunOptions {
+  readonly perf: boolean;
+  readonly concurrency: number;
+  readonly parallelSubmissionLimit?: number;
+}
+
+export function defaultConcurrency(cpuCount = cpus().length): number {
+  return Math.min(8, Math.max(2, cpuCount));
+}
+
+export function resolveWorkflowRunOptions(
+  input: WorkflowRunOptions = {},
+  env: Record<string, string | undefined> = process.env,
+): ResolvedWorkflowRunOptions {
+  const concurrency = clampInteger(input.concurrency ?? parseInteger(env.PI_WORKFLOW_CONCURRENCY), 1, 64, defaultConcurrency());
+  const parallelSubmissionLimit = optionalClampedInteger(
+    input.parallelSubmissionLimit ?? parseInteger(env.PI_WORKFLOW_PARALLEL_SUBMISSION_LIMIT),
+    1,
+    10_000,
+  );
+  return {
+    ...input,
+    perf: input.perf ?? env.PI_WORKFLOW_PERF === "1",
+    concurrency,
+    parallelSubmissionLimit,
+  };
+}
+
+function optionalClampedInteger(value: number | undefined, min: number, max: number): number | undefined {
+  if (value === undefined) return undefined;
+  return clampInteger(value, min, max, value);
+}
+
+function clampInteger(value: number | undefined, min: number, max: number, fallback: number): number {
+  if (value === undefined || !Number.isFinite(value)) return fallback;
+  return Math.min(max, Math.max(min, Math.trunc(value)));
+}
+
+function parseInteger(value: string | undefined): number | undefined {
+  if (value === undefined || value.trim() === "") return undefined;
+  const parsed = Number(value);
+  return Number.isFinite(parsed) ? Math.trunc(parsed) : undefined;
+}