From 1698100ea2a28e0de3396222765f57ec2f25b7e5 Mon Sep 17 00:00:00 2001
From: daniel-lxs <ricciodaniel98@gmail.com>
Date: Fri, 12 Sep 2025 16:22:19 -0500
Subject: [PATCH] fix: improve type safety in VertexAI custom base URL
 implementation

- Replace 'any' type with proper VertexOptions type in anthropic-vertex.ts
- Use isVertex property instead of fragile constructor.name comparison in gemini.ts
- Improve test assertions clarity by checking undefined explicitly
---
 src/api/providers/__tests__/anthropic-vertex.spec.ts |  7 ++-----
 src/api/providers/__tests__/vertex.spec.ts           |  9 ++-------
 src/api/providers/anthropic-vertex.ts                |  9 ++++++++-
 src/api/providers/gemini.ts                          | 10 ++++------
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts
index 31e1152ea2..49fe41e872 100644
--- a/src/api/providers/__tests__/anthropic-vertex.spec.ts
+++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts
@@ -893,11 +893,8 @@ describe("VertexHandler", () => {
 			})
 
 			// Verify that AnthropicVertex was called without baseURL
-			expect(AnthropicVertex).toHaveBeenCalledWith(
-				expect.not.objectContaining({
-					baseURL: expect.anything(),
-				}),
-			)
+			const callArgs = (AnthropicVertex as any).mock.calls[0][0]
+			expect(callArgs.baseURL).toBeUndefined()
 		})
 	})
 })
diff --git a/src/api/providers/__tests__/vertex.spec.ts b/src/api/providers/__tests__/vertex.spec.ts
index 64325d3052..d140a58130 100644
--- a/src/api/providers/__tests__/vertex.spec.ts
+++ b/src/api/providers/__tests__/vertex.spec.ts
@@ -184,13 +184,8 @@ describe("VertexHandler", () => {
 			await handler.completePrompt("Test prompt")
 
 			// Verify that httpOptions is undefined when no custom URL
-			expect(mockGenerateContent).toHaveBeenCalledWith(
-				expect.objectContaining({
-					config: expect.objectContaining({
-						httpOptions: undefined,
-					}),
-				}),
-			)
+			const callArgs = mockGenerateContent.mock.calls[0][0]
+			expect(callArgs.config.httpOptions).toBeUndefined()
 		})
 	})
 })
diff --git a/src/api/providers/anthropic-vertex.ts b/src/api/providers/anthropic-vertex.ts
index 73e9fe97e1..94a4cd3d59 100644
--- a/src/api/providers/anthropic-vertex.ts
+++ b/src/api/providers/anthropic-vertex.ts
@@ -34,7 +34,14 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
 		const projectId = this.options.vertexProjectId ?? "not-provided"
 		const region = this.options.vertexRegion ?? "us-east5"
 
-		const baseOptions: any = {
+		type VertexOptions = {
+			projectId: string
+			region: string
+			baseURL?: string
+			googleAuth?: GoogleAuth
+		}
+
+		const baseOptions: VertexOptions = {
 			projectId,
 			region,
 		}
diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts
index 31290921ee..823510208d 100644
--- a/src/api/providers/gemini.ts
+++ b/src/api/providers/gemini.ts
@@ -27,6 +27,7 @@ type GeminiHandlerOptions = ApiHandlerOptions & {
 
 export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
+	private isVertex: boolean
 
 	private client: GoogleGenAI
 
@@ -34,6 +35,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		super()
 
 		this.options = options
+		this.isVertex = isVertex ?? false
 
 		const project = this.options.vertexProjectId ?? "not-provided"
 		const location = this.options.vertexRegion ?? "not-provided"
@@ -79,8 +81,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		}
 
 		// Use vertexBaseUrl if this is a Vertex handler, otherwise use googleGeminiBaseUrl
-		const baseUrl =
-			this.constructor.name === "VertexHandler" ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
+		const baseUrl = this.isVertex ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
 
 		const config: GenerateContentConfig = {
 			systemInstruction,
@@ -225,10 +226,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 				tools.push({ googleSearch: {} })
 			}
 			// Use vertexBaseUrl if this is a Vertex handler, otherwise use googleGeminiBaseUrl
-			const baseUrl =
-				this.constructor.name === "VertexHandler"
-					? this.options.vertexBaseUrl
-					: this.options.googleGeminiBaseUrl
+			const baseUrl = this.isVertex ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
 
 			const promptConfig: GenerateContentConfig = {
 				httpOptions: baseUrl ? { baseUrl } : undefined,