Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 230 additions & 0 deletions src/api/providers/__tests__/claude-code.spec.ts
Original file line number Diff line number Diff line change
@@ -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)
})
})
81 changes: 75 additions & 6 deletions src/api/providers/claude-code.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
})

Expand All @@ -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) => {
Expand Down Expand Up @@ -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")) {
Expand All @@ -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)
}
}

Expand Down Expand Up @@ -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
}
}
Expand Down
14 changes: 10 additions & 4 deletions src/integrations/claude-code/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading