@tshu-w/pi-agents 0.1.1 → 0.2.0

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 { 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,6 +64,10 @@ 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));
@@ -70,9 +75,11 @@ export function programScope(
 				agents.accept(agentId, text, "write", { fromOwner: true, notification: false });
 				return undefined;
 			}
-			const record = await agents.request(agentId, text, delivery, schema);
+			const { record, timedOut } = await agents.request(agentId, text, delivery, 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));

package/src/programs/tool.ts

@@ -1,15 +1,22 @@
-import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
 import { StringEnum } from "@earendil-works/pi-ai";
+import { Box } from "@earendil-works/pi-tui";
+import { writeSync } from "node:fs";
+import { readFile } from "node:fs/promises";
+import * as path from "node:path";
 import { Type } from "typebox";
-import { SUBMIT_RESULT_TOOL_NAME, type Limits } from "../agents/agents.ts";
-import { nodes, type AgentNode } from "../agents/registry.ts";
+import { settlesBy, steerQueued, SUBMIT_RESULT_TOOL_NAME, type Limits } from "../agents/agents.ts";
+import { nodes, trackTreeWork, type AgentNode } from "../agents/registry.ts";
 import { rejectFields } from "../params.ts";
+import { agentSession } from "../roots/rpc.ts";
 import { programScope } from "./scope.ts";
 import type { PiCodemode } from "./codemode.ts";
-import { executeProgram, PROGRAM_TOOL_NAME, type ProgramRunOptions } from "./execute.ts";
+import { CLEANUP_TIMEOUT_MS, executeProgram, PROGRAM_TOOL_NAME, type ProgramRunOptions } from "./execute.ts";
+import { commandSession, runDetached } from "./host-patches.ts";
 import { programDescription, programLoadout, programRenderers } from "./loadout.ts";
 import { randomUUID } from "node:crypto";
 import { CLEANUP_ERROR, programListLine, Programs, renderProgramWait } from "./programs.ts";
+import { PROGRAM_MESSAGE, runOutput } from "./run-output.mjs";
 import { agentGlobals, withAgentPrefix } from "./sandbox.ts";
 
 const PROGRAM_DESCRIPTION = `Run JavaScript that composes tool calls and Agents; only its output and return value reach the caller. \`run\` runs a Program in the foreground, or in the background with \`background\`; \`wait\` waits for background Programs and returns their results; \`list\` lists background Programs; \`stop\` stops a running Program. When a background Program ends while its caller is not waiting for it, the caller receives a notification; \`wait\` returns the result.
@@ -18,6 +25,19 @@ Program code is the body of an async function. It calls the caller's tools throu
 
 \`program\` replaces the \`codemode\` tool; wherever \`codemode\` is mentioned, use \`program\`.`;
 
