@tshu-w/pi-agents 0.1.1 → 0.2.0

package/src/agents/registry.ts

@@ -1,4 +1,5 @@
 /** The Agents loaded in this process, shared across trees: who sees whom, how targets resolve, and how messages are addressed. */
+import type { Usage } from "@earendil-works/pi-ai";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import type { Agents } from "./agents.ts";
 import { TreeScheduler } from "./scheduler.ts";
@@ -16,12 +17,14 @@ export interface AgentNode {
 	/** The node whose tree this node sees: its root, or the Program its Agents belong to. */
 	scopeId: string;
 	ownerId?: string;
-	/** A Program that owns Agents; no Agent sees it, but the user sees it in the task panel and the Agent viewer. */
+	/** A Program that owns Agents; no Agent sees it, but the user sees it in the task panel and the task viewer. */
 	program?: boolean;
 	/** Whether a Program runs in the background. */
 	background?: boolean;
 	/** When a Program started. */
 	createdAt?: string;
+	/** What the replies of a Program's Agents and the Agents under them have cost since it started. */
+	spent?: Spend;
 	name(): string | undefined;
 	cwd(): string;
 	busy(): boolean;
@@ -29,6 +32,8 @@ export interface AgentNode {
 	agents: Agents;
 	/** Delivers a message to this node; used only for root Agents. */
 	receive(text: string, delivery: Delivery, messageId?: string): void;
+	/** A root Agent's turn was stopped: messages from other Agents become writes until the user's next input. */
+	stopped?: boolean;
 }
 
 export interface Entry {
@@ -86,6 +91,26 @@ export function trackTreeWork(rootId: string): () => void {
 	};
 }
 
+/** What replies cost: dollars, input and output tokens, and whether a reply had no price. */
+export interface Spend {
+	cost: number;
+	tokens: number;
+	unpriced: boolean;
+}
+
+export function addSpend(spend: Spend, usage: Usage): void {
+	spend.cost += usage.cost.total;
+	spend.tokens += usage.input + usage.output;
+	if (usage.cost.total === 0 && usage.input + usage.output + usage.cacheRead + usage.cacheWrite > 0) spend.unpriced = true;
+}
+
+/** Adds a reply of an Agent owned by `ownerId` to the Programs above it. */
+export function addProgramSpend(ownerId: string, usage: Usage): void {
+	for (let node = shared.nodes.get(ownerId); node; node = node.ownerId === undefined ? undefined : shared.nodes.get(node.ownerId)) {
+		if (node.spent) addSpend(node.spent, usage);
+	}
+}
+
 /** The Programs a node runs, oldest first. */
 export function programsOf(node: AgentNode): AgentNode[] {
 	return [...shared.nodes.values()]
@@ -169,7 +194,9 @@ export function treeMetadata(ctx: ExtensionContext): TreeMetadata | undefined {
 export function treeEntries(scopeId: string): Entry[] {
 	const root = shared.nodes.get(scopeId);
 	if (!root) return [];
-	const entries: Entry[] = root.program ? [] : [{
+	// A Program's Agent heads its own scope; it stays owned, so messages to it go through its owner.
+	const owned = root.ownerId === undefined ? undefined : shared.nodes.get(root.ownerId)?.agents.owned().find((agent) => agent.id === root.id);
+	const entries: Entry[] = root.program ? [] : [owned ? { ...owned, ownerId: root.ownerId } : {
 		id: root.id,
 		name: root.name(),
 		cwd: root.cwd(),
@@ -236,7 +263,7 @@ export function messageText(from: { id: string; name?: string }, fromOwner: bool
 
 /**
  * `messageId` marks a message from a root Agent in another process, so it is added once; `user`
- * marks the user's input from the Agent viewer.
+ * marks the user's input from the task viewer.
  */
 export function customMessage(text: string, delivery: Delivery, messageId?: string, user?: boolean) {
 	return { customType: MESSAGE_TYPE, content: text, display: true, details: { delivery, ...(messageId ? { messageId } : {}), ...(user ? { user } : {}) } };

package/src/agents/search.ts

@@ -4,6 +4,8 @@ import { entryLine, MESSAGE_TAG, MESSAGE_TYPE, type Entry } from "./registry.ts"
 interface SearchText {
 	time: number;
 	text: string;
+	/** A user message or an input, not a summary, write, or notification. */
+	request?: boolean;
 }
 
 /** Parsed texts by Session file, reused while the file is unchanged. */
@@ -32,9 +34,10 @@ function sessionTexts(sessionFile: string | undefined): SearchText[] {
 			const text = typeof content === "string"
 				? content
 				: Array.isArray(content) ? content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n") : "";
-			if (text) texts.push({ time, text });
+			if (text) texts.push({ time, text, request: true });
 		} else if (entry.type === "custom_message" && entry.customType === MESSAGE_TYPE && typeof entry.content === "string") {
-			texts.push({ time, text: entry.content });
+			const request = entry.details?.delivery !== "write" && (entry.content.startsWith(`<${MESSAGE_TAG} `) || entry.details?.user === true);
+			texts.push({ time, text: entry.content, request });
 		}
 	}
 	cache.set(sessionFile, { stamp, texts });
@@ -85,9 +88,23 @@ export function searchEntries(entries: Entry[], query: string | undefined): List
 	return items.sort((left, right) => right.time - left.time).map(({ time: _time, ...item }) => item);
 }
 
+/** An entry with its latest matching excerpt, else its latest request, and its time. */
 export function listLine(item: ListItem, ids: Iterable<string>): string {
 	const line = entryLine(item.entry, ids);
-	return item.match ? `${line}\n  ${formatTime(item.match.time)}  ${item.match.excerpt}` : line;
+	const detail = item.match ?? latestRequest(item.entry.sessionFile);
+	return detail ? `${line}\n  ${formatTime(detail.time)}  ${detail.excerpt}` : line;
+}
+
+function latestRequest(sessionFile: string | undefined): { time: number; excerpt: string } | undefined {
+	const text = sessionTexts(sessionFile).findLast((text) => text.request);
+	if (!text) return undefined;
+	const flat = messageBody(text.text).replace(/\s+/g, " ").trim();
+	return { time: text.time, excerpt: flat.length > 120 ? `${flat.slice(0, 120)}…` : flat };
+}
+
+/** A message's body without its `agent-message` element. */
+function messageBody(content: string): string {
+	return content.startsWith(`<${MESSAGE_TAG} `) ? content.slice(content.indexOf("\n") + 1, content.lastIndexOf("\n")) : content;
 }
 
 /**
@@ -103,7 +120,7 @@ export function latestActivity(sessionFile: string | undefined): string | undefi
 		let entry: Record<string, any>;
 		try { entry = JSON.parse(line); } catch { continue; }
 		if (input === undefined && entry.type === "custom_message" && entry.customType === MESSAGE_TYPE && typeof entry.content === "string") {
-			input = entry.content.startsWith(`<${MESSAGE_TAG} `) ? entry.content.slice(entry.content.indexOf("\n") + 1, entry.content.lastIndexOf("\n")) : entry.content;
+			input = messageBody(entry.content);
 		}
 		if (input !== undefined && entry.type === "message" && entry.message?.role === "assistant" && Array.isArray(entry.message.content)) {
 			activity = replyActivity(entry.message.content) ?? activity;

package/src/agents/tool.ts

@@ -4,9 +4,12 @@ import { Type } from "typebox";
 import { boundBlocks, boundText } from "../output.ts";
 import { rejectFields } from "../params.ts";
 import { renderTextResult, renderToolCall, startDuration, type DurationState } from "../render-call.ts";
+import { agentSession } from "../roots/rpc.ts";
 import type { RootRuntime } from "../roots/runtime.ts";
 import {
+	steerQueued,
 	THINKING_LEVELS,
+	WAIT_INTERRUPTED,
 	type Limits,
 	type WaitOutcome,
 	type WaitResult,
@@ -34,9 +37,10 @@ function resultDetails({ agent, input, history }: WaitResult, result?: string) {
 
 /** Renders wait results; returns the text and the results the caller has now read. */
 function renderWait(outcome: WaitOutcome, ids: string[]): { text: string; read: WaitResult[]; details: Record<string, unknown> } {
-	const { results, pending } = outcome;
-	const pendingLine = pending.length > 0 ? `Still pending: ${pending.map((agent) => label(agent, ids)).join(", ")}` : "";
-	if (results.length === 0) return { text: pendingLine || "No results.", read: [], details: { results: [], pending: pending.map((agent) => agent.id) } };
+	const { results, pending, interrupted } = outcome;
+	const pendingLine = pending.length > 0 ? `Still pending: ${pending.map((agent) => label(agent, ids)).join(", ")}${interrupted ? `\n${WAIT_INTERRUPTED}` : ""}` : "";
+	const flag = interrupted ? { interrupted } : {};
+	if (results.length === 0) return { text: pendingLine || "No results.", read: [], details: { results: [], pending: pending.map((agent) => agent.id), ...flag } };
 	const only = results[0]!;
 	if (results.length === 1 && pending.length === 0 && !only.history && only.input.state === "completed") {
 		const bounded = boundText(only.input.result ?? "", "pi-agents-wait");
@@ -56,7 +60,7 @@ function renderWait(outcome: WaitOutcome, ids: string[]): { text: string; read:
 	const omitted = results.slice(bounded.shown);
 	if (omitted.length > 0) text += `\n\n[Results omitted: ${omitted.map((result) => label(result.agent, ids)).join(", ")}. Use wait with fewer targets.]`;
 	if (pendingLine) text += `\n\n${pendingLine}`;
-	return { text, read, details: { results: details, pending: pending.map((agent) => agent.id), ...bounded.details } };
+	return { text, read, details: { results: details, pending: pending.map((agent) => agent.id), ...flag, ...bounded.details } };
 }
 
 /** Registers the `agent` tool. */
@@ -168,7 +172,8 @@ export function registerAgentTool(
 					const selected = targets?.map((target) => agents.ownedTarget(target));
 					// An empty partial result lets the row show the elapsed time.
 					onUpdate?.({ content: [], details: undefined });
-					const outcome = await agents.wait(selected, params.history ?? 0, params.timeout ?? 30, signal);
+					const session = agentSession(current.sessionManager);
+					const outcome = await agents.wait(selected, params.history ?? 0, params.timeout ?? 30, signal, () => steerQueued(session?.agent));
 					const rendered = renderWait(outcome, ids());
 					agents.markRead(rendered.read);
 					return { content: [{ type: "text", text: rendered.text }], details: rendered.details };

package/src/daemon/client.mjs

@@ -42,13 +42,16 @@ export async function ensureDaemon(paths, { sessionDir, signal }) {
       if (code !== 0) launchError = new Error(`Daemon exited (${code}). See ${paths.daemonLog}`);
     });
     child.unref();
+    // The new daemon sets its socket's mode just after binding; a mode that stays wrong is an error.
+    let modeWait = 0;
     for (;;) {
       signal?.throwIfAborted();
       if (launchError) throw launchError;
       try {
         return await request(paths.daemon, { action: 'status' }, { signal, timeoutMs: 1000 });
       } catch (failure) {
-        if (!['ENOENT', 'ECONNREFUSED', ...RETRY].includes(failure.code)) throw failure;
+        if (failure.code === 'SOCKET_MODE' && ++modeWait > 20) throw failure;
+        if (!['ENOENT', 'ECONNREFUSED', 'SOCKET_MODE', ...RETRY].includes(failure.code)) throw failure;
         // Another daemon holds the lock; it answers once its socket is up.
         if (child.exitCode === 0) break;
       }

package/src/host/client.mjs

@@ -16,7 +16,7 @@ const ALT_SCREEN = /\x1b\[\?1049([hl])/g;
 // Splits terminal input into key sequences: CSI and SS3 sequences, Alt+key, or a single character.
 const KEYS = /\x1b\[[0-?]*[ -/]*[@-~]|\x1bO[\s\S]|\x1b?[\s\S]/g;
 
-function hostStatus(path) {
+export function hostStatus(path) {
   return new Promise((resolve, reject) => {
     const socket = net.connect(path);
     socket.setTimeout(1000, () => socket.destroy(Object.assign(new Error('Host status timed out'), { code: 'ETIMEDOUT' })));
@@ -50,6 +50,19 @@ export async function listHosts(paths) {
   return hosts;
 }
 
+/**
+ * Running hosts once `found` matches one or no host `starting` remains, waiting at most `timeoutMs`:
+ * a host whose Pi has not reported its Session yet may be starting the Session a caller looks for.
+ */
+export async function settledHosts(paths, found, { starting = host => !host.session, timeoutMs = START_TIMEOUT_MS } = {}) {
+  const deadline = performance.now() + timeoutMs;
+  for (;;) {
+    const hosts = await listHosts(paths);
+    if (hosts.some(found) || !hosts.some(starting) || performance.now() > deadline) return hosts;
+    await delay(100);
+  }
+}
+
 /**
  * Starts a detached host running `pi args` in `cwd` and resolves to its socket once it accepts connections.
  * With `waitForTerminal`, Pi starts when the first terminal attaches.
@@ -88,9 +101,10 @@ export async function startHost(paths, { pi, args, cwd, cols = 120, rows = 40, e
 /**
  * Attaches this process's terminal to a host until the host detaches it or Pi exits.
  * The detach key detaches here, before Pi sees it, so it works whatever Pi shows.
+ * `onAttach` runs once the host has attached the terminal and told Pi.
  * Resolves to { reason: 'detach' | 'takeover' | 'exited' | 'closed', session }.
  */
-export function attach(socketPath, { isDetach, stdin = process.stdin, stdout = process.stdout }) {
+export function attach(socketPath, { isDetach, onAttach, stdin = process.stdin, stdout = process.stdout }) {
   return new Promise((resolve, reject) => {
     const socket = net.connect(socketPath);
     const size = () => ({ cols: stdout.columns, rows: stdout.rows });
@@ -104,7 +118,7 @@ export function attach(socketPath, { isDetach, stdin = process.stdin, stdout = p
       if (at > 0) socket.write(frame(FRAME.data, Buffer.from(keys.slice(0, at).join(''), encoding)));
       socket.write(frame(FRAME.detached));
     };
-    let result, connected = false, altScreen = false;
+    let result, connected = false, altScreen = false, attached = false;
     function finish() {
       stdout.off('resize', onResize);
       stdin.off('data', onInput);
@@ -113,7 +127,8 @@ export function attach(socketPath, { isDetach, stdin = process.stdin, stdout = p
         stdin.pause();
         // An exiting Pi restores the terminal itself.
         // Leaving the alternate screen returns the cursor to the line below the command; otherwise it ends a line.
-        if (result?.reason !== 'exited') stdout.write(altScreen ? `\x1b[?1049l${RESET}` : `${RESET}\r\n`);
+        // Terminals keep a keyboard mode stack per screen, so the reset comes before leaving it.
+        if (result?.reason !== 'exited') stdout.write(altScreen ? `${RESET}\x1b[?1049l` : `${RESET}\r\n`);
       }
       resolve(result ?? { reason: 'closed' });
     }
@@ -129,6 +144,11 @@ export function attach(socketPath, { isDetach, stdin = process.stdin, stdout = p
       if (type === FRAME.data) {
         for (const [, mode] of body.toString('latin1').matchAll(ALT_SCREEN)) altScreen = mode === 'h';
         stdout.write(body);
+        // The host sends the screen only after it has told Pi that a terminal is attached.
+        if (!attached) {
+          attached = true;
+          onAttach?.();
+        }
       }
       else if (type === FRAME.detached) result = { ...json(body) };
       else if (type === FRAME.exited) result = { reason: 'exited', ...json(body) };

package/src/host/extension.ts

@@ -3,7 +3,7 @@ import net from "node:net";
 import { treeMetadata } from "../agents/registry.ts";
 import { FRAME, frame, json, readFrames } from "./frames.mjs";
 
-const IDLE_EXIT_MS = 5 * 60_000;
+const IDLE_EXIT_MS = 10 * 60_000;
 
 // Only the host's own child runs in it; children of Pi inherit the variables but have another parent.
 // They stay in the environment so a Pi that replaces itself with execve finds its host again.
@@ -12,20 +12,21 @@ const hostSocket = process.env.PI_AGENTS_HOST && process.env.PI_AGENTS_HOST_PID
 export interface Host {
 	socket: string;
 	attached(): boolean;
-	onChange(listener: () => void): void;
+	onChange(listener: () => void): () => void;
 }
 
 /**
  * Connects Pi to the host it runs in: reports the current Session, detaches on request,
- * disables suspend, and exits after the tree has been idle and detached for a while.
+ * disables suspend, and exits after the tree has been idle and detached, with no RPC client, for a while.
  * Returns the host, if Pi runs in one.
  */
-export function hostExtension(pi: ExtensionAPI, { treeIdle }: { treeIdle: () => boolean }): Host | undefined {
+export function hostExtension(pi: ExtensionAPI, { treeIdle, rpcClients }: { treeIdle: () => boolean; rpcClients: () => number }): Host | undefined {
 	if (!hostSocket) return;
 	const listeners: Array<() => void> = [];
 
 	let ctx: ExtensionContext | undefined;
-	let attached = true;
+	// Whether a terminal was attached when Pi started, until the host reports a change.
+	let attached = process.env.PI_AGENTS_HOST_ATTACHED !== "0";
 	let idleSince: number | undefined;
 	let socket: net.Socket | undefined;
 	const connect = (path: string): net.Socket => {
@@ -43,32 +44,55 @@ export function hostExtension(pi: ExtensionAPI, { treeIdle }: { treeIdle: () =>
 	const report = () => {
 		if (!ctx) return;
 		const file = ctx.sessionManager.getSessionFile();
-		send({ type: "session", session: { id: ctx.sessionManager.getSessionId(), name: pi.getSessionName(), file } });
+		send({ type: "session", session: { id: ctx.sessionManager.getSessionId(), name: pi.getSessionName(), file, cwd: ctx.cwd } });
 	};
 
 	// Pi suspends by stopping its process group; a host has no shell to resume it, so resume at once.
-	process.on("SIGTSTP", () => process.kill(process.pid, "SIGCONT"));
-
-	const timer = setInterval(() => {
-		if (!ctx || attached || !ctx.isIdle() || ctx.hasPendingMessages() || !treeIdle()) {
+	const resume = () => process.kill(process.pid, "SIGCONT");
+	const checkIdle = () => {
+		if (!ctx || attached || !ctx.isIdle() || ctx.hasPendingMessages() || !treeIdle() || rpcClients() > 0) {
 			idleSince = undefined;
 			return;
 		}
 		idleSince ??= Date.now();
 		if (Date.now() - idleSince >= IDLE_EXIT_MS) ctx.shutdown();
-	}, 1000);
-	timer.unref();
+	};
+	let timer: ReturnType<typeof setInterval> | undefined;
 
 	pi.on("session_start", (_event, current) => {
 		// Owned Agents run in the same process; only the root speaks for the host.
 		if (treeMetadata(current)) return;
 		ctx = current;
 		socket ??= connect(hostSocket);
+		if (!timer) {
+			process.on("SIGTSTP", resume);
+			timer = setInterval(checkIdle, 1000);
+			timer.unref();
+		}
 		report();
 	});
 	pi.on("session_info_changed", () => report());
-	pi.on("session_shutdown", () => { ctx = undefined; });
+	// A reload or another Session loads this extension again, which takes over.
+	pi.on("session_shutdown", () => {
+		ctx = undefined;
+		idleSince = undefined;
+		clearInterval(timer);
+		timer = undefined;
+		process.off("SIGTSTP", resume);
+		socket?.destroy();
+		socket = undefined;
+	});
 
 	pi.registerCommand("detach", { description: "Detach this terminal; Pi keeps running", handler: async () => send({ type: "detach" }) });
-	return { socket: hostSocket, attached: () => attached, onChange: (listener) => { listeners.push(listener); } };
+	return {
+		socket: hostSocket,
+		attached: () => attached,
+		onChange: (listener) => {
+			listeners.push(listener);
+			return () => {
+				const at = listeners.indexOf(listener);
+				if (at >= 0) listeners.splice(at, 1);
+			};
+		},
+	};
 }