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"); }