fix: Handle long Claude code messages (#5072)

packages/types/src/providers/claude-code.ts

@@ -5,9 +5,29 @@ import { anthropicModels } from "./anthropic.js"
 export type ClaudeCodeModelId = keyof typeof claudeCodeModels
 export const claudeCodeDefaultModelId: ClaudeCodeModelId = "claude-sonnet-4-20250514"
 export const claudeCodeModels = {
-	"claude-sonnet-4-20250514": anthropicModels["claude-sonnet-4-20250514"],
-	"claude-opus-4-20250514": anthropicModels["claude-opus-4-20250514"],
-	"claude-3-7-sonnet-20250219": anthropicModels["claude-3-7-sonnet-20250219"],
-	"claude-3-5-sonnet-20241022": anthropicModels["claude-3-5-sonnet-20241022"],
-	"claude-3-5-haiku-20241022": anthropicModels["claude-3-5-haiku-20241022"],
+	"claude-sonnet-4-20250514": {
+		...anthropicModels["claude-sonnet-4-20250514"],
+		supportsImages: false,
+		supportsPromptCache: false,
+	},
+	"claude-opus-4-20250514": {
+		...anthropicModels["claude-opus-4-20250514"],
+		supportsImages: false,
+		supportsPromptCache: false,
+	},
+	"claude-3-7-sonnet-20250219": {
+		...anthropicModels["claude-3-7-sonnet-20250219"],
+		supportsImages: false,
+		supportsPromptCache: false,
+	},
+	"claude-3-5-sonnet-20241022": {
+		...anthropicModels["claude-3-5-sonnet-20241022"],
+		supportsImages: false,
+		supportsPromptCache: false,
+	},
+	"claude-3-5-haiku-20241022": {
+		...anthropicModels["claude-3-5-haiku-20241022"],
+		supportsImages: false,
+		supportsPromptCache: false,
+	},
 } as const satisfies Record<string, ModelInfo>

src/activate/__tests__/registerCommands.spec.ts

@@ -16,6 +16,15 @@ vi.mock("vscode", () => ({
 	window: {
 		createTextEditorDecorationType: vi.fn().mockReturnValue({ dispose: vi.fn() }),
 	},
+	workspace: {
+		workspaceFolders: [
+			{
+				uri: {
+					fsPath: "/mock/workspace",
+				},
+			},
+		],
+	},
 }))
 
 vi.mock("../../core/webview/ClineProvider")

src/api/providers/__tests__/claude-code.spec.ts

@@ -1,230 +1,503 @@
 import { describe, test, expect, vi, beforeEach } from "vitest"
 import { ClaudeCodeHandler } from "../claude-code"
 import { ApiHandlerOptions } from "../../../shared/api"
+import { ClaudeCodeMessage } from "../../../integrations/claude-code/types"
 
 // Mock the runClaudeCode function
 vi.mock("../../../integrations/claude-code/run", () => ({
 	runClaudeCode: vi.fn(),
 }))
 
+// Mock the message filter
+vi.mock("../../../integrations/claude-code/message-filter", () => ({
+	filterMessagesForClaudeCode: vi.fn((messages) => messages),
+}))
+
 const { runClaudeCode } = await import("../../../integrations/claude-code/run")
+const { filterMessagesForClaudeCode } = await import("../../../integrations/claude-code/message-filter")
 const mockRunClaudeCode = vi.mocked(runClaudeCode)
-
-// Mock the EventEmitter for the process
-class MockEventEmitter {
-	private handlers: { [event: string]: ((...args: any[]) => void)[] } = {}
-
-	on(event: string, handler: (...args: any[]) => void) {
-		if (!this.handlers[event]) {
-			this.handlers[event] = []
-		}
-		this.handlers[event].push(handler)
-	}
-
-	emit(event: string, ...args: any[]) {
-		if (this.handlers[event]) {
-			this.handlers[event].forEach((handler) => handler(...args))
-		}
-	}
-}
+const mockFilterMessages = vi.mocked(filterMessagesForClaudeCode)
 
 describe("ClaudeCodeHandler", () => {
 	let handler: ClaudeCodeHandler
-	let mockProcess: any
 
 	beforeEach(() => {
+		vi.clearAllMocks()
 		const options: ApiHandlerOptions = {
 			claudeCodePath: "claude",
 			apiModelId: "claude-3-5-sonnet-20241022",
 		}
 		handler = new ClaudeCodeHandler(options)
+	})
 
-		const mainEmitter = new MockEventEmitter()
-		mockProcess = {
-			stdout: new MockEventEmitter(),
-			stderr: new MockEventEmitter(),
-			on: mainEmitter.on.bind(mainEmitter),
-			emit: mainEmitter.emit.bind(mainEmitter),
+	test("should create handler with correct model configuration", () => {
+		const model = handler.getModel()
+		expect(model.id).toBe("claude-3-5-sonnet-20241022")
+		expect(model.info.supportsImages).toBe(false)
+		expect(model.info.supportsPromptCache).toBe(false)
+	})
+
+	test("should use default model when invalid model provided", () => {
+		const options: ApiHandlerOptions = {
+			claudeCodePath: "claude",
+			apiModelId: "invalid-model",
 		}
+		const handlerWithInvalidModel = new ClaudeCodeHandler(options)
+		const model = handlerWithInvalidModel.getModel()
 
-		mockRunClaudeCode.mockReturnValue(mockProcess)
+		expect(model.id).toBe("claude-sonnet-4-20250514") // default model
 	})
 
-	test("should handle thinking content properly", async () => {
+	test("should filter messages and call runClaudeCode", async () => {
 		const systemPrompt = "You are a helpful assistant"
 		const messages = [{ role: "user" as const, content: "Hello" }]
+		const filteredMessages = [{ role: "user" as const, content: "Hello (filtered)" }]
+
+		mockFilterMessages.mockReturnValue(filteredMessages)
+
+		// Mock empty async generator
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			// Empty generator for basic test
+		}
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
 
-		// Start the stream
 		const stream = handler.createMessage(systemPrompt, messages)
-		const streamGenerator = stream[Symbol.asyncIterator]()
-
-		// Simulate thinking content response
-		const thinkingResponse = {
-			type: "assistant",
-			message: {
-				id: "msg_123",
-				type: "message",
-				role: "assistant",
-				model: "claude-3-5-sonnet-20241022",
-				content: [
-					{
-						type: "thinking",
-						thinking: "I need to think about this carefully...",
-						signature: "abc123",
-					},
-				],
-				stop_reason: null,
-				stop_sequence: null,
-				usage: {
-					input_tokens: 10,
-					output_tokens: 20,
-					service_tier: "standard" as const,
-				},
-			},
-			session_id: "session_123",
-		}
-
-		// Emit the thinking response and wait for processing
-		setImmediate(() => {
-			mockProcess.stdout.emit("data", JSON.stringify(thinkingResponse) + "\n")
-			setImmediate(() => {
-				mockProcess.emit("close", 0)
-			})
+
+		// Need to start iterating to trigger the call
+		const iterator = stream[Symbol.asyncIterator]()
+		await iterator.next()
+
+		// Verify message filtering was called
+		expect(mockFilterMessages).toHaveBeenCalledWith(messages)
+
+		// Verify runClaudeCode was called with filtered messages
+		expect(mockRunClaudeCode).toHaveBeenCalledWith({
+			systemPrompt,
+			messages: filteredMessages,
+			path: "claude",
+			modelId: "claude-3-5-sonnet-20241022",
 		})
+	})
+
+	test("should handle thinking content properly", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+
+		// Mock async generator that yields thinking content
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "thinking",
+							thinking: "I need to think about this carefully...",
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const results = []
 
-		// Get the result
-		const result = await streamGenerator.next()
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
 
-		expect(result.done).toBe(false)
-		expect(result.value).toEqual({
+		expect(results).toHaveLength(1)
+		expect(results[0]).toEqual({
 			type: "reasoning",
 			text: "I need to think about this carefully...",
 		})
 	})
 
-	test("should handle mixed content types", async () => {
+	test("should handle redacted thinking content", async () => {
 		const systemPrompt = "You are a helpful assistant"
 		const messages = [{ role: "user" as const, content: "Hello" }]
 
-		const stream = handler.createMessage(systemPrompt, messages)
-		const streamGenerator = stream[Symbol.asyncIterator]()
-
-		// Simulate mixed content response
-		const mixedResponse = {
-			type: "assistant",
-			message: {
-				id: "msg_123",
-				type: "message",
-				role: "assistant",
-				model: "claude-3-5-sonnet-20241022",
-				content: [
-					{
-						type: "thinking",
-						thinking: "Let me think about this...",
+		// Mock async generator that yields redacted thinking content
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "redacted_thinking",
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
 					},
-					{
-						type: "text",
-						text: "Here's my response!",
-					},
-				],
-				stop_reason: null,
-				stop_sequence: null,
-				usage: {
-					input_tokens: 10,
-					output_tokens: 20,
-					service_tier: "standard" as const,
-				},
-			},
-			session_id: "session_123",
-		}
-
-		// Emit the mixed response and wait for processing
-		setImmediate(() => {
-			mockProcess.stdout.emit("data", JSON.stringify(mixedResponse) + "\n")
-			setImmediate(() => {
-				mockProcess.emit("close", 0)
-			})
+				} as any,
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const results = []
+
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		expect(results).toHaveLength(1)
+		expect(results[0]).toEqual({
+			type: "reasoning",
+			text: "[Redacted thinking block]",
 		})
+	})
+
+	test("should handle mixed content types", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+
+		// Mock async generator that yields mixed content
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "thinking",
+							thinking: "Let me think about this...",
+						},
+						{
+							type: "text",
+							text: "Here's my response!",
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+		}
 
-		// Get the first result (thinking)
-		const thinkingResult = await streamGenerator.next()
-		expect(thinkingResult.done).toBe(false)
-		expect(thinkingResult.value).toEqual({
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const results = []
+
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		expect(results).toHaveLength(2)
+		expect(results[0]).toEqual({
 			type: "reasoning",
 			text: "Let me think about this...",
 		})
-
-		// Get the second result (text)
-		const textResult = await streamGenerator.next()
-		expect(textResult.done).toBe(false)
-		expect(textResult.value).toEqual({
+		expect(results[1]).toEqual({
 			type: "text",
 			text: "Here's my response!",
 		})
 	})
 
-	test("should handle stop_reason with thinking content in error messages", async () => {
+	test("should handle string chunks from generator", async () => {
 		const systemPrompt = "You are a helpful assistant"
 		const messages = [{ role: "user" as const, content: "Hello" }]
 
+		// Mock async generator that yields string chunks
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield "This is a string chunk"
+			yield "Another string chunk"
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
 		const stream = handler.createMessage(systemPrompt, messages)
-		const streamGenerator = stream[Symbol.asyncIterator]()
-
-		// Simulate error response with thinking content
-		const errorResponse = {
-			type: "assistant",
-			message: {
-				id: "msg_123",
-				type: "message",
-				role: "assistant",
-				model: "claude-3-5-sonnet-20241022",
-				content: [
-					{
-						type: "thinking",
-						thinking: "This is an error scenario",
-					},
-				],
-				stop_reason: "max_tokens",
-				stop_sequence: null,
-				usage: {
-					input_tokens: 10,
-					output_tokens: 20,
-					service_tier: "standard" as const,
-				},
-			},
-			session_id: "session_123",
-		}
-
-		// Emit the error response and wait for processing
-		setImmediate(() => {
-			mockProcess.stdout.emit("data", JSON.stringify(errorResponse) + "\n")
-			setImmediate(() => {
-				mockProcess.emit("close", 0)
-			})
+		const results = []
+
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		expect(results).toHaveLength(2)
+		expect(results[0]).toEqual({
+			type: "text",
+			text: "This is a string chunk",
+		})
+		expect(results[1]).toEqual({
+			type: "text",
+			text: "Another string chunk",
 		})
+	})
+
+	test("should handle usage and cost tracking with paid usage", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+
+		// Mock async generator with init, assistant, and result messages
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			// Init message indicating paid usage
+			yield {
+				type: "system" as const,
+				subtype: "init" as const,
+				session_id: "session_123",
+				tools: [],
+				mcp_servers: [],
+				apiKeySource: "/login managed key",
+			}
+
+			// Assistant message
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "text",
+							text: "Hello there!",
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+						cache_read_input_tokens: 5,
+						cache_creation_input_tokens: 3,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+
+			// Result message
+			yield {
+				type: "result" as const,
+				subtype: "success" as const,
+				total_cost_usd: 0.05,
+				is_error: false,
+				duration_ms: 1000,
+				duration_api_ms: 800,
+				num_turns: 1,
+				result: "success",
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const results = []
 
-		// Should throw error with thinking content
-		await expect(streamGenerator.next()).rejects.toThrow("This is an error scenario")
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		// Should have text chunk and usage chunk
+		expect(results).toHaveLength(2)
+		expect(results[0]).toEqual({
+			type: "text",
+			text: "Hello there!",
+		})
+		expect(results[1]).toEqual({
+			type: "usage",
+			inputTokens: 10,
+			outputTokens: 20,
+			cacheReadTokens: 5,
+			cacheWriteTokens: 3,
+			totalCost: 0.05, // Paid usage, so cost is included
+		})
 	})
 
-	test("should handle incomplete JSON in buffer on process close", async () => {
+	test("should handle usage tracking with subscription (free) usage", async () => {
 		const systemPrompt = "You are a helpful assistant"
 		const messages = [{ role: "user" as const, content: "Hello" }]
 
+		// Mock async generator with subscription usage
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			// Init message indicating subscription usage
+			yield {
+				type: "system" as const,
+				subtype: "init" as const,
+				session_id: "session_123",
+				tools: [],
+				mcp_servers: [],
+				apiKeySource: "none", // Subscription usage
+			}
+
+			// Assistant message
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "text",
+							text: "Hello there!",
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+
+			// Result message
+			yield {
+				type: "result" as const,
+				subtype: "success" as const,
+				total_cost_usd: 0.05,
+				is_error: false,
+				duration_ms: 1000,
+				duration_api_ms: 800,
+				num_turns: 1,
+				result: "success",
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
 		const stream = handler.createMessage(systemPrompt, messages)
-		const streamGenerator = stream[Symbol.asyncIterator]()
-
-		// Simulate incomplete JSON data followed by process close
-		setImmediate(() => {
-			// Send incomplete JSON (missing closing brace)
-			mockProcess.stdout.emit("data", '{"type":"assistant","message":{"id":"msg_123"')
-			setImmediate(() => {
-				mockProcess.emit("close", 0)
-			})
+		const results = []
+
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		// Should have text chunk and usage chunk
+		expect(results).toHaveLength(2)
+		expect(results[0]).toEqual({
+			type: "text",
+			text: "Hello there!",
 		})
+		expect(results[1]).toEqual({
+			type: "usage",
+			inputTokens: 10,
+			outputTokens: 20,
+			cacheReadTokens: 0,
+			cacheWriteTokens: 0,
+			totalCost: 0, // Subscription usage, so cost is 0
+		})
+	})
+
+	test("should handle API errors properly", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+
+		// Mock async generator that yields an API error
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "text",
+							text: 'API Error: 400 {"error":{"message":"Invalid model name"}}',
+						},
+					],
+					stop_reason: "stop_sequence",
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const iterator = stream[Symbol.asyncIterator]()
+
+		// Should throw an error
+		await expect(iterator.next()).rejects.toThrow()
+	})
+
+	test("should log warning for unsupported tool_use content", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+		const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {})
+
+		// Mock async generator that yields tool_use content
+		const mockGenerator = async function* (): AsyncGenerator<ClaudeCodeMessage | string> {
+			yield {
+				type: "assistant" as const,
+				message: {
+					id: "msg_123",
+					type: "message",
+					role: "assistant",
+					model: "claude-3-5-sonnet-20241022",
+					content: [
+						{
+							type: "tool_use",
+							id: "tool_123",
+							name: "test_tool",
+							input: { test: "data" },
+						},
+					],
+					stop_reason: null,
+					stop_sequence: null,
+					usage: {
+						input_tokens: 10,
+						output_tokens: 20,
+					},
+				} as any,
+				session_id: "session_123",
+			}
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockGenerator())
+
+		const stream = handler.createMessage(systemPrompt, messages)
+		const results = []
+
+		for await (const chunk of stream) {
+			results.push(chunk)
+		}
+
+		// Should log error for unsupported tool_use
+		expect(consoleSpy).toHaveBeenCalledWith(expect.stringContaining("tool_use is not supported yet"))
 
-		// Should complete without throwing, incomplete JSON should be discarded
-		const result = await streamGenerator.next()
-		expect(result.done).toBe(true)
+		consoleSpy.mockRestore()
 	})
 })

src/api/providers/claude-code.ts

@@ -3,7 +3,7 @@ import { claudeCodeDefaultModelId, type ClaudeCodeModelId, claudeCodeModels } fr
 import { type ApiHandler } from ".."
 import { ApiStreamUsageChunk, type ApiStream } from "../transform/stream"
 import { runClaudeCode } from "../../integrations/claude-code/run"
-import { ClaudeCodeMessage } from "../../integrations/claude-code/types"
+import { filterMessagesForClaudeCode } from "../../integrations/claude-code/message-filter"
 import { BaseProvider } from "./base-provider"
 import { t } from "../../i18n"
 import { ApiHandlerOptions } from "../../shared/api"
@@ -17,61 +17,16 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 	}
 
 	override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
+		// Filter out image blocks since Claude Code doesn't support them
+		const filteredMessages = filterMessagesForClaudeCode(messages)
+
 		const claudeProcess = runClaudeCode({
 			systemPrompt,
-			messages,
+			messages: filteredMessages,
 			path: this.options.claudeCodePath,
 			modelId: this.getModel().id,
 		})
 
-		const dataQueue: string[] = []
-		let processError = null
-		let errorOutput = ""
-		let exitCode: number | null = null
-		let buffer = ""
-
-		claudeProcess.stdout.on("data", (data) => {
-			buffer += data.toString()
-			const lines = buffer.split("\n")
-
-			// Keep the last line in buffer as it might be incomplete
-			buffer = lines.pop() || ""
-
-			// Process complete lines
-			for (const line of lines) {
-				const trimmedLine = line.trim()
-				if (trimmedLine !== "") {
-					dataQueue.push(trimmedLine)
-				}
-			}
-		})
-
-		claudeProcess.stderr.on("data", (data) => {
-			errorOutput += data.toString()
-		})
-
-		claudeProcess.on("close", (code) => {
-			exitCode = code
-			// Process any remaining data in buffer
-			const trimmedBuffer = buffer.trim()
-			if (trimmedBuffer) {
-				// Validate that the remaining buffer looks like valid JSON before processing
-				if (this.isLikelyValidJSON(trimmedBuffer)) {
-					dataQueue.push(trimmedBuffer)
-				} else {
-					console.warn(
-						"Discarding incomplete JSON data on process close:",
-						trimmedBuffer.substring(0, 100) + (trimmedBuffer.length > 100 ? "..." : ""),
-					)
-				}
-				buffer = ""
-			}
-		})
-
-		claudeProcess.on("error", (error) => {
-			processError = error
-		})
-
 		// Usage is included with assistant messages,
 		// but cost is included in the result chunk
 		let usage: ApiStreamUsageChunk = {
@@ -82,72 +37,74 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 			cacheWriteTokens: 0,
 		}
 
-		while (exitCode !== 0 || dataQueue.length > 0) {
-			if (dataQueue.length === 0) {
-				await new Promise((resolve) => setImmediate(resolve))
-			}
+		let isPaidUsage = true
 
-			if (exitCode !== null && exitCode !== 0) {
-				if (errorOutput) {
-					throw new Error(
-						t("common:errors.claudeCode.processExitedWithError", {
-							exitCode,
-							output: errorOutput.trim(),
-						}),
-					)
-				}
-				throw new Error(t("common:errors.claudeCode.processExited", { exitCode }))
-			}
-
-			const data = dataQueue.shift()
-			if (!data) {
-				continue
-			}
-
-			const chunk = this.attemptParseChunk(data)
-
-			if (!chunk) {
+		for await (const chunk of claudeProcess) {
+			if (typeof chunk === "string") {
 				yield {
 					type: "text",
-					text: data || "",
+					text: chunk,
 				}
 
 				continue
 			}
 
 			if (chunk.type === "system" && chunk.subtype === "init") {
+				// Based on my tests, subscription usage sets the `apiKeySource` to "none"
+				isPaidUsage = chunk.apiKeySource !== "none"
 				continue
 			}
 
 			if (chunk.type === "assistant" && "message" in chunk) {
 				const message = chunk.message
 
-				if (message.stop_reason !== null && message.stop_reason !== "tool_use") {
-					const firstContent = message.content[0]
-					const errorMessage =
-						this.getContentText(firstContent) ||
-						t("common:errors.claudeCode.stoppedWithReason", { reason: message.stop_reason })
+				if (message.stop_reason !== null) {
+					const content = "text" in message.content[0] ? message.content[0] : undefined
 
-					if (errorMessage.includes("Invalid model name")) {
-						throw new Error(errorMessage + `\n\n${t("common:errors.claudeCode.apiKeyModelPlanMismatch")}`)
-					}
+					const isError = content && content.text.startsWith(`API Error`)
+					if (isError) {
+						// Error messages are formatted as: `API Error: <<status code>> <<json>>`
+						const errorMessageStart = content.text.indexOf("{")
+						const errorMessage = content.text.slice(errorMessageStart)
+
+						const error = this.attemptParse(errorMessage)
+						if (!error) {
+							throw new Error(content.text)
+						}
+
+						if (error.error.message.includes("Invalid model name")) {
+							throw new Error(
+								content.text + `\n\n${t("common:errors.claudeCode.apiKeyModelPlanMismatch")}`,
+							)
+						}
 
-					throw new Error(errorMessage)
+						throw new Error(errorMessage)
+					}
 				}
 
 				for (const content of message.content) {
-					if (content.type === "text") {
-						yield {
-							type: "text",
-							text: content.text,
-						}
-					} else if (content.type === "thinking") {
-						yield {
-							type: "reasoning",
-							text: content.thinking,
-						}
-					} else {
-						console.warn("Unsupported content type:", content)
+					switch (content.type) {
+						case "text":
+							yield {
+								type: "text",
+								text: content.text,
+							}
+							break
+						case "thinking":
+							yield {
+								type: "reasoning",
+								text: content.thinking || "",
+							}
+							break
+						case "redacted_thinking":
+							yield {
+								type: "reasoning",
+								text: "[Redacted thinking block]",
+							}
+							break
+						case "tool_use":
+							console.error(`tool_use is not supported yet. Received: ${JSON.stringify(content)}`)
+							break
 					}
 				}
 
@@ -161,16 +118,10 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 			}
 
 			if (chunk.type === "result" && "result" in chunk) {
-				// Only use the cost from the CLI if provided
-				// Don't calculate cost as it may be $0 for subscription users
-				usage.totalCost = chunk.cost_usd ?? 0
+				usage.totalCost = isPaidUsage ? chunk.total_cost_usd : 0
 
 				yield usage
 			}
-
-			if (processError) {
-				throw processError
-			}
 		}
 	}
 
@@ -187,53 +138,10 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 		}
 	}
 
-	private getContentText(content: any): string | undefined {
-		if (!content) return undefined
-		switch (content.type) {
-			case "text":
-				return content.text
-			case "thinking":
-				return content.thinking
-			default:
-				return undefined
-		}
-	}
-
-	private isLikelyValidJSON(data: string): boolean {
-		// Basic validation to check if the data looks like it could be valid JSON
-		const trimmed = data.trim()
-		if (!trimmed) return false
-
-		// Must start and end with appropriate JSON delimiters
-		const startsCorrectly = trimmed.startsWith("{") || trimmed.startsWith("[")
-		const endsCorrectly = trimmed.endsWith("}") || trimmed.endsWith("]")
-
-		if (!startsCorrectly || !endsCorrectly) return false
-
-		// Check for balanced braces/brackets (simple heuristic)
-		let braceCount = 0
-		let bracketCount = 0
-		for (const char of trimmed) {
-			if (char === "{") braceCount++
-			else if (char === "}") braceCount--
-			else if (char === "[") bracketCount++
-			else if (char === "]") bracketCount--
-		}
-
-		return braceCount === 0 && bracketCount === 0
-	}
-
-	// TODO: Validate instead of parsing
-	private attemptParseChunk(data: string): ClaudeCodeMessage | null {
+	private attemptParse(str: string) {
 		try {
-			return JSON.parse(data)
-		} catch (error) {
-			console.error(
-				"Error parsing chunk:",
-				error,
-				"Data:",
-				data.substring(0, 100) + (data.length > 100 ? "..." : ""),
-			)
+			return JSON.parse(str)
+		} catch (err) {
 			return null
 		}
 	}

src/integrations/claude-code/__tests__/message-filter.spec.ts

@@ -0,0 +1,263 @@
+import { describe, test, expect } from "vitest"
+import { filterMessagesForClaudeCode } from "../message-filter"
+import type { Anthropic } from "@anthropic-ai/sdk"
+
+describe("filterMessagesForClaudeCode", () => {
+	test("should pass through string messages unchanged", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: "Hello, this is a simple text message",
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual(messages)
+	})
+
+	test("should pass through text-only content blocks unchanged", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "This is a text block",
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual(messages)
+	})
+
+	test("should replace image blocks with text placeholders", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Here's an image:",
+					},
+					{
+						type: "image",
+						source: {
+							type: "base64",
+							media_type: "image/png",
+							data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+						},
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual([
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Here's an image:",
+					},
+					{
+						type: "text",
+						text: "[Image (base64): image/png not supported by Claude Code]",
+					},
+				],
+			},
+		])
+	})
+
+	test("should handle image blocks with unknown source types", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: [
+					{
+						type: "image",
+						source: undefined as any,
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual([
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "[Image (unknown): unknown not supported by Claude Code]",
+					},
+				],
+			},
+		])
+	})
+
+	test("should handle mixed content with multiple images", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Compare these images:",
+					},
+					{
+						type: "image",
+						source: {
+							type: "base64",
+							media_type: "image/jpeg",
+							data: "base64data1",
+						},
+					},
+					{
+						type: "text",
+						text: "and",
+					},
+					{
+						type: "image",
+						source: {
+							type: "base64",
+							media_type: "image/gif",
+							data: "base64data2",
+						},
+					},
+					{
+						type: "text",
+						text: "What do you think?",
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual([
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Compare these images:",
+					},
+					{
+						type: "text",
+						text: "[Image (base64): image/jpeg not supported by Claude Code]",
+					},
+					{
+						type: "text",
+						text: "and",
+					},
+					{
+						type: "text",
+						text: "[Image (base64): image/gif not supported by Claude Code]",
+					},
+					{
+						type: "text",
+						text: "What do you think?",
+					},
+				],
+			},
+		])
+	})
+
+	test("should handle multiple messages with images", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: "First message with text only",
+			},
+			{
+				role: "assistant",
+				content: [
+					{
+						type: "text",
+						text: "I can help with that.",
+					},
+				],
+			},
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Here's an image:",
+					},
+					{
+						type: "image",
+						source: {
+							type: "base64",
+							media_type: "image/png",
+							data: "imagedata",
+						},
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual([
+			{
+				role: "user",
+				content: "First message with text only",
+			},
+			{
+				role: "assistant",
+				content: [
+					{
+						type: "text",
+						text: "I can help with that.",
+					},
+				],
+			},
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Here's an image:",
+					},
+					{
+						type: "text",
+						text: "[Image (base64): image/png not supported by Claude Code]",
+					},
+				],
+			},
+		])
+	})
+
+	test("should preserve other content block types unchanged", () => {
+		const messages: Anthropic.Messages.MessageParam[] = [
+			{
+				role: "user",
+				content: [
+					{
+						type: "text",
+						text: "Regular text",
+					},
+					// This would be some other content type that's not an image
+					{
+						type: "tool_use" as any,
+						id: "tool_123",
+						name: "test_tool",
+						input: { test: "data" },
+					},
+				],
+			},
+		]
+
+		const result = filterMessagesForClaudeCode(messages)
+
+		expect(result).toEqual(messages)
+	})
+})

