@oh-my-pi/pi-coding-agent 15.10.11 → 15.11.0

package/src/thinking.ts

@@ -71,6 +71,13 @@ export function toReasoningEffort(level: ThinkingLevel | undefined): Effort | un
 	return level;
 }
 
+/**
+ * True when a selector explicitly requests provider-side reasoning disablement.
+ */
+export function shouldDisableReasoning(level: ThinkingLevel | undefined): boolean {
+	return level === ThinkingLevel.Off;
+}
+
 /**
  * Resolves a selector against the current model while preserving explicit "off".
  */

package/src/tiny/title-client.ts

@@ -39,6 +39,17 @@ export interface TinyTitleDownloadOptions {
 	onProgress?: (event: TinyTitleProgressEvent) => void;
 }
 
+/**
+ * Per-request controls for {@link TinyTitleClient.generate}.
+ *
+ * Carries the optional abort signal and title-system-prompt override used by
+ * callers that customize automatic session-title generation.
+ */
+export interface TinyTitleGenerateOptions {
+	signal?: AbortSignal;
+	systemPrompt?: string;
+}
+
 // Cold-starting the worker subprocess from a compiled binary (decompress + module
 // graph load) is slow on contended CI runners — the macos-15-intel release smoke
 // blew past 5s while arm64/linux/win passed. The probe only needs to prove the
