pi-workflow-engine 1.1.0 → 1.3.0

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

@@ -4,18 +4,75 @@ import { pathToFileURL } from "node:url";
 import { getAgentDir } from "@earendil-works/pi-coding-agent";
 import type { LoadedWorkflow, WorkflowSourceIdentity } from "./types.ts";
 import type { PerfSink } from "./perf.ts";
+import { compileInlineWorkflow, inlineScriptFingerprint } from "./inline-workflow.ts";
 import { loadWorkflow, parseWorkflowModule } from "./workflow-module.ts";
 import { BUILTIN_SOURCE_ROOT, BUILTIN_WORKFLOW_DEFINITIONS, BUILTIN_WORKFLOW_FILES } from "./workflows.ts";
 import { captureSourceTreeFingerprint, type FingerprintCapture } from "./tree-fingerprint.ts";
 import { unknownErrorMessage } from "./unknown-error.ts";
+import { readBoundedScriptFile } from "./workflow-script-file.ts";
 
 export interface DiscoverWorkflowsOptions {
   readonly refresh?: boolean;
   readonly perf?: PerfSink;
   readonly userWorkflowDir?: string;
+  /**
+   * Project root whose `.pi/workflows/*.js` saved scripts load right after the built-ins. Callers omit
+   * it for an untrusted project, which then contributes no workflows.
+   */
+  readonly projectDir?: string;
+}
+
+/**
+ * Built-ins and `.ts` drop-ins hash source trees and import modules, so they load once per directory pair
+ * until a refresh. Saved scripts are not cached: see `discoverWorkflows`.
+ */
+interface CachedModules {
+  readonly builtins: readonly LoadedWorkflow[];
+  readonly repoDynamic: readonly LoadedWorkflow[];
+  readonly userDynamic: readonly LoadedWorkflow[];
+}
+
+const moduleCache = new Map<string, CachedModules>();
+
+/** Saved scripts are re-read on every discovery, so each distinct skip is reported once, not on every lookup. */
+const reportedScriptSkips = new Set<string>();
+
+function isWorkflowFileCandidate(name: string, extension: string): boolean {
+  return name.endsWith(extension) && !name.startsWith("_") && !name.startsWith(".");
+}
+
+/**
+ * Load every `*.js` saved workflow in the inline script format from a directory. Unlike `.ts` drop-ins
+ * they are compiled with the injected `Type`, so they share pi's typebox, and a content fingerprint is
+ * their whole identity, so they replay on resume.
+ */
+async function loadScriptDir(dir: string): Promise<LoadedWorkflow[]> {
+  let entries: string[];
+  try {
+    entries = await readdir(dir);
+  } catch {
+    return [];
+  }
+  const workflows: LoadedWorkflow[] = [];
+  for (const name of entries.filter((entry) => isWorkflowFileCandidate(entry, ".js")).sort()) {
+    const path = join(dir, name);
+    try {
+      // Discovery runs on every lookup, so a linked FIFO or huge file must not stall it.
+      const source = await readBoundedScriptFile(path, path);
+      const compiled = compileInlineWorkflow(source);
+      workflows.push(loadWorkflow(compiled, { kind: "file", path, root: dir, fingerprint: inlineScriptFingerprint(source), selfContained: true }));
+    } catch (error) {
+      const note = `[workflow-engine] skipped ${name}: ${unknownErrorMessage(error)}`;
+      const key = `${path}\0${note}`;
+      if (!reportedScriptSkips.has(key)) {
+        reportedScriptSkips.add(key);
+        console.error(note);
+      }
+    }
+  }
+  return workflows;
 }
 