+
+interface ProgramCall {
+	path: string;
+	args?: string;
+	code: string;
+}
+
+interface ProgramMessageDetails {
+	call: ProgramCall;
+	calls?: unknown[];
+	isError: boolean;
+}
+
 const AGENT_API = `Agent API:
 \`\`\`ts
 type JsonSchema = boolean | Record<string, unknown>;
@@ -38,11 +58,13 @@ interface SendOptions {
   deliverAs?: "followUp" | "steer" | "write";
   /** \`followUp\` only. The Agent must submit a value matching it. */
   schema?: JsonSchema;
+  /** \`followUp\` only. Seconds before the input is aborted. Default: none. */
+  timeout?: number;
 }
 interface AgentHandle {
   readonly id: string;
   readonly name: string;
-  /** Resolves with the answer, or with the submitted value when \`schema\` is set; a write resolves with undefined once accepted. Throws if the input fails or is aborted. */
+  /** Resolves with the answer, or with the submitted value when \`schema\` is set; a write resolves with undefined once accepted. Throws if the input fails, is aborted, or times out. */
   send(message: string, options?: SendOptions): Promise<unknown>;
   /** Aborts the Agent's current turn and queued inputs; the Agent stays usable. */
   abort(): Promise<void>;
@@ -63,6 +85,8 @@ export function registerProgramTools(
 		programs: () => Programs | undefined;
 	},
 ): void {
+	const renderers = programRenderers(piCodemode);
+
 	/** Runs a Program whose Agents go offline when it ends. */
 	const runProgram = async (
 		id: string,
@@ -71,26 +95,32 @@ export function registerProgramTools(
 		code: string,
 		signal: AbortSignal,
 		current: Parameters<typeof executeProgram>[3],
-		options: Pick<ProgramRunOptions, "timeout" | "onUpdate" | "background" | "abort">,
+		{ background, args, ...options }: Pick<ProgramRunOptions, "timeout" | "onUpdate" | "detached" | "abort"> & { background?: boolean; args?: string },
 	) => {
-		const scope = programScope(id, caller, pi, current, { agentDir, limits: settings, extensions: settings.extensions, background: options.background === true });
+		const scope = programScope(id, caller, pi, current, { agentDir, limits: settings, extensions: settings.extensions, background: background === true });
 		let run: Awaited<ReturnType<typeof executeProgram>> | undefined;
 		try {
 			run = await executeProgram(toolCallId, code, signal, current, {
 				...options,
-				globals: agentGlobals(scope.host),
-				prepare: withAgentPrefix,
+				globals: (track) => agentGlobals(scope.host, track),
+				// On the first line, so line numbers in errors match the script as written.
+				prepare: (source) => withAgentPrefix(args === undefined ? source : `const args = ${JSON.stringify(args)};${source}`),
 				appendEntry: (customType, data) => pi.appendEntry(customType, data),
 				getToolNamespace: (name) => pi.getAllTools().find((tool) => tool.name === name)?.namespace,
 			});
 		} finally {
+			// Tool calls the script left running and the Program's Agents stop within the same timeout.
+			const deadline = Date.now() + CLEANUP_TIMEOUT_MS;
+			const stopped = await Promise.all([
+				run ? settlesBy(run.toolCalls, deadline) : true,
+				scope.close().then(() => true, () => false),
+			]);
 			// A cleanup failure ends the result, which keeps the Program's output and calls.
-			await scope.close().catch(() => {
-				if (!run) return;
+			if (run && stopped.includes(false)) {
 				run.result.content.push({ type: "text", text: CLEANUP_ERROR });
 				run.result.isError = true;
 				run.abandoned = true;
-			});
+			}
 		}
 		return run;
 	};
@@ -114,6 +144,7 @@ export function registerProgramTools(
 		},
 	});
 
+	const codemodeReach = piCodemode.createCodemodeToolDefinition({}).reachesTools;
 	pi.registerTool({
 		name: PROGRAM_TOOL_NAME,
 		label: "Program",
@@ -121,7 +152,10 @@ export function registerProgramTools(
 		exposure: "model-only",
 		defaultActive: false,
 		prepareLoadout: programLoadout(pi, piCodemode, PROGRAM_DESCRIPTION, AGENT_API),
-		...programRenderers(piCodemode),
+		// Only `run` reaches tools, as a `codemode` script does; extensions that register tools late wait on it.
+		// MCP sees a run as a `codemode` call, without `action`.
+		...(codemodeReach && { reachesTools: (input: { action?: string; code?: unknown }) => (input.action ?? "run") === "run" ? codemodeReach({ code: input.code }) : [] }),
+		...renderers,
 		description: programDescription(piCodemode, PROGRAM_DESCRIPTION, AGENT_API),
 		promptSnippet: "Run JavaScript that composes tool calls and Agents",
 		promptGuidelines: [
@@ -149,7 +183,7 @@ export function registerProgramTools(
 					if (!params.code?.trim()) throw new Error("run requires code");
 					const code = params.code;
 					if (params.background) {
-						const record = programs.start((id, programSignal, abort) => runProgram(id, self, toolCallId, code, programSignal, current, { timeout: params.timeout, background: true, abort }));
+						const record = programs.start(code, (id, programSignal, abort, update) => runProgram(id, self, toolCallId, code, programSignal, current, { timeout: params.timeout, background: true, detached: true, abort, onUpdate: update }));
 						return { content: [{ type: "text", text: `Program ${programs.label(record)} started.` }], details: { id: record.id } };
 					}
 					const { result } = await self.agents.whileSuspended(() => runProgram(randomUUID(), self, toolCallId, code, signal ?? new AbortController().signal, current, {
@@ -164,11 +198,12 @@ export function registerProgramTools(
 					if (timeout < 10 || timeout > 3600) throw new Error("wait timeout must be between 10 and 3600 seconds");
 					const selected = programs.select(targets);
 					onUpdate?.({ content: [], details: undefined });
-					const waitFor = () => programs!.wait(selected, timeout, signal);
+					const session = agentSession(current.sessionManager);
+					const waitFor = () => programs!.wait(selected, timeout, signal, () => steerQueued(session?.agent));
 					const outcome = programs.running(selected) ? await self.agents.whileSuspended(waitFor, signal, true) : await waitFor();
 					const rendered = renderProgramWait(outcome, programs.ids());
 					programs.markReturned(rendered.returned);
-					return { content: rendered.content, details: { running: outcome.running.map((record) => record.id) } };
+					return { content: rendered.content, details: { running: outcome.running.map((record) => record.id), ...(outcome.interrupted && { interrupted: true }) } };
 				}
 				case "list": {
 					rejectFields(params, "list", ["code", "background", "timeout", "target"]);
@@ -195,4 +230,125 @@ export function registerProgramTools(
 			}
 		},
 	});
+
+	/** A run as Pi shows a `run` call: the script, then its calls and output. */
+	const renderRun = (call: ProgramCall, result: { content: unknown[]; details?: unknown }, isPartial: boolean, isError: boolean, expanded: boolean, outputPad: number, theme: Theme) => {
+		const component = new Box(outputPad, 1, (text) => theme.bg(isPartial ? "toolPendingBg" : isError ? "toolErrorBg" : "toolSuccessBg", text));
+		const context = { args: call, state: {}, executionStarted: true, expanded, isError, showImages: false, isPartial } as never;
+		component.addChild(renderers.renderCall!(call, theme, context));
+		component.addChild(renderers.renderResult!(result as never, { expanded, isPartial }, theme, context));
+		return component;
+	};
+
+	/** Runs `/program <path> [args]` in the foreground: shown above the editor while it runs, and added to the context like a `!` command. */
+	const runCommand = async (text: string, current: ExtensionCommandContext) => {
+		const self = requireNode(current);
+		const session = commandSession()!;
+		const [, file, args = ""] = /^(\S+)\s*([\s\S]*)$/.exec(text.trim()) ?? [];
+		if (!file) throw new Error("Usage: /program <path> [args]");
+		const code = await readFile(path.resolve(current.cwd, file), "utf8");
+		const toolCallId = `program-command-${randomUUID()}`;
+		const call: ProgramCall = { path: file, ...(args ? { args } : {}), code };
+		const controller = new AbortController();
+		// Pi's Esc cancels running `!` commands while no reply is streaming; it stops the Program likewise.
+		const stop = { abort: () => controller.abort(new Error("Program stopped")) };
+		session._bashAbortControllers.add(stop);
+		const release = trackTreeWork(self.rootId);
+		const widget = `${PROGRAM_MESSAGE}:${toolCallId}`;
+		const shown = current.mode === "tui";
+		let partial: { content: unknown[]; details?: unknown; isError?: boolean } = { content: [] };
+		let ended = false;
+		let tui: { requestRender(): void } | undefined;
+		if (shown) {
+			current.ui.setWidget(widget, (ui, theme) => {
+				tui = ui;
+				const outputPad = pi.getSettings().outputPad === 0 ? 0 : 1;
+				return { render: (width: number) => renderRun(call, partial, !ended, partial.isError === true, current.ui.getToolsExpanded(), outputPad, theme).render(width), invalidate() {} };
+			});
+		}
+		let run: Awaited<ReturnType<typeof runProgram>>;
+		try {
+			// Hooks see the run as a `run` call, so MCP waits for the servers the script names.
+			const blocked = await runDetached<{ block?: boolean; reason?: string } | undefined>(session, controller.signal, stop.abort, () => session._beforeToolCall({ toolCall: { name: PROGRAM_TOOL_NAME, id: toolCallId }, args: { action: "run", code } }));
+			if (blocked?.block) throw new Error(blocked.reason || "Tool execution was blocked");
+			const ctx = session._extensionRunner.createToolContext(toolCallId, controller.signal);
+			run = await runProgram(randomUUID(), self, toolCallId, code, controller.signal, ctx, {
+				detached: true,
+				abort: stop.abort,
+				args,
+				onUpdate: (partialResult) => {
+					partial = partialResult;
+					tui?.requestRender();
+				},
+			});
+		} catch (error) {
+			if (shown) current.ui.setWidget(widget, undefined);
+			throw error;
+		} finally {
+			session._bashAbortControllers.delete(stop);
+			release();
+		}
+		const { result } = run;
+		const details: ProgramMessageDetails = { call, calls: (result.details as { calls?: unknown[] } | undefined)?.calls, isError: result.isError === true };
+		if (shown) {
+			// During a reply, the message comes after the reply; until then the widget shows the result.
+			partial = result;
+			ended = true;
+			tui?.requestRender();
+			const unsubscribe = session.subscribe((event: { type: string; message?: { details?: unknown } }) => {
+				if (event.type !== "message_start" || event.message?.details !== details) return;
+				unsubscribe();
+				current.ui.setWidget(widget, undefined);
+			});
+		}
+		const send = () => pi.sendMessage({ customType: PROGRAM_MESSAGE, content: [{ type: "text", text: `Ran \`/program ${text.trim()}\`\n` }, ...result.content], display: true, details }, { triggerTurn: false });
+		// Like a `!` command that ends during a reply, the run joins the context once the reply has ended.
+		if (session.isStreaming) {
+			const unsubscribe = session.subscribe((event: { type: string }) => {
+				if (event.type !== "agent_settled") return;
+				unsubscribe();
+				send();
+			});
+		} else send();
+		if (current.mode !== "print") return;
+		// A Session in a host prints through its print client instead (src/host/rpc-client.mjs).
+		const body = runOutput(result.content);
+		if (result.isError) {
+			process.stderr.write(`${body}\n`);
+			process.exitCode = 1;
+		} else await writeStdout(`${body}\n`);
+	};
+
+	pi.registerCommand(PROGRAM_TOOL_NAME, {
+		description: "Run the Program script in a file, with the text after the path as `args`",
+		handler: async (text, current) => {
+			try {
+				await runCommand(text, current);
+			} catch (error) {
+				// Pi reports the error; `pi -p` also exits with status 1.
+				if (current.mode === "print") process.exitCode = 1;
+				throw error;
+			}
+		},
+	});
+
+	pi.registerMessageRenderer<ProgramMessageDetails>(PROGRAM_MESSAGE, (message, { expanded, outputPad }, theme) => {
+		const { details } = message;
+		if (!details) return undefined;
+		const content = (message.content as unknown[]).slice(1);
+		return renderRun(details.call, { content, details: { calls: details.calls } }, false, details.isError, expanded, outputPad, theme);
+	});
+}
+
+/** Writes all of `text` to stdout, which print mode sends to stderr and Pi makes non-blocking. */
+async function writeStdout(text: string): Promise<void> {
+	let rest = Buffer.from(text);
+	while (rest.length > 0) {
+		try {
+			rest = rest.subarray(writeSync(1, rest));
+		} catch (error) {
+			if ((error as NodeJS.ErrnoException).code !== "EAGAIN") throw error;
+			await new Promise((resolve) => setTimeout(resolve, 10));
+		}
+	}
 }

