feat: vertex/gemini prompt caching (#2996)

.changeset/curly-frogs-pull.md

@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+Use gemini provider when using Gemini on vertex ai

package-lock.json

@@ -12,7 +12,6 @@
 				"@anthropic-ai/sdk": "^0.37.0",
 				"@anthropic-ai/vertex-sdk": "^0.7.0",
 				"@aws-sdk/client-bedrock-runtime": "^3.779.0",
-				"@google-cloud/vertexai": "^1.9.3",
 				"@google/genai": "^0.9.0",
 				"@mistralai/mistralai": "^1.3.6",
 				"@modelcontextprotocol/sdk": "^1.7.0",
@@ -5772,18 +5771,6 @@
 				"node": "^12.22.0 || ^14.17.0 || >=16.0.0"
 			}
 		},
-		"node_modules/@google-cloud/vertexai": {
-			"version": "1.9.3",
-			"resolved": "https://registry.npmjs.org/@google-cloud/vertexai/-/vertexai-1.9.3.tgz",
-			"integrity": "sha512-35o5tIEMLW3JeFJOaaMNR2e5sq+6rpnhrF97PuAxeOm0GlqVTESKhkGj7a5B5mmJSSSU3hUfIhcQCRRsw4Ipzg==",
-			"license": "Apache-2.0",
-			"dependencies": {
-				"google-auth-library": "^9.1.0"
-			},
-			"engines": {
-				"node": ">=18.0.0"
-			}
-		},
 		"node_modules/@google/genai": {
 			"version": "0.9.0",
 			"resolved": "https://registry.npmjs.org/@google/genai/-/genai-0.9.0.tgz",

package.json

@@ -404,7 +404,6 @@
 		"@anthropic-ai/sdk": "^0.37.0",
 		"@anthropic-ai/vertex-sdk": "^0.7.0",
 		"@aws-sdk/client-bedrock-runtime": "^3.779.0",
-		"@google-cloud/vertexai": "^1.9.3",
 		"@google/genai": "^0.9.0",
 		"@mistralai/mistralai": "^1.3.6",
 		"@modelcontextprotocol/sdk": "^1.7.0",

src/api/index.ts

@@ -8,6 +8,7 @@ import { AnthropicHandler } from "./providers/anthropic"
 import { AwsBedrockHandler } from "./providers/bedrock"
 import { OpenRouterHandler } from "./providers/openrouter"
 import { VertexHandler } from "./providers/vertex"
+import { AnthropicVertexHandler } from "./providers/anthropic-vertex"
 import { OpenAiHandler } from "./providers/openai"
 import { OllamaHandler } from "./providers/ollama"
 import { LmStudioHandler } from "./providers/lmstudio"
@@ -45,6 +46,7 @@ export interface ApiHandler {
 
 export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
 	const { apiProvider, ...options } = configuration
+
 	switch (apiProvider) {
 		case "anthropic":
 			return new AnthropicHandler(options)
@@ -55,7 +57,11 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
 		case "bedrock":
 			return new AwsBedrockHandler(options)
 		case "vertex":
-			return new VertexHandler(options)
+			if (options.apiModelId?.startsWith("claude")) {
+				return new AnthropicVertexHandler(options)
+			} else {
+				return new VertexHandler(options)
+			}
 		case "openai":
 			return new OpenAiHandler(options)
 		case "ollama":

src/api/providers/__tests__/anthropic-vertex.test.ts

@@ -0,0 +1,816 @@
+// npx jest src/api/providers/__tests__/anthropic-vertex.test.ts
+
+import { Anthropic } from "@anthropic-ai/sdk"
+import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
+
+import { ApiStreamChunk } from "../../transform/stream"
+
+import { AnthropicVertexHandler } from "../anthropic-vertex"
+
+jest.mock("@anthropic-ai/vertex-sdk", () => ({
+	AnthropicVertex: jest.fn().mockImplementation(() => ({
+		messages: {
+			create: jest.fn().mockImplementation(async (options) => {
+				if (!options.stream) {
+					return {
+						id: "test-completion",
+						content: [{ type: "text", text: "Test response" }],
+						role: "assistant",
+						model: options.model,
+						usage: {
+							input_tokens: 10,
+							output_tokens: 5,
+						},
+					}
+				}
+				return {
+					async *[Symbol.asyncIterator]() {
+						yield {
+							type: "message_start",
+							message: {
+								usage: {
+									input_tokens: 10,
+									output_tokens: 5,
+								},
+							},
+						}
+						yield {
+							type: "content_block_start",
+							content_block: {
+								type: "text",
+								text: "Test response",
+							},
+						}
+					},
+				}
+			}),
+		},
+	})),
+}))
+
+describe("VertexHandler", () => {
+	let handler: AnthropicVertexHandler
+
+	describe("constructor", () => {
+		it("should initialize with provided config for Claude", () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			expect(AnthropicVertex).toHaveBeenCalledWith({
+				projectId: "test-project",
+				region: "us-central1",
+			})
+		})
+	})
+
+	describe("createMessage", () => {
+		const mockMessages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: "Hello",
+			},
+			{
+				role: "assistant",
+				content: "Hi there!",
+			},
+		]
+
+		const systemPrompt = "You are a helpful assistant"
+
+		it("should handle streaming responses correctly for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "message_start",
+					message: {
+						usage: {
+							input_tokens: 10,
+							output_tokens: 0,
+						},
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "text",
+						text: "Hello",
+					},
+				},
+				{
+					type: "content_block_delta",
+					delta: {
+						type: "text_delta",
+						text: " world!",
+					},
+				},
+				{
+					type: "message_delta",
+					usage: {
+						output_tokens: 5,
+					},
+				},
+			]
+
+			// Setup async iterator for mock stream
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+			const chunks: ApiStreamChunk[] = []
+
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks.length).toBe(4)
+			expect(chunks[0]).toEqual({
+				type: "usage",
+				inputTokens: 10,
+				outputTokens: 0,
+			})
+			expect(chunks[1]).toEqual({
+				type: "text",
+				text: "Hello",
+			})
+			expect(chunks[2]).toEqual({
+				type: "text",
+				text: " world!",
+			})
+			expect(chunks[3]).toEqual({
+				type: "usage",
+				inputTokens: 0,
+				outputTokens: 5,
+			})
+
+			expect(mockCreate).toHaveBeenCalledWith({
+				model: "claude-3-5-sonnet-v2@20241022",
+				max_tokens: 8192,
+				temperature: 0,
+				system: [
+					{
+						type: "text",
+						text: "You are a helpful assistant",
+						cache_control: { type: "ephemeral" },
+					},
+				],
+				messages: [
+					{
+						role: "user",
+						content: [
+							{
+								type: "text",
+								text: "Hello",
+								cache_control: { type: "ephemeral" },
+							},
+						],
+					},
+					{
+						role: "assistant",
+						content: "Hi there!",
+					},
+				],
+				stream: true,
+			})
+		})
+
+		it("should handle multiple content blocks with line breaks for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "text",
+						text: "First line",
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 1,
+					content_block: {
+						type: "text",
+						text: "Second line",
+					},
+				},
+			]
+
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+			const chunks: ApiStreamChunk[] = []
+
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks.length).toBe(3)
+			expect(chunks[0]).toEqual({
+				type: "text",
+				text: "First line",
+			})
+			expect(chunks[1]).toEqual({
+				type: "text",
+				text: "\n",
+			})
+			expect(chunks[2]).toEqual({
+				type: "text",
+				text: "Second line",
+			})
+		})
+
+		it("should handle API errors for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockError = new Error("Vertex API error")
+			const mockCreate = jest.fn().mockRejectedValue(mockError)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+
+			await expect(async () => {
+				for await (const _chunk of stream) {
+					// Should throw before yielding any chunks
+				}
+			}).rejects.toThrow("Vertex API error")
+		})
+
+		it("should handle prompt caching for supported models for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "message_start",
+					message: {
+						usage: {
+							input_tokens: 10,
+							output_tokens: 0,
+							cache_creation_input_tokens: 3,
+							cache_read_input_tokens: 2,
+						},
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "text",
+						text: "Hello",
+					},
+				},
+				{
+					type: "content_block_delta",
+					delta: {
+						type: "text_delta",
+						text: " world!",
+					},
+				},
+				{
+					type: "message_delta",
+					usage: {
+						output_tokens: 5,
+					},
+				},
+			]
+
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, [
+				{
+					role: "user",
+					content: "First message",
+				},
+				{
+					role: "assistant",
+					content: "Response",
+				},
+				{
+					role: "user",
+					content: "Second message",
+				},
+			])
+
+			const chunks: ApiStreamChunk[] = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			// Verify usage information
+			const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
+			expect(usageChunks).toHaveLength(2)
+			expect(usageChunks[0]).toEqual({
+				type: "usage",
+				inputTokens: 10,
+				outputTokens: 0,
+				cacheWriteTokens: 3,
+				cacheReadTokens: 2,
+			})
+			expect(usageChunks[1]).toEqual({
+				type: "usage",
+				inputTokens: 0,
+				outputTokens: 5,
+			})
+
+			// Verify text content
+			const textChunks = chunks.filter((chunk) => chunk.type === "text")
+			expect(textChunks).toHaveLength(2)
+			expect(textChunks[0].text).toBe("Hello")
+			expect(textChunks[1].text).toBe(" world!")
+
+			// Verify cache control was added correctly
+			expect(mockCreate).toHaveBeenCalledWith(
+				expect.objectContaining({
+					system: [
+						{
+							type: "text",
+							text: "You are a helpful assistant",
+							cache_control: { type: "ephemeral" },
+						},
+					],
+					messages: [
+						expect.objectContaining({
+							role: "user",
+							content: [
+								{
+									type: "text",
+									text: "First message",
+									cache_control: { type: "ephemeral" },
+								},
+							],
+						}),
+						expect.objectContaining({
+							role: "assistant",
+							content: "Response",
+						}),
+						expect.objectContaining({
+							role: "user",
+							content: [
+								{
+									type: "text",
+									text: "Second message",
+									cache_control: { type: "ephemeral" },
+								},
+							],
+						}),
+					],
+				}),
+			)
+		})
+
+		it("should handle cache-related usage metrics for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "message_start",
+					message: {
+						usage: {
+							input_tokens: 10,
+							output_tokens: 0,
+							cache_creation_input_tokens: 5,
+							cache_read_input_tokens: 3,
+						},
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "text",
+						text: "Hello",
+					},
+				},
+			]
+
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+			const chunks: ApiStreamChunk[] = []
+
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			// Check for cache-related metrics in usage chunk
+			const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
+			expect(usageChunks.length).toBeGreaterThan(0)
+			expect(usageChunks[0]).toHaveProperty("cacheWriteTokens", 5)
+			expect(usageChunks[0]).toHaveProperty("cacheReadTokens", 3)
+		})
+	})
+
+	describe("thinking functionality", () => {
+		const mockMessages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: "Hello",
+			},
+		]
+
+		const systemPrompt = "You are a helpful assistant"
+
+		it("should handle thinking content blocks and deltas for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "message_start",
+					message: {
+						usage: {
+							input_tokens: 10,
+							output_tokens: 0,
+						},
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "thinking",
+						thinking: "Let me think about this...",
+					},
+				},
+				{
+					type: "content_block_delta",
+					delta: {
+						type: "thinking_delta",
+						thinking: " I need to consider all options.",
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 1,
+					content_block: {
+						type: "text",
+						text: "Here's my answer:",
+					},
+				},
+			]
+
+			// Setup async iterator for mock stream
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+			const chunks: ApiStreamChunk[] = []
+
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			// Verify thinking content is processed correctly
+			const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
+			expect(reasoningChunks).toHaveLength(2)
+			expect(reasoningChunks[0].text).toBe("Let me think about this...")
+			expect(reasoningChunks[1].text).toBe(" I need to consider all options.")
+
+			// Verify text content is processed correctly
+			const textChunks = chunks.filter((chunk) => chunk.type === "text")
+			expect(textChunks).toHaveLength(2) // One for the text block, one for the newline
+			expect(textChunks[0].text).toBe("\n")
+			expect(textChunks[1].text).toBe("Here's my answer:")
+		})
+
+		it("should handle multiple thinking blocks with line breaks for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockStream = [
+				{
+					type: "content_block_start",
+					index: 0,
+					content_block: {
+						type: "thinking",
+						thinking: "First thinking block",
+					},
+				},
+				{
+					type: "content_block_start",
+					index: 1,
+					content_block: {
+						type: "thinking",
+						thinking: "Second thinking block",
+					},
+				},
+			]
+
+			const asyncIterator = {
+				async *[Symbol.asyncIterator]() {
+					for (const chunk of mockStream) {
+						yield chunk
+					}
+				},
+			}
+
+			const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
+			;(handler["client"].messages as any).create = mockCreate
+
+			const stream = handler.createMessage(systemPrompt, mockMessages)
+			const chunks: ApiStreamChunk[] = []
+
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks.length).toBe(3)
+			expect(chunks[0]).toEqual({
+				type: "reasoning",
+				text: "First thinking block",
+			})
+			expect(chunks[1]).toEqual({
+				type: "reasoning",
+				text: "\n",
+			})
+			expect(chunks[2]).toEqual({
+				type: "reasoning",
+				text: "Second thinking block",
+			})
+		})
+	})
+
+	describe("completePrompt", () => {
+		it("should complete prompt successfully for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const result = await handler.completePrompt("Test prompt")
+			expect(result).toBe("Test response")
+			expect(handler["client"].messages.create).toHaveBeenCalledWith({
+				model: "claude-3-5-sonnet-v2@20241022",
+				max_tokens: 8192,
+				temperature: 0,
+				messages: [
+					{
+						role: "user",
+						content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }],
+					},
+				],
+				stream: false,
+			})
+		})
+
+		it("should handle API errors for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockError = new Error("Vertex API error")
+			const mockCreate = jest.fn().mockRejectedValue(mockError)
+			;(handler["client"].messages as any).create = mockCreate
+
+			await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
+				"Vertex completion error: Vertex API error",
+			)
+		})
+
+		it("should handle non-text content for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockCreate = jest.fn().mockResolvedValue({
+				content: [{ type: "image" }],
+			})
+			;(handler["client"].messages as any).create = mockCreate
+
+			const result = await handler.completePrompt("Test prompt")
+			expect(result).toBe("")
+		})
+
+		it("should handle empty response for Claude", async () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockCreate = jest.fn().mockResolvedValue({
+				content: [{ type: "text", text: "" }],
+			})
+			;(handler["client"].messages as any).create = mockCreate
+
+			const result = await handler.completePrompt("Test prompt")
+			expect(result).toBe("")
+		})
+	})
+
+	describe("getModel", () => {
+		it("should return correct model info for Claude", () => {
+			handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const modelInfo = handler.getModel()
+			expect(modelInfo.id).toBe("claude-3-5-sonnet-v2@20241022")
+			expect(modelInfo.info).toBeDefined()
+			expect(modelInfo.info.maxTokens).toBe(8192)
+			expect(modelInfo.info.contextWindow).toBe(200_000)
+		})
+
+		it("honors custom maxTokens for thinking models", () => {
+			const handler = new AnthropicVertexHandler({
+				apiKey: "test-api-key",
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				modelMaxTokens: 32_768,
+				modelMaxThinkingTokens: 16_384,
+			})
+
+			const result = handler.getModel()
+			expect(result.maxTokens).toBe(32_768)
+			expect(result.thinking).toEqual({ type: "enabled", budget_tokens: 16_384 })
+			expect(result.temperature).toBe(1.0)
+		})
+
+		it("does not honor custom maxTokens for non-thinking models", () => {
+			const handler = new AnthropicVertexHandler({
+				apiKey: "test-api-key",
+				apiModelId: "claude-3-7-sonnet@20250219",
+				modelMaxTokens: 32_768,
+				modelMaxThinkingTokens: 16_384,
+			})
+
+			const result = handler.getModel()
+			expect(result.maxTokens).toBe(8192)
+			expect(result.thinking).toBeUndefined()
+			expect(result.temperature).toBe(0)
+		})
+	})
+
+	describe("thinking model configuration", () => {
+		it("should configure thinking for models with :thinking suffix", () => {
+			const thinkingHandler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+				modelMaxTokens: 16384,
+				modelMaxThinkingTokens: 4096,
+			})
+
+			const modelInfo = thinkingHandler.getModel()
+
+			// Verify thinking configuration
+			expect(modelInfo.id).toBe("claude-3-7-sonnet@20250219")
+			expect(modelInfo.thinking).toBeDefined()
+			const thinkingConfig = modelInfo.thinking as { type: "enabled"; budget_tokens: number }
+			expect(thinkingConfig.type).toBe("enabled")
+			expect(thinkingConfig.budget_tokens).toBe(4096)
+			expect(modelInfo.temperature).toBe(1.0) // Thinking requires temperature 1.0
+		})
+
+		it("should calculate thinking budget correctly", () => {
+			// Test with explicit thinking budget
+			const handlerWithBudget = new AnthropicVertexHandler({
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+				modelMaxTokens: 16384,
+				modelMaxThinkingTokens: 5000,
+			})
+
+			expect((handlerWithBudget.getModel().thinking as any).budget_tokens).toBe(5000)
+
+			// Test with default thinking budget (80% of max tokens)
+			const handlerWithDefaultBudget = new AnthropicVertexHandler({
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+				modelMaxTokens: 10000,
+			})
+
+			expect((handlerWithDefaultBudget.getModel().thinking as any).budget_tokens).toBe(8000) // 80% of 10000
+
+			// Test with minimum thinking budget (should be at least 1024)
+			const handlerWithSmallMaxTokens = new AnthropicVertexHandler({
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+				modelMaxTokens: 1000, // This would result in 800 tokens for thinking, but minimum is 1024
+			})
+
+			expect((handlerWithSmallMaxTokens.getModel().thinking as any).budget_tokens).toBe(1024)
+		})
+
+		it("should pass thinking configuration to API", async () => {
+			const thinkingHandler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-7-sonnet@20250219:thinking",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+				modelMaxTokens: 16384,
+				modelMaxThinkingTokens: 4096,
+			})
+
+			const mockCreate = jest.fn().mockImplementation(async (options) => {
+				if (!options.stream) {
+					return {
+						id: "test-completion",
+						content: [{ type: "text", text: "Test response" }],
+						role: "assistant",
+						model: options.model,
+						usage: { input_tokens: 10, output_tokens: 5 },
+					}
+				}
+				return {
+					async *[Symbol.asyncIterator]() {
+						yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }
+					},
+				}
+			})
+			;(thinkingHandler["client"].messages as any).create = mockCreate
+
+			await thinkingHandler
+				.createMessage("You are a helpful assistant", [{ role: "user", content: "Hello" }])
+				.next()
+
+			expect(mockCreate).toHaveBeenCalledWith(
+				expect.objectContaining({
+					thinking: { type: "enabled", budget_tokens: 4096 },
+					temperature: 1.0, // Thinking requires temperature 1.0
+				}),
+			)
+		})
+	})
+})

src/api/providers/anthropic-vertex.ts

@@ -0,0 +1,227 @@
+import { Anthropic } from "@anthropic-ai/sdk"
+import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
+import { GoogleAuth, JWTInput } from "google-auth-library"
+
+import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api"
+import { ApiStream } from "../transform/stream"
+import { safeJsonParse } from "../../shared/safeJsonParse"
+
+import { getModelParams, SingleCompletionHandler } from "../index"
+import { BaseProvider } from "./base-provider"
+import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./constants"
+import { formatMessageForCache } from "../transform/vertex-caching"
+
+// https://docs.anthropic.com/en/api/claude-on-vertex-ai
+export class AnthropicVertexHandler extends BaseProvider implements SingleCompletionHandler {
+	protected options: ApiHandlerOptions
+	private client: AnthropicVertex
+
+	constructor(options: ApiHandlerOptions) {
+		super()
+
+		this.options = options
+
+		// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
+		const projectId = this.options.vertexProjectId ?? "not-provided"
+		const region = this.options.vertexRegion ?? "us-east5"
+
+		if (this.options.vertexJsonCredentials) {
+			this.client = new AnthropicVertex({
+				projectId,
+				region,
+				googleAuth: new GoogleAuth({
+					scopes: ["https://www.googleapis.com/auth/cloud-platform"],
+					credentials: safeJsonParse<JWTInput>(this.options.vertexJsonCredentials, undefined),
+				}),
+			})
+		} else if (this.options.vertexKeyFile) {
+			this.client = new AnthropicVertex({
+				projectId,
+				region,
+				googleAuth: new GoogleAuth({
+					scopes: ["https://www.googleapis.com/auth/cloud-platform"],
+					keyFile: this.options.vertexKeyFile,
+				}),
+			})
+		} else {
+			this.client = new AnthropicVertex({ projectId, region })
+		}
+	}
+
+	override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
+		let {
+			id,
+			info: { supportsPromptCache },
+			temperature,
+			maxTokens,
+			thinking,
+		} = this.getModel()
+
+		// Find indices of user messages that we want to cache
+		// We only cache the last two user messages to stay within the 4-block limit
+		// (1 block for system + 1 block each for last two user messages = 3 total)
+		const userMsgIndices = supportsPromptCache
+			? messages.reduce((acc, msg, i) => (msg.role === "user" ? [...acc, i] : acc), [] as number[])
+			: []
+
+		const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
+		const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
+
+		/**
+		 * Vertex API has specific limitations for prompt caching:
+		 * 1. Maximum of 4 blocks can have cache_control
+		 * 2. Only text blocks can be cached (images and other content types cannot)
+		 * 3. Cache control can only be applied to user messages, not assistant messages
+		 *
+		 * Our caching strategy:
+		 * - Cache the system prompt (1 block)
+		 * - Cache the last text block of the second-to-last user message (1 block)
+		 * - Cache the last text block of the last user message (1 block)
+		 * This ensures we stay under the 4-block limit while maintaining effective caching
+		 * for the most relevant context.
+		 */
+		const params: Anthropic.Messages.MessageCreateParamsStreaming = {
+			model: id,
+			max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
+			temperature,
+			thinking,
+			// Cache the system prompt if caching is enabled.
+			system: supportsPromptCache
+				? [{ text: systemPrompt, type: "text" as const, cache_control: { type: "ephemeral" } }]
+				: systemPrompt,
+			messages: messages.map((message, index) => {
+				// Only cache the last two user messages.
+				const shouldCache =
+					supportsPromptCache && (index === lastUserMsgIndex || index === secondLastMsgUserIndex)
+				return formatMessageForCache(message, shouldCache)
+			}),
+			stream: true,
+		}
+
+		const stream = await this.client.messages.create(params)
+
+		for await (const chunk of stream) {
+			switch (chunk.type) {
+				case "message_start": {
+					const usage = chunk.message!.usage
+
+					yield {
+						type: "usage",
+						inputTokens: usage.input_tokens || 0,
+						outputTokens: usage.output_tokens || 0,
+						cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
+						cacheReadTokens: usage.cache_read_input_tokens || undefined,
+					}
+
+					break
+				}
+				case "message_delta": {
+					yield {
+						type: "usage",
+						inputTokens: 0,
+						outputTokens: chunk.usage!.output_tokens || 0,
+					}
+
+					break
+				}
+				case "content_block_start": {
+					switch (chunk.content_block!.type) {
+						case "text": {
+							if (chunk.index! > 0) {
+								yield { type: "text", text: "\n" }
+							}
+
+							yield { type: "text", text: chunk.content_block!.text }
+							break
+						}
+						case "thinking": {
+							if (chunk.index! > 0) {
+								yield { type: "reasoning", text: "\n" }
+							}
+
+							yield { type: "reasoning", text: (chunk.content_block as any).thinking }
+							break
+						}
+					}
+
+					break
+				}
+				case "content_block_delta": {
+					switch (chunk.delta!.type) {
+						case "text_delta": {
+							yield { type: "text", text: chunk.delta!.text }
+							break
+						}
+						case "thinking_delta": {
+							yield { type: "reasoning", text: (chunk.delta as any).thinking }
+							break
+						}
+					}
+
+					break
+				}
+			}
+		}
+	}
+
+	getModel() {
+		const modelId = this.options.apiModelId
+		let id = modelId && modelId in vertexModels ? (modelId as VertexModelId) : vertexDefaultModelId
+		const info: ModelInfo = vertexModels[id]
+
+		// The `:thinking` variant is a virtual identifier for thinking-enabled
+		// models (similar to how it's handled in the Anthropic provider.)
+		if (id.endsWith(":thinking")) {
+			id = id.replace(":thinking", "") as VertexModelId
+		}
+
+		return {
+			id,
+			info,
+			...getModelParams({ options: this.options, model: info, defaultMaxTokens: ANTHROPIC_DEFAULT_MAX_TOKENS }),
+		}
+	}
+
+	async completePrompt(prompt: string) {
+		try {
+			let {
+				id,
+				info: { supportsPromptCache },
+				temperature,
+				maxTokens = ANTHROPIC_DEFAULT_MAX_TOKENS,
+				thinking,
+			} = this.getModel()
+
+			const params: Anthropic.Messages.MessageCreateParamsNonStreaming = {
+				model: id,
+				max_tokens: maxTokens,
+				temperature,
+				thinking,
+				messages: [
+					{
+						role: "user",
+						content: supportsPromptCache
+							? [{ type: "text" as const, text: prompt, cache_control: { type: "ephemeral" } }]
+							: prompt,
+					},
+				],
+				stream: false,
+			}
+
+			const response = await this.client.messages.create(params)
+			const content = response.content[0]
+
+			if (content.type === "text") {
+				return content.text
+			}
+
+			return ""
+		} catch (error) {
+			if (error instanceof Error) {
+				throw new Error(`Vertex completion error: ${error.message}`)
+			}
+
+			throw error
+		}
+	}
+}

src/api/providers/gemini.ts

@@ -5,11 +5,13 @@ import {
 	type GenerateContentParameters,
 	type Content,
 } from "@google/genai"
+import type { JWTInput } from "google-auth-library"
 import NodeCache from "node-cache"
 
-import { SingleCompletionHandler } from "../"
-import type { ApiHandlerOptions, GeminiModelId, ModelInfo } from "../../shared/api"
-import { geminiDefaultModelId, geminiModels } from "../../shared/api"
+import { ApiHandlerOptions, ModelInfo, GeminiModelId, geminiDefaultModelId, geminiModels } from "../../shared/api"
+import { safeJsonParse } from "../../shared/safeJsonParse"
+
+import { SingleCompletionHandler } from "../index"
 import {
 	convertAnthropicContentToGemini,
 	convertAnthropicMessageToGemini,
@@ -27,6 +29,10 @@ type CacheEntry = {
 	count: number
 }
 
+type GeminiHandlerOptions = ApiHandlerOptions & {
+	isVertex?: boolean
+}
+
 export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
 
@@ -34,10 +40,35 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 	private contentCaches: NodeCache
 	private isCacheBusy = false
 
-	constructor(options: ApiHandlerOptions) {
+	constructor({ isVertex, ...options }: GeminiHandlerOptions) {
 		super()
+
 		this.options = options
-		this.client = new GoogleGenAI({ apiKey: options.geminiApiKey ?? "not-provided" })
+
+		const project = this.options.vertexProjectId ?? "not-provided"
+		const location = this.options.vertexRegion ?? "not-provided"
+		const apiKey = this.options.geminiApiKey ?? "not-provided"
+
+		this.client = this.options.vertexJsonCredentials
+			? new GoogleGenAI({
+					vertexai: true,
+					project,
+					location,
+					googleAuthOptions: {
+						credentials: safeJsonParse<JWTInput>(this.options.vertexJsonCredentials, undefined),
+					},
+				})
+			: this.options.vertexKeyFile
+				? new GoogleGenAI({
+						vertexai: true,
+						project,
+						location,
+						googleAuthOptions: { keyFile: this.options.vertexKeyFile },
+					})
+				: isVertex
+					? new GoogleGenAI({ vertexai: true, project, location })
+					: new GoogleGenAI({ apiKey })
+
 		this.contentCaches = new NodeCache({ stdTTL: 5 * 60, checkperiod: 5 * 60 })
 	}
 
@@ -170,14 +201,14 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 	}
 
 	override getModel() {
-		let id = this.options.apiModelId ? (this.options.apiModelId as GeminiModelId) : geminiDefaultModelId
-		let info: ModelInfo = geminiModels[id]
+		let id = this.options.apiModelId ?? geminiDefaultModelId
+		let info: ModelInfo = geminiModels[id as GeminiModelId]
 
 		if (id?.endsWith(":thinking")) {
-			id = id.slice(0, -":thinking".length) as GeminiModelId
+			id = id.slice(0, -":thinking".length)
 
-			if (geminiModels[id]) {
-				info = geminiModels[id]
+			if (geminiModels[id as GeminiModelId]) {
+				info = geminiModels[id as GeminiModelId]
 
 				return {
 					id,

src/api/providers/vertex.ts

@@ -1,490 +1,39 @@
-import { Anthropic } from "@anthropic-ai/sdk"
-import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
-import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming"
+import { ApiHandlerOptions, ModelInfo, VertexModelId, vertexDefaultModelId, vertexModels } from "../../shared/api"
 
-import { VertexAI } from "@google-cloud/vertexai"
-
-import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api"
-import { ApiStream } from "../transform/stream"
-import { convertAnthropicMessageToVertexGemini } from "../transform/vertex-gemini-format"
-import { BaseProvider } from "./base-provider"
-
-import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./constants"
-import { getModelParams, SingleCompletionHandler } from "../"
-import { GoogleAuth } from "google-auth-library"
-
-// Types for Vertex SDK
-
-/**
- * Vertex API has specific limitations for prompt caching:
- * 1. Maximum of 4 blocks can have cache_control
- * 2. Only text blocks can be cached (images and other content types cannot)
- * 3. Cache control can only be applied to user messages, not assistant messages
- *
- * Our caching strategy:
- * - Cache the system prompt (1 block)
- * - Cache the last text block of the second-to-last user message (1 block)
- * - Cache the last text block of the last user message (1 block)
- * This ensures we stay under the 4-block limit while maintaining effective caching
- * for the most relevant context.
- */
-
-interface VertexTextBlock {
-	type: "text"
-	text: string
-	cache_control?: { type: "ephemeral" }
-}
-
-interface VertexImageBlock {
-	type: "image"
-	source: {
-		type: "base64"
-		media_type: "image/jpeg" | "image/png" | "image/gif" | "image/webp"
-		data: string
-	}
-}
-
-type VertexContentBlock = VertexTextBlock | VertexImageBlock
-
-interface VertexUsage {
-	input_tokens?: number
-	output_tokens?: number
-	cache_creation_input_tokens?: number
-	cache_read_input_tokens?: number
-}
-
-interface VertexMessage extends Omit<Anthropic.Messages.MessageParam, "content"> {
-	content: string | VertexContentBlock[]
-}
-
-interface VertexMessageResponse {
-	content: Array<{ type: "text"; text: string }>
-}
-
-interface VertexMessageStreamEvent {
-	type: "message_start" | "message_delta" | "content_block_start" | "content_block_delta"
-	message?: {
-		usage: VertexUsage
-	}
-	usage?: {
-		output_tokens: number
-	}
-	content_block?:
-		| {
-				type: "text"
-				text: string
-		  }
-		| {
-				type: "thinking"
-				thinking: string
-		  }
-	index?: number
-	delta?:
-		| {
-				type: "text_delta"
-				text: string
-		  }
-		| {
-				type: "thinking_delta"
-				thinking: string
-		  }
-}
-
-// https://docs.anthropic.com/en/api/claude-on-vertex-ai
-export class VertexHandler extends BaseProvider implements SingleCompletionHandler {
-	MODEL_CLAUDE = "claude"
-	MODEL_GEMINI = "gemini"
-
-	protected options: ApiHandlerOptions
-	private anthropicClient: AnthropicVertex
-	private geminiClient: VertexAI
-	private modelType: string
+import { SingleCompletionHandler } from "../index"
+import { GeminiHandler } from "./gemini"
 
+export class VertexHandler extends GeminiHandler implements SingleCompletionHandler {
 	constructor(options: ApiHandlerOptions) {
-		super()
-		this.options = options
-
-		if (this.options.apiModelId?.startsWith(this.MODEL_CLAUDE)) {
-			this.modelType = this.MODEL_CLAUDE
-		} else if (this.options.apiModelId?.startsWith(this.MODEL_GEMINI)) {
-			this.modelType = this.MODEL_GEMINI
-		} else {
-			throw new Error(`Unknown model ID: ${this.options.apiModelId}`)
-		}
-
-		if (this.options.vertexJsonCredentials) {
-			this.anthropicClient = new AnthropicVertex({
-				projectId: this.options.vertexProjectId ?? "not-provided",
-				// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
-				region: this.options.vertexRegion ?? "us-east5",
-				googleAuth: new GoogleAuth({
-					scopes: ["https://www.googleapis.com/auth/cloud-platform"],
-					credentials: JSON.parse(this.options.vertexJsonCredentials),
-				}),
-			})
-		} else if (this.options.vertexKeyFile) {
-			this.anthropicClient = new AnthropicVertex({
-				projectId: this.options.vertexProjectId ?? "not-provided",
-				// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
-				region: this.options.vertexRegion ?? "us-east5",
-				googleAuth: new GoogleAuth({
-					scopes: ["https://www.googleapis.com/auth/cloud-platform"],
-					keyFile: this.options.vertexKeyFile,
-				}),
-			})
-		} else {
-			this.anthropicClient = new AnthropicVertex({
-				projectId: this.options.vertexProjectId ?? "not-provided",
-				// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
-				region: this.options.vertexRegion ?? "us-east5",
-			})
-		}
-
-		if (this.options.vertexJsonCredentials) {
-			this.geminiClient = new VertexAI({
-				project: this.options.vertexProjectId ?? "not-provided",
-				location: this.options.vertexRegion ?? "us-east5",
-				googleAuthOptions: {
-					credentials: JSON.parse(this.options.vertexJsonCredentials),
-				},
-			})
-		} else if (this.options.vertexKeyFile) {
-			this.geminiClient = new VertexAI({
-				project: this.options.vertexProjectId ?? "not-provided",
-				location: this.options.vertexRegion ?? "us-east5",
-				googleAuthOptions: {
-					keyFile: this.options.vertexKeyFile,
-				},
-			})
-		} else {
-			this.geminiClient = new VertexAI({
-				project: this.options.vertexProjectId ?? "not-provided",
-				location: this.options.vertexRegion ?? "us-east5",
-			})
-		}
+		super({ ...options, isVertex: true })
 	}
 
-	private formatMessageForCache(message: Anthropic.Messages.MessageParam, shouldCache: boolean): VertexMessage {
-		// Assistant messages are kept as-is since they can't be cached
-		if (message.role === "assistant") {
-			return message as VertexMessage
-		}
+	override getModel() {
+		let id = this.options.apiModelId ?? vertexDefaultModelId
+		let info: ModelInfo = vertexModels[id as VertexModelId]
 
-		// For string content, we convert to array format with optional cache control
-		if (typeof message.content === "string") {
-			return {
-				...message,
-				content: [
-					{
-						type: "text" as const,
-						text: message.content,
-						// For string content, we only have one block so it's always the last
-						...(shouldCache && { cache_control: { type: "ephemeral" } }),
-					},
-				],
-			}
-		}
+		if (id?.endsWith(":thinking")) {
+			id = id.slice(0, -":thinking".length) as VertexModelId
 
-		// For array content, find the last text block index once before mapping
-		const lastTextBlockIndex = message.content.reduce(
-			(lastIndex, content, index) => (content.type === "text" ? index : lastIndex),
-			-1,
-		)
-
-		// Then use this pre-calculated index in the map function
-		return {
-			...message,
-			content: message.content.map((content, contentIndex) => {
-				// Images and other non-text content are passed through unchanged
-				if (content.type === "image") {
-					return content as VertexImageBlock
-				}
-
-				// Check if this is the last text block using our pre-calculated index
-				const isLastTextBlock = contentIndex === lastTextBlockIndex
+			if (vertexModels[id as VertexModelId]) {
+				info = vertexModels[id as VertexModelId]
 
 				return {
-					type: "text" as const,
-					text: (content as { text: string }).text,
-					...(shouldCache && isLastTextBlock && { cache_control: { type: "ephemeral" } }),
-				}
-			}),
-		}
-	}
-
-	private async *createGeminiMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
-		const model = this.geminiClient.getGenerativeModel({
-			model: this.getModel().id,
-			systemInstruction: systemPrompt,
-		})
-
-		const result = await model.generateContentStream({
-			contents: messages.map(convertAnthropicMessageToVertexGemini),
-			generationConfig: {
-				maxOutputTokens: this.getModel().info.maxTokens ?? undefined,
-				temperature: this.options.modelTemperature ?? 0,
-			},
-		})
-
-		for await (const chunk of result.stream) {
-			if (chunk.candidates?.[0]?.content?.parts) {
-				for (const part of chunk.candidates[0].content.parts) {
-					if (part.text) {
-						yield {
-							type: "text",
-							text: part.text,
-						}
-					}
+					id,
+					info,
+					thinkingConfig: this.options.modelMaxThinkingTokens
+						? { thinkingBudget: this.options.modelMaxThinkingTokens }
+						: undefined,
+					maxOutputTokens: this.options.modelMaxTokens ?? info.maxTokens ?? undefined,
 				}
 			}
 		}
 
-		const response = await result.response
-
-		yield {
-			type: "usage",
-			inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
-			outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
-		}
-	}
-
-	private async *createClaudeMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
-		const model = this.getModel()
-		let { id, temperature, maxTokens, thinking } = model
-		const useCache = model.info.supportsPromptCache
-
-		// Find indices of user messages that we want to cache
-		// We only cache the last two user messages to stay within the 4-block limit
-		// (1 block for system + 1 block each for last two user messages = 3 total)
-		const userMsgIndices = useCache
-			? messages.reduce((acc, msg, i) => (msg.role === "user" ? [...acc, i] : acc), [] as number[])
-			: []
-		const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
-		const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
-
-		// Create the stream with appropriate caching configuration
-		const params = {
-			model: id,
-			max_tokens: maxTokens,
-			temperature,
-			thinking,
-			// Cache the system prompt if caching is enabled
-			system: useCache
-				? [
-						{
-							text: systemPrompt,
-							type: "text" as const,
-							cache_control: { type: "ephemeral" },
-						},
-					]
-				: systemPrompt,
-			messages: messages.map((message, index) => {
-				// Only cache the last two user messages
-				const shouldCache = useCache && (index === lastUserMsgIndex || index === secondLastMsgUserIndex)
-				return this.formatMessageForCache(message, shouldCache)
-			}),
-			stream: true,
+		if (!info) {
+			id = vertexDefaultModelId
+			info = vertexModels[vertexDefaultModelId]
 		}
 
-		const stream = (await this.anthropicClient.messages.create(
-			params as Anthropic.Messages.MessageCreateParamsStreaming,
-		)) as unknown as AnthropicStream<VertexMessageStreamEvent>
-
-		// Process the stream chunks
-		for await (const chunk of stream) {
-			switch (chunk.type) {
-				case "message_start": {
-					const usage = chunk.message!.usage
-					yield {
-						type: "usage",
-						inputTokens: usage.input_tokens || 0,
-						outputTokens: usage.output_tokens || 0,
-						cacheWriteTokens: usage.cache_creation_input_tokens,
-						cacheReadTokens: usage.cache_read_input_tokens,
-					}
-					break
-				}
-				case "message_delta": {
-					yield {
-						type: "usage",
-						inputTokens: 0,
-						outputTokens: chunk.usage!.output_tokens || 0,
-					}
-					break
-				}
-				case "content_block_start": {
-					switch (chunk.content_block!.type) {
-						case "text": {
-							if (chunk.index! > 0) {
-								yield {
-									type: "text",
-									text: "\n",
-								}
-							}
-							yield {
-								type: "text",
-								text: chunk.content_block!.text,
-							}
-							break
-						}
-						case "thinking": {
-							if (chunk.index! > 0) {
-								yield {
-									type: "reasoning",
-									text: "\n",
-								}
-							}
-							yield {
-								type: "reasoning",
-								text: (chunk.content_block as any).thinking,
-							}
-							break
-						}
-					}
-					break
-				}
-				case "content_block_delta": {
-					switch (chunk.delta!.type) {
-						case "text_delta": {
-							yield {
-								type: "text",
-								text: chunk.delta!.text,
-							}
-							break
-						}
-						case "thinking_delta": {
-							yield {
-								type: "reasoning",
-								text: (chunk.delta as any).thinking,
-							}
-							break
-						}
-					}
-					break
-				}
-			}
-		}
-	}
-
-	override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
-		switch (this.modelType) {
-			case this.MODEL_CLAUDE: {
-				yield* this.createClaudeMessage(systemPrompt, messages)
-				break
-			}
-			case this.MODEL_GEMINI: {
-				yield* this.createGeminiMessage(systemPrompt, messages)
-				break
-			}
-			default: {
-				throw new Error(`Invalid model type: ${this.modelType}`)
-			}
-		}
-	}
-
-	getModel() {
-		const modelId = this.options.apiModelId
-		let id = modelId && modelId in vertexModels ? (modelId as VertexModelId) : vertexDefaultModelId
-		const info: ModelInfo = vertexModels[id]
-
-		// The `:thinking` variant is a virtual identifier for thinking-enabled
-		// models (similar to how it's handled in the Anthropic provider.)
-		if (id.endsWith(":thinking")) {
-			id = id.replace(":thinking", "") as VertexModelId
-		}
-
-		return {
-			id,
-			info,
-			...getModelParams({ options: this.options, model: info, defaultMaxTokens: ANTHROPIC_DEFAULT_MAX_TOKENS }),
-		}
-	}
-
-	private async completePromptGemini(prompt: string) {
-		try {
-			const model = this.geminiClient.getGenerativeModel({
-				model: this.getModel().id,
-			})
-
-			const result = await model.generateContent({
-				contents: [{ role: "user", parts: [{ text: prompt }] }],
-				generationConfig: {
-					temperature: this.options.modelTemperature ?? 0,
-				},
-			})
-
-			let text = ""
-			result.response.candidates?.forEach((candidate) => {
-				candidate.content.parts.forEach((part) => {
-					text += part.text
-				})
-			})
-
-			return text
-		} catch (error) {
-			if (error instanceof Error) {
-				throw new Error(`Vertex completion error: ${error.message}`)
-			}
-			throw error
-		}
-	}
-
-	private async completePromptClaude(prompt: string) {
-		try {
-			let { id, info, temperature, maxTokens, thinking } = this.getModel()
-			const useCache = info.supportsPromptCache
-
-			const params: Anthropic.Messages.MessageCreateParamsNonStreaming = {
-				model: id,
-				max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
-				temperature,
-				thinking,
-				system: "", // No system prompt needed for single completions
-				messages: [
-					{
-						role: "user",
-						content: useCache
-							? [
-									{
-										type: "text" as const,
-										text: prompt,
-										cache_control: { type: "ephemeral" },
-									},
-								]
-							: prompt,
-					},
-				],
-				stream: false,
-			}
-
-			const response = (await this.anthropicClient.messages.create(params)) as unknown as VertexMessageResponse
-			const content = response.content[0]
-
-			if (content.type === "text") {
-				return content.text
-			}
-
-			return ""
-		} catch (error) {
-			if (error instanceof Error) {
-				throw new Error(`Vertex completion error: ${error.message}`)
-			}
-
-			throw error
-		}
-	}
-
-	async completePrompt(prompt: string) {
-		switch (this.modelType) {
-			case this.MODEL_CLAUDE: {
-				return this.completePromptClaude(prompt)
-			}
-			case this.MODEL_GEMINI: {
-				return this.completePromptGemini(prompt)
-			}
-			default: {
-				throw new Error(`Invalid model type: ${this.modelType}`)
-			}
-		}
+		return { id, info }
 	}
 }

src/api/transform/__tests__/vertex-gemini-format.test.ts

@@ -1,338 +0,0 @@
-// npx jest src/api/transform/__tests__/vertex-gemini-format.test.ts
-
-import { Anthropic } from "@anthropic-ai/sdk"
-
-import { convertAnthropicMessageToVertexGemini } from "../vertex-gemini-format"
-
-describe("convertAnthropicMessageToVertexGemini", () => {
-	it("should convert a simple text message", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: "Hello, world!",
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [{ text: "Hello, world!" }],
-		})
-	})
-
-	it("should convert assistant role to model role", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "assistant",
-			content: "I'm an assistant",
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "model",
-			parts: [{ text: "I'm an assistant" }],
-		})
-	})
-
-	it("should convert a message with text blocks", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{ type: "text", text: "First paragraph" },
-				{ type: "text", text: "Second paragraph" },
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [{ text: "First paragraph" }, { text: "Second paragraph" }],
-		})
-	})
-
-	it("should convert a message with an image", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{ type: "text", text: "Check out this image:" },
-				{
-					type: "image",
-					source: {
-						type: "base64",
-						media_type: "image/jpeg",
-						data: "base64encodeddata",
-					},
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [
-				{ text: "Check out this image:" },
-				{
-					inlineData: {
-						data: "base64encodeddata",
-						mimeType: "image/jpeg",
-					},
-				},
-			],
-		})
-	})
-
-	it("should throw an error for unsupported image source type", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "image",
-					source: {
-						type: "url", // Not supported
-						url: "https://example.com/image.jpg",
-					} as any,
-				},
-			],
-		}
-
-		expect(() => convertAnthropicMessageToVertexGemini(anthropicMessage)).toThrow("Unsupported image source type")
-	})
-
-	it("should convert a message with tool use", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "assistant",
-			content: [
-				{ type: "text", text: "Let me calculate that for you." },
-				{
-					type: "tool_use",
-					id: "calc-123",
-					name: "calculator",
-					input: { operation: "add", numbers: [2, 3] },
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "model",
-			parts: [
-				{ text: "Let me calculate that for you." },
-				{
-					functionCall: {
-						name: "calculator",
-						args: { operation: "add", numbers: [2, 3] },
-					},
-				},
-			],
-		})
-	})
-
-	it("should convert a message with tool result as string", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{ type: "text", text: "Here's the result:" },
-				{
-					type: "tool_result",
-					tool_use_id: "calculator-123",
-					content: "The result is 5",
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [
-				{ text: "Here's the result:" },
-				{
-					functionResponse: {
-						name: "calculator",
-						response: {
-							name: "calculator",
-							content: "The result is 5",
-						},
-					},
-				},
-			],
-		})
-	})
-
-	it("should handle empty tool result content", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "tool_result",
-					tool_use_id: "calculator-123",
-					content: null as any, // Empty content
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		// Should skip the empty tool result
-		expect(result).toEqual({
-			role: "user",
-			parts: [],
-		})
-	})
-
-	it("should convert a message with tool result as array with text only", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "tool_result",
-					tool_use_id: "search-123",
-					content: [
-						{ type: "text", text: "First result" },
-						{ type: "text", text: "Second result" },
-					],
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [
-				{
-					functionResponse: {
-						name: "search",
-						response: {
-							name: "search",
-							content: "First result\n\nSecond result",
-						},
-					},
-				},
-			],
-		})
-	})
-
-	it("should convert a message with tool result as array with text and images", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "tool_result",
-					tool_use_id: "search-123",
-					content: [
-						{ type: "text", text: "Search results:" },
-						{
-							type: "image",
-							source: {
-								type: "base64",
-								media_type: "image/png",
-								data: "image1data",
-							},
-						},
-						{
-							type: "image",
-							source: {
-								type: "base64",
-								media_type: "image/jpeg",
-								data: "image2data",
-							},
-						},
-					],
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [
-				{
-					functionResponse: {
-						name: "search",
-						response: {
-							name: "search",
-							content: "Search results:\n\n(See next part for image)",
-						},
-					},
-				},
-				{
-					inlineData: {
-						data: "image1data",
-						mimeType: "image/png",
-					},
-				},
-				{
-					inlineData: {
-						data: "image2data",
-						mimeType: "image/jpeg",
-					},
-				},
-			],
-		})
-	})
-
-	it("should convert a message with tool result containing only images", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "tool_result",
-					tool_use_id: "imagesearch-123",
-					content: [
-						{
-							type: "image",
-							source: {
-								type: "base64",
-								media_type: "image/png",
-								data: "onlyimagedata",
-							},
-						},
-					],
-				},
-			],
-		}
-
-		const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
-
-		expect(result).toEqual({
-			role: "user",
-			parts: [
-				{
-					functionResponse: {
-						name: "imagesearch",
-						response: {
-							name: "imagesearch",
-							content: "\n\n(See next part for image)",
-						},
-					},
-				},
-				{
-					inlineData: {
-						data: "onlyimagedata",
-						mimeType: "image/png",
-					},
-				},
-			],
-		})
-	})
-
-	it("should throw an error for unsupported content block type", () => {
-		const anthropicMessage: Anthropic.Messages.MessageParam = {
-			role: "user",
-			content: [
-				{
-					type: "unknown_type", // Unsupported type
-					data: "some data",
-				} as any,
-			],
-		}
-
-		expect(() => convertAnthropicMessageToVertexGemini(anthropicMessage)).toThrow(
-			"Unsupported content block type: unknown_type",
-		)
-	})
-})

