From 530113eb7169af329c319db2178b6c996ae46426 Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 29 Jun 2026 13:51:42 +0000 Subject: [PATCH] fix(providers): supported gemini api key search Enabled the Gemini web search provider to use standard Google developer API credentials when Cloud Code Assist OAuth is absent. Added developer API request coverage for native Google Search grounding and preserved existing OAuth request serialization. Fixes #3810 --- packages/coding-agent/CHANGELOG.md | 2 + .../src/web/search/providers/gemini.ts | 456 +++++++++++------- .../test/tools/web-search-gemini.test.ts | 58 ++- 3 files changed, 327 insertions(+), 189 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d231182bd..cc8fc7247 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,8 @@ ### Fixed +- 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)) ## [16.2.5] - 2026-06-28 diff --git a/packages/coding-agent/src/web/search/providers/gemini.ts b/packages/coding-agent/src/web/search/providers/gemini.ts index 6e7b25eba..cf0824fe4 100644 --- a/packages/coding-agent/src/web/search/providers/gemini.ts +++ b/packages/coding-agent/src/web/search/providers/gemini.ts @@ -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): Promise { + const answerParts: string[] = []; + const sources: SearchSource[] = []; + const citations: SearchCitation[] = []; + const searchQueries: string[] = []; + const seenUrls = new Set(); + 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 { 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 { + const normalizedSystemPrompt = systemPrompt?.toWellFormed(); + const requestBody: Record = { + contents: [ + { + role: "user", + parts: [{ text: query }], + }, + ], + tools: buildGeminiRequestTools(toolParams), + ...(normalizedSystemPrompt && { + systemInstruction: { + parts: [{ text: normalizedSystemPrompt }], + }, + }), + }; + + if (maxOutputTokens !== undefined || temperature !== undefined) { + const generationConfig: Record = {}; + 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(); - 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 { 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 { diff --git a/packages/coding-agent/test/tools/web-search-gemini.test.ts b/packages/coding-agent/test/tools/web-search-gemini.test.ts index 978a56e89..56638a7d6 100644 --- a/packages/coding-agent/test/tools/web-search-gemini.test.ts +++ b/packages/coding-agent/test/tools/web-search-gemini.test.ts @@ -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; body: Record | 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) : 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 });