diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 57b37854b..3062ec17d 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -10,6 +10,7 @@ - Added `ApiKeyResolver` and `ApiKey` auth helpers, including `isApiKeyResolver`, `isAuthRetryableError`, `resolveApiKeyOnce`, and `withAuth`, and exported them from the package root - Added support for a function-valued `apiKey` in `SimpleStreamOptions` so a single stream request can refresh or rotate credentials during retry - Added `forceRefresh` credential option to `AuthStorage.getApiKey` and `rotateSessionCredential` support for session-level credential rotation after auth failures +- Added `AuthStorage.resolver(provider, options)` method that builds an `ApiKeyResolver` implementing the a/b/c auth-retry policy directly on the storage instance ### Changed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index e4c80bd39..b76c2bbd6 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -590,36 +590,6 @@ export interface InvalidateCredentialMatchingOptions { sessionId?: string; } -/** - * Build an {@link ApiKeyResolver} backed directly by an {@link AuthStorage}, - * implementing the central a/b/c auth-retry policy for consumers that hold an - * AuthStorage rather than a higher-level registry (e.g. web-search providers): - * - * - initial (`error: undefined`) → resolve the session credential. - * - step (b) `!lastChance` → force-refresh the SAME session-sticky credential. - * - step (c) `lastChance` → rotate to a sibling credential, then re-resolve. - */ -export function createAuthStorageResolver( - storage: Pick, - provider: string, - options?: { sessionId?: string; baseUrl?: string; modelId?: string }, -): ApiKeyResolver { - const { sessionId, baseUrl, modelId } = options ?? {}; - return async ({ lastChance, error, signal }) => { - if (error === undefined) { - return storage.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); - } - if (lastChance) { - // Rotate to a sibling credential on an account constraint. No - // retry-after handling here — `markUsageLimitReached` owns the block - // duration and the outer retry layer owns the no-sibling wait. - await storage.rotateSessionCredential(provider, sessionId, { error, signal }); - return storage.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); - } - return storage.getApiKey(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); - }; -} - function isAbortSignalOption( value: InvalidateCredentialMatchingOptions | AbortSignal | undefined, ): value is AbortSignal { @@ -3787,6 +3757,31 @@ export class AuthStorage { return hasSibling; } + /** + * Build an {@link ApiKeyResolver} backed by this storage, implementing the + * central a/b/c auth-retry policy: + * + * - initial (`error: undefined`) → resolve the session credential. + * - step (b) `!lastChance` → force-refresh the SAME session-sticky credential. + * - step (c) `lastChance` → rotate to a sibling credential, then re-resolve. + * + * Used by web-search providers and other consumers that hold an AuthStorage + * directly (no ModelRegistry in scope). + */ + resolver(provider: string, options?: { sessionId?: string; baseUrl?: string; modelId?: string }): ApiKeyResolver { + const { sessionId, baseUrl, modelId } = options ?? {}; + return async ({ lastChance, error, signal }) => { + if (error === undefined) { + return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); + } + if (lastChance) { + await this.rotateSessionCredential(provider, sessionId, { error, signal }); + return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); + } + return this.getApiKey(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); + }; + } + // ─── Auth Broker integration ──────────────────────────────────────────── /** diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 83325513c..b6cb7ca6e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,7 +9,9 @@ ### Changed -- Changed non-agent API operations (title and commit-message generation) to use session-aware API key resolution with refresh-first auth retries, refreshing the active credential before rotating to another account +- Changed non-agent API operations (title and commit-message generation, image generation, web search, eval `llm()`, auto-thinking classifier, memory consolidation) to use session-aware API key resolution with auth retries via `registry.resolver()` / `authStorage.resolver()`, refreshing the active credential before rotating to another account +- Changed image generation to wrap every provider fetch branch in `withAuth`, so 401 / usage-limit errors trigger credential force-refresh and rotation for authStorage-backed providers (OpenAI-hosted, antigravity, xai-oauth) while env-only providers (openrouter, gemini) stay single-attempt +- Changed web-search providers using `authStorage.getApiKey` (anthropic, exa, tavily, parallel, synthetic, zai, kimi) to wrap HTTP calls in `withAuth` for automatic credential rotation on 401 / usage-limit errors - Changed the directory grouping for `find`, `search`, `ast_grep`, `ast_edit`, and `lsp` diagnostics from a single flat `# dir/` heading per immediate directory to a multi-level tree that folds the common path prefix into one heading. Previously every group repeated the full directory path — so results rooted outside cwd printed the absolute prefix (e.g. `/Users/me/proj/`) on every heading and nested directories were never collapsed. Now a single-child directory chain folds into one heading (`# packages/pkg/src/`, including an absolute root for out-of-cwd results), subdirectories nest one `#` deeper (`## nested/` → `### child.ts`), and each directory's own files are listed before its subdirectories. TUI hyperlink reconstruction tracks the nested directory stack across the whole output so file and code-frame links keep resolving to the correct absolute paths. - Changed the plan-mode approval surface from an inline transcript block plus a separate bottom selector into a single fullscreen overlay (like `/copy`) and overhauled its navigation. The overlay now renders the plan per-section through `ScrollView` (line-level ↑/↓ scroll, Shift+↑/↓ to scroll faster, PgUp/PgDn, g/G) with no stray per-line `…`, and — when the terminal is wide enough and the plan has ≥2 headings — shows a compact VS Code-style Contents sidebar (the redundant plan-title heading is omitted). Focus moves between regions with Tab/Shift+Tab (and flows at the edges: Down past the last section or the bottom of the body drops into the approval options; Up steps back), while the sidebar glows to track the scrolled section. The sidebar can fast-jump between sections, delete a section (with `u` undo), and annotate sections with feedback (`a`) that is submitted back to the model when the operator picks "Refine plan"; ←/→ always drive the model-tier slider, Enter confirms, the external-editor key opens the plan, and Esc cancels. The overlay borrows the terminal's alternate screen buffer for its lifetime (`fullscreen` overlay), so the transcript stays put on the normal screen instead of bleeding through scrollback behind the modal. - Changed the interactive controllers (command, MCP, selector, extension-UI, event), debug panels, and the status/error/warning helpers to render chat output through `ctx.present(...)` instead of appending to `chatContainer` and calling `ui.requestRender()` directly; transcript rebuilds dispose live blocks via `ctx.resetTranscript()` so animated blocks' timers stop on reset. diff --git a/packages/coding-agent/src/auto-thinking/classifier.ts b/packages/coding-agent/src/auto-thinking/classifier.ts index fc029123a..ccae82f73 100644 --- a/packages/coding-agent/src/auto-thinking/classifier.ts +++ b/packages/coding-agent/src/auto-thinking/classifier.ts @@ -15,7 +15,7 @@ */ import { type AssistantMessage, completeSimple, Effort, type Model } from "@oh-my-pi/pi-ai"; import { prompt } from "@oh-my-pi/pi-utils"; -import { reuseInitialApiKey } from "../config/api-key-resolver"; + import type { ModelRegistry } from "../config/model-registry"; import { resolveRoleSelection } from "../config/model-resolver"; import type { Settings } from "../config/settings"; @@ -83,7 +83,7 @@ async function classifyOnline(input: string, deps: ClassifyDifficultyDeps): Prom messages: [{ role: "user", content: input, timestamp: Date.now() }], }, { - apiKey: reuseInitialApiKey(apiKey, deps.registry, model.provider, { + apiKey: deps.registry.resolver(model.provider, { sessionId: deps.sessionId, baseUrl: model.baseUrl, }), diff --git a/packages/coding-agent/src/commit/model-selection.ts b/packages/coding-agent/src/commit/model-selection.ts index 00fadcc15..135bb6be7 100644 --- a/packages/coding-agent/src/commit/model-selection.ts +++ b/packages/coding-agent/src/commit/model-selection.ts @@ -1,6 +1,6 @@ import type { ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import type { Api, ApiKey, Model } from "@oh-my-pi/pi-ai"; -import { type ApiKeyResolverRegistry, reuseInitialApiKey } from "../config/api-key-resolver"; +import type { ApiKeyResolverRegistry } from "../config/api-key-resolver"; import { MODEL_ROLE_IDS } from "../config/model-registry"; import { type ModelLookupRegistry, @@ -47,7 +47,7 @@ export async function resolvePrimaryModel( } return { model, - apiKey: reuseInitialApiKey(apiKey, modelRegistry, model.provider, { baseUrl: model.baseUrl }), + apiKey: modelRegistry.resolver(model.provider, { baseUrl: model.baseUrl }), thinkingLevel: resolved?.thinkingLevel, }; } @@ -65,7 +65,7 @@ export async function resolveSmolModel( if (apiKey) { return { model: resolvedSmol.model, - apiKey: reuseInitialApiKey(apiKey, modelRegistry, resolvedSmol.model.provider, { + apiKey: modelRegistry.resolver(resolvedSmol.model.provider, { baseUrl: resolvedSmol.model.baseUrl, }), thinkingLevel: resolvedSmol.thinkingLevel, @@ -81,7 +81,7 @@ export async function resolveSmolModel( if (apiKey) { return { model: candidate, - apiKey: reuseInitialApiKey(apiKey, modelRegistry, candidate.provider, { baseUrl: candidate.baseUrl }), + apiKey: modelRegistry.resolver(candidate.provider, { baseUrl: candidate.baseUrl }), }; } } diff --git a/packages/coding-agent/src/config/api-key-resolver.ts b/packages/coding-agent/src/config/api-key-resolver.ts index fa039aa49..5204c5496 100644 --- a/packages/coding-agent/src/config/api-key-resolver.ts +++ b/packages/coding-agent/src/config/api-key-resolver.ts @@ -19,25 +19,23 @@ export interface ApiKeyResolverRegistry { options?: { baseUrl?: string; forceRefresh?: boolean; signal?: AbortSignal }, ): Promise; authStorage: Pick; + /** + * Build an {@link ApiKeyResolver} implementing the central a/b/c auth-retry + * policy: initial → resolve; step (b) → force-refresh same account; step (c) + * → rotate to a sibling credential, then re-resolve. + * + * The resolver is stateless (safe to reuse across requests). Callers that + * need the initial key for a guard can call `resolveApiKeyOnce(resolver)`. + */ + resolver(provider: string, options?: ApiKeyResolverOptions): ApiKeyResolver; } /** - * Build an {@link ApiKeyResolver} backed by the model registry's auth storage, - * implementing the central a/b/c auth-retry policy for every non-agent network - * consumer (utility completions, image generation, web search): - * - * - initial (`error: undefined`) → resolve the session credential (cheap; may - * return a locally-cached not-yet-expired token). - * - step (b) `!lastChance` → force-refresh the SAME session-sticky credential - * (a peer/broker may have rotated its token out from under our cached copy). - * - step (c) `lastChance` → rotate to a sibling credential (usage-limit block - * vs credential invalidation, by error class), then re-resolve. - * - * Stateless: nothing is captured beyond the registry/provider/options, so the - * same resolver is safe to reuse across attempts and requests. + * Default implementation of {@link ApiKeyResolverRegistry.resolver}. + * Also usable standalone for structural registries that don't carry the method. */ export function createApiKeyResolver( - registry: ApiKeyResolverRegistry, + registry: Pick, provider: string, options: ApiKeyResolverOptions = {}, ): ApiKeyResolver { @@ -58,20 +56,3 @@ export function createApiKeyResolver( return registry.getApiKeyForProvider(provider, sessionId, { baseUrl, forceRefresh: true, signal }); }; } - -/** - * Wrap an already-resolved `initialKey` in an {@link ApiKeyResolver}: the - * initial step returns that key verbatim (so a caller that resolved it eagerly - * — e.g. to preserve a metadata / guard ordering, or to short-circuit when no - * credential exists — does not pay for a second resolve), while retry steps - * fall through to {@link createApiKeyResolver}'s force-refresh / rotate policy. - */ -export function reuseInitialApiKey( - initialKey: string | undefined, - registry: ApiKeyResolverRegistry, - provider: string, - options: ApiKeyResolverOptions = {}, -): ApiKeyResolver { - const retry = createApiKeyResolver(registry, provider, options); - return ctx => (ctx.error === undefined ? initialKey : retry(ctx)); -} diff --git a/packages/coding-agent/src/config/model-registry.ts b/packages/coding-agent/src/config/model-registry.ts index 01b21eb04..6b0a86a23 100644 --- a/packages/coding-agent/src/config/model-registry.ts +++ b/packages/coding-agent/src/config/model-registry.ts @@ -95,12 +95,14 @@ const STARTUP_MODEL_CACHE_PROVIDER_IDS: readonly string[] = [ ...SPECIAL_MODEL_MANAGER_PROVIDER_IDS, ]; +import type { ApiKeyResolver } from "@oh-my-pi/pi-ai"; import { registerOAuthProvider, unregisterOAuthProviders } from "@oh-my-pi/pi-ai/utils/oauth"; import type { OAuthCredentials, OAuthLoginCallbacks } from "@oh-my-pi/pi-ai/utils/oauth/types"; import { isRecord, logger } from "@oh-my-pi/pi-utils"; import { parseModelString, resolveProviderModelReference } from "../config/model-resolver"; import { isValidThemeColor, type ThemeColor } from "../modes/theme/theme"; import type { AuthStorage, OAuthCredential } from "../session/auth-storage"; +import { type ApiKeyResolverOptions, createApiKeyResolver } from "./api-key-resolver"; import { type ConfigError, ConfigFile } from "./config-file"; import { buildCanonicalModelIndex, @@ -2385,6 +2387,15 @@ export class ModelRegistry { }); } + /** + * Build an {@link ApiKeyResolver} for this provider, implementing the + * central a/b/c auth-retry policy. Callers that need the initial key for + * a guard can call `resolveApiKeyOnce(resolver)`. + */ + resolver(provider: string, options?: ApiKeyResolverOptions): ApiKeyResolver { + return createApiKeyResolver(this, provider, options); + } + async #peekApiKeyForProvider(provider: string): Promise { if (this.#keylessProviders.has(provider) && !this.authStorage.hasAuth(provider)) { return kNoAuth; diff --git a/packages/coding-agent/src/eval/llm-bridge.ts b/packages/coding-agent/src/eval/llm-bridge.ts index 484290352..1f37dc2c9 100644 --- a/packages/coding-agent/src/eval/llm-bridge.ts +++ b/packages/coding-agent/src/eval/llm-bridge.ts @@ -15,7 +15,7 @@ import { instrumentedCompleteSimple, resolveTelemetry } from "@oh-my-pi/pi-agent import { type Api, Effort, getSupportedEfforts, type Model, type Tool } from "@oh-my-pi/pi-ai"; import * as z from "zod/v4"; import { extractTextContent, extractToolCall, parseJsonPayload } from "../commit/utils"; -import { reuseInitialApiKey } from "../config/api-key-resolver"; + import { expandRoleAlias, formatModelString, resolveModelFromString } from "../config/model-resolver"; import type { ToolSession } from "../tools"; import { ToolError } from "../tools/tool-errors"; @@ -145,7 +145,7 @@ export async function runEvalLlm(args: unknown, options: EvalLlmBridgeOptions): tools, }, { - apiKey: reuseInitialApiKey(apiKey, registry, model.provider, { + apiKey: registry.resolver(model.provider, { sessionId: options.session.getSessionId?.() ?? undefined, baseUrl: model.baseUrl, }), diff --git a/packages/coding-agent/src/memories/index.ts b/packages/coding-agent/src/memories/index.ts index 5ddafbba6..38205c5e2 100644 --- a/packages/coding-agent/src/memories/index.ts +++ b/packages/coding-agent/src/memories/index.ts @@ -5,7 +5,7 @@ import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { type ApiKey, clampThinkingLevelForModel, completeSimple, Effort, type Model } from "@oh-my-pi/pi-ai"; import { getAgentDbPath, getMemoriesDir, logger, parseJsonlLenient, prompt } from "@oh-my-pi/pi-utils"; -import { reuseInitialApiKey } from "../config/api-key-resolver"; + import type { ModelRegistry } from "../config/model-registry"; import { resolveModelRoleValue } from "../config/model-resolver"; import type { Settings } from "../config/settings"; @@ -272,7 +272,7 @@ async function runPhase1(options: { const result = await runStage1Job({ claim, model: phase1Model, - apiKey: reuseInitialApiKey(phase1ApiKey, modelRegistry, phase1Model.provider, { + apiKey: modelRegistry.resolver(phase1Model.provider, { sessionId: session.sessionId, baseUrl: phase1Model.baseUrl, }), @@ -432,7 +432,7 @@ async function runPhase2(options: { const consolidated = await runConsolidationModel({ memoryRoot, model: phase2Model, - apiKey: reuseInitialApiKey(phase2ApiKey, modelRegistry, phase2Model.provider, { + apiKey: modelRegistry.resolver(phase2Model.provider, { sessionId: session.sessionId, baseUrl: phase2Model.baseUrl, }), diff --git a/packages/coding-agent/src/mnemopi/backend.ts b/packages/coding-agent/src/mnemopi/backend.ts index 9af906894..7eb6d531d 100644 --- a/packages/coding-agent/src/mnemopi/backend.ts +++ b/packages/coding-agent/src/mnemopi/backend.ts @@ -5,7 +5,7 @@ import { Mnemopi } from "@oh-my-pi/pi-mnemopi"; import { BankManager } from "@oh-my-pi/pi-mnemopi/core"; import { type DiagnosticSummary, inspectDatabase } from "@oh-my-pi/pi-mnemopi/diagnose"; import { logger } from "@oh-my-pi/pi-utils"; -import { reuseInitialApiKey } from "../config/api-key-resolver"; + import type { ModelRegistry } from "../config/model-registry"; import { resolveRoleSelection } from "../config/model-resolver"; import type { MemoryBackend, MemoryBackendStartOptions } from "../memory-backend/types"; @@ -335,7 +335,7 @@ async function resolveMnemopiProviderOptions( messages: [{ role: "user", content: prompt, timestamp: Date.now() }], }, { - apiKey: reuseInitialApiKey(apiKey, modelRegistry, model.provider, { + apiKey: modelRegistry.resolver(model.provider, { sessionId, baseUrl: model.baseUrl, }), diff --git a/packages/coding-agent/src/tools/image-gen.ts b/packages/coding-agent/src/tools/image-gen.ts index 521fb94e6..ee90c5ca0 100644 --- a/packages/coding-agent/src/tools/image-gen.ts +++ b/packages/coding-agent/src/tools/image-gen.ts @@ -1,6 +1,6 @@ import * as os from "node:os"; import * as path from "node:path"; -import { getAntigravityUserAgent, getEnvApiKey, type Model } from "@oh-my-pi/pi-ai"; +import { type ApiKey, getAntigravityUserAgent, getEnvApiKey, type Model, withAuth } from "@oh-my-pi/pi-ai"; import { CODEX_BASE_URL, getCodexAccountId, @@ -20,6 +20,7 @@ import { } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; import packageJson from "../../package.json" with { type: "json" }; + import { isAuthenticated, type ModelRegistry } from "../config/model-registry"; import type { CustomTool } from "../extensibility/custom-tools/types"; import { ohMyPiXAIUserAgent, resolveXAIHttpCredentials } from "../lib/xai-http"; @@ -864,7 +865,10 @@ async function generateOpenAIHostedImage( if (!response.ok) { const errorText = await response.text(); - throw new Error(`OpenAI image request failed (${response.status}): ${getOpenAIResponseErrorMessage(errorText)}`); + throw Object.assign( + new Error(`OpenAI image request failed (${response.status}): ${getOpenAIResponseErrorMessage(errorText)}`), + { status: response.status }, + ); } const contentType = response.headers.get("content-type") ?? ""; @@ -1037,13 +1041,16 @@ export const imageGenTool: CustomTool generateOpenAIHostedImage(key, hostedModel, params, resolvedImages, requestSignal, sessionId), + { signal: requestSignal }, ); if (parsed.images.length === 0) { @@ -1088,38 +1095,57 @@ export const imageGenTool: CustomTool { + // On a retry the resolver yields the raw stored credential JSON + // ({ token, projectId }); the initial seed is the already-parsed + // access token. Tolerate both, falling back to the seed projectId. + const rotated = parseAntigravityCredentials(key); + const bearer = rotated?.accessToken ?? key; + const projectId = rotated?.projectId ?? apiKey.projectId!; + const requestBody = buildAntigravityRequest( + prompt, + model, + projectId, + params.aspect_ratio, + params.image_size, + resolvedImages, + ); + + const resp = await fetch(`${ANTIGRAVITY_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`, { + method: "POST", + headers: { + Authorization: `Bearer ${bearer}`, + "Content-Type": "application/json", + Accept: "text/event-stream", + "User-Agent": getAntigravityUserAgent(), + }, + body: JSON.stringify(requestBody), + signal: requestSignal, + }); + + if (!resp.ok) { + const errorText = await resp.text(); + let message = errorText; + try { + const parsedErr = JSON.parse(errorText) as { error?: { message?: string } }; + message = parsedErr.error?.message ?? message; + } catch { + // Keep raw text. + } + throw Object.assign(new Error(`Antigravity image request failed (${resp.status}): ${message}`), { + status: resp.status, + }); + } + return resp; + }, + { signal: requestSignal }, + ); const parsed = await parseAntigravitySseForImage(response, requestSignal); const responseText = parsed.text.length > 0 ? parsed.text.join(" ") : undefined; @@ -1191,28 +1217,41 @@ export const imageGenTool: CustomTool { + const resp = await fetch(`${xaiCreds.baseURL}${xaiEndpoint}`, { + method: "POST", + headers: { + Authorization: `Bearer ${key}`, + "Content-Type": "application/json", + "User-Agent": ohMyPiXAIUserAgent(), + }, + body: JSON.stringify(xaiBody), + signal: requestSignal, + }); + const rawText = await resp.text(); + if (!resp.ok) { + let message = rawText; + try { + const parsedErr = JSON.parse(rawText) as { error?: { message?: string } }; + message = parsedErr.error?.message ?? message; + } catch { + // Keep raw text. + } + throw Object.assign(new Error(`xAI image request failed (${resp.status}): ${message}`), { + status: resp.status, + }); + } + return rawText; + }, + { signal: requestSignal }, + ); const xaiData = JSON.parse(xaiRawText) as { data?: Array<{ b64_json?: string; url?: string }>; @@ -1269,30 +1308,34 @@ export const imageGenTool: CustomTool { + const resp = await fetch("https://openrouter.ai/api/v1/chat/completions", { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${key}`, + "HTTP-Referer": "https://omp.sh/", + "X-OpenRouter-Title": "Oh-My-Pi", + "X-OpenRouter-Categories": "cli-agent", + }, + body: JSON.stringify(requestBody), + signal: requestSignal, + }); + const text = await resp.text(); + if (!resp.ok) { + let message = text; + try { + const parsed = JSON.parse(text) as { error?: { message?: string } }; + message = parsed.error?.message ?? message; + } catch { + // Keep raw text. + } + throw Object.assign(new Error(`OpenRouter image request failed (${resp.status}): ${message}`), { + status: resp.status, + }); } - throw new Error(`OpenRouter image request failed (${response.status}): ${message}`); - } + return text; + }); const data = JSON.parse(rawText) as OpenRouterResponse; const message = data.choices?.[0]?.message; @@ -1360,30 +1403,34 @@ export const imageGenTool: CustomTool { + const resp = await fetch( + `https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(model)}:generateContent`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + "x-goog-api-key": key, + }, + body: JSON.stringify(requestBody), + signal: requestSignal, }, - body: JSON.stringify(requestBody), - signal: requestSignal, - }, - ); - - const rawText = await response.text(); - if (!response.ok) { - let message = rawText; - try { - const parsed = JSON.parse(rawText) as { error?: { message?: string } }; - message = parsed.error?.message ?? message; - } catch { - // Keep raw text. + ); + const text = await resp.text(); + if (!resp.ok) { + let message = text; + try { + const parsed = JSON.parse(text) as { error?: { message?: string } }; + message = parsed.error?.message ?? message; + } catch { + // Keep raw text. + } + throw Object.assign(new Error(`Gemini image request failed (${resp.status}): ${message}`), { + status: resp.status, + }); } - throw new Error(`Gemini image request failed (${response.status}): ${message}`); - } + return text; + }); const data = JSON.parse(rawText) as GeminiGenerateContentResponse; const responseParts = combineParts(data); diff --git a/packages/coding-agent/src/tools/inspect-image.ts b/packages/coding-agent/src/tools/inspect-image.ts index 3759e5ed1..e7db9fc3c 100644 --- a/packages/coding-agent/src/tools/inspect-image.ts +++ b/packages/coding-agent/src/tools/inspect-image.ts @@ -4,6 +4,7 @@ import { type Api, completeSimple, type Model } from "@oh-my-pi/pi-ai"; import { prompt } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; import { extractTextContent } from "../commit/utils"; + import { expandRoleAlias, resolveModelFromString } from "../config/model-resolver"; import inspectImageDescription from "../prompts/tools/inspect-image.md" with { type: "text" }; import inspectImageSystemPromptTemplate from "../prompts/tools/inspect-image-system.md" with { type: "text" }; @@ -136,7 +137,13 @@ export class InspectImageTool implements AgentTool { 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); diff --git a/packages/coding-agent/src/web/search/providers/exa.ts b/packages/coding-agent/src/web/search/providers/exa.ts index bbca2a5b7..7c0642cbd 100644 --- a/packages/coding-agent/src/web/search/providers/exa.ts +++ b/packages/coding-agent/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 { + // 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[] = []; diff --git a/packages/coding-agent/src/web/search/providers/kimi.ts b/packages/coding-agent/src/web/search/providers/kimi.ts index 2bc8d53b1..08401dbe6 100644 --- a/packages/coding-agent/src/web/search/providers/kimi.ts +++ b/packages/coding-agent/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 { +): Promise { 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 { - 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 ?? []) { diff --git a/packages/coding-agent/src/web/search/providers/parallel.ts b/packages/coding-agent/src/web/search/providers/parallel.ts index 19f1b8164..876ef4a8c 100644 --- a/packages/coding-agent/src/web/search/providers/parallel.ts +++ b/packages/coding-agent/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( diff --git a/packages/coding-agent/src/web/search/providers/synthetic.ts b/packages/coding-agent/src/web/search/providers/synthetic.ts index c531a9a65..12d4a6ca1 100644 --- a/packages/coding-agent/src/web/search/providers/synthetic.ts +++ b/packages/coding-agent/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 { - 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 ?? []) { diff --git a/packages/coding-agent/src/web/search/providers/tavily.ts b/packages/coding-agent/src/web/search/providers/tavily.ts index a32a27856..9b5e0a9f3 100644 --- a/packages/coding-agent/src/web/search/providers/tavily.ts +++ b/packages/coding-agent/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 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 ?? []) { diff --git a/packages/coding-agent/src/web/search/providers/zai.ts b/packages/coding-agent/src/web/search/providers/zai.ts index 4795038b8..2f39ff682 100644 --- a/packages/coding-agent/src/web/search/providers/zai.ts +++ b/packages/coding-agent/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 { - 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); diff --git a/packages/coding-agent/test/commit-model-selection-role-thinking.test.ts b/packages/coding-agent/test/commit-model-selection-role-thinking.test.ts index 71d8134ec..f5cffc8f5 100644 --- a/packages/coding-agent/test/commit-model-selection-role-thinking.test.ts +++ b/packages/coding-agent/test/commit-model-selection-role-thinking.test.ts @@ -39,7 +39,8 @@ describe("commit role thinking selection", () => { getAvailable: () => [defaultModel, commitModel], getApiKey: async () => "test-key", getApiKeyForProvider: async () => "test-key", - authStorage: { rotateSessionCredential: async () => false }, + authStorage: { rotateSessionCredential: async () => false as const }, + resolver: () => async () => "test-key", }; const primary = await resolvePrimaryModel(undefined, settings, registry); diff --git a/packages/coding-agent/test/title-generator.test.ts b/packages/coding-agent/test/title-generator.test.ts index dc1a6b4a4..46069a57f 100644 --- a/packages/coding-agent/test/title-generator.test.ts +++ b/packages/coding-agent/test/title-generator.test.ts @@ -29,6 +29,9 @@ function createRegistry(model: Model) { return { getAvailable: () => [model], getApiKey: async () => "test-key", + getApiKeyForProvider: async () => "test-key", + authStorage: { rotateSessionCredential: async () => false }, + resolver: () => async () => "test-key", } as never; } diff --git a/packages/coding-agent/test/tools/image-gen.test.ts b/packages/coding-agent/test/tools/image-gen.test.ts index 74b5182db..2ecdaa8bc 100644 --- a/packages/coding-agent/test/tools/image-gen.test.ts +++ b/packages/coding-agent/test/tools/image-gen.test.ts @@ -62,6 +62,8 @@ describe("imageGenTool", () => { modelRegistry: { getApiKey: async () => "test-openai-key", getApiKeyForProvider: async () => undefined, + authStorage: { rotateSessionCredential: async () => false }, + resolver: () => async () => "test-openai-key", } as unknown as ModelRegistry, model, isIdle: () => true, @@ -126,7 +128,9 @@ describe("imageGenTool", () => { getAll: () => [], authStorage: { hasNonEnvCredential: (provider: string) => provider === "xai-oauth", + rotateSessionCredential: async () => false, }, + resolver: () => async () => "test-xai-token", } as unknown as ModelRegistry, model: undefined, isIdle: () => true, diff --git a/packages/coding-agent/test/tools/inspect-image.test.ts b/packages/coding-agent/test/tools/inspect-image.test.ts index fe27a38cb..8345c1865 100644 --- a/packages/coding-agent/test/tools/inspect-image.test.ts +++ b/packages/coding-agent/test/tools/inspect-image.test.ts @@ -68,6 +68,9 @@ function createSession( modelRegistry: { getAvailable: () => availableModels, getApiKey: async () => apiKey, + getApiKeyForProvider: async () => apiKey, + authStorage: { rotateSessionCredential: async () => false }, + resolver: () => async () => apiKey, } as unknown as NonNullable, }; } diff --git a/packages/coding-agent/test/tools/search-renderer.test.ts b/packages/coding-agent/test/tools/search-renderer.test.ts index 13f1c7f32..38d04c0ec 100644 --- a/packages/coding-agent/test/tools/search-renderer.test.ts +++ b/packages/coding-agent/test/tools/search-renderer.test.ts @@ -52,7 +52,7 @@ describe("searchToolRenderer", () => { }, ); const renderedLines = sanitizeText(collapsed.render(200).join("\n")).split("\n"); - const bodyLines = renderedLines.slice(1); + const bodyLines = renderedLines.slice(1, -1); expect(renderedLines[0]).toContain("truncated"); expect(bodyLines).toHaveLength(6); @@ -93,7 +93,7 @@ describe("searchToolRenderer", () => { { pattern: "Flag" }, ); const renderedLines = sanitizeText(collapsed.render(240).join("\n")).split("\n"); - const bodyLines = renderedLines.slice(1); + const bodyLines = renderedLines.slice(1, -1); expect(bodyLines).toHaveLength(6); expect(bodyLines.some(line => line.includes("const firstFlag = true;"))).toBe(true); diff --git a/packages/coding-agent/test/tools/web-search-parallel.test.ts b/packages/coding-agent/test/tools/web-search-parallel.test.ts index 2fa5f6d55..1886417cb 100644 --- a/packages/coding-agent/test/tools/web-search-parallel.test.ts +++ b/packages/coding-agent/test/tools/web-search-parallel.test.ts @@ -30,6 +30,12 @@ describe("Parallel web search", () => { hasAuth() { return Boolean(process.env.PARALLEL_API_KEY); }, + resolver(_provider: string) { + return async () => process.env.PARALLEL_API_KEY ?? undefined; + }, + async rotateSessionCredential() { + return false; + }, } as unknown as AuthStorage; let capturedRequestBody: unknown; diff --git a/packages/coding-agent/test/tools/web-search-tavily.test.ts b/packages/coding-agent/test/tools/web-search-tavily.test.ts index a260179c7..bea6429f1 100644 --- a/packages/coding-agent/test/tools/web-search-tavily.test.ts +++ b/packages/coding-agent/test/tools/web-search-tavily.test.ts @@ -21,6 +21,12 @@ describe("Tavily web search provider", () => { hasAuth() { return Boolean(process.env.TAVILY_API_KEY); }, + resolver(_provider: string) { + return async () => process.env.TAVILY_API_KEY ?? undefined; + }, + async rotateSessionCredential() { + return false; + }, } as unknown as AuthStorage; function makeParams(query: string) {