src/integrations/claude-code/__tests__/run.spec.ts

@@ -0,0 +1,37 @@
+import { describe, test, expect, vi, beforeEach } from "vitest"
+
+// Mock vscode workspace
+vi.mock("vscode", () => ({
+	workspace: {
+		workspaceFolders: [
+			{
+				uri: {
+					fsPath: "/test/workspace",
+				},
+			},
+		],
+	},
+}))
+
+describe("runClaudeCode", () => {
+	beforeEach(() => {
+		vi.clearAllMocks()
+	})
+
+	test("should export runClaudeCode function", async () => {
+		const { runClaudeCode } = await import("../run")
+		expect(typeof runClaudeCode).toBe("function")
+	})
+
+	test("should be an async generator function", async () => {
+		const { runClaudeCode } = await import("../run")
+		const options = {
+			systemPrompt: "You are a helpful assistant",
+			messages: [{ role: "user" as const, content: "Hello" }],
+		}
+
+		const result = runClaudeCode(options)
+		expect(Symbol.asyncIterator in result).toBe(true)
+		expect(typeof result[Symbol.asyncIterator]).toBe("function")
+	})
+})

src/integrations/claude-code/message-filter.ts

@@ -0,0 +1,35 @@
+import type { Anthropic } from "@anthropic-ai/sdk"
+
+/**
+ * Filters out image blocks from messages since Claude Code doesn't support images.
+ * Replaces image blocks with text placeholders similar to how VSCode LM provider handles it.
+ */
+export function filterMessagesForClaudeCode(
+	messages: Anthropic.Messages.MessageParam[],
+): Anthropic.Messages.MessageParam[] {
+	return messages.map((message) => {
+		// Handle simple string messages
+		if (typeof message.content === "string") {
+			return message
+		}
+
+		// Handle complex message structures
+		const filteredContent = message.content.map((block) => {
+			if (block.type === "image") {
+				// Replace image blocks with text placeholders
+				const sourceType = block.source?.type || "unknown"
+				const mediaType = block.source?.media_type || "unknown"
+				return {
+					type: "text" as const,
+					text: `[Image (${sourceType}): ${mediaType} not supported by Claude Code]`,
+				}
+			}
+			return block
+		})
+
+		return {
+			...message,
+			content: filteredContent,
+		}
+	})
+}