@@ -46,6 +57,14 @@ export interface TinyTitleDownloadOptions {
 // generous bound removes the flake without weakening the check.
 const SMOKE_TEST_TIMEOUT_MS = 30_000;
 
+function normalizeTinyTitleGenerateOptions(
+	options: AbortSignal | TinyTitleGenerateOptions | undefined,
+): TinyTitleGenerateOptions {
+	if (!options) return {};
+	if ("aborted" in options && "addEventListener" in options) return { signal: options };
+	return options;
+}
+
 /**
  * Hidden subcommand on the main CLI that boots the tiny-model worker in the
  * spawned subprocess. Kept in sync with the dispatch in `cli.ts`.
@@ -295,9 +314,16 @@ export class TinyTitleClient {
 		return () => this.#progressListeners.delete(listener);
 	}
 
-	async generate(modelKey: string, message: string, signal?: AbortSignal): Promise<string | null> {
+	async generate(modelKey: string, message: string, signal?: AbortSignal): Promise<string | null>;
+	async generate(modelKey: string, message: string, options?: TinyTitleGenerateOptions): Promise<string | null>;
+	async generate(
+		modelKey: string,
+		message: string,
+		optionsOrSignal?: AbortSignal | TinyTitleGenerateOptions,
+	): Promise<string | null> {
+		const options = normalizeTinyTitleGenerateOptions(optionsOrSignal);
 		if (!isTinyTitleLocalModelKey(modelKey)) return null;
-		if (signal?.aborted) return null;
+		if (options.signal?.aborted) return null;
 
 		try {
 			const worker = this.#ensureWorker();
@@ -310,12 +336,15 @@ export class TinyTitleClient {
 				this.#pending.delete(id);
 				pending.resolve(null);
 			};
-			signal?.addEventListener("abort", abort, { once: true });
+			options.signal?.addEventListener("abort", abort, { once: true });
 			try {
-				worker.send({ type: "generate", id, modelKey, message });
+				const request: TinyTitleWorkerInbound = options.systemPrompt
+					? { type: "generate", id, modelKey, message, systemPrompt: options.systemPrompt }
+					: { type: "generate", id, modelKey, message };
+				worker.send(request);
 				return await promise;
 			} finally {
-				signal?.removeEventListener("abort", abort);
+				options.signal?.removeEventListener("abort", abort);
 				this.#pending.delete(id);
 			}
 		} catch (error) {

package/src/tiny/title-protocol.ts

@@ -29,7 +29,7 @@ export interface TinyTitleProgressEvent {
 
 export type TinyTitleWorkerInbound =
 	| { type: "ping"; id: string }
-	| { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string }
+	| { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string; systemPrompt?: string }
 	| { type: "complete"; id: string; modelKey: TinyLocalModelKey; prompt: string; maxTokens?: number }
 	| { type: "download"; id: string; modelKey: TinyLocalModelKey };
 

package/src/tiny/worker.ts

@@ -436,9 +436,10 @@ async function loadPipeline(
 	return loaded;
 }
 
-function buildPrompt(generator: TextGenerationPipeline, message: string): string {
+function buildPrompt(generator: TextGenerationPipeline, message: string, systemPrompt?: string): string {
+	const selectedSystemPrompt = systemPrompt?.trim() || TINY_TITLE_SYSTEM_PROMPT;
 	const chat = [
-		{ role: "system", content: TINY_TITLE_SYSTEM_PROMPT },
+		{ role: "system", content: selectedSystemPrompt },
 		{ role: "user", content: formatTitleUserMessage(message) },
 	];
 	const chatTemplateOptions = {
@@ -464,9 +465,10 @@ async function generateTitle(
 	requestId: string,
 	modelKey: TinyTitleLocalModelKey,
 	message: string,
+	systemPrompt?: string,
 ): Promise<string | null> {
 	const generator = await loadPipeline(modelKey, transport, requestId);
-	const promptText = buildPrompt(generator, message);
+	const promptText = buildPrompt(generator, message, systemPrompt);
 	const transformers = await loadTransformers(transport, requestId, modelKey);
 	const output = (await generator(promptText, {
 		max_new_tokens: TITLE_MAX_NEW_TOKENS,
@@ -548,7 +550,7 @@ async function handleQueuedRequest(
 			transport.send({ type: "completion", id: request.id, text });
 			return;
 		}
-		const title = await generateTitle(transport, request.id, request.modelKey, request.message);
+		const title = await generateTitle(transport, request.id, request.modelKey, request.message, request.systemPrompt);
 		transport.send({ type: "title", id: request.id, title });
 	} catch (error) {
 		transport.send({ type: "error", id: request.id, error: errorText(error) });

package/src/tools/ask.ts

@@ -104,7 +104,7 @@ const RECOMMENDED_SUFFIX = " (Recommended)";
 const TIMEOUT_DETECTION_TOLERANCE_MS = 1_000;
 
 function getDoneOptionLabel(): string {
-	return `${theme.symbol("tool.ask")} Done selecting`;
+	return `${theme.status.success} Done selecting`;
 }
 
 /** Add "(Recommended)" suffix to the option at the given index if not already present */
@@ -694,7 +694,9 @@ function normalizeRenderQuestions(raw: unknown): NonNullable<AskRenderArgs["ques
 /** Render a custom free-text answer as a status line plus indented continuation rows. */
 function renderCustomInputLines(uiTheme: Theme, customInput: string): string[] {
 	const lines = customInput.split("\n");
-	const out: string[] = [` ${uiTheme.styledSymbol("tool.ask", "accent")} ${uiTheme.fg("toolOutput", lines[0] ?? "")}`];
+	const out: string[] = [
+		` ${uiTheme.styledSymbol("status.success", "success")} ${uiTheme.fg("toolOutput", lines[0] ?? "")}`,
+	];
 	for (let i = 1; i < lines.length; i++) out.push(`   ${uiTheme.fg("toolOutput", lines[i])}`);
 	return out;
 }

package/src/tools/bash.ts

@@ -17,7 +17,7 @@ import { truncateToVisualLines } from "../modes/components/visual-truncate";
 import { highlightCode, type Theme } from "../modes/theme/theme";
 import bashDescription from "../prompts/tools/bash.md" with { type: "text" };
 import type { ClientBridgeTerminalExitStatus, ClientBridgeTerminalOutput } from "../session/client-bridge";
-import { DEFAULT_MAX_BYTES, streamTailUpdates, TailBuffer } from "../session/streaming-output";
+import { DEFAULT_MAX_BYTES, enforceInlineByteCap, streamTailUpdates, TailBuffer } from "../session/streaming-output";
 import { renderStatusLine } from "../tui";
 import { CachedOutputBlock, markFramedBlockComponent } from "../tui/output-block";
 import { getSixelLineMask } from "../utils/sixel";
@@ -429,7 +429,7 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
 		}
 	}
 
-	#buildCompletedResult(
+	async #buildCompletedResult(
 		result: BashResult | BashInteractiveResult,
 		timeoutSec: number,
 		options: {
@@ -438,7 +438,7 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
 			terminalId?: string;
 			wallTimeMs?: number;
 		} = {},