-const discoveryCache = new Map<string, Map<string, LoadedWorkflow>>();
 /** Best-effort dynamic load of every `*.ts` workflow in a directory. */
 async function loadDir(
   dir: string,
@@ -29,9 +86,7 @@ async function loadDir(
   } catch {
     return [];
   }
-  const candidates = entries.filter(
-    (name) => !excludeFiles.has(name) && name.endsWith(".ts") && !name.startsWith("_") && !name.startsWith("."),
-  );
+  const candidates = entries.filter((name) => !excludeFiles.has(name) && isWorkflowFileCandidate(name, ".ts"));
   // The shipped workflows/ dir holds only excluded built-ins; skip the tree hash when nothing will be imported.
   if (candidates.length === 0) return [];
 
@@ -56,72 +111,86 @@ async function loadDir(
 }
 
 /**
- * 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.
+ * All available workflows by name. The first source to define a name wins: the static registry, so
+ * the bundled example is always the verified one even if a same-named file is dropped in, then the
+ * project's saved scripts, the extension's `.ts` drop-ins, the user's saved scripts, and the user's
+ * `.ts` drop-ins.
+ *
+ * Saved scripts are read and compiled on every call (small files, no import), so one written or edited
+ * mid-session resolves by name with its current content; built-ins and `.ts` drop-ins come from the cache
+ * unless `refresh` is set.
  */
 export async function discoverWorkflows(repoDir: string, options: DiscoverWorkflowsOptions = {}): Promise<Map<string, LoadedWorkflow>> {
   const userWorkflowDir = options.userWorkflowDir ?? join(getAgentDir(), "workflows");
+  const projectWorkflowDir = options.projectDir === undefined ? undefined : join(options.projectDir, ".pi", "workflows");
+  return await timed(options.perf, "discovery.total_ms", async () => {
+    const [modules, projectScripts, userScripts] = await Promise.all([
+      cachedModules(repoDir, userWorkflowDir, options),
+      projectWorkflowDir === undefined
+        ? []
+        : timed(options.perf, "discovery.project_scripts_ms", () => loadScriptDir(projectWorkflowDir)),
+      timed(options.perf, "discovery.user_scripts_ms", () => loadScriptDir(userWorkflowDir)),
+    ]);
+    const byName = new Map<string, LoadedWorkflow>();
+    for (const mod of [...modules.builtins, ...projectScripts, ...modules.repoDynamic, ...userScripts, ...modules.userDynamic]) {
+      if (!byName.has(mod.meta.name)) byName.set(mod.meta.name, mod);
+    }
+    return byName;
+  });
+}
+
+async function cachedModules(repoDir: string, userWorkflowDir: string, options: DiscoverWorkflowsOptions): Promise<CachedModules> {
   const cacheKey = `${repoDir}\0${userWorkflowDir}`;
-  const cached = discoveryCache.get(cacheKey);
+  const cached = moduleCache.get(cacheKey);
   if (cached && !options.refresh) {
     options.perf?.counter("discovery.cache_hit");
-    return new Map(cached);
+    return cached;
   }
 
-  const byName = await timed(options.perf, "discovery.total_ms", async () => {
-    const next = new Map<string, LoadedWorkflow>();
-    const builtinSource = await captureSourceTreeFingerprint(BUILTIN_SOURCE_ROOT);
-    for (const definition of BUILTIN_WORKFLOW_DEFINITIONS) {
-      const source: WorkflowSourceIdentity =
-        builtinSource.kind === "verified"
-          ? {
-              kind: "file",
-              path: definition.path,
-              root: BUILTIN_SOURCE_ROOT,
-              fingerprint: builtinSource.fingerprint,
-            }
-          : {
-              kind: "unverifiable",
-              reason: `built-in workflow source capture failed: ${builtinSource.reason}`,
-            };
-      next.set(definition.module.meta.name, loadWorkflow(definition.module, source));
-    }
-
-    const repoWorkflowDir = join(repoDir, "workflows");
-    // The shipped extension dir is both the built-in root and the drop-in root; hash that tree once.
-    const repoProvenance =
-      resolve(repoDir) === resolve(BUILTIN_SOURCE_ROOT) ? async () => builtinSource : () => captureSourceTreeFingerprint(repoDir);
-    const [repoDynamic, userDynamic] = await Promise.all([
-      timed(options.perf, "discovery.repo_dir_ms", () =>
-        loadDir(
-          repoWorkflowDir,
-          () => ({
+  const builtinSource = await captureSourceTreeFingerprint(BUILTIN_SOURCE_ROOT);
+  const builtins = BUILTIN_WORKFLOW_DEFINITIONS.map((definition) => {
+    const source: WorkflowSourceIdentity =
+      builtinSource.kind === "verified"
+        ? {
+            kind: "file",
+            path: definition.path,
+            root: BUILTIN_SOURCE_ROOT,
+            fingerprint: builtinSource.fingerprint,
+          }
+        : {
             kind: "unverifiable",
-            reason: "dynamic workflow module graphs are not loaded from an immutable source snapshot",
-          }),
-          BUILTIN_WORKFLOW_FILES,
-          repoProvenance,
-        ),
-      ),
-      timed(options.perf, "discovery.user_dir_ms", () =>
-        loadDir(userWorkflowDir, () => ({
+            reason: `built-in workflow source capture failed: ${builtinSource.reason}`,
+          };
+    return loadWorkflow(definition.module, source);
+  });
+
+  const repoWorkflowDir = join(repoDir, "workflows");
+  // The shipped extension dir is both the built-in root and the drop-in root; hash that tree once.
+  const repoProvenance =
+    resolve(repoDir) === resolve(BUILTIN_SOURCE_ROOT) ? async () => builtinSource : () => captureSourceTreeFingerprint(repoDir);
+  const [repoDynamic, userDynamic] = await Promise.all([
+    timed(options.perf, "discovery.repo_dir_ms", () =>
+      loadDir(
+        repoWorkflowDir,
+        () => ({
           kind: "unverifiable",
-          reason: "dynamic user workflow dependencies do not have a declared trusted source root",
-        })),
+          reason: "dynamic workflow module graphs are not loaded from an immutable source snapshot",
+        }),
+        BUILTIN_WORKFLOW_FILES,
+        repoProvenance,
       ),
-    ]);
-
-    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;
-  });
+    ),
+    timed(options.perf, "discovery.user_dir_ms", () =>
+      loadDir(userWorkflowDir, () => ({
+        kind: "unverifiable",
+        reason: "dynamic user workflow dependencies do not have a declared trusted source root",
+      })),
+    ),
+  ]);
 
-  discoveryCache.set(cacheKey, byName);
-  return new Map(byName);
+  const modules: CachedModules = { builtins, repoDynamic, userDynamic };
+  moduleCache.set(cacheKey, modules);
+  return modules;
 }
 
 async function timed<T>(perf: PerfSink | undefined, name: string, fn: () => Promise<T>): Promise<T> {

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

@@ -1,8 +1,10 @@
-import { CustomEditor, type ExtensionAPI, type ExtensionContext, type Theme } from "@earendil-works/pi-coding-agent";
+import { CustomEditor, type BeforeAgentStartEvent, type ExtensionAPI, type ExtensionContext, type Theme } from "@earendil-works/pi-coding-agent";
 import type { EditorComponent } from "@earendil-works/pi-tui";
 import { completeCurrentArgument } from "./command-completions.ts";
 import type { DynamaxShortcuts } from "./dynamax-shortcuts.ts";
+import { appendSectionToForcedPrompt } from "./prompt-sections.ts";
 import { sessionKey } from "./session-identity.ts";
+import { WORKFLOW_TOOL_NAME } from "./types.ts";
 import { unknownErrorMessage } from "./unknown-error.ts";
 import { dot, GLYPH } from "./ui/kit.ts";
 import { setWorkflowInspectorShortcut } from "./ui/workflow-widget.ts";
@@ -20,8 +22,19 @@ export interface DynamaxState {
   turnActive: boolean;
 }
 
+/** The output-token pool a `+500k`-style directive in the user's prompt sets for that turn's workflow runs. */
+export interface DynamaxTurnBudget {
+  readonly total: number;
+  /** Tokens not yet spent by this turn's synchronous runs; never negative. */
+  remaining(): number;
+  /** Count a finished synchronous run's output tokens against the pool. */
+  spend(outputTokens: number): void;
+}
+
 export interface DynamaxRuntime {
   state: DynamaxState;
+  /** Set by a budget directive in a user prompt, cleared when that run settles. */
+  turnBudget?: DynamaxTurnBudget;
 }
 
 export type DynamaxRuntimeStore = Map<string, DynamaxRuntime>;
@@ -35,9 +48,18 @@ export interface DynamaxRegistrationOptions {
 export interface DynamaxHandle {
   /** Opt the next agent run in. The `input` hook ignores prompts this extension sends (pi tags them `source: "extension"`). */
   markOneShot(ctx: ExtensionContext): void;
+  /** The current turn's output-token pool, when the user's prompt set one. */
+  turnBudget(ctx: ExtensionContext): DynamaxTurnBudget | undefined;
+  /** Whether Dynamax covers this session now: sticky, a pending one-shot, or a one-shot turn in progress. */
+  isActive(ctx: ExtensionContext): boolean;
 }
 
 export const DYNAMAX_TOKEN_PATTERN = /(^|[^A-Za-z0-9_])dynamax([^A-Za-z0-9_]|$)/i;
+/** `+500k`, `+1.5m`, `+2M`, `+250000`: after the start or whitespace, before the end, whitespace, or punctuation. */
+export const BUDGET_DIRECTIVE_PATTERN = /(?:^|\s)\+(\d+(?:\.\d+)?)([km]?)(?=$|\s|[.,;:!?)\]}"'])/i;
+/** Smaller values are ignored, so a `+1` or `+2` in ordinary prose never caps a run at a few tokens. */
+export const BUDGET_DIRECTIVE_MIN_TOKENS = 1_000;
+const BUDGET_DIRECTIVE_MULTIPLIERS: Readonly<Record<string, number>> = { "": 1, k: 1_000, m: 1_000_000 };
 export const DYNAMAX_STATUS_KEY = "dynamax";
 
 type EditorFactory = NonNullable<ReturnType<ExtensionContext["ui"]["getEditorComponent"]>>;
@@ -53,29 +75,31 @@ function disposeDynamaxDecorations(decorations: Set<DynamaxEditorDecoration>): v
   decorations.clear();
 }
 
+/** With no argument the command toggles sticky mode, so these name only the explicit forms. */
 const DYNAMAX_ACTION_COMPLETIONS = [
-  { value: "on", description: "Keep Dynamax enabled for this session" },
+  { value: "on", description: "Make workflow orchestration the default for this session" },
   { value: "off", description: "Disable sticky and pending Dynamax modes" },
   { value: "status", description: "Show the current Dynamax state" },
 ] as const;
 
-export const ADAPTIVE_WORKFLOW_GUIDANCE = `
-Adaptive multi-pass workflows are optional. Use a simple single-pass fan-out when it is sufficient. When the first pass may expose gaps, conflicts, weak claims, or missing evidence:
-- run the bounded first-pass agents;
-- give their surviving results to a structured gap-analysis agent so the LLM decides what needs follow-up;
-- put a hard maxItems bound on LLM-authored task arrays and defensively slice them before fan-out;
-- use ordinary TypeScript conditionals or bounded loops to launch follow-up agents only when gaps exist;
-- synthesize the first-pass and follow-up results together.
-Do not generate a second pass when the first pass is sufficient, and do not invent iteration, quorum, graph, reduction, or retry primitives for this pattern.`;
-
 /** System prompt section pi wraps as `<dynamax>`; adding or dropping it appends a section patch instead of rewriting the prompt. */
 export const DYNAMAX_SECTION = "dynamax";
 
-/** Only the opt-in itself: the authoring rules already reach the model as the workflow tool's guidelines. */
+/**
+ * Only the opt-in itself: the workflow tool's guidelines send the model to the workflow-authoring
+ * skill, which carries the authoring rules and patterns.
+ */
 export const DYNAMAX_REMINDER =
   "The user opted into dynamax multi-agent orchestration for this request, so the workflow tool is permitted. " +
   "Run a registered workflow by name, or author an inline workflow script when none fits, following the workflow tool's guidelines.";
 
+/** Sticky mode is a standing default for the session, not a one-request permission. */
+export const DYNAMAX_STICKY_REMINDER =
+  "Dynamax is on for this session. Orchestrate every substantive task with the workflow tool by default: " +
+  "scout inline first to find the work list, then run one single-phase workflow per phase (understand, design, implement, review) " +
+  "and read each result before choosing the next. Lean toward adversarial verification of findings. " +
+  "Work without a workflow only for conversational turns and trivial or already-verified edits.";
+
 export function createDynamaxState(): DynamaxState {
   return { sticky: false, oneShotPending: false, turnActive: false };
 }
@@ -97,6 +121,46 @@ export function hasDynamaxToken(text: string): boolean {
   return DYNAMAX_TOKEN_PATTERN.test(text);
 }
 
+/** The output tokens the first budget directive in `text` asks for, if any. */
+export function parseBudgetDirective(text: string): number | undefined {
+  const match = BUDGET_DIRECTIVE_PATTERN.exec(text);
+  if (!match) return undefined;
+  const multiplier = BUDGET_DIRECTIVE_MULTIPLIERS[(match[2] ?? "").toLowerCase()] ?? 1;
+  // Round, not floor: 2.3 * 1000 is 2299.999... in floating point.
+  const tokens = Math.round(Number(match[1]) * multiplier);
+  return Number.isSafeInteger(tokens) && tokens >= BUDGET_DIRECTIVE_MIN_TOKENS ? tokens : undefined;
+}
+
+export function createDynamaxTurnBudget(total: number): DynamaxTurnBudget {
+  let spent = 0;
+  return {
+    total,
+    remaining: () => Math.max(0, total - spent),
+    spend(outputTokens) {
+      if (Number.isFinite(outputTokens) && outputTokens > 0) spent += outputTokens;
+    },
+  };
+}
+
+export type DynamaxTurnBudgetResolution =
+  | { readonly kind: "budget"; readonly budget: number | undefined }
+  | { readonly kind: "spent"; readonly total: number };
+
+/**
+ * A tool run's budget. Only the user sets one: without a turn pool this adds none, and the pool caps
+ * every run in its turn at the remainder, clamped to the run budget range. A pool with less than `min`
+ * left is spent.
+ */
+export function resolveTurnBudget(
+  pool: DynamaxTurnBudget | undefined,
+  range: { readonly min: number; readonly max: number },
+): DynamaxTurnBudgetResolution {
+  if (pool === undefined) return { kind: "budget", budget: undefined };
+  const remaining = pool.remaining();
+  if (remaining < range.min) return { kind: "spent", total: pool.total };
+  return { kind: "budget", budget: Math.min(range.max, Math.floor(remaining)) };
+}
+
 export function markDynamaxOneShot(state: DynamaxState): void {
   state.oneShotPending = true;
 }
@@ -136,12 +200,21 @@ export function dynamaxStatusText(state: DynamaxState, theme: Theme): string {
 /**
  * Add the opt-in section for this run and consume a pending one-shot. pi rebuilds `sections` for every
  * run, so an inactive run simply leaves it out and pi records the removal as a patch.
+ *
+ * Sticky mode covers every run, including one a custom message started (`sendMessage` with
+ * `triggerTurn`, such as a background completion). Only a user prompt consumes the one-shot: the
+ * opt-in belongs to the request the user typed it in.
  */
-export function applyDynamaxPromptSection(sections: Record<string, string>, state: DynamaxState): boolean {
-  if (!state.sticky && !state.oneShotPending) return false;
-  state.oneShotPending = false;
+export function applyDynamaxPromptSection(
+  sections: Record<string, string>,
+  state: DynamaxState,
+  messageType: BeforeAgentStartEvent["messageType"] = "user",
+): boolean {
+  const oneShot = messageType === "user" && state.oneShotPending;
+  if (!state.sticky && !oneShot) return false;
+  if (oneShot) state.oneShotPending = false;
   state.turnActive = true;
-  sections[DYNAMAX_SECTION] = DYNAMAX_REMINDER;
+  sections[DYNAMAX_SECTION] = state.sticky ? DYNAMAX_STICKY_REMINDER : DYNAMAX_REMINDER;
   return true;
 }
 
@@ -230,23 +303,30 @@ export function registerDynamax(pi: ExtensionAPI, shortcuts: DynamaxShortcuts, o
   });
 
   pi.on("input", (event, ctx) => {
-    if (event.source !== "extension" && hasDynamaxToken(event.text)) markOneShot(ctx);
+    if (event.source === "extension") return { action: "continue" };
+    if (hasDynamaxToken(event.text)) markOneShot(ctx);
+    const budget = parseBudgetDirective(event.text);
+    if (budget !== undefined) getDynamaxRuntime(runtimes, ctx).turnBudget = createDynamaxTurnBudget(budget);
     return { action: "continue" };
   });
 
   pi.on("before_agent_start", (event, ctx) => {
-    const runtime = getDynamaxRuntime(runtimes, ctx);
     const options = event.systemPromptOptions;
-    if (!applyDynamaxPromptSection(options.sections, runtime.state)) return undefined;
+    // Subagent sessions load this extension but withhold the workflow tool, so a `dynamax` in a
+    // subagent prompt must neither permit that tool nor consume the one-shot.
+    if (!options.selectedTools.includes(WORKFLOW_TOOL_NAME)) return undefined;
+    const runtime = getDynamaxRuntime(runtimes, ctx);
+    if (!applyDynamaxPromptSection(options.sections, runtime.state, event.messageType)) return undefined;
     updateDynamaxSurfaces(ctx, runtime);
-    // An earlier extension forced the whole prompt, and pi renders no sections then; append to its text instead.
-    if (options.forceSystemPrompt === undefined) return undefined;
-    return { systemPrompt: `${options.forceSystemPrompt}\n\n<${DYNAMAX_SECTION}>\n${DYNAMAX_REMINDER}\n</${DYNAMAX_SECTION}>` };
+    return appendSectionToForcedPrompt(options, DYNAMAX_SECTION);
   });
 
-  pi.on("agent_end", (_event, ctx) => {
+  // Not agent_end: pi emits that per agent loop, and one run continues in fresh loops after a
+  // provider retry, overflow compaction, or queued follow-up. agent_settled fires once per run.
+  pi.on("agent_settled", (_event, ctx) => {
     const runtime = getDynamaxRuntime(runtimes, ctx);
     if (!runtime.state.sticky) runtime.state.turnActive = false;
+    runtime.turnBudget = undefined;
     updateDynamaxSurfaces(ctx, runtime);
   });
 
@@ -260,15 +340,12 @@ export function registerDynamax(pi: ExtensionAPI, shortcuts: DynamaxShortcuts, o
   }
 
   pi.registerCommand("workflow:dynamax", {
-    description: "Toggle Dynamax workflow orchestration: /workflow:dynamax [on|off|status]",
+    description: "Toggle sticky Dynamax orchestration for this session: /workflow:dynamax [on|off|status]",
     getArgumentCompletions: (argumentPrefix) => completeCurrentArgument(argumentPrefix, DYNAMAX_ACTION_COMPLETIONS),
     handler: async (args, ctx) => {
       const runtime = getDynamaxRuntime(runtimes, ctx);
-      const action = args.trim().toLowerCase();
-      if (action === "") {
-        ctx.ui.notify(`Dynamax ${describeDynamaxState(runtime.state)}. Usage: /workflow:dynamax [on|off|status]`, "info");
-        return;
-      }
+      const typed = args.trim().toLowerCase();
+      const action = typed === "" ? (runtime.state.sticky ? "off" : "on") : typed;
       if (action === "on") {
         enableDynamaxSticky(runtime.state);
         updateDynamaxSurfaces(ctx, runtime);
@@ -285,11 +362,15 @@ export function registerDynamax(pi: ExtensionAPI, shortcuts: DynamaxShortcuts, o
         ctx.ui.notify(`Dynamax ${describeDynamaxState(runtime.state)}`, "info");
         return;
       }
-      ctx.ui.notify("Usage: /workflow:dynamax [on|off|status]", "warning");
+      ctx.ui.notify("Usage: /workflow:dynamax [on|off|status] (no argument toggles sticky mode)", "warning");
     },
   });
 
-  return { markOneShot };
+  return {
+    markOneShot,
+    turnBudget: (ctx) => getDynamaxRuntime(runtimes, ctx).turnBudget,
+    isActive: (ctx) => isDynamaxActive(getDynamaxRuntime(runtimes, ctx).state),
+  };
 }
 
 export function updateDynamaxSurfaces(ctx: Pick<ExtensionContext, "hasUI" | "ui">, runtime: DynamaxRuntime): void {