Merge pull request #18 from liuxiaopai-ai/fix/persist-structured-messages

Code review approved: high quality bug fix, backward compatible design, comprehensive tests (10 cases).

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<MessageContentBlock, { type: 'tool_result' }>;
+    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<MessageContentBlock, { type: 'text' }> => 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',
+    ]);
+  });
+});

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<string>, 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<string>, 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<string>, 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<MessageContentBlock, { type: 'text' }> => 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<MessageContentBlock, { type: 'text' }> => b.type === 'text')
+            .map((b) => b.text)
+            .join('')
+            .trim();
+      if (content) {
+        addMessage(sessionId, 'assistant', content);
+      }
     }
   }
 }