-	): AgentToolResult<BashToolDetails> {
+	): Promise<AgentToolResult<BashToolDetails>> {
 		const exitCode = result.exitCode;
 		const failedExit = exitCode !== undefined && exitCode !== 0;
 
@@ -472,7 +472,17 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
 		if (failedExit) {
 			details.exitCode = exitCode;
 		}
-		const resultBuilder = toolResult(details).text(outputText).truncationFromSummary(result, { direction: "tail" });
+		// Final defense at the tool-result boundary: no bash path (client bridge,
+		// head-retention spill, minimizer miss) may emit more than
+		// ~DEFAULT_MAX_BYTES inline. No-op for already-bounded output.
+		const cappedOutputText = await enforceInlineByteCap(outputText, {
+			label: "bash output",
+			saveArtifact: full => saveBashOriginalArtifact(this.session, full),
+		});
+
+		const resultBuilder = toolResult(details)
+			.text(cappedOutputText)
+			.truncationFromSummary(result, { direction: "tail" });
 		if (failedExit) resultBuilder.error();
 		return resultBuilder.done();
 	}
@@ -560,7 +570,7 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
 						onMinimizedSave: originalText => saveBashOriginalArtifact(this.session, originalText),
 					});
 					const wallTimeMs = performance.now() - wallTimeStart;
-					const finalResult = this.#buildCompletedResult(result, options.timeoutSec, {
+					const finalResult = await this.#buildCompletedResult(result, options.timeoutSec, {
 						requestedTimeoutSec: options.requestedTimeoutSec,
 						notices: options.notices ?? [],
 						wallTimeMs,
