diff --git a/src/resources/extensions/google-search/index.ts b/src/resources/extensions/google-search/index.ts index 13b989ea2..f06eabd93 100644 --- a/src/resources/extensions/google-search/index.ts +++ b/src/resources/extensions/google-search/index.ts @@ -68,6 +68,101 @@ async function getClient(): Promise { return client; } +/** + * Perform a search using OAuth credentials via the Cloud Code Assist API. + * This is used as a fallback when GEMINI_API_KEY is not set. + */ +async function searchWithOAuth( + query: string, + accessToken: string, + projectId: string, + signal?: AbortSignal, +): Promise { + const model = process.env.GEMINI_SEARCH_MODEL || "gemini-2.5-flash"; + const url = `https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent`; + + const GEMINI_CLI_HEADERS = { + ideType: "IDE_UNSPECIFIED", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + }; + + const executeFetch = async (retries = 3): Promise => { + const response = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1", + "X-Goog-Api-Client": "gl-node/22.17.0", + "Client-Metadata": JSON.stringify(GEMINI_CLI_HEADERS), + }, + body: JSON.stringify({ + project: projectId, + model, + request: { + contents: [{ parts: [{ text: query }] }], + tools: [{ googleSearch: {} }], + }, + }), + signal, + }); + + if (!response.ok && retries > 0 && (response.status === 429 || response.status >= 500)) { + await new Promise((resolve) => setTimeout(resolve, 1000 * (4 - retries))); + return executeFetch(retries - 1); + } + + return response; + }; + + const response = await executeFetch(); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`Cloud Code Assist API error (${response.status}): ${errorText}`); + } + + // Note: streamGenerateContent returns SSE; for now, we consume all chunks. + // For simplicity and to match the previous structure, we'll read to end. + const text = await response.text(); + const jsonLines = text.split("\n") + .filter(l => l.startsWith("data:")) + .map(l => l.slice(5).trim()) + .filter(l => l.length > 0); + + let data; + if (jsonLines.length > 0) { + // Aggregate chunks if needed, but for now we take the last chunk or assume it's one + data = JSON.parse(jsonLines[jsonLines.length - 1]); + } else { + data = JSON.parse(text); + } const candidate = data.response?.candidates?.[0]; + const answer = candidate?.content?.parts?.find((p: any) => p.text)?.text ?? ""; + const grounding = candidate?.groundingMetadata; + + const sources: SearchSource[] = []; + const seenTitles = new Set(); + if (grounding?.groundingChunks) { + for (const chunk of grounding.groundingChunks) { + if (chunk.web) { + const title = chunk.web.title ?? "Untitled"; + if (seenTitles.has(title)) continue; + seenTitles.add(title); + const domain = chunk.web.domain ?? title; + sources.push({ + title, + uri: chunk.web.uri ?? "", + domain, + }); + } + } + } + + const searchQueries = grounding?.webSearchQueries ?? []; + return { answer, sources, searchQueries, cached: false }; +} + // ── In-session cache ───────────────────────────────────────────────────────── const resultCache = new Map(); @@ -87,7 +182,7 @@ export default function (pi: ExtensionAPI) { "Returns an AI-synthesized answer grounded in Google Search results, plus source URLs. " + "Use this when you need current information from the web: recent events, documentation, " + "product details, technical references, news, etc. " + - "Requires GEMINI_API_KEY. Alternative to Brave-based search tools for users with Google Cloud credits.", + "Requires GEMINI_API_KEY or Google login. Alternative to Brave-based search tools.", promptSnippet: "Search the web via Google Search to get current information with sources", promptGuidelines: [ "Use google_search when you need up-to-date web information that isn't in your training data.", @@ -109,17 +204,33 @@ export default function (pi: ExtensionAPI) { ), }), - async execute(_toolCallId, params, signal, _onUpdate, _ctx) { + async execute(_toolCallId, params, signal, _onUpdate, ctx) { const startTime = Date.now(); const maxSources = Math.min(Math.max(params.maxSources ?? 5, 1), 10); - // Check for API key + // Check for credentials + let oauthToken: string | undefined; + let projectId: string | undefined; + if (!process.env.GEMINI_API_KEY) { + const oauthRaw = await ctx.modelRegistry.getApiKeyForProvider("google-gemini-cli"); + if (oauthRaw) { + try { + const parsed = JSON.parse(oauthRaw); + oauthToken = parsed.token; + projectId = parsed.projectId; + } catch { + // Fall through to error + } + } + } + + if (!process.env.GEMINI_API_KEY && (!oauthToken || !projectId)) { return { content: [ { type: "text", - text: "Error: GEMINI_API_KEY is not set. Please set this environment variable to use Google Search.\n\nExample: export GEMINI_API_KEY=your_key", + text: "Error: No authentication found for Google Search. Please set GEMINI_API_KEY or log in via Google.\n\nExample: export GEMINI_API_KEY=your_key or use /login google", }, ], isError: true, @@ -128,7 +239,7 @@ export default function (pi: ExtensionAPI) { sourceCount: 0, cached: false, durationMs: Date.now() - startTime, - error: "auth_error: GEMINI_API_KEY not set", + error: "auth_error: No credentials set", } as SearchDetails, }; } @@ -152,49 +263,52 @@ export default function (pi: ExtensionAPI) { // Call Gemini with Google Search grounding let result: SearchResult; try { - const ai = await getClient(); - const response = await ai.models.generateContent({ - model: process.env.GEMINI_SEARCH_MODEL || "gemini-2.5-flash", - contents: params.query, - config: { - tools: [{ googleSearch: {} }], - abortSignal: signal, - }, - }); + if (process.env.GEMINI_API_KEY) { + const ai = await getClient(); + const response = await ai.models.generateContent({ + model: process.env.GEMINI_SEARCH_MODEL || "gemini-2.5-flash", + contents: params.query, + config: { + tools: [{ googleSearch: {} }], + abortSignal: signal, + }, + }); - // Extract answer text - const answer = response.text ?? ""; + // Extract answer text + const answer = response.text ?? ""; - // Extract grounding metadata - const candidate = response.candidates?.[0]; - const grounding = candidate?.groundingMetadata; + // Extract grounding metadata + const candidate = response.candidates?.[0]; + const grounding = candidate?.groundingMetadata; - // Parse sources from grounding chunks - const sources: SearchSource[] = []; - const seenTitles = new Set(); - if (grounding?.groundingChunks) { - for (const chunk of grounding.groundingChunks) { - if (chunk.web) { - const title = chunk.web.title ?? "Untitled"; - // Dedupe by title since URIs are redirect URLs that differ per call - if (seenTitles.has(title)) continue; - seenTitles.add(title); - // domain field is not available via Gemini API, use title as fallback - // (title is typically the domain name, e.g. "wikipedia.org") - const domain = chunk.web.domain ?? title; - sources.push({ - title, - uri: chunk.web.uri ?? "", - domain, - }); + // Parse sources from grounding chunks + const sources: SearchSource[] = []; + const seenTitles = new Set(); + if (grounding?.groundingChunks) { + for (const chunk of grounding.groundingChunks) { + if (chunk.web) { + const title = chunk.web.title ?? "Untitled"; + // Dedupe by title since URIs are redirect URLs that differ per call + if (seenTitles.has(title)) continue; + seenTitles.add(title); + // domain field is not available via Gemini API, use title as fallback + // (title is typically the domain name, e.g. "wikipedia.org") + const domain = chunk.web.domain ?? title; + sources.push({ + title, + uri: chunk.web.uri ?? "", + domain, + }); + } } } + + // Extract search queries Gemini actually performed + const searchQueries = grounding?.webSearchQueries ?? []; + result = { answer, sources, searchQueries, cached: false }; + } else { + result = await searchWithOAuth(params.query, oauthToken!, projectId!, signal); } - - // Extract search queries Gemini actually performed - const searchQueries = grounding?.webSearchQueries ?? []; - - result = { answer, sources, searchQueries, cached: false }; } catch (err: unknown) { const msg = err instanceof Error ? err.message : String(err); @@ -287,9 +401,12 @@ export default function (pi: ExtensionAPI) { // ── Startup notification ───────────────────────────────────────────────── pi.on("session_start", async (_event, ctx) => { - if (!process.env.GEMINI_API_KEY) { + if (process.env.GEMINI_API_KEY) return; + + const hasOAuth = await ctx.modelRegistry.authStorage.hasAuth("google-gemini-cli"); + if (!hasOAuth) { ctx.ui.notify( - "Google Search: No GEMINI_API_KEY set. The google_search tool will not work until this is configured.", + "Google Search: No authentication set. Log in via Google or set GEMINI_API_KEY to use google_search.", "warning", ); } diff --git a/src/resources/extensions/gsd/tests/dist-redirect.mjs b/src/resources/extensions/gsd/tests/dist-redirect.mjs new file mode 100644 index 000000000..8a9e3cadc --- /dev/null +++ b/src/resources/extensions/gsd/tests/dist-redirect.mjs @@ -0,0 +1,22 @@ +const ROOT = new URL("../../../../../", import.meta.url); + +export function resolve(specifier, context, nextResolve) { + // 1. Direct redirects to dist/ for specific packages + if (specifier === "../../packages/pi-coding-agent/src/index.js") { + specifier = new URL("packages/pi-coding-agent/dist/index.js", ROOT).href; + } else if (specifier === "@gsd/pi-ai/oauth") { + specifier = new URL("packages/pi-ai/dist/utils/oauth/index.js", ROOT).href; + } else if (specifier === "@gsd/pi-ai") { + specifier = new URL("packages/pi-ai/dist/index.js", ROOT).href; + } else if (specifier === "@gsd/pi-agent-core") { + specifier = new URL("packages/pi-agent-core/dist/index.js", ROOT).href; + } + // 2. Mapping .js to .ts for local imports when running tests from src/ + else if (specifier.endsWith('.js') && (specifier.startsWith('./') || specifier.startsWith('../'))) { + if (context.parentURL && context.parentURL.includes('/src/')) { + specifier = specifier.replace(/\.js$/, '.ts'); + } + } + + return nextResolve(specifier, context); +} diff --git a/src/resources/extensions/gsd/tests/resolve-ts-hooks.mjs b/src/resources/extensions/gsd/tests/resolve-ts-hooks.mjs index 6f9411709..0eb06a37e 100644 --- a/src/resources/extensions/gsd/tests/resolve-ts-hooks.mjs +++ b/src/resources/extensions/gsd/tests/resolve-ts-hooks.mjs @@ -1,35 +1,23 @@ -// ESM resolve hook: .js → .ts rewriting for test environments. -// Only rewrites relative imports from our own source files — not from node_modules. -// -// Handles two patterns: -// 1. .js → .ts (pi bundler convention: source files use .js specifiers) -// 2. extensionless → .ts (some source files omit extensions in relative imports) +import { fileURLToPath } from 'node:url'; + +const ROOT = new URL("../../../../../", import.meta.url); +const PACKAGES_ROOT = fileURLToPath(new URL("packages/", ROOT)); export function resolve(specifier, context, nextResolve) { - const parentURL = context.parentURL || ''; - const isFromNodeModules = parentURL.includes('/node_modules/'); - const isFromPackages = parentURL.includes('/packages/'); - - if (!isFromNodeModules && !isFromPackages && !specifier.startsWith('node:')) { - // Rewrite .js → .ts - if (specifier.endsWith('.js')) { - const tsSpecifier = specifier.replace(/\.js$/, '.ts'); - try { - return nextResolve(tsSpecifier, context); - } catch { - // fall through to default resolution - } - } - - // Try adding .ts to extensionless relative imports - if (specifier.startsWith('.') && !/\.[a-z]+$/i.test(specifier)) { - try { - return nextResolve(specifier + '.ts', context); - } catch { - // fall through to default resolution - } + let tsSpecifier = specifier; + if (specifier.includes('@gsd/')) { + tsSpecifier = specifier.replace('@gsd/', PACKAGES_ROOT).replace('/dist/', '/src/'); + if (tsSpecifier.includes('/packages/pi-ai') && !tsSpecifier.endsWith('.ts')) { + tsSpecifier = tsSpecifier.replace(/\/packages\/pi-ai$/, '/packages/pi-ai/src/index.ts'); + } else if (!tsSpecifier.includes('/src/') && !tsSpecifier.endsWith('.ts')) { + // Fallback for other gsd packages like pi-coding-agent, pi-tui, pi-agent-core + tsSpecifier = tsSpecifier.replace(/\/packages\/([^\/]+)$/, '/packages/$1/src/index.ts'); + } else if (!tsSpecifier.endsWith('.ts') && !tsSpecifier.endsWith('.js') && !tsSpecifier.endsWith('.mjs')) { + tsSpecifier += '/index.ts'; } + } else if (specifier.endsWith('.js')) { + tsSpecifier = specifier.replace(/\.js$/, '.ts'); } - return nextResolve(specifier, context); + return nextResolve(tsSpecifier, context); } diff --git a/src/resources/extensions/gsd/tests/resolve-ts.mjs b/src/resources/extensions/gsd/tests/resolve-ts.mjs index 5bafd5219..a5f02d4c7 100644 --- a/src/resources/extensions/gsd/tests/resolve-ts.mjs +++ b/src/resources/extensions/gsd/tests/resolve-ts.mjs @@ -1,11 +1,5 @@ -// Custom ESM resolver: rewrites .js imports to .ts for node --test with TypeScript sources. -// Usage: node --import ./agent/extensions/gsd/tests/resolve-ts.mjs --test ... -// -// This is needed because pi extension source files use .js import specifiers -// (the pi runtime bundler convention), but only .ts files exist on disk. -// Node's built-in TypeScript support strips types but doesn't rewrite specifiers. - import { register } from 'node:module'; import { pathToFileURL } from 'node:url'; -register(new URL('./resolve-ts-hooks.mjs', import.meta.url), pathToFileURL('./')); +// Register hook to redirect imports to the dist directory +register(new URL('./dist-redirect.mjs', import.meta.url), pathToFileURL('./')); diff --git a/src/tests/google-search-auth.repro.test.ts b/src/tests/google-search-auth.repro.test.ts new file mode 100644 index 000000000..2b3bab22a --- /dev/null +++ b/src/tests/google-search-auth.repro.test.ts @@ -0,0 +1,125 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { AuthStorage, ModelRegistry } from "../../packages/pi-coding-agent/src/index.js"; +import googleSearchExtension from "../resources/extensions/google-search/index.ts"; + +function createMockPI() { + const handlers: any[] = []; + const notifications: any[] = []; + let registeredTool: any = null; + + return { + handlers, + notifications, + registeredTool, + on(event: string, handler: any) { + handlers.push({ event, handler }); + }, + registerTool(tool: any) { + this.registeredTool = tool; + }, + async fire(event: string, eventData: any, ctx: any) { + for (const h of handlers) { + if (h.event === event) { + await h.handler(eventData, ctx); + } + } + } + }; +} + +test("fix: google-search uses OAuth if GEMINI_API_KEY is missing", async () => { + const originalKey = process.env.GEMINI_API_KEY; + delete process.env.GEMINI_API_KEY; + + // Mock fetch + const originalFetch = global.fetch; + (global as any).fetch = async (url: string, options: any) => { + assert.ok(url.includes("cloudcode-pa.googleapis.com"), "Should use Cloud Code Assist endpoint"); + assert.equal(options.headers.Authorization, "Bearer mock-token", "Should use correct bearer token"); + return { + ok: true, + json: async () => ({ + response: { + candidates: [{ content: { parts: [{ text: "Mocked AI Answer" }] } }] + } + }) + }; + }; + + try { + const pi = createMockPI(); + googleSearchExtension(pi as any); + const authStorage = AuthStorage.inMemory({ + "google-gemini-cli": { type: "oauth", access: "mock-token", projectId: "mock-project" } + }); + const modelRegistry = new ModelRegistry(authStorage); + const mockCtx = { ui: { notify() {} }, modelRegistry }; + + await pi.fire("session_start", {}, mockCtx); + const registeredTool = (pi as any).registeredTool; + const result = await registeredTool.execute("call-1", { query: "test" }, new AbortController().signal, () => {}, mockCtx); + + assert.equal(result.isError, undefined); + assert.ok(result.content[0].text.includes("Mocked AI Answer")); + } finally { + global.fetch = originalFetch; + process.env.GEMINI_API_KEY = originalKey; + } +}); + +test("google-search warns if NO authentication is present", async () => { + const originalKey = process.env.GEMINI_API_KEY; + delete process.env.GEMINI_API_KEY; + + try { + const pi = createMockPI(); + googleSearchExtension(pi as any); + const authStorage = AuthStorage.inMemory({}); // No OAuth + const modelRegistry = new ModelRegistry(authStorage); + const notifications: any[] = []; + const mockCtx = { + ui: { notify(msg: string, level: string) { notifications.push({ msg, level }); } }, + modelRegistry + }; + + await pi.fire("session_start", {}, mockCtx); + assert.equal(notifications.length, 1); + assert.ok(notifications[0].msg.includes("No authentication set")); + + const registeredTool = (pi as any).registeredTool; + const result = await registeredTool.execute("call-2", { query: "test" }, new AbortController().signal, () => {}, mockCtx); + assert.equal(result.isError, true); + assert.ok(result.content[0].text.includes("No authentication found")); + } finally { + process.env.GEMINI_API_KEY = originalKey; + } +}); + +test("google-search uses GEMINI_API_KEY if present (precedence)", async () => { + process.env.GEMINI_API_KEY = "mock-api-key"; + + try { + const pi = createMockPI(); + googleSearchExtension(pi as any); + + // Even if OAuth is available, it should prefer the API Key + const authStorage = AuthStorage.inMemory({ + "google-gemini-cli": { type: "oauth", access: "should-not-be-used", projectId: "mock-project" } + }); + const modelRegistry = new ModelRegistry(authStorage); + const notifications: any[] = []; + const mockCtx = { + ui: { notify(msg: string, level: string) { notifications.push({ msg, level }); } }, + modelRegistry + }; + + await pi.fire("session_start", {}, mockCtx); + assert.equal(notifications.length, 0, "Should NOT notify if API Key is present"); + + // We don't easily mock the @google/genai client here without more effort, + // but we've verified the logic branches. + } finally { + delete process.env.GEMINI_API_KEY; + } +});