@prometheus-ai/agent 0.5.3 → 0.5.8

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

@@ -8,12 +8,12 @@
  * sibling SQLite store and never POSTs the broker sentinel to a Google token
  * endpoint.
  */
+import { type AuthStorage, type FetchImpl, type OAuthAccess, withOAuthAccess } from "@prometheus-ai/ai";
 import {
 	ANTIGRAVITY_SYSTEM_INSTRUCTION,
-	type AuthStorage,
 	getAntigravityUserAgent,
 	getGeminiCliHeaders,
-} from "@prometheus-ai/ai";
+} from "@prometheus-ai/catalog/wire/gemini-headers";
 import { fetchWithRetry } from "@prometheus-ai/utils";
 
 import type { SearchCitation, SearchResponse, SearchSource } from "../../../web/search/types";
@@ -51,6 +51,7 @@ export interface GeminiSearchParams extends GeminiToolParams {
 	signal?: AbortSignal;
 	authStorage: AuthStorage;
 	sessionId?: string;
+	fetch?: FetchImpl;
 }
 
 export function buildGeminiRequestTools(params: GeminiToolParams): Array<Record<string, Record<string, unknown>>> {
@@ -71,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;
 }
@@ -156,6 +161,7 @@ async function callGeminiSearch(
 	maxOutputTokens: number | undefined,
 	temperature: number | undefined,
 	toolParams: GeminiToolParams,
+	fetchImpl: FetchImpl | undefined,
 	signal: AbortSignal | undefined,
 ): Promise<{
 	answer: string;
@@ -175,8 +181,8 @@ async function callGeminiSearch(
 				requestId: `agent-${crypto.randomUUID()}`,
 			}
 		: {
-				userAgent: "prometheus-coding-agent",
-				requestId: `pi-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
+				userAgent: "prometheus-agent",
+				requestId: `prometheus-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
 			};
 
 	const normalizedSystemPrompt = systemPrompt?.toWellFormed();
@@ -237,6 +243,7 @@ async function callGeminiSearch(
 
 	const response = await fetchWithRetry(urlFor, {
 		...buildInit(),
+		fetch: fetchImpl,
 		maxAttempts: MAX_RETRIES + 1,
 		defaultDelayMs: attempt => BASE_DELAY_MS * 2 ** attempt,
 		maxDelayMs: RATE_LIMIT_BUDGET_MS,
@@ -387,25 +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 'prometheus /login google-gemini-cli' or 'prometheus /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.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;
@@ -450,6 +474,7 @@ export class GeminiProvider extends SearchProvider {
 			signal: params.signal,
 			authStorage: params.authStorage,
 			sessionId: params.sessionId,
+			fetch: params.fetch,
 		});
 	}
 }

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

@@ -5,7 +5,7 @@
  * cleaned content.
  */
 
-import { type AuthStorage, getEnvApiKey } from "@prometheus-ai/ai";
+import { type AuthStorage, type FetchImpl, getEnvApiKey } from "@prometheus-ai/ai";
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import type { SearchParams } from "./base";
@@ -13,11 +13,13 @@ import { SearchProvider } from "./base";
 import { classifyProviderHttpError, withHardTimeout } from "./utils";
 
 const JINA_SEARCH_URL = "https://s.jina.ai";
+type SearchParamsWithFetch = SearchParams & { fetch?: FetchImpl };
 
 export interface JinaSearchParams {
 	query: string;
 	num_results?: number;
 	signal?: AbortSignal;
+	fetch?: FetchImpl;
 }
 
 interface JinaSearchResult {
@@ -34,9 +36,14 @@ export function findApiKey(): string | null {
 }
 
 /** Call Jina Reader search API. */
-async function callJinaSearch(apiKey: string, query: string, signal?: AbortSignal): Promise<JinaSearchResponse> {
+async function callJinaSearch(
+	apiKey: string,
+	query: string,
+	signal?: AbortSignal,
+	fetchImpl: FetchImpl = fetch,
+): Promise<JinaSearchResponse> {
 	const requestUrl = `${JINA_SEARCH_URL}/${encodeURIComponent(query)}`;
-	const response = await fetch(requestUrl, {
+	const response = await fetchImpl(requestUrl, {
 		headers: {
 			Accept: "application/json",
 			Authorization: `Bearer ${apiKey}`,
@@ -62,7 +69,7 @@ export async function searchJina(params: JinaSearchParams): Promise<SearchRespon
 		throw new Error("JINA_API_KEY not found. Set it in environment or .env file.");
 	}
 
-	const response = await callJinaSearch(apiKey, params.query, params.signal);
+	const response = await callJinaSearch(apiKey, params.query, params.signal, params.fetch);
 	const sources: SearchSource[] = [];
 
 	for (const result of response) {
@@ -91,11 +98,14 @@ export class JinaProvider extends SearchProvider {
 		return !!findApiKey();
 	}
 
-	search(params: SearchParams): Promise<SearchResponse> {
+	search(params: SearchParamsWithFetch): Promise<SearchResponse> {
+		const fetchImpl = params.fetch;
+
 		return searchJina({
 			query: params.query,
 			num_results: params.numSearchResults ?? params.limit,
 			signal: params.signal,
+			fetch: fetchImpl,
 		});
 	}
 }

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

@@ -3,7 +3,7 @@
  *
  * Thin wrapper that adapts shared Kagi API utilities to SearchResponse shape.
  */
-import type { AuthStorage } from "@prometheus-ai/ai";
+import type { AuthStorage, FetchImpl } from "@prometheus-ai/ai";
 import type { SearchResponse } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import { KagiApiError, searchWithKagi } from "../../kagi";
@@ -12,6 +12,8 @@ import type { SearchParams } from "./base";
 import { SearchProvider } from "./base";
 import { classifyProviderHttpError, toSearchSources } from "./utils";
 
+type SearchParamsWithFetch = SearchParams & { fetch?: FetchImpl };
+
 const DEFAULT_NUM_RESULTS = 10;
 const MAX_NUM_RESULTS = 40;
 
@@ -23,6 +25,7 @@ export async function searchKagi(params: {
 	signal?: AbortSignal;
 	authStorage: AuthStorage;
 	sessionId?: string;
+	fetch?: FetchImpl;
 }): Promise<SearchResponse> {
 	const numResults = clampNumResults(params.num_results, DEFAULT_NUM_RESULTS, MAX_NUM_RESULTS);
 
@@ -34,6 +37,7 @@ export async function searchKagi(params: {
 				recency: params.recency,
 				sessionId: params.sessionId,
 				signal: params.signal,
+				fetch: params.fetch,
 			},
 			params.authStorage,
 		);
@@ -66,7 +70,9 @@ export class KagiProvider extends SearchProvider {
 		return authStorage.hasAuth("kagi");
 	}
 
-	search(params: SearchParams): Promise<SearchResponse> {
+	search(params: SearchParamsWithFetch): Promise<SearchResponse> {
+		const fetchImpl = params.fetch;
+
 		return searchKagi({
 			query: params.query,
 			num_results: params.numSearchResults ?? params.limit,
@@ -74,6 +80,7 @@ export class KagiProvider extends SearchProvider {
 			signal: params.signal,
 			authStorage: params.authStorage,
 			sessionId: params.sessionId,
+			fetch: fetchImpl,
 		});
 	}
 }

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 "@prometheus-ai/ai";
+import { type ApiKey, type AuthStorage, type FetchImpl, withAuth } from "@prometheus-ai/ai";
 import { $env } from "@prometheus-ai/utils";
 
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
@@ -14,6 +14,8 @@ import type { SearchParams } from "./base";
 import { SearchProvider } from "./base";
 import { classifyProviderHttpError, withHardTimeout } from "./utils";
 
+type SearchParamsWithFetch = SearchParams & { fetch?: FetchImpl };
+
 const KIMI_SEARCH_URL = "https://api.kimi.com/coding/v1/search";
 
 const DEFAULT_NUM_RESULTS = 10;
@@ -27,6 +29,7 @@ export interface KimiSearchParams {
 	signal?: AbortSignal;
 	authStorage: AuthStorage;
 	sessionId?: string;
+	fetch?: FetchImpl;
 }
 
 interface KimiSearchResult {
@@ -54,27 +57,40 @@ 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(
 	apiKey: string,
-	params: { query: string; limit: number; includeContent: boolean; signal?: AbortSignal },
+	params: {
+		query: string;
+		limit: number;
+		includeContent: boolean;
+		signal?: AbortSignal;
+		fetch?: FetchImpl;
+	},
 ): Promise<{ response: KimiSearchResponse; requestId?: string }> {
-	const response = await fetch(resolveBaseUrl(), {
+	const fetchImpl = params.fetch ?? fetch;
+	const response = await fetchImpl(resolveBaseUrl(), {
 		method: "POST",
 		headers: {
 			Accept: "application/json",
@@ -108,20 +124,26 @@ 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 'prometheus /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,
+				fetch: params.fetch,
+			}),
+		{ signal: params.signal },
+	);
 	const sources: SearchSource[] = [];
 
 	for (const result of response.search_results ?? []) {
@@ -159,13 +181,16 @@ export class KimiProvider extends SearchProvider {
 		);
 	}
 
-	search(params: SearchParams): Promise<SearchResponse> {
+	search(params: SearchParamsWithFetch): Promise<SearchResponse> {
+		const fetchImpl = params.fetch;
+
 		return searchKimi({
 			query: params.query,
 			num_results: params.numSearchResults ?? params.limit,
 			signal: params.signal,
 			authStorage: params.authStorage,
 			sessionId: params.sessionId,
+			fetch: fetchImpl,
 		});
 	}
 }

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

@@ -1,4 +1,4 @@
-import { type AuthStorage, getEnvApiKey } from "@prometheus-ai/ai";
+import { type ApiKey, type AuthStorage, type FetchImpl, getEnvApiKey, withAuth } from "@prometheus-ai/ai";
 import type { SearchResponse } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import { ParallelApiError, type ParallelSearchResult, type ParallelSearchSource } from "../../parallel";
@@ -112,6 +112,7 @@ async function searchWithAuthStorage(
 	queries: string[],
 	params: {
 		signal?: AbortSignal;
+		fetch?: FetchImpl;
 	},
 	authStorage: AuthStorage,
 	sessionId?: string,
@@ -123,30 +124,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 (params.fetch ?? 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(
@@ -154,6 +166,7 @@ export async function searchParallel(
 		query: string;
 		num_results?: number;
 		signal?: AbortSignal;
+		fetch?: FetchImpl;
 	},
 	authStorage: AuthStorage,
 	sessionId?: string,
@@ -166,6 +179,7 @@ export async function searchParallel(
 			[params.query],
 			{
 				signal: params.signal,
+				fetch: params.fetch,
 			},
 			authStorage,
 			sessionId,
@@ -202,6 +216,7 @@ export class ParallelProvider extends SearchProvider {
 				query: params.query,
 				num_results: params.numSearchResults ?? params.limit,
 				signal: params.signal,
+				fetch: params.fetch,
 			},
 			params.authStorage,
 			params.sessionId,