From 503b2b7156e313002fbc1d12685eae9cec0759b7 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 12 Jun 2026 02:17:06 +0200 Subject: [PATCH] fix(coding-agent): routed web-search OAuth bearers through withOAuthAccess --- .../src/web/search/providers/codex.ts | 80 ++++++++++--------- .../src/web/search/providers/gemini.ts | 64 ++++++++++----- .../src/web/search/providers/perplexity.ts | 32 +++++--- 3 files changed, 105 insertions(+), 71 deletions(-) diff --git a/packages/coding-agent/src/web/search/providers/codex.ts b/packages/coding-agent/src/web/search/providers/codex.ts index e2bd085e3..bca571379 100644 --- a/packages/coding-agent/src/web/search/providers/codex.ts +++ b/packages/coding-agent/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 { - 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 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"); } - } - } + const auth = { accessToken: access.accessToken, accountId }; - if (!result) { - throw lastError ?? new Error("Codex search failed without returning a result"); - } + 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; diff --git a/packages/coding-agent/src/web/search/providers/gemini.ts b/packages/coding-agent/src/web/search/providers/gemini.ts index 8a3780a51..3d1d8c290 100644 --- a/packages/coding-agent/src/web/search/providers/gemini.ts +++ b/packages/coding-agent/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 { +): Promise { 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 { - 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; diff --git a/packages/coding-agent/src/web/search/providers/perplexity.ts b/packages/coding-agent/src/web/search/providers/perplexity.ts index ca7d8fc93..731a33fdd 100644 --- a/packages/coding-agent/src/web/search/providers/perplexity.ts +++ b/packages/coding-agent/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 { +): Promise { 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",