diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 1fb272604d3..754c88d9081 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -145,6 +145,7 @@ export class ClineProvider private recentTasksCache?: string[] private pendingOperations: Map = new Map() private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds + private static readonly DRAFT_MESSAGE_KEY = "roo.draftMessage" private cloudOrganizationsCache: CloudOrganizationMembership[] | null = null private cloudOrganizationsCacheTimestamp: number | null = null @@ -550,6 +551,43 @@ export class ClineProvider this.log(`[clearAllPendingEditOperations] Cleared all pending operations`) } + // Draft Message Management + + /** + * Save draft message to workspace state + * @param text - The draft message text + * @param images - Array of base64 image data URLs + */ + public async saveDraftMessage(text: string, images: string[]): Promise { + // Don't save empty drafts + if (!text.trim() && images.length === 0) { + await this.clearDraftMessage() + return + } + + const draft = { + text, + images, + timestamp: Date.now(), + } + await this.context.workspaceState.update(ClineProvider.DRAFT_MESSAGE_KEY, draft) + } + + /** + * Get draft message from workspace state + * @returns The saved draft or undefined + */ + public getDraftMessage(): { text: string; images: string[]; timestamp: number } | undefined { + return this.context.workspaceState.get(ClineProvider.DRAFT_MESSAGE_KEY) + } + + /** + * Clear draft message from workspace state + */ + public async clearDraftMessage(): Promise { + await this.context.workspaceState.update(ClineProvider.DRAFT_MESSAGE_KEY, undefined) + } + /* VSCode extensions use the disposable pattern to clean up resources when the sidebar/editor tab is closed by the user or system. This applies to event listening, commands, interacting with the UI, etc. - https://vscode-docs.readthedocs.io/en/stable/extensions/patterns-and-principles/ diff --git a/src/core/webview/__tests__/webviewMessageHandler.draftMessage.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.draftMessage.spec.ts new file mode 100644 index 00000000000..bd1c8b61680 --- /dev/null +++ b/src/core/webview/__tests__/webviewMessageHandler.draftMessage.spec.ts @@ -0,0 +1,150 @@ +// npx vitest run src/core/webview/__tests__/webviewMessageHandler.draftMessage.spec.ts + +import { webviewMessageHandler } from "../webviewMessageHandler" +import type { ClineProvider } from "../ClineProvider" + +vi.mock("vscode", () => ({ + window: { + showInformationMessage: vi.fn(), + showErrorMessage: vi.fn(), + }, + workspace: { + workspaceFolders: [{ uri: { fsPath: "/mock/workspace" } }], + }, +})) + +describe("webviewMessageHandler - Draft Message", () => { + let mockClineProvider: ClineProvider + + beforeEach(() => { + vi.clearAllMocks() + + mockClineProvider = { + saveDraftMessage: vi.fn().mockResolvedValue(undefined), + getDraftMessage: vi.fn(), + clearDraftMessage: vi.fn().mockResolvedValue(undefined), + postMessageToWebview: vi.fn(), + getState: vi.fn().mockResolvedValue({ + apiConfiguration: {}, + }), + contextProxy: { + context: { + extensionPath: "/mock/extension/path", + globalStorageUri: { fsPath: "/mock/global/storage" }, + }, + setValue: vi.fn(), + getValue: vi.fn(), + }, + log: vi.fn(), + } as unknown as ClineProvider + }) + + describe("saveDraftMessage", () => { + it("should save draft message with text and images", async () => { + await webviewMessageHandler(mockClineProvider, { + type: "saveDraftMessage", + text: "Test draft message", + images: ["data:image/png;base64,abc123"], + }) + + expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("Test draft message", [ + "data:image/png;base64,abc123", + ]) + }) + + it("should save draft message with only text", async () => { + await webviewMessageHandler(mockClineProvider, { + type: "saveDraftMessage", + text: "Text only draft", + }) + + expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("Text only draft", []) + }) + + it("should save draft message with only images", async () => { + await webviewMessageHandler(mockClineProvider, { + type: "saveDraftMessage", + images: ["data:image/png;base64,image1", "data:image/png;base64,image2"], + }) + + expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("", [ + "data:image/png;base64,image1", + "data:image/png;base64,image2", + ]) + }) + + it("should handle empty text and images", async () => { + await webviewMessageHandler(mockClineProvider, { + type: "saveDraftMessage", + }) + + expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("", []) + }) + }) + + describe("getDraftMessage", () => { + it("should return saved draft with text and images", async () => { + const mockDraft = { + text: "Saved draft", + images: ["data:image/png;base64,savedImage"], + timestamp: 1234567890, + } + ;(mockClineProvider.getDraftMessage as ReturnType).mockReturnValue(mockDraft) + + await webviewMessageHandler(mockClineProvider, { + type: "getDraftMessage", + }) + + expect(mockClineProvider.getDraftMessage).toHaveBeenCalled() + expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({ + type: "draftMessage", + text: "Saved draft", + images: ["data:image/png;base64,savedImage"], + }) + }) + + it("should handle when no draft is saved", async () => { + ;(mockClineProvider.getDraftMessage as ReturnType).mockReturnValue(undefined) + + await webviewMessageHandler(mockClineProvider, { + type: "getDraftMessage", + }) + + expect(mockClineProvider.getDraftMessage).toHaveBeenCalled() + expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({ + type: "draftMessage", + text: undefined, + images: undefined, + }) + }) + + it("should handle draft with only text", async () => { + const mockDraft = { + text: "Text only", + images: [], + timestamp: 1234567890, + } + ;(mockClineProvider.getDraftMessage as ReturnType).mockReturnValue(mockDraft) + + await webviewMessageHandler(mockClineProvider, { + type: "getDraftMessage", + }) + + expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({ + type: "draftMessage", + text: "Text only", + images: [], + }) + }) + }) + + describe("clearDraftMessage", () => { + it("should clear the draft message", async () => { + await webviewMessageHandler(mockClineProvider, { + type: "clearDraftMessage", + }) + + expect(mockClineProvider.clearDraftMessage).toHaveBeenCalled() + }) + }) +}) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 21a515b6107..c9a1ea957d5 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -3100,6 +3100,28 @@ export const webviewMessageHandler = async ( break } + /** + * Draft Message Persistence + */ + + case "saveDraftMessage": + await provider.saveDraftMessage(message.text || "", message.images || []) + break + + case "getDraftMessage": { + const draft = provider.getDraftMessage() + await provider.postMessageToWebview({ + type: "draftMessage", + text: draft?.text, + images: draft?.images, + }) + break + } + + case "clearDraftMessage": + await provider.clearDraftMessage() + break + default: { // console.log(`Unhandled message type: ${message.type}`) // diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index b9424f0bf21..47da08b2726 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -131,6 +131,7 @@ export interface ExtensionMessage { | "interactionRequired" | "browserSessionUpdate" | "browserSessionNavigate" + | "draftMessage" text?: string payload?: any // Add a generic payload for now, can refine later // Checkpoint warning message diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index eeae8e70cb0..fcf0de63d87 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -175,6 +175,9 @@ export interface WebviewMessage { | "browserPanelDidLaunch" | "openDebugApiHistory" | "openDebugUiHistory" + | "saveDraftMessage" + | "getDraftMessage" + | "clearDraftMessage" text?: string editedMessageContent?: string tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "cloud" diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 243e583613b..53102d9b1d8 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -137,6 +137,29 @@ const ChatViewComponent: React.ForwardRefRenderFunction([]) + // Debounced draft save to prevent excessive writes + const saveDraftDebounced = useMemo( + () => + 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]) + + // Request saved draft on mount + useEffect(() => { + vscode.postMessage({ type: "getDraftMessage" }) + }, []) + // We need to hold on to the ask because useEffect > lastMessage will always // let us know when an ask comes in and handle it, but by the time // handleMessage is called, the last message might not be the ask anymore @@ -557,8 +580,20 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + setInputValue(value) + saveDraftDebounced(value, selectedImages) + }, + [selectedImages, saveDraftDebounced], + ) + /** * Handles sending messages to the extension * @param text - The message text to send @@ -840,6 +875,17 @@ const ChatViewComponent: React.ForwardRefRenderFunction 0) { + setSelectedImages(message.images) + } + } + break } // textAreaRef.current is not explicitly required here since React // guarantees that ref will be stable across re-renders, and we're @@ -1568,7 +1614,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction