pi-squire 0.2.2 → 0.3.0

package/src/limits.ts

@@ -0,0 +1,242 @@
+// Usage-limit marks: which providers and models are out of quota, kept in one
+// JSON file in Pi's agent directory so every Pi session on the machine sees
+// them. A delegate records a mark when its run ends on a usage-limit error;
+// the parent refuses to launch on a marked model until the mark clears.
+import { mkdirSync, readFileSync, renameSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
+
+const MINUTE = 60_000;
+const HOUR: number = 60 * MINUTE;
+const FIRST_COOLDOWN: number = 5 * MINUTE;
+const MAX_COOLDOWN: number = 6 * HOUR;
+/** An expired mark stays this long, only so the next hit escalates. */
+const KEEP_EXPIRED: number = 24 * HOUR;
+const MAX_REASON = 300;
+const MAX_DELEGATIONS = 10;
+
+export type Mark = {
+	/** `provider` for a whole provider, `provider/model` for one model. */
+	scope: string;
+	reason: string;
+	/** Epoch milliseconds. */
+	recordedAt: number;
+	/** Epoch milliseconds. */
+	clearsAt: number;
+	source: "reactive" | "proactive";
+	hits: number;
+	/** The last delegation ids that hit this mark. */
+	delegations: string[];
+};
+
+/** The slice of Pi's model registry this module reads. */
+export type Models = { find(provider: string, id: string): { cost?: { input?: number; output?: number } } | undefined };
+
+export const limitsFile = (): string => join(getAgentDir(), "pi-squire-limits.json");
+
+function isMark(value: unknown): value is Mark {
+	const mark = value as Partial<Mark> | null;
+	return (
+		typeof mark === "object" &&
+		mark !== null &&
+		typeof mark.scope === "string" &&
+		typeof mark.reason === "string" &&
+		Number.isFinite(mark.recordedAt) &&
+		Number.isFinite(mark.clearsAt) &&
+		(mark.source === "reactive" || mark.source === "proactive") &&
+		Number.isFinite(mark.hits) &&
+		Array.isArray(mark.delegations)
+	);
+}
+
+/** Every mark still worth keeping, active or recently expired. Never throws. */
+export function readMarks(now = Date.now()): Mark[] {
+	try {
+		const parsed: unknown = JSON.parse(readFileSync(limitsFile(), "utf8"));
+		return Array.isArray(parsed) ? parsed.filter(isMark).filter((mark) => mark.clearsAt + KEEP_EXPIRED > now) : [];
+	} catch {
+		return [];
+	}
+}
+
+function writeMarks(marks: Mark[]): void {
+	const file = limitsFile();
+	mkdirSync(getAgentDir(), { recursive: true });
+	const temp = `${file}.${process.pid}.tmp`;
+	writeFileSync(temp, JSON.stringify(marks, null, 2));
+	renameSync(temp, file);
+}
+
+/** Whether a model is free: an id ending in `-free`, or zero cost in the registry. */
+export function isFree(models: Models, provider: string, id: string): boolean {
+	if (id.endsWith("-free")) return true;
+	const cost = models.find(provider, id)?.cost;
+	return cost !== undefined && cost.input === 0 && cost.output === 0;
+}
+
+/** The active mark that covers provider/id, if any. A provider mark never covers a free model. */
+export function activeMark(models: Models, provider: string, id: string, now = Date.now()): Mark | undefined {
+	const active = readMarks(now).filter((mark) => mark.clearsAt > now);
+	return (
+		active.find((mark) => mark.scope === `${provider}/${id}`) ??
+		(isFree(models, provider, id) ? undefined : active.find((mark) => mark.scope === provider))
+	);
+}
+
+const LIMIT_PATTERN =
+	/GoUsageLimitError|FreeUsageLimitError|Monthly usage limit reached|insufficient_quota|quota exceeded|subscription_sharing_usage_limit_exceeded|usage_limit_reached|usage limit reached|rate_limit_error/i;
+
+/** The reset time an error text states: "try again in N units", "retry after", or a `resets_at` field. */
+function statedReset(message: string, now: number): number | undefined {
+	const relative = /(?:try again in|retry after)\s+(\d+(?:\.\d+)?)\s*(second|minute|hour)s?/i.exec(message);
+	if (relative !== null) {
+		return (
+			now + Number(relative[1]) * { second: 1000, minute: MINUTE, hour: HOUR }[relative[2].toLowerCase() as "second"]
+		);
+	}
+	const text = /resets_?at\W{1,4}([0-9][0-9T:.+\-Z]*)/i.exec(message)?.[1];
+	if (text === undefined) return undefined;
+	return /^\d+(\.\d+)?$/.test(text) ? Number(text) * (Number(text) > 1e12 ? 1 : 1000) : Date.parse(text);
+}
+
+export type Response = { status: number; headers: Record<string, string> };
+
+const UNIT_MS: Record<string, number> = { ms: 1, s: 1000, m: MINUTE, h: HOUR };
+
+/** A duration string such as `1s` or `6m0s`, in milliseconds. */
+function duration(text: string): number | undefined {
+	const parts = [...text.matchAll(/(\d+(?:\.\d+)?)(ms|s|m|h)/g)];
+	return parts.length === 0 || parts.map((part) => part[0]).join("") !== text.trim()
+		? undefined
+		: parts.reduce((sum, part) => sum + Number(part[1]) * UNIT_MS[part[2]], 0);
+}
+
+/**
+ * The reset time a failed response's headers state, if any, in the order
+ * `retry-after`, `anthropic-ratelimit-*-reset`, `x-ratelimit-reset-*`. Only a
+ * non-2xx response counts, and only a time in the future.
+ */
+export function headerReset(response: Response | undefined, now = Date.now()): number | undefined {
+	if (response === undefined || (response.status >= 200 && response.status < 300)) return undefined;
+	const headers = Object.entries(response.headers).map(([name, value]) => [name.toLowerCase(), String(value)] as const);
+	const future = (times: (number | undefined)[]): number | undefined => {
+		const found = times.filter((time): time is number => time !== undefined && time > now);
+		return found.length === 0 ? undefined : Math.max(...found);
+	};
+	const retry = headers.find(([name]) => name === "retry-after")?.[1].trim();
+	let retryTime: number | undefined;
+	if (retry !== undefined) retryTime = /^\d+$/.test(retry) ? now + Number(retry) * 1000 : Date.parse(retry);
+	const retryAt = future([retryTime]);
+	if (retryAt !== undefined) return retryAt;
+	const anthropic = future(
+		headers.filter(([name]) => /^anthropic-ratelimit-.*-reset$/.test(name)).map(([, value]) => Date.parse(value)),
+	);
+	if (anthropic !== undefined) return anthropic;
+	return future(
+		headers
+			.filter(([name]) => name === "x-ratelimit-reset-requests" || name === "x-ratelimit-reset-tokens")
+			.map(([, value]) => {
+				const ms = duration(value);
+				return ms === undefined ? undefined : now + ms;
+			}),
+	);
+}
+
+/**
+ * Whether an error message (Pi's `<status>: <body>`) is a usage limit, and
+ * the reset time the text states, if any. Overloaded and 5xx errors never are.
+ */
+export function classify(message: string, now = Date.now()): { resetAt?: number } | undefined {
+	if (/overloaded/i.test(message) || /^\s*5\d\d\b/.test(message)) return undefined;
+	if (!LIMIT_PATTERN.test(message) && !/^\s*429\b/.test(message)) return undefined;
+	const resetAt = statedReset(message, now);
+	return resetAt !== undefined && resetAt > now ? { resetAt } : {};
+}
+
+/**
+ * Record a reactive mark for a delegate's usage-limit failure. The scope is
+ * the provider, or only the model for a free model or `FreeUsageLimitError`.
+ */
+export function recordLimit(
+	models: Models,
+	hit: { provider: string; model: string; message: string; delegation: string; response?: Response },
+	now = Date.now(),
+): void {
+	const limit = classify(hit.message, now);
+	if (limit === undefined) return;
+	const single = isFree(models, hit.provider, hit.model) || /FreeUsageLimitError/i.test(hit.message);
+	const scope = single ? `${hit.provider}/${hit.model}` : hit.provider;
+	const marks = readMarks(now);
+	const old = marks.find((mark) => mark.scope === scope);
+	const hits = (old?.hits ?? 0) + 1;
+	const mark: Mark = {
+		scope,
+		reason: hit.message.slice(0, MAX_REASON),
+		recordedAt: now,
+		clearsAt:
+			headerReset(hit.response, now) ?? limit.resetAt ?? now + Math.min(FIRST_COOLDOWN * 2 ** (hits - 1), MAX_COOLDOWN),
+		source: "reactive",
+		hits,
+		delegations: [...(old?.delegations ?? []), hit.delegation].slice(-MAX_DELEGATIONS),
+	};
+	writeMarks([...marks.filter((item) => item !== old), mark]);
+}
+
+/**
+ * Record a proactive mark: a quota reading says `provider` is used up. It
+ * clears at `clearsAt`, or after the escalating cooldown when none is known.
+ * A failed write is swallowed; the mark is returned either way.
+ */
+export function recordProactive(
+	provider: string,
+	reason: string,
+	clearsAt: number | undefined,
+	now = Date.now(),
+): Mark {
+	const marks = readMarks(now);
+	const old = marks.find((mark) => mark.scope === provider);
+	const hits = (old?.hits ?? 0) + 1;
+	const mark: Mark = {
+		scope: provider,
+		reason: reason.slice(0, MAX_REASON),
+		recordedAt: now,
+		clearsAt: clearsAt ?? now + Math.min(FIRST_COOLDOWN * 2 ** (hits - 1), MAX_COOLDOWN),
+		source: "proactive",
+		hits,
+		delegations: old?.delegations ?? [],
+	};
+	try {
+		writeMarks([...marks.filter((item) => item !== old), mark]);
+	} catch {
+		// The reading is real even if the file is not writable: the caller blocks on the returned mark.
+	}
+	return mark;
+}
+
+/**
+ * Remove marks for `target`: `all`, a provider (its mark and every model mark
+ * under it) or `provider/model` (that mark and the provider mark covering it).
+ * Removing also forgets the escalation history. Returns the scopes of the
+ * active marks removed; when none, the file is left alone.
+ */
+export function clearMarks(target: string, now = Date.now()): string[] {
+	const marks = readMarks(now);
+	const hit = (mark: Mark): boolean =>
+		target === "all" ||
+		mark.scope === target ||
+		(target.includes("/") ? mark.scope === target.slice(0, target.indexOf("/")) : mark.scope.startsWith(`${target}/`));
+	const cleared = marks.filter((mark) => hit(mark) && mark.clearsAt > now).map((mark) => mark.scope);
+	if (cleared.length > 0) writeMarks(marks.filter((mark) => !hit(mark)));
+	return cleared;
+}
+
+/** Remove the mark with exactly this scope, and with it the escalation history. False when it could not be written. */
+export function removeMark(scope: string, now = Date.now()): boolean {
+	try {
+		writeMarks(readMarks(now).filter((mark) => mark.scope !== scope));
+		return true;
+	} catch {
+		return false;
+	}
+}

package/src/results.ts

@@ -9,6 +9,8 @@ import { join } from "node:path";
 
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 
+import { readMarks } from "./limits.ts";
+
 /** The session entry type both delegation records and results use. */
 export const CUSTOM_TYPE = "delegate";
 /** The footer entry that counts delegations whose result has not arrived. */
@@ -157,6 +159,18 @@ export function findSession(dir: string | undefined, id: string): string | undef
 	return name === undefined ? undefined : join(dir, name);
 }
 