package/src/roots/guard.mjs

@@ -3,8 +3,14 @@ import { closeSync, openSync, readFileSync, readSync, writeFileSync } from 'node
 import { acquireLock, canonicalPath, SessionOccupiedError } from './locks.mjs';
 import { resolvePath } from './paths.mjs';
 
-/** The ID in a Session file's header, its first parseable entry. Reads only up to that entry. */
+/** The ID in a Session file's header. */
 export function readSessionId(path) {
+  const header = readSessionHeader(path);
+  return typeof header?.id === 'string' ? header.id : undefined;
+}
+
+/** A Session file's header, its first parseable entry. Reads only up to that entry. */
+export function readSessionHeader(path) {
   let fd;
   try { fd = openSync(path, 'r'); }
   catch (error) { if (error.code === 'ENOENT') return; throw error; }
@@ -19,7 +25,7 @@ export function readSessionId(path) {
       for (const line of data.subarray(0, end).toString('utf8').split('\n')) {
         let entry;
         try { entry = JSON.parse(line); } catch { continue; }
-        if (entry) return entry.type === 'session' && typeof entry.id === 'string' ? entry.id : undefined;
+        if (entry) return entry.type === 'session' ? entry : undefined;
       }
       if (!bytes) return;
       rest = data.subarray(end);

package/src/roots/queue.ts

@@ -0,0 +1,52 @@
+/**
+ * Agent messages queued in a root Session outlive the user's actions on Pi's queue. Relies on
+ * private members of `AgentSession`; check them when Pi is updated.
+ */
+import type { AgentSession } from "@earendil-works/pi-coding-agent";
+import { takeQueued } from "../agents/agents.ts";
+import { nodes } from "../agents/registry.ts";
+
+const PATCHED = Symbol.for("pi-agents:queue-patched");
+
+type Patchable = Record<PropertyKey, any>;
+
+/**
+ * Clearing the queue, as Pi's dequeue does, leaves Agent messages queued; stopping the turn adds
+ * them to the Session without starting a turn, as later ones are until the user's next input.
+ * Owned Agents withdraw their queued inputs on abort.
+ */
+export function installQueuePatches(Session: typeof AgentSession): void {
+	const session = Session.prototype as unknown as Patchable;
+	if (session[PATCHED]) return;
+	session[PATCHED] = true;
+	// Sessions without pi-agents, as other extensions may create, keep Pi's behavior.
+	const root = (self: AgentSession) => {
+		const node = nodes.get(self.sessionManager.getSessionId());
+		return node !== undefined && node.ownerId === undefined;
+	};
+	const clearQueue = session.clearQueue;
+	session.clearQueue = function (this: AgentSession) {
+		if (!root(this)) return clearQueue.call(this);
+		const kept = takeQueued(this.agent);
+		const cleared = clearQueue.call(this);
+		for (const message of kept) {
+			if (message.details.delivery === "followUp") this.agent.followUp(message);
+			else this.agent.steer(message);
+		}
+		return cleared;
+	};
+	const abort = session.abort;
+	session.abort = async function (this: AgentSession) {
+		if (!root(this)) return abort.call(this);
+		const node = nodes.get(this.sessionManager.getSessionId())!;
+		if (this.isStreaming) node.stopped = true;
+		const add = async () => {
+			for (const message of takeQueued(this.agent)) await this.sendCustomMessage(message, { triggerTurn: false });
+		};
+		await add();
+		const result = await abort.call(this);
+		// Messages that arrive while the turn stops are queued; add them too.
+		await add();
+		return result;
+	};
+}

package/src/roots/report.ts

@@ -13,6 +13,8 @@ export interface ReportHooks {
 	treeIdle(): boolean;
 	/** The tree's running owned Agents and background Programs. */
 	liveWork(): number;
+	/** Whether `/tasks` is open, whose list and viewer Pi counts as prompts though they ask nothing. */
+	browsing(): boolean;
 	/** The host this Pi runs in, if any. */
 	host?: { socket: string; attached(): boolean; onChange(listener: () => void): void };
 }
@@ -45,7 +47,7 @@ export function createReporter(pi: ExtensionAPI, hooks: ReportHooks): void {
 	let stream: { send(message: object): void; close(): void } | undefined;
 	let generation = 0;
 	let timer: ReturnType<typeof setTimeout> | undefined;
-	let dialogs = 0;
+	let dialog = false;
 	let tools = new Map<string, string>();
 	let turn: { at: number; failed: boolean } | undefined;
 	let reply: string | undefined;
@@ -67,8 +69,8 @@ export function createReporter(pi: ExtensionAPI, hooks: ReportHooks): void {
 			title,
 			host: hooks.host?.socket,
 			attached: hooks.host ? hooks.host.attached() : true,
-			working: !current.isIdle() || current.hasPendingMessages() || !hooks.treeIdle() || dialogs > 0,
-			blocked: dialogs > 0,
+			working: !current.isIdle() || current.hasPendingMessages() || !hooks.treeIdle(),
+			blocked: dialog,
 			tool: [...tools.values()].at(-1),
 			reply,
 			agents: hooks.liveWork(),
@@ -127,7 +129,7 @@ export function createReporter(pi: ExtensionAPI, hooks: ReportHooks): void {
 		ctx = undefined;
 		unsubscribe?.();
 		unsubscribe = undefined;
-		dialogs = 0;
+		dialog = false;
 		tools = new Map();
 		turn = reply = title = undefined;
 	});
@@ -141,8 +143,12 @@ export function createReporter(pi: ExtensionAPI, hooks: ReportHooks): void {
 	pi.on("tool_execution_end", (event) => {
 		if (tools.delete(event.toolCallId)) changed();
 	});
-	pi.on("ui_prompt_start", () => { dialogs += 1; changed(); });
-	pi.on("ui_prompt_end", () => { dialogs = Math.max(0, dialogs - 1); changed(); });
+	pi.on("ui_prompt_start", () => {
+		if (hooks.browsing()) return;
+		dialog = true;
+		changed();
+	});
+	pi.on("ui_prompt_end", () => { dialog = false; changed(); });
 	pi.on("agent_settled", (event) => {
 		if (!ctx) return;
 		const last = ctx.sessionManager.getBranch().findLast((entry) => entry.type === "message" && entry.message.role === "assistant");