@@ -1212,6 +1222,22 @@ export function createShellRenderer<TArgs>(config: ShellRendererConfig<TArgs>) {
 			const details = result.details;
 			const outputBlock = new CachedOutputBlock();
 
+			// Per-instance cache for the expensive inner lines computation. Mirrors
+			// the eval-renderer pattern (`eval-render.ts:709-752`): without this,
+			// every TUI repaint (one per keystroke when a long transcript is on
+			// screen) re-runs `split` / `replaceTabs` / `truncateToVisualLines` over
+			// the whole stored output for every bash row in scrollback. With a
+			// 50KB-tail bash result times hundreds of rows, that re-rendering is
+			// what pinned the main thread in issue #2081 and made keystrokes feel
+			// like the CPU was at 100%. The cache key includes every render input
+			// that materially affects the produced lines.
+			let cachedWidth: number | undefined;
+			let cachedPreviewLines: number | undefined;
+			let cachedExpanded: boolean | undefined;
+			let cachedRawOutput: string | undefined;
+			let cachedIsPartial: boolean | undefined;
+			let cachedLines: readonly string[] | undefined;
+
 			return markFramedBlockComponent({
 				render: (width: number): readonly string[] => {
 					// REACTIVE: read mutable options at render time
@@ -1223,6 +1249,19 @@ export function createShellRenderer<TArgs>(config: ShellRendererConfig<TArgs>) {
 					// Strip the LLM-facing notice appended by wrappedExecute so we don't
 					// double-print it alongside the styled warning line below.
 					const rawOutput = renderContext?.output ?? result.content?.find(c => c.type === "text")?.text ?? "";
+
+					const isPartial = options.isPartial === true;
+
+					if (
+						cachedLines !== undefined &&
+						cachedWidth === width &&
+						cachedPreviewLines === previewLines &&
+						cachedExpanded === expanded &&
+						cachedRawOutput === rawOutput &&
+						cachedIsPartial === isPartial
+					) {
+						return cachedLines;
+					}
 					const strippedOutput = stripOutputNotice(rawOutput, details?.meta);
 					const withoutExit = stripExitCodeNotice(strippedOutput, details?.exitCode);
 					const withoutWall = stripWallTimeNotice(withoutExit, details?.wallTimeMs);
@@ -1299,15 +1338,13 @@ export function createShellRenderer<TArgs>(config: ShellRendererConfig<TArgs>) {
 					if (timeoutLine) outputLines.push(timeoutLine);
 					if (warningLine) outputLines.push(warningLine);
 
-					return outputBlock.render(
+					const framed = outputBlock.render(
 						{
 							header,
-							state: options.isPartial ? "pending" : isError ? "error" : "success",
+							state: isPartial ? "pending" : isError ? "error" : "success",
 							sections: [
 								{
-									lines: options.isPartial
-										? capPreviewLines(cmdLines ?? [], uiTheme, { expanded })
-										: (cmdLines ?? []),
+									lines: isPartial ? capPreviewLines(cmdLines ?? [], uiTheme, { expanded }) : (cmdLines ?? []),
 								},
 								{ label: uiTheme.fg("toolTitle", "Output"), lines: outputLines },
 							],
@@ -1315,9 +1352,23 @@ export function createShellRenderer<TArgs>(config: ShellRendererConfig<TArgs>) {
 						},
 						uiTheme,
 					);
+
+					cachedWidth = width;
+					cachedPreviewLines = previewLines;
+					cachedExpanded = expanded;
+					cachedRawOutput = rawOutput;
+					cachedIsPartial = isPartial;
+					cachedLines = framed;
+					return framed;
 				},
 				invalidate: () => {
 					outputBlock.invalidate();
+					cachedLines = undefined;
+					cachedWidth = undefined;
+					cachedPreviewLines = undefined;
+					cachedExpanded = undefined;
+					cachedRawOutput = undefined;
+					cachedIsPartial = undefined;
 				},
 			});
 		},

package/src/tools/browser/tab-worker.ts

@@ -10,6 +10,7 @@ import type {
 	ElementHandle,
 	ElementScreenshotOptions,
 	HTTPResponse,
+	ImageFormat,
 	KeyInput,
 	Page,
 	SerializedAXNode,
@@ -436,6 +437,19 @@ export function describeScreenshot(opts?: ScreenshotOptions): string {
 	return "tab.screenshot()";
 }
 
+/** Map an explicit save path's extension to a puppeteer capture format (default png). */
+export function imageFormatForPath(filePath: string): ImageFormat {
+	switch (path.extname(filePath).toLowerCase()) {
+		case ".webp":
+			return "webp";
+		case ".jpg":
+		case ".jpeg":
+			return "jpeg";
+		default:
+			return "png";
+	}
+}
+
 /** Summarize still-running helpers (oldest first) so a cell timeout names what stalled. */
 export function describeInflight(inflight: Map<number, InflightOp>): string {
 	const now = Date.now();
@@ -931,6 +945,12 @@ export class WorkerCore {
 	): Promise<ScreenshotResult> {
 		const page = this.#requirePage();
 		const fullPage = opts.selector ? false : (opts.fullPage ?? false);
+		// An explicit save path picks the full-res capture format: puppeteer encodes
+		// png/jpeg/webp natively, so `save: "shot.webp"` gets real WebP bytes instead
+		// of PNG bytes hiding behind a .webp name. Unknown/missing extensions stay PNG.
+		const explicitPath = opts.save ? resolveToCwd(opts.save, session.cwd) : undefined;
+		const captureType = explicitPath ? imageFormatForPath(explicitPath) : "png";
+		const captureMime = `image/${captureType}` as const;
 		let buffer: Buffer;
 		if (opts.selector) {
 			const handle = (await untilAborted(signal, () =>
@@ -951,24 +971,23 @@ export class WorkerCore {
 				).catch(() => undefined);
 				// scrollIntoView:false skips the same IntersectionObserver check inside screenshot();
 				// captureBeyondViewport (puppeteer's default) still renders the clipped region.
-				const shotOpts: ElementScreenshotOptions = { type: "png", scrollIntoView: false };
+				const shotOpts: ElementScreenshotOptions = { type: captureType, scrollIntoView: false };
 				buffer = (await untilAborted(signal, () => handle.screenshot(shotOpts))) as Buffer;
 			} finally {
 				await handle.dispose().catch(() => undefined);
 			}
 		} else {
-			buffer = (await untilAborted(signal, () => page.screenshot({ type: "png", fullPage }))) as Buffer;
+			buffer = (await untilAborted(signal, () => page.screenshot({ type: captureType, fullPage }))) as Buffer;
 		}
 		const resized = await resizeImage(
-			{ type: "image", data: buffer.toBase64(), mimeType: "image/png" },
+			{ type: "image", data: buffer.toBase64(), mimeType: captureMime },
 			{ maxWidth: 1024, maxHeight: 1024, maxBytes: 150 * 1024, jpegQuality: 70 },
 		);
-		const explicitPath = opts.save ? resolveToCwd(opts.save, session.cwd) : undefined;
 		const saveFullRes = !!(explicitPath || session.browserScreenshotDir);
 		const savedBuffer = saveFullRes ? buffer : resized.buffer;
-		const savedMimeType = saveFullRes ? "image/png" : resized.mimeType;
-		// Auto-generated names must match the bytes we actually write: full-res is always
-		// PNG, but the resized buffer is whichever of PNG/JPEG/WebP encoded smallest.
+		const savedMimeType = saveFullRes ? captureMime : resized.mimeType;
+		// Names must match the bytes we actually write: full-res follows the capture
+		// format, the resized buffer is whichever of PNG/JPEG/WebP encoded smallest.
 		const ext = savedMimeType === "image/webp" ? "webp" : savedMimeType === "image/jpeg" ? "jpg" : "png";
 		const dest =
 			explicitPath ??

package/src/tools/browser.ts

@@ -3,6 +3,7 @@ import { prompt, untilAborted } from "@oh-my-pi/pi-utils";
 import * as z from "zod/v4";
 import browserDescription from "../prompts/tools/browser.md" with { type: "text" };
 import type { ToolSession } from "../sdk";
+import { enforceInlineByteCap } from "../session/streaming-output";
 import { truncateForPrompt } from "./approval";
 import { acquireBrowser, type BrowserHandle, type BrowserKind, type BrowserKindTag } from "./browser/registry";
 import type { Observation, ScreenshotResult } from "./browser/tab-protocol";
@@ -271,11 +272,37 @@ export class BrowserTool implements AgentTool<typeof browserSchema, BrowserToolD
 			.filter((c): c is { type: "text"; text: string } => c.type === "text")
 			.map(c => c.text)
 			.join("\n");
-		details.result = textOnly;
+		// Final defense at the tool-result boundary: a single run can display
+		// tens of KB (large JSON returns, dumped observations). Cap the combined
+		// text inline; the full text stays recoverable via the artifact footer
+		// when allocation succeeds.
+		const cappedText = await enforceInlineByteCap(textOnly, {
+			label: "browser output",
+			saveArtifact: full => saveBrowserOutputArtifact(this.session, full),
+		});
+		details.result = cappedText;
+		if (cappedText !== textOnly) {
+			const nonText = content.filter(c => c.type !== "text");
+			return toolResult(details)
+				.content([...nonText, { type: "text", text: cappedText }])
+				.done();
+		}
 		return toolResult(details).content(content).done();
 	}
 }
 
+/** Persist over-cap browser run output as a session artifact; mirrors the bash minimizer's save path. */
+async function saveBrowserOutputArtifact(session: ToolSession, fullText: string): Promise<string | undefined> {
+	try {
+		const alloc = await session.allocateOutputArtifact?.("browser-original");
+		if (!alloc?.path || !alloc.id) return undefined;
+		await Bun.write(alloc.path, fullText);
+		return alloc.id;
+	} catch {
+		return undefined;
+	}
+}
+
 function describeBrowser(handle: BrowserHandle): string {
 	switch (handle.kind.kind) {
 		case "headless":

package/src/tools/find.ts

@@ -4,7 +4,7 @@ import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallb
 import * as natives from "@oh-my-pi/pi-natives";
 import type { Component } from "@oh-my-pi/pi-tui";
 import { Text } from "@oh-my-pi/pi-tui";
-import { isEnoent, prompt, untilAborted } from "@oh-my-pi/pi-utils";
+import { formatGroupedPaths, isEnoent, prompt, untilAborted } from "@oh-my-pi/pi-utils";
 import * as z from "zod/v4";
 import type { RenderResultOptions } from "../extensibility/custom-tools/types";
 import { InternalUrlRouter } from "../internal-urls";
@@ -13,7 +13,6 @@ import findDescription from "../prompts/tools/find.md" with { type: "text" };
 import { type TruncationResult, truncateHead } from "../session/streaming-output";
 import { Ellipsis, fileHyperlink, renderFileList, renderStatusLine, renderTreeList, truncateToWidth } from "../tui";
 import type { ToolSession } from ".";
-import { buildPathTree, walkPathTree } from "./grouped-file-output";
 import { applyListLimit } from "./list-limit";
 import { formatFullOutputReference, type OutputMeta } from "./output-meta";
 import {
@@ -54,30 +53,6 @@ const DEFAULT_GLOB_TIMEOUT_MS = 5000;
 const MIN_GLOB_TIMEOUT_MS = 500;
 const MAX_GLOB_TIMEOUT_MS = 60_000;
 
-/**
- * Group find matches into a multi-level directory tree so the model doesn't pay
- * repeated tokens for shared path prefixes. Single-child directory chains fold
- * into one header (`# a/b/c/`), so a common prefix — including an absolute root
- * for out-of-cwd results — collapses to a single line. Each level adds one `#`;
- * files are listed bare under the deepest directory header that owns them.
- *
- * Order follows the input (mtime-desc for native glob): a directory appears when
- * its first member is emitted, and a node's own files precede its subdirectories.
- */
-export function formatFindGroupedOutput(paths: readonly string[]): string {
-	if (paths.length === 0) return "";
-	const tree = buildPathTree(paths.map(entry => ({ path: entry, isDir: entry.endsWith("/") })));
-	const lines: string[] = [];
-	for (const event of walkPathTree(tree)) {
-		if (event.kind === "dir") {
-			lines.push(`${"#".repeat(event.depth + 1)} ${event.name}/`);
-		} else {
-			lines.push(event.name);
-		}
-	}
-	return lines.join("\n");
-}
-
 export interface FindToolDetails {
 	truncation?: TruncationResult;
 	resultLimitReached?: number;
@@ -270,7 +245,7 @@ export class FindTool implements AgentTool<typeof findSchema, FindToolDetails> {
 				const listLimit = applyListLimit(files, { limit: effectiveLimit });
 				const limited = listLimit.items;
 				const limitMeta = listLimit.meta;
-				const baseOutput = formatFindGroupedOutput(limited);
+				const baseOutput = formatGroupedPaths(limited);
 				const trailingNotes: string[] = [];
 				if (notice) trailingNotes.push(notice);
 				if (missingPathsNote) trailingNotes.push(missingPathsNote);

package/src/tools/grouped-file-output.ts

@@ -1,123 +1,6 @@
 import path from "node:path";
 
-const URL_LIKE_PATH_RE = /^[a-z][a-z0-9+.-]*:\/\//i;
-
-function isUrlLikePath(filePath: string): boolean {
-	return URL_LIKE_PATH_RE.test(filePath);
-}
-
-// =============================================================================
-// Multi-level path tree
-// =============================================================================
-//
-// File listings (grep / ast-grep / ast-edit / lsp diagnostics / find) used to
-// group by the *immediate* parent directory and print the full directory path in
-// every header. For results spread across a deep tree — or rooted outside cwd,
-// where paths stay absolute — that repeated the shared prefix on every line. The
-// tree below folds single-child directory chains (so the common prefix collapses
-// into one header) and nests the rest, charging the model one token per path
-// segment instead of one per file.
-
-interface PathTreeNode {
-	/** Direct file leaves, in first-seen order. */
-	files: Array<{ name: string; key: string }>;
-	/** Dedup set for `files` (a glob can surface the same path twice on retry). */
-	fileNames: Set<string>;
-	/** Child directories, in first-seen order. */
-	subdirs: Array<{ name: string; node: PathTreeNode }>;
-	/** Dedup index for `subdirs`. */
-	dirIndex: Map<string, PathTreeNode>;
-}
-
-export interface PathTreeInput {
-	/** Path string; absolute, cwd-relative, or url-like. Backslashes are normalized. */
-	path: string;
-	/** Whether the leaf itself is a directory (trailing-slash match from find). */
-	isDir: boolean;
-	/** Opaque key carried onto file events for section lookup. Defaults to `path`. */
-	key?: string;
-}
-
-/** One node emitted while walking the tree: a folded directory or a file leaf. */
-export interface GroupedTreeEvent {
-	kind: "dir" | "file";
-	/** 0-based nesting depth (root children are depth 0). */
-	depth: number;
-	/** Folded chain for dirs (e.g. `a/b/c`, no trailing slash); basename for files. */
-	name: string;
-	/** File key for `kind === "file"`; empty string for directories. */
-	key: string;
-}
-
-function createNode(): PathTreeNode {
-	return { files: [], fileNames: new Set(), subdirs: [], dirIndex: new Map() };
-}
-
-function addFile(node: PathTreeNode, name: string, key: string): void {
-	if (node.fileNames.has(name)) return;
-	node.fileNames.add(name);
-	node.files.push({ name, key });
-}
-
-/**
- * Build a directory tree from a flat list of paths. URL-like entries are kept
- * whole as root-level file leaves (they have no meaningful directory structure).
- * Absolute paths carry a leading empty segment so they share a common `/` root
- * and fold like any other prefix.
- */
-export function buildPathTree(entries: Iterable<PathTreeInput>): PathTreeNode {
-	const root = createNode();
-	for (const { path: rawPath, isDir, key } of entries) {
-		const normalized = rawPath.replace(/\\/g, "/");
-		const fileKey = key ?? rawPath;
-		if (isUrlLikePath(normalized)) {
-			addFile(root, normalized, fileKey);
-			continue;
-		}
-		const trimmed = normalized.endsWith("/") ? normalized.slice(0, -1) : normalized;
-		if (trimmed.length === 0) continue;
-		const segments = trimmed.split("/");
-		const dirCount = isDir ? segments.length : segments.length - 1;
-		let node = root;
-		for (let i = 0; i < dirCount; i++) {
-			const segment = segments[i]!;
-			let child = node.dirIndex.get(segment);
-			if (!child) {
-				child = createNode();
-				node.dirIndex.set(segment, child);
-				node.subdirs.push({ name: segment, node: child });
-			}
-			node = child;
-		}
-		if (!isDir) {
-			addFile(node, segments[segments.length - 1]!, fileKey);
-		}
-	}
-	return root;
-}
-
-/**
- * Depth-first walk yielding directory and file events. Directories collapse their
- * single-child chains (`a` → `a/b` → `a/b/c`) so a shared prefix becomes one
- * header. Each node's direct files are emitted before its subdirectories, keeping
- * a file unambiguously attached to the header above it.
- */
-export function* walkPathTree(node: PathTreeNode, depth = 0): Generator<GroupedTreeEvent> {
-	for (const file of node.files) {
-		yield { kind: "file", depth, name: file.name, key: file.key };
-	}
-	for (const subdir of node.subdirs) {
-		let dirNode = subdir.node;
-		const parts = [subdir.name];
-		while (dirNode.files.length === 0 && dirNode.subdirs.length === 1) {
-			const only = dirNode.subdirs[0]!;
-			parts.push(only.name);
-			dirNode = only.node;
-		}
-		yield { kind: "dir", depth, name: parts.join("/"), key: "" };
-		yield* walkPathTree(dirNode, depth + 1);
-	}
-}
+import { buildPathTree, isUrlLikePath, type PathTreeInput, walkPathTree } from "@oh-my-pi/pi-utils";
 
 // =============================================================================
 // Grouped file output (grep / ast-grep / ast-edit / lsp diagnostics)

package/src/tools/image-gen.ts

@@ -472,8 +472,13 @@ function parseAntigravityCredentials(raw: string): ParsedAntigravityCredentials
 	return null;
 }
 
-async function findAntigravityCredentials(modelRegistry: ModelRegistry): Promise<ImageApiKey | null> {
-	const apiKey = await modelRegistry.getApiKeyForProvider("google-antigravity");
+async function findAntigravityCredentials(
+	modelRegistry: ModelRegistry,
+	sessionId?: string,
+): Promise<ImageApiKey | null> {
+	const apiKey = await modelRegistry.getApiKeyForProvider("google-antigravity", sessionId, {
+		modelId: DEFAULT_ANTIGRAVITY_MODEL,
+	});
 	if (!apiKey) return null;
 
 	const parsed = parseAntigravityCredentials(apiKey);
@@ -523,7 +528,7 @@ async function findImageApiKey(
 		if (openAI) return openAI;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "antigravity" && modelRegistry) {
-		const antigravity = await findAntigravityCredentials(modelRegistry);
+		const antigravity = await findAntigravityCredentials(modelRegistry, sessionId);
 		if (antigravity) return antigravity;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "gemini") {
@@ -547,7 +552,7 @@ async function findImageApiKey(
 	if (openAI) return openAI;
 
 	if (modelRegistry) {
-		const antigravity = await findAntigravityCredentials(modelRegistry);
+		const antigravity = await findAntigravityCredentials(modelRegistry, sessionId);
 		if (antigravity) return antigravity;
 	}
 
@@ -1052,6 +1057,7 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 				const hostedKey: ApiKey = ctx.modelRegistry.resolver(hostedModel.provider, {
 					sessionId,
 					baseUrl: hostedModel.baseUrl,
+					modelId: hostedModel.id,
 				});
 
 				const parsed = await withAuth(
@@ -1113,6 +1119,7 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 				const prompt = assemblePrompt(params);
 				const antigravityKey: ApiKey = ctx.modelRegistry.resolver("google-antigravity", {
 					sessionId,
+					modelId: DEFAULT_ANTIGRAVITY_MODEL,
 				});
 
 				const response = await withAuth(