diff --git a/src/api/providers/__tests__/claude-code.spec.ts b/src/api/providers/__tests__/claude-code.spec.ts index af1bf809c4..5f06665b6d 100644 --- a/src/api/providers/__tests__/claude-code.spec.ts +++ b/src/api/providers/__tests__/claude-code.spec.ts @@ -563,4 +563,134 @@ describe("ClaudeCodeHandler", () => { consoleSpy.mockRestore() }) + + test("should parse string chunks that are JSON assistant messages", async () => { + const systemPrompt = "You are a helpful assistant" + const messages = [{ role: "user" as const, content: "Hello" }] + + // Mock async generator that yields a string containing JSON assistant message + const mockGenerator = async function* (): AsyncGenerator { + // Yield a string that's actually a JSON assistant message + yield JSON.stringify({ + type: "assistant", + message: { + id: "msg_123", + type: "message", + role: "assistant", + model: "claude-3-5-sonnet-20241022", + content: [ + { + type: "text", + text: "This is a response from a JSON string", + }, + ], + stop_reason: null, + stop_sequence: null, + usage: { + input_tokens: 10, + output_tokens: 20, + }, + }, + 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 parse the JSON and yield the text content + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "text", + text: "This is a response from a JSON string", + }) + }) + + test("should handle malformed JSON strings gracefully", async () => { + const systemPrompt = "You are a helpful assistant" + const messages = [{ role: "user" as const, content: "Hello" }] + + // Mock async generator that yields malformed JSON + const mockGenerator = async function* (): AsyncGenerator { + // Yield a malformed JSON string with escaped newlines + yield '{"type":"assistant","message":{"id":"msg_123"\\n\\n\\n"content":[{"type":"text","text":"Malformed"}]}' + } + + mockRunClaudeCode.mockReturnValue(mockGenerator()) + + const stream = handler.createMessage(systemPrompt, messages) + const results = [] + + for await (const chunk of stream) { + results.push(chunk) + } + + // Should yield the malformed string as text since it can't be parsed + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "text", + text: '{"type":"assistant","message":{"id":"msg_123"\\n\\n\\n"content":[{"type":"text","text":"Malformed"}]}', + }) + }) + + test("should handle string chunks with thinking content when parsed from JSON", async () => { + const systemPrompt = "You are a helpful assistant" + const messages = [{ role: "user" as const, content: "Hello" }] + + // Mock async generator that yields a string containing JSON with thinking + const mockGenerator = async function* (): AsyncGenerator { + yield JSON.stringify({ + 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 answer", + }, + ], + stop_reason: null, + stop_sequence: null, + usage: { + input_tokens: 10, + output_tokens: 20, + }, + }, + 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 parse and yield both thinking and text content + expect(results).toHaveLength(2) + expect(results[0]).toEqual({ + type: "reasoning", + text: "Let me think about this...", + }) + expect(results[1]).toEqual({ + type: "text", + text: "Here's my answer", + }) + }) }) diff --git a/src/api/providers/claude-code.ts b/src/api/providers/claude-code.ts index dfafb78aab..99248f1900 100644 --- a/src/api/providers/claude-code.ts +++ b/src/api/providers/claude-code.ts @@ -54,6 +54,15 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler { for await (const chunk of claudeProcess) { if (typeof chunk === "string") { + // Try to parse string chunks that might be JSON assistant messages + const parsedChunk = this.attemptParseAssistantMessage(chunk) + if (parsedChunk) { + // Process as assistant message + yield* this.processAssistantMessage(parsedChunk, usage, isPaidUsage) + continue + } + + // If not a JSON message, yield as text yield { type: "text", text: chunk, @@ -69,64 +78,7 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler { } if (chunk.type === "assistant" && "message" in chunk) { - const message = chunk.message - - if (message.stop_reason !== null) { - const content = "text" in message.content[0] ? message.content[0] : undefined - - const isError = content && content.text.startsWith(`API Error`) - if (isError) { - // Error messages are formatted as: `API Error: <> <>` - 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) - } - } - - for (const content of message.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 - } - } - - usage.inputTokens += message.usage.input_tokens - usage.outputTokens += message.usage.output_tokens - usage.cacheReadTokens = (usage.cacheReadTokens || 0) + (message.usage.cache_read_input_tokens || 0) - usage.cacheWriteTokens = - (usage.cacheWriteTokens || 0) + (message.usage.cache_creation_input_tokens || 0) - + yield* this.processAssistantMessage(chunk, usage, isPaidUsage) continue } @@ -172,4 +124,85 @@ export class ClaudeCodeHandler extends BaseProvider implements ApiHandler { return null } } + + private *processAssistantMessage( + chunk: any, + usage: ApiStreamUsageChunk, + isPaidUsage: boolean, + ): Generator { + const message = chunk.message + + if (message.stop_reason !== null) { + const content = "text" in message.content[0] ? message.content[0] : undefined + + const isError = content && content.text.startsWith(`API Error`) + if (isError) { + // Error messages are formatted as: `API Error: <> <>` + 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) + } + } + + for (const content of message.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 + } + } + + usage.inputTokens += message.usage.input_tokens + usage.outputTokens += message.usage.output_tokens + usage.cacheReadTokens = (usage.cacheReadTokens || 0) + (message.usage.cache_read_input_tokens || 0) + usage.cacheWriteTokens = (usage.cacheWriteTokens || 0) + (message.usage.cache_creation_input_tokens || 0) + } + + private attemptParseAssistantMessage(str: string): any { + // Only try to parse if it looks like a JSON assistant message + if (!str.trim().startsWith('{"type":"assistant"')) { + return null + } + + try { + const parsed = JSON.parse(str) + // Validate it has the expected structure + if (parsed.type === "assistant" && parsed.message) { + return parsed + } + return null + } catch (err) { + // If parsing fails, log the error for debugging but don't throw + console.error("Failed to parse potential assistant message:", err) + return null + } + } }