src/api/transform/vertex-caching.ts

@@ -0,0 +1,70 @@
+import { Anthropic } from "@anthropic-ai/sdk"
+
+interface VertexTextBlock {
+	type: "text"
+	text: string
+	cache_control?: { type: "ephemeral" }
+}
+
+interface VertexImageBlock {
+	type: "image"
+	source: {
+		type: "base64"
+		media_type: "image/jpeg" | "image/png" | "image/gif" | "image/webp"
+		data: string
+	}
+}
+
+type VertexContentBlock = VertexTextBlock | VertexImageBlock
+
+interface VertexMessage extends Omit<Anthropic.Messages.MessageParam, "content"> {
+	content: string | VertexContentBlock[]
+}
+
+export function formatMessageForCache(message: Anthropic.Messages.MessageParam, shouldCache: boolean): VertexMessage {
+	// Assistant messages are kept as-is since they can't be cached
+	if (message.role === "assistant") {
+		return message as VertexMessage
+	}
+
+	// For string content, we convert to array format with optional cache control
+	if (typeof message.content === "string") {
+		return {
+			...message,
+			content: [
+				{
+					type: "text" as const,
+					text: message.content,
+					// For string content, we only have one block so it's always the last
+					...(shouldCache && { cache_control: { type: "ephemeral" } }),
+				},
+			],
+		}
+	}
+
+	// For array content, find the last text block index once before mapping
+	const lastTextBlockIndex = message.content.reduce(
+		(lastIndex, content, index) => (content.type === "text" ? index : lastIndex),
+		-1,
+	)
+
+	// Then use this pre-calculated index in the map function.
+	return {
+		...message,
+		content: message.content.map((content, contentIndex) => {
+			// Images and other non-text content are passed through unchanged.
+			if (content.type === "image") {
+				return content as VertexImageBlock
+			}
+
+			// Check if this is the last text block using our pre-calculated index.
+			const isLastTextBlock = contentIndex === lastTextBlockIndex
+
+			return {
+				type: "text" as const,
+				text: (content as { text: string }).text,
+				...(shouldCache && isLastTextBlock && { cache_control: { type: "ephemeral" } }),
+			}
+		}),
+	}
+}

src/api/transform/vertex-gemini-format.ts

@@ -1,83 +0,0 @@
-import { Anthropic } from "@anthropic-ai/sdk"
-import { Content, FunctionCallPart, FunctionResponsePart, InlineDataPart, Part, TextPart } from "@google-cloud/vertexai"
-
-function convertAnthropicContentToVertexGemini(content: Anthropic.Messages.MessageParam["content"]): Part[] {
-	if (typeof content === "string") {
-		return [{ text: content } as TextPart]
-	}
-
-	return content.flatMap((block) => {
-		switch (block.type) {
-			case "text":
-				return { text: block.text } as TextPart
-			case "image":
-				if (block.source.type !== "base64") {
-					throw new Error("Unsupported image source type")
-				}
-				return {
-					inlineData: {
-						data: block.source.data,
-						mimeType: block.source.media_type,
-					},
-				} as InlineDataPart
-			case "tool_use":
-				return {
-					functionCall: {
-						name: block.name,
-						args: block.input,
-					},
-				} as FunctionCallPart
-			case "tool_result":
-				const name = block.tool_use_id.split("-")[0]
-				if (!block.content) {
-					return []
-				}
-				if (typeof block.content === "string") {
-					return {
-						functionResponse: {
-							name,
-							response: {
-								name,
-								content: block.content,
-							},
-						},
-					} as FunctionResponsePart
-				} else {
-					// The only case when tool_result could be array is when the tool failed and we're providing ie user feedback potentially with images
-					const textParts = block.content.filter((part) => part.type === "text")
-					const imageParts = block.content.filter((part) => part.type === "image")
-					const text = textParts.length > 0 ? textParts.map((part) => part.text).join("\n\n") : ""
-					const imageText = imageParts.length > 0 ? "\n\n(See next part for image)" : ""
-					return [
-						{
-							functionResponse: {
-								name,
-								response: {
-									name,
-									content: text + imageText,
-								},
-							},
-						} as FunctionResponsePart,
-						...imageParts.map(
-							(part) =>
-								({
-									inlineData: {
-										data: part.source.data,
-										mimeType: part.source.media_type,
-									},
-								}) as InlineDataPart,
-						),
-					]
-				}
-			default:
-				throw new Error(`Unsupported content block type: ${(block as any).type}`)
-		}
-	})
-}
-
-export function convertAnthropicMessageToVertexGemini(message: Anthropic.Messages.MessageParam): Content {
-	return {
-		role: message.role === "assistant" ? "model" : "user",
-		parts: convertAnthropicContentToVertexGemini(message.content),
-	}
-}

