@wahack/pi-coding-agent 15.11.0

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

@@ -0,0 +1,460 @@
+/**
+ * Google Gemini Web Search Provider
+ *
+ * Uses Gemini's Google Search grounding via Cloud Code Assist API.
+ * Auth is resolved through `AuthStorage.getOAuthAccess(...)` for both
+ * `google-gemini-cli` (stable prod) and `google-antigravity` (daily sandbox)
+ * — the broker is the sole refresh authority, so this module never opens a
+ * 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 {
+	ANTIGRAVITY_SYSTEM_INSTRUCTION,
+	getAntigravityUserAgent,
+	getGeminiCliHeaders,
+} from "@oh-my-pi/pi-catalog/wire/gemini-headers";
+import { fetchWithRetry } from "@oh-my-pi/pi-utils";
+
+import type { SearchCitation, SearchResponse, SearchSource } from "../../../web/search/types";
+import { SearchProviderError } from "../../../web/search/types";
+import type { SearchParams } from "./base";
+import { SearchProvider } from "./base";
+import { classifyProviderHttpError, withHardTimeout } from "./utils";
+
+const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
+const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com";
+const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
+const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT] as const;
+const DEFAULT_MODEL = "gemini-2.5-flash";
+const MAX_RETRIES = 3;
+const BASE_DELAY_MS = 1000;
+const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000;
+
+const GEMINI_PROVIDERS = ["google-gemini-cli", "google-antigravity"] as const;
+type GeminiProviderId = (typeof GEMINI_PROVIDERS)[number];
+
+interface GeminiToolParams {
+	google_search?: Record<string, unknown>;
+	code_execution?: Record<string, unknown>;
+	url_context?: Record<string, unknown>;
+}
+
+export interface GeminiSearchParams extends GeminiToolParams {
+	query: string;
+	system_prompt?: string;
+	num_results?: number;
+	/** Maximum output tokens. */
+	max_output_tokens?: number;
+	/** Sampling temperature (0–1). Lower = more focused/factual. */
+	temperature?: number;
+	signal?: AbortSignal;
+	authStorage: AuthStorage;
+	sessionId?: string;
+	fetch?: FetchImpl;
+}
+
+export function buildGeminiRequestTools(params: GeminiToolParams): Array<Record<string, Record<string, unknown>>> {
+	const tools: Array<Record<string, Record<string, unknown>>> = [{ googleSearch: params.google_search ?? {} }];
+	if (params.code_execution !== undefined) {
+		tools.push({ codeExecution: params.code_execution });
+	}
+	if (params.url_context !== undefined) {
+		tools.push({ urlContext: params.url_context });
+	}
+	return tools;
+}
+
+/** Resolved auth for a Gemini API request. */
+interface GeminiAuth {
+	accessToken: string;
+	projectId: string;
+	isAntigravity: boolean;
+}
+
+/**
+ * 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.
+ */
+export async function findGeminiAuth(
+	authStorage: AuthStorage,
+	sessionId: string | undefined,
+	signal: AbortSignal | undefined,
+): Promise<GeminiAuth | 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 null;
+}
+
+function hasGeminiOAuth(authStorage: AuthStorage): boolean {
+	return GEMINI_PROVIDERS.some((provider: GeminiProviderId) => authStorage.hasOAuth(provider));
+}
+
+/** Cloud Code Assist API response types */
+interface GeminiGroundingChunk {
+	web?: {
+		uri?: string;
+		title?: string;
+	};
+}
+
+interface GeminiGroundingSupport {
+	segment?: {
+		startIndex?: number;
+		endIndex?: number;
+		text?: string;
+	};
+	groundingChunkIndices?: number[];
+	confidenceScores?: number[];
+}
+
+interface GeminiGroundingMetadata {
+	groundingChunks?: GeminiGroundingChunk[];
+	groundingSupports?: GeminiGroundingSupport[];
+	webSearchQueries?: string[];
+}
+
+interface CloudCodeResponseChunk {
+	response?: {
+		candidates?: Array<{
+			content?: {
+				role: string;
+				parts?: Array<{ text?: string }>;
+			};
+			finishReason?: string;
+			groundingMetadata?: GeminiGroundingMetadata;
+		}>;
+		usageMetadata?: {
+			promptTokenCount?: number;
+			candidatesTokenCount?: number;
+			totalTokenCount?: number;
+		};
+		modelVersion?: string;
+	};
+}
+
+/**
+ * Calls the Cloud Code Assist API with Google Search grounding enabled.
+ *
+ * If a request returns a refreshable auth failure (401/403/auth-flavoured 400),
+ * we ask AuthStorage to invalidate + refresh the credential and retry once.
+ * Provider-direct refresh helpers are intentionally not used: AuthStorage owns
+ * the single-flight refresh and broker round-trip.
+ */
+async function callGeminiSearch(
+	auth: GeminiAuth,
+	query: string,
+	systemPrompt: string | undefined,
+	maxOutputTokens: number | undefined,
+	temperature: number | undefined,
+	toolParams: GeminiToolParams,
+	fetchImpl: FetchImpl | undefined,
+	signal: AbortSignal | undefined,
+): Promise<{
+	answer: string;
+	sources: SearchSource[];
+	citations: SearchCitation[];
+	searchQueries: string[];
+	model: string;
+	usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
+}> {
+	const endpoints = auth.isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT];
+	const headers = auth.isAntigravity ? { "User-Agent": getAntigravityUserAgent() } : getGeminiCliHeaders();
+
+	const requestMetadata = auth.isAntigravity
+		? {
+				requestType: "agent",
+				userAgent: "antigravity",
+				requestId: `agent-${crypto.randomUUID()}`,
+			}
+		: {
+				userAgent: "pi-coding-agent",
+				requestId: `pi-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
+			};
+
+	const normalizedSystemPrompt = systemPrompt?.toWellFormed();
+	const systemInstructionParts: Array<{ text: string }> = [
+		...(auth.isAntigravity
+			? [
+					{ text: ANTIGRAVITY_SYSTEM_INSTRUCTION },
+					{ text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` },
+				]
+			: []),
+		...(normalizedSystemPrompt ? [{ text: normalizedSystemPrompt }] : []),
+	];
+
+	const requestBody: Record<string, unknown> = {
+		project: auth.projectId,
+		model: DEFAULT_MODEL,
+		request: {
+			contents: [
+				{
+					role: "user",
+					parts: [{ text: query }],
+				},
+			],
+			tools: buildGeminiRequestTools(toolParams),
+			...(systemInstructionParts.length > 0 && {
+				systemInstruction: {
+					...(auth.isAntigravity ? { role: "user" } : {}),
+					parts: systemInstructionParts,
+				},
+			}),
+		},
+		...requestMetadata,
+	};
+
+	if (maxOutputTokens !== undefined || temperature !== undefined) {
+		const generationConfig: Record<string, number> = {};
+		if (maxOutputTokens !== undefined) {
+			generationConfig.maxOutputTokens = maxOutputTokens;
+		}
+		if (temperature !== undefined) {
+			generationConfig.temperature = temperature;
+		}
+		(requestBody.request as Record<string, unknown>).generationConfig = generationConfig;
+	}
+	const buildInit = (): RequestInit => ({
+		method: "POST",
+		headers: {
+			Authorization: `Bearer ${auth.accessToken}`,
+			"Content-Type": "application/json",
+			Accept: "text/event-stream",
+			...headers,
+		},
+		body: JSON.stringify(requestBody),
+		signal: withHardTimeout(signal),
+	});
+	const urlFor = (attempt: number) =>
+		`${endpoints[Math.min(attempt, endpoints.length - 1)]}/v1internal:streamGenerateContent?alt=sse`;
+
+	const response = await fetchWithRetry(urlFor, {
+		...buildInit(),
+		fetch: fetchImpl,
+		maxAttempts: MAX_RETRIES + 1,
+		defaultDelayMs: attempt => BASE_DELAY_MS * 2 ** attempt,
+		maxDelayMs: RATE_LIMIT_BUDGET_MS,
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		const classified = classifyProviderHttpError("gemini", response.status, errorText);
+		if (classified) throw classified;
+		throw new SearchProviderError(
+			"gemini",
+			`Gemini Cloud Code API error (${response.status}): ${errorText}`,
+			response.status,
+		);
+	}
+
+	if (!response.body) {
+		throw new SearchProviderError("gemini", "Gemini API returned no response body", 500);
+	}
+
+	// Parse SSE stream
+	const answerParts: string[] = [];
+	const sources: SearchSource[] = [];
+	const citations: SearchCitation[] = [];
+	const searchQueries: string[] = [];
+	const seenUrls = new Set<string>();
+	let model = DEFAULT_MODEL;
+	let usage: { inputTokens: number; outputTokens: number; totalTokens: number } | undefined;
+
+	const reader = response.body.getReader();
+	const decoder = new TextDecoder();
+	let buffer = "";
+
+	try {
+		while (true) {
+			const { done, value } = await reader.read();
+			if (done) break;
+
+			buffer += decoder.decode(value, { stream: true });
+			const lines = buffer.split("\n");
+			buffer = lines.pop() || "";
+
+			for (const line of lines) {
+				if (!line.startsWith("data:")) continue;
+
+				const jsonStr = line.slice(5).trim();
+				if (!jsonStr) continue;
+
+				let chunk: CloudCodeResponseChunk;
+				try {
+					chunk = JSON.parse(jsonStr) as CloudCodeResponseChunk;
+				} catch {
+					continue;
+				}
+
+				const responseData = chunk.response;
+				if (!responseData) continue;
+
+				const candidate = responseData.candidates?.[0];
+
+				// Extract text content
+				if (candidate?.content?.parts) {
+					for (const part of candidate.content.parts) {
+						if (part.text) {
+							answerParts.push(part.text);
+						}
+					}
+				}
+
+				// Extract grounding metadata
+				const groundingMetadata = candidate?.groundingMetadata;
+				if (groundingMetadata) {
+					// Extract sources from grounding chunks
+					if (groundingMetadata.groundingChunks) {
+						for (const grChunk of groundingMetadata.groundingChunks) {
+							if (grChunk.web?.uri) {
+								const sourceUrl = grChunk.web.uri;
+								if (!seenUrls.has(sourceUrl)) {
+									seenUrls.add(sourceUrl);
+									sources.push({
+										title: grChunk.web.title ?? sourceUrl,
+										url: sourceUrl,
+									});
+								}
+							}
+						}
+					}
+
+					// Extract citations from grounding supports
+					if (groundingMetadata.groundingSupports && groundingMetadata.groundingChunks) {
+						for (const support of groundingMetadata.groundingSupports) {
+							const citedText = support.segment?.text;
+							const chunkIndices = support.groundingChunkIndices ?? [];
+
+							for (const idx of chunkIndices) {
+								const grChunk = groundingMetadata.groundingChunks[idx];
+								if (grChunk?.web?.uri) {
+									citations.push({
+										url: grChunk.web.uri,
+										title: grChunk.web.title ?? grChunk.web.uri,
+										citedText,
+									});
+								}
+							}
+						}
+					}
+
+					// Extract search queries
+					if (groundingMetadata.webSearchQueries) {
+						for (const q of groundingMetadata.webSearchQueries) {
+							if (!searchQueries.includes(q)) {
+								searchQueries.push(q);
+							}
+						}
+					}
+				}
+
+				// Extract usage metadata
+				if (responseData.usageMetadata) {
+					usage = {
+						inputTokens: responseData.usageMetadata.promptTokenCount ?? 0,
+						outputTokens: responseData.usageMetadata.candidatesTokenCount ?? 0,
+						totalTokens: responseData.usageMetadata.totalTokenCount ?? 0,
+					};
+				}
+
+				// Extract model version
+				if (responseData.modelVersion) {
+					model = responseData.modelVersion;
+				}
+			}
+		}
+	} finally {
+		reader.releaseLock();
+	}
+
+	return {
+		answer: answerParts.join(""),
+		sources,
+		citations,
+		searchQueries,
+		model,
+		usage,
+	};
+}
+
+/**
+ * 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) {
+		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,
+	);
+
+	let sources = result.sources;
+
+	if (params.num_results && sources.length > params.num_results) {
+		sources = sources.slice(0, params.num_results);
+	}
+
+	return {
+		provider: "gemini",
+		answer: result.answer || undefined,
+		sources,
+		citations: result.citations.length > 0 ? result.citations : undefined,
+		searchQueries: result.searchQueries.length > 0 ? result.searchQueries : undefined,
+		usage: result.usage,
+		model: result.model,
+	};
+}
+
+/** Search provider for Google Gemini web search. */
+export class GeminiProvider extends SearchProvider {
+	readonly id = "gemini";
+	readonly label = "Gemini";
+
+	isAvailable(authStorage: AuthStorage): boolean {
+		// Cheap, in-memory check — avoids driving the refresh pipeline during
+		// the provider-chain probe. `searchGemini` calls `getOAuthAccess` which
+		// will refresh lazily on the actual request.
+		return hasGeminiOAuth(authStorage);
+	}
+
+	search(params: SearchParams): Promise<SearchResponse> {
+		return searchGemini({
+			query: params.query,
+			system_prompt: params.systemPrompt,
+			num_results: params.numSearchResults ?? params.limit,
+			max_output_tokens: params.maxOutputTokens,
+			temperature: params.temperature,
+			google_search: params.googleSearch,
+			code_execution: params.codeExecution,
+			url_context: params.urlContext,
+			signal: params.signal,
+			authStorage: params.authStorage,
+			sessionId: params.sessionId,
+			fetch: params.fetch,
+		});
+	}
+}

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

