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

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);