@oh-my-pi/pi-coding-agent 15.11.3 → 15.11.4

package/src/web/kagi.ts

@@ -6,7 +6,7 @@
  * through the shared {@link AuthStorage} broker (Bearer token), and responses
  * are categorized result buckets rather than the legacy flat object array.
  */
-import type { AuthStorage, FetchImpl } from "@oh-my-pi/pi-ai";
+import { type AuthStorage, type FetchImpl, withAuth } from "@oh-my-pi/pi-ai";
 import { withHardTimeout } from "./search/providers/utils";
 
 const KAGI_SEARCH_URL = "https://kagi.com/api/v1/search";
@@ -173,14 +173,6 @@ export interface KagiSearchResult {
 	answer?: string;
 }
 
-export async function findKagiApiKey(
-	authStorage: AuthStorage,
-	sessionId?: string,
-	signal?: AbortSignal,
-): Promise<string | null> {
-	return (await authStorage.getApiKey("kagi", sessionId, { signal })) ?? null;
-}
-
 /**
  * Compute a YYYY-MM-DD date string `recency` units before now, in UTC.
  * UTC keeps the recency window deterministic regardless of host timezone and
@@ -247,27 +239,34 @@ export async function searchWithKagi(
 	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'.");
-	}
-
 	const fetchImpl = options.fetch ?? fetch;
+	const body = JSON.stringify(buildRequestBody(query, options));
+
+	const response = await withAuth(
+		authStorage.resolver("kagi", { sessionId: options.sessionId }),
+		async apiKey => {
+			const res = await fetchImpl(KAGI_SEARCH_URL, {
+				method: "POST",
+				headers: {
+					Authorization: `Bearer ${apiKey}`,
+					"Content-Type": "application/json",
+					Accept: "application/json",
+				},
+				body,
+				signal: withHardTimeout(options.signal),
+			});
+
+			if (!res.ok) {
+				throw parseKagiErrorResponse(res.status, await res.text());
+			}
 
-	const response = await fetchImpl(KAGI_SEARCH_URL, {
-		method: "POST",
-		headers: {
-			Authorization: `Bearer ${apiKey}`,
-			"Content-Type": "application/json",
-			Accept: "application/json",
+			return res;
 		},
-		body: JSON.stringify(buildRequestBody(query, options)),
-		signal: withHardTimeout(options.signal),
-	});
-
-	if (!response.ok) {
-		throw parseKagiErrorResponse(response.status, await response.text());
-	}
+		{
+			signal: options.signal,
+			missingKeyMessage: "Kagi credentials not found. Set KAGI_API_KEY or login with 'omp /login kagi'.",
+		},
+	);
 
 	const payload = (await response.json()) as KagiSearchResponse;
 	if (payload.error && payload.error.length > 0) {

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

@@ -7,7 +7,7 @@
  * SQLite store, never POSTs the broker sentinel to an OpenAI token endpoint.
  */
 import * as os from "node:os";
-import type { AuthStorage, FetchImpl } from "@oh-my-pi/pi-ai";
+import { type AuthStorage, type FetchImpl, type OAuthAccess, withOAuthAccess } from "@oh-my-pi/pi-ai";
 import { decodeJwt } from "@oh-my-pi/pi-ai/oauth/openai-codex";
 import { getBundledModels } from "@oh-my-pi/pi-catalog/models";
 import { $env, readSseJson } from "@oh-my-pi/pi-utils";
