@tshu-w/pi-agents 0.1.1 → 0.2.1

package/src/roots/transport.mjs

@@ -1,9 +1,10 @@
 import net from 'node:net';
-import { StringDecoder } from 'node:string_decoder';
 import { chmodSync, lstatSync, mkdirSync } from 'node:fs';
 import { dirname, resolve } from 'node:path';
 
 const MAX_BYTES = 1024 * 1024;
+// Pi's events on an rpc stream carry whole messages, images included.
+const RPC_MAX_BYTES = 256 * 1024 * 1024;
 const MAX_SOCKET_PATH_BYTES = 103;
 const MAX_PENDING = 64;
 // Each running Pi keeps a stream open to the daemon.
@@ -39,9 +40,9 @@ function safeDirectory(path, create = false) {
 function privateSocket(path) {
   safeDirectory(dirname(path));
   const info = lstatSync(path);
-  if (!info.isSocket() || info.uid !== process.getuid() || (info.mode & 0o777) !== 0o600) {
-    throw fault('UNSAFE_PATH', 'Socket must be owned by the current user with mode 0600');
-  }
+  if (!info.isSocket() || info.uid !== process.getuid()) throw fault('UNSAFE_PATH', 'Socket must be owned by the current user');
+  // A listener sets the mode just after it binds; a client waiting for a new listener may retry this.
+  if ((info.mode & 0o777) !== 0o600) throw fault('SOCKET_MODE', 'Socket must have mode 0600');
 }
 
 const record = value => value !== null && typeof value === 'object' && !Array.isArray(value);
@@ -51,7 +52,7 @@ function validate(value) {
   if (!record(value)) throw fault('INVALID_REQUEST', 'Expected a request object');
   if (value.action === 'status') return { action: 'status' };
   if (value.action === 'stream') {
-    if (value.role !== 'session' && value.role !== 'workbench') throw fault('INVALID_REQUEST', 'stream role must be session or workbench');
+    if (!['session', 'workbench', 'rpc', 'print'].includes(value.role)) throw fault('INVALID_REQUEST', 'stream role must be session, workbench, rpc or print');
     return { action: 'stream', role: value.role };
   }
   const m = value.message, l = value.launch;
@@ -75,12 +76,14 @@ function validate(value) {
   }, ...(l === undefined ? {} : { launch: { command: l.command, args: [...l.args], env: { ...l.env } } }) };
 }
 
-function encode(value) {
+function encode(value, limit = MAX_BYTES) {
   const bytes = Buffer.from(`${JSON.stringify(value)}\n`);
-  if (bytes.length > MAX_BYTES) throw fault('FRAME_TOO_LARGE', 'JSONL frame exceeds 1 MiB');
+  if (bytes.length > limit) throw fault('FRAME_TOO_LARGE', `JSONL frame exceeds ${limit / 1024 / 1024} MiB`);
   return bytes;
 }
 
+const streamLimit = role => role === 'rpc' || role === 'print' ? RPC_MAX_BYTES : MAX_BYTES;
+
 // Consume exactly one bounded JSONL frame; never dispatch additional requests.
 function readFrame(socket, done) {
   let chunks = [], size = 0, finished = false;
@@ -108,28 +111,33 @@ function readFrame(socket, done) {
 }
 
 // JSONL messages after a stream's opening frame.
-function readLines(socket, onLine, onInvalid) {
-  const decoder = new StringDecoder('utf8');
-  let buffer = '';
+function readLines(socket, onLine, onInvalid, limit) {
+  // An unfinished line is kept as chunks, so each byte is scanned once.
+  let pending = [], size = 0;
   socket.on('data', chunk => {
-    buffer += decoder.write(chunk);
-    for (let newline; (newline = buffer.indexOf('\n')) >= 0;) {
-      const line = buffer.slice(0, newline);
-      buffer = buffer.slice(newline + 1);
+    let start = 0;
+    for (let newline; (newline = chunk.indexOf(10, start)) >= 0; start = newline + 1) {
+      const line = Buffer.concat([...pending, chunk.subarray(start, newline)]).toString('utf8');
+      pending = [];
+      size = 0;
       let value;
       try { value = JSON.parse(line); } catch { return onInvalid(); }
       onLine(value);
     }
-    if (Buffer.byteLength(buffer) > MAX_BYTES) onInvalid();
+    if (start < chunk.length) {
+      pending.push(chunk.subarray(start));
+      size += chunk.length - start;
+    }
+    if (size > limit) onInvalid();
   });
 }
 
-function channel(socket, signal) {
+function channel(socket, signal, limit) {
   const listeners = [];
-  readLines(socket, value => { for (const listener of listeners) listener(value); }, () => socket.destroy());
+  readLines(socket, value => { for (const listener of listeners) listener(value); }, () => socket.destroy(), limit);
   return {
     signal,
-    send(value) { if (!socket.destroyed && !socket.writableEnded) socket.write(encode(value)); },
+    send(value) { if (!socket.destroyed && !socket.writableEnded) socket.write(encode(value, limit)); },
     onMessage(listener) { listeners.push(listener); },
     close() { socket.destroy(); },
   };
@@ -145,7 +153,8 @@ function errorFrame(error) {
 
 /**
  * FIFO by default. Handlers receive a connection-lifetime signal, not an acceptance rollback.
- * `stream` takes over the connections that open a stream; they keep the server from being idle.
+ * `stream` takes over the connections that open a stream; they keep the server from being idle. It may
+ * refuse a stream by throwing; its result joins the acceptance, which reaches the client before anything it sends.
  * @param {string} socketPath
  * @param {{ status: Function, accept: Function, stream?: Function, serialize?: boolean }} handlers
  */
@@ -189,8 +198,17 @@ export async function listen(socketPath, { status, accept, stream, serialize = t
       if (input.action === 'stream') {
         if (!stream) return respond(fault('INVALID_REQUEST', 'Streams are not supported here'));
         clearTimeout(timer);
-        socket.write(encode({ ok: true, result: { pid: process.pid } }));
-        return stream(input, channel(socket, signal));
+        const opened = channel(socket, signal, streamLimit(input.role));
+        const send = opened.send;
+        const early = [];
+        opened.send = value => early.push(value);
+        let result;
+        try { result = stream(input, opened); }
+        catch (failure) { return respond(failure); }
+        socket.write(encode({ ok: true, result: { pid: process.pid, ...result } }));
+        opened.send = send;
+        for (const value of early) send(value);
+        return;
       }
       if (pending >= MAX_PENDING) return respond(fault('SERVER_BUSY', 'Request queue is full'));
       pending++;
@@ -332,10 +350,10 @@ export function openStream(socketPath, input, { onMessage, onClose }) {
       opened = true;
       resolvePromise({
         result: response.result,
-        send(message) { if (!socket.destroyed && !socket.writableEnded) socket.write(encode(message)); },
+        send(message) { if (!socket.destroyed && !socket.writableEnded) socket.write(encode(message, streamLimit(value.role))); },
         close() { socket.destroy(); },
         unref() { socket.unref(); },
       });
-    }, () => socket.destroy());
+    }, () => socket.destroy(), streamLimit(value.role));
   });
 }

package/src/ui/message.ts

@@ -0,0 +1,24 @@
+import { getMarkdownTheme, UserMessageComponent, type MessageRenderer } from "@earendil-works/pi-coding-agent";
+import { Box, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
+import { messageBody, type Sender } from "../agents/registry.ts";
+
+/**
+ * Collapsed, a message from another Agent shows its sender after the label, as Pi shows a collapsed skill, and the
+ * user's input from the task viewer shows as a user message. Expanded, or for any other message, Pi shows the text the
+ * model received.
+ */
+export const renderAgentMessage: MessageRenderer = (message, { expanded, outputPad }, theme) => {
+	if (expanded || typeof message.content !== "string") return undefined;
+	const details = message.details as { user?: boolean; sender?: Sender } | undefined;
+	if (details?.user) return new UserMessageComponent(message.content, getMarkdownTheme(), outputPad);
+	const body = details?.sender && messageBody(message.content, details.sender);
+	if (body === undefined) return undefined;
+	const { id, name } = details!.sender!;
+	const box = new Box(outputPad, 1, (text) => theme.bg("customMessageBg", text));
+	const label = theme.fg("customMessageLabel", `\x1b[1m[${message.customType}]\x1b[22m`);
+	const sender = name === undefined ? theme.fg("customMessageText", ` from ${id}`) : theme.fg("customMessageText", ` from ${name}`) + theme.fg("dim", ` (${id})`);
+	box.addChild(new Text(label + sender, 0, 0));
+	box.addChild(new Spacer(1));
+	box.addChild(new Markdown(body, 0, 0, getMarkdownTheme(), { color: (text) => theme.fg("customMessageText", text) }));
+	return box;
+};

package/src/ui/panel.ts

@@ -1,16 +1,34 @@
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { Container, TruncatedText } from "@earendil-works/pi-tui";
-import { label, nodes, onTreeChange, programLabel, programsOf, type AgentNode } from "../agents/registry.ts";
+import { label, nodes, onTreeChange, programLabel, programsOf, type AgentNode, type Spend } from "../agents/registry.ts";
 
 const WIDGET_KEY = "pi-agents";
 
 const MAX_LINES = 10;
 
+/** Time since `since` as Pi's bash tool shows durations, in whole seconds: `45s`, `1m 32s`, `1h 5m 3s`. */
+function elapsed(since: number, now: number): string {
+	const seconds = Math.max(0, Math.floor((now - since) / 1000));
+	const minutes = Math.floor(seconds / 60);
+	if (minutes === 0) return `${seconds}s`;
+	return minutes < 60 ? `${minutes}m ${seconds % 60}s` : `${Math.floor(minutes / 60)}h ${minutes % 60}m ${seconds % 60}s`;
+}
+
+/** `$2.10`, or the input and output tokens as Pi's footer shows them, `7.3M tok`, when a reply had no price. */
+function cost({ cost, tokens, unpriced }: Spend): string {
+	if (!unpriced) return `$${cost.toFixed(2)}`;
+	if (tokens < 1000) return `${tokens} tok`;
+	if (tokens < 10_000) return `${(tokens / 1000).toFixed(1)}k tok`;
+	if (tokens < 1_000_000) return `${Math.round(tokens / 1000)}k tok`;
+	return tokens < 10_000_000 ? `${(tokens / 1_000_000).toFixed(1)}M tok` : `${Math.round(tokens / 1_000_000)}M tok`;
+}
+
 /**
  * The task panel: the Agent's live owned Agents and running background Programs as a tree, in the
- * order they started, with their Programs' live Agents under them. Empty when none.
+ * order they started, with their Programs' live Agents under them, and how long each has run and what
+ * it has cost. Empty when none.
  */
-export function panelLines(self: AgentNode): string[] {
+export function panelLines(self: AgentNode, now = Date.now()): string[] {
 	const rows: Array<{ line: string; parent?: number; shown: boolean }> = [];
 	const ids = self.agents.ids();
 	let live = 0;
@@ -26,7 +44,9 @@ export function panelLines(self: AgentNode): string[] {
 		for (const child of children) {
 			const index = rows.length;
 			if ("program" in child) {
-				rows.push({ line: `${"  ".repeat(level)}${programLabel(child.program)}  running`, parent, shown: false });
+				const since = Date.parse(child.program.createdAt ?? "");
+				const time = since ? [elapsed(since, now), ...(child.program.spent ? [cost(child.program.spent)] : [])] : [];
+				rows.push({ line: `${"  ".repeat(level)}${[programLabel(child.program), "running", ...time].join("  ")}`, parent, shown: false });
 				live += 1;
 				show(index);
 				visit(child.program, level + 1, index);
@@ -35,7 +55,8 @@ export function panelLines(self: AgentNode): string[] {
 			const { entry } = child;
 			const own = node.agents.counts(entry.id);
 			const pending = own?.pending ? [`${own.pending} pending`] : [];
-			rows.push({ line: `${"  ".repeat(level)}Agent ${[label(entry, ids), entry.state, ...pending].join("  ")}`, parent, shown: false });
+			const time = own?.since ? [elapsed(own.since, now), ...(own.spent ? [cost(own.spent)] : [])] : [];
+			rows.push({ line: `${"  ".repeat(level)}Agent ${[label(entry, ids), entry.state, ...time, ...pending].join("  ")}`, parent, shown: false });
 			if (entry.state !== "idle") {
 				live += 1;
 				show(index);
@@ -48,21 +69,30 @@ export function panelLines(self: AgentNode): string[] {
 	const lines = rows.filter((row) => row.shown).map((row) => row.line);
 	if (lines.length === 0) return [];
 	if (lines.length > MAX_LINES - 1) lines.splice(MAX_LINES - 2, Infinity, `+${lines.length - (MAX_LINES - 2)} more`);
-	return [`Tasks (${live} live, /tasks to hide)`, ...lines.map((line) => `  ${line}`)];
+	return [`Tasks (${live} live, /tasks to view)`, ...lines.map((line) => `  ${line}`)];
 }
 
-/** Keeps the task panel above the editor current while the Session is loaded; `toggle` hides or shows it. */
+/**
+ * Keeps the task panel above the editor current while the Session is loaded, and its times each second
+ * while it shows; `toggle` hides or shows it.
+ */
 export function createPanel(getNode: () => AgentNode | undefined) {
 	let ctx: ExtensionContext | undefined;
 	let hidden = false;
 	let scheduled = false;
 	let unsubscribe: (() => void) | undefined;
+	let ticker: ReturnType<typeof setInterval> | undefined;
 
 	const render = () => {
 		scheduled = false;
 		const node = getNode();
 		if (!ctx?.hasUI || !node) return;
 		const lines = hidden ? [] : panelLines(node);
+		if (lines.length > 0) ticker ??= setInterval(render, 1000);
+		else {
+			clearInterval(ticker);
+			ticker = undefined;
+		}
 		ctx.ui.setWidget(WIDGET_KEY, lines.length === 0 ? undefined : (_tui, theme) => {
 			const panel = new Container();
 			lines.forEach((line, index) => panel.addChild(new TruncatedText(theme.fg(index === 0 ? "accent" : "muted", line), 0, 0)));
@@ -87,6 +117,8 @@ export function createPanel(getNode: () => AgentNode | undefined) {
 			render();
 		},
 		stop() {
+			clearInterval(ticker);
+			ticker = undefined;
 			unsubscribe?.();
 			unsubscribe = undefined;
 			if (ctx?.hasUI) ctx.ui.setWidget(WIDGET_KEY, undefined);

package/src/ui/queue.ts

@@ -0,0 +1,47 @@
+import type { AgentSession, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { TruncatedText, type TUI } from "@earendil-works/pi-tui";
+import { messageBody, MESSAGE_TYPE, type Sender } from "../agents/registry.ts";
+
+const WIDGET_KEY = "pi-agents-queue";
+
+type Queued = { role: string; customType?: string; content?: unknown; details?: { delivery?: string; sender?: Sender } };
+
+/** One line for an Agent message: its sender and body. */
+function summary({ content, details }: Queued): string {
+	const text = typeof content === "string" ? content : "";
+	const sender = details?.sender;
+	const body = sender && messageBody(text, sender);
+	return (body === undefined ? text : `[${sender!.name ?? sender!.id}] ${body}`).replace(/\s+/g, " ").trim();
+}
+
+/** Agent messages waiting in the Session's queues, which Pi does not list as it lists the user's. */
+export function queuedLines(agent: AgentSession["agent"]): string[] {
+	const queues = agent as unknown as Record<"steeringQueue" | "followUpQueue", { messages: Queued[] }>;
+	const lines = (label: string, messages: Queued[]) => messages
+		.filter((message) => message.role === "custom" && message.customType === MESSAGE_TYPE)
+		// A write waits in the steering queue but asks for nothing, so it is labeled apart from a steer.
+		.map((message) => `${message.details?.delivery === "write" ? "Write" : label}: ${summary(message)}`);
+	return [...lines("Steering", queues.steeringQueue.messages), ...lines("Follow-up", queues.followUpQueue.messages)];
+}
+
+/**
+ * Lists the Agent messages queued in `session` under Pi's own queued messages. The list is read from
+ * the queues on each render, so it follows them as they are consumed or cleared; `update` renders it
+ * after a message is queued.
+ */
+export function showQueue(ctx: ExtensionContext, session: () => AgentSession | undefined) {
+	let tui: TUI | undefined;
+	if (ctx.hasUI) {
+		ctx.ui.setWidget(WIDGET_KEY, (current, theme) => {
+			tui = current;
+			return {
+				render: (width: number) => {
+					const agent = session()?.agent;
+					return agent ? queuedLines(agent).flatMap((line) => new TruncatedText(theme.fg("dim", line), 1, 0).render(width)) : [];
+				},
+				invalidate() {},
+			};
+		});
+	}
+	return { update: () => tui?.requestRender() };
+}