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.
2026-09-22 05:05:20 +08:00 · 2026-04-02 00:22:33 +11:00
parent 196dc51edf
commit ddafdbcbce
8 changed files with 477 additions and 239 deletions
@@ -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<ChatMessageInputRef>(null);
+	const [observedValue, setObservedValue] = useState("");
+
+	useLayoutEffect(() => {
+		setObservedValue(inputRef.current?.getValue() ?? "");
+	}, []);
+
+	return (
+		<>
+			<div data-testid="observed-value">{observedValue}</div>
+			<ChatMessageInput
+				ref={inputRef}
+				initialValue={initialValue}
+				aria-label="Chat message input"
+			/>
+		</>
+	);
+};
+
+const QueuedReplacementHarness: FC<{
+	initialValue: string;
+	replacementValue: string;
+}> = ({ initialValue, replacementValue }) => {
+	const inputRef = useRef<ChatMessageInputRef>(null);
+	const [observedValue, setObservedValue] = useState("");
+
+	useLayoutEffect(() => {
+		inputRef.current?.setValue(replacementValue);
+		setObservedValue(inputRef.current?.getValue() ?? "");
+	}, [replacementValue]);
+
+	return (
+		<>
+			<div data-testid="observed-value">{observedValue}</div>
+			<ChatMessageInput
+				ref={inputRef}
+				initialValue={initialValue}
+				aria-label="Chat message input"
+			/>
+		</>
+	);
+};
+
+describe("ChatMessageInput", () => {
+	it("returns the initial draft before the editor visually hydrates", async () => {
+		render(<InitialValueHarness initialValue="persisted draft" />);
+
+		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(
+			<QueuedReplacementHarness
+				initialValue="persisted draft"
+				replacementValue="queued replacement"
+			/>,
+		);
+
+		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(<ChatMessageInput ref={inputRef} aria-label="Chat message input" />);
+
+		await waitFor(() => {
+			expect(inputRef.current).not.toBeNull();
+		});
+
+		inputRef.current?.insertText("typed content");
+
+		await waitFor(() => {
+			expect(inputRef.current?.getValue()).toBe("typed content");
+		});
+	});
+});
@@ -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<LexicalEditor | null>(null);
+		const lastKnownValueRef = useRef(initialValue ?? "");
+		const pendingReplacementRef = useRef<string | null>(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 (
-			<LexicalComposer initialConfig={initialConfig} key={initialValue}>
+			<LexicalComposer initialConfig={initialConfig}>
 				<div
 					className={cn(
 						"grid w-full rounded-md bg-transparent text-base placeholder:text-content-secondary focus-visible:outline-none whitespace-pre-wrap break-words [&>*]:col-start-1 [&>*]:row-start-1",
@@ -639,7 +627,6 @@ const ChatMessageInput = memo(
 					/>
 					<EnterKeyPlugin onEnter={disabled ? undefined : onEnter} />
 					<ContentChangePlugin onChange={handleContentChange} />
-					<ValueSyncPlugin initialValue={initialValue} />
 					<InsertTextPlugin onEditorReady={handleEditorReady} />
 					<EditableStatePlugin disabled={!!disabled} />
 					{autoFocus && <AutoFocusPlugin />}
@@ -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<typeof vi.fn>;
+	clear: ReturnType<typeof vi.fn>;
+	focus: ReturnType<typeof vi.fn>;
+	getValue: ReturnType<typeof vi.fn>;
+	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<ChatMessageInputRef>();
-		const inputValueRef: import("react").RefObject<string> = { 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();
-	});
 });
@@ -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<void>;
 	onDeleteQueuedMessage: (id: number) => Promise<void>;
 	chatInputRef: React.RefObject<ChatMessageInputRef | null>;
-	inputValueRef: React.RefObject<string>;
 }) {
-	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<string>).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<number | null>(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<ChatMessageInputRef | null>(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}
@@ -58,7 +58,6 @@ const buildEditing = (
 	overrides: Partial<ComponentProps<typeof AgentChatPageView>["editing"]> = {},
 ) => ({
 	chatInputRef: { current: null },
-	editorInitialValue: "",
 	editingMessageId: null as number | null,
 	editingFileBlocks: [] as readonly ChatMessagePart[],
 	handleEditUserMessage: fn(),
@@ -112,6 +111,7 @@ const StoryAgentChatPageView: FC<StoryProps> = ({ 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
 		/>
@@ -31,7 +31,6 @@ type ChatStoreHandle = ReturnType<typeof useChatStore>["store"];

 interface EditingState {
 	chatInputRef: RefObject<ChatMessageInputRef | null>;
-	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<AgentChatPageViewProps> = ({
 	store,
 	editing,
 	pendingEditMessageId,
+	initialInputValue,
 	effectiveSelectedModel,
 	setSelectedModel,
 	modelOptions,
@@ -335,7 +338,7 @@ export const AgentChatPageView: FC<AgentChatPageViewProps> = ({
 							modelSelectorPlaceholder={modelSelectorPlaceholder}
 							isModelCatalogLoading={isModelCatalogLoading}
 							inputRef={editing.chatInputRef}
-							initialValue={editing.editorInitialValue}
+							initialInputValue={initialInputValue}
 							onContentChange={editing.handleContentChange}
 							editingQueuedMessageID={editing.editingQueuedMessageID}
 							onStartQueueEdit={editing.handleStartQueueEdit}
@@ -75,7 +75,7 @@ interface AgentChatInputProps {
 	isLoading: boolean;
 	// Ref for the Lexical editor, exposed for imperative access.
 	inputRef?: React.Ref<ChatMessageInputRef>;
-	// 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;
@@ -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<ChatMessageInputRef>;
-	initialValue?: string;
+	initialInputValue?: string;
 	onContentChange?: (content: string) => void;
 	editingQueuedMessageID: number | null;
 	onStartQueueEdit: (
@@ -170,7 +170,7 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
 	modelSelectorPlaceholder,
 	isModelCatalogLoading = false,
 	inputRef,
-	initialValue,
+	initialInputValue,
 	onContentChange,
 	editingQueuedMessageID,
 	onStartQueueEdit,
@@ -315,7 +315,7 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
 			previewUrls={previewUrls}
 			textContents={textContents}
 			inputRef={inputRef}
-			initialValue={initialValue}
+			initialValue={initialInputValue}
 			onContentChange={onContentChange}
 			queuedMessages={queuedMessages}
 			onDeleteQueuedMessage={onDeleteQueuedMessage}