@@ -287,12 +287,12 @@ async function findCodexAuth(
 	authStorage: AuthStorage,
 	sessionId: string | undefined,
 	signal: AbortSignal | undefined,
-): Promise<{ accessToken: string; accountId: string } | null> {
+): Promise<{ access: OAuthAccess; accountId: string } | null> {
 	const access = await authStorage.getOAuthAccess("openai-codex", sessionId, { signal });
 	if (!access) return null;
 	const accountId = access.accountId ?? getAccountIdFromJwt(access.accessToken);
 	if (!accountId) return null;
-	return { accessToken: access.accessToken, accountId };
+	return { access, accountId };
 }
 
 /**
@@ -495,8 +495,8 @@ async function callCodexSearch(
  *   `gpt-5-codex-mini` first on ChatGPT accounts, which OpenAI rejects.
  */
 export async function searchCodex(params: SearchParams): Promise<SearchResponse> {
-	const auth = await findCodexAuth(params.authStorage, params.sessionId, params.signal);
-	if (!auth) {
+	const seed = await findCodexAuth(params.authStorage, params.sessionId, params.signal);
+	if (!seed) {
 		throw new Error(
 			"No Codex OAuth credentials found. Login with 'omp /login openai-codex' to enable Codex web search.",
 		);
@@ -505,42 +505,44 @@ export async function searchCodex(params: SearchParams): Promise<SearchResponse>
 	const configuredModel = getConfiguredModel();
 	const modelCandidates = configuredModel ? [configuredModel] : getDefaultModelCandidates();
 
-	let result:
-		| {
-				answer: string;
-				sources: SearchSource[];
-				model: string;
-				requestId: string;
-				usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
-		  }
-		| undefined;
-	let lastError: unknown;
-
-	for (let index = 0; index < modelCandidates.length; index += 1) {
-		const modelId = modelCandidates[index];
-		if (!modelId) continue;
-
-		try {
-			result = await callCodexSearch(auth, params.query, {
-				signal: params.signal,
-				systemPrompt: params.systemPrompt,
-				searchContextSize: "high",
-				modelId,
-				fetch: params.fetch,
-			});
-			break;
-		} catch (error) {
-			lastError = error;
-			const isLastCandidate = index === modelCandidates.length - 1;
-			if (configuredModel || isLastCandidate || !shouldRetryWithNextDefaultModel(error)) {
-				throw error;
+	const result = await withOAuthAccess(
+		params.authStorage,
+		"openai-codex",
+		async access => {
+			// Derive ALL auth material from the access this attempt received —
+			// a refreshed/rotated credential carries a different bearer and
+			// ChatGPT account id than the seed.
+			const accountId = access.accountId ?? getAccountIdFromJwt(access.accessToken);
+			if (!accountId) {
+				throw new Error("Codex OAuth credential is missing a ChatGPT account id");
 			}
-		}
-	}
-
-	if (!result) {
-		throw lastError ?? new Error("Codex search failed without returning a result");
-	}
+			const auth = { accessToken: access.accessToken, accountId };
+
+			let lastError: unknown;
+			for (let index = 0; index < modelCandidates.length; index += 1) {
+				const modelId = modelCandidates[index];
+				if (!modelId) continue;
+
+				try {
+					return await callCodexSearch(auth, params.query, {
+						signal: params.signal,
+						systemPrompt: params.systemPrompt,
+						searchContextSize: "high",
+						modelId,
+						fetch: params.fetch,
+					});
+				} catch (error) {
+					lastError = error;
+					const isLastCandidate = index === modelCandidates.length - 1;
+					if (configuredModel || isLastCandidate || !shouldRetryWithNextDefaultModel(error)) {
+						throw error;
+					}
+				}
+			}
+			throw lastError ?? new Error("Codex search failed without returning a result");
+		},
+		{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
+	);
 
 	let sources = result.sources;
 

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

@@ -8,7 +8,7 @@
  * sibling SQLite store and never POSTs the broker sentinel to a Google token
  * endpoint.
  */
-import type { AuthStorage, FetchImpl } from "@oh-my-pi/pi-ai";
+import { type AuthStorage, type FetchImpl, type OAuthAccess, withOAuthAccess } from "@oh-my-pi/pi-ai";
 import {
 	ANTIGRAVITY_SYSTEM_INSTRUCTION,
 	getAntigravityUserAgent,
@@ -72,25 +72,29 @@ interface GeminiAuth {
 	isAntigravity: boolean;
 }
 
+/** First configured Gemini OAuth provider plus its pre-resolved access. */
+interface GeminiAuthSeed {
+	provider: GeminiProviderId;
+	access: OAuthAccess;
+	projectId: string;
+}
+
 /**
  * Walks the configured Gemini OAuth providers in deterministic order and
  * returns the first one that yields a usable access token + projectId via
  * {@link AuthStorage.getOAuthAccess}. AuthStorage handles refresh + broker
  * routing internally; this helper never touches refresh tokens directly.
+ * The resolved access seeds `withOAuthAccess` so the happy path resolves once.
  */
 export async function findGeminiAuth(
 	authStorage: AuthStorage,
 	sessionId: string | undefined,
 	signal: AbortSignal | undefined,
-): Promise<GeminiAuth | null> {
+): Promise<GeminiAuthSeed | null> {
 	for (const provider of GEMINI_PROVIDERS) {
 		const access = await authStorage.getOAuthAccess(provider, sessionId, { signal });
 		if (!access?.accessToken || !access.projectId) continue;
-		return {
-			accessToken: access.accessToken,
-			projectId: access.projectId,
-			isAntigravity: provider === "google-antigravity",
-		};
+		return { provider, access, projectId: access.projectId };
 	}
 	return null;
 }
@@ -390,26 +394,42 @@ async function callGeminiSearch(
  * Executes a web search using Google Gemini with Google Search grounding.
  */
 export async function searchGemini(params: GeminiSearchParams): Promise<SearchResponse> {
-	const auth = await findGeminiAuth(params.authStorage, params.sessionId, params.signal);
-	if (!auth) {
+	const seed = await findGeminiAuth(params.authStorage, params.sessionId, params.signal);
+	if (!seed) {
 		throw new Error(
 			"No Gemini OAuth credentials found. Login with 'omp /login google-gemini-cli' or 'omp /login google-antigravity' to enable Gemini web search.",
 		);
 	}
 
-	const result = await callGeminiSearch(
-		auth,
-		params.query,
-		params.system_prompt,
-		params.max_output_tokens,
-		params.temperature,
-		{
-			google_search: params.google_search,
-			code_execution: params.code_execution,
-			url_context: params.url_context,
-		},
-		params.fetch,
-		params.signal,
+	const isAntigravity = seed.provider === "google-antigravity";
+	const result = await withOAuthAccess(
+		params.authStorage,
+		seed.provider,
+		access =>
+			// Derive bearer + projectId from the access this attempt received; a
+			// re-resolved access may omit projectId, in which case the seed's
+			// project is still the right tenant for the credential. The
+			// `fetchWithRetry` transport backoff stays INSIDE this attempt — auth
+			// retry wraps transport retry.
+			callGeminiSearch(
+				{
+					accessToken: access.accessToken,
+					projectId: access.projectId ?? seed.projectId,
+					isAntigravity,
+				},
+				params.query,
+				params.system_prompt,
+				params.max_output_tokens,
+				params.temperature,
+				{
+					google_search: params.google_search,
+					code_execution: params.code_execution,
+					url_context: params.url_context,
+				},
+				params.fetch,
+				params.signal,
+			),
+		{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
 	);
 
 	let sources = result.sources;

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

@@ -8,7 +8,7 @@
  * - Anonymous via `www.perplexity.ai/rest/sse/perplexity_ask`
  */
 
-import { type AuthStorage, type FetchImpl, getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { type AuthStorage, type FetchImpl, getEnvApiKey, type OAuthAccess, withOAuthAccess } from "@oh-my-pi/pi-ai";
 import { $env, readSseJson } from "@oh-my-pi/pi-utils";
 import type {
 	PerplexityMessageOutput,
@@ -43,7 +43,7 @@ type PerplexityAuth =
 	  }
 	| {
 			type: "oauth";
-			token: string;
+			access: OAuthAccess;
 	  }
 	| {
 			type: "cookies";
@@ -302,11 +302,11 @@ function jwtExpiryMs(token: string): number | undefined {
 	}
 }
 
-async function findOAuthToken(
+async function findOAuthAccess(
 	authStorage: AuthStorage,
 	sessionId: string | undefined,
 	signal: AbortSignal | undefined,
-): Promise<string | null> {
+): Promise<OAuthAccess | null> {
 	try {
 		// `getOAuthAccess` returns the raw OAuth bearer only — runtime/config
 		// api_key overrides and stored api_key credentials are intentionally
@@ -314,12 +314,12 @@ async function findOAuthToken(
 		// `www.perplexity.ai` session/SSE endpoint.
 		const access = await authStorage.getOAuthAccess("perplexity", sessionId, { signal });
 		const token = access?.accessToken;
-		if (!token) return null;
+		if (!access || !token) return null;
 		// Trust the JWT's own `exp` claim if it has one; otherwise treat as
 		// non-expiring. Perplexity session JWTs commonly omit `exp`.
 		const jwtExpiry = jwtExpiryMs(token);
 		if (jwtExpiry !== undefined && jwtExpiry <= Date.now() + OAUTH_EXPIRY_BUFFER_MS) return null;
-		return token;
+		return access;
 	} catch {
 		return null;
 	}
@@ -339,9 +339,9 @@ async function findPerplexityAuth(
 	const apiKey = findApiKey();
 
 	// 2. OAuth/session bearer from AuthStorage.
-	const oauthToken = await findOAuthToken(authStorage, sessionId, signal);
-	if (oauthToken) {
-		return { type: "oauth", token: oauthToken };
+	const oauthAccess = await findOAuthAccess(authStorage, sessionId, signal);
+	if (oauthAccess) {
+		return { type: "oauth", access: oauthAccess };
 	}
 
 	// 3. PERPLEXITY_API_KEY env var
@@ -646,7 +646,19 @@ export async function searchPerplexity(params: PerplexitySearchParams): Promise<
 	const auth = await findPerplexityAuth(params.authStorage, params.sessionId, params.signal);
 
 	if (auth.type !== "api_key") {
-		const askResult = await callPerplexityAsk(auth, params);
+		// OAuth bearer mode routes the whole authenticated unit (the ask
+		// session/SSE request) through the central auth-retry policy so a 401 or
+		// usage-limit force-refreshes, then rotates to a sibling credential.
+		// Cookie/env/anonymous modes have no rotatable credential — untouched.
+		const askResult =
+			auth.type === "oauth"
+				? await withOAuthAccess(
+						params.authStorage,
+						"perplexity",
+						access => callPerplexityAsk({ type: "oauth", token: access.accessToken }, params),
+						{ sessionId: params.sessionId, signal: params.signal, seed: auth.access },
+					)
+				: await callPerplexityAsk(auth, params);
 		return applySourceLimit(
 			{
 				provider: "perplexity",