@oh-my-pi/pi-coding-agent 15.3.2 → 15.4.1

package/src/tools/search.ts

@@ -9,6 +9,7 @@ import { prompt, untilAborted } from "@oh-my-pi/pi-utils";
 import * as z from "zod/v4";
 import { getFileReadCache } from "../edit/file-read-cache";
 import type { RenderResultOptions } from "../extensibility/custom-tools/types";
+import { computeFileHash, formatHashlineHeader } from "../hashline/hash";
 import type { Theme } from "../modes/theme/theme";
 import searchDescription from "../prompts/tools/search.md" with { type: "text" };
 import { DEFAULT_MAX_COLUMN, type TruncationResult, truncateHead } from "../session/streaming-output";
@@ -38,13 +39,13 @@ import {
 import { ToolError } from "./tool-errors";
 import { toolResult } from "./tool-result";
 
+const searchPathEntrySchema = z.string().describe("file, directory, glob, or internal URL to search");
 const searchSchema = z
 	.object({
 		pattern: z.string().describe("regex pattern"),
 		paths: z
-			.array(z.string().describe("file, directory, glob, or internal URL to search"))
-			.min(1)
-			.describe("files, directories, globs, or internal URLs to search"),
+			.union([searchPathEntrySchema, z.array(searchPathEntrySchema).min(1)])
+			.describe("file, directory, glob, internal URL, or array of those to search"),
 		i: z.boolean().optional().describe("case-insensitive search"),
 		gitignore: z.boolean().optional().describe("respect gitignore"),
 		skip: z
@@ -55,6 +56,9 @@ const searchSchema = z
 	.strict();
 
 export type SearchToolInput = z.infer<typeof searchSchema>;
+export function toPathList(input: string | string[] | undefined): string[] {
+	return typeof input === "string" ? [input] : (input ?? []);
+}
 
 /** Maximum number of distinct files surfaced in a single response. The
  * agent paginates further pages via `skip`. */
@@ -236,7 +240,7 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 		_onUpdate?: AgentToolUpdateCallback<SearchToolDetails>,
 		_toolContext?: AgentToolContext,
 	): Promise<AgentToolResult<SearchToolDetails>> {
-		const { pattern, paths, i, gitignore, skip } = params;
+		const { pattern, paths: rawPaths, i, gitignore, skip } = params;
 
 		return untilAborted(signal, async () => {
 			const normalizedPattern = pattern.trim();
@@ -248,6 +252,7 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 			if (normalizedSkip < 0 || !Number.isFinite(normalizedSkip)) {
 				throw new ToolError("Skip must be a non-negative number");
 			}
+			const paths = toPathList(rawPaths);
 			for (const entry of paths) {
 				if (containsTopLevelComma(entry)) {
 					throw new ToolError('paths is an array — pass ["a", "b"] not ["a,b"]');
@@ -303,7 +308,6 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 				}
 				const { globFilter } = scope;
 				const baseDisplayMode = resolveFileDisplayMode(this.session);
-				const immutableDisplayMode = resolveFileDisplayMode(this.session, { immutable: true });
 
 				const effectiveOutputMode = GrepOutputMode.Content;
 				// Multi-scope = more than one file may match. We fetch up to
@@ -485,14 +489,27 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 					matchesByFile.get(relativePath)!.push(match);
 				}
 				const displayLines: string[] = [];
+				const hashContexts = new Map<string, { absolutePath: string; fileHash: string }>();
+				if (baseDisplayMode.hashLines) {
+					for (const relativePath of fileList) {
+						if (archiveDisplaySet.has(relativePath)) continue;
+						const absoluteFilePath = path.resolve(this.session.cwd, relativePath);
+						if (immutableSourcePaths.has(absoluteFilePath)) continue;
+						try {
+							const fullText = await Bun.file(absoluteFilePath).text();
+							const fileHash = computeFileHash(fullText);
+							hashContexts.set(relativePath, { absolutePath: absoluteFilePath, fileHash });
+						} catch {
+							// Best-effort: if the file disappeared between grep and render, fall back to plain line output.
+						}
+					}
+				}
 				const renderMatchesForFile = (relativePath: string): { model: string[]; display: string[] } => {
 					const modelOut: string[] = [];
 					const displayOut: string[] = [];
 					const fileMatches = matchesByFile.get(relativePath) ?? [];
-					const absoluteFilePath = path.resolve(this.session.cwd, relativePath);
-					const useHashLines = immutableSourcePaths.has(absoluteFilePath)
-						? immutableDisplayMode.hashLines
-						: baseDisplayMode.hashLines;
+					const hashContext = hashContexts.get(relativePath);
+					const useHashLines = hashContext !== undefined;
 					const lineNumberWidth = fileMatches.reduce((width, match) => {
 						let nextWidth = Math.max(width, String(match.lineNumber).length);
 						for (const ctx of match.contextBefore ?? []) {
@@ -533,17 +550,21 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 						}
 						fileMatchCounts.set(relativePath, (fileMatchCounts.get(relativePath) ?? 0) + 1);
 					}
-					if (cacheEntries.length > 0 && !archiveDisplaySet.has(relativePath)) {
-						getFileReadCache(this.session).recordSparse(path.resolve(searchPath, relativePath), cacheEntries);
+					if (cacheEntries.length > 0 && hashContext) {
+						getFileReadCache(this.session).recordSparse(hashContext.absolutePath, cacheEntries, {
+							fileHash: hashContext.fileHash,
+						});
 					}
 					return { model: modelOut, display: displayOut };
 				};
 				if (isDirectory) {
 					const grouped = formatGroupedFiles(fileList, relativePath => {
 						const rendered = renderMatchesForFile(relativePath);
+						const hashContext = hashContexts.get(relativePath);
 						return {
 							modelLines: rendered.model,
 							displayLines: rendered.display,
+							headerSuffix: hashContext ? `#${hashContext.fileHash}` : "",
 							skip: rendered.model.length === 0,
 						};
 					});
@@ -552,6 +573,15 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 				} else {
 					for (const relativePath of fileList) {
 						const rendered = renderMatchesForFile(relativePath);
+						if (rendered.model.length === 0) continue;
+						if (outputLines.length > 0) {
+							outputLines.push("");
+							displayLines.push("");
+						}
+						const hashContext = hashContexts.get(relativePath);
+						if (hashContext) {
+							outputLines.push(formatHashlineHeader(relativePath, hashContext.fileHash));
+						}
 						outputLines.push(...rendered.model);
 						displayLines.push(...rendered.display);
 					}
@@ -607,7 +637,7 @@ export class SearchTool implements AgentTool<typeof searchSchema, SearchToolDeta
 
 interface SearchRenderArgs {
 	pattern: string;
-	paths?: string[];
+	paths?: string | string[];
 	i?: boolean;
 	gitignore?: boolean;
 	skip?: number;
@@ -618,8 +648,9 @@ const COLLAPSED_TEXT_LIMIT = PREVIEW_LIMITS.COLLAPSED_LINES * 2;
 export const searchToolRenderer = {
 	inline: true,
 	renderCall(args: SearchRenderArgs, _options: RenderResultOptions, uiTheme: Theme): Component {
+		const paths = toPathList(args.paths);
 		const meta: string[] = [];
-		if (args.paths?.length) meta.push(`in ${args.paths.join(", ")}`);
+		if (paths.length) meta.push(`in ${paths.join(", ")}`);
 		if (args.i) meta.push("case:insensitive");
 		if (args.gitignore === false) meta.push("gitignore:false");
 		if (args.skip !== undefined && args.skip > 0) meta.push(`skip:${args.skip}`);
@@ -745,11 +776,12 @@ export const searchToolRenderer = {
 							let contextDir = searchBase ?? "";
 							return group.map(line => {
 								if (line.startsWith("## ")) {
-									// Strip optional ` (suffix)` like ` (3 replacements)` before resolving.
+									// Strip optional ` (suffix)` and `#hash` before resolving.
 									const fileName = line
 										.slice(3)
 										.trimEnd()
-										.replace(/\s+\([^)]*\)\s*$/, "");
+										.replace(/\s+\([^)]*\)\s*$/, "")
+										.replace(/#[0-9a-f]+$/, "");
 									const absPath = contextDir && fileName ? path.join(contextDir, fileName) : undefined;
 									const styled = uiTheme.fg("dim", line);
 									return absPath ? fileHyperlink(absPath, styled) : styled;
@@ -760,7 +792,7 @@ export const searchToolRenderer = {
 										.trimEnd()
 										.replace(/\s+\([^)]*\)\s*$/, "");
 									const isDirectory = raw.endsWith("/");
-									const name = raw.replace(/\/$/, "");
+									const name = isDirectory ? raw.replace(/\/$/, "") : raw.replace(/#[0-9a-f]+$/, "");
 									if (isDirectory) {
 										if (searchBase) {
 											contextDir = name === "." ? searchBase : path.join(searchBase, name);

package/src/tools/write.ts

@@ -74,8 +74,8 @@ export interface WriteToolDetails {
 /**
  * Strip hashline display prefixes from write content.
  *
- * Only active when hashline edit mode is enabled — the model sees `LINE+ID|`
- * prefixes in read output and sometimes copies them into write content.
+ * Only active when hashline edit mode is enabled — the model sees `¶PATH#HASH`
+ * headers plus `LINE:` prefixes in read output and sometimes copies them into write content.
  */
 function stripWriteContent(session: ToolSession, content: string): { text: string; stripped: boolean } {
 	if (!resolveFileDisplayMode(session).hashLines) {
@@ -658,7 +658,7 @@ export class WriteTool implements AgentTool<typeof writeSchema, WriteToolDetails
 		context?: AgentToolContext,
 	): Promise<AgentToolResult<WriteToolDetails>> {
 		return untilAborted(signal, async () => {
-			// Strip hashline display prefixes (LINE+ID|) if the model copied them from read output
+			// Strip hashline display prefixes (¶PATH#HASH + LINE:) if the model copied them from read output
 			const { text: cleanContent, stripped } = stripWriteContent(this.session, content);
 			const internalRouter = InternalUrlRouter.instance();
 			if (internalRouter.canHandle(path)) {

package/src/tools/yield.ts

@@ -8,15 +8,13 @@ import type { TSchema } from "@oh-my-pi/pi-ai/types";
 import {
 	dereferenceJsonSchema,
 	isValidJsonSchema,
-	type JsonSchemaValidationIssue,
 	type JsonSchemaValidationResult,
 	sanitizeSchemaForStrictMode,
 	tryEnforceStrictSchema,
-	validateJsonSchemaValue,
 } from "@oh-my-pi/pi-ai/utils/schema";
 import { subprocessToolRegistry } from "../task/subprocess-tool-registry";
 import type { ToolSession } from ".";
-import { jtdToJsonSchema, normalizeSchema } from "./jtd-to-json-schema";
+import { buildOutputValidator, formatAllValidationIssues } from "./output-schema-validator";
 
 export interface YieldDetails {
 	data: unknown;
@@ -34,16 +32,6 @@ function formatSchema(schema: unknown): string {
 	}
 }
 
-function formatJsonSchemaIssues(issues: ReadonlyArray<JsonSchemaValidationIssue> | undefined): string {
-	if (!issues || issues.length === 0) return "Unknown schema validation error.";
-	return issues
-		.map(issue => {
-			const path = issue.path.length === 0 ? "" : `${issue.path.map(seg => String(seg)).join("/")}: `;
-			return `${path}${issue.message}`;
-		})
-		.join("; ");
-}
-
 function looseRecordSchema(description: string): Record<string, unknown> {
 	return {
 		type: "object",
@@ -100,6 +88,15 @@ function wrapYieldParameters(dataSchema: Record<string, unknown>): Record<string
 	};
 }
 
+/**
+ * Max consecutive schema-validation failures before the yield tool overrides validation
+ * and lets non-conforming data through. The override is a safety net for schemas the
+ * JTD→JSON-Schema converter cannot fully express; it should not be reached during normal
+ * model retries. Three matches the existing "3 reminders" pattern elsewhere in the agent
+ * runtime.
+ */
+const MAX_SCHEMA_RETRIES = 3;
+
 export class YieldTool implements AgentTool<TSchema, YieldDetails> {
 	readonly name = "yield";
 	readonly label = "Submit Result";
@@ -120,21 +117,14 @@ export class YieldTool implements AgentTool<TSchema, YieldDetails> {
 		let parameters: TSchema;
 
 		try {
-			const schemaResult = normalizeSchema(session.outputSchema);
-			const normalizedSchema =
-				schemaResult.normalized !== undefined ? jtdToJsonSchema(schemaResult.normalized) : undefined;
-			let schemaError = schemaResult.error;
-
-			if (!schemaError && normalizedSchema === false) {
-				schemaError = "boolean false schema rejects all outputs";
-			}
-
-			if (normalizedSchema !== undefined && normalizedSchema !== false && !schemaError) {
-				if (!isValidJsonSchema(normalizedSchema)) {
-					schemaError = "invalid JSON schema";
-				} else {
-					validate = value => validateJsonSchemaValue(normalizedSchema, value);
-				}
+			const {
+				validator,
+				jsonSchema: normalizedSchema,
+				normalized,
+				error: schemaError,
+			} = buildOutputValidator(session.outputSchema);
+			if (validator) {
+				validate = value => validator.validate(value);
 			}
 
 			const schemaHint = formatSchema(normalizedSchema ?? session.outputSchema);
@@ -142,21 +132,15 @@ export class YieldTool implements AgentTool<TSchema, YieldDetails> {
 				? `Structured JSON output (output schema invalid; accepting unconstrained object): ${schemaError}`
 				: `Structured output matching the schema:\n${schemaHint}`;
 			let sanitizedSchema: Record<string, unknown> | undefined;
-			if (
-				!schemaError &&
-				normalizedSchema != null &&
-				typeof normalizedSchema === "object" &&
-				!Array.isArray(normalizedSchema)
-			) {
-				const normalizedRecord = normalizedSchema as Record<string, unknown>;
-				const strictProbe = tryEnforceStrictSchema(normalizedRecord);
+			if (!schemaError && normalizedSchema !== undefined) {
+				const strictProbe = tryEnforceStrictSchema(normalizedSchema);
 				if (strictProbe.strict) {
-					sanitizedSchema = sanitizeSchemaForStrictMode(normalizedRecord);
+					sanitizedSchema = sanitizeSchemaForStrictMode(normalizedSchema);
 				} else {
-					sanitizedSchema = normalizedRecord;
+					sanitizedSchema = normalizedSchema;
 					this.strict = false;
 				}
-			} else if (!schemaError && normalizedSchema === true) {
+			} else if (!schemaError && normalized === true) {
 				sanitizedSchema = {};
 				this.strict = false;
 			}
@@ -229,8 +213,15 @@ export class YieldTool implements AgentTool<TSchema, YieldDetails> {
 				const parsed = this.#validate(data);
 				if (!parsed.success) {
 					this.#schemaValidationFailures++;
-					if (this.#schemaValidationFailures <= 1) {
-						throw new Error(`Output does not match schema: ${formatJsonSchemaIssues(parsed.issues)}`);
+					if (this.#schemaValidationFailures <= MAX_SCHEMA_RETRIES) {
+						const remaining = MAX_SCHEMA_RETRIES - this.#schemaValidationFailures;
+						const retryHint =
+							remaining > 0
+								? ` Call yield again with the corrected shape — ${remaining} retry attempt(s) remain before the schema constraint is dropped.`
+								: " Call yield again with the corrected shape — this is the final retry before the schema constraint is dropped.";
+						throw new Error(
+							`Output does not match schema: ${formatAllValidationIssues(parsed.issues)}.${retryHint}`,
+						);
 					}
 					schemaValidationOverridden = true;
 				}

package/src/utils/edit-mode.ts

@@ -1,6 +1,6 @@
 import { $env } from "@oh-my-pi/pi-utils";
 
-export type EditMode = "replace" | "patch" | "hashline" | "vim" | "apply_patch";
+export type EditMode = "replace" | "patch" | "hashline" | "apply_patch";
 
 export const DEFAULT_EDIT_MODE: EditMode = "hashline";
 
@@ -9,7 +9,6 @@ const EDIT_MODE_IDS = {
 	hashline: "hashline",
 	patch: "patch",
 	replace: "replace",
-	vim: "vim",
 } as const satisfies Record<string, EditMode>;
 
 export const EDIT_MODES = Object.keys(EDIT_MODE_IDS) as EditMode[];

package/src/utils/file-mentions.ts

@@ -12,7 +12,7 @@ import type { ImageContent } from "@oh-my-pi/pi-ai";
 import { glob } from "@oh-my-pi/pi-natives";
 import { fuzzyMatch } from "@oh-my-pi/pi-tui";
 import { formatAge, formatBytes, readImageMetadata } from "@oh-my-pi/pi-utils";
-import { formatHashLines } from "../hashline/hash";
+import { computeFileHash, formatHashlineHeader, formatNumberedLines } from "../hashline/hash";
 import type { FileMentionMessage } from "../session/messages";
 import {
 	DEFAULT_MAX_BYTES,
@@ -356,7 +356,7 @@ export async function generateFileMentionMessages(
 			const content = await Bun.file(absolutePath).text();
 			let { output, lineCount } = buildTextOutput(content);
 			if (options?.useHashLines) {
-				output = formatHashLines(output);
+				output = `${formatHashlineHeader(resolvedPath, computeFileHash(content))}\n${formatNumberedLines(output)}`;
 			}
 			files.push({ path: resolvedPath, content: output, lineCount });
 		} catch {

package/src/web/kagi.ts

@@ -1,5 +1,5 @@
-import { getEnvApiKey } from "@oh-my-pi/pi-ai";
-import { findCredential, withHardTimeout } from "./search/providers/utils";
+import type { AuthStorage } from "@oh-my-pi/pi-ai";
+import { withHardTimeout } from "./search/providers/utils";
 
 const KAGI_SEARCH_URL = "https://kagi.com/api/v0/search";
 
@@ -97,6 +97,7 @@ function parseKagiErrorResponse(statusCode: number, responseText: string): KagiA
 
 export interface KagiSearchOptions {
 	limit?: number;
+	sessionId?: string;
 	signal?: AbortSignal;
 }
 
@@ -113,8 +114,12 @@ export interface KagiSearchResult {
 	relatedQuestions: string[];
 }
 
-export async function findKagiApiKey(): Promise<string | null> {
-	return findCredential(getEnvApiKey("kagi"), "kagi");
+export async function findKagiApiKey(
+	authStorage: AuthStorage,
+	sessionId?: string,
+	signal?: AbortSignal,
+): Promise<string | null> {
+	return (await authStorage.getApiKey("kagi", sessionId, { signal })) ?? null;
 }
 
 function getAuthHeaders(apiKey: string): Record<string, string> {
@@ -124,8 +129,12 @@ function getAuthHeaders(apiKey: string): Record<string, string> {
 	};
 }
 
-export async function searchWithKagi(query: string, options: KagiSearchOptions = {}): Promise<KagiSearchResult> {
-	const apiKey = await findKagiApiKey();
+export async function searchWithKagi(
+	query: string,
+	options: KagiSearchOptions = {},
+	authStorage: AuthStorage,
+): Promise<KagiSearchResult> {
+	const apiKey = await findKagiApiKey(authStorage, options.sessionId, options.signal);
 	if (!apiKey) {
 		throw new KagiApiError("Kagi credentials not found. Set KAGI_API_KEY or login with 'omp /login kagi'.");
 	}

package/src/web/parallel.ts

@@ -1,4 +1,5 @@
 import { getEnvApiKey } from "@oh-my-pi/pi-ai";
+import type { AgentStorage } from "../session/agent-storage";
 import { findCredential, withHardTimeout } from "./search/providers/utils";
 
 const PARALLEL_API_URL = "https://api.parallel.ai";
@@ -73,8 +74,8 @@ export class ParallelApiError extends Error {
 	}
 }
 
-export async function findParallelApiKey(): Promise<string | null> {
-	return findCredential(getEnvApiKey("parallel"), "parallel");
+export function findParallelApiKey(storage: AgentStorage | null | undefined): string | null {
+	return findCredential(storage, getEnvApiKey("parallel"), "parallel");
 }
 
 export function getParallelExtractContent(document: ParallelExtractDocument): string {
@@ -284,9 +285,10 @@ function parseExtractPayload(payload: unknown): ParallelExtractResult {
 export async function searchWithParallel(
 	objective: string,
 	queries: string[],
-	options: ParallelSearchOptions = {},
+	options: ParallelSearchOptions,
+	storage: AgentStorage | null | undefined,
 ): Promise<ParallelSearchResult> {
-	const apiKey = await findParallelApiKey();
+	const apiKey = findParallelApiKey(storage);
 	if (!apiKey) {
 		throw new ParallelApiError(
 			"Parallel credentials not found. Set PARALLEL_API_KEY or login with 'omp /login parallel'.",
@@ -316,9 +318,10 @@ export async function searchWithParallel(
 
 export async function extractWithParallel(
 	urls: string[],
-	options: ParallelExtractOptions = {},
+	options: ParallelExtractOptions,
+	storage: AgentStorage | null | undefined,
 ): Promise<ParallelExtractResult> {
-	const apiKey = await findParallelApiKey();
+	const apiKey = findParallelApiKey(storage);
 	if (!apiKey) {
 		throw new ParallelApiError(
 			"Parallel credentials not found. Set PARALLEL_API_KEY or login with 'omp /login parallel'.",

package/src/web/scrapers/types.ts

@@ -4,6 +4,7 @@
 import { ptree } from "@oh-my-pi/pi-utils";
 import type TurndownService from "turndown";
 
+import type { AgentStorage } from "../../session/agent-storage";
 import { ToolAbortError } from "../../tools/tool-errors";
 
 export { formatNumber } from "@oh-my-pi/pi-utils";
@@ -19,7 +20,12 @@ export interface RenderResult {
 	notes: string[];
 }
 
-export type SpecialHandler = (url: string, timeout: number, signal?: AbortSignal) => Promise<RenderResult | null>;
+export type SpecialHandler = (
+	url: string,
+	timeout: number,
+	signal?: AbortSignal,
+	storage?: AgentStorage | null,
+) => Promise<RenderResult | null>;
 
 export const MAX_OUTPUT_CHARS = 500_000;
 export const MAX_BYTES = 50 * 1024 * 1024;

package/src/web/scrapers/youtube.ts

@@ -3,6 +3,7 @@ import * as os from "node:os";
 import * as path from "node:path";
 import { ptree, Snowflake } from "@oh-my-pi/pi-utils";
 import { settings } from "../../config/settings";
+import type { AgentStorage } from "../../session/agent-storage";
 import { throwIfAborted } from "../../tools/tool-errors";
 import { ensureTool } from "../../utils/tools-manager";
 import { extractWithParallel, findParallelApiKey, getParallelExtractContent } from "../parallel";
@@ -101,6 +102,7 @@ export const handleYouTube: SpecialHandler = async (
 	url: string,
 	timeout: number,
 	userSignal?: AbortSignal,
+	storage?: AgentStorage | null,
 ): Promise<RenderResult | null> => {
 	throwIfAborted(userSignal);
 	const yt = parseYouTubeUrl(url);
@@ -112,14 +114,18 @@ export const handleYouTube: SpecialHandler = async (
 	const videoUrl = `https://www.youtube.com/watch?v=${yt.videoId}`;
 
 	// Prefer Parallel extract when credentials are available
-	if (settings.get("providers.parallelFetch") && (await findParallelApiKey())) {
+	if (settings.get("providers.parallelFetch") && findParallelApiKey(storage)) {
 		try {
-			const parallelResult = await extractWithParallel([videoUrl], {
-				objective: "Extract the main content of this YouTube video page",
-				excerpts: true,
-				fullContent: false,
-				signal,
-			});
+			const parallelResult = await extractWithParallel(
+				[videoUrl],
+				{
+					objective: "Extract the main content of this YouTube video page",
+					excerpts: true,
+					fullContent: false,
+					signal,
+				},
+				storage,
+			);
 			const firstDocument = parallelResult.results[0];
 			if (firstDocument) {
 				const content = getParallelExtractContent(firstDocument);

package/src/web/search/index.ts

@@ -3,15 +3,16 @@
  *
  * Single tool supporting Anthropic, Perplexity, Exa, Brave, Jina, Kimi, Gemini, Codex, Tavily, Kagi, Z.AI, SearXNG, and Synthetic
  * providers with provider-specific parameters exposed conditionally.
- *
  */
 import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
+import type { AuthStorage } from "@oh-my-pi/pi-ai";
 import { prompt } from "@oh-my-pi/pi-utils";
 import * as z from "zod/v4";
 import type { CustomTool, CustomToolContext, RenderResultOptions } from "../../extensibility/custom-tools/types";
 import type { Theme } from "../../modes/theme/theme";
 import webSearchSystemPrompt from "../../prompts/system/web-search.md" with { type: "text" };
 import webSearchDescription from "../../prompts/tools/web-search.md" with { type: "text" };
+import { discoverAuthStorage } from "../../sdk";
 import type { ToolSession } from "../../tools";
 import { formatAge } from "../../tools/render-utils";
 import { throwIfAborted } from "../../tools/tool-errors";
@@ -114,18 +115,25 @@ function formatForLLM(response: SearchResponse): string {
 	return parts.join("\n");
 }
 
+interface ExecuteSearchOptions {
+	authStorage: AuthStorage;
+	sessionId?: string;
+	signal?: AbortSignal;
+}
+
 /** Execute web search */
 async function executeSearch(
 	_toolCallId: string,
 	params: SearchQueryParams,
-	signal?: AbortSignal,
+	options: ExecuteSearchOptions,
 ): Promise<{ content: Array<{ type: "text"; text: string }>; details: SearchRenderDetails }> {
+	const { authStorage, sessionId, signal } = options;
 	const providers =
 		params.provider && params.provider !== "auto"
-			? await getSearchProvider(params.provider).then(provider =>
-					provider.isAvailable() ? [provider] : resolveProviderChain("auto"),
+			? await getSearchProvider(params.provider).then(async provider =>
+					(await provider.isAvailable(authStorage)) ? [provider] : resolveProviderChain(authStorage, "auto"),
 				)
-			: await resolveProviderChain();
+			: await resolveProviderChain(authStorage);
 	if (providers.length === 0) {
 		const message = "No web search provider configured.";
 		return {
@@ -148,6 +156,8 @@ async function executeSearch(
 				numSearchResults: params.num_search_results,
 				temperature: params.temperature,
 				signal,
+				authStorage,
+				sessionId,
 			});
 
 			const text = formatForLLM(response);
@@ -190,18 +200,27 @@ async function executeSearch(
 
 /**
  * Execute a web search query for CLI/testing workflows.
+ *
+ * `authStorage` may be omitted; in that case we discover one via the standard
+ * factory (`discoverAuthStorage`), which honours `OMP_AUTH_BROKER_URL` and
+ * otherwise opens the local SQLite credential store.
  */
 export async function runSearchQuery(
 	params: SearchQueryParams,
+	options: { authStorage?: AuthStorage; sessionId?: string; signal?: AbortSignal } = {},
 ): Promise<{ content: Array<{ type: "text"; text: string }>; details: SearchRenderDetails }> {
-	return executeSearch("cli-web-search", params);
+	const authStorage = options.authStorage ?? (await discoverAuthStorage());
+	return executeSearch("cli-web-search", params, {
+		authStorage,
+		sessionId: options.sessionId,
+		signal: options.signal,
+	});
 }
 
 /**
  * Web search tool implementation.
  *
  * Supports Anthropic, Perplexity, Exa, Brave, Jina, Kimi, Gemini, Codex, Z.AI, SearXNG, and Synthetic providers with automatic fallback.
- * Session is accepted for interface consistency but not used.
  */
 export class WebSearchTool implements AgentTool<typeof webSearchSchema, SearchRenderDetails> {
 	readonly name = "web_search";
@@ -212,7 +231,10 @@ export class WebSearchTool implements AgentTool<typeof webSearchSchema, SearchRe
 	readonly loadMode = "discoverable";
 	readonly summary = "Search the web for up-to-date information";
 
-	constructor(_session: ToolSession) {
+	#session: ToolSession;
+
+	constructor(session: ToolSession) {
+		this.#session = session;
 		this.description = prompt.render(webSearchDescription);
 	}
 
@@ -223,7 +245,9 @@ export class WebSearchTool implements AgentTool<typeof webSearchSchema, SearchRe
 		_onUpdate?: AgentToolUpdateCallback<SearchRenderDetails>,
 		_context?: AgentToolContext,
 	): Promise<AgentToolResult<SearchRenderDetails>> {
-		return executeSearch(_toolCallId, params, signal);
+		const authStorage = this.#session.authStorage ?? (await discoverAuthStorage());
+		const sessionId = this.#session.getSessionId?.() ?? undefined;
+		return executeSearch(_toolCallId, params, { authStorage, sessionId, signal });
 	}
 }
 
@@ -238,10 +262,12 @@ export const webSearchCustomTool: CustomTool<typeof webSearchSchema, SearchRende
 		toolCallId: string,
 		params: SearchToolParams,
 		_onUpdate,
-		_ctx: CustomToolContext,
+		ctx: CustomToolContext,
 		signal?: AbortSignal,
 	) {
-		return executeSearch(toolCallId, params, signal);
+		const authStorage = ctx.modelRegistry?.authStorage ?? (await discoverAuthStorage());
+		const sessionId = ctx.sessionManager.getSessionId();
+		return executeSearch(toolCallId, params, { authStorage, sessionId, signal });
 	},
 
 	renderCall(args: SearchToolParams, options: RenderResultOptions, theme: Theme) {