From 7ddd0cba058cf074f6b8e3115bffb19ec8c3bf4d Mon Sep 17 00:00:00 2001 From: claude_code Date: Mon, 22 Jun 2026 20:00:35 +0300 Subject: [PATCH] feat(ai-chat): include in-progress streaming turn in chat export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The "Copy chat" export read only persisted DB rows (messageRows), so an assistant reply that was still streaming — and the user message that triggered it — were absent from the export until the turn finished and the messages query was refetched. ChatThread now mirrors its live useChat snapshot ({ messages, isStreaming }) into a parent-owned ref; the effect clears the ref on unmount so a thread switch can't leak its tail into the next chat. AiChatWindow.handleCopy computes the not-yet-persisted live tail (messages whose id is absent from messageRows, only while streaming) and passes it to buildChatMarkdown as `pending`. buildChatMarkdown appends pending messages after the persisted rows (continuing the heading numbering), flags the streaming assistant message with an "still being generated" note, and reuses an extracted renderMessageParts helper so persisted and pending rendering stay identical. Co-Authored-By: Claude Opus 4.8 --- .../ai-chat/components/ai-chat-window.tsx | 27 ++++ .../ai-chat/components/chat-thread.tsx | 29 ++++- .../ai-chat/utils/chat-markdown.test.ts | 123 ++++++++++++++++++ .../features/ai-chat/utils/chat-markdown.ts | 115 +++++++++++----- 4 files changed, 257 insertions(+), 37 deletions(-) diff --git a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx index 3d8cf55d..5fda9f8c 100644 --- a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx +++ b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx @@ -7,6 +7,7 @@ import { useState, } from "react"; import { generateId } from "ai"; +import { type UIMessage } from "@ai-sdk/react"; import { Group, Loader, Tooltip } from "@mantine/core"; import { IconArrowsDiagonal, @@ -171,6 +172,14 @@ export default function AiChatWindow() { const { data: messageRows, isLoading: messagesLoading } = useAiChatMessagesQuery(activeChatId ?? undefined); + // Live snapshot of the active thread's useChat state, kept up to date by + // ChatThread. Lets the export include the in-progress (not-yet-persisted) + // streaming turn. A ref avoids re-rendering this window on every token. + const liveThreadRef = useRef<{ messages: UIMessage[]; isStreaming: boolean }>({ + messages: [], + isStreaming: false, + }); + // The page the user is currently viewing. AiChatWindow lives in a pathless // parent layout route, so useParams() can't see :pageSlug. Match the full // pathname against the authenticated page route instead so "the current page" @@ -259,10 +268,27 @@ export default function AiChatWindow() { // feedback. const handleCopy = useCallback(() => { if (!activeChatId || !messageRows || messageRows.length === 0) return; + // While the active thread is streaming, the current user message and the + // in-progress assistant reply are NOT yet in messageRows (the persisted + // query is only refetched after the turn finishes). Pull the live tail — + // messages whose id is not among the persisted rows — and append them, + // flagging the streaming assistant message as still generating. + const live = liveThreadRef.current; + const rowIds = new Set(messageRows.map((r) => r.id)); + const pending = live.isStreaming + ? live.messages + .filter((m) => !rowIds.has(m.id)) + .map((m) => ({ + role: m.role, + parts: (m.parts ?? []) as { type: string; text?: string }[], + generating: m.role === "assistant", + })) + : []; const markdown = buildChatMarkdown({ title: activeChat?.title ?? null, chatId: activeChatId, rows: messageRows, + pending, t, }); clipboard.copy(markdown); @@ -657,6 +683,7 @@ export default function AiChatWindow() { onRolePicked={(role) => setSelectedRoleId(role.id)} assistantName={currentRole?.name} onTurnFinished={onTurnFinished} + liveStateRef={liveThreadRef} /> )} diff --git a/apps/client/src/features/ai-chat/components/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index f14cd732..5b7e48eb 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -1,4 +1,11 @@ -import { useCallback, useMemo, useRef, useState } from "react"; +import { + useCallback, + useEffect, + useMemo, + useRef, + useState, + type MutableRefObject, +} from "react"; import { generateId } from "ai"; import { ActionIcon, Box, Group, Stack, Text } from "@mantine/core"; import { IconClockHour4, IconX } from "@tabler/icons-react"; @@ -52,6 +59,12 @@ interface ChatThreadProps { /** Called when a turn finishes; the parent refreshes the chat list and, for * a new chat, adopts the freshly created chat id. */ onTurnFinished: () => void; + /** Parent-owned ref that this thread keeps updated with its live useChat + * snapshot (full message list + streaming flag), so the header's + * "Copy chat" export can include the in-progress, not-yet-persisted + * assistant message. A ref (not state) avoids re-rendering the parent on + * every streamed delta. */ + liveStateRef?: MutableRefObject<{ messages: UIMessage[]; isStreaming: boolean }>; } /** @@ -90,6 +103,7 @@ export default function ChatThread({ onRolePicked, assistantName, onTurnFinished, + liveStateRef, }: ChatThreadProps) { const { t } = useTranslation(); @@ -249,6 +263,19 @@ export default function ChatThread({ const isStreaming = status === "submitted" || status === "streaming"; + // Mirror the live useChat snapshot into the parent-owned ref so the export + // (handled in AiChatWindow) can include the in-progress streaming turn. The + // cleanup clears the ref on unmount so a thread torn down by `key` on chat + // switch can't leak its (possibly still-streaming) tail into the next chat's + // export before the new thread's effect repopulates the ref. + useEffect(() => { + if (!liveStateRef) return; + liveStateRef.current = { messages, isStreaming }; + return () => { + liveStateRef.current = { messages: [], isStreaming: false }; + }; + }, [liveStateRef, messages, isStreaming]); + // Classify the turn error into a heading + detail so the banner names the cause // (connection reset, timeout, rate limit, context overflow, quota, ...) instead // of a generic "Something went wrong". diff --git a/apps/client/src/features/ai-chat/utils/chat-markdown.test.ts b/apps/client/src/features/ai-chat/utils/chat-markdown.test.ts index 57f4c2e4..79eb6023 100644 --- a/apps/client/src/features/ai-chat/utils/chat-markdown.test.ts +++ b/apps/client/src/features/ai-chat/utils/chat-markdown.test.ts @@ -315,3 +315,126 @@ describe("buildChatMarkdown — token totals", () => { expect(md).toContain("- Total tokens: 99"); }); }); + +describe("buildChatMarkdown — pending / in-progress messages", () => { + it("continues the heading numbering after the persisted rows", () => { + const md = buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [row({ role: "user", content: "persisted" })], + pending: [ + { + role: "user", + parts: [{ type: "text", text: "live question" }], + generating: false, + }, + { + role: "assistant", + parts: [{ type: "text", text: "live answer" }], + generating: true, + }, + ], + t, + }); + expect(md).toContain("## 1. You"); + expect(md).toContain("## 2. You"); + expect(md).toContain("## 3. AI agent"); + expect(md).toContain("live question"); + expect(md).toContain("live answer"); + }); + + it("flags a generating assistant pending message as still being generated", () => { + const md = buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [row({ role: "user", content: "persisted" })], + pending: [ + { + role: "assistant", + parts: [{ type: "text", text: "partial reply" }], + generating: true, + }, + ], + t, + }); + expect(md).toContain("partial reply"); + expect(md).toContain("still being generated"); + }); + + it("renders a non-generating user pending message without the note", () => { + const md = buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [row({ role: "user", content: "persisted" })], + pending: [ + { + role: "user", + parts: [{ type: "text", text: "my live message" }], + generating: false, + }, + ], + t, + }); + expect(md).toContain("my live message"); + expect(md).not.toContain("still being generated"); + }); + + it("includes the pending messages in the metadata message count", () => { + const md = buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [ + row({ role: "user", content: "a" }), + row({ role: "assistant", content: "b" }), + ], + pending: [ + { + role: "user", + parts: [{ type: "text", text: "c" }], + generating: false, + }, + { + role: "assistant", + parts: [{ type: "text", text: "d" }], + generating: true, + }, + ], + t, + }); + // 2 persisted rows + 2 pending = 4. + expect(md).toContain("- Messages: 4"); + }); + + it("emits the heading and note for a generating assistant with empty parts", () => { + expect(() => + buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [row({ role: "user", content: "persisted" })], + pending: [ + { + role: "assistant", + parts: [], + generating: true, + }, + ], + t, + }), + ).not.toThrow(); + const md = buildChatMarkdown({ + title: "t", + chatId: "c", + rows: [row({ role: "user", content: "persisted" })], + pending: [ + { + role: "assistant", + parts: [], + generating: true, + }, + ], + t, + }); + expect(md).toContain("## 2. AI agent"); + expect(md).toContain("still being generated"); + }); +}); diff --git a/apps/client/src/features/ai-chat/utils/chat-markdown.ts b/apps/client/src/features/ai-chat/utils/chat-markdown.ts index 784692e7..f54d012a 100644 --- a/apps/client/src/features/ai-chat/utils/chat-markdown.ts +++ b/apps/client/src/features/ai-chat/utils/chat-markdown.ts @@ -26,6 +26,10 @@ interface BuildChatMarkdownArgs { title: string | null; chatId: string; rows: IAiChatMessageRow[]; + /** In-progress, not-yet-persisted live messages (the current streaming + * turn) to append after the persisted rows. `generating: true` adds a + * note that the message is still being produced. */ + pending?: PendingMessage[]; t: Translate; } @@ -35,6 +39,13 @@ interface TextLikePart { text?: string; } +/** A live, not-yet-persisted message (current streaming turn) to append. */ +interface PendingMessage { + role: "user" | "assistant" | string; + parts: TextLikePart[]; + generating: boolean; +} + /** * Stringify an arbitrary tool input/output value for a fenced block. Strings * pass through as-is; everything else is pretty-printed JSON, falling back to @@ -72,12 +83,55 @@ function rowTokens(usage: { ); } +/** Render one message's UIMessage parts into an array of Markdown blocks + * (text blocks + tool blocks). Mirrors MessageItem's part handling. */ +function renderMessageParts(parts: TextLikePart[], t: Translate): string[] { + const out: string[] = []; + + for (const part of parts) { + if (part.type === "text") { + const text = (part.text ?? "").trim(); + // Skip empty/whitespace-only text parts (matches MessageItem). + if (text.length > 0) out.push(text); + continue; + } + + const isToolPart = + part.type.startsWith("tool-") || part.type === "dynamic-tool"; + if (!isToolPart) continue; + + const tp = part as unknown as ToolUiPart; + const name = getToolName(tp); + const { key, values } = toolLabelKey(name); + const label = t(key, values); + const state = toolRunState(tp.state); + + const toolLines: string[] = [ + `**Tool: ${label}** (\`${name}\`) — ${state}`, + ]; + if (tp.input !== undefined) { + toolLines.push("Input:"); + toolLines.push(fence(stringify(tp.input), "json")); + } + if (tp.output !== undefined) { + toolLines.push("Output:"); + toolLines.push(fence(stringify(tp.output), "json")); + } + if (tp.errorText) { + toolLines.push(`**Error:** ${tp.errorText}`); + } + out.push(toolLines.join("\n\n")); + } + + return out; +} + /** * Serialize a chat to a Markdown string. Pure (apart from `new Date()` for the * export timestamp), so it is straightforward to unit-test. */ export function buildChatMarkdown(args: BuildChatMarkdownArgs): string { - const { title, chatId, rows, t } = args; + const { title, chatId, rows, pending, t } = args; const blocks: string[] = []; const heading = (title ?? "").trim() || t("Untitled chat"); @@ -91,7 +145,7 @@ export function buildChatMarkdown(args: BuildChatMarkdownArgs): string { const meta = [ `- Chat ID: \`${chatId}\``, `- Exported: ${new Date().toISOString()}`, - `- Messages: ${rows.length}`, + `- Messages: ${rows.length + (pending?.length ?? 0)}`, ]; if (totalTokens > 0) meta.push(`- Total tokens: ${totalTokens}`); blocks.push(meta.join("\n")); @@ -112,40 +166,7 @@ export function buildChatMarkdown(args: BuildChatMarkdownArgs): string { ? (row.metadata.parts as TextLikePart[]) : [{ type: "text", text: row.content ?? "" }]; - for (const part of parts) { - if (part.type === "text") { - const text = (part.text ?? "").trim(); - // Skip empty/whitespace-only text parts (matches MessageItem). - if (text.length > 0) blocks.push(text); - continue; - } - - const isToolPart = - part.type.startsWith("tool-") || part.type === "dynamic-tool"; - if (!isToolPart) continue; - - const tp = part as unknown as ToolUiPart; - const name = getToolName(tp); - const { key, values } = toolLabelKey(name); - const label = t(key, values); - const state = toolRunState(tp.state); - - const toolLines: string[] = [ - `**Tool: ${label}** (\`${name}\`) — ${state}`, - ]; - if (tp.input !== undefined) { - toolLines.push("Input:"); - toolLines.push(fence(stringify(tp.input), "json")); - } - if (tp.output !== undefined) { - toolLines.push("Output:"); - toolLines.push(fence(stringify(tp.output), "json")); - } - if (tp.errorText) { - toolLines.push(`**Error:** ${tp.errorText}`); - } - blocks.push(toolLines.join("\n\n")); - } + blocks.push(...renderMessageParts(parts, t)); if (row.metadata?.error) { blocks.push(`**⚠️ Error:** ${row.metadata.error}`); @@ -160,6 +181,28 @@ export function buildChatMarkdown(args: BuildChatMarkdownArgs): string { } }); + // Append the in-progress, not-yet-persisted live messages (the current + // streaming turn) after the persisted rows. Heading numbering CONTINUES from + // the persisted rows. A `generating` assistant gets a note that the captured + // response is partial; pending messages carry no usage/token footer yet. + (pending ?? []).forEach((message, p) => { + blocks.push("---"); + + const num = rows.length + p + 1; + const roleLabel = message.role === "assistant" ? t("AI agent") : t("You"); + blocks.push(`## ${num}. ${roleLabel}`); + + blocks.push(...renderMessageParts(message.parts, t)); + + // A generating assistant may have empty/no parts yet — still emit the + // heading (above) and this note so the export shows the in-progress turn. + if (message.generating === true) { + blocks.push( + "_⏳ This message is still being generated — the export captured a partial, in-progress response._", + ); + } + }); + // Blank line between blocks so the Markdown renders cleanly. return blocks.join("\n\n"); }