@tshu-w/pi-agents 0.1.1 → 0.2.1

package/src/programs/programs.ts

@@ -1,12 +1,11 @@
 import { randomUUID } from "node:crypto";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { deferred, settlesBy, waitForChange, type Deferred } from "../agents/agents.ts";
+import { deferred, settlesBy, STEER_POLL_MS, WAIT_INTERRUPTED, waitForChange, type Deferred } from "../agents/agents.ts";
 import { shortId, trackTreeWork } from "../agents/registry.ts";
-import { boundBlocks, boundText } from "../output.ts";
-import type { ProgramFiles, ProgramOutcome, ProgramRunResult } from "./execute.ts";
+import { boundBlocks, boundText, type Truncation } from "../output.ts";
+import { CLEANUP_TIMEOUT_MS, PROGRAM_TOOL_NAME, type ProgramFiles, type ProgramOutcome, type ProgramRunResult } from "./execute.ts";
 
 const PROGRAM_ENTRY = "pi-agents-program";
-const CLEANUP_TIMEOUT_MS = 15_000;
 /** Ends the result of a Program some of whose work did not stop. */
 export const CLEANUP_ERROR = "Cleanup error: some of the Program's work did not stop within the timeout and may still be running.";
 
@@ -22,13 +21,25 @@ export interface ProgramRecord {
 	isError?: boolean;
 	/** Files the Program's tool calls used, for compaction. */
 	files?: ProgramFiles;
+	/** The tool calls its result lists, which the task viewer shows after a restart. */
+	calls?: unknown[];
 	/** Whether `wait` has returned the result. */
 	returned: boolean;
 	notified: boolean;
 }
 
