feat: enable grounding features for Vertex AI (#6777)

* feat: enable grounding features for Vertex AI - Add enableUrlContext and enableGrounding options to vertexSchema in provider-settings.ts - Update Vertex.tsx UI component to include URL context and grounding checkboxes - Add comprehensive tests for the new grounding features in Vertex.spec.tsx - Reuse existing localization strings from Gemini provider Fixes #6776 * fix: address review feedback - hide grounding features in welcome view and only show for Gemini models --------- Co-authored-by: Roo Code <roomote@roocode.com> Co-authored-by: daniel-lxs <ricciodaniel98@gmail.com>
2026-09-07 08:26:51 +00:00 · 2025-08-21 13:02:45 -07:00 · 2025-08-21 13:02:45 -07:00 · 0bc1183c71
commit 0bc1183c71
parent 4664955127
4 changed files with 308 additions and 53 deletions
--- a/packages/types/src/provider-settings.ts
+++ b/packages/types/src/provider-settings.ts
@ -163,6 +163,8 @@ const vertexSchema = apiModelIdProviderModelSchema.extend({
 	vertexJsonCredentials: z.string().optional(),
 	vertexProjectId: z.string().optional(),
 	vertexRegion: z.string().optional(),
+	enableUrlContext: z.boolean().optional(),
+	enableGrounding: z.boolean().optional(),
 })

 const openAiSchema = baseProviderSettingsSchema.extend({
--- a/webview-ui/src/components/settings/ApiOptions.tsx
+++ b/webview-ui/src/components/settings/ApiOptions.tsx
@ -479,7 +479,11 @@ const ApiOptions = ({
 			)}

 			{selectedProvider === "vertex" && (
-				<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={setApiConfigurationField} />
+				<Vertex
+					apiConfiguration={apiConfiguration}
+					setApiConfigurationField={setApiConfigurationField}
+					fromWelcomeView={fromWelcomeView}
+				/>
 			)}

 			{selectedProvider === "gemini" && (
--- a/webview-ui/src/components/settings/providers/Vertex.tsx
+++ b/webview-ui/src/components/settings/providers/Vertex.tsx
@ -1,4 +1,5 @@
 import { useCallback } from "react"
+import { Checkbox } from "vscrui"
 import { VSCodeLink, VSCodeTextField } from "@vscode/webview-ui-toolkit/react"

 import { type ProviderSettings, VERTEX_REGIONS } from "@roo-code/types"
@ -11,9 +12,10 @@ import { inputEventTransform } from "../transforms"
 type VertexProps = {
 	apiConfiguration: ProviderSettings
 	setApiConfigurationField: (field: keyof ProviderSettings, value: ProviderSettings[keyof ProviderSettings]) => void
+	fromWelcomeView?: boolean
 }

-export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexProps) => {
+export const Vertex = ({ apiConfiguration, setApiConfigurationField, fromWelcomeView }: VertexProps) => {
 	const { t } = useAppTranslation()

 	const handleInputChange = useCallback(
@ -91,6 +93,30 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexPro
 					</SelectContent>
 				</Select>
 			</div>
+
+			{!fromWelcomeView && apiConfiguration.apiModelId?.startsWith("gemini") && (
+				<div className="mt-6">
+					<Checkbox
+						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 mt-1.5">
+						{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 mt-1.5">
+						{t("settings:providers.geminiParameters.groundingSearch.description")}
+					</div>
+				</div>
+			)}
 		</>
 	)
 }
--- a/webview-ui/src/components/settings/providers/tests/Vertex.spec.tsx
+++ b/webview-ui/src/components/settings/providers/tests/Vertex.spec.tsx
@ -1,61 +1,284 @@
-// Tests for VERTEX_REGIONS "global" region handling
+import { render, screen } from "@testing-library/react"
+import userEvent from "@testing-library/user-event"
+import { Vertex } from "../Vertex"
+import type { ProviderSettings } from "@roo-code/types"
+import { VERTEX_REGIONS } from "@roo-code/types"

-import { describe, it, expect } from "vitest"
-import { VERTEX_REGIONS } from "../../../../../../packages/types/src/providers/vertex"
+vi.mock("@vscode/webview-ui-toolkit/react", () => ({
+	VSCodeTextField: ({ children, value, onInput, type }: any) => (
+		<div>
+			{children}
+			<input type={type} value={value} onChange={(e) => onInput(e)} />
+		</div>
+	),
+	VSCodeLink: ({ children, href }: any) => <a href={href}>{children}</a>,
+}))

-describe("VERTEX_REGIONS", () => {
-	it('should include the "global" region as the first entry', () => {
-		expect(VERTEX_REGIONS[0]).toEqual({ value: "global", label: "global" })
+vi.mock("vscrui", () => ({
+	Checkbox: ({ children, checked, onChange, "data-testid": testId }: any) => (
+		<label data-testid={testId}>
+			<input type="checkbox" checked={checked} onChange={(e) => onChange(e.target.checked)} />
+			{children}
+		</label>
+	),
+}))
+
+vi.mock("@src/i18n/TranslationContext", () => ({
+	useAppTranslation: () => ({ t: (key: string) => key }),
+}))
+
+vi.mock("@src/components/ui", () => ({
+	Select: ({ children, value, onValueChange }: any) => (
+		<div data-value={value} data-onvaluechange={onValueChange}>
+			{children}
+		</div>
+	),
+	SelectContent: ({ children }: any) => <div>{children}</div>,
+	SelectItem: ({ children, value }: any) => <div data-value={value}>{children}</div>,
+	SelectTrigger: ({ children }: any) => <div>{children}</div>,
+	SelectValue: ({ placeholder }: any) => <div>{placeholder}</div>,
+}))
+
+describe("Vertex", () => {
+	const defaultApiConfiguration: ProviderSettings = {
+		vertexKeyFile: "",
+		vertexJsonCredentials: "",
+		vertexProjectId: "",
+		vertexRegion: "",
+		enableUrlContext: false,
+		enableGrounding: false,
+		apiModelId: "gemini-2.0-flash-001",
+	}
+
+	const mockSetApiConfigurationField = vi.fn()
+
+	beforeEach(() => {
+		vi.clearAllMocks()
 	})

-	it('should contain "global" region exactly once', () => {
-		const globalRegions = VERTEX_REGIONS.filter((r: { value: string; label: string }) => r.value === "global")
-		expect(globalRegions).toHaveLength(1)
+	describe("VERTEX_REGIONS", () => {
+		it('should include the "global" region as the first entry', () => {
+			expect(VERTEX_REGIONS[0]).toEqual({ value: "global", label: "global" })
+		})
+
+		it('should contain "global" region exactly once', () => {
+			const globalRegions = VERTEX_REGIONS.filter((r: { value: string; label: string }) => r.value === "global")
+			expect(globalRegions).toHaveLength(1)
+		})
+
+		it('should contain all expected regions including "global"', () => {
+			// The expected list is the imported VERTEX_REGIONS itself
+			expect(VERTEX_REGIONS).toEqual([
+				{ value: "global", label: "global" },
+				{ value: "us-central1", label: "us-central1" },
+				{ value: "us-east1", label: "us-east1" },
+				{ value: "us-east4", label: "us-east4" },
+				{ value: "us-east5", label: "us-east5" },
+				{ value: "us-west1", label: "us-west1" },
+				{ value: "us-west2", label: "us-west2" },
+				{ value: "us-west3", label: "us-west3" },
+				{ value: "us-west4", label: "us-west4" },
+				{ value: "northamerica-northeast1", label: "northamerica-northeast1" },
+				{ value: "northamerica-northeast2", label: "northamerica-northeast2" },
+				{ value: "southamerica-east1", label: "southamerica-east1" },
+				{ value: "europe-west1", label: "europe-west1" },
+				{ value: "europe-west2", label: "europe-west2" },
+				{ value: "europe-west3", label: "europe-west3" },
+				{ value: "europe-west4", label: "europe-west4" },
+				{ value: "europe-west6", label: "europe-west6" },
+				{ value: "europe-central2", label: "europe-central2" },
+				{ value: "asia-east1", label: "asia-east1" },
+				{ value: "asia-east2", label: "asia-east2" },
+				{ value: "asia-northeast1", label: "asia-northeast1" },
+				{ value: "asia-northeast2", label: "asia-northeast2" },
+				{ value: "asia-northeast3", label: "asia-northeast3" },
+				{ value: "asia-south1", label: "asia-south1" },
+				{ value: "asia-south2", label: "asia-south2" },
+				{ value: "asia-southeast1", label: "asia-southeast1" },
+				{ value: "asia-southeast2", label: "asia-southeast2" },
+				{ value: "australia-southeast1", label: "australia-southeast1" },
+				{ value: "australia-southeast2", label: "australia-southeast2" },
+				{ value: "me-west1", label: "me-west1" },
+				{ value: "me-central1", label: "me-central1" },
+				{ value: "africa-south1", label: "africa-south1" },
+			])
+		})
+
+		it('should contain "asia-east1" region exactly once', () => {
+			const asiaEast1Regions = VERTEX_REGIONS.filter(
+				(r: { value: string; label: string }) => r.value === "asia-east1" && r.label === "asia-east1",
+			)
+			expect(asiaEast1Regions).toHaveLength(1)
+			expect(asiaEast1Regions[0]).toEqual({ value: "asia-east1", label: "asia-east1" })
+		})
 	})

-	it('should contain all expected regions including "global"', () => {
-		// The expected list is the imported VERTEX_REGIONS itself
-		expect(VERTEX_REGIONS).toEqual([
-			{ value: "global", label: "global" },
-			{ value: "us-central1", label: "us-central1" },
-			{ value: "us-east1", label: "us-east1" },
-			{ value: "us-east4", label: "us-east4" },
-			{ value: "us-east5", label: "us-east5" },
-			{ value: "us-west1", label: "us-west1" },
-			{ value: "us-west2", label: "us-west2" },
-			{ value: "us-west3", label: "us-west3" },
-			{ value: "us-west4", label: "us-west4" },
-			{ value: "northamerica-northeast1", label: "northamerica-northeast1" },
-			{ value: "northamerica-northeast2", label: "northamerica-northeast2" },
-			{ value: "southamerica-east1", label: "southamerica-east1" },
-			{ value: "europe-west1", label: "europe-west1" },
-			{ value: "europe-west2", label: "europe-west2" },
-			{ value: "europe-west3", label: "europe-west3" },
-			{ value: "europe-west4", label: "europe-west4" },
-			{ value: "europe-west6", label: "europe-west6" },
-			{ value: "europe-central2", label: "europe-central2" },
-			{ value: "asia-east1", label: "asia-east1" },
-			{ value: "asia-east2", label: "asia-east2" },
-			{ value: "asia-northeast1", label: "asia-northeast1" },
-			{ value: "asia-northeast2", label: "asia-northeast2" },
-			{ value: "asia-northeast3", label: "asia-northeast3" },
-			{ value: "asia-south1", label: "asia-south1" },
-			{ value: "asia-south2", label: "asia-south2" },
-			{ value: "asia-southeast1", label: "asia-southeast1" },
-			{ value: "asia-southeast2", label: "asia-southeast2" },
-			{ value: "australia-southeast1", label: "australia-southeast1" },
-			{ value: "australia-southeast2", label: "australia-southeast2" },
-			{ value: "me-west1", label: "me-west1" },
-			{ value: "me-central1", label: "me-central1" },
-			{ value: "africa-south1", label: "africa-south1" },
-		])
+	describe("URL Context Checkbox", () => {
+		it("should render URL context checkbox unchecked by default for Gemini models", () => {
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+				/>,
+			)
+
+			const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
+			const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+			expect(checkbox.checked).toBe(false)
+		})
+
+		it("should NOT render URL context checkbox for non-Gemini models", () => {
+			const apiConfiguration = { ...defaultApiConfiguration, apiModelId: "claude-3-opus@20240229" }
+			render(
+				<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
+			)
+
+			const urlContextCheckbox = screen.queryByTestId("checkbox-url-context")
+			expect(urlContextCheckbox).toBeNull()
+		})
+
+		it("should NOT render URL context checkbox when fromWelcomeView is true", () => {
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+					fromWelcomeView={true}
+				/>,
+			)
+
+			const urlContextCheckbox = screen.queryByTestId("checkbox-url-context")
+			expect(urlContextCheckbox).toBeNull()
+		})
+
+		it("should render URL context checkbox checked when enableUrlContext is true for Gemini models", () => {
+			const apiConfiguration = {
+				...defaultApiConfiguration,
+				enableUrlContext: true,
+				apiModelId: "gemini-2.0-flash-001",
+			}
+			render(
+				<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
+			)
+
+			const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
+			const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+			expect(checkbox.checked).toBe(true)
+		})
+
+		it("should call setApiConfigurationField with correct parameters when URL context checkbox is toggled", async () => {
+			const user = userEvent.setup()
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+				/>,
+			)
+
+			const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
+			const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+
+			await user.click(checkbox)
+
+			expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableUrlContext", true)
+		})
 	})

-	it('should contain "asia-east1" region exactly once', () => {
-		const asiaEast1Regions = VERTEX_REGIONS.filter(
-			(r: { value: string; label: string }) => r.value === "asia-east1" && r.label === "asia-east1",
-		)
-		expect(asiaEast1Regions).toHaveLength(1)
-		expect(asiaEast1Regions[0]).toEqual({ value: "asia-east1", label: "asia-east1" })
+	describe("Grounding with Google Search Checkbox", () => {
+		it("should render grounding search checkbox unchecked by default for Gemini models", () => {
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+				/>,
+			)
+
+			const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
+			const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+			expect(checkbox.checked).toBe(false)
+		})
+
+		it("should NOT render grounding search checkbox for non-Gemini models", () => {
+			const apiConfiguration = { ...defaultApiConfiguration, apiModelId: "claude-3-opus@20240229" }
+			render(
+				<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
+			)
+
+			const groundingCheckbox = screen.queryByTestId("checkbox-grounding-search")
+			expect(groundingCheckbox).toBeNull()
+		})
+
+		it("should NOT render grounding search checkbox when fromWelcomeView is true", () => {
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+					fromWelcomeView={true}
+				/>,
+			)
+
+			const groundingCheckbox = screen.queryByTestId("checkbox-grounding-search")
+			expect(groundingCheckbox).toBeNull()
+		})
+
+		it("should render grounding search checkbox checked when enableGrounding is true for Gemini models", () => {
+			const apiConfiguration = {
+				...defaultApiConfiguration,
+				enableGrounding: true,
+				apiModelId: "gemini-2.0-flash-001",
+			}
+			render(
+				<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
+			)
+
+			const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
+			const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+			expect(checkbox.checked).toBe(true)
+		})
+
+		it("should call setApiConfigurationField with correct parameters when grounding search checkbox is toggled", async () => {
+			const user = userEvent.setup()
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+				/>,
+			)
+
+			const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
+			const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+
+			await user.click(checkbox)
+
+			expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableGrounding", true)
+		})
+	})
+
+	describe("Both checkboxes interaction", () => {
+		it("should be able to toggle both checkboxes independently", async () => {
+			const user = userEvent.setup()
+			render(
+				<Vertex
+					apiConfiguration={defaultApiConfiguration}
+					setApiConfigurationField={mockSetApiConfigurationField}
+				/>,
+			)
+
+			const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
+			const urlCheckbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+
+			const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
+			const groundCheckbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
+
+			// Toggle URL context
+			await user.click(urlCheckbox)
+			expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableUrlContext", true)
+
+			// Toggle grounding
+			await user.click(groundCheckbox)
+			expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableGrounding", true)
+
+			// Both should have been called
+			expect(mockSetApiConfigurationField).toHaveBeenCalledTimes(2)
+		})
 	})
 })