@oh-my-pi/pi-coding-agent 15.10.0 → 15.10.1

package/src/tui/status-line.ts

@@ -7,6 +7,9 @@ import { formatStatusIcon } from "../tools/render-utils";
 
 export interface StatusLineOptions {
 	icon?: ToolUIStatus;
+	/** Pre-rendered glyph that replaces the status icon (e.g. a magnifier for
+	 * search-family tools). Takes precedence over `icon`. */
+	iconOverride?: string;
 	spinnerFrame?: number;
 	title: string;
 	titleColor?: ThemeColor;
@@ -27,7 +30,8 @@ function flattenForHeader(text: string): string {
 }
 
 export function renderStatusLine(options: StatusLineOptions, theme: Theme): string {
-	const icon = options.icon ? formatStatusIcon(options.icon, theme, options.spinnerFrame) : "";
+	const icon =
+		options.iconOverride ?? (options.icon ? formatStatusIcon(options.icon, theme, options.spinnerFrame) : "");
 	const titleColor = options.titleColor ?? "accent";
 	const title = theme.fg(titleColor, flattenForHeader(options.title));
 	let line = icon ? `${icon} ${title}` : title;

package/src/utils/commit-message-generator.ts

@@ -6,6 +6,7 @@ import type { ThinkingLevel } from "@oh-my-pi/pi-agent-core";
 import type { Api, Model } from "@oh-my-pi/pi-ai";
 import { completeSimple } from "@oh-my-pi/pi-ai";
 import { logger, prompt } from "@oh-my-pi/pi-utils";
+
 import type { ModelRegistry } from "../config/model-registry";
 import { resolveModelRoleValue } from "../config/model-resolver";
 import type { Settings } from "../config/settings";
@@ -110,7 +111,14 @@ export async function generateCommitMessage(
 					systemPrompt: [COMMIT_SYSTEM_PROMPT],
 					messages: [{ role: "user", content: userMessage, timestamp: Date.now() }],
 				},
-				{ apiKey, maxTokens, reasoning: toReasoningEffort(candidate.thinkingLevel) },
+				{
+					apiKey: registry.resolver(candidate.model.provider, {
+						sessionId,
+						baseUrl: candidate.model.baseUrl,
+					}),
+					maxTokens,
+					reasoning: toReasoningEffort(candidate.thinkingLevel),
+				},
 			);
 
 			if (response.stopReason === "error") {

package/src/utils/enhanced-paste.ts

@@ -0,0 +1,202 @@
+import type { ImageContent } from "@oh-my-pi/pi-ai";
+
+const OSC5522_PREFIX = "\x1b]5522;";
+const OSC_TERMINATOR_ST = "\x1b\\";
+const OSC_TERMINATOR_BEL = "\x07";
+const PASTE_EVENT_NAME_BASE64 = Buffer.from("Paste event", "utf8").toString("base64");
+
+const IMAGE_MIME_PRIORITY = ["image/png", "image/jpeg", "image/webp", "image/gif"] as const;
+const TEXT_MIME_TYPE = "text/plain";
+
+type PasteReadKind = "image" | "text";
+
+export interface Osc5522Packet {
+	metadata: Map<string, string>;
+	payload: string;
+}
+
+interface PasteListingState {
+	phase: "listing";
+	mimes: string[];
+	pw?: string;
+	loc?: string;
+}
+
+interface PasteReadState {
+	phase: "reading";
+	kind: PasteReadKind;
+	mimeType: string;
+	chunks: string[];
+}
+
+type PasteState = PasteListingState | PasteReadState;
+
+export interface EnhancedPasteHandlers {
+	write(data: string): void;
+	pasteText(text: string): void;
+	pasteImage(image: ImageContent): void | Promise<void>;
+	showStatus(message: string): void;
+}
+
+export function isOsc5522Packet(data: string): boolean {
+	return data.startsWith(OSC5522_PREFIX) && (data.endsWith(OSC_TERMINATOR_ST) || data.endsWith(OSC_TERMINATOR_BEL));
+}
+
+function decodeBase64Utf8(value: string): string | undefined {
+	try {
+		return Buffer.from(value, "base64").toString("utf8");
+	} catch {
+		return undefined;
+	}
+}
+
+function parseMetadata(raw: string): Map<string, string> {
+	const metadata = new Map<string, string>();
+	for (const part of raw.split(":")) {
+		const eq = part.indexOf("=");
+		if (eq <= 0) continue;
+		metadata.set(part.slice(0, eq), part.slice(eq + 1));
+	}
+	return metadata;
+}
+
+export function parseOsc5522Packet(data: string): Osc5522Packet | undefined {
+	if (!isOsc5522Packet(data)) return undefined;
+	const bodyEnd = data.endsWith(OSC_TERMINATOR_BEL) ? data.length - 1 : data.length - OSC_TERMINATOR_ST.length;
+	const body = data.slice(OSC5522_PREFIX.length, bodyEnd);
+	const separator = body.indexOf(";");
+	const metadataRaw = separator === -1 ? body : body.slice(0, separator);
+	const payload = separator === -1 ? "" : body.slice(separator + 1);
+	return { metadata: parseMetadata(metadataRaw), payload };
+}
+
+function choosePasteMime(mimes: readonly string[]): { kind: PasteReadKind; mimeType: string } | undefined {
+	for (const mimeType of IMAGE_MIME_PRIORITY) {
+		if (mimes.includes(mimeType)) return { kind: "image", mimeType };
+	}
+	return mimes.includes(TEXT_MIME_TYPE) ? { kind: "text", mimeType: TEXT_MIME_TYPE } : undefined;
+}
+
+export class EnhancedPasteController {
+	#state: PasteState | undefined;
+	#handlers: EnhancedPasteHandlers;
+
+	constructor(handlers: EnhancedPasteHandlers) {
+		this.#handlers = handlers;
+	}
+
+	enable(): void {
+		this.#handlers.write("\x1b[?5522h");
+	}
+
+	disable(): void {
+		this.#handlers.write("\x1b[?5522l");
+		this.#state = undefined;
+	}
+
+	handleInput(data: string): boolean {
+		const packet = parseOsc5522Packet(data);
+		if (!packet) return false;
+		void this.#handlePacket(packet);
+		return true;
+	}
+
+	async #handlePacket(packet: Osc5522Packet): Promise<void> {
+		const type = packet.metadata.get("type");
+		if (type !== "read") return;
+
+		const status = packet.metadata.get("status");
+		if (status === "OK") {
+			this.#handleOk(packet);
+			return;
+		}
+		if (status === "DATA") {
+			this.#handleData(packet);
+			return;
+		}
+		if (status === "DONE") {
+			await this.#handleDone();
+			return;
+		}
+		if (status) {
+			this.#state = undefined;
+			this.#handlers.showStatus(`Enhanced paste failed: ${status}`);
+		}
+	}
+
+	#handleOk(packet: Osc5522Packet): void {
+		if (this.#state?.phase === "reading") return;
+		const loc = packet.metadata.get("loc");
+		this.#state = {
+			phase: "listing",
+			mimes: [],
+			pw: packet.metadata.get("pw"),
+			loc: loc === "primary" ? loc : undefined,
+		};
+	}
+
+	#handleData(packet: Osc5522Packet): void {
+		const state = this.#state;
+		if (!state) return;
+		const encodedMime = packet.metadata.get("mime");
+		if (!encodedMime) return;
+		const mimeType = decodeBase64Utf8(encodedMime);
+		if (!mimeType) return;
+
+		if (state.phase === "listing") {
+			state.mimes.push(mimeType);
+			return;
+		}
+
+		if (state.mimeType === mimeType && packet.payload) {
+			state.chunks.push(packet.payload);
+		}
+	}
+
+	async #handleDone(): Promise<void> {
+		const state = this.#state;
+		if (!state) return;
+		if (state.phase === "listing") {
+			this.#finishListing(state);
+			return;
+		}
+		this.#state = undefined;
+		const bytes = Buffer.concat(state.chunks.map(chunk => Buffer.from(chunk, "base64")));
+		if (bytes.byteLength === 0) {
+			this.#handlers.showStatus("Clipboard paste was empty");
+			return;
+		}
+		if (state.kind === "text") {
+			this.#handlers.pasteText(bytes.toString("utf8"));
+			return;
+		}
+		await this.#handlers.pasteImage({
+			type: "image",
+			data: bytes.toString("base64"),
+			mimeType: state.mimeType,
+		});
+	}
+
+	#finishListing(state: PasteListingState): void {
+		const selected = choosePasteMime(state.mimes);
+		if (!selected) {
+			this.#state = undefined;
+			this.#handlers.showStatus("Clipboard paste has no supported text or image data");
+			return;
+		}
+
+		this.#state = {
+			phase: "reading",
+			kind: selected.kind,
+			mimeType: selected.mimeType,
+			chunks: [],
+		};
+
+		const metadata = [`type=read`, `mime=${Buffer.from(selected.mimeType, "utf8").toString("base64")}`];
+		if (state.loc) metadata.push(`loc=${state.loc}`);
+		if (state.pw) {
+			metadata.push(`pw=${state.pw}`, `name=${PASTE_EVENT_NAME_BASE64}`);
+		}
+		this.#handlers.write(`${OSC5522_PREFIX}${metadata.join(":")}${OSC_TERMINATOR_ST}`);
+	}
+}