src/core/Cline.ts

@@ -953,12 +953,14 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 		if (mcpEnabled ?? true) {
 			const provider = this.providerRef.deref()
+
 			if (!provider) {
 				throw new Error("Provider reference lost during view transition")
 			}
 
 			// Wait for MCP hub initialization through McpServerManager
 			mcpHub = await McpServerManager.getInstance(provider.context, provider)
+
 			if (!mcpHub) {
 				throw new Error("Failed to get MCP hub from server manager")
 			}
@@ -980,12 +982,16 @@ export class Cline extends EventEmitter<ClineEvents> {
 			browserToolEnabled,
 			language,
 		} = (await this.providerRef.deref()?.getState()) ?? {}
+
 		const { customModes } = (await this.providerRef.deref()?.getState()) ?? {}
+
 		const systemPrompt = await (async () => {
 			const provider = this.providerRef.deref()
+
 			if (!provider) {
 				throw new Error("Provider not available")
 			}
+
 			return SYSTEM_PROMPT(
 				provider.context,
 				this.cwd,
@@ -1008,7 +1014,10 @@ export class Cline extends EventEmitter<ClineEvents> {
 		// If the previous API request's total token usage is close to the context window, truncate the conversation history to free up space for the new request
 		if (previousApiReqIndex >= 0) {
 			const previousRequest = this.clineMessages[previousApiReqIndex]?.text
-			if (!previousRequest) return
+
+			if (!previousRequest) {
+				return
+			}
 
 			const {
 				tokensIn = 0,
@@ -1135,11 +1144,14 @@ export class Cline extends EventEmitter<ClineEvents> {
 					"api_req_failed",
 					error.message ?? JSON.stringify(serializeError(error), null, 2),
 				)
+
 				if (response !== "yesButtonClicked") {
 					// this will never happen since if noButtonClicked, we will clear current task, aborting this instance
 					throw new Error("API request failed")
 				}
+
 				await this.say("api_req_retried")
+
 				// delegate generator output from the recursive call
 				yield* this.attemptApiRequest(previousApiReqIndex)
 				return
@@ -1903,8 +1915,13 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 			return didEndLoop // will always be false for now
 		} catch (error) {
-			// this should never happen since the only thing that can throw an error is the attemptApiRequest, which is wrapped in a try catch that sends an ask where if noButtonClicked, will clear current task and destroy this instance. However to avoid unhandled promise rejection, we will end this loop which will end execution of this instance (see startTask)
-			return true // needs to be true so parent loop knows to end task
+			// This should never happen since the only thing that can throw an
+			// error is the attemptApiRequest, which is wrapped in a try catch
+			// that sends an ask where if noButtonClicked, will clear current
+			// task and destroy this instance. However to avoid unhandled
+			// promise rejection, we will end this loop which will end execution
+			// of this instance (see `startTask`).
+			return true // Needs to be true so parent loop knows to end task.
 		}
 	}
 

src/shared/api.ts

@@ -497,7 +497,8 @@ export const vertexModels = {
 		maxTokens: 65_535,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
+		isPromptCacheOptional: true,
 		inputPrice: 2.5,
 		outputPrice: 15,
 	},
@@ -521,7 +522,8 @@ export const vertexModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
+		isPromptCacheOptional: true,
 		inputPrice: 0.15,
 		outputPrice: 0.6,
 	},
@@ -545,7 +547,8 @@ export const vertexModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
+		isPromptCacheOptional: true,
 		inputPrice: 0.075,
 		outputPrice: 0.3,
 	},

src/shared/safeJsonParse.ts

@@ -1,11 +1,14 @@
 /**
  * Safely parses JSON without crashing on invalid input
+ *
  * @param jsonString The string to parse
  * @param defaultValue Value to return if parsing fails
  * @returns Parsed JSON object or defaultValue if parsing fails
  */
 export function safeJsonParse<T>(jsonString: string | null | undefined, defaultValue?: T): T | undefined {
-	if (!jsonString) return defaultValue
+	if (!jsonString) {
+		return defaultValue
+	}
 
 	try {
 		return JSON.parse(jsonString) as T

webview-ui/src/components/chat/ChatRow.tsx

@@ -6,9 +6,9 @@ import { VSCodeBadge, VSCodeButton } from "@vscode/webview-ui-toolkit/react"
 
 import { ClineApiReqInfo, ClineAskUseMcpServer, ClineMessage, ClineSayTool } from "@roo/shared/ExtensionMessage"
 import { splitCommandOutput, COMMAND_OUTPUT_STRING } from "@roo/shared/combineCommandSequences"
+import { safeJsonParse } from "@roo/shared/safeJsonParse"
 
 import { useCopyToClipboard } from "@src/utils/clipboard"
-import { safeJsonParse } from "@src/utils/json"
 import { useExtensionState } from "@src/context/ExtensionStateContext"
 import { findMatchingResourceOrTemplate } from "@src/utils/mcp"
 import { vscode } from "@src/utils/vscode"