Added nono profiles & orca integration - Added some extensions to start working on
This commit is contained in:
@@ -0,0 +1,145 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { mkdtempSync, readFileSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { test } from "node:test";
|
||||
import {
|
||||
boundedArtifactTranscript,
|
||||
createWorkflowPersistence,
|
||||
persistWorkflowJson,
|
||||
} from "./artifacts.ts";
|
||||
import {
|
||||
emptyUsage,
|
||||
type TranscriptEntry,
|
||||
type WorkflowDetails,
|
||||
} from "./model.ts";
|
||||
|
||||
function workflowDetails(): WorkflowDetails {
|
||||
return {
|
||||
runId: "wf_fixture",
|
||||
sessionId: "session_fixture",
|
||||
background: false,
|
||||
status: "running",
|
||||
startedAt: 1,
|
||||
phases: [],
|
||||
agents: [],
|
||||
};
|
||||
}
|
||||
|
||||
test("artifact transcript keeps the initial prompt, marker, and newest entries", () => {
|
||||
const prompt = `initial:${"p".repeat(70)}`;
|
||||
const transcript = [
|
||||
{ role: "user" as const, text: prompt },
|
||||
...Array.from({ length: 5 }, (_, index) => ({
|
||||
role: "assistant" as const,
|
||||
text: `entry-${index}:${String(index).repeat(70)}`,
|
||||
})),
|
||||
];
|
||||
|
||||
const bounded = boundedArtifactTranscript(transcript, {
|
||||
maxBytes: 256,
|
||||
entryMaxBytes: 80,
|
||||
});
|
||||
|
||||
assert.equal(bounded[0]?.role, "user");
|
||||
assert.equal(bounded[0]?.text, prompt);
|
||||
assert.match(bounded[1]?.text ?? "", /artifact transcript truncated/);
|
||||
assert.equal(bounded.at(-1)?.text, transcript.at(-1)?.text);
|
||||
assert.equal(
|
||||
bounded.some((entry) => entry.text.startsWith("entry-0:")),
|
||||
false,
|
||||
);
|
||||
assert.ok(
|
||||
bounded.reduce(
|
||||
(total, entry) => total + Buffer.byteLength(entry.text, "utf8"),
|
||||
0,
|
||||
) <= 256,
|
||||
);
|
||||
});
|
||||
|
||||
test("live artifact persistence includes current agents and transcripts", () => {
|
||||
const directory = mkdtempSync(join(tmpdir(), "pi-workflow-artifacts-"));
|
||||
try {
|
||||
const details = workflowDetails();
|
||||
details.agents.push({
|
||||
index: 1,
|
||||
label: "running-fixture",
|
||||
state: "running",
|
||||
startedAt: 2,
|
||||
preview: "working",
|
||||
usage: emptyUsage(),
|
||||
transcript: [
|
||||
{ role: "user", text: "current prompt" },
|
||||
{
|
||||
role: "tool",
|
||||
name: "fixture",
|
||||
toolCallId: "call-fixture",
|
||||
text: "{}",
|
||||
startedAt: 10,
|
||||
finishedAt: 25,
|
||||
durationMs: 15,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
persistWorkflowJson(directory, details);
|
||||
|
||||
const workflow = JSON.parse(
|
||||
readFileSync(join(directory, "workflow.json"), "utf8"),
|
||||
) as WorkflowDetails;
|
||||
const transcripts = JSON.parse(
|
||||
readFileSync(join(directory, "transcripts.json"), "utf8"),
|
||||
) as Record<string, TranscriptEntry[]>;
|
||||
assert.equal(workflow.agents.length, 1);
|
||||
assert.equal(workflow.agents[0]?.label, "running-fixture");
|
||||
assert.equal(transcripts["1"]?.[0]?.text, "current prompt");
|
||||
assert.deepEqual(
|
||||
{
|
||||
toolCallId: transcripts["1"]?.[1]?.toolCallId,
|
||||
startedAt: transcripts["1"]?.[1]?.startedAt,
|
||||
finishedAt: transcripts["1"]?.[1]?.finishedAt,
|
||||
durationMs: transcripts["1"]?.[1]?.durationMs,
|
||||
},
|
||||
{
|
||||
toolCallId: "call-fixture",
|
||||
startedAt: 10,
|
||||
finishedAt: 25,
|
||||
durationMs: 15,
|
||||
},
|
||||
);
|
||||
} finally {
|
||||
rmSync(directory, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("workflow checkpoints throttle updates and support immediate/final flushes", async () => {
|
||||
const details = workflowDetails();
|
||||
const snapshots: WorkflowDetails[] = [];
|
||||
const persistence = createWorkflowPersistence("fixture", details, {
|
||||
intervalMs: 15,
|
||||
persist: (_runDir, current) => snapshots.push(structuredClone(current)),
|
||||
});
|
||||
|
||||
details.currentPhase = "Scan";
|
||||
persistence.checkpoint();
|
||||
details.currentPhase = "Review";
|
||||
persistence.checkpoint();
|
||||
assert.equal(snapshots.length, 0);
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(snapshots.length, 1);
|
||||
assert.equal(snapshots[0]?.currentPhase, "Review");
|
||||
|
||||
details.status = "completed";
|
||||
persistence.checkpoint({ immediate: true });
|
||||
assert.equal(snapshots.length, 2);
|
||||
assert.equal(snapshots[1]?.status, "completed");
|
||||
|
||||
details.finishedAt = 3;
|
||||
persistence.flush();
|
||||
assert.equal(snapshots.length, 3);
|
||||
assert.equal(snapshots[2]?.finishedAt, 3);
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
assert.equal(snapshots.length, 3);
|
||||
});
|
||||
@@ -0,0 +1,179 @@
|
||||
import type { TranscriptEntry, WorkflowDetails } from "./model.ts";
|
||||
import {
|
||||
safeStringify,
|
||||
truncateUtf8,
|
||||
writeFileAtomic,
|
||||
} from "./serialization.ts";
|
||||
import * as path from "node:path";
|
||||
|
||||
const ARTIFACT_TRANSCRIPT_MAX_BYTES = 32 * 1024;
|
||||
const ARTIFACT_TRANSCRIPT_ENTRY_MAX_BYTES = 8 * 1024;
|
||||
export const WORKFLOW_CHECKPOINT_INTERVAL_MS = 500;
|
||||
const ENTRY_TRUNCATION_MARKER = "\n[entry truncated]";
|
||||
const TRANSCRIPT_TRUNCATION_MARKER =
|
||||
"[artifact transcript truncated: older entries omitted]";
|
||||
|
||||
function textBytes(text: string) {
|
||||
return Buffer.byteLength(text, "utf8");
|
||||
}
|
||||
|
||||
function boundEntry(entry: TranscriptEntry, maxBytes: number) {
|
||||
if (textBytes(entry.text) <= maxBytes) return { ...entry };
|
||||
const markerBytes = textBytes(ENTRY_TRUNCATION_MARKER);
|
||||
const text =
|
||||
maxBytes > markerBytes
|
||||
? `${truncateUtf8(entry.text, maxBytes - markerBytes)}${ENTRY_TRUNCATION_MARKER}`
|
||||
: truncateUtf8(ENTRY_TRUNCATION_MARKER, maxBytes);
|
||||
return { ...entry, text };
|
||||
}
|
||||
|
||||
/** Keep the initial prompt plus the newest useful context within the artifact cap. */
|
||||
export function boundedArtifactTranscript(
|
||||
transcript: TranscriptEntry[],
|
||||
options: { maxBytes?: number; entryMaxBytes?: number } = {},
|
||||
) {
|
||||
if (transcript.length === 0) return [];
|
||||
const maxBytes = Math.max(
|
||||
256,
|
||||
options.maxBytes ?? ARTIFACT_TRANSCRIPT_MAX_BYTES,
|
||||
);
|
||||
const entryMaxBytes = Math.max(
|
||||
64,
|
||||
Math.min(
|
||||
maxBytes,
|
||||
options.entryMaxBytes ?? ARTIFACT_TRANSCRIPT_ENTRY_MAX_BYTES,
|
||||
),
|
||||
);
|
||||
const bounded = transcript.map((entry) => boundEntry(entry, entryMaxBytes));
|
||||
if (
|
||||
bounded.reduce((total, entry) => total + textBytes(entry.text), 0) <=
|
||||
maxBytes
|
||||
) {
|
||||
return bounded;
|
||||
}
|
||||
|
||||
const initialIndex = transcript.findIndex((entry) => entry.role === "user");
|
||||
const initial = boundEntry(
|
||||
transcript[initialIndex >= 0 ? initialIndex : 0],
|
||||
Math.min(entryMaxBytes, maxBytes - textBytes(TRANSCRIPT_TRUNCATION_MARKER)),
|
||||
);
|
||||
const marker: TranscriptEntry = {
|
||||
role: "toolResult",
|
||||
name: "transcript",
|
||||
text: TRANSCRIPT_TRUNCATION_MARKER,
|
||||
};
|
||||
let remaining = maxBytes - textBytes(initial.text) - textBytes(marker.text);
|
||||
const tail: TranscriptEntry[] = [];
|
||||
|
||||
for (
|
||||
let index = transcript.length - 1;
|
||||
index >= 0 && remaining > 0;
|
||||
index--
|
||||
) {
|
||||
if (index === initialIndex || (initialIndex < 0 && index === 0)) continue;
|
||||
const entry = boundEntry(
|
||||
transcript[index],
|
||||
Math.min(entryMaxBytes, remaining),
|
||||
);
|
||||
tail.push(entry);
|
||||
remaining -= textBytes(entry.text);
|
||||
}
|
||||
|
||||
tail.reverse();
|
||||
return [initial, marker, ...tail];
|
||||
}
|
||||
|
||||
function writeRunFile(runDir: string, name: string, content: string) {
|
||||
writeFileAtomic(path.join(runDir, name), content);
|
||||
}
|
||||
|
||||
export function persistWorkflowJson(runDir: string, details: WorkflowDetails) {
|
||||
const transcripts = Object.fromEntries(
|
||||
details.agents.map((agent) => [
|
||||
agent.index,
|
||||
boundedArtifactTranscript(agent.transcript),
|
||||
]),
|
||||
);
|
||||
writeRunFile(
|
||||
runDir,
|
||||
"transcripts.json",
|
||||
safeStringify(transcripts, { maxBytes: 2 * 1024 * 1024 }),
|
||||
);
|
||||
if (details.result !== undefined) {
|
||||
writeRunFile(
|
||||
runDir,
|
||||
"result.json",
|
||||
safeStringify(details.result, { maxBytes: 1024 * 1024 }),
|
||||
);
|
||||
}
|
||||
const compact: WorkflowDetails = {
|
||||
...details,
|
||||
...(details.result !== undefined
|
||||
? { result: "[stored in result.json]", resultArtifact: "result.json" }
|
||||
: {}),
|
||||
transcriptArtifact: "transcripts.json",
|
||||
agents: details.agents.map((agent) => ({ ...agent, transcript: [] })),
|
||||
};
|
||||
writeRunFile(
|
||||
runDir,
|
||||
"workflow.json",
|
||||
safeStringify(compact, { maxBytes: 1024 * 1024 }),
|
||||
);
|
||||
}
|
||||
|
||||
/** Coalesce live checkpoints while keeping final persistence synchronous. */
|
||||
export function createWorkflowPersistence(
|
||||
runDir: string,
|
||||
details: WorkflowDetails,
|
||||
options: {
|
||||
intervalMs?: number;
|
||||
persist?: (runDir: string, details: WorkflowDetails) => void;
|
||||
} = {},
|
||||
) {
|
||||
const intervalMs = Math.max(
|
||||
0,
|
||||
options.intervalMs ?? WORKFLOW_CHECKPOINT_INTERVAL_MS,
|
||||
);
|
||||
const persist = options.persist ?? persistWorkflowJson;
|
||||
let lastPersistedAt = Date.now();
|
||||
let dirty = false;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
|
||||
const savePending = () => {
|
||||
timer = undefined;
|
||||
if (!dirty) return;
|
||||
try {
|
||||
persist(runDir, details);
|
||||
dirty = false;
|
||||
lastPersistedAt = Date.now();
|
||||
} catch {
|
||||
// Final flush retries and reports persistence failures synchronously.
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
checkpoint(options: { immediate?: boolean } = {}) {
|
||||
dirty = true;
|
||||
if (options.immediate) {
|
||||
if (timer) clearTimeout(timer);
|
||||
timer = undefined;
|
||||
savePending();
|
||||
return;
|
||||
}
|
||||
if (timer) return;
|
||||
const delay = Math.max(0, intervalMs - (Date.now() - lastPersistedAt));
|
||||
if (delay === 0) {
|
||||
savePending();
|
||||
return;
|
||||
}
|
||||
timer = setTimeout(savePending, delay);
|
||||
},
|
||||
flush() {
|
||||
if (timer) clearTimeout(timer);
|
||||
timer = undefined;
|
||||
persist(runDir, details);
|
||||
dirty = false;
|
||||
lastPersistedAt = Date.now();
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { test } from "node:test";
|
||||
import { MAX_AGENT_CALLS, RunController } from "./controller.ts";
|
||||
|
||||
const delay = (milliseconds: number) =>
|
||||
new Promise<void>((resolve) => setTimeout(resolve, milliseconds));
|
||||
|
||||
test("RunController reserves calls synchronously and caps global fanout", async () => {
|
||||
const controller = new RunController(undefined, 4);
|
||||
let active = 0;
|
||||
let peak = 0;
|
||||
const tasks = Array.from({ length: 12 }, (_, index) =>
|
||||
controller.schedule(async () => {
|
||||
active++;
|
||||
peak = Math.max(peak, active);
|
||||
await delay(5);
|
||||
active--;
|
||||
return index;
|
||||
}),
|
||||
);
|
||||
assert.deepEqual(
|
||||
await Promise.all(tasks),
|
||||
Array.from({ length: 12 }, (_, i) => i),
|
||||
);
|
||||
assert.equal(peak, 4);
|
||||
assert.equal(await controller.settle(), true);
|
||||
});
|
||||
|
||||
test("RunController propagates invocation cancellation without aborting the run", async () => {
|
||||
const controller = new RunController(undefined, 1);
|
||||
const invocation = new AbortController();
|
||||
const pending = controller.schedule(
|
||||
(signal) =>
|
||||
new Promise<string>((resolve) => {
|
||||
signal.addEventListener("abort", () => resolve("stopped"), {
|
||||
once: true,
|
||||
});
|
||||
}),
|
||||
invocation.signal,
|
||||
);
|
||||
|
||||
invocation.abort(new Error("Workflow agent request was cancelled"));
|
||||
await assert.rejects(pending, /request was cancelled/);
|
||||
assert.equal(controller.signal.aborted, false);
|
||||
assert.equal(await controller.schedule(async () => "recovered"), "recovered");
|
||||
assert.equal(await controller.settle(), true);
|
||||
});
|
||||
|
||||
test("RunController enforces call budget and aborts queued tasks", async () => {
|
||||
const controller = new RunController(undefined, 1);
|
||||
const blocker = controller.schedule(
|
||||
(signal) =>
|
||||
new Promise<void>((resolve) =>
|
||||
signal.addEventListener("abort", () => resolve(), { once: true }),
|
||||
),
|
||||
);
|
||||
const queued = Array.from({ length: MAX_AGENT_CALLS - 1 }, () =>
|
||||
controller.schedule(async () => "queued"),
|
||||
);
|
||||
await assert.rejects(
|
||||
controller.schedule(async () => "too many"),
|
||||
/exceeded the limit/,
|
||||
);
|
||||
controller.abort();
|
||||
await blocker;
|
||||
const results = await Promise.allSettled(queued);
|
||||
assert.ok(results.every((result) => result.status === "rejected"));
|
||||
assert.equal(await controller.settle({ abort: true }), true);
|
||||
});
|
||||
@@ -0,0 +1,192 @@
|
||||
const DEFAULT_CONCURRENCY = 4;
|
||||
export const MAX_AGENT_CALLS = 32;
|
||||
export const RUN_SHUTDOWN_TIMEOUT_MS = 8_000;
|
||||
|
||||
function abortError(signal: AbortSignal) {
|
||||
return signal.reason instanceof Error
|
||||
? signal.reason
|
||||
: new Error("Workflow was aborted");
|
||||
}
|
||||
|
||||
class Semaphore {
|
||||
private active = 0;
|
||||
private readonly limit: number;
|
||||
private queue: Array<{
|
||||
resolve: () => void;
|
||||
reject: (error: Error) => void;
|
||||
signal: AbortSignal;
|
||||
onAbort: () => void;
|
||||
}> = [];
|
||||
|
||||
constructor(limit: number) {
|
||||
this.limit = limit;
|
||||
}
|
||||
|
||||
acquire(signal: AbortSignal) {
|
||||
if (signal.aborted) return Promise.reject(abortError(signal));
|
||||
if (this.active < this.limit) {
|
||||
this.active++;
|
||||
return Promise.resolve();
|
||||
}
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
const waiter = {
|
||||
resolve: () => {
|
||||
signal.removeEventListener("abort", onAbort);
|
||||
this.active++;
|
||||
resolve();
|
||||
},
|
||||
reject,
|
||||
signal,
|
||||
onAbort: () => {},
|
||||
};
|
||||
const onAbort = () => {
|
||||
const index = this.queue.indexOf(waiter);
|
||||
if (index >= 0) this.queue.splice(index, 1);
|
||||
reject(abortError(signal));
|
||||
};
|
||||
waiter.onAbort = onAbort;
|
||||
this.queue.push(waiter);
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
});
|
||||
}
|
||||
|
||||
release() {
|
||||
this.active = Math.max(0, this.active - 1);
|
||||
while (this.queue.length > 0) {
|
||||
const waiter = this.queue.shift()!;
|
||||
if (waiter.signal.aborted) {
|
||||
waiter.signal.removeEventListener("abort", waiter.onAbort);
|
||||
waiter.reject(abortError(waiter.signal));
|
||||
continue;
|
||||
}
|
||||
waiter.resolve();
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
clear() {
|
||||
const queued = this.queue;
|
||||
this.queue = [];
|
||||
for (const waiter of queued) {
|
||||
waiter.signal.removeEventListener("abort", waiter.onAbort);
|
||||
waiter.reject(abortError(waiter.signal));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Owns every agent task and the run-wide fanout/abort budget. */
|
||||
export class RunController {
|
||||
private readonly abortController = new AbortController();
|
||||
private readonly semaphore: Semaphore;
|
||||
private readonly tasks = new Set<Promise<unknown>>();
|
||||
private callCount = 0;
|
||||
private sealed = false;
|
||||
private parentAbort?: () => void;
|
||||
private parentSignal?: AbortSignal;
|
||||
|
||||
constructor(parentSignal?: AbortSignal, concurrency = DEFAULT_CONCURRENCY) {
|
||||
this.semaphore = new Semaphore(
|
||||
Math.max(1, Math.min(DEFAULT_CONCURRENCY, Math.floor(concurrency))),
|
||||
);
|
||||
if (parentSignal) {
|
||||
this.parentSignal = parentSignal;
|
||||
this.parentAbort = () => this.abort("Parent operation was aborted");
|
||||
if (parentSignal.aborted) this.parentAbort();
|
||||
else
|
||||
parentSignal.addEventListener("abort", this.parentAbort, {
|
||||
once: true,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
get signal() {
|
||||
return this.abortController.signal;
|
||||
}
|
||||
|
||||
get calls() {
|
||||
return this.callCount;
|
||||
}
|
||||
|
||||
schedule<T>(
|
||||
task: (signal: AbortSignal) => Promise<T>,
|
||||
invocationSignal?: AbortSignal,
|
||||
): Promise<T> {
|
||||
if (this.sealed) return Promise.reject(new Error("Workflow is settling"));
|
||||
if (this.signal.aborted) return Promise.reject(abortError(this.signal));
|
||||
if (this.callCount >= MAX_AGENT_CALLS) {
|
||||
return Promise.reject(
|
||||
new Error(
|
||||
`Workflow exceeded the limit of ${MAX_AGENT_CALLS} agent calls`,
|
||||
),
|
||||
);
|
||||
}
|
||||
this.callCount++;
|
||||
|
||||
const running = (async () => {
|
||||
const taskAbort = new AbortController();
|
||||
const onRunAbort = () => taskAbort.abort(this.signal.reason);
|
||||
const onInvocationAbort = () => taskAbort.abort(invocationSignal?.reason);
|
||||
this.signal.addEventListener("abort", onRunAbort, { once: true });
|
||||
invocationSignal?.addEventListener("abort", onInvocationAbort, {
|
||||
once: true,
|
||||
});
|
||||
if (this.signal.aborted) onRunAbort();
|
||||
else if (invocationSignal?.aborted) onInvocationAbort();
|
||||
|
||||
let acquired = false;
|
||||
try {
|
||||
await this.semaphore.acquire(taskAbort.signal);
|
||||
acquired = true;
|
||||
if (taskAbort.signal.aborted) throw abortError(taskAbort.signal);
|
||||
const result = await task(taskAbort.signal);
|
||||
if (invocationSignal?.aborted) throw abortError(invocationSignal);
|
||||
return result;
|
||||
} finally {
|
||||
this.signal.removeEventListener("abort", onRunAbort);
|
||||
invocationSignal?.removeEventListener("abort", onInvocationAbort);
|
||||
if (acquired) this.semaphore.release();
|
||||
}
|
||||
})();
|
||||
this.tasks.add(running);
|
||||
void running.finally(() => this.tasks.delete(running)).catch(() => {});
|
||||
return running;
|
||||
}
|
||||
|
||||
abort(reason = "Workflow was aborted") {
|
||||
if (!this.signal.aborted) this.abortController.abort(new Error(reason));
|
||||
this.semaphore.clear();
|
||||
}
|
||||
|
||||
/** Seal the task registry and wait a bounded time for every task to settle. */
|
||||
async settle(options: { abort?: boolean; timeoutMs?: number } = {}) {
|
||||
this.sealed = true;
|
||||
if (options.abort) this.abort();
|
||||
const tasks = [...this.tasks];
|
||||
if (tasks.length === 0) {
|
||||
this.detachParent();
|
||||
return true;
|
||||
}
|
||||
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const timeout = new Promise<false>((resolve) => {
|
||||
timer = setTimeout(
|
||||
() => resolve(false),
|
||||
options.timeoutMs ?? RUN_SHUTDOWN_TIMEOUT_MS,
|
||||
);
|
||||
timer.unref?.();
|
||||
});
|
||||
const settled = Promise.allSettled(tasks).then(() => true as const);
|
||||
const completed = await Promise.race([settled, timeout]);
|
||||
if (timer) clearTimeout(timer);
|
||||
this.detachParent();
|
||||
return completed;
|
||||
}
|
||||
|
||||
private detachParent() {
|
||||
if (this.parentAbort) {
|
||||
this.parentSignal?.removeEventListener("abort", this.parentAbort);
|
||||
}
|
||||
this.parentAbort = undefined;
|
||||
this.parentSignal = undefined;
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,871 @@
|
||||
/**
|
||||
* workflows: model-authored multi-agent orchestration.
|
||||
*
|
||||
* A `workflow` tool that runs a JavaScript orchestration script written inline
|
||||
* by the model. The script executes ordered phases, fanning work out to
|
||||
* isolated subagents:
|
||||
*
|
||||
* export const meta = { name, description, phases: [{ title, detail? }] }
|
||||
* phase(title) // mark runtime phase progression
|
||||
* await agent(prompt, { label?, phase?, schema?, model?, provider?, effort? })
|
||||
* await parallel([() => agent(...), ...], { concurrency? })
|
||||
* args // parsed JSON args passed with the tool call
|
||||
*
|
||||
* `agent()` always resolves to `{ ok, output, structured?, error? }` — it
|
||||
* never throws into the script. Scripts branch on `ok` explicitly.
|
||||
*
|
||||
* Runs are blocking by default (live progress in the tool block). Pass
|
||||
* `background: true` to return immediately and get a follow-up message when
|
||||
* the run finishes. Run artifacts (script, args, statuses, result) are saved
|
||||
* under `~/.pi/agent/workflows/<runId>/` for inspection; result and bounded
|
||||
* transcripts use separate artifacts, and there is no resume.
|
||||
*/
|
||||
|
||||
import { randomBytes } from "node:crypto";
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
import {
|
||||
getAgentDir,
|
||||
getMarkdownTheme,
|
||||
keyHint,
|
||||
type ExtensionAPI,
|
||||
type ExtensionContext,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import { Container, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
|
||||
import { Type, type Static } from "typebox";
|
||||
import { formatActivityStatus } from "../shared/activity-status.ts";
|
||||
import { createWorkflowPersistence, persistWorkflowJson } from "./artifacts.ts";
|
||||
import { RunController } from "./controller.ts";
|
||||
import { sessionWorkflowRunIds, showWorkflowDashboard } from "./dashboard.ts";
|
||||
import {
|
||||
extractMeta,
|
||||
prepareWorkflowScript,
|
||||
type WorkflowMeta,
|
||||
} from "./meta.ts";
|
||||
import {
|
||||
agentContext,
|
||||
aggregateUsage,
|
||||
countStates,
|
||||
emptyUsage,
|
||||
formatElapsed,
|
||||
formatUsage,
|
||||
phaseGroups,
|
||||
resultJson,
|
||||
stateSquare,
|
||||
statusColor,
|
||||
statusWord,
|
||||
SQUARE,
|
||||
type AgentRecord,
|
||||
type WorkflowDetails,
|
||||
} from "./model.ts";
|
||||
import {
|
||||
buildBackgroundWorkflowFollowUp,
|
||||
buildBackgroundWorkflowLaunchResult,
|
||||
buildWorkflowAgentPrompt,
|
||||
buildWorkflowResultMessage,
|
||||
WORKFLOW_PARAMETER_DESCRIPTIONS,
|
||||
WORKFLOW_PROMPT_GUIDELINES,
|
||||
WORKFLOW_PROMPT_SNIPPET,
|
||||
WORKFLOW_TOOL_DESCRIPTION,
|
||||
} from "./prompt.ts";
|
||||
import {
|
||||
createWorkflowResources,
|
||||
runAgent,
|
||||
type ThinkingLevel,
|
||||
type WorkflowModel,
|
||||
} from "./runner.ts";
|
||||
import { runWorkflowSandbox } from "./sandbox.ts";
|
||||
import { safeStringify, writeFileAtomic } from "./serialization.ts";
|
||||
|
||||
const PREVIEW_LENGTH = 200;
|
||||
const EMIT_INTERVAL_MS = 120;
|
||||
|
||||
const THINKING_LEVELS = [
|
||||
"off",
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max",
|
||||
] as const;
|
||||
|
||||
/** What `agent()` resolves to inside the script. */
|
||||
interface ScriptAgentResult {
|
||||
ok: boolean;
|
||||
output: string;
|
||||
structured?: unknown;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
interface AgentCallOptions {
|
||||
label?: unknown;
|
||||
phase?: unknown;
|
||||
schema?: unknown;
|
||||
model?: unknown;
|
||||
provider?: unknown;
|
||||
effort?: unknown;
|
||||
}
|
||||
|
||||
const WorkflowParams = Type.Object({
|
||||
script: Type.String({
|
||||
description: WORKFLOW_PARAMETER_DESCRIPTIONS.script,
|
||||
}),
|
||||
args: Type.Optional(
|
||||
Type.String({
|
||||
description: WORKFLOW_PARAMETER_DESCRIPTIONS.args,
|
||||
}),
|
||||
),
|
||||
background: Type.Optional(
|
||||
Type.Boolean({
|
||||
description: WORKFLOW_PARAMETER_DESCRIPTIONS.background,
|
||||
}),
|
||||
),
|
||||
});
|
||||
|
||||
type WorkflowInput = Static<typeof WorkflowParams>;
|
||||
|
||||
function errorText(error: unknown): string {
|
||||
return (error instanceof Error ? error.message : String(error)).slice(
|
||||
0,
|
||||
16 * 1024,
|
||||
);
|
||||
}
|
||||
|
||||
function summaryLine(details: WorkflowDetails): string {
|
||||
const { done, failed } = countStates(details);
|
||||
const settled = done + failed;
|
||||
return `workflow ${details.name ?? details.runId}: ${settled}/${details.agents.length} agents${
|
||||
details.currentPhase ? ` · ${details.currentPhase}` : ""
|
||||
}`;
|
||||
}
|
||||
|
||||
function writeRunFile(runDir: string, name: string, content: string) {
|
||||
writeFileAtomic(path.join(runDir, name), content);
|
||||
}
|
||||
|
||||
function compactToolDetails(details: WorkflowDetails): WorkflowDetails {
|
||||
return {
|
||||
...details,
|
||||
...(details.result !== undefined
|
||||
? {
|
||||
result: JSON.parse(
|
||||
safeStringify(details.result, { maxBytes: 64 * 1024 }),
|
||||
),
|
||||
}
|
||||
: {}),
|
||||
agents: details.agents.map((agent) => ({ ...agent, transcript: [] })),
|
||||
};
|
||||
}
|
||||
|
||||
interface RunSummary {
|
||||
runId: string;
|
||||
name?: string;
|
||||
status: string;
|
||||
done: number;
|
||||
total: number;
|
||||
startedAt: number;
|
||||
active: boolean;
|
||||
}
|
||||
|
||||
function listRuns(
|
||||
activeRuns: Map<string, WorkflowDetails>,
|
||||
sessionId: string,
|
||||
referencedRunIds: ReadonlySet<string>,
|
||||
): RunSummary[] {
|
||||
const base = path.join(getAgentDir(), "workflows");
|
||||
let names: string[] = [];
|
||||
try {
|
||||
names = fs.readdirSync(base).filter((name) => name.startsWith("wf_"));
|
||||
} catch {
|
||||
// No runs yet.
|
||||
}
|
||||
const summaries: RunSummary[] = [];
|
||||
for (const runId of names) {
|
||||
const live = activeRuns.get(runId);
|
||||
if (live) {
|
||||
const { done, failed } = countStates(live);
|
||||
summaries.push({
|
||||
runId,
|
||||
name: live.name,
|
||||
status: live.status,
|
||||
done: done + failed,
|
||||
total: live.agents.length,
|
||||
startedAt: live.startedAt,
|
||||
active: true,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
const parsed = JSON.parse(
|
||||
fs.readFileSync(path.join(base, runId, "workflow.json"), "utf8"),
|
||||
) as Partial<WorkflowDetails>;
|
||||
if (parsed.sessionId !== sessionId && !referencedRunIds.has(runId)) {
|
||||
continue;
|
||||
}
|
||||
const agents = parsed.agents ?? [];
|
||||
summaries.push({
|
||||
runId,
|
||||
name: parsed.name,
|
||||
status:
|
||||
parsed.status === "running"
|
||||
? "aborted"
|
||||
: (parsed.status ?? "unknown"),
|
||||
done: agents.filter((agent) => agent.state !== "running").length,
|
||||
total: agents.length,
|
||||
startedAt: parsed.startedAt ?? 0,
|
||||
active: false,
|
||||
});
|
||||
} catch {
|
||||
// Ignore unreadable artifacts because their session cannot be verified.
|
||||
}
|
||||
}
|
||||
return summaries.sort((a, b) => b.startedAt - a.startedAt);
|
||||
}
|
||||
|
||||
function runDetailText(
|
||||
run: RunSummary,
|
||||
activeRuns: Map<string, WorkflowDetails>,
|
||||
): string {
|
||||
const runDir = path.join(getAgentDir(), "workflows", run.runId);
|
||||
const live = activeRuns.get(run.runId);
|
||||
if (live) return buildWorkflowResultMessage(live, runDir);
|
||||
try {
|
||||
const parsed = JSON.parse(
|
||||
fs.readFileSync(path.join(runDir, "workflow.json"), "utf8"),
|
||||
) as WorkflowDetails;
|
||||
return buildWorkflowResultMessage(parsed, runDir);
|
||||
} catch {
|
||||
return `Run ${run.runId} — ${run.status}`;
|
||||
}
|
||||
}
|
||||
|
||||
export default function workflows(pi: ExtensionAPI) {
|
||||
/** Live background runs, for /workflows and shutdown cleanup. */
|
||||
const activeRuns = new Map<
|
||||
string,
|
||||
{
|
||||
details: WorkflowDetails;
|
||||
controller: RunController;
|
||||
completion?: Promise<void>;
|
||||
}
|
||||
>();
|
||||
const activeDetails = () =>
|
||||
new Map(
|
||||
[...activeRuns].map(([runId, run]) => [runId, run.details] as const),
|
||||
);
|
||||
|
||||
/** Finished counts remain visible until the dashboard acknowledges them. */
|
||||
let lastUi: ExtensionContext["ui"] | undefined;
|
||||
let completedRuns = 0;
|
||||
let failedRuns = 0;
|
||||
const updateIndicator = () => {
|
||||
const ui = lastUi;
|
||||
if (!ui) return;
|
||||
try {
|
||||
const running = activeRuns.size;
|
||||
if (running === 0 && completedRuns === 0 && failedRuns === 0) {
|
||||
ui.setStatus("workflows", undefined);
|
||||
return;
|
||||
}
|
||||
ui.setStatus(
|
||||
"workflows",
|
||||
formatActivityStatus(ui.theme, "workflows", {
|
||||
running,
|
||||
done: completedRuns,
|
||||
failed: failedRuns,
|
||||
}),
|
||||
);
|
||||
} catch {
|
||||
// UI may be unavailable.
|
||||
}
|
||||
};
|
||||
|
||||
const recordSettledRun = (status: WorkflowDetails["status"]) => {
|
||||
if (status === "completed") completedRuns += 1;
|
||||
else failedRuns += 1;
|
||||
};
|
||||
|
||||
pi.on("session_start", (_event, ctx) => {
|
||||
if (ctx.hasUI) lastUi = ctx.ui;
|
||||
updateIndicator();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async () => {
|
||||
const runs = [...activeRuns.values()];
|
||||
for (const run of runs) run.controller.abort("Session is shutting down");
|
||||
await Promise.all(
|
||||
runs.map((run) => run.controller.settle({ abort: true })),
|
||||
);
|
||||
const completions = runs
|
||||
.map((run) => run.completion)
|
||||
.filter(
|
||||
(completion): completion is Promise<void> => completion !== undefined,
|
||||
);
|
||||
if (completions.length > 0) {
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const timeout = new Promise<void>((resolve) => {
|
||||
timer = setTimeout(resolve, 8_000);
|
||||
timer.unref?.();
|
||||
});
|
||||
await Promise.race([Promise.allSettled(completions), timeout]);
|
||||
if (timer) clearTimeout(timer);
|
||||
}
|
||||
lastUi?.setStatus("workflows", undefined);
|
||||
lastUi = undefined;
|
||||
});
|
||||
|
||||
pi.registerCommand("workflows", {
|
||||
description:
|
||||
"List workflow runs (`/workflows <runId>` for one run's detail)",
|
||||
handler: async (rawArgs, ctx) => {
|
||||
const arg = rawArgs.trim();
|
||||
if (ctx.mode === "tui") {
|
||||
lastUi = ctx.ui;
|
||||
await showWorkflowDashboard(ctx, activeDetails, arg || undefined);
|
||||
// Opening the dashboard acknowledges finished runs.
|
||||
completedRuns = 0;
|
||||
failedRuns = 0;
|
||||
updateIndicator();
|
||||
return;
|
||||
}
|
||||
// Non-TUI fallback: plain text listing.
|
||||
const runs = listRuns(
|
||||
activeDetails(),
|
||||
ctx.sessionManager.getSessionId(),
|
||||
sessionWorkflowRunIds(ctx),
|
||||
);
|
||||
if (runs.length === 0) {
|
||||
ctx.ui.notify("No workflow runs yet.", "info");
|
||||
return;
|
||||
}
|
||||
if (arg) {
|
||||
const run = runs.find((r) => r.runId === arg || r.runId.endsWith(arg));
|
||||
ctx.ui.notify(
|
||||
run
|
||||
? runDetailText(run, activeDetails())
|
||||
: `No workflow run matching "${arg}".`,
|
||||
run ? "info" : "warning",
|
||||
);
|
||||
return;
|
||||
}
|
||||
const labels = runs.map(
|
||||
(r) =>
|
||||
`${r.active ? "* " : " "}${r.runId} ${r.status} ${r.name ?? ""} ${r.done}/${r.total}`,
|
||||
);
|
||||
if (!ctx.hasUI) {
|
||||
ctx.ui.notify(labels.join("\n"), "info");
|
||||
return;
|
||||
}
|
||||
const choice = await ctx.ui.select("Workflow runs", labels);
|
||||
if (!choice) return;
|
||||
const run = runs[labels.indexOf(choice)];
|
||||
if (run) ctx.ui.notify(runDetailText(run, activeDetails()), "info");
|
||||
},
|
||||
});
|
||||
|
||||
pi.registerTool({
|
||||
name: "workflow",
|
||||
label: "Workflow",
|
||||
description: WORKFLOW_TOOL_DESCRIPTION,
|
||||
promptSnippet: WORKFLOW_PROMPT_SNIPPET,
|
||||
promptGuidelines: WORKFLOW_PROMPT_GUIDELINES,
|
||||
parameters: WorkflowParams,
|
||||
|
||||
async execute(_toolCallId, params, signal, onUpdate, ctx) {
|
||||
let prepared: ReturnType<typeof prepareWorkflowScript>;
|
||||
try {
|
||||
prepared = prepareWorkflowScript(params.script);
|
||||
} catch (error) {
|
||||
throw new Error(`Workflow script failed to parse: ${errorText(error)}`);
|
||||
}
|
||||
|
||||
let args: unknown;
|
||||
if (params.args !== undefined) {
|
||||
try {
|
||||
args = JSON.parse(params.args);
|
||||
} catch {
|
||||
args = params.args;
|
||||
}
|
||||
}
|
||||
|
||||
const meta = prepared.meta;
|
||||
const runId = `wf_${randomBytes(6).toString("hex")}`;
|
||||
const runDir = path.join(getAgentDir(), "workflows", runId);
|
||||
const background = (params.background ?? false) && ctx.hasUI;
|
||||
|
||||
const details: WorkflowDetails = {
|
||||
runId,
|
||||
sessionId: ctx.sessionManager.getSessionId(),
|
||||
name: meta.name,
|
||||
description: meta.description,
|
||||
background,
|
||||
status: "running",
|
||||
startedAt: Date.now(),
|
||||
phases: [...meta.phases],
|
||||
agents: [],
|
||||
};
|
||||
|
||||
writeRunFile(runDir, "script.js", params.script);
|
||||
if (params.args !== undefined)
|
||||
writeRunFile(runDir, "args.json", params.args);
|
||||
persistWorkflowJson(runDir, details);
|
||||
const persistence = createWorkflowPersistence(runDir, details);
|
||||
|
||||
// Background runs survive Esc on the parent turn, but all runs are
|
||||
// aborted and settled during session shutdown.
|
||||
const controller = new RunController(background ? undefined : signal);
|
||||
|
||||
// Each concurrent child gets its own extension runtime. All children use
|
||||
// the parent cwd and live trust decision.
|
||||
const projectTrusted = ctx.isProjectTrusted();
|
||||
const getResources = (structured: boolean) =>
|
||||
createWorkflowResources(
|
||||
ctx.cwd,
|
||||
structured ? "structured" : "plain",
|
||||
projectTrusted,
|
||||
);
|
||||
|
||||
// Throttled progress: tool-block updates when blocking. Background
|
||||
// runs are covered by the below-editor indicator and /workflows.
|
||||
let emitTimer: ReturnType<typeof setTimeout> | undefined;
|
||||
let lastEmit = 0;
|
||||
const flush = () => {
|
||||
emitTimer = undefined;
|
||||
lastEmit = Date.now();
|
||||
if (background) return;
|
||||
onUpdate?.({
|
||||
content: [{ type: "text", text: summaryLine(details) }],
|
||||
details: compactToolDetails(details),
|
||||
});
|
||||
};
|
||||
const emit = (checkpoint = true) => {
|
||||
if (checkpoint) persistence.checkpoint();
|
||||
if (emitTimer) return;
|
||||
emitTimer = setTimeout(
|
||||
flush,
|
||||
Math.max(0, EMIT_INTERVAL_MS - (Date.now() - lastEmit)),
|
||||
);
|
||||
};
|
||||
const flushNow = () => {
|
||||
if (emitTimer) clearTimeout(emitTimer);
|
||||
flush();
|
||||
};
|
||||
|
||||
const phaseFn = (title: unknown) => {
|
||||
const text = String(title);
|
||||
details.currentPhase = text;
|
||||
if (!details.phases.some((p) => p.title === text))
|
||||
details.phases.push({ title: text });
|
||||
emit();
|
||||
};
|
||||
|
||||
let agentCounter = 0;
|
||||
const agentFn = async (
|
||||
promptValue: unknown,
|
||||
optsValue: unknown = {},
|
||||
invocationSignal?: AbortSignal,
|
||||
): Promise<ScriptAgentResult> => {
|
||||
const index = ++agentCounter;
|
||||
const opts: AgentCallOptions =
|
||||
optsValue && typeof optsValue === "object"
|
||||
? (optsValue as AgentCallOptions)
|
||||
: {};
|
||||
const label =
|
||||
typeof opts.label === "string" && opts.label.trim()
|
||||
? opts.label.trim().slice(0, 160)
|
||||
: `agent-${index}`;
|
||||
|
||||
const record: AgentRecord = {
|
||||
index,
|
||||
label,
|
||||
phase:
|
||||
typeof opts.phase === "string"
|
||||
? opts.phase.slice(0, 160)
|
||||
: details.currentPhase,
|
||||
state: "running",
|
||||
model: ctx.model?.id,
|
||||
contextWindow: ctx.model?.contextWindow,
|
||||
startedAt: Date.now(),
|
||||
preview: "",
|
||||
usage: emptyUsage(),
|
||||
transcript: [],
|
||||
};
|
||||
details.agents.push(record);
|
||||
persistence.checkpoint({ immediate: true });
|
||||
emit(false);
|
||||
|
||||
const fail = (error: string): ScriptAgentResult => {
|
||||
record.state = "error";
|
||||
record.error = error;
|
||||
record.finishedAt = Date.now();
|
||||
emit();
|
||||
return { ok: false, output: "", error };
|
||||
};
|
||||
|
||||
const prompt = buildWorkflowAgentPrompt(
|
||||
typeof promptValue === "string"
|
||||
? promptValue
|
||||
: String(promptValue ?? ""),
|
||||
);
|
||||
if (!prompt.trim())
|
||||
return fail("agent() requires a non-empty prompt string");
|
||||
if (controller.signal.aborted)
|
||||
return fail("Workflow was aborted before this agent started");
|
||||
|
||||
return controller
|
||||
.schedule(async (runSignal) => {
|
||||
// Model/provider resolution: default to the parent session's model.
|
||||
let model: WorkflowModel | undefined = ctx.model;
|
||||
if (opts.model !== undefined || opts.provider !== undefined) {
|
||||
const modelOpt =
|
||||
typeof opts.model === "string" ? opts.model : undefined;
|
||||
const providerOpt =
|
||||
typeof opts.provider === "string" ? opts.provider : undefined;
|
||||
if (!modelOpt)
|
||||
return fail(
|
||||
`agent "${label}": \`provider\` requires \`model\` as well`,
|
||||
);
|
||||
let resolved: WorkflowModel | undefined;
|
||||
if (providerOpt) {
|
||||
resolved = ctx.modelRegistry.find(providerOpt, modelOpt);
|
||||
} else {
|
||||
const slash = modelOpt.indexOf("/");
|
||||
if (slash > 0) {
|
||||
resolved = ctx.modelRegistry.find(
|
||||
modelOpt.slice(0, slash),
|
||||
modelOpt.slice(slash + 1),
|
||||
);
|
||||
}
|
||||
resolved ??= ctx.modelRegistry
|
||||
.getAll()
|
||||
.find((m) => m.id === modelOpt);
|
||||
}
|
||||
if (!resolved) {
|
||||
const requested = providerOpt
|
||||
? `${providerOpt}/${modelOpt}`
|
||||
: modelOpt;
|
||||
return fail(
|
||||
`agent "${label}": unknown model "${requested}" (use provider/id)`,
|
||||
);
|
||||
}
|
||||
model = resolved;
|
||||
}
|
||||
record.model = model?.id;
|
||||
record.contextWindow = model?.contextWindow;
|
||||
emit();
|
||||
|
||||
// Effort → thinking level; default inherits the parent session.
|
||||
let thinkingLevel: ThinkingLevel = pi.getThinkingLevel();
|
||||
if (opts.effort !== undefined) {
|
||||
const effort = String(opts.effort);
|
||||
if (!(THINKING_LEVELS as readonly string[]).includes(effort)) {
|
||||
return fail(
|
||||
`agent "${label}": invalid effort "${effort}" (use ${THINKING_LEVELS.join("|")})`,
|
||||
);
|
||||
}
|
||||
thinkingLevel = effort as ThinkingLevel;
|
||||
}
|
||||
|
||||
const resources = await getResources(opts.schema !== undefined);
|
||||
const outcome = await runAgent({
|
||||
prompt,
|
||||
schema: opts.schema,
|
||||
model,
|
||||
thinkingLevel,
|
||||
cwd: ctx.cwd,
|
||||
loader: resources.loader,
|
||||
settingsManager: resources.settingsManager,
|
||||
modelRegistry: ctx.modelRegistry,
|
||||
signal: runSignal,
|
||||
onProgress: (progress) => {
|
||||
record.preview = progress.preview.slice(0, PREVIEW_LENGTH);
|
||||
record.usage = progress.usage;
|
||||
record.model = progress.model ?? record.model;
|
||||
record.contextWindow =
|
||||
progress.contextWindow ?? record.contextWindow;
|
||||
record.transcript = progress.transcript;
|
||||
emit();
|
||||
},
|
||||
});
|
||||
|
||||
record.usage = outcome.usage;
|
||||
record.model = outcome.model ?? record.model;
|
||||
record.contextWindow =
|
||||
outcome.contextWindow ?? record.contextWindow;
|
||||
record.transcript = outcome.transcript;
|
||||
record.preview = (outcome.output || record.preview).slice(
|
||||
0,
|
||||
PREVIEW_LENGTH,
|
||||
);
|
||||
record.finishedAt = Date.now();
|
||||
record.state = outcome.ok ? "done" : "error";
|
||||
if (outcome.ok) {
|
||||
delete record.error;
|
||||
} else {
|
||||
record.error = outcome.error ?? "Agent failed";
|
||||
}
|
||||
emit();
|
||||
|
||||
return {
|
||||
ok: outcome.ok,
|
||||
output: outcome.output,
|
||||
...(outcome.structured !== undefined
|
||||
? { structured: outcome.structured }
|
||||
: {}),
|
||||
...(outcome.error !== undefined ? { error: outcome.error } : {}),
|
||||
};
|
||||
}, invocationSignal)
|
||||
.catch((error) => fail(errorText(error)));
|
||||
};
|
||||
|
||||
const runScript = async () => {
|
||||
let status: WorkflowDetails["status"] = "completed";
|
||||
try {
|
||||
details.result = await runWorkflowSandbox({
|
||||
source: prepared.source,
|
||||
args,
|
||||
cwd: ctx.cwd,
|
||||
signal: controller.signal,
|
||||
onAgent: agentFn,
|
||||
onPhase: phaseFn,
|
||||
});
|
||||
} catch (error) {
|
||||
details.error = errorText(error);
|
||||
status = controller.signal.aborted ? "aborted" : "failed";
|
||||
controller.abort("Workflow script failed");
|
||||
}
|
||||
|
||||
const settled = await controller.settle({
|
||||
abort: status !== "completed",
|
||||
});
|
||||
if (!settled) {
|
||||
status = "failed";
|
||||
details.error = details.error
|
||||
? `${details.error}; agent shutdown deadline exceeded`
|
||||
: "Agent shutdown deadline exceeded";
|
||||
}
|
||||
for (const record of details.agents) {
|
||||
if (record.state !== "running") continue;
|
||||
record.state = "error";
|
||||
record.error =
|
||||
record.error ?? "Agent did not settle before run cleanup";
|
||||
record.finishedAt = Date.now();
|
||||
}
|
||||
details.status = status;
|
||||
details.finishedAt = Date.now();
|
||||
try {
|
||||
persistence.flush();
|
||||
} catch (error) {
|
||||
details.status = "failed";
|
||||
details.error = `Artifact persistence failed: ${errorText(error)}`;
|
||||
throw new Error(details.error);
|
||||
} finally {
|
||||
flushNow();
|
||||
}
|
||||
};
|
||||
|
||||
// Registered for /workflows visibility and session_shutdown abort;
|
||||
// blocking runs are watchable live from the dashboard too.
|
||||
const activeRun = { details, controller } as {
|
||||
details: WorkflowDetails;
|
||||
controller: RunController;
|
||||
completion?: Promise<void>;
|
||||
};
|
||||
activeRuns.set(runId, activeRun);
|
||||
const completion = runScript();
|
||||
activeRun.completion = completion;
|
||||
if (ctx.hasUI) lastUi = ctx.ui;
|
||||
updateIndicator();
|
||||
|
||||
if (background) {
|
||||
void completion
|
||||
.catch((error) => {
|
||||
details.status = "failed";
|
||||
details.finishedAt = Date.now();
|
||||
details.error = details.error ?? errorText(error);
|
||||
})
|
||||
.finally(() => {
|
||||
activeRuns.delete(runId);
|
||||
recordSettledRun(details.status);
|
||||
updateIndicator();
|
||||
try {
|
||||
pi.sendUserMessage(
|
||||
buildBackgroundWorkflowFollowUp({
|
||||
runId,
|
||||
status: details.status,
|
||||
result: buildWorkflowResultMessage(details, runDir),
|
||||
}),
|
||||
{ deliverAs: "followUp" },
|
||||
);
|
||||
} catch {
|
||||
// Session may be shutting down.
|
||||
}
|
||||
});
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: buildBackgroundWorkflowLaunchResult({
|
||||
runId,
|
||||
name: details.name,
|
||||
runDir,
|
||||
}),
|
||||
},
|
||||
],
|
||||
details: compactToolDetails(details),
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
await completion;
|
||||
} finally {
|
||||
activeRuns.delete(runId);
|
||||
recordSettledRun(details.status);
|
||||
updateIndicator();
|
||||
}
|
||||
if (details.status !== "completed") {
|
||||
// Pi marks tool failures only when execute throws; returning isError is
|
||||
// ignored by the extension API.
|
||||
throw new Error(buildWorkflowResultMessage(details, runDir));
|
||||
}
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: buildWorkflowResultMessage(details, runDir),
|
||||
},
|
||||
],
|
||||
details: compactToolDetails(details),
|
||||
};
|
||||
},
|
||||
|
||||
renderCall(args: Partial<WorkflowInput>, theme) {
|
||||
const meta =
|
||||
typeof args.script === "string"
|
||||
? extractMeta(args.script)
|
||||
: { phases: [] };
|
||||
let text =
|
||||
theme.fg("toolTitle", theme.bold("workflow ")) +
|
||||
theme.fg("accent", (meta as WorkflowMeta).name ?? "(script)");
|
||||
if (args.background) text += theme.fg("dim", " (background)");
|
||||
const description = (meta as WorkflowMeta).description;
|
||||
if (description) text += `\n ${theme.fg("dim", description)}`;
|
||||
for (const phase of meta.phases.slice(0, 8)) {
|
||||
text += `\n ${theme.fg("dim", SQUARE)} ${theme.fg("accent", phase.title)}${
|
||||
phase.detail ? theme.fg("dim", ` — ${phase.detail}`) : ""
|
||||
}`;
|
||||
}
|
||||
return new Text(text, 0, 0);
|
||||
},
|
||||
|
||||
renderResult(result, { expanded }, theme) {
|
||||
const details = result.details as WorkflowDetails | undefined;
|
||||
if (!details) {
|
||||
const first = result.content[0];
|
||||
return new Text(
|
||||
first?.type === "text" ? first.text : "(no output)",
|
||||
0,
|
||||
0,
|
||||
);
|
||||
}
|
||||
|
||||
const { done, failed } = countStates(details);
|
||||
const settled = done + failed;
|
||||
const elapsed = formatElapsed(details.startedAt, details.finishedAt);
|
||||
let header =
|
||||
`${theme.fg(statusColor(details.status), SQUARE)} ${theme.fg("toolTitle", theme.bold("workflow "))}` +
|
||||
`${theme.fg("accent", details.name ?? details.runId)} ` +
|
||||
theme.fg(
|
||||
"dim",
|
||||
`${settled}/${details.agents.length} agents · ${elapsed} · `,
|
||||
) +
|
||||
theme.fg(statusColor(details.status), statusWord(details.status));
|
||||
if (failed) header += theme.fg("error", ` · ${failed} failed`);
|
||||
if (details.background) header += theme.fg("dim", " (background)");
|
||||
if (details.status === "running" && details.currentPhase) {
|
||||
header += theme.fg("muted", ` · ${details.currentPhase}`);
|
||||
}
|
||||
const totals = formatUsage(aggregateUsage(details.agents));
|
||||
|
||||
if (!expanded) {
|
||||
let text = header;
|
||||
for (const agent of details.agents) {
|
||||
const context = agentContext(agent);
|
||||
text += `\n ${stateSquare(agent.state, theme)} ${theme.fg("accent", agent.label)}${
|
||||
agent.phase ? theme.fg("dim", ` (${agent.phase})`) : ""
|
||||
}${theme.fg(
|
||||
"dim",
|
||||
`${context ? ` · ${context}` : ""} · ${formatElapsed(agent.startedAt, agent.finishedAt)}`,
|
||||
)}`;
|
||||
}
|
||||
if (totals) text += `\n ${theme.fg("dim", `Total: ${totals}`)}`;
|
||||
if (details.error)
|
||||
text += `\n ${theme.fg("error", `Error: ${details.error}`)}`;
|
||||
text += `\n${theme.fg("muted", `(${keyHint("app.tools.expand", "to expand")})`)}`;
|
||||
return new Text(text, 0, 0);
|
||||
}
|
||||
|
||||
const container = new Container();
|
||||
container.addChild(new Text(header, 0, 0));
|
||||
if (details.description) {
|
||||
container.addChild(
|
||||
new Text(theme.fg("dim", details.description), 0, 0),
|
||||
);
|
||||
}
|
||||
|
||||
for (const group of phaseGroups(details)) {
|
||||
container.addChild(new Spacer(1));
|
||||
container.addChild(
|
||||
new Text(theme.fg("muted", `─── ${group.title} ───`), 0, 0),
|
||||
);
|
||||
for (const agent of group.agents) {
|
||||
const usage = formatUsage(agent.usage, agent.model);
|
||||
const context = agentContext(agent);
|
||||
let line = `${stateSquare(agent.state, theme)} ${theme.fg("accent", agent.label)} ${theme.fg(
|
||||
"dim",
|
||||
[context, formatElapsed(agent.startedAt, agent.finishedAt)]
|
||||
.filter(Boolean)
|
||||
.join(" · "),
|
||||
)}`;
|
||||
if (usage) line += ` ${theme.fg("dim", usage)}`;
|
||||
container.addChild(new Text(line, 0, 0));
|
||||
if (agent.error) {
|
||||
container.addChild(
|
||||
new Text(` ${theme.fg("error", agent.error)}`, 0, 0),
|
||||
);
|
||||
} else if (agent.preview) {
|
||||
const preview = agent.preview.split("\n").slice(0, 2).join(" ");
|
||||
container.addChild(new Text(` ${theme.fg("dim", preview)}`, 0, 0));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (details.error) {
|
||||
container.addChild(new Spacer(1));
|
||||
container.addChild(
|
||||
new Text(theme.fg("error", `Error: ${details.error}`), 0, 0),
|
||||
);
|
||||
}
|
||||
|
||||
if (details.result !== undefined) {
|
||||
container.addChild(new Spacer(1));
|
||||
container.addChild(new Text(theme.fg("muted", "─── result ───"), 0, 0));
|
||||
container.addChild(
|
||||
new Markdown(
|
||||
`\`\`\`json\n${resultJson(details.result)}\n\`\`\``,
|
||||
0,
|
||||
0,
|
||||
getMarkdownTheme(),
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
if (totals) {
|
||||
container.addChild(new Spacer(1));
|
||||
container.addChild(new Text(theme.fg("dim", `Total: ${totals}`), 0, 0));
|
||||
}
|
||||
return container;
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { test } from "node:test";
|
||||
import { extractMeta, prepareWorkflowScript } from "./meta.ts";
|
||||
|
||||
test("metadata is decoded statically and removed from executable source", () => {
|
||||
const source = `export const meta = {
|
||||
name: "audit",
|
||||
description: "safe",
|
||||
phases: [{ title: "Scan", detail: "files" }],
|
||||
};
|
||||
return { ok: true };`;
|
||||
const prepared = prepareWorkflowScript(source);
|
||||
assert.deepEqual(prepared.meta, {
|
||||
name: "audit",
|
||||
description: "safe",
|
||||
phases: [{ title: "Scan", detail: "files" }],
|
||||
});
|
||||
assert.doesNotMatch(prepared.source, /name:\s*"audit"/);
|
||||
assert.equal(prepared.source.split("\n").length, source.split("\n").length);
|
||||
});
|
||||
|
||||
test("export-like text in strings, comments, regexes, and templates is untouched", () => {
|
||||
const source = `
|
||||
const string = "export default notSyntax";
|
||||
const template = \`export const meta = \${string}\`;
|
||||
const regex = /export\\s+default/;
|
||||
// export const fake = 1
|
||||
return { string, template, matches: regex.test(string) };
|
||||
`;
|
||||
const prepared = prepareWorkflowScript(source);
|
||||
assert.equal(prepared.source, source);
|
||||
assert.deepEqual(prepared.meta, { phases: [] });
|
||||
});
|
||||
|
||||
test("executable and unsupported metadata fail closed", () => {
|
||||
assert.throws(
|
||||
() =>
|
||||
prepareWorkflowScript(
|
||||
`export const meta = { name: (() => "executed")(), phases: [] }; return 1;`,
|
||||
),
|
||||
/only static literals/,
|
||||
);
|
||||
assert.throws(
|
||||
() => prepareWorkflowScript(`export default 1; return 1;`),
|
||||
/may only export/,
|
||||
);
|
||||
assert.deepEqual(
|
||||
extractMeta(`export const meta = { name: process.exit(), phases: [] }`),
|
||||
{ phases: [] },
|
||||
);
|
||||
});
|
||||
@@ -0,0 +1,237 @@
|
||||
import {
|
||||
parse,
|
||||
type ArrayExpression,
|
||||
type ExportNamedDeclaration,
|
||||
type Expression,
|
||||
type Identifier,
|
||||
type Literal,
|
||||
type ObjectExpression,
|
||||
type Program,
|
||||
type Property,
|
||||
type VariableDeclaration,
|
||||
} from "acorn";
|
||||
|
||||
/** Static workflow metadata and source preparation helpers. */
|
||||
|
||||
export interface WorkflowPhase {
|
||||
title: string;
|
||||
detail?: string;
|
||||
}
|
||||
|
||||
export interface WorkflowMeta {
|
||||
name?: string;
|
||||
description?: string;
|
||||
phases: WorkflowPhase[];
|
||||
}
|
||||
|
||||
export interface PreparedWorkflowScript {
|
||||
/** Script body with the metadata declaration removed. */
|
||||
source: string;
|
||||
meta: WorkflowMeta;
|
||||
}
|
||||
|
||||
function parseProgram(source: string): Program {
|
||||
return parse(source, {
|
||||
ecmaVersion: "latest",
|
||||
sourceType: "module",
|
||||
allowReturnOutsideFunction: true,
|
||||
});
|
||||
}
|
||||
|
||||
function isIdentifier(node: Expression, name?: string): node is Identifier {
|
||||
return (
|
||||
node.type === "Identifier" && (name === undefined || node.name === name)
|
||||
);
|
||||
}
|
||||
|
||||
function isLiteral(node: Expression): node is Literal {
|
||||
return node.type === "Literal";
|
||||
}
|
||||
|
||||
function isProperty(
|
||||
node: ObjectExpression["properties"][number],
|
||||
): node is Property {
|
||||
return node.type === "Property";
|
||||
}
|
||||
|
||||
function propertyName(property: Property) {
|
||||
if (property.computed || property.kind !== "init" || property.method) {
|
||||
return undefined;
|
||||
}
|
||||
if (isIdentifier(property.key)) return property.key.name;
|
||||
if (isLiteral(property.key) && typeof property.key.value === "string") {
|
||||
return property.key.value;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert an Acorn expression to data without evaluating source. Only plain
|
||||
* object/array/primitive literals are accepted. Getters, methods, spreads,
|
||||
* computed keys, templates, identifiers, and calls all fail closed.
|
||||
*/
|
||||
function literalValue(node: Expression, depth = 0): unknown {
|
||||
if (depth > 8) throw new Error("workflow metadata is nested too deeply");
|
||||
if (isLiteral(node)) {
|
||||
if (
|
||||
node.value === null ||
|
||||
typeof node.value === "string" ||
|
||||
typeof node.value === "number" ||
|
||||
typeof node.value === "boolean"
|
||||
) {
|
||||
return node.value;
|
||||
}
|
||||
throw new Error("workflow metadata contains an unsupported literal");
|
||||
}
|
||||
if (node.type === "ArrayExpression") {
|
||||
const arrayNode: ArrayExpression = node;
|
||||
return arrayNode.elements.map((element) => {
|
||||
if (!element || element.type === "SpreadElement") {
|
||||
throw new Error(
|
||||
"workflow metadata arrays cannot contain holes or spreads",
|
||||
);
|
||||
}
|
||||
return literalValue(element, depth + 1);
|
||||
});
|
||||
}
|
||||
if (node.type === "ObjectExpression") {
|
||||
const objectNode: ObjectExpression = node;
|
||||
const value: Record<string, unknown> = Object.create(null);
|
||||
for (const item of objectNode.properties) {
|
||||
if (!isProperty(item)) {
|
||||
throw new Error("workflow metadata objects cannot contain spreads");
|
||||
}
|
||||
const key = propertyName(item);
|
||||
if (key === undefined || item.shorthand) {
|
||||
throw new Error(
|
||||
"workflow metadata keys and values must be plain literals",
|
||||
);
|
||||
}
|
||||
value[key] = literalValue(item.value, depth + 1);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
throw new Error("workflow metadata must contain only static literals");
|
||||
}
|
||||
|
||||
function sanitizeMeta(value: unknown): WorkflowMeta {
|
||||
const meta: WorkflowMeta = { phases: [] };
|
||||
if (!value || typeof value !== "object") return meta;
|
||||
const raw = value as {
|
||||
name?: unknown;
|
||||
description?: unknown;
|
||||
phases?: unknown;
|
||||
};
|
||||
if (typeof raw.name === "string") meta.name = raw.name.slice(0, 160);
|
||||
if (typeof raw.description === "string") {
|
||||
meta.description = raw.description.slice(0, 2_000);
|
||||
}
|
||||
if (Array.isArray(raw.phases)) {
|
||||
for (const item of raw.phases.slice(0, 64)) {
|
||||
if (!item || typeof item !== "object") continue;
|
||||
const phase = item as { title?: unknown; detail?: unknown };
|
||||
if (typeof phase.title !== "string" || !phase.title.trim()) continue;
|
||||
meta.phases.push({
|
||||
title: phase.title.slice(0, 160),
|
||||
...(typeof phase.detail === "string"
|
||||
? { detail: phase.detail.slice(0, 2_000) }
|
||||
: {}),
|
||||
});
|
||||
}
|
||||
}
|
||||
return meta;
|
||||
}
|
||||
|
||||
function metadataDeclaration(statement: Program["body"][number]) {
|
||||
if (statement.type !== "ExportNamedDeclaration") return undefined;
|
||||
const exported: ExportNamedDeclaration = statement;
|
||||
if (
|
||||
exported.source ||
|
||||
exported.specifiers.length > 0 ||
|
||||
exported.declaration?.type !== "VariableDeclaration"
|
||||
) {
|
||||
throw new Error(
|
||||
"workflow scripts may only export a static `const meta = {...}` declaration",
|
||||
);
|
||||
}
|
||||
const declaration: VariableDeclaration = exported.declaration;
|
||||
if (declaration.kind !== "const" || declaration.declarations.length !== 1) {
|
||||
throw new Error(
|
||||
"workflow metadata must be declared as `export const meta = {...}`",
|
||||
);
|
||||
}
|
||||
const declarator = declaration.declarations[0];
|
||||
if (
|
||||
declarator.id.type !== "Identifier" ||
|
||||
declarator.id.name !== "meta" ||
|
||||
!declarator.init
|
||||
) {
|
||||
throw new Error("workflow scripts may only export a `meta` declaration");
|
||||
}
|
||||
return { exported, initializer: declarator.init };
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse a workflow as JavaScript, reject module syntax other than the metadata
|
||||
* API, statically decode metadata, and remove that declaration from the body.
|
||||
* Source inside strings, comments, regexes, and templates is never rewritten.
|
||||
*/
|
||||
export function prepareWorkflowScript(source: string): PreparedWorkflowScript {
|
||||
const program = parseProgram(source);
|
||||
let meta: WorkflowMeta = { phases: [] };
|
||||
let metadataRange: { start: number; end: number } | undefined;
|
||||
|
||||
for (const statement of program.body) {
|
||||
if (statement.type === "ImportDeclaration") {
|
||||
throw new Error("workflow scripts cannot use static imports");
|
||||
}
|
||||
if (
|
||||
statement.type === "ExportDefaultDeclaration" ||
|
||||
statement.type === "ExportAllDeclaration"
|
||||
) {
|
||||
throw new Error(
|
||||
"workflow scripts may only export a static `const meta = {...}` declaration",
|
||||
);
|
||||
}
|
||||
const declaration = metadataDeclaration(statement);
|
||||
if (!declaration) continue;
|
||||
if (metadataRange)
|
||||
throw new Error("workflow metadata may only be declared once");
|
||||
meta = sanitizeMeta(literalValue(declaration.initializer));
|
||||
metadataRange = {
|
||||
start: declaration.exported.start,
|
||||
end: declaration.exported.end,
|
||||
};
|
||||
}
|
||||
|
||||
if (!metadataRange) return { source, meta };
|
||||
// Preserve line numbers for parse/runtime errors while removing every byte of
|
||||
// executable metadata. A semicolon keeps adjacent statements separated.
|
||||
const removed = source.slice(metadataRange.start, metadataRange.end);
|
||||
const replacement = `;${removed.slice(1).replace(/[^\n\r]/g, " ")}`;
|
||||
return {
|
||||
source:
|
||||
source.slice(0, metadataRange.start) +
|
||||
replacement +
|
||||
source.slice(metadataRange.end),
|
||||
meta,
|
||||
};
|
||||
}
|
||||
|
||||
const metaCache = new Map<string, WorkflowMeta>();
|
||||
const META_CACHE_LIMIT = 32;
|
||||
|
||||
/** Safe metadata extraction for render paths. Invalid/partial scripts are empty. */
|
||||
export function extractMeta(script: string): WorkflowMeta {
|
||||
const cached = metaCache.get(script);
|
||||
if (cached) return cached;
|
||||
let meta: WorkflowMeta = { phases: [] };
|
||||
try {
|
||||
meta = prepareWorkflowScript(script).meta;
|
||||
} catch {
|
||||
// Renderers must tolerate partial model output and invalid scripts.
|
||||
}
|
||||
if (metaCache.size >= META_CACHE_LIMIT) metaCache.clear();
|
||||
metaCache.set(script, meta);
|
||||
return meta;
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
/**
|
||||
* Shared workflow run model + formatting helpers, used by the tool renderers
|
||||
* (index.ts) and the /workflows dashboard (dashboard.ts).
|
||||
*/
|
||||
|
||||
import * as os from "node:os";
|
||||
import {
|
||||
truncateHead,
|
||||
type ExtensionContext,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import { formatContextUtilization } from "../shared/context-utilization.ts";
|
||||
import { safeStringify } from "./serialization.ts";
|
||||
|
||||
export type Theme = ExtensionContext["ui"]["theme"];
|
||||
|
||||
export const RESULT_JSON_MAX_BYTES = 24 * 1024;
|
||||
export const RESULT_JSON_MAX_LINES = 600;
|
||||
|
||||
export interface AgentUsage {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheRead: number;
|
||||
cacheWrite: number;
|
||||
cost: number;
|
||||
/** Latest compaction-aware conversation occupancy, not cumulative billing. */
|
||||
contextTokens?: number;
|
||||
turns: number;
|
||||
}
|
||||
|
||||
export function emptyUsage(): AgentUsage {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
cost: 0,
|
||||
turns: 0,
|
||||
};
|
||||
}
|
||||
|
||||
export type AgentState = "running" | "done" | "error";
|
||||
export type WorkflowStatus = "running" | "completed" | "failed" | "aborted";
|
||||
|
||||
export type TranscriptRole =
|
||||
"user" | "assistant" | "thinking" | "tool" | "toolResult";
|
||||
|
||||
export interface TranscriptEntry {
|
||||
role: TranscriptRole;
|
||||
text: string;
|
||||
/** Tool name for tool calls/results. */
|
||||
name?: string;
|
||||
/** Stable tool-call identifier used to pair calls, results, and timings. */
|
||||
toolCallId?: string;
|
||||
isError?: boolean;
|
||||
/** Original message timestamp, when provided by the model/session. */
|
||||
timestamp?: number;
|
||||
/** Tool execution lifecycle timestamps, measured by the child session. */
|
||||
startedAt?: number;
|
||||
finishedAt?: number;
|
||||
durationMs?: number;
|
||||
}
|
||||
|
||||
export interface AgentRecord {
|
||||
index: number;
|
||||
label: string;
|
||||
phase?: string;
|
||||
state: AgentState;
|
||||
model?: string;
|
||||
/** Context capacity of the active model used for this agent. */
|
||||
contextWindow?: number;
|
||||
startedAt: number;
|
||||
finishedAt?: number;
|
||||
error?: string;
|
||||
preview: string;
|
||||
usage: AgentUsage;
|
||||
/** Normalized, serializable subagent conversation shown by /workflows. */
|
||||
transcript: TranscriptEntry[];
|
||||
}
|
||||
|
||||
export interface WorkflowDetails {
|
||||
runId: string;
|
||||
/** Pi session that launched this run. */
|
||||
sessionId?: string;
|
||||
name?: string;
|
||||
description?: string;
|
||||
background: boolean;
|
||||
status: WorkflowStatus;
|
||||
startedAt: number;
|
||||
finishedAt?: number;
|
||||
phases: { title: string; detail?: string }[];
|
||||
currentPhase?: string;
|
||||
agents: AgentRecord[];
|
||||
result?: unknown;
|
||||
resultArtifact?: string;
|
||||
transcriptArtifact?: string;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/** Colored square state indicator (no emojis/glyphs). */
|
||||
export const SQUARE = "■";
|
||||
|
||||
export function stateSquare(state: AgentState, theme: Theme): string {
|
||||
if (state === "done") return theme.fg("success", SQUARE);
|
||||
if (state === "error") return theme.fg("error", SQUARE);
|
||||
return theme.fg("warning", SQUARE);
|
||||
}
|
||||
|
||||
export function statusSquare(status: WorkflowStatus, theme: Theme): string {
|
||||
if (status === "completed") return theme.fg("success", SQUARE);
|
||||
if (status === "running") return theme.fg("warning", SQUARE);
|
||||
return theme.fg("error", SQUARE);
|
||||
}
|
||||
|
||||
export function statusWord(status: WorkflowStatus): string {
|
||||
return status === "completed" ? "done" : status;
|
||||
}
|
||||
|
||||
export function statusColor(
|
||||
status: WorkflowStatus,
|
||||
): "success" | "warning" | "error" {
|
||||
if (status === "completed") return "success";
|
||||
if (status === "running") return "warning";
|
||||
return "error";
|
||||
}
|
||||
|
||||
export function shortenHome(p: string): string {
|
||||
const home = os.homedir();
|
||||
return p.startsWith(home) ? `~${p.slice(home.length)}` : p;
|
||||
}
|
||||
|
||||
export function formatTokens(count: number): string {
|
||||
if (count < 1000) return count.toString();
|
||||
if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
|
||||
if (count < 1000000) return `${Math.round(count / 1000)}k`;
|
||||
return `${(count / 1000000).toFixed(1)}M`;
|
||||
}
|
||||
|
||||
export function formatUsage(usage: AgentUsage, model?: string): string {
|
||||
const parts: string[] = [];
|
||||
if (usage.turns)
|
||||
parts.push(`${usage.turns} turn${usage.turns > 1 ? "s" : ""}`);
|
||||
if (usage.input) parts.push(`${formatTokens(usage.input)} in`);
|
||||
if (usage.output) parts.push(`${formatTokens(usage.output)} out`);
|
||||
if (usage.cost) parts.push(`$${usage.cost.toFixed(4)}`);
|
||||
if (model) parts.push(model);
|
||||
return parts.join(" · ");
|
||||
}
|
||||
|
||||
/** Current per-agent context-window utilization, e.g. "7%/372k". */
|
||||
export function agentContext(agent: AgentRecord): string {
|
||||
return formatContextUtilization({
|
||||
tokens: agent.usage.contextTokens,
|
||||
contextWindow: agent.contextWindow,
|
||||
});
|
||||
}
|
||||
|
||||
export function formatElapsed(startedAt: number, finishedAt?: number): string {
|
||||
const totalSeconds = Math.max(
|
||||
0,
|
||||
Math.round(((finishedAt ?? Date.now()) - startedAt) / 1000),
|
||||
);
|
||||
const minutes = Math.floor(totalSeconds / 60);
|
||||
const seconds = totalSeconds % 60;
|
||||
return minutes > 0
|
||||
? `${minutes}m${seconds.toString().padStart(2, "0")}s`
|
||||
: `${seconds}s`;
|
||||
}
|
||||
|
||||
export function aggregateUsage(agents: AgentRecord[]): AgentUsage {
|
||||
const total = emptyUsage();
|
||||
for (const agent of agents) {
|
||||
total.input += agent.usage.input;
|
||||
total.output += agent.usage.output;
|
||||
total.cacheRead += agent.usage.cacheRead;
|
||||
total.cacheWrite += agent.usage.cacheWrite;
|
||||
total.cost += agent.usage.cost;
|
||||
total.turns += agent.usage.turns;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
export function countStates(details: WorkflowDetails) {
|
||||
let done = 0;
|
||||
let failed = 0;
|
||||
let running = 0;
|
||||
for (const agent of details.agents) {
|
||||
if (agent.state === "done") done++;
|
||||
else if (agent.state === "error") failed++;
|
||||
else running++;
|
||||
}
|
||||
return { done, failed, running };
|
||||
}
|
||||
|
||||
export interface PhaseGroup {
|
||||
title: string;
|
||||
agents: AgentRecord[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Group agents by phase in declared phase order. With `includeEmpty`, phases
|
||||
* that have no agents yet are still listed (used by the dashboard sidebar).
|
||||
*/
|
||||
export function phaseGroups(
|
||||
details: WorkflowDetails,
|
||||
includeEmpty = false,
|
||||
): PhaseGroup[] {
|
||||
const byPhase = new Map<string, AgentRecord[]>();
|
||||
for (const agent of details.agents) {
|
||||
const key = agent.phase ?? "(unphased)";
|
||||
const list = byPhase.get(key) ?? [];
|
||||
list.push(agent);
|
||||
byPhase.set(key, list);
|
||||
}
|
||||
const groups: PhaseGroup[] = [];
|
||||
for (const phase of details.phases) {
|
||||
const agents = byPhase.get(phase.title);
|
||||
if (agents || includeEmpty)
|
||||
groups.push({ title: phase.title, agents: agents ?? [] });
|
||||
byPhase.delete(phase.title);
|
||||
}
|
||||
for (const [title, agents] of byPhase) groups.push({ title, agents });
|
||||
return groups;
|
||||
}
|
||||
|
||||
export function resultJson(value: unknown): string {
|
||||
const text =
|
||||
typeof value === "string"
|
||||
? value
|
||||
: safeStringify(value, {
|
||||
maxBytes: RESULT_JSON_MAX_BYTES * 2,
|
||||
maxDepth: 16,
|
||||
maxNodes: 10_000,
|
||||
});
|
||||
const truncation = truncateHead(text ?? "", {
|
||||
maxLines: RESULT_JSON_MAX_LINES,
|
||||
maxBytes: RESULT_JSON_MAX_BYTES,
|
||||
});
|
||||
return truncation.truncated
|
||||
? `${truncation.content}\n…[result truncated; bounded result artifact in result.json]`
|
||||
: truncation.content;
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
import {
|
||||
countStates,
|
||||
formatElapsed,
|
||||
resultJson,
|
||||
shortenHome,
|
||||
type WorkflowDetails,
|
||||
} from "./model.ts";
|
||||
|
||||
/** Model-facing schema descriptions for workflow source, arguments, and background mode. */
|
||||
export const WORKFLOW_PARAMETER_DESCRIPTIONS = {
|
||||
script:
|
||||
"JavaScript workflow script. May start with `export const meta = {...}`, then use phase(), agent(), parallel(), args, and a final `return`.",
|
||||
args: "Optional JSON string exposed to the script as `args` (parsed when valid JSON, otherwise passed through as the raw string).",
|
||||
background:
|
||||
"Run in the background: the tool returns a run id immediately and you receive a follow-up message when the workflow finishes. Defaults to false (blocking with live progress).",
|
||||
};
|
||||
|
||||
/** Defines the workflow DSL, constraints, reliability guidance, and model-authored task examples. */
|
||||
export const WORKFLOW_TOOL_DESCRIPTION = [
|
||||
"The workflow tool is only to be called when the user says 'ultracode' or specifically requests a workflow run.",
|
||||
"Run a multi-agent workflow from a JavaScript orchestration script you write inline. Use this when a task benefits from fanning work out across several isolated subagents in ordered phases (research fan-out, per-file review, verify-then-synthesize pipelines).",
|
||||
"The script runs as an async function body with these primitives:",
|
||||
"• export const meta = { name, description, phases: [{ title, detail? }] } — metadata for the progress UI. Declare all phases up front.",
|
||||
"• phase(title) — mark the current phase at runtime (use titles from meta.phases).",
|
||||
"• await agent(prompt, { label?, phase?, schema?, model?, provider?, effort? }) — run ONE subagent in an isolated context and wait for it. Always resolves to { ok, output, structured?, error? }. Check `ok` before using the result. When you pass a JSON `schema`, `structured` holds the validated object on success. `model`/`provider` override the session model; `effort` sets the thinking level (off|minimal|low|medium|high|xhigh|max). Children receive normal built-ins and trust-appropriate extensions, settings, skills, and AGENTS.md context, but cannot recursively orchestrate or ask the user.",
|
||||
"• await parallel([() => agent(...), () => agent(...)], { concurrency? }) — run zero-argument agent thunks concurrently and return results in order. Concurrency is globally capped at 4 for the run.",
|
||||
"• args — the parsed value of the `args` tool parameter (or undefined).",
|
||||
"Workflow JavaScript runs in a restricted, killable child with no imports, eval, timers, filesystem, network, or process APIs. A run may make at most 32 agent calls and has no overall deadline. Each agent must receive its first assistant response event within 45 seconds so silent provider requests fail clearly; after that, agent() has no wall-clock deadline. Each individual child tool call times out independently after 3 minutes, becomes an error tool result, and leaves the agent loop free to recover. Use map/filter/if/await/template strings to orchestrate, and `return` a JSON-serializable aggregate.",
|
||||
"Pass a `schema` to agent() whenever a later step branches on the result, so you get typed fields instead of prose. There is no resume: a failed run is simply re-run. Artifacts are saved under ~/.pi/agent/workflows/<runId>/ for inspection.",
|
||||
"Example:",
|
||||
"export const meta = { name: 'reliability-review', description: 'Review modules for reliability risks, then report', phases: [{ title: 'Scan' }, { title: 'Report' }] }",
|
||||
"const FINDINGS = { type: 'object', properties: { issues: { type: 'array', items: { type: 'string' } }, ok: { type: 'boolean' } }, required: ['issues', 'ok'] }",
|
||||
"phase('Scan')",
|
||||
"const scans = await parallel(args.files.map((f) => () => agent(`Review ${f} for correctness and reliability risks.`, { label: `scan:${f}`, phase: 'Scan', schema: FINDINGS })))",
|
||||
"const findings = scans.filter((r) => r.ok).map((r) => r.structured)",
|
||||
"phase('Report')",
|
||||
"const report = await agent(`Summarize these findings: ${JSON.stringify(findings)}`, { label: 'report', phase: 'Report' })",
|
||||
"return { findings, report: report.ok ? report.output : report.error }",
|
||||
].join("\n");
|
||||
|
||||
/** Adds workflow orchestration primitives and background execution to the model's tool prompt. */
|
||||
export const WORKFLOW_PROMPT_SNIPPET =
|
||||
"Orchestrate isolated subagents from an inline JS script: phase()/agent()/parallel() with structured outputs and optional background execution";
|
||||
|
||||
/** Guides the model on appropriate workflow fan-out and mandatory agent result checks. */
|
||||
export const WORKFLOW_PROMPT_GUIDELINES = [
|
||||
"Use workflow when a task needs several subagents with phase dependencies or dynamic fan-out; keep single small delegations in the main session.",
|
||||
"In workflow scripts, agent() never throws — always check `.ok` on its result before using `.output`/`.structured`.",
|
||||
];
|
||||
|
||||
/** Marks and forwards a workflow script's agent() task as an isolated child-model prompt. */
|
||||
export function buildWorkflowAgentPrompt(prompt: string) {
|
||||
return prompt;
|
||||
}
|
||||
|
||||
/** Instructs structured workflow children to terminate with exactly one structured_output call. */
|
||||
export const STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION =
|
||||
"When your task is complete, call the `structured_output` tool exactly once as your final action, with fields matching the required schema. Do not write any other text after it.";
|
||||
|
||||
/** Describes the terminating structured_output tool and its final-action contract. */
|
||||
export const STRUCTURED_OUTPUT_TOOL_DESCRIPTION =
|
||||
"Return your final result as structured data matching the required schema. Call this exactly once, as your last action; do not write any other text after it.";
|
||||
|
||||
/** Builds the workflow completion report returned to the parent model. */
|
||||
export function buildWorkflowResultMessage(
|
||||
details: WorkflowDetails,
|
||||
runDir: string,
|
||||
) {
|
||||
const { done, failed } = countStates(details);
|
||||
const elapsed = formatElapsed(details.startedAt, details.finishedAt);
|
||||
const lines = [
|
||||
`Workflow ${details.name ? `"${details.name}"` : details.runId} ${details.status} — ` +
|
||||
`${done}/${details.agents.length} agents ok${failed ? `, ${failed} failed` : ""} ` +
|
||||
`across ${details.phases.length} phase(s) in ${elapsed}.`,
|
||||
`Run dir: ${shortenHome(runDir)}`,
|
||||
];
|
||||
if (details.error) lines.push(`Error: ${details.error}`);
|
||||
if (details.agents.length > 0) {
|
||||
lines.push("", "Agents:");
|
||||
for (const agent of details.agents) {
|
||||
const status =
|
||||
agent.state === "done"
|
||||
? "ok"
|
||||
: agent.state === "error"
|
||||
? "FAILED"
|
||||
: "running";
|
||||
lines.push(
|
||||
`- [${agent.label}]${agent.phase ? ` (${agent.phase})` : ""} ${status}` +
|
||||
(agent.error ? ` — ${agent.error}` : ""),
|
||||
);
|
||||
}
|
||||
}
|
||||
if (details.result !== undefined)
|
||||
lines.push("", "Result:", resultJson(details.result));
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
/** Builds the follow-up user message that delivers a settled background workflow to the parent model. */
|
||||
export function buildBackgroundWorkflowFollowUp(options: {
|
||||
runId: string;
|
||||
status: WorkflowDetails["status"];
|
||||
result: string;
|
||||
}) {
|
||||
return `[Background workflow ${options.runId} ${options.status}]\n\n${options.result}`;
|
||||
}
|
||||
|
||||
/** Builds the background-launch result and tells the parent model where progress and artifacts appear. */
|
||||
export function buildBackgroundWorkflowLaunchResult(options: {
|
||||
runId: string;
|
||||
name?: string;
|
||||
runDir: string;
|
||||
}) {
|
||||
return [
|
||||
`Workflow ${options.name ? `"${options.name}"` : options.runId} launched in background (run ${options.runId}).`,
|
||||
`Artifacts: ${shortenHome(options.runDir)}`,
|
||||
"You'll receive a follow-up message when it finishes; /workflows shows progress.",
|
||||
].join("\n");
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { test } from "node:test";
|
||||
import type {
|
||||
AgentSession,
|
||||
AgentSessionEventListener,
|
||||
ToolDefinition,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import { Type } from "typebox";
|
||||
import {
|
||||
createFirstResponseWatchdog,
|
||||
guardWorkflowChildTools,
|
||||
recordToolExecutionTiming,
|
||||
transcriptFromMessages,
|
||||
type ToolExecutionTiming,
|
||||
} from "./runner.ts";
|
||||
|
||||
const zeroUsage = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
};
|
||||
|
||||
function parallelToolMessages(): AgentSession["messages"] {
|
||||
return [
|
||||
{ role: "user", content: "run both", timestamp: 900 },
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "call-a",
|
||||
name: "first",
|
||||
arguments: { value: 1 },
|
||||
},
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "call-b",
|
||||
name: "second",
|
||||
arguments: { value: 2 },
|
||||
},
|
||||
],
|
||||
api: "openai-responses",
|
||||
provider: "fixture",
|
||||
model: "fixture",
|
||||
usage: zeroUsage,
|
||||
stopReason: "toolUse",
|
||||
timestamp: 950,
|
||||
},
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "call-a",
|
||||
toolName: "first",
|
||||
content: [{ type: "text", text: "first result" }],
|
||||
isError: false,
|
||||
timestamp: 1_040,
|
||||
},
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "call-b",
|
||||
toolName: "second",
|
||||
content: [{ type: "text", text: "second result" }],
|
||||
isError: false,
|
||||
timestamp: 1_041,
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
test("completed parallel tool calls pair lifecycle timings with calls and results", () => {
|
||||
const timings = new Map<string, ToolExecutionTiming>();
|
||||
recordToolExecutionTiming(
|
||||
timings,
|
||||
{
|
||||
type: "tool_execution_start",
|
||||
toolCallId: "call-a",
|
||||
toolName: "first",
|
||||
args: { value: 1 },
|
||||
},
|
||||
1_000,
|
||||
);
|
||||
recordToolExecutionTiming(
|
||||
timings,
|
||||
{
|
||||
type: "tool_execution_start",
|
||||
toolCallId: "call-b",
|
||||
toolName: "second",
|
||||
args: { value: 2 },
|
||||
},
|
||||
1_002,
|
||||
);
|
||||
// Parallel calls can finish in a different order than their result messages.
|
||||
recordToolExecutionTiming(
|
||||
timings,
|
||||
{
|
||||
type: "tool_execution_end",
|
||||
toolCallId: "call-b",
|
||||
toolName: "second",
|
||||
result: { content: [{ type: "text", text: "second result" }] },
|
||||
isError: false,
|
||||
},
|
||||
1_012,
|
||||
);
|
||||
recordToolExecutionTiming(
|
||||
timings,
|
||||
{
|
||||
type: "tool_execution_end",
|
||||
toolCallId: "call-a",
|
||||
toolName: "first",
|
||||
result: { content: [{ type: "text", text: "first result" }] },
|
||||
isError: false,
|
||||
},
|
||||
1_030,
|
||||
);
|
||||
|
||||
const transcript = transcriptFromMessages(parallelToolMessages(), timings);
|
||||
const toolEntries = transcript.filter((entry) => entry.role === "tool");
|
||||
const resultEntries = transcript.filter(
|
||||
(entry) => entry.role === "toolResult",
|
||||
);
|
||||
|
||||
for (const entries of [toolEntries, resultEntries]) {
|
||||
assert.deepEqual(
|
||||
entries.map(({ toolCallId, startedAt, finishedAt, durationMs }) => ({
|
||||
toolCallId,
|
||||
startedAt,
|
||||
finishedAt,
|
||||
durationMs,
|
||||
})),
|
||||
[
|
||||
{
|
||||
toolCallId: "call-a",
|
||||
startedAt: 1_000,
|
||||
finishedAt: 1_030,
|
||||
durationMs: 30,
|
||||
},
|
||||
{
|
||||
toolCallId: "call-b",
|
||||
startedAt: 1_002,
|
||||
finishedAt: 1_012,
|
||||
durationMs: 10,
|
||||
},
|
||||
],
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
test("in-flight aborted tool calls retain start timing without completion", () => {
|
||||
const timings = new Map<string, ToolExecutionTiming>();
|
||||
recordToolExecutionTiming(
|
||||
timings,
|
||||
{
|
||||
type: "tool_execution_start",
|
||||
toolCallId: "call-a",
|
||||
toolName: "first",
|
||||
args: { value: 1 },
|
||||
},
|
||||
2_000,
|
||||
);
|
||||
|
||||
const transcript = transcriptFromMessages(
|
||||
parallelToolMessages().slice(0, 2),
|
||||
timings,
|
||||
);
|
||||
const first = transcript.find((entry) => entry.toolCallId === "call-a");
|
||||
|
||||
assert.equal(first?.startedAt, 2_000);
|
||||
assert.equal(first?.finishedAt, undefined);
|
||||
assert.equal(first?.durationMs, undefined);
|
||||
assert.equal(
|
||||
transcript.some((entry) => entry.role === "toolResult"),
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
test("first-response watchdog aborts a silent provider request", async () => {
|
||||
let aborted = false;
|
||||
const watchdog = createFirstResponseWatchdog(
|
||||
async () => {
|
||||
aborted = true;
|
||||
},
|
||||
{ timeoutMs: 10, model: "fixture-model" },
|
||||
);
|
||||
|
||||
await assert.rejects(
|
||||
watchdog.waitFor(new Promise<never>(() => {})),
|
||||
/no assistant response event for fixture-model within 10 ms.*stalled/i,
|
||||
);
|
||||
assert.equal(aborted, true);
|
||||
});
|
||||
|
||||
test("first assistant response disarms the watchdog without limiting the run", async () => {
|
||||
const watchdog = createFirstResponseWatchdog(
|
||||
async () => {
|
||||
throw new Error("watchdog should have been disarmed");
|
||||
},
|
||||
{ timeoutMs: 10 },
|
||||
);
|
||||
watchdog.markResponse();
|
||||
|
||||
const result = await watchdog.waitFor(
|
||||
new Promise<string>((resolve) => setTimeout(() => resolve("done"), 20)),
|
||||
);
|
||||
assert.equal(result, "done");
|
||||
});
|
||||
|
||||
test("workflow children guard structured, normal, and dynamically registered tools", async () => {
|
||||
const structuredResult = {
|
||||
content: [{ type: "text" as const, text: "recorded" }],
|
||||
details: { value: "fixture" },
|
||||
terminate: true,
|
||||
};
|
||||
const structured = {
|
||||
name: "structured_output",
|
||||
label: "Structured Output",
|
||||
description: "fixture",
|
||||
parameters: Type.Object({}),
|
||||
async execute() {
|
||||
return structuredResult;
|
||||
},
|
||||
} satisfies ToolDefinition;
|
||||
const definitions = new Map<string, ToolDefinition>([
|
||||
[structured.name, structured],
|
||||
]);
|
||||
let listener: AgentSessionEventListener | undefined;
|
||||
const session = {
|
||||
getAllTools: () => [...definitions.keys()].map((name) => ({ name })),
|
||||
getToolDefinition: (name: string) => definitions.get(name),
|
||||
subscribe(next: AgentSessionEventListener) {
|
||||
listener = next;
|
||||
return () => {
|
||||
listener = undefined;
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const unsubscribe = guardWorkflowChildTools(session, 10);
|
||||
assert.equal(await structured.execute(), structuredResult);
|
||||
|
||||
let dynamicSignal: AbortSignal | undefined;
|
||||
const dynamic = {
|
||||
name: "dynamic_fixture",
|
||||
label: "Dynamic Fixture",
|
||||
description: "fixture",
|
||||
parameters: Type.Object({}),
|
||||
async execute(
|
||||
_toolCallId: string,
|
||||
_params: Record<string, never>,
|
||||
signal?: AbortSignal,
|
||||
) {
|
||||
dynamicSignal = signal;
|
||||
return new Promise<never>(() => {});
|
||||
},
|
||||
} satisfies ToolDefinition;
|
||||
const originalDynamicExecute = dynamic.execute;
|
||||
definitions.set(dynamic.name, dynamic);
|
||||
listener?.({ type: "agent_start" });
|
||||
assert.notEqual(dynamic.execute, originalDynamicExecute);
|
||||
|
||||
await assert.rejects(
|
||||
dynamic.execute("fixture", {}, undefined),
|
||||
/Tool call "dynamic_fixture" timed out after 10 ms\./,
|
||||
);
|
||||
assert.equal(dynamicSignal?.aborted, true);
|
||||
unsubscribe();
|
||||
});
|
||||
@@ -0,0 +1,650 @@
|
||||
/**
|
||||
* Workflow subagent runner.
|
||||
*
|
||||
* Each `agent()` call in a workflow script becomes one isolated in-process
|
||||
* AgentSession created here: in-memory session, normal trust-aware resources
|
||||
* and extensions, recursive orchestration/user-prompt tools denied, and an
|
||||
* optional one-shot `structured_output` tool when a schema is supplied.
|
||||
*
|
||||
* `runAgent()` never throws: every failure mode (session creation, provider
|
||||
* errors, aborts, missing structured output) settles into an `AgentOutcome`.
|
||||
*/
|
||||
|
||||
import {
|
||||
createAgentSession,
|
||||
DefaultResourceLoader,
|
||||
defineTool,
|
||||
SessionManager,
|
||||
SettingsManager,
|
||||
type AgentSession,
|
||||
type AgentSessionEvent,
|
||||
type AgentSessionEventListener,
|
||||
type ExtensionAPI,
|
||||
type ExtensionContext,
|
||||
type ToolDefinition,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import { Type, type TSchema } from "typebox";
|
||||
import {
|
||||
bindChildSessionExtensions,
|
||||
childToolPolicy,
|
||||
createChildResources,
|
||||
shutdownAndDisposeChildSession,
|
||||
} from "../shared/child-session.ts";
|
||||
import { createToolCallTimeoutGuard } from "../shared/tool-call-timeout.ts";
|
||||
import { emptyUsage, type AgentUsage, type TranscriptEntry } from "./model.ts";
|
||||
import {
|
||||
buildWorkflowAgentPrompt,
|
||||
STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION,
|
||||
STRUCTURED_OUTPUT_TOOL_DESCRIPTION,
|
||||
} from "./prompt.ts";
|
||||
import { safeStringify, truncateUtf8 } from "./serialization.ts";
|
||||
|
||||
const AGENT_OUTPUT_MAX_BYTES = 64 * 1024;
|
||||
export const FIRST_RESPONSE_TIMEOUT_MS = 45_000;
|
||||
const TRANSCRIPT_ENTRY_MAX_BYTES = 16 * 1024;
|
||||
const TRANSCRIPT_TOTAL_MAX_BYTES = 256 * 1024;
|
||||
const TRANSCRIPT_MAX_ENTRIES = 200;
|
||||
|
||||
export type WorkflowModel = NonNullable<ExtensionContext["model"]>;
|
||||
export type ThinkingLevel = ReturnType<ExtensionAPI["getThinkingLevel"]>;
|
||||
type AgentMessage = AgentSession["messages"][number];
|
||||
type ToolTimingEvent = Extract<
|
||||
AgentSessionEvent,
|
||||
{ type: "tool_execution_start" | "tool_execution_end" }
|
||||
>;
|
||||
|
||||
export interface ToolExecutionTiming {
|
||||
startedAt?: number;
|
||||
finishedAt?: number;
|
||||
durationMs?: number;
|
||||
}
|
||||
|
||||
export interface AgentOutcome {
|
||||
ok: boolean;
|
||||
/** Final assistant text (may be empty when only structured output was produced). */
|
||||
output: string;
|
||||
/** Captured structured_output payload when a schema was supplied. */
|
||||
structured?: unknown;
|
||||
error?: string;
|
||||
aborted: boolean;
|
||||
usage: AgentUsage;
|
||||
model?: string;
|
||||
contextWindow?: number;
|
||||
transcript: TranscriptEntry[];
|
||||
}
|
||||
|
||||
export interface AgentProgress {
|
||||
preview: string;
|
||||
usage: AgentUsage;
|
||||
model?: string;
|
||||
contextWindow?: number;
|
||||
transcript: TranscriptEntry[];
|
||||
}
|
||||
|
||||
export interface RunAgentOptions {
|
||||
prompt: string;
|
||||
schema?: unknown;
|
||||
model?: WorkflowModel;
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
cwd: string;
|
||||
loader: DefaultResourceLoader;
|
||||
settingsManager: SettingsManager;
|
||||
modelRegistry: ExtensionContext["modelRegistry"];
|
||||
signal?: AbortSignal;
|
||||
onProgress?: (progress: AgentProgress) => void;
|
||||
/** Test-only override for the per-tool execution timeout. */
|
||||
toolCallTimeoutMs?: number;
|
||||
/** Test-only override for the first assistant response-event timeout. */
|
||||
firstResponseTimeoutMs?: number;
|
||||
}
|
||||
|
||||
/** Build a fresh extension runtime for each concurrent workflow child. */
|
||||
export function createWorkflowResources(
|
||||
cwd: string,
|
||||
variant: "plain" | "structured",
|
||||
projectTrusted: boolean,
|
||||
) {
|
||||
return createChildResources({
|
||||
cwd,
|
||||
projectTrusted,
|
||||
...(variant === "structured"
|
||||
? { appendSystemPrompt: [STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION] }
|
||||
: {}),
|
||||
});
|
||||
}
|
||||
|
||||
interface WorkflowToolSession {
|
||||
getAllTools(): Array<{ name: string }>;
|
||||
getToolDefinition(name: string): ToolDefinition | undefined;
|
||||
subscribe(listener: AgentSessionEventListener): () => void;
|
||||
}
|
||||
|
||||
/** Guard current tools and tools registered by extensions at later agent starts. */
|
||||
export function guardWorkflowChildTools(
|
||||
session: WorkflowToolSession,
|
||||
timeoutMs?: number,
|
||||
) {
|
||||
const guard = createToolCallTimeoutGuard(timeoutMs);
|
||||
guard.apply(session);
|
||||
return session.subscribe((event) => {
|
||||
if (event.type === "agent_start") guard.apply(session);
|
||||
});
|
||||
}
|
||||
|
||||
function isJsonSchema(value: unknown): value is TSchema {
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) return false;
|
||||
const seen = new WeakSet<object>();
|
||||
let nodes = 0;
|
||||
const validate = (current: unknown, depth: number): boolean => {
|
||||
if (++nodes > 10_000 || depth > 24) return false;
|
||||
if (
|
||||
current === null ||
|
||||
typeof current === "string" ||
|
||||
typeof current === "boolean"
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
if (typeof current === "number") return Number.isFinite(current);
|
||||
if (Array.isArray(current)) {
|
||||
return current.every((item) => validate(item, depth + 1));
|
||||
}
|
||||
if (typeof current !== "object") return false;
|
||||
if (seen.has(current)) return false;
|
||||
seen.add(current);
|
||||
return Object.keys(current).every((key) => {
|
||||
if (key === "__proto__" || key === "constructor" || key === "prototype") {
|
||||
return false;
|
||||
}
|
||||
return validate((current as Record<string, unknown>)[key], depth + 1);
|
||||
});
|
||||
};
|
||||
return validate(value, 0);
|
||||
}
|
||||
|
||||
/** Preserve the caller's full JSON Schema instead of lossy keyword conversion. */
|
||||
function jsonSchemaToTypebox(schema: unknown): TSchema {
|
||||
if (!isJsonSchema(schema)) {
|
||||
throw new Error("structured output schema must be a bounded JSON object");
|
||||
}
|
||||
return Type.Unsafe(schema);
|
||||
}
|
||||
|
||||
/**
|
||||
* One-shot terminating tool injected when a schema is supplied: the subagent
|
||||
* calls it as its final action and we capture the validated object.
|
||||
*/
|
||||
function makeStructuredOutputTool(
|
||||
schema: unknown,
|
||||
capture: (value: unknown) => void,
|
||||
): ToolDefinition {
|
||||
return defineTool({
|
||||
name: "structured_output",
|
||||
label: "Structured Output",
|
||||
description: STRUCTURED_OUTPUT_TOOL_DESCRIPTION,
|
||||
parameters: jsonSchemaToTypebox(schema),
|
||||
async execute(_toolCallId, params) {
|
||||
capture(params);
|
||||
return {
|
||||
content: [{ type: "text", text: "Recorded structured result." }],
|
||||
details: params,
|
||||
terminate: true,
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function finalOutput(messages: AgentMessage[]): string {
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const msg = messages[i];
|
||||
if (msg.role !== "assistant") continue;
|
||||
const text = msg.content
|
||||
.filter((part) => part.type === "text")
|
||||
.map((part) => part.text)
|
||||
.join("\n")
|
||||
.trim();
|
||||
if (text) return text;
|
||||
}
|
||||
return "";
|
||||
}
|
||||
|
||||
function safeJson(value: unknown): string {
|
||||
return safeStringify(value, {
|
||||
maxBytes: TRANSCRIPT_ENTRY_MAX_BYTES,
|
||||
maxDepth: 12,
|
||||
maxNodes: 2_000,
|
||||
});
|
||||
}
|
||||
|
||||
/** Record lifecycle timings without inferring completion from message timestamps. */
|
||||
export function recordToolExecutionTiming(
|
||||
timings: Map<string, ToolExecutionTiming>,
|
||||
event: ToolTimingEvent,
|
||||
observedAt = Date.now(),
|
||||
) {
|
||||
const previous = timings.get(event.toolCallId);
|
||||
if (event.type === "tool_execution_start") {
|
||||
if (previous?.startedAt !== undefined) return;
|
||||
timings.set(event.toolCallId, { ...previous, startedAt: observedAt });
|
||||
return;
|
||||
}
|
||||
if (previous?.finishedAt !== undefined) return;
|
||||
const durationMs =
|
||||
previous?.startedAt === undefined
|
||||
? undefined
|
||||
: Math.max(0, observedAt - previous.startedAt);
|
||||
timings.set(event.toolCallId, {
|
||||
...previous,
|
||||
finishedAt: observedAt,
|
||||
...(durationMs === undefined ? {} : { durationMs }),
|
||||
});
|
||||
}
|
||||
|
||||
function toolMetadata(
|
||||
toolCallId: string,
|
||||
timings: ReadonlyMap<string, ToolExecutionTiming>,
|
||||
) {
|
||||
const timing = timings.get(toolCallId);
|
||||
return {
|
||||
toolCallId: truncateUtf8(toolCallId, 1024),
|
||||
...(timing?.startedAt === undefined ? {} : { startedAt: timing.startedAt }),
|
||||
...(timing?.finishedAt === undefined
|
||||
? {}
|
||||
: { finishedAt: timing.finishedAt }),
|
||||
...(timing?.durationMs === undefined
|
||||
? {}
|
||||
: { durationMs: timing.durationMs }),
|
||||
};
|
||||
}
|
||||
|
||||
/** Convert pi messages into a compact, serializable transcript for the UI. */
|
||||
export function transcriptFromMessages(
|
||||
messages: AgentMessage[],
|
||||
toolTimings: ReadonlyMap<string, ToolExecutionTiming> = new Map(),
|
||||
): TranscriptEntry[] {
|
||||
const entries: TranscriptEntry[] = [];
|
||||
for (const message of messages) {
|
||||
if (message.role === "user") {
|
||||
const text =
|
||||
typeof message.content === "string"
|
||||
? message.content
|
||||
: message.content
|
||||
.map((part) =>
|
||||
part.type === "text" ? part.text : `[image: ${part.mimeType}]`,
|
||||
)
|
||||
.join("\n");
|
||||
if (text.trim()) {
|
||||
entries.push({ role: "user", text, timestamp: message.timestamp });
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (message.role === "assistant") {
|
||||
for (const part of message.content) {
|
||||
if (part.type === "text" && part.text.trim()) {
|
||||
entries.push({
|
||||
role: "assistant",
|
||||
text: part.text,
|
||||
timestamp: message.timestamp,
|
||||
});
|
||||
} else if (part.type === "thinking" && part.thinking.trim()) {
|
||||
entries.push({
|
||||
role: "thinking",
|
||||
text: part.thinking,
|
||||
timestamp: message.timestamp,
|
||||
});
|
||||
} else if (part.type === "toolCall") {
|
||||
entries.push({
|
||||
role: "tool",
|
||||
name: part.name,
|
||||
text: safeJson(part.arguments),
|
||||
timestamp: message.timestamp,
|
||||
...toolMetadata(part.id, toolTimings),
|
||||
});
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (message.role !== "toolResult") continue;
|
||||
const text = message.content
|
||||
.map((part) =>
|
||||
part.type === "text" ? part.text : `[image: ${part.mimeType}]`,
|
||||
)
|
||||
.join("\n");
|
||||
entries.push({
|
||||
role: "toolResult",
|
||||
name: message.toolName,
|
||||
text,
|
||||
isError: message.isError,
|
||||
timestamp: message.timestamp,
|
||||
...toolMetadata(message.toolCallId, toolTimings),
|
||||
});
|
||||
}
|
||||
const selected =
|
||||
entries.length <= TRANSCRIPT_MAX_ENTRIES
|
||||
? entries
|
||||
: [entries[0], ...entries.slice(-(TRANSCRIPT_MAX_ENTRIES - 1))];
|
||||
const bounded: TranscriptEntry[] = [];
|
||||
let totalBytes = 0;
|
||||
for (const entry of selected) {
|
||||
const remaining = TRANSCRIPT_TOTAL_MAX_BYTES - totalBytes;
|
||||
if (remaining <= 0) break;
|
||||
const text = truncateUtf8(
|
||||
entry.text,
|
||||
Math.min(TRANSCRIPT_ENTRY_MAX_BYTES, remaining),
|
||||
);
|
||||
totalBytes += Buffer.byteLength(text, "utf8");
|
||||
bounded.push({
|
||||
...entry,
|
||||
text:
|
||||
text === entry.text ? text : `${text}\n[transcript entry truncated]`,
|
||||
});
|
||||
}
|
||||
if (bounded.length < entries.length) {
|
||||
bounded.push({
|
||||
role: "toolResult",
|
||||
name: "transcript",
|
||||
text: `[transcript truncated: retained ${bounded.length} of ${entries.length} entries]`,
|
||||
});
|
||||
}
|
||||
return bounded;
|
||||
}
|
||||
|
||||
function computeUsage(messages: AgentMessage[]): AgentUsage {
|
||||
const usage = emptyUsage();
|
||||
for (const msg of messages) {
|
||||
if (msg.role !== "assistant") continue;
|
||||
usage.turns++;
|
||||
const u = msg.usage;
|
||||
if (!u) continue;
|
||||
usage.input += u.input || 0;
|
||||
usage.output += u.output || 0;
|
||||
usage.cacheRead += u.cacheRead || 0;
|
||||
usage.cacheWrite += u.cacheWrite || 0;
|
||||
usage.cost += u.cost?.total || 0;
|
||||
}
|
||||
return usage;
|
||||
}
|
||||
|
||||
function errorText(error: unknown): string {
|
||||
return (error instanceof Error ? error.message : String(error)).slice(
|
||||
0,
|
||||
16 * 1024,
|
||||
);
|
||||
}
|
||||
|
||||
function formatTimeout(timeoutMs: number) {
|
||||
return timeoutMs % 1_000 === 0
|
||||
? `${timeoutMs / 1_000} seconds`
|
||||
: `${timeoutMs} ms`;
|
||||
}
|
||||
|
||||
/** Abort a provider call that opens but never emits its first assistant event. */
|
||||
export function createFirstResponseWatchdog(
|
||||
onTimeout: () => Promise<unknown>,
|
||||
options: { timeoutMs?: number; model?: string } = {},
|
||||
) {
|
||||
const timeoutMs = options.timeoutMs ?? FIRST_RESPONSE_TIMEOUT_MS;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const timeout = new Promise<never>((_resolve, reject) => {
|
||||
timer = setTimeout(() => {
|
||||
timer = undefined;
|
||||
const model = options.model ? ` for ${options.model}` : "";
|
||||
reject(
|
||||
new Error(
|
||||
`Agent received no assistant response event${model} within ${formatTimeout(timeoutMs)}; the provider request may be stalled. Retry the workflow.`,
|
||||
),
|
||||
);
|
||||
void onTimeout().catch(() => {});
|
||||
}, timeoutMs);
|
||||
timer.unref?.();
|
||||
});
|
||||
|
||||
const cancel = () => {
|
||||
if (timer) clearTimeout(timer);
|
||||
timer = undefined;
|
||||
};
|
||||
|
||||
return {
|
||||
markResponse: cancel,
|
||||
async waitFor<T>(operation: Promise<T>) {
|
||||
try {
|
||||
return await Promise.race([operation, timeout]);
|
||||
} finally {
|
||||
cancel();
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function isAssistantResponseEvent(event: AgentSessionEvent) {
|
||||
return (
|
||||
(event.type === "message_start" ||
|
||||
event.type === "message_update" ||
|
||||
event.type === "message_end") &&
|
||||
event.message.role === "assistant"
|
||||
);
|
||||
}
|
||||
|
||||
export async function runAgent(
|
||||
options: RunAgentOptions,
|
||||
): Promise<AgentOutcome> {
|
||||
let structured: unknown;
|
||||
let customTools: ToolDefinition[] | undefined;
|
||||
let session: AgentSession | undefined;
|
||||
let unsubscribeToolTimeout: (() => void) | undefined;
|
||||
try {
|
||||
customTools =
|
||||
options.schema !== undefined
|
||||
? [
|
||||
makeStructuredOutputTool(options.schema, (value) => {
|
||||
structured = value;
|
||||
}),
|
||||
]
|
||||
: undefined;
|
||||
({ session } = await createAgentSession({
|
||||
cwd: options.cwd,
|
||||
...(options.model ? { model: options.model } : {}),
|
||||
...(options.thinkingLevel
|
||||
? { thinkingLevel: options.thinkingLevel }
|
||||
: {}),
|
||||
resourceLoader: options.loader,
|
||||
settingsManager: options.settingsManager,
|
||||
sessionManager: SessionManager.inMemory(options.cwd),
|
||||
...(customTools ? { customTools } : {}),
|
||||
...childToolPolicy(),
|
||||
}));
|
||||
await bindChildSessionExtensions(session);
|
||||
unsubscribeToolTimeout = guardWorkflowChildTools(
|
||||
session,
|
||||
options.toolCallTimeoutMs,
|
||||
);
|
||||
} catch (error) {
|
||||
unsubscribeToolTimeout?.();
|
||||
if (session) await shutdownAndDisposeChildSession(session);
|
||||
return {
|
||||
ok: false,
|
||||
output: "",
|
||||
error: `Failed to create agent session: ${errorText(error)}`,
|
||||
aborted: false,
|
||||
usage: emptyUsage(),
|
||||
model: options.model?.id,
|
||||
contextWindow: options.model?.contextWindow,
|
||||
transcript: [],
|
||||
};
|
||||
}
|
||||
|
||||
const childSession = session;
|
||||
let usage = emptyUsage();
|
||||
let modelId = childSession.model?.id ?? options.model?.id;
|
||||
let contextWindow = childSession.model?.contextWindow;
|
||||
let stopReason: string | undefined;
|
||||
let errorMessage: string | undefined;
|
||||
const toolTimings = new Map<string, ToolExecutionTiming>();
|
||||
|
||||
const sync = () => {
|
||||
const messages = childSession.messages;
|
||||
usage = computeUsage(messages);
|
||||
|
||||
const sessionModel = childSession.model;
|
||||
modelId = sessionModel?.id ?? modelId;
|
||||
contextWindow = sessionModel?.contextWindow ?? contextWindow;
|
||||
const context = childSession.getContextUsage();
|
||||
if (
|
||||
typeof context?.tokens === "number" &&
|
||||
Number.isFinite(context.tokens) &&
|
||||
context.tokens >= 0
|
||||
) {
|
||||
usage.contextTokens = context.tokens;
|
||||
}
|
||||
if (
|
||||
typeof context?.contextWindow === "number" &&
|
||||
Number.isFinite(context.contextWindow) &&
|
||||
context.contextWindow > 0
|
||||
) {
|
||||
contextWindow = context.contextWindow;
|
||||
}
|
||||
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const msg = messages[i];
|
||||
if (msg.role !== "assistant") continue;
|
||||
// Some gateways report a concrete fallback model. Prefer its registry
|
||||
// metadata when available so capacity tracks the model that served the
|
||||
// latest response rather than a hardcoded/configured guess.
|
||||
const responseMatchesSession =
|
||||
!sessionModel ||
|
||||
(msg.provider === sessionModel.provider &&
|
||||
msg.model === sessionModel.id);
|
||||
const reportedId = msg.responseModel ?? msg.model;
|
||||
const reportedModel = responseMatchesSession
|
||||
? options.modelRegistry.find(msg.provider, reportedId)
|
||||
: undefined;
|
||||
if (reportedModel) {
|
||||
modelId = reportedModel.id;
|
||||
contextWindow = reportedModel.contextWindow;
|
||||
}
|
||||
if (msg.stopReason) stopReason = msg.stopReason;
|
||||
if (msg.errorMessage) errorMessage = msg.errorMessage;
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let markFirstResponse = () => {};
|
||||
const unsubscribe = childSession.subscribe((event) => {
|
||||
if (isAssistantResponseEvent(event)) markFirstResponse();
|
||||
if (
|
||||
event.type === "tool_execution_start" ||
|
||||
event.type === "tool_execution_end"
|
||||
) {
|
||||
recordToolExecutionTiming(toolTimings, event);
|
||||
} else if (
|
||||
event.type !== "message_end" &&
|
||||
event.type !== "compaction_end"
|
||||
) {
|
||||
return;
|
||||
}
|
||||
sync();
|
||||
options.onProgress?.({
|
||||
preview: finalOutput(childSession.messages),
|
||||
usage,
|
||||
model: modelId,
|
||||
contextWindow,
|
||||
transcript: transcriptFromMessages(childSession.messages, toolTimings),
|
||||
});
|
||||
});
|
||||
|
||||
let aborted = false;
|
||||
let abortPromise: Promise<void> | undefined;
|
||||
const onAbort = () => {
|
||||
aborted = true;
|
||||
abortPromise ??= childSession.abort().catch(() => {});
|
||||
};
|
||||
if (options.signal) {
|
||||
if (options.signal.aborted) onAbort();
|
||||
else options.signal.addEventListener("abort", onAbort, { once: true });
|
||||
}
|
||||
|
||||
let output = "";
|
||||
let transcript: TranscriptEntry[] = [];
|
||||
try {
|
||||
if (!aborted) {
|
||||
const watchdog = createFirstResponseWatchdog(() => childSession.abort(), {
|
||||
timeoutMs: options.firstResponseTimeoutMs,
|
||||
model: modelId,
|
||||
});
|
||||
markFirstResponse = watchdog.markResponse;
|
||||
await watchdog.waitFor(
|
||||
childSession.prompt(buildWorkflowAgentPrompt(options.prompt)),
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
errorMessage = errorMessage ?? errorText(error);
|
||||
stopReason = stopReason ?? "error";
|
||||
} finally {
|
||||
options.signal?.removeEventListener("abort", onAbort);
|
||||
if (abortPromise) await abortPromise;
|
||||
unsubscribe();
|
||||
unsubscribeToolTimeout?.();
|
||||
sync();
|
||||
output = truncateUtf8(
|
||||
finalOutput(childSession.messages),
|
||||
AGENT_OUTPUT_MAX_BYTES,
|
||||
);
|
||||
transcript = transcriptFromMessages(childSession.messages, toolTimings);
|
||||
await shutdownAndDisposeChildSession(childSession);
|
||||
}
|
||||
|
||||
if (aborted || stopReason === "aborted") {
|
||||
return {
|
||||
ok: false,
|
||||
output,
|
||||
structured,
|
||||
error: "Agent was aborted",
|
||||
aborted: true,
|
||||
usage,
|
||||
model: modelId,
|
||||
contextWindow,
|
||||
transcript,
|
||||
};
|
||||
}
|
||||
|
||||
const failed = stopReason === "error" || errorMessage !== undefined;
|
||||
if (failed) {
|
||||
return {
|
||||
ok: false,
|
||||
output,
|
||||
structured,
|
||||
error: errorMessage ?? "Agent failed",
|
||||
aborted: false,
|
||||
usage,
|
||||
model: modelId,
|
||||
contextWindow,
|
||||
transcript,
|
||||
};
|
||||
}
|
||||
|
||||
if (options.schema !== undefined && structured === undefined) {
|
||||
return {
|
||||
ok: false,
|
||||
output,
|
||||
error:
|
||||
"Agent finished without calling structured_output; no structured result matching the schema was produced.",
|
||||
aborted: false,
|
||||
usage,
|
||||
model: modelId,
|
||||
contextWindow,
|
||||
transcript,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
output,
|
||||
structured,
|
||||
aborted: false,
|
||||
usage,
|
||||
model: modelId,
|
||||
contextWindow,
|
||||
transcript,
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,274 @@
|
||||
"use strict";
|
||||
|
||||
// This file is launched by sandbox.ts in Node permission mode. It deliberately
|
||||
// has no filesystem/network/child-process permissions and receives workflow
|
||||
// source only over a validated IPC channel.
|
||||
const vm = require("node:vm");
|
||||
const sendIpc =
|
||||
typeof process.send === "function" ? process.send.bind(process) : undefined;
|
||||
// If a future V8 escape exposes `process`, remove the convenient bridges to
|
||||
// builtins, native bindings, parent signalling, and addons before any workflow
|
||||
// source is compiled. The parent still enforces the authenticated IPC protocol.
|
||||
for (const capability of [
|
||||
"getBuiltinModule",
|
||||
"binding",
|
||||
"_linkedBinding",
|
||||
"dlopen",
|
||||
"kill",
|
||||
"abort",
|
||||
"send",
|
||||
]) {
|
||||
try {
|
||||
Object.defineProperty(process, capability, {
|
||||
value: undefined,
|
||||
writable: false,
|
||||
configurable: false,
|
||||
});
|
||||
} catch {
|
||||
// The VM boundary and permission mode remain mandatory controls.
|
||||
}
|
||||
}
|
||||
|
||||
const BOOTSTRAP = String.raw`
|
||||
(function bootstrapWorkflowApi() {
|
||||
"use strict";
|
||||
const callHost = globalThis.__hostBridge;
|
||||
delete globalThis.__hostBridge;
|
||||
let nextRequestId = 0;
|
||||
const unconsumed = new Set();
|
||||
const inFlight = new Set();
|
||||
|
||||
function deepFreeze(value, depth = 0) {
|
||||
if (!value || typeof value !== "object" || depth > 32 || Object.isFrozen(value)) return value;
|
||||
Object.freeze(value);
|
||||
for (const key of Object.keys(value)) deepFreeze(value[key], depth + 1);
|
||||
return value;
|
||||
}
|
||||
|
||||
function requestAgent(promptValue, optionsValue = {}) {
|
||||
const id = ++nextRequestId;
|
||||
unconsumed.add(id);
|
||||
let started;
|
||||
const begin = () => {
|
||||
unconsumed.delete(id);
|
||||
if (!started) {
|
||||
let payload;
|
||||
try {
|
||||
payload = JSON.stringify({
|
||||
id,
|
||||
prompt: typeof promptValue === "string" ? promptValue : String(promptValue ?? ""),
|
||||
options: optionsValue && typeof optionsValue === "object" ? optionsValue : {},
|
||||
});
|
||||
} catch (error) {
|
||||
started = Promise.reject(new Error("agent() arguments must be serializable: " + error.message));
|
||||
return started;
|
||||
}
|
||||
inFlight.add(id);
|
||||
started = callHost("agent", payload)
|
||||
.then((json) => JSON.parse(json))
|
||||
.finally(() => inFlight.delete(id));
|
||||
}
|
||||
return started;
|
||||
};
|
||||
return Object.freeze({
|
||||
then(resolve, reject) {
|
||||
return begin().then(resolve, reject);
|
||||
},
|
||||
catch(reject) {
|
||||
return begin().catch(reject);
|
||||
},
|
||||
finally(callback) {
|
||||
return begin().finally(callback);
|
||||
},
|
||||
get [Symbol.toStringTag]() {
|
||||
return "Promise";
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function mapLimited(items, concurrency, invoke) {
|
||||
const results = new Array(items.length);
|
||||
let next = 0;
|
||||
const workers = Array.from({ length: Math.min(concurrency, items.length) }, async () => {
|
||||
while (true) {
|
||||
const index = next++;
|
||||
if (index >= items.length) return;
|
||||
results[index] = await invoke(items[index]);
|
||||
}
|
||||
});
|
||||
await Promise.all(workers);
|
||||
return results;
|
||||
}
|
||||
|
||||
async function parallel(items, options = {}) {
|
||||
if (!Array.isArray(items)) throw new Error("parallel() expects an array of zero-argument agent thunks");
|
||||
const requested = options && typeof options.concurrency === "number"
|
||||
? Math.floor(options.concurrency)
|
||||
: 4;
|
||||
if (!Number.isFinite(requested) || requested < 1) {
|
||||
throw new Error("parallel(): concurrency must be a positive integer");
|
||||
}
|
||||
const concurrency = Math.min(4, requested);
|
||||
return mapLimited(items, concurrency, (item) => {
|
||||
if (typeof item !== "function") {
|
||||
throw new Error("parallel() items must be zero-argument functions");
|
||||
}
|
||||
return item();
|
||||
});
|
||||
}
|
||||
|
||||
function phase(title) {
|
||||
callHost("phase", JSON.stringify({ title: String(title) }));
|
||||
}
|
||||
|
||||
const argsEnvelope = JSON.parse(globalThis.__argsJson);
|
||||
const args = argsEnvelope.defined ? deepFreeze(argsEnvelope.value) : undefined;
|
||||
delete globalThis.__argsJson;
|
||||
const stringify = JSON.stringify;
|
||||
function serializeResult(value) {
|
||||
const seen = new WeakSet();
|
||||
return stringify(value === undefined ? null : value, (_key, item) => {
|
||||
if (typeof item === "bigint") return item.toString() + "n";
|
||||
if (item && typeof item === "object") {
|
||||
if (seen.has(item)) return "[circular]";
|
||||
seen.add(item);
|
||||
}
|
||||
return item;
|
||||
});
|
||||
}
|
||||
Object.defineProperties(globalThis, {
|
||||
agent: { value: requestAgent, writable: false, configurable: false },
|
||||
parallel: { value: parallel, writable: false, configurable: false },
|
||||
phase: { value: phase, writable: false, configurable: false },
|
||||
args: { value: args, writable: false, configurable: false },
|
||||
__workflowCheck: {
|
||||
value: Object.freeze(() => ({
|
||||
unconsumed: unconsumed.size,
|
||||
inFlight: inFlight.size,
|
||||
})),
|
||||
writable: false,
|
||||
configurable: false,
|
||||
},
|
||||
__workflowSerialize: {
|
||||
value: Object.freeze(serializeResult),
|
||||
writable: false,
|
||||
configurable: false,
|
||||
},
|
||||
});
|
||||
})();
|
||||
`;
|
||||
|
||||
let initialized = false;
|
||||
let token;
|
||||
const pendingAgents = new Map();
|
||||
|
||||
function send(message) {
|
||||
sendIpc?.({ token, ...message });
|
||||
}
|
||||
|
||||
function fail(error) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
send({ kind: "error", error: message.slice(0, 16 * 1024) });
|
||||
}
|
||||
|
||||
process.on("message", (message) => {
|
||||
if (!message || typeof message !== "object") return;
|
||||
if (!initialized) {
|
||||
if (
|
||||
message.kind !== "init" ||
|
||||
typeof message.token !== "string" ||
|
||||
typeof message.source !== "string" ||
|
||||
typeof message.argsJson !== "string"
|
||||
) {
|
||||
process.exitCode = 1;
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
token = message.token;
|
||||
run(message.source, message.argsJson);
|
||||
return;
|
||||
}
|
||||
if (message.token !== token || message.kind !== "agentResult") return;
|
||||
const pending = pendingAgents.get(message.id);
|
||||
if (!pending) return;
|
||||
pendingAgents.delete(message.id);
|
||||
if (typeof message.resultJson === "string")
|
||||
pending.resolve(message.resultJson);
|
||||
else
|
||||
pending.reject(
|
||||
new Error(
|
||||
typeof message.error === "string" ? message.error : "Agent IPC failed",
|
||||
),
|
||||
);
|
||||
});
|
||||
|
||||
function run(source, argsJson) {
|
||||
try {
|
||||
const sandbox = Object.create(null);
|
||||
sandbox.__argsJson = argsJson;
|
||||
sandbox.__hostBridge = (kind, payloadJson) => {
|
||||
if (kind === "phase") {
|
||||
send({ kind: "phase", payloadJson });
|
||||
return undefined;
|
||||
}
|
||||
if (kind !== "agent")
|
||||
return Promise.reject(new Error("Unknown workflow operation"));
|
||||
let id;
|
||||
try {
|
||||
id = JSON.parse(payloadJson).id;
|
||||
} catch {
|
||||
return Promise.reject(new Error("Invalid agent request"));
|
||||
}
|
||||
return new Promise((resolve, reject) => {
|
||||
pendingAgents.set(id, { resolve, reject });
|
||||
send({ kind: "agent", payloadJson });
|
||||
});
|
||||
};
|
||||
|
||||
const context = vm.createContext(sandbox, {
|
||||
name: "pi-workflow",
|
||||
codeGeneration: { strings: false, wasm: false },
|
||||
});
|
||||
new vm.Script(BOOTSTRAP, {
|
||||
filename: "workflow-bootstrap.js",
|
||||
}).runInContext(context, { timeout: 1000 });
|
||||
const workflow = vm.compileFunction(
|
||||
`"use strict";\nreturn (async function workflow() {\n${source}\n})();`,
|
||||
["agent", "parallel", "phase", "args"],
|
||||
{ filename: "workflow-script.js", parsingContext: context },
|
||||
);
|
||||
context.__workflowBody = workflow;
|
||||
const invoke = `
|
||||
(() => {
|
||||
const workflowBody = globalThis.__workflowBody;
|
||||
delete globalThis.__workflowBody;
|
||||
globalThis.__workflowPromise = Promise.resolve(
|
||||
workflowBody(agent, parallel, phase, args),
|
||||
).then(async (value) => {
|
||||
await Promise.resolve();
|
||||
const pending = __workflowCheck();
|
||||
if (pending.unconsumed > 0) {
|
||||
throw new Error("Workflow created " + pending.unconsumed + " unawaited agent() call(s)");
|
||||
}
|
||||
if (pending.inFlight > 0) {
|
||||
throw new Error("Workflow returned before " + pending.inFlight + " agent call(s) settled");
|
||||
}
|
||||
return __workflowSerialize(value);
|
||||
});
|
||||
})();
|
||||
`;
|
||||
new vm.Script(invoke, { filename: "workflow-invoke.js" }).runInContext(
|
||||
context,
|
||||
{ timeout: 1000 },
|
||||
);
|
||||
Promise.resolve(context.__workflowPromise)
|
||||
.then((resultJson) => {
|
||||
if (typeof resultJson !== "string")
|
||||
throw new Error("Workflow result was not serializable");
|
||||
send({ kind: "result", resultJson });
|
||||
})
|
||||
.catch(fail);
|
||||
} catch (error) {
|
||||
fail(error);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { test } from "node:test";
|
||||
import { runWorkflowSandbox } from "./sandbox.ts";
|
||||
|
||||
function run(
|
||||
source: string,
|
||||
overrides: Partial<Parameters<typeof runWorkflowSandbox>[0]> = {},
|
||||
) {
|
||||
const abort = new AbortController();
|
||||
return runWorkflowSandbox({
|
||||
source,
|
||||
args: undefined,
|
||||
cwd: process.cwd(),
|
||||
signal: abort.signal,
|
||||
onAgent: async (prompt) => ({ ok: true, output: `reply:${prompt}` }),
|
||||
onPhase: () => {},
|
||||
...overrides,
|
||||
});
|
||||
}
|
||||
|
||||
test("sandbox exposes only workflow capabilities and validates results", async () => {
|
||||
const phases: string[] = [];
|
||||
const result = await run(
|
||||
`
|
||||
phase("Gather");
|
||||
const replies = await parallel([
|
||||
() => agent("one"),
|
||||
() => agent("two"),
|
||||
], { concurrency: 99 });
|
||||
return {
|
||||
replies: replies.map((reply) => reply.output),
|
||||
processType: typeof process,
|
||||
requireType: typeof require,
|
||||
fetchType: typeof fetch,
|
||||
};
|
||||
`,
|
||||
{ onPhase: (title) => phases.push(title) },
|
||||
);
|
||||
assert.deepEqual(result, {
|
||||
replies: ["reply:one", "reply:two"],
|
||||
processType: "undefined",
|
||||
requireType: "undefined",
|
||||
fetchType: "undefined",
|
||||
});
|
||||
assert.deepEqual(phases, ["Gather"]);
|
||||
});
|
||||
|
||||
test("sandbox result serialization handles cycles and bigint", async () => {
|
||||
const result = await run(`
|
||||
const value = { count: 7n };
|
||||
value.self = value;
|
||||
return value;
|
||||
`);
|
||||
assert.deepEqual(result, { count: "7n", self: "[circular]" });
|
||||
});
|
||||
|
||||
test("sandbox rejects unawaited agent calls", async () => {
|
||||
let calls = 0;
|
||||
await assert.rejects(
|
||||
run(`agent("orphan"); return "done";`, {
|
||||
onAgent: async () => {
|
||||
calls++;
|
||||
return { ok: true, output: "unexpected" };
|
||||
},
|
||||
}),
|
||||
/unawaited agent/,
|
||||
);
|
||||
assert.equal(calls, 0);
|
||||
});
|
||||
|
||||
test("sandbox source cannot escape the host accounting wrapper", async () => {
|
||||
let calls = 0;
|
||||
await assert.rejects(
|
||||
run(
|
||||
`}), agent("orphan"), Promise.resolve("bypass"); (async function () {`,
|
||||
{
|
||||
onAgent: async () => {
|
||||
calls++;
|
||||
return { ok: true, output: "unexpected" };
|
||||
},
|
||||
},
|
||||
),
|
||||
/unawaited agent/,
|
||||
);
|
||||
assert.equal(calls, 0);
|
||||
});
|
||||
|
||||
test("sandbox VM still rejects non-yielding synchronous code", async () => {
|
||||
await assert.rejects(run(`while (true) {}`), /timed out/);
|
||||
});
|
||||
|
||||
test("workflow agent invocations have no per-request wall timer", async () => {
|
||||
let signalAborted = false;
|
||||
const result = await run(`return (await agent("delayed")).output;`, {
|
||||
onAgent: async (_prompt, _options, signal) => {
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
signalAborted = signal.aborted;
|
||||
return { ok: true, output: "completed" };
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(result, "completed");
|
||||
assert.equal(signalAborted, false);
|
||||
});
|
||||
|
||||
test("workflow cancellation aborts a pending agent request", async () => {
|
||||
const controller = new AbortController();
|
||||
let startedResolve: (() => void) | undefined;
|
||||
const started = new Promise<void>((resolve) => {
|
||||
startedResolve = resolve;
|
||||
});
|
||||
let requestAborted = false;
|
||||
const pending = run(`return await agent("pending");`, {
|
||||
signal: controller.signal,
|
||||
onAgent: async (_prompt, _options, signal) => {
|
||||
startedResolve?.();
|
||||
await new Promise<void>((resolve) => {
|
||||
signal.addEventListener(
|
||||
"abort",
|
||||
() => {
|
||||
requestAborted = true;
|
||||
resolve();
|
||||
},
|
||||
{ once: true },
|
||||
);
|
||||
});
|
||||
return { ok: false, output: "", error: "Agent was aborted" };
|
||||
},
|
||||
});
|
||||
|
||||
await started;
|
||||
controller.abort(new Error("cancel fixture"));
|
||||
await assert.rejects(pending, /Workflow was aborted/);
|
||||
assert.equal(requestAborted, true);
|
||||
});
|
||||
@@ -0,0 +1,299 @@
|
||||
import { randomBytes } from "node:crypto";
|
||||
import { spawn, type ChildProcess } from "node:child_process";
|
||||
import * as path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { safeStringify, toSerializable } from "./serialization.ts";
|
||||
|
||||
const MAX_SOURCE_BYTES = 512 * 1024;
|
||||
const MAX_ARGS_BYTES = 256 * 1024;
|
||||
const MAX_RESULT_BYTES = 1024 * 1024;
|
||||
const MAX_AGENT_MESSAGE_BYTES = 512 * 1024;
|
||||
const MAX_AGENT_REQUESTS = 32;
|
||||
|
||||
export interface SandboxAgentOptions {
|
||||
label?: unknown;
|
||||
phase?: unknown;
|
||||
schema?: unknown;
|
||||
model?: unknown;
|
||||
provider?: unknown;
|
||||
effort?: unknown;
|
||||
}
|
||||
|
||||
export interface SandboxAgentResult {
|
||||
ok: boolean;
|
||||
output: string;
|
||||
structured?: unknown;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface RunWorkflowSandboxOptions {
|
||||
source: string;
|
||||
args: unknown;
|
||||
cwd: string;
|
||||
signal: AbortSignal;
|
||||
onAgent: (
|
||||
prompt: string,
|
||||
options: SandboxAgentOptions,
|
||||
signal: AbortSignal,
|
||||
) => Promise<SandboxAgentResult>;
|
||||
onPhase: (title: string) => void;
|
||||
}
|
||||
|
||||
function byteLength(value: string) {
|
||||
return Buffer.byteLength(value, "utf8");
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function errorText(error: unknown) {
|
||||
return error instanceof Error ? error.message : String(error);
|
||||
}
|
||||
|
||||
function terminateChild(child: ChildProcess) {
|
||||
if (child.exitCode !== null || child.signalCode !== null) return;
|
||||
child.kill("SIGTERM");
|
||||
const force = setTimeout(() => {
|
||||
if (child.exitCode === null && child.signalCode === null)
|
||||
child.kill("SIGKILL");
|
||||
}, 1_000);
|
||||
force.unref?.();
|
||||
}
|
||||
|
||||
function sanitizeAgentOptions(value: unknown): SandboxAgentOptions {
|
||||
if (!isRecord(value)) return {};
|
||||
return {
|
||||
...(value.label !== undefined ? { label: value.label } : {}),
|
||||
...(value.phase !== undefined ? { phase: value.phase } : {}),
|
||||
...(value.schema !== undefined ? { schema: value.schema } : {}),
|
||||
...(value.model !== undefined ? { model: value.model } : {}),
|
||||
...(value.provider !== undefined ? { provider: value.provider } : {}),
|
||||
...(value.effort !== undefined ? { effort: value.effort } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute orchestration code in a separate, permission-restricted Node process.
|
||||
* The child can only invoke the narrow agent/phase IPC protocol and is always
|
||||
* terminated on completion, cancellation, or protocol failure. The workflow
|
||||
* itself and its agent requests have no wall-clock deadline. Active requests
|
||||
* are aborted only when the workflow is cancelled or the sandbox is cleaned up.
|
||||
*/
|
||||
export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
|
||||
if (!process.allowedNodeEnvironmentFlags.has("--permission")) {
|
||||
return Promise.reject(
|
||||
new Error("This Node runtime cannot enforce workflow child permissions"),
|
||||
);
|
||||
}
|
||||
if (byteLength(options.source) > MAX_SOURCE_BYTES) {
|
||||
return Promise.reject(
|
||||
new Error(`Workflow script exceeds the ${MAX_SOURCE_BYTES} byte limit`),
|
||||
);
|
||||
}
|
||||
|
||||
const argsJson = safeStringify(
|
||||
{ defined: options.args !== undefined, value: options.args },
|
||||
{ maxBytes: MAX_ARGS_BYTES, maxDepth: 16, maxNodes: 10_000 },
|
||||
);
|
||||
if (byteLength(argsJson) > MAX_ARGS_BYTES) {
|
||||
return Promise.reject(new Error("Workflow args exceed the IPC limit"));
|
||||
}
|
||||
|
||||
return new Promise<unknown>((resolve, reject) => {
|
||||
const workerPath = fileURLToPath(
|
||||
new URL("./sandbox-child.cjs", import.meta.url),
|
||||
);
|
||||
const child = spawn(
|
||||
process.execPath,
|
||||
[
|
||||
"--permission",
|
||||
`--allow-fs-read=${path.dirname(workerPath)}`,
|
||||
"--max-old-space-size=128",
|
||||
"--stack-size=2048",
|
||||
workerPath,
|
||||
],
|
||||
{
|
||||
cwd: options.cwd,
|
||||
env: {
|
||||
PATH: process.env.PATH ?? "",
|
||||
NODE_NO_WARNINGS: "1",
|
||||
},
|
||||
stdio: ["ignore", "ignore", "ignore", "ipc"],
|
||||
},
|
||||
);
|
||||
const token = randomBytes(24).toString("hex");
|
||||
const requestIds = new Set<number>();
|
||||
const activeAgentRequests = new Map<number, AbortController>();
|
||||
let requestCount = 0;
|
||||
let finished = false;
|
||||
|
||||
const cleanup = () => {
|
||||
for (const abortController of activeAgentRequests.values()) {
|
||||
abortController.abort(new Error("Workflow stopped"));
|
||||
}
|
||||
activeAgentRequests.clear();
|
||||
options.signal.removeEventListener("abort", onAbort);
|
||||
child.removeAllListeners("message");
|
||||
child.removeAllListeners("error");
|
||||
child.removeAllListeners("exit");
|
||||
terminateChild(child);
|
||||
};
|
||||
const finish = (error?: Error, value?: unknown) => {
|
||||
if (finished) return;
|
||||
finished = true;
|
||||
cleanup();
|
||||
if (error) reject(error);
|
||||
else resolve(value);
|
||||
};
|
||||
const onAbort = () => finish(new Error("Workflow was aborted"));
|
||||
|
||||
options.signal.addEventListener("abort", onAbort, { once: true });
|
||||
if (options.signal.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
|
||||
child.on("error", (error) => finish(error));
|
||||
child.on("exit", (code, exitSignal) => {
|
||||
if (!finished) {
|
||||
finish(
|
||||
new Error(
|
||||
`Workflow sandbox exited before completion (${exitSignal ?? code ?? "unknown"})`,
|
||||
),
|
||||
);
|
||||
}
|
||||
});
|
||||
child.on("message", (raw: unknown) => {
|
||||
if (
|
||||
!isRecord(raw) ||
|
||||
raw.token !== token ||
|
||||
typeof raw.kind !== "string"
|
||||
) {
|
||||
finish(new Error("Workflow sandbox sent an invalid IPC message"));
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "phase") {
|
||||
if (
|
||||
typeof raw.payloadJson !== "string" ||
|
||||
byteLength(raw.payloadJson) > 4096
|
||||
) {
|
||||
finish(new Error("Workflow sandbox sent an invalid phase update"));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const payload: unknown = JSON.parse(raw.payloadJson);
|
||||
if (!isRecord(payload) || typeof payload.title !== "string") {
|
||||
throw new Error("invalid title");
|
||||
}
|
||||
options.onPhase(payload.title.slice(0, 160));
|
||||
} catch {
|
||||
finish(new Error("Workflow sandbox sent an invalid phase update"));
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "agent") {
|
||||
if (
|
||||
typeof raw.payloadJson !== "string" ||
|
||||
byteLength(raw.payloadJson) > MAX_AGENT_MESSAGE_BYTES
|
||||
) {
|
||||
finish(new Error("Workflow sandbox sent an oversized agent request"));
|
||||
return;
|
||||
}
|
||||
let payload: unknown;
|
||||
try {
|
||||
payload = JSON.parse(raw.payloadJson);
|
||||
} catch {
|
||||
finish(new Error("Workflow sandbox sent malformed agent JSON"));
|
||||
return;
|
||||
}
|
||||
if (
|
||||
!isRecord(payload) ||
|
||||
!Number.isSafeInteger(payload.id) ||
|
||||
typeof payload.id !== "number" ||
|
||||
payload.id < 1 ||
|
||||
typeof payload.prompt !== "string" ||
|
||||
payload.prompt.length > 100_000 ||
|
||||
!isRecord(payload.options)
|
||||
) {
|
||||
finish(new Error("Workflow sandbox sent an invalid agent request"));
|
||||
return;
|
||||
}
|
||||
if (requestIds.has(payload.id) || ++requestCount > MAX_AGENT_REQUESTS) {
|
||||
finish(
|
||||
new Error("Workflow sandbox exceeded its agent request budget"),
|
||||
);
|
||||
return;
|
||||
}
|
||||
requestIds.add(payload.id);
|
||||
const id = payload.id;
|
||||
const abortController = new AbortController();
|
||||
const sendResult = (result: SandboxAgentResult) => {
|
||||
if (!activeAgentRequests.delete(id)) return;
|
||||
if (finished || !child.connected) return;
|
||||
const normalized = toSerializable(result, {
|
||||
maxDepth: 16,
|
||||
maxNodes: 10_000,
|
||||
maxStringBytes: 128 * 1024,
|
||||
});
|
||||
let resultJson = JSON.stringify(normalized);
|
||||
if (byteLength(resultJson) > MAX_AGENT_MESSAGE_BYTES) {
|
||||
resultJson = JSON.stringify({
|
||||
ok: false,
|
||||
output: "",
|
||||
error: "Agent result exceeded the workflow IPC output limit",
|
||||
});
|
||||
}
|
||||
child.send({ token, kind: "agentResult", id, resultJson });
|
||||
};
|
||||
activeAgentRequests.set(id, abortController);
|
||||
void options
|
||||
.onAgent(
|
||||
payload.prompt,
|
||||
sanitizeAgentOptions(payload.options),
|
||||
abortController.signal,
|
||||
)
|
||||
.then(sendResult)
|
||||
.catch((error) =>
|
||||
sendResult({ ok: false, output: "", error: errorText(error) }),
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "result") {
|
||||
if (
|
||||
typeof raw.resultJson !== "string" ||
|
||||
byteLength(raw.resultJson) > MAX_RESULT_BYTES
|
||||
) {
|
||||
finish(new Error("Workflow result exceeded the IPC limit"));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const normalized = toSerializable(JSON.parse(raw.resultJson));
|
||||
finish(undefined, JSON.parse(JSON.stringify(normalized)));
|
||||
} catch (error) {
|
||||
finish(
|
||||
new Error(`Workflow returned invalid JSON: ${errorText(error)}`),
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "error" && typeof raw.error === "string") {
|
||||
finish(new Error(raw.error.slice(0, 16 * 1024)));
|
||||
return;
|
||||
}
|
||||
finish(new Error("Workflow sandbox sent an unknown IPC message"));
|
||||
});
|
||||
|
||||
child.send(
|
||||
{
|
||||
kind: "init",
|
||||
token,
|
||||
source: options.source,
|
||||
argsJson,
|
||||
},
|
||||
(error) => {
|
||||
if (error) finish(error);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { mkdtempSync, readFileSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { test } from "node:test";
|
||||
import { safeStringify, writeFileAtomic } from "./serialization.ts";
|
||||
|
||||
test("safeStringify handles cycles, bigint, depth, and size", () => {
|
||||
const value: Record<string, unknown> = {
|
||||
bigint: 42n,
|
||||
nested: { deeper: { deepest: true } },
|
||||
large: "x".repeat(20_000),
|
||||
};
|
||||
value.self = value;
|
||||
|
||||
const text = safeStringify(value, {
|
||||
maxBytes: 2_048,
|
||||
maxDepth: 2,
|
||||
maxStringBytes: 512,
|
||||
});
|
||||
assert.ok(Buffer.byteLength(text, "utf8") <= 2_048);
|
||||
const parsed: unknown = JSON.parse(text);
|
||||
assert.ok(parsed && typeof parsed === "object");
|
||||
assert.match(text, /42n/);
|
||||
assert.match(text, /circular/);
|
||||
assert.match(text, /truncated/);
|
||||
});
|
||||
|
||||
test("atomic writes leave complete readable content", () => {
|
||||
const directory = mkdtempSync(join(tmpdir(), "pi-workflow-test-"));
|
||||
try {
|
||||
const file = join(directory, "artifact.json");
|
||||
writeFileAtomic(file, '{"value":1}');
|
||||
writeFileAtomic(file, '{"value":2}');
|
||||
assert.deepEqual(JSON.parse(readFileSync(file, "utf8")), { value: 2 });
|
||||
} finally {
|
||||
rmSync(directory, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,162 @@
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
|
||||
export interface SerializationOptions {
|
||||
maxBytes?: number;
|
||||
maxDepth?: number;
|
||||
maxNodes?: number;
|
||||
maxStringBytes?: number;
|
||||
}
|
||||
|
||||
const DEFAULT_MAX_BYTES = 1024 * 1024;
|
||||
const DEFAULT_MAX_DEPTH = 16;
|
||||
const DEFAULT_MAX_NODES = 20_000;
|
||||
const DEFAULT_MAX_STRING_BYTES = 64 * 1024;
|
||||
|
||||
function byteLength(value: string) {
|
||||
return Buffer.byteLength(value, "utf8");
|
||||
}
|
||||
|
||||
export function truncateUtf8(value: string, maxBytes: number) {
|
||||
if (maxBytes <= 0) return "";
|
||||
if (byteLength(value) <= maxBytes) return value;
|
||||
const buffer = Buffer.from(value, "utf8");
|
||||
let end = Math.min(maxBytes, buffer.length);
|
||||
while (end > 0 && (buffer[end] & 0xc0) === 0x80) end--;
|
||||
return buffer.subarray(0, end).toString("utf8");
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize arbitrary values to inert JSON data. Cycles, bigint, non-finite
|
||||
* numbers, deep trees, throwing properties, and very large strings are all
|
||||
* represented explicitly instead of making artifact persistence fail.
|
||||
*/
|
||||
export function toSerializable(
|
||||
value: unknown,
|
||||
options: SerializationOptions = {},
|
||||
): unknown {
|
||||
const maxDepth = options.maxDepth ?? DEFAULT_MAX_DEPTH;
|
||||
const maxNodes = options.maxNodes ?? DEFAULT_MAX_NODES;
|
||||
const maxStringBytes = options.maxStringBytes ?? DEFAULT_MAX_STRING_BYTES;
|
||||
const seen = new WeakMap<object, string>();
|
||||
let nodes = 0;
|
||||
|
||||
const visit = (
|
||||
current: unknown,
|
||||
depth: number,
|
||||
location: string,
|
||||
): unknown => {
|
||||
nodes++;
|
||||
if (nodes > maxNodes) return "[truncated: node limit]";
|
||||
if (depth > maxDepth) return "[truncated: depth limit]";
|
||||
if (current === null || typeof current === "boolean") return current;
|
||||
if (typeof current === "string") {
|
||||
if (byteLength(current) <= maxStringBytes) return current;
|
||||
return `${truncateUtf8(current, maxStringBytes)}\n[truncated: string limit]`;
|
||||
}
|
||||
if (typeof current === "number") {
|
||||
return Number.isFinite(current)
|
||||
? current
|
||||
: `[number: ${String(current)}]`;
|
||||
}
|
||||
if (typeof current === "bigint") return `${current.toString()}n`;
|
||||
if (typeof current === "undefined") return "[undefined]";
|
||||
if (typeof current === "symbol")
|
||||
return `[symbol: ${current.description ?? ""}]`;
|
||||
if (typeof current === "function")
|
||||
return `[function: ${current.name || "anonymous"}]`;
|
||||
if (typeof current !== "object") return String(current);
|
||||
|
||||
const prior = seen.get(current);
|
||||
if (prior) return `[circular: ${prior}]`;
|
||||
seen.set(current, location);
|
||||
|
||||
if (Array.isArray(current)) {
|
||||
return current.map((item, index) =>
|
||||
visit(item, depth + 1, `${location}[${index}]`),
|
||||
);
|
||||
}
|
||||
|
||||
if (current instanceof Date) {
|
||||
return Number.isNaN(current.getTime())
|
||||
? "[date: invalid]"
|
||||
: current.toISOString();
|
||||
}
|
||||
if (current instanceof Error) {
|
||||
return {
|
||||
name: current.name,
|
||||
message: current.message,
|
||||
...(current.stack
|
||||
? { stack: truncateUtf8(current.stack, 16 * 1024) }
|
||||
: {}),
|
||||
};
|
||||
}
|
||||
|
||||
const result: Record<string, unknown> = Object.create(null);
|
||||
let keys: string[];
|
||||
try {
|
||||
keys = Object.keys(current);
|
||||
} catch (error) {
|
||||
return `[unreadable object: ${error instanceof Error ? error.message : String(error)}]`;
|
||||
}
|
||||
for (const key of keys) {
|
||||
try {
|
||||
result[key] = visit(
|
||||
(current as Record<string, unknown>)[key],
|
||||
depth + 1,
|
||||
`${location}.${key}`,
|
||||
);
|
||||
} catch (error) {
|
||||
result[key] =
|
||||
`[unreadable property: ${error instanceof Error ? error.message : String(error)}]`;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
return visit(value, 0, "$root");
|
||||
}
|
||||
|
||||
/** Serialize to valid JSON no larger than the requested cap. */
|
||||
export function safeStringify(
|
||||
value: unknown,
|
||||
options: SerializationOptions = {},
|
||||
) {
|
||||
const maxBytes = Math.max(256, options.maxBytes ?? DEFAULT_MAX_BYTES);
|
||||
const normalized = toSerializable(value, options);
|
||||
const serialized = JSON.stringify(normalized, null, 2) ?? "null";
|
||||
if (byteLength(serialized) <= maxBytes) return serialized;
|
||||
|
||||
let previewBytes = Math.max(32, Math.floor(maxBytes / 3));
|
||||
while (previewBytes > 0) {
|
||||
const fallback = JSON.stringify(
|
||||
{
|
||||
truncated: true,
|
||||
reason: `serialized value exceeded ${maxBytes} bytes`,
|
||||
preview: truncateUtf8(serialized, previewBytes),
|
||||
},
|
||||
null,
|
||||
2,
|
||||
);
|
||||
if (byteLength(fallback) <= maxBytes) return fallback;
|
||||
previewBytes = Math.floor(previewBytes / 2);
|
||||
}
|
||||
return JSON.stringify({ truncated: true });
|
||||
}
|
||||
|
||||
/** Durable same-directory replace: readers see either the old or new file. */
|
||||
export function writeFileAtomic(filePath: string, content: string) {
|
||||
fs.mkdirSync(path.dirname(filePath), { recursive: true });
|
||||
const temporary = `${filePath}.${process.pid}.${Date.now()}.tmp`;
|
||||
try {
|
||||
fs.writeFileSync(temporary, content, { encoding: "utf8", mode: 0o600 });
|
||||
fs.renameSync(temporary, filePath);
|
||||
} catch (error) {
|
||||
try {
|
||||
fs.unlinkSync(temporary);
|
||||
} catch {
|
||||
// The original write error is more useful.
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user