diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 53102d9b1d..33b8c6c2e9 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -136,27 +136,29 @@ const ChatViewComponent: React.ForwardRefRenderFunction(null) const [sendingDisabled, setSendingDisabled] = useState(false) const [selectedImages, setSelectedImages] = useState([]) + const selectedImagesRef = useRef(selectedImages) + const draftHydratedRef = useRef(false) + + type SaveDraftFn = (text: string, images: string[]) => void + type SaveDraftDebouncedFn = debounce.DebouncedFunction // Debounced draft save to prevent excessive writes - const saveDraftDebounced = useMemo( - () => - debounce((text: string, images: string[]) => { - vscode.postMessage({ type: "saveDraftMessage", text, images }) - }, 500), - [], - ) + const saveDraftDebounced: SaveDraftDebouncedFn = useMemo(() => { + return debounce((text: string, images: string[]) => { + vscode.postMessage({ type: "saveDraftMessage", text, images }) + }, 500) + }, []) // Cleanup debounce on unmount useEffect(() => { return () => { - if (saveDraftDebounced && typeof (saveDraftDebounced as any).clear === "function") { - ;(saveDraftDebounced as any).clear() - } + saveDraftDebounced.clear() } }, [saveDraftDebounced]) // Request saved draft on mount useEffect(() => { + draftHydratedRef.current = false vscode.postMessage({ type: "getDraftMessage" }) }, []) @@ -212,6 +214,10 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + selectedImagesRef.current = selectedImages + }, [selectedImages]) + useEffect(() => { isMountedRef.current = true return () => { @@ -594,6 +600,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + if (!draftHydratedRef.current) { + return + } + saveDraftDebounced(inputValueRef.current, selectedImages) + }, [selectedImages, saveDraftDebounced]) + /** * Handles sending messages to the extension * @param text - The message text to send @@ -876,15 +890,16 @@ const ChatViewComponent: React.ForwardRefRenderFunction 0) { + if (message.images !== undefined) { setSelectedImages(message.images) } } + draftHydratedRef.current = true break } // textAreaRef.current is not explicitly required here since React diff --git a/webview-ui/src/components/chat/__tests__/ChatView.draft-message.spec.tsx b/webview-ui/src/components/chat/__tests__/ChatView.draft-message.spec.tsx new file mode 100644 index 0000000000..2de7efa729 --- /dev/null +++ b/webview-ui/src/components/chat/__tests__/ChatView.draft-message.spec.tsx @@ -0,0 +1,268 @@ +// npx vitest run src/components/chat/__tests__/ChatView.draft-message.spec.tsx + +import React from "react" +import { act, render, waitFor } from "@/utils/test-utils" +import { QueryClient, QueryClientProvider } from "@tanstack/react-query" + +import { ExtensionStateContextProvider } from "@src/context/ExtensionStateContext" +import { vscode } from "@src/utils/vscode" + +import ChatView, { type ChatViewProps } from "../ChatView" + +vi.mock("@src/utils/vscode", () => ({ + vscode: { + postMessage: vi.fn(), + }, +})) + +vi.mock("use-sound", () => ({ + default: vi.fn().mockImplementation(() => { + return [vi.fn()] + }), +})) + +vi.mock("../BrowserSessionRow", () => ({ + default: () => null, +})) + +vi.mock("../ChatRow", () => ({ + default: () => null, +})) + +vi.mock("../AutoApproveMenu", () => ({ + default: () => null, +})) + +vi.mock("../common/TelemetryBanner", () => ({ + default: () => null, +})) + +vi.mock("../Announcement", () => ({ + default: () => null, +})) + +vi.mock("../TaskHeader", () => ({ + default: () => null, +})) + +vi.mock("../SystemPromptWarning", () => ({ + default: () => null, +})) + +vi.mock("../ProfileViolationWarning", () => ({ + default: () => null, +})) + +vi.mock("../CheckpointWarning", () => ({ + CheckpointWarning: () => null, +})) + +vi.mock("../QueuedMessages", () => ({ + QueuedMessages: () => null, +})) + +vi.mock("../history/HistoryPreview", () => ({ + default: () => null, +})) + +vi.mock("@src/components/welcome/RooHero", () => ({ + default: () => null, +})) + +vi.mock("@src/components/welcome/RooTips", () => ({ + default: () => null, +})) + +vi.mock("@src/hooks/useCloudUpsell", () => ({ + useCloudUpsell: () => ({ + isOpen: false, + openUpsell: vi.fn(), + closeUpsell: vi.fn(), + handleConnect: vi.fn(), + }), +})) + +vi.mock("@src/components/cloud/CloudUpsellDialog", () => ({ + CloudUpsellDialog: () => null, +})) + +vi.mock("react-i18next", () => ({ + useTranslation: () => ({ + t: (key: string) => key, + }), + initReactI18next: { + type: "3rdParty", + init: () => {}, + }, + Trans: ({ i18nKey }: { i18nKey: string }) => <>{i18nKey}, +})) + +type MockChatTextAreaProps = { + inputValue?: string + selectedImages?: string[] + setInputValue: (value: string) => void + setSelectedImages: React.Dispatch> +} + +let lastChatTextAreaProps: MockChatTextAreaProps | undefined + +vi.mock("../ChatTextArea", () => { + // eslint-disable-next-line @typescript-eslint/no-require-imports + const mockReact = require("react") + + const ChatTextAreaComponent = mockReact.forwardRef(function MockChatTextArea( + props: MockChatTextAreaProps, + ref: React.ForwardedRef<{ focus: () => void }>, + ) { + lastChatTextAreaProps = props + mockReact.useImperativeHandle(ref, () => ({ + focus: vi.fn(), + })) + return
+ }) + + return { + default: ChatTextAreaComponent, + ChatTextArea: ChatTextAreaComponent, + } +}) + +const defaultProps: ChatViewProps = { + isHidden: false, + showAnnouncement: false, + hideAnnouncement: () => {}, +} + +const queryClient = new QueryClient() + +const renderChatView = (props: Partial = {}) => { + return render( + + + + + , + ) +} + +const hydrateState = (state: Record) => { + window.dispatchEvent( + new MessageEvent("message", { + data: { + type: "state", + state: { + version: "1.0.0", + clineMessages: [], + taskHistory: [], + shouldShowAnnouncement: false, + cloudIsAuthenticated: false, + telemetrySetting: "enabled", + mode: "code", + customModes: [], + messageQueue: [], + organizationAllowList: { allowAll: true, providers: {} }, + apiConfiguration: { apiProvider: "anthropic" }, + ...state, + }, + }, + }), + ) +} + +const sendExtensionMessage = (data: unknown) => { + window.dispatchEvent(new MessageEvent("message", { data })) +} + +describe("ChatView - draftMessage", () => { + beforeEach(() => { + vi.clearAllMocks() + lastChatTextAreaProps = undefined + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it("requests a saved draft on mount", async () => { + renderChatView() + hydrateState({}) + await waitFor(() => { + expect(vscode.postMessage).toHaveBeenCalledWith({ type: "getDraftMessage" }) + }) + }) + + it("restores draft into an empty composer", async () => { + renderChatView() + hydrateState({}) + await act(async () => { + sendExtensionMessage({ + type: "draftMessage", + text: "hello", + images: ["data:image/png;base64,abc"], + }) + }) + await waitFor(() => { + expect(lastChatTextAreaProps?.inputValue).toBe("hello") + expect(lastChatTextAreaProps?.selectedImages).toEqual(["data:image/png;base64,abc"]) + }) + }) + + it("does not overwrite non-empty composer when draftMessage arrives", async () => { + renderChatView() + hydrateState({}) + await waitFor(() => { + expect(lastChatTextAreaProps).toBeDefined() + }) + await act(async () => { + lastChatTextAreaProps?.setInputValue("keep") + }) + await waitFor(() => { + expect(lastChatTextAreaProps?.inputValue).toBe("keep") + }) + await act(async () => { + sendExtensionMessage({ + type: "draftMessage", + text: "overwrite?", + images: ["data:image/png;base64,zzz"], + }) + }) + await waitFor(() => { + expect(lastChatTextAreaProps?.inputValue).toBe("keep") + }) + }) + + it("saves draft when images change after hydration", async () => { + vi.useFakeTimers() + renderChatView() + hydrateState({}) + + await act(async () => { + await Promise.resolve() + }) + expect(vscode.postMessage).toHaveBeenCalledWith({ type: "getDraftMessage" }) + + // Mark hydration complete + await act(async () => { + sendExtensionMessage({ + type: "draftMessage", + text: "", + images: [], + }) + }) + expect(lastChatTextAreaProps).toBeDefined() + + // Change images only + vi.clearAllMocks() + await act(async () => { + lastChatTextAreaProps?.setSelectedImages(["data:image/png;base64,img1"]) + }) + await act(async () => { + vi.advanceTimersByTime(600) + }) + expect(vscode.postMessage).toHaveBeenCalledWith({ + type: "saveDraftMessage", + text: "", + images: ["data:image/png;base64,img1"], + }) + }) +})