From 960d5906b80108b4bb1ccd3b89a32bb19f919634 Mon Sep 17 00:00:00 2001 From: liuxiaopai-ai Date: Sun, 8 Feb 2026 15:59:59 +0800 Subject: [PATCH] fix: persist structured tool_use/tool_result blocks in messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The collectStreamResponse() function only collected text events, discarding tool_use and tool_result events. When a conversation was resumed, all tool call context was lost — the user only saw the text portions of Claude's responses. Changes: - Collect tool_use and tool_result SSE events as MessageContentBlock entries alongside text blocks - Flush accumulated text into a text block before each tool_use to preserve the interleaved ordering (text → tool_use → tool_result → text → ...) - Store messages with tool calls as JSON arrays of MessageContentBlock, using the existing type definitions in types/index.ts - Text-only messages remain as plain strings for backward compatibility with existing message rendering (parseMessageContent handles both) - Add 10 unit tests covering parsing, serialization, round-tripping, edge cases, and backward compatibility (all passing) This enables future improvements like: - Rendering tool call history in resumed conversations - Showing which files were read/written in previous turns - Exporting full conversation logs with tool context --- .../unit/message-persistence.test.ts | 151 ++++++++++++++++++ src/app/api/chat/route.ts | 89 +++++++++-- 2 files changed, 229 insertions(+), 11 deletions(-) create mode 100644 src/__tests__/unit/message-persistence.test.ts diff --git a/src/__tests__/unit/message-persistence.test.ts b/src/__tests__/unit/message-persistence.test.ts new file mode 100644 index 0000000..3465482 --- /dev/null +++ b/src/__tests__/unit/message-persistence.test.ts @@ -0,0 +1,151 @@ +/** + * Unit tests for structured message persistence in collectStreamResponse. + * + * Run with: npx tsx --test src/__tests__/unit/message-persistence.test.ts + * + * Tests verify that: + * 1. parseMessageContent correctly parses structured JSON content + * 2. parseMessageContent handles plain text fallback + * 3. MessageContentBlock types are correctly structured + * 4. Backward compatibility: text-only messages stay as plain strings + * 5. Mixed content (text + tool_use + tool_result) serializes correctly + */ + +import { describe, it } from 'node:test'; +import assert from 'node:assert/strict'; + +import { parseMessageContent } from '../../types'; +import type { MessageContentBlock } from '../../types'; + +describe('parseMessageContent', () => { + it('should parse plain text as a single text block', () => { + const result = parseMessageContent('Hello, world!'); + assert.deepEqual(result, [{ type: 'text', text: 'Hello, world!' }]); + }); + + it('should parse JSON array of content blocks', () => { + const blocks: MessageContentBlock[] = [ + { type: 'text', text: 'Let me check that file.' }, + { type: 'tool_use', id: 'tu_123', name: 'read_file', input: { path: '/src/index.ts' } }, + ]; + const json = JSON.stringify(blocks); + const result = parseMessageContent(json); + assert.deepEqual(result, blocks); + }); + + it('should handle tool_result blocks', () => { + const blocks: MessageContentBlock[] = [ + { type: 'text', text: 'Reading file...' }, + { type: 'tool_use', id: 'tu_456', name: 'read_file', input: { path: '/package.json' } }, + { type: 'tool_result', tool_use_id: 'tu_456', content: '{"name": "codepilot"}', is_error: false }, + { type: 'text', text: 'The package name is codepilot.' }, + ]; + const json = JSON.stringify(blocks); + const result = parseMessageContent(json); + assert.equal(result.length, 4); + assert.equal(result[0].type, 'text'); + assert.equal(result[1].type, 'tool_use'); + assert.equal(result[2].type, 'tool_result'); + assert.equal(result[3].type, 'text'); + }); + + it('should handle error tool results', () => { + const blocks: MessageContentBlock[] = [ + { type: 'tool_result', tool_use_id: 'tu_789', content: 'File not found', is_error: true }, + ]; + const json = JSON.stringify(blocks); + const result = parseMessageContent(json); + assert.equal(result.length, 1); + const block = result[0] as Extract; + assert.equal(block.is_error, true); + }); + + it('should fall back to plain text for non-JSON content', () => { + const content = 'This is markdown **bold** text with `code`'; + const result = parseMessageContent(content); + assert.deepEqual(result, [{ type: 'text', text: content }]); + }); + + it('should fall back to plain text for JSON that is not an array', () => { + const content = JSON.stringify({ key: 'value' }); + const result = parseMessageContent(content); + assert.deepEqual(result, [{ type: 'text', text: content }]); + }); + + it('should handle empty content', () => { + const result = parseMessageContent(''); + assert.deepEqual(result, [{ type: 'text', text: '' }]); + }); +}); + +describe('Structured message serialization', () => { + it('should serialize text-only messages as plain text for backward compat', () => { + const blocks: MessageContentBlock[] = [ + { type: 'text', text: 'Hello, this is a response.' }, + ]; + + // Logic from collectStreamResponse: text-only → plain string + const hasToolBlocks = blocks.some( + (b) => b.type === 'tool_use' || b.type === 'tool_result' + ); + assert.equal(hasToolBlocks, false); + + const content = blocks + .filter((b): b is Extract => b.type === 'text') + .map((b) => b.text) + .join('') + .trim(); + + assert.equal(content, 'Hello, this is a response.'); + // parseMessageContent should handle it as plain text + const parsed = parseMessageContent(content); + assert.deepEqual(parsed, [{ type: 'text', text: 'Hello, this is a response.' }]); + }); + + it('should serialize mixed content as JSON', () => { + const blocks: MessageContentBlock[] = [ + { type: 'text', text: 'Let me read that.' }, + { type: 'tool_use', id: 'tu_001', name: 'read_file', input: { path: '/src/app.ts' } }, + { type: 'tool_result', tool_use_id: 'tu_001', content: 'export default {};', is_error: false }, + { type: 'text', text: 'The file exports a default empty object.' }, + ]; + + const hasToolBlocks = blocks.some( + (b) => b.type === 'tool_use' || b.type === 'tool_result' + ); + assert.equal(hasToolBlocks, true); + + const content = JSON.stringify(blocks); + // Should round-trip correctly + const parsed = parseMessageContent(content); + assert.equal(parsed.length, 4); + assert.equal(parsed[0].type, 'text'); + assert.equal(parsed[1].type, 'tool_use'); + assert.equal(parsed[2].type, 'tool_result'); + assert.equal(parsed[3].type, 'text'); + }); + + it('should handle multiple text blocks being flushed around tool calls', () => { + const blocks: MessageContentBlock[] = [ + { type: 'text', text: 'First I will check the file structure.' }, + { type: 'tool_use', id: 'tu_a', name: 'list_files', input: { dir: '.' } }, + { type: 'tool_result', tool_use_id: 'tu_a', content: 'src/\npackage.json', is_error: false }, + { type: 'text', text: 'Now let me read package.json.' }, + { type: 'tool_use', id: 'tu_b', name: 'read_file', input: { path: 'package.json' } }, + { type: 'tool_result', tool_use_id: 'tu_b', content: '{"name":"test"}', is_error: false }, + { type: 'text', text: 'Done! The project is named "test".' }, + ]; + + const content = JSON.stringify(blocks); + const parsed = parseMessageContent(content); + assert.equal(parsed.length, 7); + + // Verify interleaved structure is preserved + const types = parsed.map((b) => b.type); + assert.deepEqual(types, [ + 'text', 'tool_use', 'tool_result', + 'text', 'tool_use', 'tool_result', + 'text', + ]); + }); +}); diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index a32c270..114d269 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -1,7 +1,7 @@ import { NextRequest } from 'next/server'; import { streamClaude } from '@/lib/claude-client'; import { addMessage, getSession, updateSessionTitle, updateSdkSessionId, getSetting } from '@/lib/db'; -import type { SendMessageRequest, SSEEvent, TokenUsage } from '@/types'; +import type { SendMessageRequest, SSEEvent, TokenUsage, MessageContentBlock } from '@/types'; export const runtime = 'nodejs'; export const dynamic = 'force-dynamic'; @@ -99,7 +99,8 @@ export async function POST(request: NextRequest) { async function collectStreamResponse(stream: ReadableStream, sessionId: string) { const reader = stream.getReader(); - let fullText = ''; + const contentBlocks: MessageContentBlock[] = []; + let currentText = ''; let tokenUsage: TokenUsage | null = null; try { @@ -115,7 +116,36 @@ async function collectStreamResponse(stream: ReadableStream, sessionId: if (event.type === 'permission_request' || event.type === 'tool_output') { // Skip permission_request and tool_output events - not saved as message content } else if (event.type === 'text') { - fullText += event.data; + currentText += event.data; + } else if (event.type === 'tool_use') { + // Flush any accumulated text before the tool use block + if (currentText.trim()) { + contentBlocks.push({ type: 'text', text: currentText }); + currentText = ''; + } + try { + const toolData = JSON.parse(event.data); + contentBlocks.push({ + type: 'tool_use', + id: toolData.id, + name: toolData.name, + input: toolData.input, + }); + } catch { + // skip malformed tool_use data + } + } else if (event.type === 'tool_result') { + try { + const resultData = JSON.parse(event.data); + contentBlocks.push({ + type: 'tool_result', + tool_use_id: resultData.tool_use_id, + content: resultData.content, + is_error: resultData.is_error || false, + }); + } catch { + // skip malformed tool_result data + } } else if (event.type === 'status') { // Capture SDK session_id from init event and persist it try { @@ -147,18 +177,55 @@ async function collectStreamResponse(stream: ReadableStream, sessionId: } } - if (fullText.trim()) { - addMessage( - sessionId, - 'assistant', - fullText.trim(), - tokenUsage ? JSON.stringify(tokenUsage) : null, + // Flush any remaining text + if (currentText.trim()) { + contentBlocks.push({ type: 'text', text: currentText }); + } + + if (contentBlocks.length > 0) { + // If the message is text-only (no tool calls), store as plain text + // for backward compatibility with existing message rendering. + // If it contains tool calls, store as structured JSON. + const hasToolBlocks = contentBlocks.some( + (b) => b.type === 'tool_use' || b.type === 'tool_result' ); + + const content = hasToolBlocks + ? JSON.stringify(contentBlocks) + : contentBlocks + .filter((b): b is Extract => b.type === 'text') + .map((b) => b.text) + .join('') + .trim(); + + if (content) { + addMessage( + sessionId, + 'assistant', + content, + tokenUsage ? JSON.stringify(tokenUsage) : null, + ); + } } } catch { // Stream reading error - best effort save - if (fullText.trim()) { - addMessage(sessionId, 'assistant', fullText.trim()); + if (currentText.trim()) { + contentBlocks.push({ type: 'text', text: currentText }); + } + if (contentBlocks.length > 0) { + const hasToolBlocks = contentBlocks.some( + (b) => b.type === 'tool_use' || b.type === 'tool_result' + ); + const content = hasToolBlocks + ? JSON.stringify(contentBlocks) + : contentBlocks + .filter((b): b is Extract => b.type === 'text') + .map((b) => b.text) + .join('') + .trim(); + if (content) { + addMessage(sessionId, 'assistant', content); + } } } }