+/** The line naming a usage-limit mark this delegation's failure recorded, if any. */
+function limitLine(id: string): string | undefined {
+	const now = Date.now();
+	const mark = readMarks(now).find((item) => item.delegations.includes(id));
+	if (mark === undefined) return undefined;
+	const scope = mark.scope.includes("/") ? `model ${mark.scope}` : `provider ${mark.scope}`;
+	const until = new Date(mark.clearsAt).toISOString();
+	return `Usage limit: the delegate hit a usage limit on ${scope}; ${
+		mark.clearsAt > now ? `it is marked until ${until} and later delegations skip it.` : `it was marked until ${until}.`
+	}`;
+}
+
 /** The one message the parent sees for a settled delegation. */
 function resultText(delegation: Delegation, payload: Inbound, session: string | undefined): string {
 	const reply = payload.envelope;
@@ -176,9 +190,11 @@ function resultText(delegation: Delegation, payload: Inbound, session: string |
 	const cut = cap(reply.body, BODY_CAP);
 	const result =
 		cut === undefined ? reply.body : `${cut}\n[delegate] Body cut at 32 KiB; the full envelope is ${payload.path}`;
+	const limit = reply.status === "failed" ? limitLine(delegation.id) : undefined;
 	return [
 		`[delegate] Result from ${delegation.agent} (${delegation.model}), delegation ${delegation.id}, request ${delegation.requestId}.`,
 		`Status: ${reply.status}`,
+		...(limit === undefined ? [] : [limit]),
 		quote,
 		`Envelope: ${payload.path}`,
 		`Delegate session: ${session ?? delegation.id}`,

package/src/tmux.ts

@@ -25,10 +25,14 @@ export interface TmuxClient {
 }
 
 /** Run one tmux command and return its stdout; never a shell. */
-function run(args: string[]): Promise<string> {
+function runTmux(args: string[], timeout: number): Promise<string> {
 	return new Promise((resolve, reject) => {
-		execFile("tmux", args, { encoding: "utf8" }, (error, stdout, stderr) => {
+		execFile("tmux", args, { encoding: "utf8", timeout }, (error, stdout, stderr) => {
 			if (error) {
+				if (error.killed) {
+					reject(new Error(`tmux ${args[0]} timed out after ${timeout / 1000} s`));
+					return;
+				}
 				const detail = typeof stderr === "string" && stderr.trim() !== "" ? stderr.trim() : error.message;
 				reject(new Error(detail));
 				return;
@@ -43,7 +47,10 @@ function run(args: string[]): Promise<string> {
  * `TMUX_PANE`, not in whatever session a client happens to have attached, so
  * it works with several clients on one server.
  */
-export function createTmuxClient(): TmuxClient {
+export function createTmuxClient(options: { timeoutMs?: number } = {}): TmuxClient {
+	const timeoutMs = options.timeoutMs ?? 10_000;
+	const run = (args: string[]) => runTmux(args, timeoutMs);
+
 	/** The id of the session this pane belongs to, when known. */
 	async function targetSession(): Promise<string | undefined> {
 		const pane = process.env.TMUX_PANE;

package/src/tracking.ts

@@ -2,9 +2,12 @@
 // and the `delegation_status` and `delegation_close` tools. A delegation's
 // state comes from `results.ts` (`deriveState`); this module only asks tmux
 // about the window and words what it is told.
+import { statSync } from "node:fs";
+
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { Type } from "typebox";
 
+import { type Mark, readMarks } from "./limits.ts";
 import {
 	type DelegationState,
 	deriveState,
@@ -15,6 +18,7 @@ import {
 	type WindowState,
 } from "./results.ts";
 import type { TmuxClient, TmuxWindow } from "./tmux.ts";
+import type { ReadingError } from "./usage.ts";
 
 /** The hook a tick emits so the provider claims waiting replies before windows are checked. */
 const SCAN = "message:scan";
@@ -97,21 +101,45 @@ function stateLine(state: DelegationState, record: Recorded): string {
 }
 
 /** One delegation as `delegation_status` reports it; its text and details both read from this. */
-type View = { record: Recorded; state: DelegationState; name: string; sessionPath: string | undefined };
+type View = {
+	record: Recorded;
+	state: DelegationState;
+	name: string;
+	sessionPath: string | undefined;
+	/** When the session file was last written; set only for a running delegation with a file. */
+	lastActive: Date | undefined;
+};
+
+/** How long ago, worded for the `state:` line: `just now`, `N min ago` or `N h M min ago`. */
+function ago(ms: number): string {
+	const minutes = Math.floor(ms / 60_000);
+	if (minutes < 1) return "just now";
+	if (minutes < 60) return `${minutes} min ago`;
+	return `${Math.floor(minutes / 60)} h ${minutes % 60} min ago`;
+}
 
 /** The view of a delegation: its state and the fields the text and details share. */
 function view(record: Recorded, windows: TmuxWindow[] | undefined, ctx: ExtensionContext): View {
 	const state = deriveState(record, windowState(windows, record));
 	const sessionPath = record.result?.sessionPath ?? findSession(ctx.sessionManager.getSessionDir(), record.id);
-	return { record, state, name: recordedName(record), sessionPath };
+	let lastActive: Date | undefined;
+	if (state === "running" && sessionPath !== undefined) {
+		try {
+			lastActive = statSync(sessionPath).mtime;
+		} catch {
+			// a failed stat counts as no file
+		}
+	}
+	return { record, state, name: recordedName(record), sessionPath, lastActive };
 }
 
 /** The text block for one delegation. */
-function viewText({ record, state, name, sessionPath }: View): string {
+function viewText({ record, state, name, sessionPath, lastActive }: View): string {
+	const seen = lastActive === undefined ? "" : ` (last active ${ago(Date.now() - lastActive.getTime())})`;
 	const lines = [
 		`${record.id} ${record.agent} (${record.model}, thinking ${record.thinking})`,
 		`  name: ${name}`,
-		`  state: ${stateLine(state, record)}`,
+		`  state: ${stateLine(state, record)}${seen}`,
 		`  window: ${record.windowName} (${record.windowId})`,
 		`  auto-exit: ${autoExitLine(record)}`,
 		`  session: ${sessionPath ?? `${record.id} (no session file yet)`}`,
@@ -121,7 +149,7 @@ function viewText({ record, state, name, sessionPath }: View): string {
 }
 
 /** The structured details for one delegation. */
-function viewDetails({ record, name, sessionPath }: View): Record<string, unknown> {
+function viewDetails({ record, name, sessionPath, lastActive }: View): Record<string, unknown> {
 	return {
 		id: record.id,
 		name,
@@ -136,9 +164,23 @@ function viewDetails({ record, name, sessionPath }: View): Record<string, unknow
 		status: record.result?.status,
 		envelopePath: record.result?.envelopePath,
 		sessionPath,
+		...(lastActive === undefined ? {} : { lastActiveAt: lastActive.toISOString() }),
 	};
 }
 
+/** The usage-limit marks in force now, as `delegation_status` reports them. */
+function activeMarks(): { scope: string; source: Mark["source"]; reason: string; clearsAt: string }[] {
+	const now = Date.now();
+	return readMarks(now)
+		.filter((mark) => mark.clearsAt > now)
+		.map(({ scope, source, reason, clearsAt }) => ({
+			scope,
+			source,
+			reason: reason.length > 120 ? `${reason.slice(0, 119)}…` : reason,
+			clearsAt: new Date(clearsAt).toISOString(),
+		}));
+}
+
 /** Every delegation, or the one named; throws when a named id is unknown. */
 function selectRecords(results: Results, id: string | undefined): Recorded[] {
 	const all = results.list();
@@ -216,12 +258,12 @@ function createPoll(
 }
 
 /** The `delegation_status` tool. */
-function statusTool(results: Results, tmux: TmuxClient): ToolDef {
+function statusTool(results: Results, tmux: TmuxClient, readingErrors: () => ReadingError[]): ToolDef {
 	return {
 		name: "delegation_status",
 		label: "Delegation status",
 		description:
-			"The delegations this session started, each reported as running (the window is open and no result has arrived yet), done (with the result envelope's status) or closed (closed with delegation_close, or its window is gone). A delegation whose window went with no reply is reported closed without a result until a late reply makes it done. Names the delegation id, agent, session name, model, thinking, window, delegate session path and result envelope path. A delegate's window being open means it is connected, not that its task is unfinished or finished; only a result means done.",
+			"The delegations this session started, each reported as running (the window is open and no result has arrived yet), done (with the result envelope's status) or closed (closed with delegation_close, or its window is gone). A delegation whose window went with no reply is reported closed without a result until a late reply makes it done. Names the delegation id, agent, session name, model, thinking, window, delegate session path and result envelope path. A running delegation also shows when its session file was last written; a long gap may mean it is stuck but does not prove it. A delegate's window being open means it is connected, not that its task is unfinished or finished; only a result means done.",
 		parameters: Type.Object({
 			id: Type.Optional(Type.String({ description: "One delegation id; omit to report every delegation." })),
 		}),
@@ -235,10 +277,25 @@ function statusTool(results: Results, tmux: TmuxClient): ToolDef {
 			const records = selectRecords(results, params.id);
 			const windows = await listWindows(tmux);
 			const views = records.map((record) => view(record, windows, ctx));
-			const details = { delegations: views.map(viewDetails) };
-			if (views.length === 0) return toolResult("No delegations are recorded in this session.", details);
+			const marks = activeMarks();
+			const errors = readingErrors();
+			const details = { delegations: views.map(viewDetails), marks, readingErrors: errors };
+			const block =
+				marks.length === 0 && errors.length === 0
+					? []
+					: [
+							[
+								"Usage-limit marks:",
+								...marks.map((m) => `- ${m.scope} (${m.source}) until ${m.clearsAt}: ${m.reason}`),
+								...errors.map((e) => `- could not read ${e.provider} quota at ${e.at}: ${e.reason}`),
+							].join("\n"),
+						];
 			const heading = `${views.length} delegation${views.length === 1 ? "" : "s"}:`;
-			return toolResult(`${heading}\n\n${views.map(viewText).join("\n\n")}`, details);
+			const parts =
+				views.length === 0
+					? ["No delegations are recorded in this session.", ...block]
+					: [`${heading}\n\n${views.map(viewText).join("\n\n")}`, ...block];
+			return toolResult(parts.join("\n\n"), details);
 		},
 	};
 }
@@ -292,11 +349,16 @@ function closeTool(results: Results, tmux: TmuxClient): ToolDef {
 }
 
 /** The poll and the two tools, over the extension API, the records and tmux. */
-export function createTracking(pi: ExtensionAPI, results: Results, tmux: TmuxClient): Tracking {
+export function createTracking(
+	pi: ExtensionAPI,
+	results: Results,
+	tmux: TmuxClient,
+	readingErrors: () => ReadingError[] = () => [],
+): Tracking {
 	const poll = createPoll(pi, results, tmux);
 	return {
 		registerTools(): void {
-			pi.registerTool(statusTool(results, tmux) as never);
+			pi.registerTool(statusTool(results, tmux, readingErrors) as never);
 			pi.registerTool(closeTool(results, tmux) as never);
 		},
 		startPolling: poll.start,