src/integrations/claude-code/run.ts

@@ -1,21 +1,115 @@
 import * as vscode from "vscode"
-import Anthropic from "@anthropic-ai/sdk"
+import type Anthropic from "@anthropic-ai/sdk"
 import { execa } from "execa"
+import { ClaudeCodeMessage } from "./types"
+import readline from "readline"
 
-export function runClaudeCode({
-	systemPrompt,
-	messages,
-	path,
-	modelId,
-}: {
+const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
+
+type ClaudeCodeOptions = {
 	systemPrompt: string
 	messages: Anthropic.Messages.MessageParam[]
 	path?: string
 	modelId?: string
-}) {
+}
+
+type ProcessState = {
+	partialData: string | null
+	error: Error | null
+	stderrLogs: string
+	exitCode: number | null
+}
+
+export async function* runClaudeCode(options: ClaudeCodeOptions): AsyncGenerator<ClaudeCodeMessage | string> {
+	const process = runProcess(options)
+
+	const rl = readline.createInterface({
+		input: process.stdout,
+	})
+
+	try {
+		const processState: ProcessState = {
+			error: null,
+			stderrLogs: "",
+			exitCode: null,
+			partialData: null,
+		}
+
+		process.stderr.on("data", (data) => {
+			processState.stderrLogs += data.toString()
+		})
+
+		process.on("close", (code) => {
+			processState.exitCode = code
+		})
+
+		process.on("error", (err) => {
+			processState.error = err
+		})
+
+		for await (const line of rl) {
+			if (processState.error) {
+				throw processState.error
+			}
+
+			if (line.trim()) {
+				const chunk = parseChunk(line, processState)
+
+				if (!chunk) {
+					continue
+				}
+
+				yield chunk
+			}
+		}
+
+		// We rely on the assistant message. If the output was truncated, it's better having a poorly formatted message
+		// from which to extract something, than throwing an error/showing the model didn't return any messages.
+		if (processState.partialData && processState.partialData.startsWith(`{"type":"assistant"`)) {
+			yield processState.partialData
+		}
+
+		const { exitCode } = await process
+		if (exitCode !== null && exitCode !== 0) {
+			const errorOutput = processState.error?.message || processState.stderrLogs?.trim()
+			throw new Error(
+				`Claude Code process exited with code ${exitCode}.${errorOutput ? ` Error output: ${errorOutput}` : ""}`,
+			)
+		}
+	} finally {
+		rl.close()
+		if (!process.killed) {
+			process.kill()
+		}
+	}
+}
+
+// We want the model to use our custom tool format instead of built-in tools.
+// Disabling built-in tools prevents tool-only responses and ensures text output.
+const claudeCodeTools = [
+	"Task",
+	"Bash",
+	"Glob",
+	"Grep",
+	"LS",
+	"exit_plan_mode",
+	"Read",
+	"Edit",
+	"MultiEdit",
+	"Write",
+	"NotebookRead",
+	"NotebookEdit",
+	"WebFetch",
+	"TodoRead",
+	"TodoWrite",
+	"WebSearch",
+].join(",")
+
+const CLAUDE_CODE_TIMEOUT = 600000 // 10 minutes
+
+function runProcess({ systemPrompt, messages, path, modelId }: ClaudeCodeOptions) {
 	const claudePath = path || "claude"
 
-	// TODO: Is it worth using sessions? Where do we store the session ID?
 	const args = [
 		"-p",
 		JSON.stringify(messages),
@@ -24,7 +118,9 @@ export function runClaudeCode({
 		"--verbose",
 		"--output-format",
 		"stream-json",
-		// Cline will handle recursive calls
+		"--disallowedTools",
+		claudeCodeTools,
+		// Roo Code will handle recursive calls
 		"--max-turns",
 		"1",
 	]
@@ -33,12 +129,49 @@ export function runClaudeCode({
 		args.push("--model", modelId)
 	}
 
-	const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
 	return execa(claudePath, args, {
 		stdin: "ignore",
 		stdout: "pipe",
 		stderr: "pipe",
-		env: process.env,
+		env: {
+			...process.env,
+			// The default is 32000. However, I've gotten larger responses, so we increase it unless the user specified it.
+			CLAUDE_CODE_MAX_OUTPUT_TOKENS: process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS || "64000",
+		},
 		cwd,
+		maxBuffer: 1024 * 1024 * 1000,
+		timeout: CLAUDE_CODE_TIMEOUT,
 	})
 }
+
+function parseChunk(data: string, processState: ProcessState) {
+	if (processState.partialData) {
+		processState.partialData += data
+
+		const chunk = attemptParseChunk(processState.partialData)
+
+		if (!chunk) {
+			return null
+		}
+
+		processState.partialData = null
+		return chunk
+	}
+
+	const chunk = attemptParseChunk(data)
+
+	if (!chunk) {
+		processState.partialData = data
+	}
+
+	return chunk
+}
+
+function attemptParseChunk(data: string): ClaudeCodeMessage | null {
+	try {
+		return JSON.parse(data)
+	} catch (error) {
+		console.error("Error parsing chunk:", error, data.length)
+		return null
+	}
+}

src/integrations/claude-code/types.ts

@@ -1,40 +1,17 @@
+import type { Anthropic } from "@anthropic-ai/sdk"
+
 type InitMessage = {
 	type: "system"
 	subtype: "init"
 	session_id: string
 	tools: string[]
 	mcp_servers: string[]
+	apiKeySource: "none" | "/login managed key" | string
 }
 
-type ClaudeCodeContent =
-	| {
-			type: "text"
-			text: string
-	  }
-	| {
-			type: "thinking"
-			thinking: string
-			signature?: string
-	  }
-
 type AssistantMessage = {
 	type: "assistant"
-	message: {
-		id: string
-		type: "message"
-		role: "assistant"
-		model: string
-		content: ClaudeCodeContent[]
-		stop_reason: string | null
-		stop_sequence: null
-		usage: {
-			input_tokens: number
-			cache_creation_input_tokens?: number
-			cache_read_input_tokens?: number
-			output_tokens: number
-			service_tier: "standard"
-		}
-	}
+	message: Anthropic.Messages.Message
 	session_id: string
 }
 
@@ -45,13 +22,12 @@ type ErrorMessage = {
 type ResultMessage = {
 	type: "result"
 	subtype: "success"
-	cost_usd: number
+	total_cost_usd: number
 	is_error: boolean
 	duration_ms: number
 	duration_api_ms: number
 	num_turns: number
 	result: string
-	total_cost: number
 	session_id: string
 }