From d4a3902fccb4799cfa0a4e69969e58b6ecf394d7 Mon Sep 17 00:00:00 2001 From: masudahiroto <96814344+masudahiroto@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:47:09 +0900 Subject: [PATCH 1/3] fix(client): make text selection work while an answer streams The markdown component built its element override map inside its render function, so every entry was a new function on each render and React remounted the whole answer instead of updating it. During streaming that happened on every token, which cleared any selection the user was making. Hoist the map and the remark plugins to module scope and memoize the component. The streaming answer and the persisted answer are separate elements, so the swap at the end of a run collapsed the selection as well. useSelectionHandoff records the selection as character offsets while the answer streams, then applies them to the persisted element. It restores nothing when the text does not match, or when the run ends without a new answer. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/components/MarkdownContent.test.tsx | 85 +++++ .../client/src/components/MarkdownContent.tsx | 306 +++++++++--------- .../UseAIChatPanel.selection.test.tsx | 144 +++++++++ .../client/src/components/UseAIChatPanel.tsx | 40 ++- .../client/src/hooks/useSelectionHandoff.ts | 65 ++++ .../client/src/utils/textSelection.test.ts | 110 +++++++ packages/client/src/utils/textSelection.ts | 97 ++++++ 7 files changed, 699 insertions(+), 148 deletions(-) create mode 100644 packages/client/src/components/MarkdownContent.test.tsx create mode 100644 packages/client/src/components/UseAIChatPanel.selection.test.tsx create mode 100644 packages/client/src/hooks/useSelectionHandoff.ts create mode 100644 packages/client/src/utils/textSelection.test.ts create mode 100644 packages/client/src/utils/textSelection.ts 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..7bb9ad11 100644 --- a/packages/client/src/components/MarkdownContent.tsx +++ b/packages/client/src/components/MarkdownContent.tsx @@ -1,5 +1,6 @@ import React from 'react'; import ReactMarkdown from 'react-markdown'; +import type { Components } from 'react-markdown'; import remarkGfm from 'remark-gfm'; interface MarkdownContentProps { @@ -7,153 +8,170 @@ interface MarkdownContentProps { } /** - * Renders markdown content with appropriate styling for the chat panel. + * Both of these must stay module-level constants. + * + * 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 them 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 REMARK_PLUGINS = [remarkGfm]; + +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/UseAIChatPanel.selection.test.tsx b/packages/client/src/components/UseAIChatPanel.selection.test.tsx new file mode 100644 index 00000000..37e34594 --- /dev/null +++ b/packages/client/src/components/UseAIChatPanel.selection.test.tsx @@ -0,0 +1,144 @@ +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 userMessage: PersistedMessage = { + id: 'msg-user', + role: 'user', + content: 'question', + createdAt: new Date(0), +}; + +const assistantMessage: PersistedMessage = { + id: 'msg-assistant', + role: 'assistant', + content: ANSWER, + createdAt: new Date(0), +}; + +/** + * 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, + ...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); + // Browsers fire this on every selection change; jsdom does not. + document.dispatchEvent(new window.Event('selectionchange')); + return; + } + throw new Error(`"${text}" not found in container`); +} + +describe('selecting the answer while it streams', () => { + beforeEach(() => { + window.getSelection()?.removeAllRanges(); + }); + + // Text the model has already written keeps its DOM nodes, so a selection over + // it holds while the rest of the answer streams in. (The paragraph currently + // being written is the exception: replacing a text node's data collapses any + // range inside it, which is how every browser behaves.) + 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('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'); + // Matching text alone would also pass with the range still sitting on the + // unmounted streaming bubble's detached nodes. + expect(container.querySelector('.markdown-answer')!.contains(selection.anchorNode)).toBe(true); + }); + + // The user selects text mid-stream and then presses stop. The notice bubble + // that abort appends is the last assistant message, but it renders as a plain + // pill with no answer wrapper, so treating it as the answer would leave the + // handoff with nothing to restore onto and drop the selection. + test('restores the selection onto the answer 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.querySelector('.markdown-answer')!.contains(selection.anchorNode)).toBe(true); + }); + + test('does not restore a selection the user never made', () => { + const { rerender } = render(); + + rerender( + + ); + + expect(window.getSelection()!.toString()).toBe(''); + }); +}); diff --git a/packages/client/src/components/UseAIChatPanel.tsx b/packages/client/src/components/UseAIChatPanel.tsx index 751b7505..516340e3 100644 --- a/packages/client/src/components/UseAIChatPanel.tsx +++ b/packages/client/src/components/UseAIChatPanel.tsx @@ -12,6 +12,7 @@ import type { SavedCommand } from '../commands/types'; import { useSlashCommands } from '../hooks/useSlashCommands'; import { useFileUpload } from '../hooks/useFileUpload'; import { useDropdownState } from '../hooks/useDropdownState'; +import { useSelectionHandoff } from '../hooks/useSelectionHandoff'; import { useTheme, useStrings } from '../theme'; import type { UseAIStrings, UseAITheme } from '../theme'; import { ToolApprovalDialog } from './ToolApprovalDialog'; @@ -282,6 +283,20 @@ export function UseAIChatPanel({ // Message hover state for save button const [hoveredMessageId, setHoveredMessageId] = useState(null); + // Wrappers around the rendered answer, used to carry a text selection from + // the streaming bubble to the persisted message that replaces it. + const streamingResponseRef = useRef(null); + const lastAnswerRef = useRef(null); + useSelectionHandoff({ loading, streamingRef: streamingResponseRef, persistedRef: lastAnswerRef }); + + // Last assistant answer on screen. Abort notices and error bubbles are system + // messages that happen to carry the assistant role, so they are matched by + // display mode rather than excluded one at a time. This is the bubble the + // streaming one hands its selection over to. + const lastAnswerId = [...displayMessages] + .reverse() + .find((m) => m.role === 'assistant' && (m.displayMode ?? 'default') === 'default')?.id; + // File upload hook - includes processing state for transformation progress const { attachments, @@ -1004,7 +1019,16 @@ export function UseAIChatPanel({ strings={strings} /> )} - + {/* display:contents keeps this purely a selection anchor: + it holds only the rendered answer, never the reasoning + block, so text offsets map 1:1 onto the streaming bubble. */} +
    + +
    ) : ( // User/tool bubbles: display-only text so transformed_file @@ -1098,7 +1122,11 @@ export function UseAIChatPanel({ strings={strings} /> )} - {streamingText && } + {streamingText && ( +
    + +
    + )} {!streamingText && ( ... )} @@ -1402,10 +1430,14 @@ export function UseAIChatPanel({