package/src/utils/title-generator.ts

@@ -6,6 +6,7 @@ import * as path from "node:path";
 import { type Api, type AssistantMessage, completeSimple, type Model, type Tool } from "@oh-my-pi/pi-ai";
 import { logger, prompt } from "@oh-my-pi/pi-utils";
 import type { ModelRegistry } from "../config/model-registry";
+
 import { resolveRoleSelection } from "../config/model-resolver";
 import type { Settings } from "../config/settings";
 import titleSystemPrompt from "../prompts/system/title-system.md" with { type: "text" };
@@ -238,7 +239,7 @@ export async function generateTitleOnline(
 				tools: [setTitleTool],
 			},
 			{
-				apiKey,
+				apiKey: registry.resolver(model.provider, { sessionId, baseUrl: model.baseUrl }),
 				maxTokens,
 				disableReasoning: true,
 				toolChoice: { type: "tool", name: SET_TITLE_TOOL_NAME },

package/src/web/search/providers/anthropic.ts

@@ -7,12 +7,14 @@
 import {
 	type AnthropicAuthConfig,
 	type AnthropicSystemBlock,
+	type ApiKey,
 	type AuthStorage,
 	buildAnthropicAuthConfig,
 	buildAnthropicSearchHeaders,
 	buildAnthropicSystemBlocks,
 	buildAnthropicUrl,
 	stripClaudeToolPrefix,
+	withAuth,
 } from "@oh-my-pi/pi-ai";
 import { $env } from "@oh-my-pi/pi-utils";
 import type {
@@ -247,18 +249,13 @@ export async function searchAnthropic(
 ): Promise<SearchResponse> {
 	const searchApiKey = $env.ANTHROPIC_SEARCH_API_KEY;
 	const searchBaseUrl = $env.ANTHROPIC_SEARCH_BASE_URL;
-	let auth: AnthropicAuthConfig | undefined;
+	const keyOrResolver: ApiKey | undefined = searchApiKey
+		? searchApiKey
+		: "authStorage" in params
+			? params.authStorage.resolver("anthropic", { sessionId: params.sessionId })
+			: undefined;
 
-	if (searchApiKey) {
-		auth = buildAnthropicAuthConfig(searchApiKey, searchBaseUrl);
-	} else if ("authStorage" in params) {
-		const apiKey = await params.authStorage.getApiKey("anthropic", params.sessionId, {
-			signal: params.signal,
-		});
-		if (apiKey) auth = buildAnthropicAuthConfig(apiKey, searchBaseUrl);
-	}
-
-	if (!auth) {
+	if (!keyOrResolver) {
 		throw new Error(
 			"No Anthropic credentials found. Set ANTHROPIC_SEARCH_API_KEY or ANTHROPIC_API_KEY, or configure Anthropic OAuth.",
 		);
@@ -267,14 +264,23 @@ export async function searchAnthropic(
 	const model = getModel();
 	const systemPrompt = "authStorage" in params ? params.systemPrompt : params.system_prompt;
 	const maxTokens = "authStorage" in params ? params.maxOutputTokens : params.max_tokens;
-	const response = await callSearch(
-		auth,
-		model,
-		params.query,
-		systemPrompt,
-		maxTokens,
-		params.temperature,
-		params.signal,
+	const response = await withAuth(
+		keyOrResolver,
+		key =>
+			callSearch(
+				buildAnthropicAuthConfig(key, searchBaseUrl),
+				model,
+				params.query,
+				systemPrompt,
+				maxTokens,
+				params.temperature,
+				params.signal,
+			),
+		{
+			signal: params.signal,
+			missingKeyMessage:
+				"No Anthropic credentials found. Set ANTHROPIC_SEARCH_API_KEY or ANTHROPIC_API_KEY, or configure Anthropic OAuth.",
+		},
 	);
 
 	const result = parseResponse(response);

package/src/web/search/providers/exa.ts

@@ -6,7 +6,7 @@
  * Requests per-result summaries via `contents.summary` and synthesizes
  * them into a combined `answer` string on the SearchResponse.
  */
-import { type AuthStorage, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, getEnvApiKey, withAuth } from "@oh-my-pi/pi-ai";
 import { settings } from "../../../config/settings";
 import { callExaTool, findApiKey, isSearchResponse } from "../../../exa/mcp-client";
 
@@ -228,11 +228,19 @@ async function callExaMcpSearch(params: ExaSearchParams): Promise<ExaSearchRespo
 
 /** Execute Exa web search */
 export async function searchExa(params: ExaSearchParams): Promise<SearchResponse> {
+	// AuthStorage-backed key takes precedence (existing behavior); probe it once
+	// so the env-key and keyless-MCP fallbacks below stay intact, then drive the
+	// authStorage path through the central force-refresh/rotate retry policy.
 	const storedKey = params.authStorage
 		? await params.authStorage.getApiKey("exa", params.sessionId, { signal: params.signal })
 		: undefined;
-	const apiKey = storedKey ?? getEnvApiKey("exa");
-	const response = apiKey ? await callExaSearch(apiKey, params) : await callExaMcpSearch(params);
+	const keyOrResolver: ApiKey | undefined =
+		storedKey && params.authStorage
+			? params.authStorage.resolver("exa", { sessionId: params.sessionId })
+			: getEnvApiKey("exa");
+	const response = keyOrResolver
+		? await withAuth(keyOrResolver, key => callExaSearch(key, params), { signal: params.signal })
+		: await callExaMcpSearch(params);
 
 	// Convert to unified SearchResponse
 	const sources: SearchSource[] = [];

package/src/web/search/providers/kimi.ts

@@ -4,7 +4,7 @@
  * Uses Moonshot Kimi Code search API to retrieve web results.
  * Endpoint: POST https://api.kimi.com/coding/v1/search
  */
-import type { AuthStorage } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, withAuth } from "@oh-my-pi/pi-ai";
 import { $env } from "@oh-my-pi/pi-utils";
 
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
@@ -54,20 +54,26 @@ function resolveBaseUrl(): string {
 	return asTrimmed($env.MOONSHOT_SEARCH_BASE_URL) ?? asTrimmed($env.KIMI_SEARCH_BASE_URL) ?? KIMI_SEARCH_URL;
 }
 
-/** Find Kimi search credentials from environment or AuthStorage. */
-async function findApiKey(
+/**
+ * Resolve the Kimi search credential. Highest precedence is the static env key;
+ * otherwise an AuthStorage-backed resolver for whichever stored provider id
+ * holds a key (`moonshot` first, then `kimi-code`), so a stale token triggers
+ * the central force-refresh / sibling-rotate retry. Returns `undefined` when
+ * neither is configured.
+ */
+async function resolveKey(
 	authStorage: AuthStorage,
 	sessionId: string | undefined,
 	signal: AbortSignal | undefined,
-): Promise<string | null> {
+): Promise<ApiKey | undefined> {
 	const envKey = asTrimmed($env.MOONSHOT_SEARCH_API_KEY) ?? asTrimmed($env.KIMI_SEARCH_API_KEY);
 	if (envKey) return envKey;
 
-	return (
-		(await authStorage.getApiKey("moonshot", sessionId, { signal })) ??
-		(await authStorage.getApiKey("kimi-code", sessionId, { signal })) ??
-		null
-	);
+	for (const provider of ["moonshot", "kimi-code"] as const) {
+		const stored = await authStorage.getApiKey(provider, sessionId, { signal });
+		if (stored) return authStorage.resolver(provider, { sessionId });
+	}
+	return undefined;
 }
 
 async function callKimiSearch(
@@ -108,20 +114,25 @@ async function callKimiSearch(
 
 /** Execute Kimi web search. */
 export async function searchKimi(params: KimiSearchParams): Promise<SearchResponse> {
-	const apiKey = await findApiKey(params.authStorage, params.sessionId, params.signal);
-	if (!apiKey) {
+	const keyOrResolver = await resolveKey(params.authStorage, params.sessionId, params.signal);
+	if (!keyOrResolver) {
 		throw new Error(
 			"Kimi search credentials not found. Set MOONSHOT_SEARCH_API_KEY, KIMI_SEARCH_API_KEY, MOONSHOT_API_KEY, or login with 'omp /login moonshot'.",
 		);
 	}
 
 	const limit = clampNumResults(params.num_results, DEFAULT_NUM_RESULTS, MAX_NUM_RESULTS);
-	const { response, requestId } = await callKimiSearch(apiKey, {
-		query: params.query,
-		limit,
-		includeContent: params.include_content ?? false,
-		signal: params.signal,
-	});
+	const { response, requestId } = await withAuth(
+		keyOrResolver,
+		key =>
+			callKimiSearch(key, {
+				query: params.query,
+				limit,
+				includeContent: params.include_content ?? false,
+				signal: params.signal,
+			}),
+		{ signal: params.signal },
+	);
 	const sources: SearchSource[] = [];
 
 	for (const result of response.search_results ?? []) {

package/src/web/search/providers/parallel.ts

@@ -1,4 +1,4 @@
-import { type AuthStorage, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, getEnvApiKey, withAuth } from "@oh-my-pi/pi-ai";
 import type { SearchResponse } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import { ParallelApiError, type ParallelSearchResult, type ParallelSearchSource } from "../../parallel";
@@ -123,30 +123,41 @@ async function searchWithAuthStorage(
 		);
 	}
 
-	const response = await fetch(PARALLEL_SEARCH_URL, {
-		method: "POST",
-		headers: {
-			Accept: "application/json",
-			"Content-Type": "application/json",
-			"x-api-key": apiKey,
-			"parallel-beta": PARALLEL_BETA_HEADER,
-		},
-		body: JSON.stringify({
-			objective,
-			search_queries: queries,
-			mode: "fast",
-			excerpts: {
-				max_chars_per_result: 10_000,
-			},
-		}),
-		signal: withHardTimeout(params.signal),
-	});
-	if (!response.ok) {
-		throw parseParallelErrorResponse(response.status, await response.text());
-	}
+	// Drive the (already-present) credential through the central force-refresh /
+	// sibling-rotate retry policy. The `ParallelApiError` thrown below carries a
+	// `statusCode`, which `withAuth`'s default classifier reads to detect a
+	// retryable 401 / usage-limit.
+	const keyOrResolver: ApiKey = authStorage.resolver("parallel", { sessionId });
+	return withAuth(
+		keyOrResolver,
+		async key => {
+			const response = await fetch(PARALLEL_SEARCH_URL, {
+				method: "POST",
+				headers: {
+					Accept: "application/json",
+					"Content-Type": "application/json",
+					"x-api-key": key,
+					"parallel-beta": PARALLEL_BETA_HEADER,
+				},
+				body: JSON.stringify({
+					objective,
+					search_queries: queries,
+					mode: "fast",
+					excerpts: {
+						max_chars_per_result: 10_000,
+					},
+				}),
+				signal: withHardTimeout(params.signal),
+			});
+			if (!response.ok) {
+				throw parseParallelErrorResponse(response.status, await response.text());
+			}
 
-	const payload: unknown = await response.json();
-	return parseSearchPayload(payload);
+			const payload: unknown = await response.json();
+			return parseSearchPayload(payload);
+		},
+		{ signal: params.signal },
+	);
 }
 
 export async function searchParallel(

package/src/web/search/providers/synthetic.ts

@@ -5,7 +5,7 @@
  * Endpoint: POST https://api.synthetic.new/v2/search
  */
 
-import { type AuthStorage, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, getEnvApiKey, withAuth } from "@oh-my-pi/pi-ai";
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import type { SearchParams } from "./base";
@@ -66,12 +66,14 @@ async function callSyntheticSearch(
 
 /** Execute Synthetic web search. */
 export async function searchSynthetic(params: SearchParams): Promise<SearchResponse> {
-	const apiKey = await findApiKey(params.authStorage, params.sessionId, params.signal);
-	if (!apiKey) {
-		throw new Error("Synthetic credentials not found. Set SYNTHETIC_API_KEY or login with 'omp /login synthetic'.");
-	}
+	const keyOrResolver: ApiKey = params.authStorage.resolver("synthetic", {
+		sessionId: params.sessionId,
+	});
 
-	const data = await callSyntheticSearch(apiKey, params.query, params.signal);
+	const data = await withAuth(keyOrResolver, key => callSyntheticSearch(key, params.query, params.signal), {
+		signal: params.signal,
+		missingKeyMessage: "Synthetic credentials not found. Set SYNTHETIC_API_KEY or login with 'omp /login synthetic'.",
+	});
 	const sources: SearchSource[] = [];
 
 	for (const result of data.results ?? []) {

package/src/web/search/providers/tavily.ts

@@ -4,7 +4,7 @@
  * Uses Tavily's agent-focused search API to return structured results with an
  * optional synthesized answer.
  */
-import { type AuthStorage, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, getEnvApiKey, withAuth } from "@oh-my-pi/pi-ai";
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import { clampNumResults, dateToAgeSeconds } from "../utils";
@@ -127,15 +127,16 @@ export async function searchTavily(params: SearchParams): Promise<SearchResponse
 		recency: params.recency,
 		signal: params.signal,
 	};
-	const apiKey = await findApiKey(params.authStorage, params.sessionId, params.signal);
-	if (!apiKey) {
-		throw new Error(
-			'Tavily credentials not found. Set TAVILY_API_KEY or configure an API key for provider "tavily".',
-		);
-	}
+	const keyOrResolver: ApiKey = params.authStorage.resolver("tavily", {
+		sessionId: params.sessionId,
+	});
 
 	const numResults = clampNumResults(tavilyParams.num_results, DEFAULT_NUM_RESULTS, MAX_NUM_RESULTS);
-	const response = await callTavilySearch(apiKey, tavilyParams);
+	const response = await withAuth(keyOrResolver, key => callTavilySearch(key, tavilyParams), {
+		signal: params.signal,
+		missingKeyMessage:
+			'Tavily credentials not found. Set TAVILY_API_KEY or configure an API key for provider "tavily".',
+	});
 	const sources: SearchSource[] = [];
 
 	for (const result of response.results ?? []) {

package/src/web/search/providers/zai.ts

@@ -4,7 +4,7 @@
  * Calls Z.AI's remote MCP server (`webSearchPrime`) and adapts results into
  * the unified SearchResponse shape used by the web search tool.
  */
-import { type AuthStorage, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type ApiKey, type AuthStorage, getEnvApiKey, withAuth } from "@oh-my-pi/pi-ai";
 import { asRecord, asString } from "../../../web/scrapers/utils";
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
@@ -278,12 +278,14 @@ function toSources(results: ZaiSearchResult[]): SearchSource[] {
 
 /** Execute Z.AI web search via remote MCP endpoint. */
 export async function searchZai(params: ZaiSearchParams): Promise<SearchResponse> {
-	const apiKey = await findApiKey(params.authStorage, params.sessionId, params.signal);
-	if (!apiKey) {
-		throw new Error("Z.AI credentials not found. Set ZAI_API_KEY or login with 'omp /login zai'.");
-	}
+	const keyOrResolver: ApiKey = params.authStorage.resolver("zai", {
+		sessionId: params.sessionId,
+	});
 
-	const rawResult = await callZaiSearch(apiKey, params);
+	const rawResult = await withAuth(keyOrResolver, key => callZaiSearch(key, params), {
+		signal: params.signal,
+		missingKeyMessage: "Z.AI credentials not found. Set ZAI_API_KEY or login with 'omp /login zai'.",
+	});
 	const payload = parseSearchPayload(rawResult);
 	let sources = toSources(payload.results);