+/** What the task viewer shows of a Program, as its tool call shows in the foreground. */
+export interface ProgramView {
+	record: ProgramRecord;
+	code?: string;
+	/** The latest partial result while it runs, then its result. */
+	result?: ProgramRunResult["result"];
+}
+
 interface Running {
 	record: ProgramRecord;
+	code?: string;
+	result?: ProgramRunResult["result"];
 	controller?: AbortController;
 	done?: Promise<void>;
 	/** Stops counting the Program as the tree's work. */
@@ -41,6 +52,8 @@ interface Running {
 export interface ProgramWaitOutcome {
 	results: ProgramRecord[];
 	running: ProgramRecord[];
+	/** A steer ended the wait. */
+	interrupted?: boolean;
 }
 
 function ended(record: ProgramRecord): boolean {
@@ -58,6 +71,7 @@ export class Programs {
 	private programs = new Map<string, Running>();
 	private change: Deferred = deferred();
 	private closing = false;
+	private viewers = new Set<() => void>();
 
 	constructor(
 		private rootId: string,
@@ -70,11 +84,27 @@ export class Programs {
 	) {}
 
 	restore(ctx: ExtensionContext): void {
+		// A Program's code is in the `run` call that started it, whose result names the Program.
+		const codes = new Map<string, string>();
+		const callIds = new Map<string, string>();
 		for (const entry of ctx.sessionManager.getEntries()) {
+			if (entry.type === "message" && entry.message.role === "assistant") {
+				for (const part of entry.message.content) {
+					const code = part.type === "toolCall" && part.name === PROGRAM_TOOL_NAME ? (part.arguments as { code?: unknown }).code : undefined;
+					if (typeof code === "string") codes.set((part as { id: string }).id, code);
+				}
+			} else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === PROGRAM_TOOL_NAME) {
+				const id = (entry.message.details as { id?: unknown } | undefined)?.id;
+				if (typeof id === "string") callIds.set(id, entry.message.toolCallId);
+			}
 			if (entry.type !== "custom" || entry.customType !== PROGRAM_ENTRY) continue;
 			const record = entry.data as ProgramRecord;
 			this.programs.set(record.id, { record, waiters: 0 });
 		}
+		for (const program of this.programs.values()) {
+			const callId = callIds.get(program.record.id);
+			program.code = callId === undefined ? undefined : codes.get(callId);
+		}
 		for (const program of this.programs.values()) {
 			const { record } = program;
 			if (ended(record)) continue;
@@ -86,19 +116,25 @@ export class Programs {
 		}
 	}
 
-	/** Starts a background Program, which `run` executes with its ID, and returns its record. */
-	start(run: (id: string, signal: AbortSignal, abort: () => void) => Promise<ProgramRunResult>): ProgramRecord {
+	/** Starts a background Program, which `run` executes with its ID and reports its partial results to `update`, and returns its record. */
+	start(code: string, run: (id: string, signal: AbortSignal, abort: () => void, update: (result: unknown) => void) => Promise<ProgramRunResult>): ProgramRecord {
 		const record: ProgramRecord = { id: randomUUID(), state: "running", startedAt: Date.now(), returned: false, notified: false };
 		const controller = new AbortController();
-		const program: Running = { record, controller, waiters: 0 };
+		const program: Running = { record, code, controller, waiters: 0 };
+		const update = (result: unknown) => {
+			if (ended(record)) return;
+			program.result = result as Running["result"];
+			this.viewed();
+		};
 		this.programs.set(record.id, program);
 		// The tree stays busy until the Program ends and its notification is delivered.
 		const release = program.release = trackTreeWork(this.rootId);
 		this.persist(record);
 		this.hooks.changed?.();
-		program.done = run(record.id, controller.signal, () => controller.abort(new Error("Program stopped"))).then(
+		program.done = run(record.id, controller.signal, () => controller.abort(new Error("Program stopped")), update).then(
 			({ outcome, result, files, abandoned }) => {
 				if (abandoned) program.abandoned = true;
+				if (!ended(record)) program.result = result;
 				this.end(program, outcome, result.content, result.isError === true, files);
 			},
 			(error: unknown) => this.end(program, "failed", [{ type: "text", text: `Script failed\nOutput:\nScript error:\n${error instanceof Error ? error.message : String(error)}` }], true),
@@ -131,6 +167,20 @@ export class Programs {
 		return [...this.programs.values()].map((program) => program.record).sort((a, b) => b.startedAt - a.startedAt);
 	}
 
+	view(id: string): ProgramView | undefined {
+		const program = this.programs.get(id);
+		if (!program) return undefined;
+		const { record, code } = program;
+		const restored = record.content && { content: record.content, details: record.calls && { calls: record.calls }, isError: record.isError };
+		return { record, code, result: program.result ?? restored as ProgramView["result"] };
+	}
+
+	/** Calls `listener` when a Program's view changes; returns a function that stops it. */
+	onView(listener: () => void): () => void {
+		this.viewers.add(listener);
+		return () => this.viewers.delete(listener);
+	}
+
 	/**
 	 * Returns whether the Program was running; it has ended when this resolves. A Program that has
 	 * not stopped within the timeout is abandoned and ends as stopped; this throws when some of its
@@ -138,6 +188,8 @@ export class Programs {
 	 */
 	async stop(program: Running, reason = "Program stopped"): Promise<boolean> {
 		if (ended(program.record)) return false;
+		// As with `abort`, the caller that stops a Program is not notified.
+		program.record.notified = true;
 		program.controller?.abort(new Error(reason));
 		if (!await settlesBy(program.done!, Date.now() + CLEANUP_TIMEOUT_MS)) {
 			program.abandoned = true;
@@ -157,22 +209,28 @@ export class Programs {
 		return programs.some(({ record }) => !ended(record));
 	}
 
-	async wait(unique: Running[], timeoutSeconds: number, signal?: AbortSignal): Promise<ProgramWaitOutcome> {
+	async wait(unique: Running[], timeoutSeconds: number, signal?: AbortSignal, steered?: () => boolean): Promise<ProgramWaitOutcome> {
 		for (const program of unique) program.waiters += 1;
 		let returned = false;
+		let interrupted = false;
 		try {
 			const deadline = Date.now() + timeoutSeconds * 1000;
 			while (unique.some(({ record }) => !ended(record))) {
 				const remaining = deadline - Date.now();
 				if (remaining <= 0) break;
+				if (steered?.()) {
+					interrupted = true;
+					break;
+				}
 				signal?.throwIfAborted();
-				await waitForChange(this.change.promise, remaining, signal);
+				await waitForChange(this.change.promise, steered ? Math.min(remaining, STEER_POLL_MS) : remaining, signal);
 			}
 			signal?.throwIfAborted();
 			returned = true;
 			return {
 				results: unique.filter(({ record }) => ended(record)).map(({ record }) => record),
 				running: unique.filter(({ record }) => !ended(record)).map(({ record }) => record),
+				...(interrupted && { interrupted }),
 			};
 		} finally {
 			for (const program of unique) {
@@ -212,6 +270,10 @@ export class Programs {
 		record.state = outcome;
 		record.content = content;
 		record.isError = isError;
+		// Keeps the calls the result's details list, from the result or the last partial one.
+		program.result = { content, details: program.result?.details, isError };
+		const calls = (program.result.details as { calls?: unknown[] } | undefined)?.calls;
+		if (calls?.length) record.calls = calls;
 		if (files && (files.read.length > 0 || files.written.length > 0 || files.edited.length > 0)) record.files = files;
 		program.controller = undefined;
 		this.persist(record);
@@ -220,6 +282,11 @@ export class Programs {
 		this.change = deferred();
 		previous.resolve();
 		this.hooks.changed?.();
+		this.viewed();
+	}
+
+	private viewed(): void {
+		for (const listener of this.viewers) listener();
 	}
 
 	private notify(record: ProgramRecord): void {
@@ -260,23 +327,64 @@ function resultBlock(record: ProgramRecord, ids: string[]): string {
 	return `<program-result id="${shortId(record.id, ids)}" status="${record.state}">\n${textOf(record.content)}\n</program-result>`;
 }
 
-/** Renders wait results; returns the content and the results the caller has now received. */
-export function renderProgramWait({ results, running }: ProgramWaitOutcome, ids: string[]): { content: Item[]; returned: ProgramRecord[]; isError?: boolean } {
-	const runningLine = running.length > 0 ? `Still running: ${running.map((record) => shortId(record.id, ids)).join(", ")}` : "";
-	if (results.length === 0) return { content: [{ type: "text", text: runningLine || "No results." }], returned: [] };
+// The header codemode puts before a Program's output; a Program that ended before its run did has no wall time.
+const SCRIPT_HEADER = /^Script (?:completed|failed)\n(?:Wall time ([\d.]+) seconds\n)?Output:\n/;
+
+export interface ProgramResultDetails {
+	id: string;
+	state: ProgramState;
+	wallTimeSeconds?: number;
+	/** The output the result shows, without the header. */
+	output?: string;
+}
+
+/** A result's details, with the output in the first `available` characters of its text, if shown. */
+function resultDetails(record: ProgramRecord, available?: number): ProgramResultDetails {
+	const text = textOf(record.content);
+	const header = SCRIPT_HEADER.exec(text);
+	return {
+		id: record.id,
+		state: record.state,
+		...(header?.[1] !== undefined && { wallTimeSeconds: Number(header[1]) }),
+		...(available !== undefined && { output: text.slice(0, available).slice(header?.[0].length ?? 0) }),
+	};
+}
+
+/** Renders wait results; returns the content, its details, and the results the caller has now received. */
+/** The line naming the results left out to keep within Pi's output limits. */
+export function omittedLine(omitted: string[], ids: string[]): string {
+	return omitted.length > 0 ? `[Results omitted: ${omitted.map((id) => shortId(id, ids)).join(", ")}. Use wait with fewer targets.]` : "";
+}
+
+/** The line naming the Programs still running, and that a steer ended the wait if one did. */
+export function runningLine(running: string[], interrupted: boolean | undefined, ids: string[]): string {
+	return running.length > 0 ? `Still running: ${running.map((id) => shortId(id, ids)).join(", ")}${interrupted ? `\n${WAIT_INTERRUPTED}` : ""}` : "";
+}
+
+export function renderProgramWait({ results, running, interrupted }: ProgramWaitOutcome, ids: string[]): { content: Item[]; details: { results: ProgramResultDetails[] } & Partial<Truncation>; returned: ProgramRecord[]; isError?: boolean } {
+	const runningText = runningLine(running.map((record) => record.id), interrupted, ids);
+	if (results.length === 0) return { content: [{ type: "text", text: runningText || "No results." }], details: { results: [] }, returned: [] };
 	const images = (shown: ProgramRecord[]) => shown.flatMap((record) => (record.content ?? []).filter((item) => item.type === "image"));
 	const only = results[0]!;
 	if (results.length === 1 && running.length === 0 && only.state === "completed") {
 		const bounded = boundText(textOf(only.content), "pi-agents-program-wait");
-		return { content: [{ type: "text", text: bounded.text }, ...images(results)], returned: results };
+		return { content: [{ type: "text", text: bounded.text }, ...images(results)], details: { results: [resultDetails(only)], ...bounded.details }, returned: results };
 	}
-	const bounded = boundBlocks(results.map((record) => resultBlock(record, ids)), "program-result", "pi-agents-program-wait");
+	const blocks = results.map((record) => resultBlock(record, ids));
+	const bounded = boundBlocks(blocks, "program-result", "pi-agents-program-wait");
+	let offset = 0;
+	const details = results.map((record, index) => {
+		const block = blocks[index]!;
+		const available = Math.max(0, bounded.kept.length - offset - block.indexOf("\n") - 1);
+		offset += block.length + 2;
+		return resultDetails(record, index < bounded.shown ? available : undefined);
+	});
 	let text = bounded.text;
 	const returned = results.slice(0, bounded.shown);
-	const omitted = results.slice(bounded.shown);
-	if (omitted.length > 0) text += `\n\n[Results omitted: ${omitted.map((record) => shortId(record.id, ids)).join(", ")}. Use wait with fewer targets.]`;
-	if (runningLine) text += `\n\n${runningLine}`;
-	return { content: [{ type: "text", text }, ...images(returned)], returned };
+	const omittedText = omittedLine(results.slice(bounded.shown).map((record) => record.id), ids);
+	if (omittedText) text += `\n\n${omittedText}`;
+	if (runningText) text += `\n\n${runningText}`;
+	return { content: [{ type: "text", text }, ...images(returned)], details: { results: details, ...bounded.details }, returned };
 }
 
 export function programListLine(record: ProgramRecord, ids: string[]): string {

package/src/programs/run-output.mjs

@@ -0,0 +1,11 @@
+/** The message a `/program` run adds to the Session: `Ran …`, then the Program's result. */
+export const PROGRAM_MESSAGE = 'pi-agents-program-run';
+
+const SCRIPT_HEADER = /^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n$/;
+
+/** What `pi -p` prints for a Program result: its text items without codemode's header. */
+export function runOutput(content) {
+  const [first, ...rest] = content;
+  const output = first?.type === 'text' && SCRIPT_HEADER.test(first.text) ? rest : content;
+  return output.flatMap(item => item.type === 'text' ? [item.text] : []).join('\n');
+}

package/src/programs/sandbox.ts

@@ -1,5 +1,6 @@
 import { randomBytes } from "node:crypto";
 import type { CodemodeTool } from "@earendil-works/pi-codemode";
+import type { CallTracker } from "./execute.ts";
 
 export interface AgentOptions {
 	name?: string;
@@ -12,6 +13,7 @@ export interface AgentOptions {
 export interface SendOptions {
 	deliverAs?: "followUp" | "steer" | "write";
 	schema?: unknown;
+	timeout?: number;
 }
 
 /** What `agent()` and its handles do on the host. */
@@ -53,9 +55,10 @@ export function withAgentPrefix(code: string): string {
 	return `${agentPrefix()}${code}`;
 }
 
-/** Globals behind `agent()`. A handle's calls wait for its Agent to be created. */
-export function agentGlobals(host: ProgramAgentHost): CodemodeTool[] {
+/** Globals behind `agent()`. A handle's calls wait for its Agent to be created; `track` records each `send()`. */
+export function agentGlobals(host: ProgramAgentHost, track: CallTracker): CodemodeTool[] {
 	const created = new Map<string, Promise<void>>();
+	const names = new Map<string, string>();
 	const ready = (id: string) => created.get(id) ?? Promise.reject(new Error(`Unknown Agent ${id}`));
 	return [
 		{
@@ -63,6 +66,7 @@ export function agentGlobals(host: ProgramAgentHost): CodemodeTool[] {
 			spread: true,
 			execute: (args) => {
 				const [id, name, options] = args as [string, string, AgentOptions];
+				names.set(id, name);
 				const creation = host.create(id, name, options ?? {});
 				created.set(id, creation);
 				return creation;
@@ -73,8 +77,17 @@ export function agentGlobals(host: ProgramAgentHost): CodemodeTool[] {
 			spread: true,
 			execute: async (args) => {
 				const [id, message, options] = args as [string, string, SendOptions | undefined];
-				await ready(id);
-				return host.send(id, message, options ?? {});
+				// Recorded from the start, so a send whose Agent fails to be created shows too.
+				const end = track(`agent ${names.get(id) ?? id}`, JSON.stringify(message) ?? "");
+				try {
+					await ready(id);
+					const value = await host.send(id, message, options ?? {});
+					end("ok");
+					return value;
+				} catch (error) {
+					end((error as { cancelled?: boolean }).cancelled ? "cancelled" : "error", error instanceof Error ? error.message : String(error));
+					throw error;
+				}
 			},
 		},
 		{

package/src/programs/scope.ts

@@ -1,9 +1,9 @@
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { Agents, settlesBy, SUBMIT_RESULT_TOOL_NAME, THINKING_LEVELS, type Limits } from "../agents/agents.ts";
-import { label, messageText, nodes, treeChanged, visibleIds, type AgentNode } from "../agents/registry.ts";
+import { label, messageText, nodes, sender, treeChanged, visibleIds, type AgentNode } from "../agents/registry.ts";
+import { CLEANUP_TIMEOUT_MS } from "./execute.ts";
 import type { AgentOptions, ProgramAgentHost, SendOptions } from "./sandbox.ts";
 
-const CLEANUP_TIMEOUT_MS = 15_000;
 const DELIVERIES = ["followUp", "steer", "write"];
 
 export interface ProgramScope {
@@ -32,6 +32,7 @@ export function programScope(
 		program: true,
 		background: options.background,
 		createdAt: new Date().toISOString(),
+		spent: { cost: 0, tokens: 0, unpriced: false },
 		agents,
 		cwd: () => ctx.cwd,
 		busy: () => true,
@@ -54,7 +55,7 @@ export function programScope(
 		},
 		send: async (agentId, message, sendOptions: SendOptions) => {
 			const delivery = sendOptions.deliverAs ?? "followUp";
-			const { schema } = sendOptions;
+			const { schema, timeout } = sendOptions;
 			if (typeof message !== "string" || !message.trim()) throw new Error("send() requires a non-empty message");
 			if (!DELIVERIES.includes(delivery)) throw new Error('send() deliverAs must be "followUp", "steer", or "write"');
 			if (schema !== undefined) {
@@ -63,16 +64,23 @@ export function programScope(
 					throw new Error("send() schema must be a JSON Schema object or boolean");
 				}
 			}
+			if (timeout !== undefined) {
+				if (delivery !== "followUp") throw new Error('send() accepts timeout only with deliverAs "followUp"');
+				if (typeof timeout !== "number" || !(timeout > 0) || !Number.isFinite(timeout)) throw new Error("send() timeout must be a positive number of seconds");
+			}
 			const agent = agents.ownedTarget(agentId).record;
 			const body = schema === undefined ? message : `${message}\n\nWhen done, call \`${SUBMIT_RESULT_TOOL_NAME}\` with a \`value\` that matches this JSON Schema:\n${JSON.stringify(schema)}`;
-			const text = messageText(from, true, delivery, body, visibleIds(agentId, false));
+			const ids = visibleIds(agentId, false);
+			const text = messageText(from, true, delivery, body, ids);
 			if (delivery === "write") {
-				agents.accept(agentId, text, "write", { fromOwner: true, notification: false });
+				agents.accept(agentId, text, "write", { fromOwner: true, notification: false, sender: sender(from, ids) });
 				return undefined;
 			}
-			const record = await agents.request(agentId, text, delivery, schema);
+			const { record, timedOut } = await agents.request(agentId, text, delivery, sender(from, ids), schema, timeout === undefined ? undefined : timeout * 1000);
 			if (record.state === "completed") return schema === undefined ? record.result : record.value;
-			throw new Error(record.state === "failed" ? `Agent ${label(agent, agents.ids())} failed: ${record.result}` : `Agent ${label(agent, agents.ids())} was aborted.`);
+			if (record.state === "failed") throw new Error(`Agent ${label(agent, agents.ids())} failed: ${record.result}`);
+			// An aborted input, timed out or not, shows as a cancelled call.
+			throw Object.assign(new Error(timedOut ? `Agent ${label(agent, agents.ids())} timed out after ${timeout} seconds.` : `Agent ${label(agent, agents.ids())} was aborted.`), { cancelled: true });
 		},
 		abort: async (agentId) => {
 			await agents.abort(agents.ownedTarget(agentId));