Skip to content

Commit 9b0ae78

Browse files
qdaxbqdaxb
committed
use ExtensionContext instead of LocalStorage
1 parent 48f46a1 commit 9b0ae78

File tree

7 files changed

+163
-219
lines changed

7 files changed

+163
-219
lines changed

‎packages/types/src/global-settings.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ export const globalSettingsSchema = z.object({
109109
hasOpenedModeSelector: z.boolean().optional(),
110110
lastModeExportPath: z.string().optional(),
111111
lastModeImportPath: z.string().optional(),
112+
chatTextDrafts: z.record(z.string(), z.string()).optional(),
112113
})
113114

114115
export type GlobalSettings = z.infer<typeof globalSettingsSchema>

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,40 @@ export const webviewMessageHandler = async (
5858
await provider.contextProxy.setValue(key, value)
5959

6060
switch (message.type) {
61+
case "updateChatTextDraft": {
62+
// Update or remove the chat text draft in globalState under the fixed key "chatTextDraft"
63+
const text = typeof message.text === "string" ? message.text : ""
64+
const drafts = (await getGlobalState("chatTextDrafts")) ?? {}
65+
if (text && text.trim()) {
66+
// Set/Update
67+
const updated = { ...drafts, chatTextDraft: text }
68+
await updateGlobalState("chatTextDrafts", updated)
69+
} else if (drafts.chatTextDraft) {
70+
// Remove if empty
71+
const { chatTextDraft: _, ...rest } = drafts
72+
await updateGlobalState("chatTextDrafts", rest)
73+
}
74+
break
75+
}
76+
case "getChatTextDraft": {
77+
// Return the chat text draft for the fixed key "chatTextDraft"
78+
const drafts = (await getGlobalState("chatTextDrafts")) ?? {}
79+
const text = drafts.chatTextDraft ?? ""
80+
await provider.postMessageToWebview({
81+
type: "chatTextDraftValue",
82+
text,
83+
})
84+
break
85+
}
86+
case "clearChatTextDraft": {
87+
// Remove the chat text draft from globalState under the fixed key "chatTextDraft"
88+
const drafts = (await getGlobalState("chatTextDrafts")) ?? {}
89+
if (drafts.chatTextDraft) {
90+
const { chatTextDraft: _, ...rest } = drafts
91+
await updateGlobalState("chatTextDrafts", rest)
92+
}
93+
break
94+
}
6195
case "webviewDidLaunch":
6296
// Load custom modes first
6397
const customModes = await provider.customModesManager.getCustomModes()

‎src/shared/ExtensionMessage.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ export interface ExtensionMessage {
105105
| "shareTaskSuccess"
106106
| "codeIndexSettingsSaved"
107107
| "codeIndexSecretStatus"
108+
| "chatTextDraftValue"
108109
text?: string
109110
payload?: any // Add a generic payload for now, can refine later
110111
action?:

‎src/shared/WebviewMessage.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,9 @@ export interface WebviewMessage {
185185
| "checkRulesDirectoryResult"
186186
| "saveCodeIndexSettingsAtomic"
187187
| "requestCodeIndexSecretStatus"
188+
| "getChatTextDraft"
189+
| "updateChatTextDraft"
190+
| "clearChatTextDraft"
188191
text?: string
189192
tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account"
190193
disabled?: boolean

‎webview-ui/src/components/chat/ChatTextArea.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
7272
const { t } = useAppTranslation()
7373

7474
// Chat draft persistence
75-
const { handleSendAndClearDraft } = useChatTextDraft("chat_textarea_draft", inputValue, setInputValue, onSend)
75+
const { handleSendAndClearDraft } = useChatTextDraft(inputValue, setInputValue, onSend)
7676

7777
const {
7878
filePaths,

‎webview-ui/src/components/chat/hooks/__tests__/useChatTextDraft.spec.ts‎

Lines changed: 80 additions & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -3,230 +3,143 @@
33
import { renderHook, act } from "@testing-library/react"
44
import { useChatTextDraft } from "../useChatTextDraft"
55
import { vi } from "vitest"
6+
import { vscode } from "@src/utils/vscode"
67

7-
describe("useChatTextDraft", () => {
8-
const draftKey = "test-draft-key"
8+
describe("useChatTextDraft (postMessage version)", () => {
99
let setInputValue: (v: string) => void
1010
let onSend: () => void
11-
12-
let getItemMock: ReturnType<typeof vi.fn>
13-
let setItemMock: ReturnType<typeof vi.fn>
14-
let removeItemMock: ReturnType<typeof vi.fn>
11+
let postMessageMock: ReturnType<typeof vi.fn>
12+
let addEventListenerMock: ReturnType<typeof vi.fn>
13+
let removeEventListenerMock: ReturnType<typeof vi.fn>
14+
let eventListener: ((event: MessageEvent) => void) | undefined
1515

1616
beforeEach(() => {
1717
setInputValue = vi.fn((_: string) => {})
1818
onSend = vi.fn()
19+
postMessageMock = vi.fn()
20+
addEventListenerMock = vi.fn((type, cb) => {
21+
if (type === "message") eventListener = cb
22+
})
23+
removeEventListenerMock = vi.fn((type, cb) => {
24+
if (type === "message" && eventListener === cb) eventListener = undefined
25+
})
1926

20-
getItemMock = vi.fn()
21-
setItemMock = vi.fn()
22-
removeItemMock = vi.fn()
27+
global.window.addEventListener = addEventListenerMock
28+
global.window.removeEventListener = removeEventListenerMock
29+
// mock vscode.postMessage
30+
vi.resetModules()
31+
vi.clearAllMocks()
32+
vscode.postMessage = postMessageMock
2333

24-
// @ts-expect-error override readonly
25-
global.localStorage = {
26-
getItem: getItemMock,
27-
setItem: setItemMock,
28-
removeItem: removeItemMock,
29-
}
3034
vi.useFakeTimers()
3135
})
3236

3337
afterEach(() => {
3438
vi.clearAllTimers()
3539
vi.useRealTimers()
3640
vi.restoreAllMocks()
41+
eventListener = undefined
3742
})
3843

39-
describe("Draft restoration on mount", () => {
40-
it("should restore draft from localStorage on mount if inputValue is empty", () => {
41-
getItemMock.mockReturnValue("restored draft")
42-
renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend))
43-
expect(getItemMock).toHaveBeenCalledWith(draftKey)
44-
expect(setInputValue).toHaveBeenCalledWith("restored draft")
45-
})
46-
47-
it("should not restore draft if inputValue is not empty", () => {
48-
getItemMock.mockReturnValue("restored draft")
49-
renderHook(() => useChatTextDraft(draftKey, "already typed", setInputValue, onSend))
50-
expect(getItemMock).toHaveBeenCalledWith(draftKey)
51-
expect(setInputValue).not.toHaveBeenCalled()
52-
})
53-
54-
it("should ignore errors from localStorage.getItem", () => {
55-
getItemMock.mockImplementation(() => {
56-
throw new Error("getItem error")
57-
})
58-
expect(() => renderHook(() => useChatTextDraft(draftKey, "", setInputValue, onSend))).not.toThrow()
59-
expect(setInputValue).not.toHaveBeenCalled()
44+
it("should send getChatTextDraft on mount and set input value when chatTextDraftValue received", () => {
45+
renderHook(() => useChatTextDraft("", setInputValue, onSend))
46+
expect(postMessageMock).toHaveBeenCalledWith({ type: "getChatTextDraft" })
47+
expect(setInputValue).not.toHaveBeenCalled()
48+
// Simulate extension host response
49+
act(() => {
50+
eventListener?.({ data: { type: "chatTextDraftValue", text: "restored draft" } } as MessageEvent)
6051
})
52+
expect(setInputValue).toHaveBeenCalledWith("restored draft")
6153
})
6254

63-
describe("Auto-save functionality with debounce", () => {
64-
it("should auto-save draft to localStorage after 3 seconds of inactivity if inputValue is non-empty", () => {
65-
renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
66-
initialProps: { value: "hello world" },
67-
})
68-
expect(setItemMock).not.toHaveBeenCalled()
69-
act(() => {
70-
vi.advanceTimersByTime(2999)
71-
})
72-
expect(setItemMock).not.toHaveBeenCalled()
73-
act(() => {
74-
vi.advanceTimersByTime(1)
75-
})
76-
expect(setItemMock).toHaveBeenCalledWith(draftKey, "hello world")
77-
})
78-
79-
it("should reset debounce timer when inputValue changes before debounce delay", () => {
80-
const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
81-
initialProps: { value: "foo" },
82-
})
83-
act(() => {
84-
vi.advanceTimersByTime(2000)
85-
})
86-
// Should not save before debounce delay
87-
expect(setItemMock).not.toHaveBeenCalled()
88-
rerender({ value: "bar" })
89-
act(() => {
90-
vi.advanceTimersByTime(2999)
91-
})
92-
expect(setItemMock).not.toHaveBeenCalled()
93-
act(() => {
94-
vi.advanceTimersByTime(1)
95-
})
96-
expect(setItemMock).toHaveBeenCalledWith(draftKey, "bar")
55+
it("should not set input value if inputValue is not empty when chatTextDraftValue received", () => {
56+
renderHook(() => useChatTextDraft("already typed", setInputValue, onSend))
57+
act(() => {
58+
eventListener?.({ data: { type: "chatTextDraftValue", text: "restored draft" } } as MessageEvent)
9759
})
60+
expect(setInputValue).not.toHaveBeenCalled()
61+
})
9862

99-
it("should remove draft from localStorage if inputValue is empty", () => {
100-
renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
101-
initialProps: { value: "" },
102-
})
103-
expect(removeItemMock).toHaveBeenCalledWith(draftKey)
63+
it("should debounce and send updateChatTextDraft with text after 2s if inputValue is non-empty", () => {
64+
renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), {
65+
initialProps: { value: "hello world" },
10466
})
105-
106-
it("should ignore errors from localStorage.setItem", () => {
107-
setItemMock.mockImplementation(() => {
108-
throw new Error("setItem error")
109-
})
110-
renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
111-
initialProps: { value: "err" },
112-
})
113-
act(() => {
114-
vi.advanceTimersByTime(5000)
115-
})
116-
expect(setItemMock).toHaveBeenCalled()
67+
expect(postMessageMock).toHaveBeenCalledWith({ type: "getChatTextDraft" })
68+
postMessageMock.mockClear()
69+
act(() => {
70+
vi.advanceTimersByTime(1999)
11771
})
118-
119-
it("should ignore errors from localStorage.removeItem", () => {
120-
removeItemMock.mockImplementation(() => {
121-
throw new Error("removeItem error")
122-
})
123-
renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
124-
initialProps: { value: "" },
125-
})
126-
expect(removeItemMock).toHaveBeenCalledWith(draftKey)
72+
expect(postMessageMock).not.toHaveBeenCalled()
73+
act(() => {
74+
vi.advanceTimersByTime(1)
12775
})
76+
expect(postMessageMock).toHaveBeenCalledWith({ type: "updateChatTextDraft", text: "hello world" })
12877
})
12978

130-
describe("Draft clearing on send", () => {
131-
it("should remove draft and call onSend when handleSendAndClearDraft is called", () => {
132-
const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend))
133-
act(() => {
134-
result.current.handleSendAndClearDraft()
135-
})
136-
expect(removeItemMock).toHaveBeenCalledWith(draftKey)
137-
expect(onSend).toHaveBeenCalled()
138-
})
139-
140-
it("should ignore errors from localStorage.removeItem on send", () => {
141-
removeItemMock.mockImplementation(() => {
142-
throw new Error("removeItem error")
143-
})
144-
const { result } = renderHook(() => useChatTextDraft(draftKey, "msg", setInputValue, onSend))
145-
act(() => {
146-
expect(() => result.current.handleSendAndClearDraft()).not.toThrow()
147-
})
148-
expect(onSend).toHaveBeenCalled()
79+
it("should reset debounce timer when inputValue changes before debounce delay", () => {
80+
const { rerender } = renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), {
81+
initialProps: { value: "foo" },
14982
})
150-
})
151-
152-
/**
153-
* @description
154-
* Complex scenario: multiple inputValue changes, ensure debounce timer cleanup and localStorage operations have no side effects.
155-
*/
156-
it("should handle rapid inputValue changes and cleanup debounce timers", () => {
157-
const { rerender } = renderHook(({ value }) => useChatTextDraft(draftKey, value, setInputValue, onSend), {
158-
initialProps: { value: "first" },
83+
act(() => {
84+
vi.advanceTimersByTime(1000)
15985
})
86+
postMessageMock.mockClear()
87+
rerender({ value: "bar" })
16088
act(() => {
161-
vi.advanceTimersByTime(2999)
89+
vi.advanceTimersByTime(1999)
16290
})
163-
expect(setItemMock).not.toHaveBeenCalled()
91+
expect(postMessageMock).not.toHaveBeenCalled()
16492
act(() => {
16593
vi.advanceTimersByTime(1)
16694
})
167-
expect(setItemMock).toHaveBeenCalledWith(draftKey, "first")
168-
rerender({ value: "second" })
169-
act(() => {
170-
vi.advanceTimersByTime(2999)
95+
expect(postMessageMock).toHaveBeenCalledWith({ type: "updateChatTextDraft", text: "bar" })
96+
})
97+
98+
it("should send clearChatTextDraft if inputValue is empty after user has input", () => {
99+
const { rerender } = renderHook(({ value }) => useChatTextDraft(value, setInputValue, onSend), {
100+
initialProps: { value: "foo" },
171101
})
172-
expect(setItemMock).toHaveBeenCalledTimes(1)
173102
act(() => {
174-
vi.advanceTimersByTime(1)
103+
vi.advanceTimersByTime(2000)
175104
})
176-
expect(setItemMock).toHaveBeenCalledWith(draftKey, "second")
105+
postMessageMock.mockClear()
177106
rerender({ value: "" })
178-
expect(removeItemMock).toHaveBeenCalledWith(draftKey)
107+
expect(postMessageMock).toHaveBeenCalledWith({ type: "clearChatTextDraft" })
179108
})
180-
it("should not save and should warn if inputValue exceeds 100KB (ASCII)", () => {
181-
const draftKey = "large-draft-key"
109+
110+
it("should send clearChatTextDraft and call onSend when handleSendAndClearDraft is called", () => {
111+
const { result } = renderHook(() => useChatTextDraft("msg", setInputValue, onSend))
112+
postMessageMock.mockClear()
113+
act(() => {
114+
result.current.handleSendAndClearDraft()
115+
})
116+
expect(postMessageMock).toHaveBeenCalledWith({ type: "clearChatTextDraft" })
117+
expect(onSend).toHaveBeenCalled()
118+
})
119+
120+
it("should not send updateChatTextDraft and should warn if inputValue exceeds 100KB (ASCII)", () => {
182121
const MAX_DRAFT_BYTES = 102400
183-
// Generate a string larger than 100KB (1 byte per char, simple ASCII)
184122
const largeStr = "a".repeat(MAX_DRAFT_BYTES + 5000)
185-
const setItemMock = vi.fn()
186-
const getItemMock = vi.fn()
187-
const removeItemMock = vi.fn()
188-
// @ts-expect-error override readonly
189-
global.localStorage = {
190-
getItem: getItemMock,
191-
setItem: setItemMock,
192-
removeItem: removeItemMock,
193-
}
194123
const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {})
195-
const setInputValue = vi.fn()
196-
const onSend = vi.fn()
197-
renderHook(() => useChatTextDraft(draftKey, largeStr, setInputValue, onSend))
124+
renderHook(() => useChatTextDraft(largeStr, setInputValue, onSend))
198125
act(() => {
199126
vi.advanceTimersByTime(3000)
200127
})
201-
expect(setItemMock).not.toHaveBeenCalled()
128+
expect(postMessageMock).not.toHaveBeenCalledWith({ type: "updateChatTextDraft", text: largeStr })
202129
expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB"))
203130
warnMock.mockRestore()
204131
})
205132

206-
it("should not save and should warn if inputValue exceeds 100KB (UTF-8 multi-byte)", () => {
207-
const draftKey = "utf8-draft-key"
208-
// Each emoji is 4 bytes in UTF-8, 汉字 is 3 bytes
133+
it("should not send updateChatTextDraft and should warn if inputValue exceeds 100KB (UTF-8 multi-byte)", () => {
209134
const emoji = "😀"
210135
const hanzi = "汉"
211-
// Compose a string: 20,000 emojis (~80KB) + 10,000 汉 (~30KB) + some ASCII
212136
const utf8Str = emoji.repeat(20000) + hanzi.repeat(10000) + "abc"
213-
const setItemMock = vi.fn()
214-
const getItemMock = vi.fn()
215-
const removeItemMock = vi.fn()
216-
// @ts-expect-error override readonly
217-
global.localStorage = {
218-
getItem: getItemMock,
219-
setItem: setItemMock,
220-
removeItem: removeItemMock,
221-
}
222137
const warnMock = vi.spyOn(console, "warn").mockImplementation(() => {})
223-
const setInputValue = vi.fn()
224-
const onSend = vi.fn()
225-
renderHook(() => useChatTextDraft(draftKey, utf8Str, setInputValue, onSend))
138+
renderHook(() => useChatTextDraft(utf8Str, setInputValue, onSend))
226139
act(() => {
227140
vi.advanceTimersByTime(3000)
228141
})
229-
expect(setItemMock).not.toHaveBeenCalled()
142+
expect(postMessageMock).not.toHaveBeenCalledWith({ type: "updateChatTextDraft", text: utf8Str })
230143
expect(warnMock).toHaveBeenCalledWith(expect.stringContaining("exceeds 100KB"))
231144
warnMock.mockRestore()
232145
})

0 commit comments

Comments
 (0)