From ddafdbcbce53e7fd43acfcdb1423a3d4145ea7cf Mon Sep 17 00:00:00 2001 From: Ethan <39577870+ethanndickson@users.noreply.github.com> Date: Thu, 2 Apr 2026 00:22:33 +1100 Subject: [PATCH] fix(site): use imperative setValue for chat message editing instead of key-based remount (#23799) Sometimes clicking **Edit** on a chat message does not populate the composer with the message text, and the edit flow had a few timing bugs around Lexical hydration. The composer was relying on `key={initialValue}` on `LexicalComposer`, so re-editing the same text could produce no state change, no remount, and an empty editor. This PR keeps the editor mounted and switches edit flows to an imperative `setValue()` API on `ChatMessageInputRef`. It also hardens that API so draft reads and writes stay correct across initial hydration: canceling edit no longer refocuses on mobile, pre-edit draft snapshots preserve persisted drafts, early `setValue()` calls buffer until the editor is ready, and `getValue()` falls back before readiness but reads live editor state after attach. --- .../ChatMessageInput.test.tsx | 98 +++++ .../ChatMessageInput/ChatMessageInput.tsx | 133 +++---- .../pages/AgentsPage/AgentChatPage.test.ts | 365 ++++++++++++------ site/src/pages/AgentsPage/AgentChatPage.tsx | 95 ++--- .../AgentsPage/AgentChatPageView.stories.tsx | 6 +- .../pages/AgentsPage/AgentChatPageView.tsx | 7 +- .../AgentsPage/components/AgentChatInput.tsx | 2 +- .../AgentsPage/components/ChatPageContent.tsx | 10 +- 8 files changed, 477 insertions(+), 239 deletions(-) create mode 100644 site/src/components/ChatMessageInput/ChatMessageInput.test.tsx diff --git a/site/src/components/ChatMessageInput/ChatMessageInput.test.tsx b/site/src/components/ChatMessageInput/ChatMessageInput.test.tsx new file mode 100644 index 0000000000..5cafc7f13a --- /dev/null +++ b/site/src/components/ChatMessageInput/ChatMessageInput.test.tsx @@ -0,0 +1,98 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import { type FC, useLayoutEffect, useRef, useState } from "react"; +import { describe, expect, it } from "vitest"; +import { ChatMessageInput, type ChatMessageInputRef } from "./ChatMessageInput"; + +const InitialValueHarness: FC<{ initialValue: string }> = ({ + initialValue, +}) => { + const inputRef = useRef(null); + const [observedValue, setObservedValue] = useState(""); + + useLayoutEffect(() => { + setObservedValue(inputRef.current?.getValue() ?? ""); + }, []); + + return ( + <> +
{observedValue}
+ + + ); +}; + +const QueuedReplacementHarness: FC<{ + initialValue: string; + replacementValue: string; +}> = ({ initialValue, replacementValue }) => { + const inputRef = useRef(null); + const [observedValue, setObservedValue] = useState(""); + + useLayoutEffect(() => { + inputRef.current?.setValue(replacementValue); + setObservedValue(inputRef.current?.getValue() ?? ""); + }, [replacementValue]); + + return ( + <> +
{observedValue}
+ + + ); +}; + +describe("ChatMessageInput", () => { + it("returns the initial draft before the editor visually hydrates", async () => { + render(); + + expect(screen.getByTestId("observed-value")).toHaveTextContent( + "persisted draft", + ); + await waitFor(() => { + expect(screen.getByTestId("chat-message-input").textContent).toBe( + "persisted draft", + ); + }); + }); + + it("queues setValue calls made before the editor is ready", async () => { + render( + , + ); + + expect(screen.getByTestId("observed-value")).toHaveTextContent( + "queued replacement", + ); + await waitFor(() => { + expect(screen.getByTestId("chat-message-input").textContent).toBe( + "queued replacement", + ); + }); + }); + + it("returns updated content even without an external onChange prop", async () => { + const inputRef = { current: null as ChatMessageInputRef | null }; + render(); + + await waitFor(() => { + expect(inputRef.current).not.toBeNull(); + }); + + inputRef.current?.insertText("typed content"); + + await waitFor(() => { + expect(inputRef.current?.getValue()).toBe("typed content"); + }); + }); +}); diff --git a/site/src/components/ChatMessageInput/ChatMessageInput.tsx b/site/src/components/ChatMessageInput/ChatMessageInput.tsx index e7b43265bd..2d18d7cb89 100644 --- a/site/src/components/ChatMessageInput/ChatMessageInput.tsx +++ b/site/src/components/ChatMessageInput/ChatMessageInput.tsx @@ -11,7 +11,6 @@ import { $createTextNode, $getRoot, $getSelection, - $insertNodes, $isRangeSelection, COMMAND_PRIORITY_HIGH, FORMAT_ELEMENT_COMMAND, @@ -98,6 +97,28 @@ function insertPlainTextIntoEditor(editor: LexicalEditor, text: string) { }); } +function replacePlainTextInEditor(editor: LexicalEditor, text: string) { + editor.update(() => { + const root = $getRoot(); + root.clear(); + const paragraph = $createParagraphNode(); + root.append(paragraph); + if (!text) { + paragraph.select(); + return; + } + paragraph.select(); + const selection = $getSelection(); + if ($isRangeSelection(selection)) { + selection.insertText(text); + return; + } + const textNode = $createTextNode(text); + paragraph.append(textNode); + textNode.selectEnd(); + }); +} + // Intercepts paste events and inserts clipboard content as plain text, // stripping any rich-text formatting. Image files and large pasted text // are forwarded to the parent via the onFilePaste callback instead. @@ -296,35 +317,6 @@ const ContentChangePlugin: FC<{ return null; }); -// Seeds the editor with an initial value on first mount. -const ValueSyncPlugin: FC<{ initialValue?: string }> = memo( - function ValueSyncPlugin({ initialValue }) { - const [editor] = useLexicalComposerContext(); - const hasInitialized = useRef(false); - - useEffect(() => { - if (!hasInitialized.current && initialValue !== undefined) { - hasInitialized.current = true; - - if (initialValue === "") { - return; - } - - editor.update(() => { - const root = $getRoot(); - root.clear(); - const paragraph = $createParagraphNode(); - const textNode = $createTextNode(initialValue); - paragraph.append(textNode); - root.append(paragraph); - }); - } - }, [editor, initialValue]); - - return null; - }, -); - // Exposes the LexicalEditor instance to the parent via a callback // so it can be stored in a ref for imperative access. const InsertTextPlugin: FC<{ @@ -361,6 +353,10 @@ type EditorContentPart = }; export interface ChatMessageInputRef { + /** + * Replace the editor's plain-text content in a single Lexical update. + */ + setValue: (text: string) => void; insertText: (text: string) => void; clear: () => void; focus: () => void; @@ -444,13 +440,37 @@ const ChatMessageInput = memo( ); const editorRef = useRef(null); + const lastKnownValueRef = useRef(initialValue ?? ""); + const pendingReplacementRef = useRef(null); + + const replaceValueOrQueue = useCallback((text: string) => { + lastKnownValueRef.current = text; + const editor = editorRef.current; + if (!editor) { + pendingReplacementRef.current = text; + return; + } + pendingReplacementRef.current = null; + replacePlainTextInEditor(editor, text); + }, []); const handleEditorReady = useCallback((editor: LexicalEditor) => { editorRef.current = editor; + const pendingReplacement = pendingReplacementRef.current; + if (pendingReplacement !== null) { + pendingReplacementRef.current = null; + replacePlainTextInEditor(editor, pendingReplacement); + return; + } + const initialText = lastKnownValueRef.current; + if (initialText) { + replacePlainTextInEditor(editor, initialText); + } }, []); const handleContentChange = useCallback( (content: string, hasFileReferences: boolean) => { + lastKnownValueRef.current = content; onChange?.(content, hasFileReferences); }, [onChange], @@ -459,51 +479,17 @@ const ChatMessageInput = memo( useImperativeHandle( ref, () => ({ + setValue: (text: string) => { + replaceValueOrQueue(text); + }, insertText: (text: string) => { const editor = editorRef.current; if (!editor) return; - editor.update(() => { - const selection = $getSelection(); - if ($isRangeSelection(selection)) { - const textNode = $createTextNode(text); - $insertNodes([textNode]); - textNode.selectEnd(); - } else { - const root = $getRoot(); - const lastChild = root.getLastChild(); - if (lastChild) { - if (lastChild.getType() === "paragraph") { - const paragraph = lastChild as ParagraphNode; - const textNode = $createTextNode(text); - paragraph.append(textNode); - textNode.selectEnd(); - } else { - const textNode = $createTextNode(text); - lastChild.insertAfter(textNode); - textNode.selectEnd(); - } - } else { - const paragraph = $createParagraphNode(); - const textNode = $createTextNode(text); - paragraph.append(textNode); - root.append(paragraph); - textNode.selectEnd(); - } - } - }); + insertPlainTextIntoEditor(editor, text); }, clear: () => { - const editor = editorRef.current; - if (!editor) return; - - editor.update(() => { - const root = $getRoot(); - root.clear(); - const paragraph = $createParagraphNode(); - root.append(paragraph); - paragraph.select(); - }); + replaceValueOrQueue(""); }, focus: () => { const editor = editorRef.current; @@ -524,7 +510,9 @@ const ChatMessageInput = memo( }, getValue: () => { const editor = editorRef.current; - if (!editor) return ""; + if (!editor) { + return lastKnownValueRef.current; + } let content = ""; editor.getEditorState().read(() => { content = $getRoot().getTextContent(); @@ -600,11 +588,11 @@ const ChatMessageInput = memo( return parts; }, }), - [], + [replaceValueOrQueue], ); return ( - +
*]:col-start-1 [&>*]:row-start-1", @@ -639,7 +627,6 @@ const ChatMessageInput = memo( /> - {autoFocus && } diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 6ef8877157..35f99cff67 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -3,154 +3,299 @@ import { createRef } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { draftInputStorageKeyPrefix, + getPersistedDraftInputValue, useConversationEditingState, } from "./AgentChatPage"; import type { ChatMessageInputRef } from "./components/AgentChatInput"; +type MockChatInputHandle = { + handle: ChatMessageInputRef; + setValue: ReturnType; + clear: ReturnType; + focus: ReturnType; + getValue: ReturnType; + currentValue: { value: string }; +}; + +const createMockChatInputHandle = (initialValue = ""): MockChatInputHandle => { + const currentValue = { value: initialValue }; + const setValue = vi.fn((text: string) => { + currentValue.value = text; + }); + const clear = vi.fn(() => { + currentValue.value = ""; + }); + const focus = vi.fn(); + const getValue = vi.fn(() => currentValue.value); + + return { + handle: { + setValue, + insertText: vi.fn(), + clear, + focus, + getValue, + addFileReference: vi.fn(), + getContentParts: vi.fn(() => []), + }, + setValue, + clear, + focus, + getValue, + currentValue, + }; +}; + +const setMobileViewport = (isMobile: boolean) => { + Object.defineProperty(window, "matchMedia", { + writable: true, + value: vi.fn((query: string): MediaQueryList => { + return { + matches: query === "(max-width: 639px)" ? isMobile : false, + media: query, + onchange: null, + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + dispatchEvent: vi.fn(() => true), + addListener: vi.fn(), + removeListener: vi.fn(), + } as MediaQueryList; + }), + }); +}; + +describe("getPersistedDraftInputValue", () => { + const chatID = "chat-abc-123"; + const expectedKey = `${draftInputStorageKeyPrefix}${chatID}`; + + beforeEach(() => { + localStorage.clear(); + setMobileViewport(false); + }); + + it("reads the initial value from localStorage for a given chatID", () => { + localStorage.setItem(expectedKey, "saved draft"); + + expect(getPersistedDraftInputValue(chatID)).toBe("saved draft"); + }); + + it("returns empty string when localStorage has no draft", () => { + expect(getPersistedDraftInputValue(chatID)).toBe(""); + }); +}); + describe("useConversationEditingState", () => { const chatID = "chat-abc-123"; const expectedKey = `${draftInputStorageKeyPrefix}${chatID}`; beforeEach(() => { localStorage.clear(); + setMobileViewport(false); }); - const renderEditing = (id: string | undefined = chatID) => { + const renderEditing = () => { const onSend = vi.fn().mockResolvedValue(undefined); const onDeleteQueuedMessage = vi.fn().mockResolvedValue(undefined); const chatInputRef = createRef(); - const inputValueRef: import("react").RefObject = { current: "" }; const hook = renderHook(() => useConversationEditingState({ - chatID: id, + chatID, onSend, onDeleteQueuedMessage, chatInputRef, - inputValueRef, }), ); - return { ...hook, onSend, onDeleteQueuedMessage }; + return { ...hook, onSend }; }; - it("reads the initial value from localStorage for a given chatID", () => { - localStorage.setItem(expectedKey, "saved draft"); - - const { result, unmount } = renderEditing(); - - expect(result.current.editorInitialValue).toBe("saved draft"); - expect(result.current.inputValueRef.current).toBe("saved draft"); - unmount(); - }); - - it("returns empty string when localStorage has no draft", () => { - const { result, unmount } = renderEditing(); - - expect(result.current.editorInitialValue).toBe(""); - expect(result.current.inputValueRef.current).toBe(""); - unmount(); - }); - - it("writes content to localStorage via handleContentChange", () => { + it("persists and removes drafts via handleContentChange", () => { const { result, unmount } = renderEditing(); act(() => { result.current.handleContentChange("work in progress"); }); - expect(localStorage.getItem(expectedKey)).toBe("work in progress"); - expect(result.current.inputValueRef.current).toBe("work in progress"); - unmount(); - }); - - it("removes the draft key when handleContentChange receives empty string", () => { - localStorage.setItem(expectedKey, "old draft"); - const { result, unmount } = renderEditing(); act(() => { result.current.handleContentChange(""); }); + expect(localStorage.getItem(expectedKey)).toBeNull(); + unmount(); + }); + + it("loads edit text into the composer and restores the prior draft on cancel without refocusing", () => { + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("work in progress"); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleEditUserMessage(7, "edited message"); + }); + + expect(result.current.editingMessageId).toBe(7); + expect(mockInput.setValue).toHaveBeenCalledWith("edited message"); + expect(mockInput.focus).toHaveBeenCalledTimes(1); + + act(() => { + result.current.handleCancelHistoryEdit(); + }); + + expect(result.current.editingMessageId).toBeNull(); + expect(mockInput.setValue).toHaveBeenLastCalledWith("work in progress"); + expect(mockInput.currentValue.value).toBe("work in progress"); + expect(mockInput.focus).toHaveBeenCalledTimes(1); + unmount(); + }); + + it("loads queue edit text into the composer and restores the prior draft on cancel without refocusing", () => { + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("work in progress"); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleStartQueueEdit(9, "queued message", []); + }); + + expect(result.current.editingQueuedMessageID).toBe(9); + expect(mockInput.setValue).toHaveBeenCalledWith("queued message"); + expect(mockInput.focus).toHaveBeenCalledTimes(1); + + act(() => { + result.current.handleCancelQueueEdit(); + }); + + expect(result.current.editingQueuedMessageID).toBeNull(); + expect(mockInput.setValue).toHaveBeenLastCalledWith("work in progress"); + expect(mockInput.currentValue.value).toBe("work in progress"); + expect(mockInput.focus).toHaveBeenCalledTimes(1); + unmount(); + }); + + it("does not force focus when replacing input values on mobile", () => { + setMobileViewport(true); + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("draft before edit"); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleEditUserMessage(7, "edited message"); + }); + expect(mockInput.focus).not.toHaveBeenCalled(); + + act(() => { + result.current.handleCancelHistoryEdit(); + }); + expect(mockInput.focus).not.toHaveBeenCalled(); + + act(() => { + result.current.handleStartQueueEdit(9, "queued message", []); + }); + expect(mockInput.focus).not.toHaveBeenCalled(); + + act(() => { + result.current.handleCancelQueueEdit(); + }); + expect(mockInput.focus).not.toHaveBeenCalled(); + unmount(); + }); + + it("falls back to the persisted draft when history edit starts before hydration", () => { + localStorage.setItem(expectedKey, "persisted draft"); + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle(""); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleEditUserMessage(7, "edited message"); + }); + + act(() => { + result.current.handleCancelHistoryEdit(); + }); + + expect(mockInput.setValue).toHaveBeenLastCalledWith("persisted draft"); + expect(mockInput.currentValue.value).toBe("persisted draft"); + unmount(); + }); + + it("falls back to the persisted draft when queue edit starts before hydration", () => { + localStorage.setItem(expectedKey, "persisted draft"); + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle(""); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleStartQueueEdit(9, "queued message", []); + }); + + act(() => { + result.current.handleCancelQueueEdit(); + }); + + expect(mockInput.setValue).toHaveBeenLastCalledWith("persisted draft"); + expect(mockInput.currentValue.value).toBe("persisted draft"); + unmount(); + }); + + it("prefers the live editor value over stale persisted draft state", () => { + localStorage.setItem(expectedKey, "stale persisted draft"); + const { result, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("live draft"); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleEditUserMessage(7, "edited message"); + }); + + act(() => { + result.current.handleCancelHistoryEdit(); + }); + + expect(mockInput.setValue).toHaveBeenLastCalledWith("live draft"); + expect(mockInput.currentValue.value).toBe("live draft"); + unmount(); + }); + + it("can load the same edit text again after send without relying on a remount", async () => { + const { result, onSend, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle(); + result.current.chatInputRef.current = mockInput.handle; + + act(() => { + result.current.handleEditUserMessage(7, "hello"); + }); + + await act(async () => { + await result.current.handleSendFromInput("hello"); + }); + + act(() => { + result.current.handleEditUserMessage(7, "hello"); + }); + + expect(onSend).toHaveBeenCalledWith("hello", undefined, 7); + expect(mockInput.setValue).toHaveBeenNthCalledWith(1, "hello"); + expect(mockInput.setValue).toHaveBeenNthCalledWith(2, "hello"); + unmount(); + }); + + it("clears the composer and persisted draft after a successful send", async () => { + localStorage.setItem(expectedKey, "draft to clear"); + const { result, onSend, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("hello"); + result.current.chatInputRef.current = mockInput.handle; + + await act(async () => { + await result.current.handleSendFromInput("hello"); + }); + + expect(onSend).toHaveBeenCalledWith("hello", undefined, undefined); + expect(mockInput.clear).toHaveBeenCalled(); + expect(mockInput.focus).toHaveBeenCalled(); expect(localStorage.getItem(expectedKey)).toBeNull(); unmount(); }); - - it("does not write a draft key when chatID is undefined", () => { - const { result, unmount } = renderEditing(undefined); - - act(() => { - result.current.handleContentChange("should not persist"); - }); - - // The ref is still updated even without persistence. - expect(result.current.inputValueRef.current).toBe("should not persist"); - // No draft for "undefined" chatID should appear. - expect( - localStorage.getItem(`${draftInputStorageKeyPrefix}undefined`), - ).toBeNull(); - unmount(); - }); - - it("calls focus on the input ref after a successful send", async () => { - const { result, onSend, unmount } = renderEditing(); - - // Attach a mock ChatMessageInputRef to the chatInputRef - const mockFocus = vi.fn(); - const mockClear = vi.fn(); - const mockInputRef = { - focus: mockFocus, - clear: mockClear, - insertText: vi.fn(), - getValue: vi.fn().mockReturnValue(""), - addFileReference: vi.fn(), - getContentParts: vi.fn().mockReturnValue([]), - }; // The hook exposes chatInputRef – assign the mock to it. - result.current.chatInputRef.current = mockInputRef; - - await act(async () => { - result.current.handleSendFromInput("hello"); - await vi.waitFor(() => { - expect(onSend).toHaveBeenCalledWith("hello", undefined, undefined); - }); - }); - - expect(mockClear).toHaveBeenCalled(); - expect(mockFocus).toHaveBeenCalled(); - unmount(); - }); - - it("initializes with the correct draft for each chatID", () => { - const chatA = "chat-aaa"; - const chatB = "chat-bbb"; - localStorage.setItem(`${draftInputStorageKeyPrefix}${chatA}`, "draft A"); - localStorage.setItem(`${draftInputStorageKeyPrefix}${chatB}`, "draft B"); - - // Each chatID should initialize with its own draft — this is - // what the key={agentId} wrapper guarantees at the component - // level (a new chatID means a full remount). - const hookA = renderEditing(chatA); - expect(hookA.result.current.editorInitialValue).toBe("draft A"); - hookA.unmount(); - - const hookB = renderEditing(chatB); - expect(hookB.result.current.editorInitialValue).toBe("draft B"); - hookB.unmount(); - }); - - it("clears the draft from localStorage on successful send", async () => { - localStorage.setItem(expectedKey, "draft to clear"); - - const { result, unmount } = renderEditing(); - - expect(localStorage.getItem(expectedKey)).toBe("draft to clear"); - - await act(async () => { - result.current.handleSendFromInput("hello"); - await vi.waitFor(() => { - expect(localStorage.getItem(expectedKey)).toBeNull(); - }); - }); - unmount(); - }); }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 5226c27577..c449bf3a29 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -1,4 +1,4 @@ -import { type FC, useEffect, useLayoutEffect, useRef, useState } from "react"; +import { type FC, useEffect, useRef, useState } from "react"; import { useInfiniteQuery, useMutation, @@ -81,6 +81,17 @@ const lastModelConfigIDStorageKey = "agents.last-model-config-id"; /** @internal Exported for testing. */ export const draftInputStorageKeyPrefix = "agents.draft-input."; +/** @internal Exported for testing. */ +export function getPersistedDraftInputValue( + chatID: string | undefined, +): string { + if (typeof window === "undefined" || !chatID) { + return ""; + } + + return localStorage.getItem(`${draftInputStorageKeyPrefix}${chatID}`) ?? ""; +} + /** @internal Exported for testing. */ export function useConversationEditingState(deps: { chatID: string | undefined; @@ -91,31 +102,35 @@ export function useConversationEditingState(deps: { ) => Promise; onDeleteQueuedMessage: (id: number) => Promise; chatInputRef: React.RefObject; - inputValueRef: React.RefObject; }) { - const { chatID, onSend, onDeleteQueuedMessage, chatInputRef, inputValueRef } = - deps; + const { chatID, onSend, onDeleteQueuedMessage, chatInputRef } = deps; const draftStorageKey = chatID ? `${draftInputStorageKeyPrefix}${chatID}` : null; - const [editorInitialValue, setEditorInitialValue] = useState(() => { + const getDraftBeforeEdit = () => { + const currentInputValue = chatInputRef.current?.getValue() ?? ""; + if (currentInputValue) { + return currentInputValue; + } + // The editor seeds its initial value after paint, so the live editor can + // still be empty while the persisted draft already exists. if (typeof window === "undefined" || !draftStorageKey) { return ""; } return localStorage.getItem(draftStorageKey) ?? ""; - }); - - // Sync the ref with the initial draft value so callers that - // read inputValueRef.current see the persisted draft. Uses a - // layout effect so the value is available before paint. - const initialSyncDone = useRef(false); - useLayoutEffect(() => { - if (!initialSyncDone.current && editorInitialValue) { - initialSyncDone.current = true; - (inputValueRef as React.MutableRefObject).current = - editorInitialValue; + }; + const replaceInputValue = (content: string) => { + chatInputRef.current?.setValue(content); + }; + const focusInputIfDesktop = () => { + if (!isMobileViewport()) { + chatInputRef.current?.focus(); } - }, [editorInitialValue, inputValueRef]); + }; + const replaceInputValueAndFocus = (content: string) => { + replaceInputValue(content); + focusInputIfDesktop(); + }; // -- History editing state -- const [editingMessageId, setEditingMessageId] = useState(null); @@ -132,24 +147,19 @@ export function useConversationEditingState(deps: { fileBlocks?: readonly ChatMessagePart[], ) => { setDraftBeforeHistoryEdit((prev) => - editingMessageId !== null ? prev : inputValueRef.current, + editingMessageId !== null ? prev : getDraftBeforeEdit(), ); setEditingMessageId(messageId); - setEditorInitialValue(text); - inputValueRef.current = text; + replaceInputValueAndFocus(text); setEditingFileBlocks(fileBlocks ?? []); }; const handleCancelHistoryEdit = () => { - setEditorInitialValue(draftBeforeHistoryEdit ?? ""); - inputValueRef.current = draftBeforeHistoryEdit ?? ""; + const draft = draftBeforeHistoryEdit ?? ""; setEditingMessageId(null); setDraftBeforeHistoryEdit(null); setEditingFileBlocks([]); - chatInputRef.current?.clear(); - if (draftBeforeHistoryEdit) { - chatInputRef.current?.insertText(draftBeforeHistoryEdit); - } + replaceInputValue(draft); }; // -- Queue editing state -- @@ -166,20 +176,19 @@ export function useConversationEditingState(deps: { fileBlocks: readonly ChatMessagePart[], ) => { setDraftBeforeQueueEdit((prev) => - editingQueuedMessageID === null ? inputValueRef.current : prev, + editingQueuedMessageID === null ? getDraftBeforeEdit() : prev, ); setEditingQueuedMessageID(id); - setEditorInitialValue(text); - inputValueRef.current = text; + replaceInputValueAndFocus(text); setEditingFileBlocks(fileBlocks); }; const handleCancelQueueEdit = () => { - setEditorInitialValue(draftBeforeQueueEdit ?? ""); - inputValueRef.current = draftBeforeQueueEdit ?? ""; + const draft = draftBeforeQueueEdit ?? ""; setEditingQueuedMessageID(null); setDraftBeforeQueueEdit(null); setEditingFileBlocks([]); + replaceInputValue(draft); }; // Wraps the parent onSend to clear local input/editing state @@ -192,10 +201,7 @@ export function useConversationEditingState(deps: { await onSend(message, fileIds, editedMessageID); // Clear input and editing state on success. chatInputRef.current?.clear(); - if (!isMobileViewport()) { - chatInputRef.current?.focus(); - } - inputValueRef.current = ""; + focusInputIfDesktop(); if (draftStorageKey) { localStorage.removeItem(draftStorageKey); } @@ -213,7 +219,6 @@ export function useConversationEditingState(deps: { }; const handleContentChange = (content: string) => { - inputValueRef.current = content; if (draftStorageKey) { if (content) { localStorage.setItem(draftStorageKey, content); @@ -224,9 +229,7 @@ export function useConversationEditingState(deps: { }; return { - inputValueRef, chatInputRef, - editorInitialValue, editingMessageId, editingFileBlocks, handleEditUserMessage, @@ -357,10 +360,10 @@ const AgentChatPage: FC = () => { >(null); const scrollToBottomRef = useRef<(() => void) | null>(null); const chatInputRef = useRef(null); - const inputValueRef = useRef( - agentId - ? (localStorage.getItem(`${draftInputStorageKeyPrefix}${agentId}`) ?? "") - : "", + // Read once on mount — agentId is stable because KeyedAgentChatPage + // remounts the entire component when the route param changes. + const [initialInputValue] = useState(() => + getPersistedDraftInputValue(agentId), ); // Right panel open/closed state is owned here so the loading @@ -607,13 +610,15 @@ const AgentChatPage: FC = () => { const handleCommit = (repoRoot: string) => { const commitPrompt = `Commit and push the working changes in ${repoRoot}. If there are unstaged files, commit them too.`; - const current = inputValueRef.current; + const current = chatInputRef.current?.getValue() ?? ""; if (current.includes(commitPrompt)) { return; } const prefix = current.trim() ? "\n\n" : ""; chatInputRef.current?.insertText(prefix + commitPrompt); - chatInputRef.current?.focus(); + if (!isMobileViewport()) { + chatInputRef.current?.focus(); + } }; // Prefer the explicit PR number from the API, and only fall back to URL @@ -863,7 +868,6 @@ const AgentChatPage: FC = () => { onSend: handleSend, onDeleteQueuedMessage: handleDeleteQueuedMessage, chatInputRef, - inputValueRef, }); const chatTitle = chatQuery.data?.title; @@ -1031,6 +1035,7 @@ const AgentChatPage: FC = () => { isArchived={isArchived} hasWorkspace={Boolean(workspaceId)} store={store} + initialInputValue={initialInputValue} editing={editing} pendingEditMessageId={pendingEditMessageId} effectiveSelectedModel={effectiveSelectedModel} diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index 4c55cc1e81..d61ba7a9d2 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -58,7 +58,6 @@ const buildEditing = ( overrides: Partial["editing"]> = {}, ) => ({ chatInputRef: { current: null }, - editorInitialValue: "", editingMessageId: null as number | null, editingFileBlocks: [] as readonly ChatMessagePart[], handleEditUserMessage: fn(), @@ -112,6 +111,7 @@ const StoryAgentChatPageView: FC = ({ editing, ...overrides }) => { hasWorkspace: true, store: createChatStore(), pendingEditMessageId: null as number | null, + initialInputValue: "", effectiveSelectedModel: defaultModelConfigID, setSelectedModel: fn(), modelOptions: defaultModelOptions, @@ -433,8 +433,8 @@ export const EditingMessage: Story = { store={buildStoreWithMessages(editingMessages)} editing={{ editingMessageId: 3, - editorInitialValue: "Now tell me a joke", }} + initialInputValue="Now tell me a joke" /> ), }; @@ -447,8 +447,8 @@ export const EditingSaving: Story = { store={buildStoreWithMessages(editingMessages)} editing={{ editingMessageId: 3, - editorInitialValue: "Now tell me a better joke", }} + initialInputValue="Now tell me a better joke" pendingEditMessageId={3} isSubmissionPending /> diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index 93e46e871e..30a079d4de 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -31,7 +31,6 @@ type ChatStoreHandle = ReturnType["store"]; interface EditingState { chatInputRef: RefObject; - editorInitialValue: string; editingMessageId: number | null; editingFileBlocks: readonly ChatMessagePart[]; handleEditUserMessage: ( @@ -67,6 +66,9 @@ interface AgentChatPageViewProps { editing: EditingState; pendingEditMessageId: number | null; + // Input configuration. + initialInputValue: string; + // Model/input configuration. effectiveSelectedModel: string; setSelectedModel: (model: string) => void; @@ -151,6 +153,7 @@ export const AgentChatPageView: FC = ({ store, editing, pendingEditMessageId, + initialInputValue, effectiveSelectedModel, setSelectedModel, modelOptions, @@ -335,7 +338,7 @@ export const AgentChatPageView: FC = ({ modelSelectorPlaceholder={modelSelectorPlaceholder} isModelCatalogLoading={isModelCatalogLoading} inputRef={editing.chatInputRef} - initialValue={editing.editorInitialValue} + initialInputValue={initialInputValue} onContentChange={editing.handleContentChange} editingQueuedMessageID={editing.editingQueuedMessageID} onStartQueueEdit={editing.handleStartQueueEdit} diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 9f8269917a..c89d630340 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -75,7 +75,7 @@ interface AgentChatInputProps { isLoading: boolean; // Ref for the Lexical editor, exposed for imperative access. inputRef?: React.Ref; - // Initial text to seed the editor with. + // Initial text to seed the editor on first mount only. initialValue?: string; // Called on every text change inside the editor. onContentChange?: (content: string) => void; diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 467cd8d8c4..34dc4955c1 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -123,10 +123,10 @@ interface ChatPageInputProps { modelOptions: readonly ModelSelectorOption[]; modelSelectorPlaceholder: string; isModelCatalogLoading?: boolean; - // Controlled input value and editing state, owned by the - // conversation component. + // Imperative editor handle plus the one-time initial draft, + // owned by the conversation component. inputRef?: React.Ref; - initialValue?: string; + initialInputValue?: string; onContentChange?: (content: string) => void; editingQueuedMessageID: number | null; onStartQueueEdit: ( @@ -170,7 +170,7 @@ export const ChatPageInput: FC = ({ modelSelectorPlaceholder, isModelCatalogLoading = false, inputRef, - initialValue, + initialInputValue, onContentChange, editingQueuedMessageID, onStartQueueEdit, @@ -315,7 +315,7 @@ export const ChatPageInput: FC = ({ previewUrls={previewUrls} textContents={textContents} inputRef={inputRef} - initialValue={initialValue} + initialValue={initialInputValue} onContentChange={onContentChange} queuedMessages={queuedMessages} onDeleteQueuedMessage={onDeleteQueuedMessage}