fix: persist structured tool_use/tool_result blocks in messages

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
This commit is contained in:
liuxiaopai-ai committed 2026-02-08 15:59:59 +08:00
1 parent 21226fc010
commit 960d5906b8
2 files changed
+229 -11

No files matched your search

+78 -11
View File
@@ -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);
}
}
}
}