Răsfoiți Sursa

fix: resolve Claude Code provider JSON parsing and reasoning block display (#5049)

Co-authored-by: Daniel Riccio <[email protected]>
Hannes Rudolph 6 luni în urmă
părinte
comite
cabf19153e

+ 230 - 0
src/api/providers/__tests__/claude-code.spec.ts

@@ -0,0 +1,230 @@
+import { describe, test, expect, vi, beforeEach } from "vitest"
+import { ClaudeCodeHandler } from "../claude-code"
+import { ApiHandlerOptions } from "../../../shared/api"
+
+// Mock the runClaudeCode function
+vi.mock("../../../integrations/claude-code/run", () => ({
+	runClaudeCode: vi.fn(),
+}))
+
+const { runClaudeCode } = await import("../../../integrations/claude-code/run")
+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))
+		}
+	}
+}
+
+describe("ClaudeCodeHandler", () => {
+	let handler: ClaudeCodeHandler
+	let mockProcess: any
+
+	beforeEach(() => {
+		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),
+		}
+
+		mockRunClaudeCode.mockReturnValue(mockProcess)
+	})
+
+	test("should handle thinking content properly", async () => {
+		const systemPrompt = "You are a helpful assistant"
+		const messages = [{ role: "user" as const, content: "Hello" }]
+
+		// 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)
+			})
+		})
+
+		// Get the result
+		const result = await streamGenerator.next()
+
+		expect(result.done).toBe(false)
+		expect(result.value).toEqual({
+			type: "reasoning",
+			text: "I need to think about this carefully...",
+		})
+	})
+
+	test("should handle mixed content types", 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...",
+					},
+					{
+						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)
+			})
+		})
+
+		// Get the first result (thinking)
+		const thinkingResult = await streamGenerator.next()
+		expect(thinkingResult.done).toBe(false)
+		expect(thinkingResult.value).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({
+			type: "text",
+			text: "Here's my response!",
+		})
+	})
+
+	test("should handle stop_reason with thinking content in error messages", 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 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)
+			})
+		})
+
+		// Should throw error with thinking content
+		await expect(streamGenerator.next()).rejects.toThrow("This is an error scenario")
+	})
+
+	test("should handle incomplete JSON in buffer on process close", 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 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)
+			})
+		})
+
+		// Should complete without throwing, incomplete JSON should be discarded
+		const result = await streamGenerator.next()
+		expect(result.done).toBe(true)
+	})
+})

+ 75 - 6
src/api/providers/claude-code.ts

@@ -28,13 +28,21 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 		let processError = null
 		let errorOutput = ""
 		let exitCode: number | null = null
+		let buffer = ""
 
 		claudeProcess.stdout.on("data", (data) => {
-			const output = data.toString()
-			const lines = output.split("\n").filter((line: string) => line.trim() !== "")
+			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) {
-				dataQueue.push(line)
+				const trimmedLine = line.trim()
+				if (trimmedLine !== "") {
+					dataQueue.push(trimmedLine)
+				}
 			}
 		})
 
@@ -44,6 +52,20 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 
 		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) => {
@@ -101,8 +123,9 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 				const message = chunk.message
 
 				if (message.stop_reason !== null && message.stop_reason !== "tool_use") {
+					const firstContent = message.content[0]
 					const errorMessage =
-						message.content[0]?.text ||
+						this.getContentText(firstContent) ||
 						t("common:errors.claudeCode.stoppedWithReason", { reason: message.stop_reason })
 
 					if (errorMessage.includes("Invalid model name")) {
@@ -118,8 +141,13 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler {
 							type: "text",
 							text: content.text,
 						}
+					} else if (content.type === "thinking") {
+						yield {
+							type: "reasoning",
+							text: content.thinking,
+						}
 					} else {
-						console.warn("Unsupported content type:", content.type)
+						console.warn("Unsupported content type:", content)
 					}
 				}
 
@@ -159,12 +187,53 @@ 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 {
 		try {
 			return JSON.parse(data)
 		} catch (error) {
-			console.error("Error parsing chunk:", error)
+			console.error(
+				"Error parsing chunk:",
+				error,
+				"Data:",
+				data.substring(0, 100) + (data.length > 100 ? "..." : ""),
+			)
 			return null
 		}
 	}

+ 10 - 4
src/integrations/claude-code/types.ts

@@ -6,10 +6,16 @@ type InitMessage = {
 	mcp_servers: string[]
 }
 
-type ClaudeCodeContent = {
-	type: "text"
-	text: string
-}
+type ClaudeCodeContent =
+	| {
+			type: "text"
+			text: string
+	  }
+	| {
+			type: "thinking"
+			thinking: string
+			signature?: string
+	  }
 
 type AssistantMessage = {
 	type: "assistant"