@@ -0,0 +1,111 @@
+/**
+ * Jina Reader Web Search Provider
+ *
+ * Uses the Jina Reader `s.jina.ai` endpoint to fetch search results with
+ * cleaned content.
+ */
+
+import { type AuthStorage, type FetchImpl, getEnvApiKey } 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";
+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 {
+	title?: string | null;
+	url?: string | null;
+	content?: string | null;
+}
+
+type JinaSearchResponse = JinaSearchResult[];
+
+/** Find JINA_API_KEY from environment or .env files. */
+export function findApiKey(): string | null {
+	return getEnvApiKey("jina") ?? null;
+}
+
+/** Call Jina Reader search API. */
+async function callJinaSearch(
+	apiKey: string,
+	query: string,
+	signal?: AbortSignal,
+	fetchImpl: FetchImpl = fetch,
+): Promise<JinaSearchResponse> {
+	const requestUrl = `${JINA_SEARCH_URL}/${encodeURIComponent(query)}`;
+	const response = await fetchImpl(requestUrl, {
+		headers: {
+			Accept: "application/json",
+			Authorization: `Bearer ${apiKey}`,
+		},
+		signal: withHardTimeout(signal),
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		const classified = classifyProviderHttpError("jina", response.status, errorText);
+		if (classified) throw classified;
+		throw new SearchProviderError("jina", `Jina API error (${response.status}): ${errorText}`, response.status);
+	}
+
+	const payload = (await response.json()) as { data?: JinaSearchResponse } | null;
+	return Array.isArray(payload?.data) ? payload.data : [];
+}
+
+/** Execute Jina web search. */
+export async function searchJina(params: JinaSearchParams): Promise<SearchResponse> {
+	const apiKey = findApiKey();
+	if (!apiKey) {
+		throw new Error("JINA_API_KEY not found. Set it in environment or .env file.");
+	}
+
+	const response = await callJinaSearch(apiKey, params.query, params.signal, params.fetch);
+	const sources: SearchSource[] = [];
+
+	for (const result of response) {
+		if (!result?.url) continue;
+		sources.push({
+			title: result.title ?? result.url,
+			url: result.url,
+			snippet: result.content ?? undefined,
+		});
+	}
+
+	const limitedSources = params.num_results ? sources.slice(0, params.num_results) : sources;
+
+	return {
+		provider: "jina",
+		sources: limitedSources,
+	};
+}
+
+/** Search provider for Jina Reader. */
+export class JinaProvider extends SearchProvider {
+	readonly id = "jina";
+	readonly label = "Jina";
+
+	isAvailable(_authStorage: AuthStorage): boolean {
+		return !!findApiKey();
+	}
+
+	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

@@ -0,0 +1,86 @@
+/**
+ * Kagi Web Search Provider
+ *
+ * Thin wrapper that adapts shared Kagi API utilities to SearchResponse shape.
+ */
+import type { AuthStorage, FetchImpl } from "@oh-my-pi/pi-ai";
+import type { SearchResponse } from "../../../web/search/types";
+import { SearchProviderError } from "../../../web/search/types";
+import { KagiApiError, searchWithKagi } from "../../kagi";
+import { clampNumResults } from "../utils";
+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;
+
+/** Execute Kagi web search. */
+export async function searchKagi(params: {
+	query: string;
+	num_results?: number;
+	recency?: SearchParams["recency"];
+	signal?: AbortSignal;
+	authStorage: AuthStorage;
+	sessionId?: string;
+	fetch?: FetchImpl;
+}): Promise<SearchResponse> {
+	const numResults = clampNumResults(params.num_results, DEFAULT_NUM_RESULTS, MAX_NUM_RESULTS);
+
+	try {
+		const result = await searchWithKagi(
+			params.query,
+			{
+				limit: numResults,
+				recency: params.recency,
+				sessionId: params.sessionId,
+				signal: params.signal,
+				fetch: params.fetch,
+			},
+			params.authStorage,
+		);
+
+		return {
+			provider: "kagi",
+			sources: toSearchSources(result.sources, numResults),
+			relatedQuestions: result.relatedQuestions.length > 0 ? result.relatedQuestions : undefined,
+			requestId: result.requestId,
+			answer: result.answer,
+		};
+	} catch (err) {
+		if (err instanceof KagiApiError) {
+			if (typeof err.statusCode === "number") {
+				const classified = classifyProviderHttpError("kagi", err.statusCode, err.message);
+				if (classified) throw classified;
+			}
+			throw new SearchProviderError("kagi", err.message, err.statusCode);
+		}
+		throw err;
+	}
+}
+
+/** Search provider for Kagi web search. */
+export class KagiProvider extends SearchProvider {
+	readonly id = "kagi";
+	readonly label = "Kagi";
+
+	isAvailable(authStorage: AuthStorage): boolean {
+		return authStorage.hasAuth("kagi");
+	}
+
+	search(params: SearchParamsWithFetch): Promise<SearchResponse> {
+		const fetchImpl = params.fetch;
+
+		return searchKagi({
+			query: params.query,
+			num_results: params.numSearchResults ?? params.limit,
+			recency: params.recency,
+			signal: params.signal,
+			authStorage: params.authStorage,
+			sessionId: params.sessionId,
+			fetch: fetchImpl,
+		});
+	}
+}