diff --git a/packages/client/src/components/MarkdownContent.test.tsx b/packages/client/src/components/MarkdownContent.test.tsx new file mode 100644 index 00000000..0b237b64 --- /dev/null +++ b/packages/client/src/components/MarkdownContent.test.tsx @@ -0,0 +1,85 @@ +import { describe, test, expect } from 'bun:test'; +import React from 'react'; +import { render } from '@testing-library/react'; +import { MarkdownContent } from './MarkdownContent'; + +/** + * Streaming appends to `content` many times per second. If a re-render + * recreates the DOM instead of updating it, the browser drops whatever the + * user has selected, so they cannot copy an answer while it is being written. + * These tests pin the DOM nodes of already-rendered text down across updates. + */ +describe('MarkdownContent during streaming', () => { + const stream = (steps: string[]) => { + const { container, rerender } = render(); + const firstBlock = container.firstElementChild; + const firstText = firstBlock?.firstChild; + // Without these the node-identity assertions below hold vacuously when + // nothing renders at all (null === null). + expect(firstBlock).not.toBeNull(); + expect(firstText).not.toBeNull(); + + let previousText = container.textContent; + for (const step of steps.slice(1)) { + rerender(); + expect(container.firstElementChild).toBe(firstBlock!); + expect(container.firstElementChild?.firstChild).toBe(firstText!); + // Node identity alone is also satisfied by a component that stops + // updating entirely, so each step has to visibly change the text. + expect(container.textContent).not.toBe(previousText); + previousText = container.textContent; + } + + return container; + }; + + test('keeps the growing paragraph node while text is appended to it', () => { + const container = stream(['Hello wo', 'Hello world', 'Hello world and more']); + + expect(container.textContent).toBe('Hello world and more'); + }); + + test('keeps earlier paragraphs when a new paragraph starts', () => { + const container = stream([ + 'First paragraph.', + 'First paragraph.\n\nSec', + 'First paragraph.\n\nSecond paragraph.', + ]); + + expect([...container.querySelectorAll('p')].map((p) => p.textContent)).toEqual([ + 'First paragraph.', + 'Second paragraph.', + ]); + }); + + test('keeps earlier text when a list is being written', () => { + const container = stream([ + 'Intro text.', + 'Intro text.\n\n- item on', + 'Intro text.\n\n- item one\n- item t', + 'Intro text.\n\n- item one\n- item two', + ]); + + expect([...container.querySelectorAll('li')].map((li) => li.textContent)).toEqual([ + 'item one', + 'item two', + ]); + }); + + test('keeps earlier text when a table is being written', () => { + const container = stream([ + 'Intro text.', + 'Intro text.\n\n| a | b |', + 'Intro text.\n\n| a | b |\n| - | - |', + 'Intro text.\n\n| a | b |\n| - | - |\n| 1 | 2 |', + ]); + + expect(container.querySelector('table')).not.toBeNull(); + }); + + test('keeps earlier text when inline markdown completes mid-stream', () => { + const container = stream(['Hello **bo', 'Hello **bold', 'Hello **bold**', 'Hello **bold** tail']); + + expect(container.querySelector('strong')?.textContent).toBe('bold'); + }); +}); diff --git a/packages/client/src/components/MarkdownContent.tsx b/packages/client/src/components/MarkdownContent.tsx index 958dccd9..b6d3e960 100644 --- a/packages/client/src/components/MarkdownContent.tsx +++ b/packages/client/src/components/MarkdownContent.tsx @@ -1,159 +1,177 @@ import React from 'react'; import ReactMarkdown from 'react-markdown'; +import type { Components } from 'react-markdown'; import remarkGfm from 'remark-gfm'; interface MarkdownContentProps { content: string; } +const REMARK_PLUGINS = [remarkGfm]; + /** - * Renders markdown content with appropriate styling for the chat panel. + * This map must stay a module-level constant. + * + * react-markdown re-creates its element tree on every render, and React + * reconciles that tree by component identity. If the `components` map were + * built inside the render function, every entry would be a brand new function + * on each render, so React would unmount and remount the whole subtree instead + * of updating it. During streaming that happens on every token: the browser + * throws away and rebuilds the entire answer many times per second, which + * destroys any text selection the user is making and makes the scroll position + * jump. Keeping it stable lets React update text nodes in place, so a + * selection survives the stream. */ -export function MarkdownContent({ content }: MarkdownContentProps) { - return ( -

{children}

, - // Ensure last paragraph has no margin - h1: ({ children }) =>

{children}

, - h2: ({ children }) =>

{children}

, - h3: ({ children }) =>

{children}

, - ul: ({ children }) =>
    {children}
, - ol: ({ children }) =>
    {children}
, - li: ({ children }) =>
  • {children}
    • , - code: ({ className, children, ...props }) => { - // Check if this is inline code or a code block - const isInline = !className; - if (isInline) { - return ( - - {children} - - ); - } - return ( - - {children} - - ); - }, - pre: ({ children }) => ( - 
    -            {children}
-
    - ), - blockquote: ({ children }) => ( -
    - {children} -
    - ), - a: ({ children, href }) => ( - - {children} - - ), - hr: () => ( -
    - ), - table: ({ children }) => ( -
    - - {children} -
    -
    - ), - th: ({ children }) => ( - - {children} - - ), - td: ({ children }) => ( - - {children} - - ), - // Render images as links to prevent automatic HTTP requests. - // tags fire GET requests on render, which could be exploited - // via prompt injection to exfiltrate sensitive data through URLs. - img: ({ src, alt }) => ( - - {alt || 'Image'} - - ), +const MARKDOWN_COMPONENTS: Components = { + // Override default element rendering for better chat styling + p: ({ children }) =>

    {children}

    , + // Ensure last paragraph has no margin + h1: ({ children }) =>

    {children}

    , + h2: ({ children }) =>

    {children}

    , + h3: ({ children }) =>

    {children}

    , + ul: ({ children }) =>
      {children}
    , + ol: ({ children }) =>
      {children}
    , + li: ({ children }) =>
  • {children}
    • , + code: ({ className, children, ...props }) => { + // Check if this is inline code or a code block + const isInline = !className; + if (isInline) { + return ( + + {children} + + ); + } + return ( + + {children} + + ); + }, + pre: ({ children }) => ( + 
    +      {children}
+
    + ), + blockquote: ({ children }) => ( +
    + {children} +
    + ), + a: ({ children, href }) => ( + + {children} + + ), + hr: () => ( +
    + ), + table: ({ children }) => ( +
    + + {children} +
    +
    + ), + th: ({ children }) => ( + + {children} + + ), + td: ({ children }) => ( + + {children} + + ), + // Render images as links to prevent automatic HTTP requests. + // tags fire GET requests on render, which could be exploited + // via prompt injection to exfiltrate sensitive data through URLs. + img: ({ src, alt }) => ( + + {alt || 'Image'} + + ), +}; + +/** + * Renders markdown content with appropriate styling for the chat panel. + * + * Memoized on `content`: the chat panel re-renders on every streaming token, + * and without this every message in the history would be re-parsed each time. + */ +export const MarkdownContent = React.memo(function MarkdownContent({ content }: MarkdownContentProps) { + return ( + {content} ); -} +}); diff --git a/packages/client/src/components/UseAIChat.test.tsx b/packages/client/src/components/UseAIChat.test.tsx index c21dab48..892ed735 100644 --- a/packages/client/src/components/UseAIChat.test.tsx +++ b/packages/client/src/components/UseAIChat.test.tsx @@ -13,6 +13,7 @@ function createContextValue(overrides: Partial = {}): ChatUI messages: [], streamingText: '', streamingReasoning: '', + streamingMessageId: null, suggestions: [], fileUploadConfig: undefined, history: { @@ -73,6 +74,42 @@ describe('UseAIChat', () => { expect(getByText('Thinking...')).toBeInTheDocument(); }); + // Without the id the streaming bubble and the persisted one are different + // elements, and a selection made while the answer streams is lost on + // completion. See UseAIChatPanel.selection.test.tsx. + test('renders the streaming answer under the id it will be persisted with', () => { + const ctx = createContextValue({ + loading: true, + streamingText: 'partial answer', + streamingMessageId: 'msg-assistant', + }); + + const { container, rerender } = render( + <__UseAIChatContext.Provider value={ctx}> + + + ); + + const streamingBubble = container.querySelector('[data-testid="chat-message-assistant-streaming"]'); + expect(streamingBubble).not.toBeNull(); + + const persisted: PersistedMessage = { + id: 'msg-assistant', + role: 'assistant', + content: 'partial answer, completed', + createdAt: new Date(0), + }; + rerender( + <__UseAIChatContext.Provider + value={createContextValue({ messages: [persisted], streamingText: '', streamingMessageId: null })} + > + + + ); + + expect(container.querySelector('[data-testid="chat-message-assistant"]')).toBe(streamingBubble); + }); + describe('save command feature toggle', () => { const userMessage: PersistedMessage = { id: 'msg-1', diff --git a/packages/client/src/components/UseAIChat.tsx b/packages/client/src/components/UseAIChat.tsx index 0182e384..1a35bc87 100644 --- a/packages/client/src/components/UseAIChat.tsx +++ b/packages/client/src/components/UseAIChat.tsx @@ -30,6 +30,12 @@ export interface ChatUIContextValue { streamingText: string; /** Currently streaming reasoning text from extended thinking */ streamingReasoning: string; + /** + * Id the streaming answer will be persisted under, or null when no run is + * streaming into the displayed chat. Passing it to the panel lets the + * streaming bubble and the persisted one be the same element. + */ + streamingMessageId: string | null; /** Aggregated suggestions from all useAI hooks */ suggestions: string[]; /** File upload configuration */ @@ -201,6 +207,7 @@ export function UseAIChat({ floating = false, submitMode }: UseAIChatProps) { connected: ctx.connected, streamingText: ctx.streamingText, streamingReasoning: ctx.streamingReasoning, + streamingMessageId: ctx.streamingMessageId, currentChatId: ctx.history.currentId, onNewChat: ctx.history.create, onLoadChat: ctx.history.load, diff --git a/packages/client/src/components/UseAIChatPanel.selection.test.tsx b/packages/client/src/components/UseAIChatPanel.selection.test.tsx new file mode 100644 index 00000000..6028b4fa --- /dev/null +++ b/packages/client/src/components/UseAIChatPanel.selection.test.tsx @@ -0,0 +1,208 @@ +import React from 'react'; +import { describe, test, expect, beforeEach } from 'bun:test'; +import { render } from '@testing-library/react'; +import { UseAIChatPanel, type UseAIChatPanelProps } from './UseAIChatPanel'; +import type { PersistedMessage } from '../providers/chatRepository/types'; + +const ANSWER = 'The first paragraph.\n\nThe second paragraph.'; +const STREAMING_ID = 'msg-assistant'; + +const userMessage: PersistedMessage = { + id: 'msg-user', + role: 'user', + content: 'question', + createdAt: new Date(0), +}; + +/** The persisted answer carries the id that was allocated when the run started. */ +const assistantMessage: PersistedMessage = { + id: STREAMING_ID, + role: 'assistant', + content: ANSWER, + createdAt: new Date(0), + traceId: 'trace-1', +}; + +/** + * Aborting a run that already produced text persists two messages: the partial + * answer, then this notice. See persistFinalResponse in useServerEvents. + */ +const abortNotice: PersistedMessage = { + id: 'msg-abort-notice', + role: 'assistant', + content: 'Generation stopped.', + createdAt: new Date(0), + displayMode: 'info', +}; + +function panelProps(overrides: Partial = {}): UseAIChatPanelProps { + return { + onSendMessage: () => {}, + messages: [userMessage], + loading: true, + connected: true, + streamingText: ANSWER, + streamingMessageId: STREAMING_ID, + ...overrides, + }; +} + +/** Selects `text` inside the first text node that contains it. */ +function selectText(container: HTMLElement, text: string): void { + const walker = document.createTreeWalker(container, window.NodeFilter.SHOW_TEXT); + while (walker.nextNode()) { + const node = walker.currentNode as Text; + const index = node.data.indexOf(text); + if (index === -1) continue; + + const range = document.createRange(); + range.setStart(node, index); + range.setEnd(node, index + text.length); + const selection = window.getSelection()!; + selection.removeAllRanges(); + selection.addRange(range); + return; + } + throw new Error(`"${text}" not found in container`); +} + +/** Assistant bubbles, streaming or persisted. The two carry different test ids. */ +function assistantBubbles(container: HTMLElement): HTMLElement[] { + return Array.from( + container.querySelectorAll( + '[data-testid="chat-message-assistant"], [data-testid="chat-message-assistant-streaming"]' + ) + ); +} + +describe('streaming answer rendered as a provisional message', () => { + beforeEach(() => { + window.getSelection()?.removeAllRanges(); + }); + + test('renders the streamed text inside a normal assistant message bubble', () => { + const { container } = render(); + const bubbles = assistantBubbles(container); + expect(bubbles).toHaveLength(1); + expect(bubbles[0].textContent).toContain('The second paragraph.'); + }); + + test('keeps the selection as more text arrives', () => { + const { container, rerender } = render(); + selectText(container, 'first paragraph'); + + rerender(); + + expect(window.getSelection()!.toString()).toBe('first paragraph'); + }); + + test('the persisted answer reuses the streaming bubble element', () => { + const { container, rerender } = render(); + const streamingBubble = assistantBubbles(container)[0]; + + rerender( + + ); + + const bubbles = assistantBubbles(container); + expect(bubbles).toHaveLength(1); + expect(bubbles[0]).toBe(streamingBubble); + }); + + // The E2E suites wait on `chat-message-assistant` and read the answer from + // `chat-message-content`, both meaning the answer is complete. The streaming + // bubble must satisfy neither. + test('takes the persisted test ids only once the answer is persisted', () => { + const { container, rerender } = render(); + expect(container.querySelectorAll('[data-testid="chat-message-assistant"]')).toHaveLength(0); + expect(container.querySelectorAll('[data-testid="chat-message-assistant-streaming"]')).toHaveLength(1); + const streamingBubble = assistantBubbles(container)[0]; + expect(streamingBubble.querySelector('[data-testid="chat-message-content"]')).toBeNull(); + expect(streamingBubble.querySelector('[data-testid="chat-message-content-streaming"]')).not.toBeNull(); + + rerender( + + ); + + expect(container.querySelectorAll('[data-testid="chat-message-assistant"]')).toHaveLength(1); + expect(container.querySelectorAll('[data-testid="chat-message-assistant-streaming"]')).toHaveLength(0); + expect(container.querySelectorAll('[data-testid="chat-message-content-streaming"]')).toHaveLength(0); + expect(assistantBubbles(container)[0].querySelector('[data-testid="chat-message-content"]')).not.toBeNull(); + }); + + test('keeps the selection when the streamed answer becomes a persisted message', () => { + const { container, rerender } = render(); + selectText(container, 'second paragraph'); + + rerender( + + ); + + const selection = window.getSelection()!; + expect(selection.toString()).toBe('second paragraph'); + expect(container.contains(selection.anchorNode)).toBe(true); + }); + + test('keeps the selection when the user stops the run', () => { + const { container, rerender } = render(); + selectText(container, 'second paragraph'); + + rerender( + + ); + + const selection = window.getSelection()!; + expect(selection.toString()).toBe('second paragraph'); + expect(container.contains(selection.anchorNode)).toBe(true); + }); + + // saveAIResponse appends the persisted message before the run UI state + // resets, so for one render both the persisted answer and the streaming text + // exist. Only one bubble may show. + test('shows one bubble while the persisted answer and the streaming text overlap', () => { + const { container } = render( + + ); + expect(assistantBubbles(container)).toHaveLength(1); + }); + + test('shows the plain loading indicator before any text streams', () => { + const { container } = render(); + expect(assistantBubbles(container)).toHaveLength(0); + expect(container.querySelector('.dots')).not.toBeNull(); + }); + + test('shows the loading indicator inside the bubble while only reasoning streams', () => { + const { container } = render( + + ); + const bubbles = assistantBubbles(container); + expect(bubbles).toHaveLength(1); + expect(bubbles[0].querySelector('.dots')).not.toBeNull(); + }); + + test('does not show a timestamp or feedback buttons while streaming', () => { + const { container } = render( {} })} />); + const bubble = assistantBubbles(container)[0]; + expect(bubble.querySelector('[data-testid="feedback-buttons"]')).toBeNull(); + expect(bubble.querySelector('[data-testid="message-timestamp"]')).toBeNull(); + }); + + test('falls back to a bubble without a stable key when no streaming id is given', () => { + const { container } = render(); + expect(assistantBubbles(container)).toHaveLength(1); + }); +}); diff --git a/packages/client/src/components/UseAIChatPanel.tsx b/packages/client/src/components/UseAIChatPanel.tsx index 751b7505..8ffdf2f2 100644 --- a/packages/client/src/components/UseAIChatPanel.tsx +++ b/packages/client/src/components/UseAIChatPanel.tsx @@ -138,6 +138,18 @@ function fileChipInfo(part: PersistedContentPart): { name: string; size: number return null; } +/** + * A message as shown in the panel. `streaming` marks the provisional entry + * for the answer that is still arriving; it has no timestamp, feedback + * buttons or persisted reasoning parts yet. + */ +type DisplayMessage = PersistedMessage & { streaming?: boolean }; + +/** Key for the provisional streaming message when no persisted id is known. */ +const PROVISIONAL_MESSAGE_ID = 'streaming-answer'; +/** Placeholder; the provisional message never shows a timestamp. */ +const PROVISIONAL_CREATED_AT = new Date(0); + /** * Props for the chat panel component. */ @@ -157,6 +169,15 @@ export interface UseAIChatPanelProps { streamingText?: string; /** Currently streaming reasoning text from extended thinking */ streamingReasoning?: string; + /** + * Id the streaming answer will be persisted under. While set, the streaming + * answer renders as a provisional message with this id, so when the persisted + * answer arrives under the same id React updates the bubble in place instead + * of unmounting it. Without it the provisional bubble remounts on completion. + * @default null + * @example "msg_1723972800000_k3j9x2a" + */ + streamingMessageId?: string | null; currentChatId?: string | null; onNewChat?: () => Promise; onLoadChat?: (chatId: string) => Promise; @@ -230,6 +251,7 @@ export function UseAIChatPanel({ connected, streamingText = '', streamingReasoning = '', + streamingMessageId = null, currentChatId, onNewChat, onLoadChat, @@ -271,6 +293,27 @@ export function UseAIChatPanel({ // a unified bubble to the user. const displayMessages = mergeAssistantMessagesForDisplay(messages); + // The streaming answer is shown as a provisional assistant message under the + // id it will be persisted with. When the persisted message arrives it takes + // the same key, so React updates the existing bubble instead of replacing it, + // and a text selection inside it survives. saveAIResponse appends the + // persisted message one render before the streaming state clears, so the + // provisional entry is skipped once a message with that id exists. + const provisionalMessage: DisplayMessage | null = + (streamingText || streamingReasoning) && + !(streamingMessageId && displayMessages.some((m) => m.id === streamingMessageId)) + ? { + id: streamingMessageId ?? PROVISIONAL_MESSAGE_ID, + role: 'assistant', + content: streamingText, + createdAt: PROVISIONAL_CREATED_AT, + streaming: true, + } + : null; + const renderedMessages: DisplayMessage[] = provisionalMessage + ? [...displayMessages, provisionalMessage] + : displayMessages; + const [input, setInput] = useState(''); const chatHistoryDropdown = useDropdownState(); const agentDropdown = useDropdownState(); @@ -869,7 +912,7 @@ export function UseAIChatPanel({ )} - {displayMessages.map((message) => { + {renderedMessages.map((message) => { // Info notices (e.g. the abort "generation stopped" bubble) are // display-only system messages. Render a compact, centered pill — // no reasoning dropdown, markdown, feedback, or hover affordances. @@ -900,10 +943,24 @@ export function UseAIChatPanel({ ); } + // The provisional bubble is deliberately kept out of + // `chat-message-assistant` and `chat-message-content`: the E2E suites + // wait on those test ids to mean "the answer is done", and a bubble + // that appears with the first token would satisfy that wait + // mid-stream. Only the attributes change when the answer is + // persisted, so the elements themselves, and any selection inside + // them, are untouched. + const messageTestId = message.streaming + ? 'chat-message-assistant-streaming' + : `chat-message-${message.role}`; + const contentTestId = message.streaming + ? 'chat-message-content-streaming' + : 'chat-message-content'; + return (
    )}
    - {message.reasoningParts && message.reasoningParts.length > 0 && ( + {message.streaming && streamingReasoning && ( + + )} + {!message.streaming && message.reasoningParts && message.reasoningParts.length > 0 && ( )} - + {message.streaming && !streamingText ? ( + ... + ) : ( + + )} ) : ( // User/tool bubbles: display-only text so transformed_file @@ -1019,7 +1089,7 @@ export function UseAIChatPanel({ })}
    {/* Feedback buttons - only for assistant messages with traceId */} - {message.role === 'assistant' && message.traceId && feedbackEnabled && onFeedback && ( + {message.role === 'assistant' && !message.streaming && message.traceId && feedbackEnabled && onFeedback && (
    )} -
    - {message.createdAt.toLocaleTimeString([], { - hour: '2-digit', - minute: '2-digit' - })} -
    + {!message.streaming && ( +
    + {message.createdAt.toLocaleTimeString([], { + hour: '2-digit', + minute: '2-digit' + })} +
    + )}
    ); })} - {loading && ( + {loading && !provisionalMessage && (
    - {streamingText || streamingReasoning ? ( - <> - {streamingReasoning && ( - - )} - {streamingText && } - {!streamingText && ( - ... - )} - - ) : fileProcessing && fileProcessing.status === 'processing' ? ( + {fileProcessing && fileProcessing.status === 'processing' ? (
    {strings.input.processingFile} {fileProcessing.progress != null && ( diff --git a/packages/client/src/hooks/useChatManagement.ts b/packages/client/src/hooks/useChatManagement.ts index 7b3c44ae..542a0db5 100644 --- a/packages/client/src/hooks/useChatManagement.ts +++ b/packages/client/src/hooks/useChatManagement.ts @@ -318,13 +318,21 @@ export function useChatManagement({ } }, [repository, setMessages]); - /** Saves an AI response to storage and updates UI. */ + /** + * Saves an AI response to storage and updates UI. + * + * `messageId` lets the caller choose the persisted id. The streaming answer + * is rendered under the id allocated at RUN_STARTED, so persisting the final + * answer under that same id lets React reuse the streaming bubble instead of + * mounting a new one. + */ const saveAIResponse = useCallback(async ( content: string, displayMode?: MessageDisplayMode, traceId?: string, turnMessages?: PersistedMessage[], - reasoningParts?: ReasoningPart[] + reasoningParts?: ReasoningPart[], + messageId?: string ): Promise => { const currentChatIdValue = currentChatIdSnapshot.current; const pendingChatIdValue = pendingChatIdSnapshot.current; @@ -350,7 +358,7 @@ export function useChatManagement({ } const finalMessage: PersistedMessage = { - id: generateMessageId(), + id: messageId ?? generateMessageId(), role: 'assistant', content, createdAt: new Date(), diff --git a/packages/client/src/hooks/useServerEvents.streamingId.test.tsx b/packages/client/src/hooks/useServerEvents.streamingId.test.tsx new file mode 100644 index 00000000..fb4581f7 --- /dev/null +++ b/packages/client/src/hooks/useServerEvents.streamingId.test.tsx @@ -0,0 +1,105 @@ +import { describe, it, expect, mock } from 'bun:test'; +import { renderHook, act } from '@testing-library/react'; +import { useServerEvents } from './useServerEvents'; +import { defaultStrings } from '../theme/strings'; +import { ErrorCode, EventType } from '../types'; +import type { UseToolSystemReturn } from './useToolSystem'; +import type { UseAIClient } from '../client'; + +function makeToolSystem(): UseToolSystemReturn { + return { + aggregatedToolsRef: { current: {} }, + executeToolCall: mock(async () => undefined), + storePendingToolCall: mock(), + handleApprovalRequest: mock(), + } as unknown as UseToolSystemReturn; +} + +function makeClient(currentMessageContent: string) { + return { + messages: [], + currentMessageContent, + currentReasoningBlocks: [], + currentRunId: 'run-1', + finalizeRun: mock(() => {}), + currentToolCalls: new Map(), + } as unknown as UseAIClient; +} + +/** + * The streaming answer and the persisted answer render under the same React + * key, so the id the answer will be saved with must be known while it streams. + */ +describe('useServerEvents — streaming message id', () => { + it('has no id before a run starts', () => { + const { result } = renderHook(() => + useServerEvents({ toolSystem: makeToolSystem(), saveAIResponse: mock(async () => {}), strings: defaultStrings }), + ); + expect(result.current.streamingMessageId).toBeNull(); + }); + + it('assigns an id at RUN_STARTED and persists the answer under that id at RUN_FINISHED', async () => { + const saveAIResponse = mock(async () => {}); + const { result } = renderHook(() => + useServerEvents({ toolSystem: makeToolSystem(), saveAIResponse, strings: defaultStrings }), + ); + const client = makeClient('Hello'); + + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.RUN_STARTED, threadId: 't', runId: 'run-1' }); + }); + const id = result.current.streamingMessageId; + expect(id).toEqual(expect.any(String)); + + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.TEXT_MESSAGE_CONTENT, messageId: 'm', delta: 'Hello' }); + await result.current.handleServerEvent(client, { type: EventType.RUN_FINISHED, threadId: 't', runId: 'run-1' }); + }); + + expect(saveAIResponse).toHaveBeenCalledTimes(1); + const [content, displayMode, , , , messageId] = saveAIResponse.mock.calls[0]; + expect(content).toBe('Hello'); + expect(displayMode).toBeUndefined(); + expect(messageId).toBe(id); + expect(result.current.streamingMessageId).toBeNull(); + }); + + it('persists the partial answer under the streaming id when the user stops the run', async () => { + const saveAIResponse = mock(async () => {}); + const { result } = renderHook(() => + useServerEvents({ toolSystem: makeToolSystem(), saveAIResponse, strings: defaultStrings }), + ); + const client = makeClient('Partial'); + + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.RUN_STARTED, threadId: 't', runId: 'run-1' }); + }); + const id = result.current.streamingMessageId; + + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.RUN_ERROR, message: ErrorCode.ABORTED }); + }); + + expect(saveAIResponse).toHaveBeenCalledTimes(2); + expect(saveAIResponse.mock.calls[0][5]).toBe(id); + // The abort notice is a separate bubble and must not reuse the answer's id. + expect(saveAIResponse.mock.calls[1][1]).toBe('info'); + expect(saveAIResponse.mock.calls[1][5]).toBeUndefined(); + }); + + it('uses a fresh id for each run', async () => { + const { result } = renderHook(() => + useServerEvents({ toolSystem: makeToolSystem(), saveAIResponse: mock(async () => {}), strings: defaultStrings }), + ); + const client = makeClient(''); + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.RUN_STARTED, threadId: 't', runId: 'run-1' }); + }); + const first = result.current.streamingMessageId; + await act(async () => { + await result.current.handleServerEvent(client, { type: EventType.RUN_FINISHED, threadId: 't', runId: 'run-1' }); + await result.current.handleServerEvent(client, { type: EventType.RUN_STARTED, threadId: 't', runId: 'run-2' }); + }); + expect(result.current.streamingMessageId).not.toBe(first); + }); +}); diff --git a/packages/client/src/hooks/useServerEvents.ts b/packages/client/src/hooks/useServerEvents.ts index 6c8a293d..6c5b44f3 100644 --- a/packages/client/src/hooks/useServerEvents.ts +++ b/packages/client/src/hooks/useServerEvents.ts @@ -13,6 +13,7 @@ import type { } from '../types'; import { EventType, ErrorCode, TOOL_APPROVAL_REQUEST } from '../types'; import type { UseAIClient } from '../client'; +import { generateMessageId } from '../providers/chatRepository/types'; import type { UseToolSystemReturn } from './useToolSystem'; import type { UseAIStrings } from '../theme'; import type { PersistedMessage, MessageDisplayMode } from '../providers/chatRepository/types'; @@ -22,7 +23,14 @@ export interface UseServerEventsOptions { /** Tool system for executing tools and looking up tool metadata */ toolSystem: UseToolSystemReturn; /** Saves an AI response to chat storage */ - saveAIResponse: (content: string, displayMode?: MessageDisplayMode, traceId?: string, turnMessages?: PersistedMessage[], reasoningParts?: ReasoningPart[]) => Promise; + saveAIResponse: ( + content: string, + displayMode?: MessageDisplayMode, + traceId?: string, + turnMessages?: PersistedMessage[], + reasoningParts?: ReasoningPart[], + messageId?: string, + ) => Promise; /** UI strings for error messages and tool execution fallbacks */ strings: UseAIStrings; } @@ -38,6 +46,14 @@ export interface UseServerEventsReturn { setLoading: React.Dispatch>; /** Current streaming text from the AI response */ streamingText: string; + /** + * Id the streaming answer will be persisted under, allocated at RUN_STARTED + * and cleared when the run ends. The chat panel renders the streaming answer + * as a provisional message with this id, so the persisted answer replaces it + * under the same React key. + * @example "msg_1723972800000_k3j9x2a" + */ + streamingMessageId: string | null; /** Clear streaming text (e.g., when starting a new message) */ clearStreamingText: () => void; /** Currently executing tool info for UI display, or null */ @@ -81,6 +97,9 @@ export function useServerEvents({ const [loading, setLoading] = useState(false); const [streamingText, setStreamingText] = useState(''); const [streamingReasoning, setStreamingReasoning] = useState(''); + const [streamingMessageId, setStreamingMessageId] = useState(null); + // Mirror of streamingMessageId for the stable event handler. + const streamingMessageIdRef = useRef(null); const streamingChatIdRef = useRef(null); // Mirror of `loading` for use from stable callbacks (handleDisconnect) where @@ -144,7 +163,7 @@ export function useServerEvents({ // races and the second save clobbers the first — dropping the turn's // tool context (only surfaces after reload, since in-memory state uses a // functional setMessages update that keeps both). - await saveAIResponseRef.current(content, undefined, opts.traceId, turnMessages, reasoningParts); + await saveAIResponseRef.current(content, undefined, opts.traceId, turnMessages, reasoningParts, streamingMessageIdRef.current ?? undefined); await saveAIResponseRef.current(notice, 'info'); } else { // No trailing text: skip the empty placeholder bubble. Attach @@ -157,13 +176,15 @@ export function useServerEvents({ // RUN_FINISHED: only persist when the AI produced a final text response. if (content) { - await saveAIResponseRef.current(content, undefined, opts.traceId, turnMessages, reasoningParts); + await saveAIResponseRef.current(content, undefined, opts.traceId, turnMessages, reasoningParts, streamingMessageIdRef.current ?? undefined); } }, []); const resetRunUiState = useCallback(() => { setStreamingText(''); setStreamingReasoning(''); + setStreamingMessageId(null); + streamingMessageIdRef.current = null; streamingChatIdRef.current = null; // Clear executingTool in case TOOL_CALL_END was never received // (e.g., stream truncated by token limit, or aborted mid-tool). @@ -183,6 +204,9 @@ export function useServerEvents({ runIdAtRunStartRef.current = client.currentRunId ?? undefined; hasTextFromPriorStepRef.current = false; setStreamingReasoning(''); + const messageId = generateMessageId(); + streamingMessageIdRef.current = messageId; + setStreamingMessageId(messageId); } else if (event.type === EventType.REASONING_MESSAGE_START) { // Add paragraph separator between reasoning from different steps setStreamingReasoning(prev => prev ? prev + '\n\n' : prev); @@ -282,12 +306,8 @@ export function useServerEvents({ const message = strs.errors[ErrorCode.CONNECTION_LOST] || strs.errors[ErrorCode.UNKNOWN_ERROR]; saveAIResponseRef.current(message, 'error'); - setStreamingText(''); - setStreamingReasoning(''); - streamingChatIdRef.current = null; - setExecutingTool(null); - setLoading(false); - }, []); + resetRunUiState(); + }, [resetRunUiState]); // Compute display value for UI const executingTool = executingToolRaw ? { @@ -298,6 +318,7 @@ export function useServerEvents({ loading, setLoading, streamingText, + streamingMessageId, clearStreamingText, executingTool, streamingChatIdRef, diff --git a/packages/client/src/providers/useAIProvider.tsx b/packages/client/src/providers/useAIProvider.tsx index 08ebde52..005922b8 100644 --- a/packages/client/src/providers/useAIProvider.tsx +++ b/packages/client/src/providers/useAIProvider.tsx @@ -769,6 +769,8 @@ export function UseAIProvider({ ? serverEvents.streamingText : ''; const effectiveStreamingReasoning = serverEvents.streamingChatIdRef.current === chatManagement.displayedChatId ? serverEvents.streamingReasoning : ''; + const effectiveStreamingMessageId = serverEvents.streamingChatIdRef.current === chatManagement.displayedChatId + ? serverEvents.streamingMessageId : null; const chatUIContextValue: ChatUIContextValue = { connected, @@ -778,6 +780,7 @@ export function UseAIProvider({ messages, streamingText: effectiveStreamingText, streamingReasoning: effectiveStreamingReasoning, + streamingMessageId: effectiveStreamingMessageId, suggestions: promptState.aggregatedSuggestions, fileUploadConfig, fileProcessing: fileProcessingState, @@ -833,6 +836,7 @@ export function UseAIProvider({ connected, streamingText: effectiveStreamingText, streamingReasoning: effectiveStreamingReasoning, + streamingMessageId: effectiveStreamingMessageId, currentChatId: chatManagement.displayedChatId, onNewChat: chatManagement.createNewChat, onLoadChat: chatManagement.loadChat,