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 { 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'; const DEVELOPER_SSE_RESPONSE_WITHOUT_MODEL = '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}}\n\n'; const ORIGINAL_GEMINI_SEARCH_MODEL = Bun.env.GEMINI_SEARCH_MODEL; type CapturedRequest = { url: string; headers: Record; body: Record | null; }; describe("searchGemini tools serialization", () => { let capturedRequest: CapturedRequest | null = null; const fakeAuthStorage = { async getOAuthAccess() { return { accessToken: "test-access-token", projectId: "test-project", }; }, hasOAuth() { return true; }, } as unknown as AuthStorage; 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) => { 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(responseText, { status: 200, headers: { "Content-Type": "text/event-stream" }, }), ); }; } afterEach(() => { capturedRequest = null; if (ORIGINAL_GEMINI_SEARCH_MODEL === undefined) { delete Bun.env.GEMINI_SEARCH_MODEL; } else { Bun.env.GEMINI_SEARCH_MODEL = ORIGINAL_GEMINI_SEARCH_MODEL; } }); function makeParams(query: string) { return { query, authStorage: fakeAuthStorage, systemPrompt: "Gemini test prompt", } 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("normalizes query directive aliases to canonical Google forms in the grounding request", async () => { const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("k8s domain:kubernetes.io since:2024"), fetch: fetchMock, }); expect(capturedRequest).not.toBeNull(); const request = capturedRequest?.body?.request as Record; expect(request).toMatchObject({ contents: [{ role: "user", parts: [{ text: "k8s site:kubernetes.io after:2024-01-01" }] }], }); }); it("leaves directive-free queries untouched in the developer API request", async () => { const fetchMock = mockGeminiFetch(DEVELOPER_SSE_RESPONSE); await searchGemini({ ...makeParams("plain query with no operators"), authStorage: apiKeyAuthStorage, fetch: fetchMock, }); expect(capturedRequest).not.toBeNull(); expect(capturedRequest?.body).toMatchObject({ contents: [{ role: "user", parts: [{ text: "plain query with no operators" }] }], }); }); it("uses configured developer API model and reports it when modelVersion is absent", async () => { const fetchMock = mockGeminiFetch(DEVELOPER_SSE_RESPONSE_WITHOUT_MODEL); const response = await searchGemini({ ...makeParams("developer api configured"), authStorage: apiKeyAuthStorage, geminiModel: "gemini-3.5-flash", fetch: fetchMock, }); expect(capturedRequest?.url).toBe( "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.5-flash:streamGenerateContent?alt=sse", ); expect(response.model).toBe("gemini-3.5-flash"); }); it("uses configured OAuth model in the Cloud Code request body", async () => { const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("oauth configured"), geminiModel: "gemini-3.5-flash", fetch: fetchMock, }); expect(capturedRequest?.body).toMatchObject({ model: "gemini-3.5-flash", }); }); it("lets GEMINI_SEARCH_MODEL override the configured Gemini model", async () => { Bun.env.GEMINI_SEARCH_MODEL = "gemini-2.5-pro"; const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("env configured"), geminiModel: "gemini-3.5-flash", fetch: fetchMock, }); expect(capturedRequest?.body).toMatchObject({ model: "gemini-2.5-pro", }); }); it("sends default googleSearch tool when no passthrough payloads are provided", async () => { const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("default tools"), fetch: fetchMock }); expect(capturedRequest).not.toBeNull(); expect(capturedRequest?.body?.request).toMatchObject({ tools: [{ googleSearch: {} }], }); expect(capturedRequest?.body).toMatchObject({ model: "gemini-2.5-flash", }); }); it("passes through googleSearch payload into googleSearch tool", async () => { const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("google payload"), google_search: { dynamicRetrievalConfig: { mode: "MODE_DYNAMIC" } }, fetch: fetchMock, }); expect(capturedRequest).not.toBeNull(); expect(capturedRequest?.body?.request).toMatchObject({ tools: [{ googleSearch: { dynamicRetrievalConfig: { mode: "MODE_DYNAMIC" } } }], }); }); it("includes codeExecution and urlContext tools when provided", async () => { const fetchMock = mockGeminiFetch(); await searchGemini({ ...makeParams("extended tools"), code_execution: {}, url_context: { allowedDomains: ["example.com"] }, fetch: fetchMock, }); expect(capturedRequest).not.toBeNull(); expect(capturedRequest?.body?.request).toMatchObject({ tools: [{ googleSearch: {} }, { codeExecution: {} }, { urlContext: { allowedDomains: ["example.com"] } }], }); }); });