Feat: Adding Gemini tools - URL Context and Grounding with Google Search

packages/types/src/provider-settings.ts

@@ -164,6 +164,8 @@ const lmStudioSchema = baseProviderSettingsSchema.extend({
 const geminiSchema = apiModelIdProviderModelSchema.extend({
 	geminiApiKey: z.string().optional(),
 	googleGeminiBaseUrl: z.string().optional(),
+	enableUrlContext: z.boolean().optional(),
+	enableGrounding: z.boolean().optional(),
 })
 
 const geminiCliSchema = apiModelIdProviderModelSchema.extend({

src/api/providers/__tests__/gemini-handler.spec.ts

@@ -0,0 +1,36 @@
+import { describe, it, expect, vi } from "vitest"
+import { GeminiHandler } from "../gemini"
+import type { ApiHandlerOptions } from "../../../shared/api"
+
+describe("GeminiHandler backend support", () => {
+	it("passes tools for URL context and grounding in config", async () => {
+		const options = {
+			apiProvider: "gemini",
+			enableUrlContext: true,
+			enableGrounding: true,
+		} as ApiHandlerOptions
+		const handler = new GeminiHandler(options)
+		const stub = vi.fn().mockReturnValue((async function* () {})())
+		// @ts-ignore access private client
+		handler["client"].models.generateContentStream = stub
+		await handler.createMessage("instr", [] as any).next()
+		const config = stub.mock.calls[0][0].config
+		expect(config.tools).toEqual([{ urlContext: {} }, { googleSearch: {} }])
+	})
+
+	it("completePrompt passes config overrides without tools when URL context and grounding disabled", async () => {
+		const options = {
+			apiProvider: "gemini",
+			enableUrlContext: false,
+			enableGrounding: false,
+		} as ApiHandlerOptions
+		const handler = new GeminiHandler(options)
+		const stub = vi.fn().mockResolvedValue({ text: "ok" })
+		// @ts-ignore access private client
+		handler["client"].models.generateContent = stub
+		const res = await handler.completePrompt("hi")
+		expect(res).toBe("ok")
+		const promptConfig = stub.mock.calls[0][0].config
+		expect(promptConfig.tools).toBeUndefined()
+	})
+})

src/api/providers/gemini.ts

@@ -4,6 +4,7 @@ import {
 	type GenerateContentResponseUsageMetadata,
 	type GenerateContentParameters,
 	type GenerateContentConfig,
+	type GroundingMetadata,
 } from "@google/genai"
 import type { JWTInput } from "google-auth-library"
 
@@ -67,72 +68,101 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
 
-		const config: GenerateContentConfig = {
+		const tools: GenerateContentConfig["tools"] = []
+		if (this.options.enableUrlContext) {
+			tools.push({ urlContext: {} })
+		}
+		if (this.options.enableGrounding) {
+			tools.push({ googleSearch: {} })
+		}
+		const rawConfig = {
 			systemInstruction,
 			httpOptions: this.options.googleGeminiBaseUrl ? { baseUrl: this.options.googleGeminiBaseUrl } : undefined,
 			thinkingConfig,
-			maxOutputTokens: this.options.modelMaxTokens ?? maxTokens ?? undefined,
 			temperature: this.options.modelTemperature ?? 0,
+			...(tools.length > 0 ? { tools } : {}),
 		}
+		const config = rawConfig as unknown as GenerateContentConfig
 
 		const params: GenerateContentParameters = { model, contents, config }
 
-		const result = await this.client.models.generateContentStream(params)
+		try {
+			const result = await this.client.models.generateContentStream(params)
 
-		let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
+			let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
+			let pendingGroundingMetadata: GroundingMetadata | undefined
 
-		for await (const chunk of result) {
-			// Process candidates and their parts to separate thoughts from content
-			if (chunk.candidates && chunk.candidates.length > 0) {
-				const candidate = chunk.candidates[0]
-				if (candidate.content && candidate.content.parts) {
-					for (const part of candidate.content.parts) {
-						if (part.thought) {
-							// This is a thinking/reasoning part
-							if (part.text) {
-								yield { type: "reasoning", text: part.text }
-							}
-						} else {
-							// This is regular content
-							if (part.text) {
-								yield { type: "text", text: part.text }
+			for await (const chunk of result) {
+				// Process candidates and their parts to separate thoughts from content
+				if (chunk.candidates && chunk.candidates.length > 0) {
+					const candidate = chunk.candidates[0]
+
+					if (candidate.groundingMetadata) {
+						pendingGroundingMetadata = candidate.groundingMetadata
+					}
+
+					if (candidate.content && candidate.content.parts) {
+						for (const part of candidate.content.parts) {
+							if (part.thought) {
+								// This is a thinking/reasoning part
+								if (part.text) {
+									yield { type: "reasoning", text: part.text }
+								}
+							} else {
+								// This is regular content
+								if (part.text) {
+									yield { type: "text", text: part.text }
+								}
 							}
 						}
 					}
 				}
-			}
 
-			// Fallback to the original text property if no candidates structure
-			else if (chunk.text) {
-				yield { type: "text", text: chunk.text }
+				// Fallback to the original text property if no candidates structure
+				else if (chunk.text) {
+					yield { type: "text", text: chunk.text }
+				}
+
+				if (chunk.usageMetadata) {
+					lastUsageMetadata = chunk.usageMetadata
+				}
 			}
 
-			if (chunk.usageMetadata) {
-				lastUsageMetadata = chunk.usageMetadata
+			if (pendingGroundingMetadata) {
+				const citations = this.extractCitationsOnly(pendingGroundingMetadata)
+				if (citations) {
+					yield { type: "text", text: `\n\nSources: ${citations}` }
+				}
 			}
-		}
 
-		if (lastUsageMetadata) {
-			const inputTokens = lastUsageMetadata.promptTokenCount ?? 0
-			const outputTokens = lastUsageMetadata.candidatesTokenCount ?? 0
-			const cacheReadTokens = lastUsageMetadata.cachedContentTokenCount
-			const reasoningTokens = lastUsageMetadata.thoughtsTokenCount
-
-			yield {
-				type: "usage",
-				inputTokens,
-				outputTokens,
-				cacheReadTokens,
-				reasoningTokens,
-				totalCost: this.calculateCost({ info, inputTokens, outputTokens, cacheReadTokens }),
+			if (lastUsageMetadata) {
+				const inputTokens = lastUsageMetadata.promptTokenCount ?? 0
+				const outputTokens = lastUsageMetadata.candidatesTokenCount ?? 0
+				const cacheReadTokens = lastUsageMetadata.cachedContentTokenCount
+				const reasoningTokens = lastUsageMetadata.thoughtsTokenCount
+
+				yield {
+					type: "usage",
+					inputTokens,
+					outputTokens,
+					cacheReadTokens,
+					reasoningTokens,
+					totalCost: this.calculateCost({ info, inputTokens, outputTokens, cacheReadTokens }),
+				}
 			}
+		} catch (error) {
+			if (error instanceof Error) {
+				throw new Error(`Gemini Generate Context Stream error: ${error.message}`)
+			}
+
+			throw error
 		}
 	}
 
 	override getModel() {
 		const modelId = this.options.apiModelId
 		let id = modelId && modelId in geminiModels ? (modelId as GeminiModelId) : geminiDefaultModelId
-		const info: ModelInfo = geminiModels[id]
+		let info: ModelInfo = geminiModels[id]
 		const params = getModelParams({ format: "gemini", modelId: id, model: info, settings: this.options })
 
 		// The `:thinking` suffix indicates that the model is a "Hybrid"
@@ -142,22 +172,67 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		return { id: id.endsWith(":thinking") ? id.replace(":thinking", "") : id, info, ...params }
 	}
 
+	private extractCitationsOnly(groundingMetadata?: GroundingMetadata): string | null {
+		const chunks = groundingMetadata?.groundingChunks
+
+		if (!chunks) {
+			return null
+		}
+
+		const citationLinks = chunks
+			.map((chunk, i) => {
+				const uri = chunk.web?.uri
+				if (uri) {
+					return `[${i + 1}](${uri})`
+				}
+				return null
+			})
+			.filter((link): link is string => link !== null)
+
+		if (citationLinks.length > 0) {
+			return citationLinks.join(", ")
+		}
+
+		return null
+	}
+
 	async completePrompt(prompt: string): Promise<string> {
 		try {
 			const { id: model } = this.getModel()
 
+			const tools: GenerateContentConfig["tools"] = []
+			if (this.options.enableUrlContext) {
+				tools.push({ urlContext: {} })
+			}
+			if (this.options.enableGrounding) {
+				tools.push({ googleSearch: {} })
+			}
+			const rawPromptConfig = {
+				httpOptions: this.options.googleGeminiBaseUrl
+					? { baseUrl: this.options.googleGeminiBaseUrl }
+					: undefined,
+				temperature: this.options.modelTemperature ?? 0,
+				...(tools.length > 0 ? { tools } : {}),
+			}
+			const promptConfig = rawPromptConfig as unknown as GenerateContentConfig
+
 			const result = await this.client.models.generateContent({
 				model,
 				contents: [{ role: "user", parts: [{ text: prompt }] }],
-				config: {
-					httpOptions: this.options.googleGeminiBaseUrl
-						? { baseUrl: this.options.googleGeminiBaseUrl }
-						: undefined,
-					temperature: this.options.modelTemperature ?? 0,
-				},
+				config: promptConfig,
 			})
 
-			return result.text ?? ""
+			let text = result.text ?? ""
+
+			const candidate = result.candidates?.[0]
+			if (candidate?.groundingMetadata) {
+				const citations = this.extractCitationsOnly(candidate.groundingMetadata)
+				if (citations) {
+					text += `\n\nSources: ${citations}`
+				}
+			}
+
+			return text
 		} catch (error) {
 			if (error instanceof Error) {
 				throw new Error(`Gemini completion error: ${error.message}`)

src/core/sliding-window/__tests__/sliding-window.spec.ts

@@ -250,7 +250,6 @@ describe("Sliding Window", () => {
 			{ role: "assistant", content: "Fourth message" },
 			{ role: "user", content: "Fifth message" },
 		]
-
 		it("should not truncate if tokens are below max tokens threshold", async () => {
 			const modelInfo = createModelInfo(100000, 30000)
 			const dynamicBuffer = modelInfo.contextWindow * TOKEN_BUFFER_PERCENTAGE // 10000

webview-ui/src/components/settings/ApiOptions.tsx

@@ -443,7 +443,11 @@ const ApiOptions = ({
 			)}
 
 			{selectedProvider === "gemini" && (
-				<Gemini apiConfiguration={apiConfiguration} setApiConfigurationField={setApiConfigurationField} />
+				<Gemini
+					apiConfiguration={apiConfiguration}
+					setApiConfigurationField={setApiConfigurationField}
+					currentModelId={selectedModelId}
+				/>
 			)}
 
 			{selectedProvider === "openai" && (

webview-ui/src/components/settings/providers/Gemini.tsx

@@ -12,6 +12,7 @@ import { inputEventTransform } from "../transforms"
 type GeminiProps = {
 	apiConfiguration: ProviderSettings
 	setApiConfigurationField: (field: keyof ProviderSettings, value: ProviderSettings[keyof ProviderSettings]) => void
+	currentModelId?: string
 }
 
 export const Gemini = ({ apiConfiguration, setApiConfigurationField }: GeminiProps) => {
@@ -50,12 +51,13 @@ export const Gemini = ({ apiConfiguration, setApiConfigurationField }: GeminiPro
 					{t("settings:providers.getGeminiApiKey")}
 				</VSCodeButtonLink>
 			)}
+
 			<div>
 				<Checkbox
+					data-testid="checkbox-custom-base-url"
 					checked={googleGeminiBaseUrlSelected}
 					onChange={(checked: boolean) => {
 						setGoogleGeminiBaseUrlSelected(checked)
-
 						if (!checked) {
 							setApiConfigurationField("googleGeminiBaseUrl", "")
 						}
@@ -71,6 +73,27 @@ export const Gemini = ({ apiConfiguration, setApiConfigurationField }: GeminiPro
 						className="w-full mt-1"
 					/>
 				)}
+
+				<Checkbox
+					className="mt-6"
+					data-testid="checkbox-url-context"
+					checked={!!apiConfiguration.enableUrlContext}
+					onChange={(checked: boolean) => setApiConfigurationField("enableUrlContext", checked)}>
+					{t("settings:providers.geminiParameters.urlContext.title")}
+				</Checkbox>
+				<div className="text-sm text-vscode-descriptionForeground mb-3">
+					{t("settings:providers.geminiParameters.urlContext.description")}
+				</div>
+
+				<Checkbox
+					data-testid="checkbox-grounding-search"
+					checked={!!apiConfiguration.enableGrounding}
+					onChange={(checked: boolean) => setApiConfigurationField("enableGrounding", checked)}>
+					{t("settings:providers.geminiParameters.groundingSearch.title")}
+				</Checkbox>
+				<div className="text-sm text-vscode-descriptionForeground mb-3">
+					{t("settings:providers.geminiParameters.groundingSearch.description")}
+				</div>
 			</div>
 		</>
 	)