@tshu-w/pi-agents 0.1.1 → 0.2.1

package/src/programs/execute.ts

@@ -10,6 +10,8 @@ import { loadCodemode, loadPiCodemode } from "./codemode.ts";
 import type { DetachedCallOptions } from "./host-patches.ts";
 
 export const PROGRAM_TOOL_NAME = "program";
+/** How long a Program's work may take to stop once it ends. */
+export const CLEANUP_TIMEOUT_MS = 15_000;
 
 export type ProgramOutcome = "completed" | "failed" | "stopped";
 
@@ -20,28 +22,70 @@ export interface ProgramFiles {
 	edited: string[];
 }
 
+/** A call in codemode's `details.calls`. */
+interface CallRecord {
+	id: string;
+	name: string;
+	args: string;
+	status: "running" | "ok" | "error" | "cancelled";
+	durationMs?: number;
+	error?: string;
+}
+
+/**
+ * Records calls that a Program's own globals make, which codemode does not see, beside its tool calls.
+ * `start` returns the function that ends the call; a call still running when the script ends is cancelled.
+ */
+export type CallTracker = (name: string, args: string) => (status: "ok" | "error" | "cancelled", error?: string) => void;
+
+const ARGS_PREVIEW_CHARS = 200;
+const ERROR_PREVIEW_CHARS = 500;
+const NOTE_NAMES = 10;
+const NOTE_NAME_CHARS = 60;
+
+const preview = (text: string, max: number) => text.length > max ? `${text.slice(0, max - 3)}...` : text;
+
+/** Names of calls for a note: at most `NOTE_NAMES` of them, each cut to `NOTE_NAME_CHARS`, since scripts name Agents. */
+function callNames(calls: CallRecord[], status: boolean): string {
+	const shown = calls.slice(0, NOTE_NAMES).map((call) => {
+		const name = preview(call.name, NOTE_NAME_CHARS);
+		return status ? `${name} (${call.status})` : name;
+	});
+	return calls.length > NOTE_NAMES ? `${shown.join(", ")}, and ${calls.length - NOTE_NAMES} more` : shown.join(", ");
+}
+
+/** Adds a note after a blank line to the last text item, or as a new item after images. */
+function appendNote(content: AgentToolResult<unknown>["content"], note: string): void {
+	const last = content.at(-1);
+	if (last?.type === "text") last.text += `\n\n${note}`;
+	else content.push({ type: "text", text: note });
+}
+
 export interface ProgramRunResult {
 	outcome: ProgramOutcome;
 	result: AgentToolResult<unknown> & { isError?: boolean };
 	files: ProgramFiles;
 	/** Whether some of the Program's work did not stop during cleanup. */
 	abandoned?: boolean;
+	/** Settles once the tool calls the script started have ended, which may be after it ends. */
+	toolCalls: Promise<unknown>;
 }
 
 export interface ProgramRunOptions {
 	/** Seconds until the Program stops; the script's `timeout_ms` option applies as well. */
 	timeout?: number;
-	/** Globals added to the script, such as the host side of `agent()`. */
-	globals?: CodemodeTool[];
+	/** Globals added to the script, such as the host side of `agent()`; `track` records their calls. */
+	globals?: (track: CallTracker) => CodemodeTool[];
 	/** Rewrites the code before it runs, such as prepending the `agent()` prelude. */
 	prepare?: (code: string) => string;
 	onUpdate?: AgentToolUpdateCallback<unknown>;
 	/**
-	 * Whether the Program runs after its tool call returns. Its calls then get their own IDs,
-	 * since Pi records nested calls only until the calling tool call returns.
+	 * Whether the Program runs outside a tool call: after its tool call returns, or from `/program`.
+	 * Its calls then get their own IDs, since Pi records nested calls only until the calling tool
+	 * call returns.
 	 */
-	background?: boolean;
-	/** Stops a background Program; `ctx.abort()` of its calls' hooks and tools calls it. */
+	detached?: boolean;
+	/** Stops a detached Program; `ctx.abort()` of its calls' hooks and tools calls it. */
 	abort?: () => void;
 	appendEntry(customType: string, data: unknown): void;
 	getToolNamespace(name: string): { name: string } | undefined;
@@ -53,10 +97,32 @@ interface SandboxExtras {
 	timeoutMs: number;
 	result?: CodemodeResult;
 	used: boolean;
+	/** Called when the script ends, before the sandbox closes. */
+	ended(): void;
 }
 
 const sandboxExtras = new AsyncLocalStorage<SandboxExtras>();
-let sandboxPatched = false;
+let sandboxPatch: Promise<void> | undefined;
+
+async function patchSandbox(): Promise<void> {
+	const proto = (await loadCodemode()).CodemodeSandbox.prototype as unknown as Record<string, any>;
+	const execute = proto.execute;
+	proto.execute = function (this: Record<string, any>, code: string, options: Record<string, unknown> = {}) {
+		const current = sandboxExtras.getStore();
+		if (!current || current.used) return execute.call(this, code, options);
+		current.used = true;
+		for (const global of current.globals) this.globalsByName.set(global.name, global);
+		const timeoutMs = Math.min(current.timeoutMs, (options.timeoutMs as number | undefined) ?? this.timeoutMs);
+		return execute.call(this, current.prepare(code), { ...options, timeoutMs }).then((result: CodemodeResult) => {
+			current.result = result;
+			current.ended();
+			return result;
+		}, (error: unknown) => {
+			current.ended();
+			throw error;
+		});
+	};
+}
 
 /**
  * Pi's executor creates and runs one sandbox per script. Its first execution in `extras`' async
@@ -64,22 +130,8 @@ let sandboxPatched = false;
  * do not.
  */
 async function withSandboxExtras<T>(extras: SandboxExtras, run: () => Promise<T>): Promise<T> {
-	if (!sandboxPatched) {
-		sandboxPatched = true;
-		const proto = (await loadCodemode()).CodemodeSandbox.prototype as unknown as Record<string, any>;
-		const execute = proto.execute;
-		proto.execute = function (this: Record<string, any>, code: string, options: Record<string, unknown> = {}) {
-			const current = sandboxExtras.getStore();
-			if (!current || current.used) return execute.call(this, code, options);
-			current.used = true;
-			for (const global of current.globals) this.globalsByName.set(global.name, global);
-			const timeoutMs = Math.min(current.timeoutMs, (options.timeoutMs as number | undefined) ?? this.timeoutMs);
-			return execute.call(this, current.prepare(code), { ...options, timeoutMs }).then((result: CodemodeResult) => {
-				current.result = result;
-				return result;
-			});
-		};
-	}
+	// Programs started together wait for the same patch.
+	await (sandboxPatch ??= patchSandbox());
 	return sandboxExtras.run(extras, run);
 }
 
@@ -102,30 +154,94 @@ export async function executeProgram(
 		if (typeof path === "string" && list && !list.includes(path)) list.push(path);
 	};
 	let detachedCalls = 0;
+	const toolCalls: Promise<unknown>[] = [];
+	const executeTool = (name: string, args: unknown, callOptions: { signal?: AbortSignal }) => {
+		if (!options.detached) return ctx.executeTool(name, args, callOptions);
+		addFile(name, args);
+		const detachedOptions = { ...callOptions, detachedCallerId: `${toolCallId}:${++detachedCalls}`, detachedAbort: () => options.abort?.() } satisfies DetachedCallOptions;
+		return ctx.executeTool(name, args, detachedOptions as never);
+	};
 	const programCtx = Object.create(ctx, {
 		tools: { get: () => ctx.tools.filter((tool) => tool.name !== PROGRAM_TOOL_NAME) },
 		executeTool: {
 			value: (name: string, args: unknown, callOptions: { signal?: AbortSignal } = {}) => {
-				if (!options.background) return ctx.executeTool(name, args, callOptions);
-				addFile(name, args);
-				const detachedOptions = { ...callOptions, detachedCallerId: `${toolCallId}:${++detachedCalls}`, detachedAbort: () => options.abort?.() } satisfies DetachedCallOptions;
-				return ctx.executeTool(name, args, detachedOptions as never);
+				const call = executeTool(name, args, callOptions);
+				toolCalls.push(call);
+				return call;
 			},
 		},
 	}) as ExtensionToolContext;
+	// The globals' calls, each placed after the codemode calls seen when it started.
+	const own: Array<{ record: CallRecord; after: number; startedAt: number }> = [];
+	let seen: CallRecord[] = [];
+	let last: Parameters<AgentToolUpdateCallback<unknown>>[0] = { content: [], details: undefined };
+	let scriptEnded = false;
+	const merged = (): CallRecord[] => {
+		const calls: CallRecord[] = [];
+		let next = 0;
+		for (const { record, after } of own) {
+			while (next < after && next < seen.length) calls.push({ ...seen[next++]! });
+			calls.push({ ...record });
+		}
+		while (next < seen.length) calls.push({ ...seen[next++]! });
+		return calls;
+	};
+	const publish = () => options.onUpdate?.({ ...last, details: { ...(last.details as object | undefined), calls: merged() } });
+	const onUpdate: AgentToolUpdateCallback<unknown> = (update) => {
+		last = update;
+		seen = (update.details as { calls?: CallRecord[] } | undefined)?.calls ?? seen;
+		publish();
+	};
+	const finish = (entry: (typeof own)[number], status: "ok" | "error" | "cancelled", error?: string) => {
+		if (entry.record.status !== "running") return;
+		entry.record.status = status;
+		entry.record.durationMs = performance.now() - entry.startedAt;
+		if (error !== undefined) entry.record.error = preview(error, ERROR_PREVIEW_CHARS);
+		publish();
+	};
+	const track: CallTracker = (name, args) => {
+		const entry = { record: { id: `${toolCallId}/${name}-${own.length + 1}`, name, args: preview(args, ARGS_PREVIEW_CHARS), status: "running" as const }, after: seen.length, startedAt: performance.now() };
+		own.push(entry);
+		if (scriptEnded) finish(entry, "cancelled");
+		else publish();
+		return (status, error) => finish(entry, status, error);
+	};
+	// Calls still running when the script ends are cut off by it.
+	let cutOff: CallRecord[] = [];
 	const extras: SandboxExtras = {
-		globals: options.globals ?? [],
+		globals: options.globals?.(track) ?? [],
 		prepare: options.prepare ?? ((code) => code),
 		timeoutMs: options.timeout === undefined ? Number.POSITIVE_INFINITY : options.timeout * 1000,
 		used: false,
+		ended: () => {
+			scriptEnded = true;
+			const running = own.filter((entry) => entry.record.status === "running");
+			for (const entry of running) finish(entry, "cancelled");
+			cutOff = running.map((entry) => entry.record);
+		},
 	};
-	const result = await withSandboxExtras(extras, () => executeCodemode(toolCallId, { code: source }, signal, options.onUpdate, programCtx, {
-		models: true,
-		appendEntry: options.appendEntry,
-		getToolNamespace: options.getToolNamespace,
-	}));
+	let result: Awaited<ReturnType<typeof executeCodemode>>;
+	try {
+		result = await withSandboxExtras(extras, () => executeCodemode(toolCallId, { code: source }, signal, onUpdate, programCtx, {
+			models: true,
+			appendEntry: options.appendEntry,
+			getToolNamespace: options.getToolNamespace,
+		}));
+	} finally {
+		if (!scriptEnded) extras.ended();
+	}
 	const outcome: ProgramOutcome = extras.result?.ok
 		? "completed"
 		: extras.result?.error.kind === "timeout" || extras.result?.error.kind === "aborted" ? "stopped" : "failed";
-	return { outcome, result, files };
+	// codemode marks calls cut off by the script ending only in its result, after the last update.
+	seen = (result.details as { calls?: CallRecord[] } | undefined)?.calls ?? seen;
+	if (own.length > 0) result.details = { ...(result.details as object | undefined), calls: merged() };
+	// codemode lists the tool calls made before a failure; the globals' calls get a line of their own.
+	if (outcome !== "completed" && own.length > 0) {
+		appendNote(result.content, `Agent calls made before the failure (they are not undone): ${callNames(own.map((entry) => entry.record), true)}`);
+	} else if (outcome === "completed") {
+		const stopped = [...seen.filter((call) => call.status === "cancelled"), ...cutOff];
+		if (stopped.length > 0) appendNote(result.content, `Calls still running when the script returned were cancelled; their effects are not undone: ${callNames(stopped, false)}`);
+	}
+	return { outcome, result, files, toolCalls: Promise.allSettled(toolCalls) };
 }

package/src/programs/host-patches.ts

@@ -16,7 +16,10 @@ interface DetachedCall {
 	active: boolean;
 }
 
-const detached = new AsyncLocalStorage<DetachedCall>();
+// These two are shared with later copies of this module, such as after a reload, since the patches stay those of the first copy.
+const detached = ((globalThis as Record<symbol, unknown>)[Symbol.for("pi-agents:detached")] ??= new AsyncLocalStorage<DetachedCall>()) as AsyncLocalStorage<DetachedCall>;
+/** The Session whose extension command is running, including work the command started. */
+const inCommand = ((globalThis as Record<symbol, unknown>)[Symbol.for("pi-agents:in-command")] ??= new AsyncLocalStorage<object>()) as AsyncLocalStorage<object>;
 /** Set while the built-in MCP extension runs, including work it started. */
 const inMcp = new AsyncLocalStorage<true>();
 
@@ -30,11 +33,54 @@ export interface DetachedCallOptions {
 
 type Patchable = Record<PropertyKey, any>;
 
+/** The Session running the current extension command. */
+export function commandSession(): Patchable | undefined {
+	return inCommand.getStore() as Patchable | undefined;
+}
+
+/** Runs `work` as a detached call: `ctx.signal` and `ctx.abort()` of the Session's hooks and tools act on `signal` and `abort`. */
+export async function runDetached<T>(session: Patchable, signal: AbortSignal | undefined, abort: (() => void) | undefined, work: () => Promise<T>): Promise<T> {
+	const extensions = session._extensionRunner as Patchable;
+	const getSignal = extensions.getSignalFn;
+	if (!getSignal[PATCHED]) {
+		const patched = () => {
+			const call = detached.getStore();
+			return call?.active ? call.signal : getSignal();
+		};
+		(patched as Patchable)[PATCHED] = true;
+		extensions.getSignalFn = patched;
+	}
+	const abortTurn = extensions.abortFn;
+	if (!abortTurn[PATCHED]) {
+		const patched = () => {
+			const call = detached.getStore();
+			return call?.active ? call.abort?.() : abortTurn();
+		};
+		(patched as Patchable)[PATCHED] = true;
+		extensions.abortFn = patched;
+	}
+	const call: DetachedCall = { signal, abort, active: true };
+	try {
+		return await detached.run(call, work);
+	} finally {
+		// Work the call started, such as a turn it queued, keeps the store but not the signal.
+		call.active = false;
+	}
+}
+
+/** Stands in for the assistant message Pi passes to hooks of nested calls, which a command's Session may lack. */
+function commandMessage() {
+	const usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } };
+	return { role: "assistant", content: [], api: "none", provider: "none", model: "none", usage, stopReason: "stop", timestamp: Date.now() };
+}
+
 /**
  * - A call with `detachedCallerId` runs under that ID: Pi numbers nested calls per caller and
  *   resets the numbers when the calling tool call returns and when the turn ends, which would
  *   repeat IDs across the calls of a background Program. While it runs, `ctx.signal` and
  *   `ctx.abort()` of hooks and tools act on the call instead of the caller's turn.
+ * - An extension command knows its Session through `commandSession()`, and its detached calls run
+ *   before the Session has an assistant message, which Pi requires of nested calls.
  * - The built-in MCP extension sees `program` as `codemode`: it activates `program` when its tools
  *   need a script tool and neither is active, counts an active `program` as reaching its tools,
  *   and makes a Program wait for the servers its code names, as it does for `codemode` scripts.
@@ -47,36 +93,19 @@ export function installHostPatches(
 	const sessionProto = session.prototype as unknown as Patchable;
 	if (!sessionProto[PATCHED]) {
 		sessionProto[PATCHED] = true;
+		const tryCommand = sessionProto._tryExecuteExtensionCommand;
+		sessionProto._tryExecuteExtensionCommand = function (this: Patchable, text: string) {
+			return inCommand.run(this, () => tryCommand.call(this, text));
+		};
+		const findAssistant = sessionProto._findLastAssistantMessage;
+		sessionProto._findLastAssistantMessage = function (this: Patchable) {
+			return findAssistant.call(this) ?? (inCommand.getStore() === this && detached.getStore()?.active ? commandMessage() : undefined);
+		};
 		const execute = sessionProto._executeNestedToolCall;
 		sessionProto._executeNestedToolCall = async function (this: Patchable, callerId: string, name: string, args: unknown, options: Record<string, unknown> = {}) {
 			const { detachedCallerId, detachedAbort, ...rest } = options;
 			if (typeof detachedCallerId !== "string") return execute.call(this, callerId, name, args, options);
-			const extensions = this._extensionRunner as Patchable;
-			const getSignal = extensions.getSignalFn;
-			if (!getSignal[PATCHED]) {
-				const patched = () => {
-					const call = detached.getStore();
-					return call?.active ? call.signal : getSignal();
-				};
-				(patched as Patchable)[PATCHED] = true;
-				extensions.getSignalFn = patched;
-			}
-			const abort = extensions.abortFn;
-			if (!abort[PATCHED]) {
-				const patched = () => {
-					const call = detached.getStore();
-					return call?.active ? call.abort?.() : abort();
-				};
-				(patched as Patchable)[PATCHED] = true;
-				extensions.abortFn = patched;
-			}
-			const call: DetachedCall = { signal: rest.signal as AbortSignal | undefined, abort: detachedAbort as (() => void) | undefined, active: true };
-			try {
-				return await detached.run(call, () => execute.call(this, detachedCallerId, name, args, rest));
-			} finally {
-				// Work the call started, such as a turn it queued, keeps the store but not the signal.
-				call.active = false;
-			}
+			return runDetached(this, rest.signal as AbortSignal | undefined, detachedAbort as (() => void) | undefined, () => execute.call(this, detachedCallerId, name, args, rest));
 		};
 	}
 	const runnerProto = runner.prototype as unknown as Patchable;

package/src/programs/loadout.ts

@@ -2,7 +2,10 @@ import { highlightCode, type ExtensionAPI, type ToolDefinition } from "@earendil
 import { Text } from "@earendil-works/pi-tui";
 import { CODEMODE_TOOL_NAME, type PiCodemode } from "./codemode.ts";
 import { PROGRAM_TOOL_NAME } from "./execute.ts";
-import { formatToolCall, moreLines, renderTextResult, startDuration, type DurationState } from "../render-call.ts";
+import { shortId } from "../agents/registry.ts";
+import { formatDuration, formatToolCall, moreLines, renderResultList, renderTextResult, startDuration, type DurationState } from "../render-call.ts";
+import type { Truncation } from "../output.ts";
+import { omittedLine, runningLine, type ProgramResultDetails } from "./programs.ts";
 
 /** The first line of the `codemode` description, which the `program` intro replaces. */
 const CODEMODE_FIRST_LINE = "Run JavaScript code to orchestrate/compose tool calls\n";
