From 82090a2ab46b5256a42d35f7ddd6688243d5cc10 Mon Sep 17 00:00:00 2001 From: axb Date: Tue, 1 Jul 2025 19:10:44 +0800 Subject: [PATCH 1/3] support chat draft --- .../src/components/chat/ChatTextArea.tsx | 29 ++++---- .../components/chat/hooks/useChatTextDraft.ts | 72 +++++++++++++++++++ 2 files changed, 89 insertions(+), 12 deletions(-) create mode 100644 webview-ui/src/components/chat/hooks/useChatTextDraft.ts diff --git a/webview-ui/src/components/chat/ChatTextArea.tsx b/webview-ui/src/components/chat/ChatTextArea.tsx index 51279062d2..e8aec7b124 100644 --- a/webview-ui/src/components/chat/ChatTextArea.tsx +++ b/webview-ui/src/components/chat/ChatTextArea.tsx @@ -9,6 +9,7 @@ import { ExtensionMessage } from "@roo/ExtensionMessage" import { vscode } from "@/utils/vscode" import { useExtensionState } from "@/context/ExtensionStateContext" +import { useChatTextDraft } from "./hooks/useChatTextDraft" import { useAppTranslation } from "@/i18n/TranslationContext" import { ContextMenuOptionType, @@ -69,6 +70,10 @@ const ChatTextArea = forwardRef( ref, ) => { const { t } = useAppTranslation() + + // Chat draft persistence + const { handleSendAndClearDraft } = useChatTextDraft("chat_textarea_draft", inputValue, setInputValue, onSend) + const { filePaths, openedTabs, @@ -389,7 +394,7 @@ const ChatTextArea = forwardRef( if (!sendingDisabled) { // Reset history navigation state when sending resetHistoryNavigation() - onSend() + handleSendAndClearDraft() } } @@ -438,22 +443,22 @@ const ChatTextArea = forwardRef( } }, [ - sendingDisabled, - onSend, showContextMenu, - searchQuery, + handleHistoryNavigation, selectedMenuIndex, - handleMentionSelect, - selectedType, + searchQuery, inputValue, - cursorPosition, - setInputValue, - justDeletedSpaceAfterMention, + selectedType, queryItems, - allModes, fileSearchResults, - handleHistoryNavigation, + allModes, + handleMentionSelect, + sendingDisabled, resetHistoryNavigation, + handleSendAndClearDraft, + cursorPosition, + justDeletedSpaceAfterMention, + setInputValue, ], ) @@ -1151,7 +1156,7 @@ const ChatTextArea = forwardRef( iconClass="codicon-send" title={t("chat:sendMessage")} disabled={sendingDisabled} - onClick={onSend} + onClick={handleSendAndClearDraft} /> diff --git a/webview-ui/src/components/chat/hooks/useChatTextDraft.ts b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts new file mode 100644 index 0000000000..3c113bc94d --- /dev/null +++ b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts @@ -0,0 +1,72 @@ +import { useCallback, useEffect, useRef } from "react" + +/** + * Hook for chat textarea draft persistence (localStorage). + * Handles auto-save, restore on mount, and clear on send. + * @param draftKey localStorage key for draft persistence + * @param inputValue current textarea value + * @param setInputValue setter for textarea value + * @param onSend send callback + */ +export function useChatTextDraft( + draftKey: string, + inputValue: string, + setInputValue: (value: string) => void, + onSend: () => void, +) { + const saveDraftTimerRef = useRef(null) + + // Restore draft on mount + useEffect(() => { + try { + const draft = localStorage.getItem(draftKey) + if (draft && !inputValue) { + setInputValue(draft) + } + } catch (_) { + // ignore + } + // Only run on initial mount + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []) + + // Periodically save draft + useEffect(() => { + if (saveDraftTimerRef.current) { + clearInterval(saveDraftTimerRef.current) + } + if (inputValue && inputValue.trim()) { + saveDraftTimerRef.current = setInterval(() => { + try { + localStorage.setItem(draftKey, inputValue) + } catch (_) { + // ignore + } + }, 5000) + } else { + // Remove draft if no content + try { + localStorage.removeItem(draftKey) + } catch (_) { + // ignore + } + } + return () => { + if (saveDraftTimerRef.current) { + clearInterval(saveDraftTimerRef.current) + } + } + }, [inputValue, draftKey]) + + // Clear draft after send + const handleSendAndClearDraft = useCallback(() => { + try { + localStorage.removeItem(draftKey) + } catch (_) { + // ignore + } + onSend() + }, [onSend, draftKey]) + + return { handleSendAndClearDraft } +} From 48f46a19759a5e76f9030dd6d4b90cd0cc100cc6 Mon Sep 17 00:00:00 2001 From: axb Date: Wed, 2 Jul 2025 00:15:18 +0800 Subject: [PATCH 2/3] use debounced save instead of interval --- .../hooks/__tests__/useChatTextDraft.spec.ts | 233 ++++++++++++++++++ .../components/chat/hooks/useChatTextDraft.ts | 62 +++-- 2 files changed, 277 insertions(+), 18 deletions(-) create mode 100644 webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts diff --git a/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts b/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts new file mode 100644 index 0000000000..6e1df73a41 --- /dev/null +++ b/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts @@ -0,0 +1,233 @@ +// npx vitest webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts + +import { renderHook, act } from "@testing-library/react" +import { useChatTextDraft } from "../useChatTextDraft" +import { vi } from "vitest" + +describe("useChatTextDraft", () => { + const draftKey = "test-draft-key" + let setInputValue: (v: string) => void + let onSend: () => void + + let getItemMock: ReturnType + let setItemMock: ReturnType + let removeItemMock: ReturnType + + beforeEach(() => { + setInputValue = vi.fn((_: string) => {}) + onSend = vi.fn() + + getItemMock = vi.fn() + setItemMock = vi.fn() + removeItemMock = vi.fn() + + // @ts-expect-error override readonly + global.localStorage = { + getItem: getItemMock, + setItem: setItemMock, + removeItem: removeItemMock, + } + vi.useFakeTimers() + }) + + afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + vi.restoreAllMocks() + }) + + describe("Draft restoration on mount", () => { + it("should restore draft from localStorage on mount if inputValue is empty", () => { + getItemMock.mockReturnValue("restored draft") + renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend)) + expect(getItemMock).toHaveBeenCalledWith(draftKey) + expect(setInputValue).toHaveBeenCalledWith("restored draft") + }) + + it("should not restore draft if inputValue is not empty", () => { + getItemMock.mockReturnValue("restored draft") + renderHook(() => useChatTextDraft(draftKey, "already typed", setInputValue, onSend)) + expect(getItemMock).toHaveBeenCalledWith(draftKey) + expect(setInputValue).not.toHaveBeenCalled() + }) + + it("should ignore errors from localStorage.getItem", () => { + getItemMock.mockImplementation(() => { + throw new Error("getItem error") + }) + expect(() => renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend))).not.toThrow() + expect(setInputValue).not.toHaveBeenCalled() + }) + }) + + describe("Auto-save functionality with debounce", () => { + it("should auto-save draft to localStorage after 3 seconds of inactivity if inputValue is non-empty", () => { + renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "hello world" }, + }) + expect(setItemMock).not.toHaveBeenCalled() + act(() => { + vi.advanceTimersByTime(2999) + }) + expect(setItemMock).not.toHaveBeenCalled() + act(() => { + vi.advanceTimersByTime(1) + }) + expect(setItemMock).toHaveBeenCalledWith(draftKey, "hello world") + }) + + it("should reset debounce timer when inputValue changes before debounce delay", () => { + const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "foo" }, + }) + act(() => { + vi.advanceTimersByTime(2000) + }) + // Should not save before debounce delay + expect(setItemMock).not.toHaveBeenCalled() + rerender({ value: "bar" }) + act(() => { + vi.advanceTimersByTime(2999) + }) + expect(setItemMock).not.toHaveBeenCalled() + act(() => { + vi.advanceTimersByTime(1) + }) + expect(setItemMock).toHaveBeenCalledWith(draftKey, "bar") + }) + + it("should remove draft from localStorage if inputValue is empty", () => { + renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "" }, + }) + expect(removeItemMock).toHaveBeenCalledWith(draftKey) + }) + + it("should ignore errors from localStorage.setItem", () => { + setItemMock.mockImplementation(() => { + throw new Error("setItem error") + }) + renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "err" }, + }) + act(() => { + vi.advanceTimersByTime(5000) + }) + expect(setItemMock).toHaveBeenCalled() + }) + + it("should ignore errors from localStorage.removeItem", () => { + removeItemMock.mockImplementation(() => { + throw new Error("removeItem error") + }) + renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "" }, + }) + expect(removeItemMock).toHaveBeenCalledWith(draftKey) + }) + }) + + describe("Draft clearing on send", () => { + it("should remove draft and call onSend when handleSendAndClearDraft is called", () => { + const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend)) + act(() => { + result.current.handleSendAndClearDraft() + }) + expect(removeItemMock).toHaveBeenCalledWith(draftKey) + expect(onSend).toHaveBeenCalled() + }) + + it("should ignore errors from localStorage.removeItem on send", () => { + removeItemMock.mockImplementation(() => { + throw new Error("removeItem error") + }) + const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend)) + act(() => { + expect(() => result.current.handleSendAndClearDraft()).not.toThrow() + }) + expect(onSend).toHaveBeenCalled() + }) + }) + + /** + * @description + * Complex scenario: multiple inputValue changes, ensure debounce timer cleanup and localStorage operations have no side effects. + */ + it("should handle rapid inputValue changes and cleanup debounce timers", () => { + const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { + initialProps: { value: "first" }, + }) + act(() => { + vi.advanceTimersByTime(2999) + }) + expect(setItemMock).not.toHaveBeenCalled() + act(() => { + vi.advanceTimersByTime(1) + }) + expect(setItemMock).toHaveBeenCalledWith(draftKey, "first") + rerender({ value: "second" }) + act(() => { + vi.advanceTimersByTime(2999) + }) + expect(setItemMock).toHaveBeenCalledTimes(1) + act(() => { + vi.advanceTimersByTime(1) + }) + expect(setItemMock).toHaveBeenCalledWith(draftKey, "second") + rerender({ value: "" }) + expect(removeItemMock).toHaveBeenCalledWith(draftKey) + }) + it("should not save and should warn if inputValue exceeds 100KB (ASCII)", () => { + const draftKey = "large-draft-key" + const MAX_DRAFT_BYTES = 102400 + // Generate a string larger than 100KB (1 byte per char, simple ASCII) + const largeStr = "a".repeat(MAX_DRAFT_BYTES + 5000) + const setItemMock = vi.fn() + const getItemMock = vi.fn() + const removeItemMock = vi.fn() + // @ts-expect-error override readonly + global.localStorage = { + getItem: getItemMock, + setItem: setItemMock, + removeItem: removeItemMock, + } + const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {}) + const setInputValue = vi.fn() + const onSend = vi.fn() + renderHook(() => useChatTextDraft(draftKey, largeStr, setInputValue, onSend)) + act(() => { + vi.advanceTimersByTime(3000) + }) + expect(setItemMock).not.toHaveBeenCalled() + expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB")) + warnMock.mockRestore() + }) + + it("should not save and should warn if inputValue exceeds 100KB (UTF-8 multi-byte)", () => { + const draftKey = "utf8-draft-key" + // Each emoji is 4 bytes in UTF-8, 汉字 is 3 bytes + const emoji = "😀" + const hanzi = "汉" + // Compose a string: 20,000 emojis (~80KB) + 10,000 汉 (~30KB) + some ASCII + const utf8Str = emoji.repeat(20000) + hanzi.repeat(10000) + "abc" + const setItemMock = vi.fn() + const getItemMock = vi.fn() + const removeItemMock = vi.fn() + // @ts-expect-error override readonly + global.localStorage = { + getItem: getItemMock, + setItem: setItemMock, + removeItem: removeItemMock, + } + const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {}) + const setInputValue = vi.fn() + const onSend = vi.fn() + renderHook(() => useChatTextDraft(draftKey, utf8Str, setInputValue, onSend)) + act(() => { + vi.advanceTimersByTime(3000) + }) + expect(setItemMock).not.toHaveBeenCalled() + expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB")) + warnMock.mockRestore() + }) +}) diff --git a/webview-ui/src/components/chat/hooks/useChatTextDraft.ts b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts index 3c113bc94d..f699d36db2 100644 --- a/webview-ui/src/components/chat/hooks/useChatTextDraft.ts +++ b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts @@ -1,5 +1,8 @@ import { useCallback, useEffect, useRef } from "react" +// Debounce delay for saving chat draft (ms) +export const CHAT_DRAFT_SAVE_DEBOUNCE_MS = 3000 + /** * Hook for chat textarea draft persistence (localStorage). * Handles auto-save, restore on mount, and clear on send. @@ -14,8 +17,6 @@ export function useChatTextDraft( setInputValue: (value: string) => void, onSend: () => void, ) { - const saveDraftTimerRef = useRef(null) - // Restore draft on mount useEffect(() => { try { @@ -23,37 +24,61 @@ export function useChatTextDraft( if (draft && !inputValue) { setInputValue(draft) } - } catch (_) { - // ignore + } catch (err) { + // Log localStorage getItem failure for debugging + console.warn(`[useChatTextDraft] Failed to restore draft from localStorage for key "${draftKey}":`, err) } // Only run on initial mount // eslint-disable-next-line react-hooks/exhaustive-deps }, []) - // Periodically save draft + // Debounced save draft + const debounceTimerRef = useRef(null) + useEffect(() => { - if (saveDraftTimerRef.current) { - clearInterval(saveDraftTimerRef.current) + if (debounceTimerRef.current) { + clearTimeout(debounceTimerRef.current) } if (inputValue && inputValue.trim()) { - saveDraftTimerRef.current = setInterval(() => { + debounceTimerRef.current = setTimeout(() => { try { - localStorage.setItem(draftKey, inputValue) - } catch (_) { - // ignore + // Limit draft size to 100KB (102400 bytes) + const MAX_DRAFT_BYTES = 102400 + // Fast pre-check: if character count is much greater than max bytes, skip encoding + if (inputValue.length > MAX_DRAFT_BYTES * 2) { + // In UTF-8, Chinese chars are usually 2-3 bytes, English 1 byte + console.warn( + `[useChatTextDraft] Draft for key "${draftKey}" is too long (chars=${inputValue.length}), not saving to localStorage.`, + ) + } else { + const encoder = new TextEncoder() + const bytes = encoder.encode(inputValue) + if (bytes.length > MAX_DRAFT_BYTES) { + // Do not save if draft exceeds 100KB + console.warn( + `[useChatTextDraft] Draft for key "${draftKey}" exceeds 100KB, not saving to localStorage.`, + ) + } else { + localStorage.setItem(draftKey, inputValue) + } + } + } catch (err) { + // Log localStorage setItem failure for debugging + console.warn(`[useChatTextDraft] Failed to save draft to localStorage for key "${draftKey}":`, err) } - }, 5000) + }, CHAT_DRAFT_SAVE_DEBOUNCE_MS) } else { // Remove draft if no content try { localStorage.removeItem(draftKey) - } catch (_) { - // ignore + } catch (err) { + // Log localStorage removeItem failure for debugging + console.warn(`[useChatTextDraft] Failed to remove draft from localStorage for key "${draftKey}":`, err) } } return () => { - if (saveDraftTimerRef.current) { - clearInterval(saveDraftTimerRef.current) + if (debounceTimerRef.current) { + clearTimeout(debounceTimerRef.current) } } }, [inputValue, draftKey]) @@ -62,8 +87,9 @@ export function useChatTextDraft( const handleSendAndClearDraft = useCallback(() => { try { localStorage.removeItem(draftKey) - } catch (_) { - // ignore + } catch (err) { + // Log localStorage removeItem failure for debugging + console.warn(`[useChatTextDraft] Failed to remove draft from localStorage for key "${draftKey}":`, err) } onSend() }, [onSend, draftKey]) From 9b0ae787b49f14c3fc1748599ae5fc869e0e9436 Mon Sep 17 00:00:00 2001 From: axb Date: Wed, 2 Jul 2025 12:17:10 +0800 Subject: [PATCH 3/3] use ExtensionContext instead of LocalStorage --- packages/types/src/global-settings.ts | 1 + src/core/webview/webviewMessageHandler.ts | 34 +++ src/shared/ExtensionMessage.ts | 1 + src/shared/WebviewMessage.ts | 3 + .../src/components/chat/ChatTextArea.tsx | 2 +- .../hooks/__tests__/useChatTextDraft.spec.ts | 247 ++++++------------ .../components/chat/hooks/useChatTextDraft.ts | 94 +++---- 7 files changed, 163 insertions(+), 219 deletions(-) diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index 6a0dda1cf2..a820f56c87 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -109,6 +109,7 @@ export const globalSettingsSchema = z.object({ hasOpenedModeSelector: z.boolean().optional(), lastModeExportPath: z.string().optional(), lastModeImportPath: z.string().optional(), + chatTextDrafts: z.record(z.string(), z.string()).optional(), }) export type GlobalSettings = z.infer diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 0ec14ca27e..e4148cc4d5 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -58,6 +58,40 @@ export const webviewMessageHandler = async ( await provider.contextProxy.setValue(key, value) switch (message.type) { + case "updateChatTextDraft": { + // Update or remove the chat text draft in globalState under the fixed key "chatTextDraft" + const text = typeof message.text === "string" ? message.text : "" + const drafts = (await getGlobalState("chatTextDrafts")) ?? {} + if (text && text.trim()) { + // Set/Update + const updated = { ...drafts, chatTextDraft: text } + await updateGlobalState("chatTextDrafts", updated) + } else if (drafts.chatTextDraft) { + // Remove if empty + const { chatTextDraft: _, ...rest } = drafts + await updateGlobalState("chatTextDrafts", rest) + } + break + } + case "getChatTextDraft": { + // Return the chat text draft for the fixed key "chatTextDraft" + const drafts = (await getGlobalState("chatTextDrafts")) ?? {} + const text = drafts.chatTextDraft ?? "" + await provider.postMessageToWebview({ + type: "chatTextDraftValue", + text, + }) + break + } + case "clearChatTextDraft": { + // Remove the chat text draft from globalState under the fixed key "chatTextDraft" + const drafts = (await getGlobalState("chatTextDrafts")) ?? {} + if (drafts.chatTextDraft) { + const { chatTextDraft: _, ...rest } = drafts + await updateGlobalState("chatTextDrafts", rest) + } + break + } case "webviewDidLaunch": // Load custom modes first const customModes = await provider.customModesManager.getCustomModes() diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index 9db0889c88..66043cc4c4 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -105,6 +105,7 @@ export interface ExtensionMessage { | "shareTaskSuccess" | "codeIndexSettingsSaved" | "codeIndexSecretStatus" + | "chatTextDraftValue" text?: string payload?: any // Add a generic payload for now, can refine later action?: diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index a50e30b67e..679ccdca63 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -185,6 +185,9 @@ export interface WebviewMessage { | "checkRulesDirectoryResult" | "saveCodeIndexSettingsAtomic" | "requestCodeIndexSecretStatus" + | "getChatTextDraft" + | "updateChatTextDraft" + | "clearChatTextDraft" text?: string tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account" disabled?: boolean diff --git a/webview-ui/src/components/chat/ChatTextArea.tsx b/webview-ui/src/components/chat/ChatTextArea.tsx index e8aec7b124..71b3682871 100644 --- a/webview-ui/src/components/chat/ChatTextArea.tsx +++ b/webview-ui/src/components/chat/ChatTextArea.tsx @@ -72,7 +72,7 @@ const ChatTextArea = forwardRef( const { t } = useAppTranslation() // Chat draft persistence - const { handleSendAndClearDraft } = useChatTextDraft("chat_textarea_draft", inputValue, setInputValue, onSend) + const { handleSendAndClearDraft } = useChatTextDraft(inputValue, setInputValue, onSend) const { filePaths, diff --git a/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts b/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts index 6e1df73a41..e705a18141 100644 --- a/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts +++ b/webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts @@ -3,30 +3,34 @@ import { renderHook, act } from "@testing-library/react" import { useChatTextDraft } from "../useChatTextDraft" import { vi } from "vitest" +import { vscode } from "@src/utils/vscode" -describe("useChatTextDraft", () => { - const draftKey = "test-draft-key" +describe("useChatTextDraft (postMessage version)", () => { let setInputValue: (v: string) => void let onSend: () => void - - let getItemMock: ReturnType - let setItemMock: ReturnType - let removeItemMock: ReturnType + let postMessageMock: ReturnType + let addEventListenerMock: ReturnType + let removeEventListenerMock: ReturnType + let eventListener: ((event: MessageEvent) => void) | undefined beforeEach(() => { setInputValue = vi.fn((_: string) => {}) onSend = vi.fn() + postMessageMock = vi.fn() + addEventListenerMock = vi.fn((type, cb) => { + if (type === "message") eventListener = cb + }) + removeEventListenerMock = vi.fn((type, cb) => { + if (type === "message" && eventListener === cb) eventListener = undefined + }) - getItemMock = vi.fn() - setItemMock = vi.fn() - removeItemMock = vi.fn() + global.window.addEventListener = addEventListenerMock + global.window.removeEventListener = removeEventListenerMock + // mock vscode.postMessage + vi.resetModules() + vi.clearAllMocks() + vscode.postMessage = postMessageMock - // @ts-expect-error override readonly - global.localStorage = { - getItem: getItemMock, - setItem: setItemMock, - removeItem: removeItemMock, - } vi.useFakeTimers() }) @@ -34,199 +38,108 @@ describe("useChatTextDraft", () => { vi.clearAllTimers() vi.useRealTimers() vi.restoreAllMocks() + eventListener = undefined }) - describe("Draft restoration on mount", () => { - it("should restore draft from localStorage on mount if inputValue is empty", () => { - getItemMock.mockReturnValue("restored draft") - renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend)) - expect(getItemMock).toHaveBeenCalledWith(draftKey) - expect(setInputValue).toHaveBeenCalledWith("restored draft") - }) - - it("should not restore draft if inputValue is not empty", () => { - getItemMock.mockReturnValue("restored draft") - renderHook(() => useChatTextDraft(draftKey, "already typed", setInputValue, onSend)) - expect(getItemMock).toHaveBeenCalledWith(draftKey) - expect(setInputValue).not.toHaveBeenCalled() - }) - - it("should ignore errors from localStorage.getItem", () => { - getItemMock.mockImplementation(() => { - throw new Error("getItem error") - }) - expect(() => renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend))).not.toThrow() - expect(setInputValue).not.toHaveBeenCalled() + it("should send getChatTextDraft on mount and set input value when chatTextDraftValue received", () => { + renderHook(() => useChatTextDraft("", setInputValue, onSend)) + expect(postMessageMock).toHaveBeenCalledWith({ type: "getChatTextDraft" }) + expect(setInputValue).not.toHaveBeenCalled() + // Simulate extension host response + act(() => { + eventListener?.({ data: { type: "chatTextDraftValue", text: "restored draft" } } as MessageEvent) }) + expect(setInputValue).toHaveBeenCalledWith("restored draft") }) - describe("Auto-save functionality with debounce", () => { - it("should auto-save draft to localStorage after 3 seconds of inactivity if inputValue is non-empty", () => { - renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "hello world" }, - }) - expect(setItemMock).not.toHaveBeenCalled() - act(() => { - vi.advanceTimersByTime(2999) - }) - expect(setItemMock).not.toHaveBeenCalled() - act(() => { - vi.advanceTimersByTime(1) - }) - expect(setItemMock).toHaveBeenCalledWith(draftKey, "hello world") - }) - - it("should reset debounce timer when inputValue changes before debounce delay", () => { - const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "foo" }, - }) - act(() => { - vi.advanceTimersByTime(2000) - }) - // Should not save before debounce delay - expect(setItemMock).not.toHaveBeenCalled() - rerender({ value: "bar" }) - act(() => { - vi.advanceTimersByTime(2999) - }) - expect(setItemMock).not.toHaveBeenCalled() - act(() => { - vi.advanceTimersByTime(1) - }) - expect(setItemMock).toHaveBeenCalledWith(draftKey, "bar") + it("should not set input value if inputValue is not empty when chatTextDraftValue received", () => { + renderHook(() => useChatTextDraft("already typed", setInputValue, onSend)) + act(() => { + eventListener?.({ data: { type: "chatTextDraftValue", text: "restored draft" } } as MessageEvent) }) + expect(setInputValue).not.toHaveBeenCalled() + }) - it("should remove draft from localStorage if inputValue is empty", () => { - renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "" }, - }) - expect(removeItemMock).toHaveBeenCalledWith(draftKey) + it("should debounce and send updateChatTextDraft with text after 2s if inputValue is non-empty", () => { + renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), { + initialProps: { value: "hello world" }, }) - - it("should ignore errors from localStorage.setItem", () => { - setItemMock.mockImplementation(() => { - throw new Error("setItem error") - }) - renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "err" }, - }) - act(() => { - vi.advanceTimersByTime(5000) - }) - expect(setItemMock).toHaveBeenCalled() + expect(postMessageMock).toHaveBeenCalledWith({ type: "getChatTextDraft" }) + postMessageMock.mockClear() + act(() => { + vi.advanceTimersByTime(1999) }) - - it("should ignore errors from localStorage.removeItem", () => { - removeItemMock.mockImplementation(() => { - throw new Error("removeItem error") - }) - renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "" }, - }) - expect(removeItemMock).toHaveBeenCalledWith(draftKey) + expect(postMessageMock).not.toHaveBeenCalled() + act(() => { + vi.advanceTimersByTime(1) }) + expect(postMessageMock).toHaveBeenCalledWith({ type: "updateChatTextDraft", text: "hello world" }) }) - describe("Draft clearing on send", () => { - it("should remove draft and call onSend when handleSendAndClearDraft is called", () => { - const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend)) - act(() => { - result.current.handleSendAndClearDraft() - }) - expect(removeItemMock).toHaveBeenCalledWith(draftKey) - expect(onSend).toHaveBeenCalled() - }) - - it("should ignore errors from localStorage.removeItem on send", () => { - removeItemMock.mockImplementation(() => { - throw new Error("removeItem error") - }) - const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend)) - act(() => { - expect(() => result.current.handleSendAndClearDraft()).not.toThrow() - }) - expect(onSend).toHaveBeenCalled() + it("should reset debounce timer when inputValue changes before debounce delay", () => { + const { rerender } = renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), { + initialProps: { value: "foo" }, }) - }) - - /** - * @description - * Complex scenario: multiple inputValue changes, ensure debounce timer cleanup and localStorage operations have no side effects. - */ - it("should handle rapid inputValue changes and cleanup debounce timers", () => { - const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), { - initialProps: { value: "first" }, + act(() => { + vi.advanceTimersByTime(1000) }) + postMessageMock.mockClear() + rerender({ value: "bar" }) act(() => { - vi.advanceTimersByTime(2999) + vi.advanceTimersByTime(1999) }) - expect(setItemMock).not.toHaveBeenCalled() + expect(postMessageMock).not.toHaveBeenCalled() act(() => { vi.advanceTimersByTime(1) }) - expect(setItemMock).toHaveBeenCalledWith(draftKey, "first") - rerender({ value: "second" }) - act(() => { - vi.advanceTimersByTime(2999) + expect(postMessageMock).toHaveBeenCalledWith({ type: "updateChatTextDraft", text: "bar" }) + }) + + it("should send clearChatTextDraft if inputValue is empty after user has input", () => { + const { rerender } = renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), { + initialProps: { value: "foo" }, }) - expect(setItemMock).toHaveBeenCalledTimes(1) act(() => { - vi.advanceTimersByTime(1) + vi.advanceTimersByTime(2000) }) - expect(setItemMock).toHaveBeenCalledWith(draftKey, "second") + postMessageMock.mockClear() rerender({ value: "" }) - expect(removeItemMock).toHaveBeenCalledWith(draftKey) + expect(postMessageMock).toHaveBeenCalledWith({ type: "clearChatTextDraft" }) }) - it("should not save and should warn if inputValue exceeds 100KB (ASCII)", () => { - const draftKey = "large-draft-key" + + it("should send clearChatTextDraft and call onSend when handleSendAndClearDraft is called", () => { + const { result } = renderHook(() => useChatTextDraft("msg", setInputValue, onSend)) + postMessageMock.mockClear() + act(() => { + result.current.handleSendAndClearDraft() + }) + expect(postMessageMock).toHaveBeenCalledWith({ type: "clearChatTextDraft" }) + expect(onSend).toHaveBeenCalled() + }) + + it("should not send updateChatTextDraft and should warn if inputValue exceeds 100KB (ASCII)", () => { const MAX_DRAFT_BYTES = 102400 - // Generate a string larger than 100KB (1 byte per char, simple ASCII) const largeStr = "a".repeat(MAX_DRAFT_BYTES + 5000) - const setItemMock = vi.fn() - const getItemMock = vi.fn() - const removeItemMock = vi.fn() - // @ts-expect-error override readonly - global.localStorage = { - getItem: getItemMock, - setItem: setItemMock, - removeItem: removeItemMock, - } const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {}) - const setInputValue = vi.fn() - const onSend = vi.fn() - renderHook(() => useChatTextDraft(draftKey, largeStr, setInputValue, onSend)) + renderHook(() => useChatTextDraft(largeStr, setInputValue, onSend)) act(() => { vi.advanceTimersByTime(3000) }) - expect(setItemMock).not.toHaveBeenCalled() + expect(postMessageMock).not.toHaveBeenCalledWith({ type: "updateChatTextDraft", text: largeStr }) expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB")) warnMock.mockRestore() }) - it("should not save and should warn if inputValue exceeds 100KB (UTF-8 multi-byte)", () => { - const draftKey = "utf8-draft-key" - // Each emoji is 4 bytes in UTF-8, 汉字 is 3 bytes + it("should not send updateChatTextDraft and should warn if inputValue exceeds 100KB (UTF-8 multi-byte)", () => { const emoji = "😀" const hanzi = "汉" - // Compose a string: 20,000 emojis (~80KB) + 10,000 汉 (~30KB) + some ASCII const utf8Str = emoji.repeat(20000) + hanzi.repeat(10000) + "abc" - const setItemMock = vi.fn() - const getItemMock = vi.fn() - const removeItemMock = vi.fn() - // @ts-expect-error override readonly - global.localStorage = { - getItem: getItemMock, - setItem: setItemMock, - removeItem: removeItemMock, - } const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {}) - const setInputValue = vi.fn() - const onSend = vi.fn() - renderHook(() => useChatTextDraft(draftKey, utf8Str, setInputValue, onSend)) + renderHook(() => useChatTextDraft(utf8Str, setInputValue, onSend)) act(() => { vi.advanceTimersByTime(3000) }) - expect(setItemMock).not.toHaveBeenCalled() + expect(postMessageMock).not.toHaveBeenCalledWith({ type: "updateChatTextDraft", text: utf8Str }) expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB")) warnMock.mockRestore() }) diff --git a/webview-ui/src/components/chat/hooks/useChatTextDraft.ts b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts index f699d36db2..465a1b077c 100644 --- a/webview-ui/src/components/chat/hooks/useChatTextDraft.ts +++ b/webview-ui/src/components/chat/hooks/useChatTextDraft.ts @@ -1,79 +1,72 @@ import { useCallback, useEffect, useRef } from "react" +import { vscode } from "@src/utils/vscode" -// Debounce delay for saving chat draft (ms) -export const CHAT_DRAFT_SAVE_DEBOUNCE_MS = 3000 +export const CHAT_DRAFT_SAVE_DEBOUNCE_MS = 2000 /** - * Hook for chat textarea draft persistence (localStorage). - * Handles auto-save, restore on mount, and clear on send. - * @param draftKey localStorage key for draft persistence + * Hook for chat textarea draft persistence (extension globalState). + * Handles auto-save, restore on mount, and clear on send via postMessage. * @param inputValue current textarea value * @param setInputValue setter for textarea value * @param onSend send callback */ -export function useChatTextDraft( - draftKey: string, - inputValue: string, - setInputValue: (value: string) => void, - onSend: () => void, -) { - // Restore draft on mount +export function useChatTextDraft(inputValue: string, setInputValue: (value: string) => void, onSend: () => void) { + // Restore draft from extension host on mount useEffect(() => { - try { - const draft = localStorage.getItem(draftKey) - if (draft && !inputValue) { - setInputValue(draft) + const handleDraftValue = (event: MessageEvent) => { + const msg = event.data + if (msg && msg.type === "chatTextDraftValue") { + if (typeof msg.text === "string" && msg.text && !inputValue) { + setInputValue(msg.text) + } } - } catch (err) { - // Log localStorage getItem failure for debugging - console.warn(`[useChatTextDraft] Failed to restore draft from localStorage for key "${draftKey}":`, err) } - // Only run on initial mount + window.addEventListener("message", handleDraftValue) + // Request draft from extension host + vscode.postMessage({ type: "getChatTextDraft" }) + return () => { + window.removeEventListener("message", handleDraftValue) + } // eslint-disable-next-line react-hooks/exhaustive-deps }, []) - // Debounced save draft + // Debounced save draft to extension host const debounceTimerRef = useRef(null) + const hasHadUserInput = useRef(false) + useEffect(() => { if (debounceTimerRef.current) { clearTimeout(debounceTimerRef.current) } + const MAX_DRAFT_BYTES = 102400 if (inputValue && inputValue.trim()) { + hasHadUserInput.current = true debounceTimerRef.current = setTimeout(() => { try { - // Limit draft size to 100KB (102400 bytes) - const MAX_DRAFT_BYTES = 102400 // Fast pre-check: if character count is much greater than max bytes, skip encoding if (inputValue.length > MAX_DRAFT_BYTES * 2) { - // In UTF-8, Chinese chars are usually 2-3 bytes, English 1 byte - console.warn( - `[useChatTextDraft] Draft for key "${draftKey}" is too long (chars=${inputValue.length}), not saving to localStorage.`, - ) - } else { - const encoder = new TextEncoder() - const bytes = encoder.encode(inputValue) - if (bytes.length > MAX_DRAFT_BYTES) { - // Do not save if draft exceeds 100KB - console.warn( - `[useChatTextDraft] Draft for key "${draftKey}" exceeds 100KB, not saving to localStorage.`, - ) - } else { - localStorage.setItem(draftKey, inputValue) - } + console.warn(`[useChatTextDraft] Draft is too long (chars=${inputValue.length}), not saving.`) + return + } + const encoder = new TextEncoder() + const bytes = encoder.encode(inputValue) + if (bytes.length > MAX_DRAFT_BYTES) { + console.warn(`[useChatTextDraft] Draft exceeds 100KB, not saving.`) + return } + vscode.postMessage({ type: "updateChatTextDraft", text: inputValue }) } catch (err) { - // Log localStorage setItem failure for debugging - console.warn(`[useChatTextDraft] Failed to save draft to localStorage for key "${draftKey}":`, err) + console.warn(`[useChatTextDraft] Failed to save draft:`, err) } }, CHAT_DRAFT_SAVE_DEBOUNCE_MS) } else { - // Remove draft if no content - try { - localStorage.removeItem(draftKey) - } catch (err) { - // Log localStorage removeItem failure for debugging - console.warn(`[useChatTextDraft] Failed to remove draft from localStorage for key "${draftKey}":`, err) + if (hasHadUserInput.current) { + try { + vscode.postMessage({ type: "clearChatTextDraft" }) + } catch (err) { + console.warn(`[useChatTextDraft] Failed to clear draft:`, err) + } } } return () => { @@ -81,18 +74,17 @@ export function useChatTextDraft( clearTimeout(debounceTimerRef.current) } } - }, [inputValue, draftKey]) + }, [inputValue]) // Clear draft after send const handleSendAndClearDraft = useCallback(() => { try { - localStorage.removeItem(draftKey) + vscode.postMessage({ type: "clearChatTextDraft" }) } catch (err) { - // Log localStorage removeItem failure for debugging - console.warn(`[useChatTextDraft] Failed to remove draft from localStorage for key "${draftKey}":`, err) + console.warn(`[useChatTextDraft] Failed to clear draft:`, err) } onSend() - }, [onSend, draftKey]) + }, [onSend]) return { handleSendAndClearDraft } }