Merge PR #3811: fix(providers): support Gemini API key web search (@roboomp)

2026-06-29 16:45:46 +02:00
parent 17c4aa0391 530113eb71
commit 6259792bd7
3 changed files with 327 additions and 189 deletions
@@ -15,6 +15,8 @@
 - Fixed large resumed sessions dominated by repeated compaction summaries exhausting memory on Ctrl+C flush, and improved fork session previews that start with developer or assistant turns. ([#3789](https://github.com/can1357/oh-my-pi/issues/3789))

 - Fixed `omp search` applying `providers.webSearch` and `providers.webSearchExclude` before resolving its implicit provider chain. ([#3793](https://github.com/can1357/oh-my-pi/issues/3793))
+- Fixed Gemini web search using only Cloud Code Assist OAuth by allowing standard `GEMINI_API_KEY` developer API credentials to run native Google Search grounding. ([#3810](https://github.com/can1357/oh-my-pi/issues/3810))
+
 - Fixed the bash interceptor blocking `echo` / `printf` redirects to `/dev/null`, `/dev/tty`, `/dev/stdout`, and `/dev/stderr` device sinks while still directing real file writes to the write tool. ([#3763](https://github.com/can1357/oh-my-pi/issues/3763))
 - Fixed the `edit` tool persisting unbounded full-file `oldText` / `newText` snapshots in tool-result `details`, inflating per-turn session JSONL lines (hundreds of KB per edit on large files). `details.oldText`/`details.newText` are now pruned when their combined length exceeds 32 KB; the visible diff, path, line, and diagnostic metadata are preserved, and ACP `diff` content still flows for smaller edits. ([#3786](https://github.com/can1357/oh-my-pi/issues/3786))

@@ -23,6 +23,8 @@ import { SearchProvider } from "./base";
 import { classifyProviderHttpError, withHardTimeout } from "./utils";

 const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
+const DEVELOPER_API_PROVIDER = "google";
+const DEVELOPER_API_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta";
 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;
@@ -80,6 +82,15 @@ interface GeminiAuthSeed {
 	projectId: string;
 }

+interface GeminiSearchResult {
+	answer: string;
+	sources: SearchSource[];
+	citations: SearchCitation[];
+	searchQueries: string[];
+	model: string;
+	usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
+}
+
 /**
 * Walks the configured Gemini OAuth providers in deterministic order and
 * returns the first one that yields a usable access token + projectId via
@@ -128,22 +139,141 @@ interface GeminiGroundingMetadata {
 	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;
+interface GeminiModelResponse {
+	candidates?: Array<{
+		content?: {
+			role: string;
+			parts?: Array<{ text?: string }>;
 		};
-		modelVersion?: string;
+		finishReason?: string;
+		groundingMetadata?: GeminiGroundingMetadata;
+	}>;
+	usageMetadata?: {
+		promptTokenCount?: number;
+		candidatesTokenCount?: number;
+		totalTokenCount?: number;
+	};
+	modelVersion?: string;
+}
+
+interface CloudCodeResponseChunk {
+	response?: GeminiModelResponse;
+}
+
+async function parseGeminiSearchStream(body: ReadableStream<Uint8Array>): Promise<GeminiSearchResult> {
+	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 = 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 & GeminiModelResponse;
+				try {
+					chunk = JSON.parse(jsonStr) as CloudCodeResponseChunk & GeminiModelResponse;
+				} catch {
+					continue;
+				}
+
+				const responseData = chunk.response ?? chunk;
+				const candidate = responseData.candidates?.[0];
+
+				if (candidate?.content?.parts) {
+					for (const part of candidate.content.parts) {
+						if (part.text) {
+							answerParts.push(part.text);
+						}
+					}
+				}
+
+				const groundingMetadata = candidate?.groundingMetadata;
+				if (groundingMetadata) {
+					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,
+									});
+								}
+							}
+						}
+					}
+
+					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,
+									});
+								}
+							}
+						}
+					}
+
+					if (groundingMetadata.webSearchQueries) {
+						for (const q of groundingMetadata.webSearchQueries) {
+							if (!searchQueries.includes(q)) {
+								searchQueries.push(q);
+							}
+						}
+					}
+				}
+
+				if (responseData.usageMetadata) {
+					usage = {
+						inputTokens: responseData.usageMetadata.promptTokenCount ?? 0,
+						outputTokens: responseData.usageMetadata.candidatesTokenCount ?? 0,
+						totalTokens: responseData.usageMetadata.totalTokenCount ?? 0,
+					};
+				}
+
+				if (responseData.modelVersion) {
+					model = responseData.modelVersion;
+				}
+			}
+		}
+	} finally {
+		reader.releaseLock();
+	}
+
+	return {
+		answer: answerParts.join(""),
+		sources,
+		citations,
+		searchQueries,
+		model,
+		usage,
 	};
 }

@@ -165,14 +295,7 @@ async function callGeminiSearch(
 	fetchImpl: FetchImpl | undefined,
 	signal: AbortSignal | undefined,
 	mode?: "auto" | "production" | "sandbox",
-): Promise<{
-	answer: string;
-	sources: SearchSource[];
-	citations: SearchCitation[];
-	searchQueries: string[];
-	model: string;
-	usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
-}> {
+): Promise<GeminiSearchResult> {
 	let endpoints: string[];
 	if (auth.isAntigravity) {
 		const m = mode ?? "auto";
@@ -287,13 +410,75 @@ async function callGeminiSearch(
 		throw new SearchProviderError("gemini", `Gemini Cloud Code API error (${status}): ${errorText}`, status);
 	}

+	if (!response.body) {
+		throw new SearchProviderError("gemini", "Gemini API returned no response body", 500);
+	}
+
+	return parseGeminiSearchStream(response.body);
+}
+
+async function callGeminiDeveloperSearch(
+	apiKey: string,
+	query: string,
+	systemPrompt: string | undefined,
+	maxOutputTokens: number | undefined,
+	temperature: number | undefined,
+	toolParams: GeminiToolParams,
+	fetchImpl: FetchImpl | undefined,
+	signal: AbortSignal | undefined,
+): Promise<GeminiSearchResult> {
+	const normalizedSystemPrompt = systemPrompt?.toWellFormed();
+	const requestBody: Record<string, unknown> = {
+		contents: [
+			{
+				role: "user",
+				parts: [{ text: query }],
+			},
+		],
+		tools: buildGeminiRequestTools(toolParams),
+		...(normalizedSystemPrompt && {
+			systemInstruction: {
+				parts: [{ text: normalizedSystemPrompt }],
+			},
+		}),
+	};
+
+	if (maxOutputTokens !== undefined || temperature !== undefined) {
+		const generationConfig: Record<string, number> = {};
+		if (maxOutputTokens !== undefined) {
+			generationConfig.maxOutputTokens = maxOutputTokens;
+		}
+		if (temperature !== undefined) {
+			generationConfig.temperature = temperature;
+		}
+		requestBody.generationConfig = generationConfig;
+	}
+
+	const response = await fetchWithRetry(
+		() => `${DEVELOPER_API_ENDPOINT}/models/${DEFAULT_MODEL}:streamGenerateContent?alt=sse`,
+		{
+			method: "POST",
+			headers: {
+				"x-goog-api-key": apiKey,
+				"Content-Type": "application/json",
+				Accept: "text/event-stream",
+			},
+			body: JSON.stringify(requestBody),
+			signal: withHardTimeout(signal),
+			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}`,
+			`Gemini Developer API error (${response.status}): ${errorText}`,
 			response.status,
 		);
 	}
@@ -302,130 +487,7 @@ async function callGeminiSearch(
 		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,
-	};
+	return parseGeminiSearchStream(response.body);
 }

 /**
@@ -433,43 +495,64 @@ async function callGeminiSearch(
 */
 export async function searchGemini(params: GeminiSearchParams): Promise<SearchResponse> {
 	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.",
+	let result: GeminiSearchResult;
+
+	if (seed) {
+		const isAntigravity = seed.provider === "google-antigravity";
+		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
+				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,
+					params.antigravityEndpointMode,
+				),
+			{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
+		);
+	} else {
+		const apiKey = await params.authStorage.getApiKey(DEVELOPER_API_PROVIDER, params.sessionId, {
+			signal: params.signal,
+		});
+		if (!apiKey) {
+			throw new Error(
+				"No Gemini credentials found. Set GEMINI_API_KEY, configure an API key for provider \"google\", or login with 'omp /login google-gemini-cli' / 'omp /login google-antigravity' to enable Gemini web search.",
+			);
+		}
+		result = await callGeminiDeveloperSearch(
+			apiKey,
+			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
-			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,
-				params.antigravityEndpointMode,
-			),
-		{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
-	);
-
 	let sources = result.sources;

 	if (params.num_results && sources.length > params.num_results) {
@@ -494,9 +577,12 @@ export class GeminiProvider extends SearchProvider {

 	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);
+		// the provider-chain probe. `searchGemini` refreshes OAuth lazily on the
+		// actual request and resolves developer API keys through AuthStorage.
+		return (
+			GEMINI_PROVIDERS.some((provider: GeminiProviderId) => authStorage.hasOAuth(provider)) ||
+			authStorage.hasAuth(DEVELOPER_API_PROVIDER)
+		);
 	}

 	search(params: SearchParams): Promise<SearchResponse> {
@@ -1,12 +1,16 @@
 import { afterEach, describe, expect, it } from "bun:test";
 import type { AuthStorage } from "@oh-my-pi/pi-ai";
 import type { FetchImpl } from "@oh-my-pi/pi-ai/types";
-import { searchGemini } from "@oh-my-pi/pi-coding-agent/web/search/providers/gemini";
+import { GeminiProvider, searchGemini } from "@oh-my-pi/pi-coding-agent/web/search/providers/gemini";

 const SSE_RESPONSE =
 	'data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"Gemini answer"}]}}],"modelVersion":"gemini-2.5-flash"}}\n\n';
+const DEVELOPER_SSE_RESPONSE =
+	'data: {"candidates":[{"content":{"role":"model","parts":[{"text":"Developer answer"}]},"groundingMetadata":{"webSearchQueries":["latest Bun version"],"groundingChunks":[{"web":{"uri":"https://bun.sh","title":"Bun"}}],"groundingSupports":[{"segment":{"text":"Developer answer"},"groundingChunkIndices":[0]}]}}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":4,"totalTokenCount":7},"modelVersion":"gemini-2.5-flash"}\n\n';

 type CapturedRequest = {
+	url: string;
+	headers: Record<string, string>;
 	body: Record<string, unknown> | null;
 };

@@ -25,14 +29,32 @@ describe("searchGemini tools serialization", () => {
 		},
 	} as unknown as AuthStorage;

-	function mockGeminiFetch(): FetchImpl {
+	const apiKeyAuthStorage = {
+		async getOAuthAccess() {
+			return undefined;
+		},
+		hasOAuth() {
+			return false;
+		},
+		hasAuth(provider: string) {
+			return provider === "google";
+		},
+		async getApiKey(provider: string) {
+			return provider === "google" ? "test-gemini-api-key" : undefined;
+		},
+	} as unknown as AuthStorage;
+
+	function mockGeminiFetch(responseText = SSE_RESPONSE): FetchImpl {
 		capturedRequest = null;
-		return (_url, init) => {
+		return (url, init) => {
+			const headers = new Headers(init?.headers);
 			capturedRequest = {
+				url: String(url),
+				headers: Object.fromEntries(headers.entries()),
 				body: init?.body ? (JSON.parse(init.body as string) as Record<string, unknown>) : null,
 			};
 			return Promise.resolve(
-				new Response(SSE_RESPONSE, {
+				new Response(responseText, {
 					status: 200,
 					headers: { "Content-Type": "text/event-stream" },
 				}),
@@ -52,6 +74,34 @@ describe("searchGemini tools serialization", () => {
 		} as const;
 	}

+	it("treats a standard Google developer API key as available", () => {
+		const provider = new GeminiProvider();
+		expect(provider.isAvailable(apiKeyAuthStorage)).toBe(true);
+	});
+
+	it("routes API key auth through the developer API with Google Search grounding", async () => {
+		const fetchMock = mockGeminiFetch(DEVELOPER_SSE_RESPONSE);
+		const response = await searchGemini({
+			...makeParams("developer api"),
+			authStorage: apiKeyAuthStorage,
+			fetch: fetchMock,
+		});
+
+		expect(capturedRequest).not.toBeNull();
+		expect(capturedRequest?.url).toBe(
+			"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
+		);
+		expect(capturedRequest?.headers["x-goog-api-key"]).toBe("test-gemini-api-key");
+		expect(capturedRequest?.body).toMatchObject({
+			tools: [{ googleSearch: {} }],
+		});
+		expect(response).toMatchObject({
+			answer: "Developer answer",
+			sources: [{ title: "Bun", url: "https://bun.sh" }],
+			searchQueries: ["latest Bun version"],
+			usage: { inputTokens: 3, outputTokens: 4, totalTokens: 7 },
+		});
+	});
 	it("sends default googleSearch tool when no passthrough payloads are provided", async () => {
 		const fetchMock = mockGeminiFetch();
 		await searchGemini({ ...makeParams("default tools"), fetch: fetchMock });