@@ -91,9 +94,21 @@ export function programRenderers(piCodemode: PiCodemode): Renderers {
 			component.setText(text);
 			return component;
 		},
-		// `wait` shows how long it has waited; the other results render like `codemode`.
+		// `wait` shows how long it has waited, and when collapsed, each result as a line with the first lines of its output;
+		// the other results render like `codemode`.
 		renderResult: (result, options, theme, context) => (context.args as { action?: string }).action === "wait"
-			? renderTextResult(result, options, theme, context)
+			? renderTextResult(result, options, theme, context, () => {
+				const { results = [], running = [], interrupted, ...truncated } = (result.details ?? {}) as { results?: ProgramResultDetails[]; running?: string[]; interrupted?: boolean } & Partial<Truncation>;
+				const only = results[0];
+				// As `renderProgramWait` returns them, no results, or a single completed result alone, show as they are.
+				if (!only || (results.length === 1 && running.length === 0 && only.state === "completed")) return undefined;
+				const ids = [...results.map((entry) => entry.id), ...running];
+				const notes = [omittedLine(results.flatMap((entry) => entry.output === undefined ? [entry.id] : []), ids), runningLine(running, interrupted, ids)];
+				return renderResultList(results.flatMap(({ id, state, wallTimeSeconds, output }) => output === undefined ? [] : [{
+					line: [`Program ${shortId(id, ids)}`, state, ...(wallTimeSeconds === undefined ? [] : [formatDuration(wallTimeSeconds * 1000)])].join("  "),
+					body: output,
+				}]), notes, truncated, theme);
+			})
 			: piCodemode.renderResult(result, options, theme, context) as Text,
 	};
 }