diff --git a/webview-ui/src/__tests__/App.spec.tsx b/webview-ui/src/__tests__/App.spec.tsx index 4b4cbcd5fe..391e28e05e 100644 --- a/webview-ui/src/__tests__/App.spec.tsx +++ b/webview-ui/src/__tests__/App.spec.tsx @@ -6,6 +6,8 @@ import { render, screen, act, cleanup } from "@/utils/test-utils" import AppWithProviders from "../App" import { vscode } from "../utils/vscode" +let nextChatViewInstance = 0 + vi.mock("@src/utils/vscode", () => ({ vscode: { postMessage: vi.fn(), @@ -29,8 +31,9 @@ vi.mock("@src/utils/TelemetryClient", () => ({ vi.mock("@src/components/chat/ChatView", () => ({ __esModule: true, default: function ChatView({ isHidden }: { isHidden: boolean }) { + const [instance] = React.useState(() => ++nextChatViewInstance) return ( -
+
Chat View
) @@ -161,6 +164,7 @@ vi.mock("process.env", () => ({ describe("App", () => { beforeEach(() => { vi.clearAllMocks() + nextChatViewInstance = 0 window.removeEventListener("message", () => {}) // Set up default mock return value @@ -206,6 +210,35 @@ describe("App", () => { expect(chatView.getAttribute("data-hidden")).toBe("false") }, 10000) + it("preserves chat state across task and tab changes", async () => { + const state = { + didHydrateState: true, + showWelcome: false, + shouldShowAnnouncement: false, + experiments: {}, + language: "en", + telemetrySetting: "enabled", + currentTaskId: "task-a" as string | undefined, + } + mockUseExtensionState.mockImplementation(() => state) + + const { rerender } = render() + const firstInstance = screen.getByTestId("chat-view").getAttribute("data-instance") + + act(() => triggerMessage("settingsButtonClicked")) + expect(await screen.findByTestId("settings-view")).toBeInTheDocument() + expect(screen.getByTestId("chat-view").getAttribute("data-instance")).toBe(firstInstance) + + state.currentTaskId = "task-b" + rerender() + const secondInstance = screen.getByTestId("chat-view").getAttribute("data-instance") + expect(secondInstance).toBe(firstInstance) + + state.currentTaskId = undefined + rerender() + expect(screen.getByTestId("chat-view").getAttribute("data-instance")).toBe(secondInstance) + }) + it.each(["sidebar", "editor"])( "reports direct pointer and keyboard interaction in the %s webview", (renderContext) => { diff --git a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx index bbad494170..7fa4e81131 100644 --- a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx +++ b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx @@ -6,6 +6,7 @@ import { mockVscodePostMessage, renderWithExtensionState, waitFor, + screen, act, fireEvent, } from "@/utils/test-utils" @@ -49,6 +50,8 @@ vi.mock("use-sound", () => ({ }), })) +let nextRowInstance = 0 + // Mock components that use ESM dependencies vi.mock("../ChatRow", () => ({ default: function MockChatRow({ @@ -60,6 +63,7 @@ vi.mock("../ChatRow", () => ({ onSuggestionClick?: (suggestion: SuggestionItem, event?: React.MouseEvent) => void isFollowUpAnswered?: boolean }) { + const [instance] = React.useState(() => ++nextRowInstance) if (message.type === "ask" && message.ask === "followup" && message.text) { try { const followUp = JSON.parse(message.text) as { suggest?: SuggestionItem[] } @@ -81,7 +85,11 @@ vi.mock("../ChatRow", () => ({ } } - return
{JSON.stringify(message)}
+ return ( +
+ {JSON.stringify(message)} +
+ ) }, })) @@ -721,6 +729,19 @@ describe("ChatView - Virtualization Configuration", () => { mockVirtuosoState.lastConfig = null }) + it("preserves the row instance when a partial message completes", async () => { + renderChatView() + const task: ClineMessage = { type: "say", say: "task", ts: 1, text: "Task" } + const partial: ClineMessage = { type: "say", say: "text", ts: 2, text: "Streaming", partial: true } + mockPostMessage({ clineMessages: [task, partial] }) + const row = await screen.findByTestId("chat-row") + const instance = row.getAttribute("data-instance") + expect(instance).not.toBeNull() + mockPostMessage({ clineMessages: [task, { ...partial, text: "Completed", partial: false }] }) + await waitFor(() => expect(screen.getByTestId("chat-row")).toHaveTextContent("Completed")) + expect(screen.getByTestId("chat-row").getAttribute("data-instance")).toBe(instance) + }) + it("keeps the off-screen render buffer tight for chat rows", async () => { renderChatView() @@ -751,6 +772,9 @@ describe("ChatView - Virtualization Configuration", () => { expect(mockVirtuosoState.lastConfig?.defaultItemHeight).toBe(180) expect(mockVirtuosoState.lastConfig?.increaseViewportBy).toEqual({ top: 600, bottom: 800 }) expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs })).toBe(`${rowTs}-1`) + expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs, partial: true })).toBe( + `${rowTs}-1`, + ) }) }) diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index c4fa00f479..7c352a2193 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -1,5 +1,5 @@ import { providerIdentifiers } from "@roo-code/types" -import React, { createContext, useCallback, useEffect, useState } from "react" +import React, { createContext, useCallback, useEffect, useRef, useState } from "react" import { type ProviderSettings, @@ -12,6 +12,7 @@ import { type OrganizationAllowList, type CloudOrganizationMembership, type ExtensionMessage, + type ClineMessage, type ExtensionState, type MarketplaceInstalledMetadata, type SkillMetadata, @@ -35,6 +36,7 @@ import { experimentDefault } from "@roo/experiments" import { vscode } from "@src/utils/vscode" import { convertTextMateToHljs } from "@src/utils/textMateToHljs" +import { StringCache } from "@src/utils/stringCache" export interface ExtensionStateContextType extends ExtensionState { historyPreviewCollapsed?: boolean // Add the new state property @@ -185,6 +187,12 @@ export const mergeExtensionState = (prevState: ExtensionState, newState: Partial ) { rest.clineMessages = prevState.clineMessages rest.clineMessagesSeq = prevState.clineMessagesSeq + rest.currentTaskId = prevState.currentTaskId + } else if (newState.clineMessages !== undefined) { + rest.currentTaskId = newState.currentTaskId + } else { + // Pushes without messages cannot establish task identity. + rest.currentTaskId = prevState.currentTaskId } // Note that we completely replace the previous apiConfiguration and customSupportPrompts objects @@ -288,9 +296,23 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode initialState?: ExtensionStateProviderInitialState }> = ({ children, initialState }) => { - const [state, setState] = useState(() => - mergeExtensionState(createInitialExtensionState(), initialState ?? {}), - ) + const [historyCache] = useState(() => new StringCache()) + const [messageCache] = useState(() => new StringCache((value) => !("partial" in value) || value.partial !== true)) + const messagesRef = useRef(initialState?.clineMessages ?? []) + const currentTaskIdRef = useRef(initialState?.currentTaskId) + const clineMessagesSeqRef = useRef(initialState?.clineMessagesSeq) + + const [state, setState] = useState(() => { + const mergedState = mergeExtensionState(createInitialExtensionState(), initialState ?? {}) + for (const clineMessage of mergedState.clineMessages) { + messageCache.intern(clineMessage) + } + historyCache.intern(mergedState.taskHistory) + if (mergedState.currentTaskItem) { + historyCache.intern(mergedState.currentTaskItem) + } + return mergedState + }) const [didHydrateState, setDidHydrateState] = useState(false) const [showWelcome, setShowWelcome] = useState(false) @@ -346,6 +368,53 @@ export const ExtensionStateContextProvider: React.FC<{ switch (message.type) { case "state": { const newState = message.state ?? {} + + const messagesAreStale = + newState.clineMessages !== undefined && + newState.clineMessagesSeq !== undefined && + clineMessagesSeqRef.current !== undefined && + newState.clineMessagesSeq <= clineMessagesSeqRef.current + + if (!messagesAreStale && newState.clineMessages !== undefined) { + if (newState.currentTaskId !== currentTaskIdRef.current) { + messageCache.clear() + messagesRef.current = [] + currentTaskIdRef.current = newState.currentTaskId + } + } + + if (newState.clineMessages && !messagesAreStale) { + const previousMessages = new Map(messagesRef.current.map((item) => [item.ts, item])) + for (const clineMessage of newState.clineMessages) { + const previous = previousMessages.get(clineMessage.ts) + if ( + previous && + Object.keys(clineMessage).length === Object.keys(previous).length && + Object.entries(clineMessage).every( + ([key, value]) => previous[key as keyof ClineMessage] === value, + ) + ) { + Object.assign(clineMessage, previous) + } else { + messageCache.intern(clineMessage) + } + } + messagesRef.current = newState.clineMessages + } + if ( + newState.clineMessagesSeq !== undefined && + (clineMessagesSeqRef.current === undefined || + newState.clineMessagesSeq > clineMessagesSeqRef.current) + ) { + clineMessagesSeqRef.current = newState.clineMessagesSeq + } + if (newState.taskHistory) { + historyCache.clear() + historyCache.intern(newState.taskHistory) + } + if (newState.currentTaskItem) { + historyCache.intern(newState.currentTaskItem) + } setState((prevState) => mergeExtensionState(prevState, newState)) setShowWelcome(!checkExistKey(newState.apiConfiguration, newState.zooCodeIsAuthenticated)) setDidHydrateState(true) @@ -410,6 +479,13 @@ export const ExtensionStateContextProvider: React.FC<{ } case "messageUpdated": { const clineMessage = message.clineMessage! + const messageIndex = messagesRef.current.findIndex((item) => item.ts === clineMessage.ts) + if (messageIndex !== -1) { + messageCache.intern(clineMessage) + messagesRef.current = messagesRef.current.map((item, index) => + index === messageIndex ? clineMessage : item, + ) + } setState((prevState) => { // worth noting it will never be possible for a more up-to-date message to be sent here or in normal messages post since the presentAssistantContent function uses lock const lastIndex = findLastIndex(prevState.clineMessages, (msg) => msg.ts === clineMessage.ts) @@ -474,6 +550,8 @@ export const ExtensionStateContextProvider: React.FC<{ case "taskHistoryUpdated": { // Efficiently update just the task history without replacing entire state if (message.taskHistory !== undefined) { + historyCache.clear() + historyCache.intern(message.taskHistory) setState((prevState) => ({ ...prevState, taskHistory: message.taskHistory!, @@ -486,6 +564,7 @@ export const ExtensionStateContextProvider: React.FC<{ if (!item) { break } + historyCache.intern(item) setState((prevState) => { const existingIndex = prevState.taskHistory.findIndex((h) => h.id === item.id) let nextHistory: typeof prevState.taskHistory @@ -508,7 +587,7 @@ export const ExtensionStateContextProvider: React.FC<{ } } }, - [setListApiConfigMeta], + [historyCache, messageCache, setListApiConfigMeta], ) useEffect(() => { diff --git a/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx b/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx index dc33af38d1..f6b4e2ce3a 100644 --- a/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx +++ b/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx @@ -7,6 +7,7 @@ import { type ExperimentId, type ExtensionState, type ClineMessage, + type ExtensionMessage, type MarketplaceItem, type MarketplaceInstalledMetadata, type RouterModels, @@ -21,6 +22,7 @@ import { mergeExtensionState, createInitialExtensionState, } from "../ExtensionStateContext" +import { StringCache } from "@/utils/stringCache" const TestComponent = () => { const { allowedCommands, setAllowedCommands, soundEnabled, showRooIgnoredFiles, setShowRooIgnoredFiles } = @@ -112,6 +114,296 @@ const InitialStateTestComponent = () => { } describe("ExtensionStateContext", () => { + afterEach(() => vi.restoreAllMocks()) + it.each(["initial", "state", "messageUpdated", "taskHistoryUpdated", "taskHistoryItemUpdated"] as const)( + "interns accepted payloads through %s", + (entry) => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + const payload: ClineMessage = { ts: 1, type: "say", say: "text", text: "accepted" } + const item = { id: "task-a", number: 1, ts: 1, task: "history", tokensIn: 0, tokensOut: 0, totalCost: 0 } + render( + + + , + ) + if (entry === "initial") { + expect(internSpy).toHaveBeenCalledWith(payload) + const cache = internSpy.mock.contexts[0] + expect(cache).toBeInstanceOf(StringCache) + if (!(cache instanceof StringCache)) throw new Error("Missing cache") + expect(typeof cache["shouldIntern"]).toBe("function") + expect(cache.size).toBeGreaterThan(0) + return + } + const caches = internSpy.mock.contexts.filter((cache): cache is StringCache => cache instanceof StringCache) + const targetCache = caches.find( + (cache) => (cache["shouldIntern"] !== undefined) === (entry === "state" || entry === "messageUpdated"), + ) + if (!targetCache) throw new Error("Missing target cache") + const sizeBefore = targetCache.size + internSpy.mockClear() + const messages: Record, ExtensionMessage> = { + state: { type: "state", state: { clineMessages: [{ ...payload, text: "changed" }] } }, + messageUpdated: { type: "messageUpdated", clineMessage: { ...payload, text: "updated" } }, + taskHistoryUpdated: { type: "taskHistoryUpdated", taskHistory: [item] }, + taskHistoryItemUpdated: { type: "taskHistoryItemUpdated", taskHistoryItem: item }, + } + act(() => window.dispatchEvent(new MessageEvent("message", { data: messages[entry] }))) + expect(internSpy.mock.contexts).toEqual([targetCache]) + expect(targetCache.size).toBeGreaterThan(sizeBefore) + expect(internSpy).toHaveBeenCalledWith( + entry === "state" + ? expect.objectContaining({ text: "changed" }) + : entry === "messageUpdated" + ? { ...payload, text: "updated" } + : entry === "taskHistoryUpdated" + ? [item] + : item, + ) + }, + ) + + it.each(["task switch", "message update"] as const)( + "keeps the message mirror synchronized after %s", + (scenario) => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + const payload: ClineMessage = { ts: 1, type: "say", say: "text", text: "seed" } + render( + + + , + ) + const dispatch = (data: ExtensionMessage) => + act(() => window.dispatchEvent(new MessageEvent("message", { data }))) + if (scenario === "task switch") { + internSpy.mockClear() + dispatch({ type: "state", state: { currentTaskId: "task-b", clineMessages: [{ ...payload }] } }) + expect(internSpy).toHaveBeenCalledWith(payload) + } else { + const updated = { ...payload, text: "updated" } + dispatch({ type: "messageUpdated", clineMessage: updated }) + internSpy.mockClear() + dispatch({ type: "state", state: { currentTaskId: "task-a", clineMessages: [{ ...updated }] } }) + expect(internSpy).not.toHaveBeenCalledWith(updated) + } + }, + ) + + it("ignores task identity in pushes without messages", () => { + const initial = { ...createInitialExtensionState(), currentTaskId: "task-a" } + expect(mergeExtensionState(initial, { currentTaskId: "old" }).currentTaskId).toBe("task-a") + const internSpy = vi.spyOn(StringCache.prototype, "intern") + const payload: ClineMessage = { ts: 1, type: "say", say: "text", text: "seed" } + render( + + + , + ) + const cache = internSpy.mock.contexts[0] + if (!(cache instanceof StringCache)) throw new Error("Missing cache") + const size = cache.size + act(() => + window.dispatchEvent( + new MessageEvent("message", { data: { type: "state", state: { currentTaskId: "old" } } }), + ), + ) + expect(cache.size).toBe(size) + internSpy.mockClear() + act(() => + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { currentTaskId: "task-a", clineMessages: [{ ...payload }] } }, + }), + ), + ) + expect(internSpy).not.toHaveBeenCalledWith(payload) + }) + + it("does not retain partial messages or re-intern unchanged snapshot messages", () => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + const payload: ClineMessage = { ts: 1, type: "say", say: "text", text: "streaming", partial: true } + render( + + + , + ) + const cache = internSpy.mock.contexts[0] + if (!(cache instanceof StringCache)) throw new Error("Missing message cache") + expect(cache?.size).toBe(0) + internSpy.mockClear() + act(() => + window.dispatchEvent( + new MessageEvent("message", { data: { type: "state", state: { clineMessages: [{ ...payload }] } } }), + ), + ) + expect(internSpy).not.toHaveBeenCalledWith(payload) + expect(cache?.size).toBe(0) + }) + + it("clears a serialized absent task id and preserves task identity on stale snapshots", () => { + const initial = { ...createInitialExtensionState(), currentTaskId: "task-a", clineMessagesSeq: 5 } + const cleared = mergeExtensionState( + initial, + JSON.parse(JSON.stringify({ currentTaskId: undefined, clineMessages: [], clineMessagesSeq: 6 })), + ) + expect(cleared.currentTaskId).toBeUndefined() + const stale = mergeExtensionState(initial, { + currentTaskId: "old-task", + clineMessages: [], + clineMessagesSeq: 4, + }) + expect(stale.currentTaskId).toBe("task-a") + }) + + it("preserves the message cache on stale task pushes and rebuilds replaced history", () => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + render( + + + , + ) + const messageCache = internSpy.mock.contexts[0] + const historyCache = internSpy.mock.contexts[1] + if (!(messageCache instanceof StringCache) || !(historyCache instanceof StringCache)) + throw new Error("Missing caches") + const messageSize = messageCache.size + historyCache.intern("removed-history") + act(() => + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { currentTaskId: "old", clineMessages: [], clineMessagesSeq: 4 } }, + }), + ), + ) + expect(messageCache.size).toBe(messageSize) + act(() => + window.dispatchEvent( + new MessageEvent("message", { data: { type: "taskHistoryUpdated", taskHistory: [] } }), + ), + ) + expect(historyCache.size).toBe(0) + expect(messageCache.size).toBe(messageSize) + historyCache.intern("removed-again") + act(() => + window.dispatchEvent(new MessageEvent("message", { data: { type: "state", state: { taskHistory: [] } } })), + ) + expect(historyCache.size).toBe(0) + }) + + it("clears the task-scoped message cache only on explicit task transitions", () => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + const clearSpy = vi.spyOn(StringCache.prototype, "clear") + render( + + + , + ) + const messageCache = internSpy.mock.contexts[0] + const historyCache = internSpy.mock.contexts[1] + if (!(messageCache instanceof StringCache) || !(historyCache instanceof StringCache)) { + throw new Error("Missing provider caches") + } + expect(messageCache?.size).toBeGreaterThan(0) + historyCache?.intern("history-seed") + const historySize = historyCache?.size + clearSpy.mockClear() + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { data: { type: "state", state: { currentTaskId: "task-a" } } }), + ) + window.dispatchEvent(new MessageEvent("message", { data: { type: "state", state: { language: "en" } } })) + }) + expect(clearSpy).not.toHaveBeenCalled() + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { + data: JSON.parse( + JSON.stringify({ type: "state", state: { currentTaskId: undefined, clineMessages: [] } }), + ), + }), + ) + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { currentTaskId: "task-b", clineMessages: [] } }, + }), + ) + }) + expect(clearSpy).toHaveBeenCalledTimes(2) + expect(messageCache?.size).toBe(0) + expect(historyCache?.size).toBe(historySize) + }) + + it("does not intern messages from stale sequenced state updates", () => { + const internSpy = vi.spyOn(StringCache.prototype, "intern") + render( + + + , + ) + internSpy.mockClear() + const staleMessage = { ts: 1, type: "say", say: "text", text: "stale" } as ClineMessage + const currentMessage = { ts: 2, type: "say", say: "text", text: "current" } as ClineMessage + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { clineMessages: [staleMessage], clineMessagesSeq: 5 } }, + }), + ) + }) + expect(internSpy).not.toHaveBeenCalledWith(staleMessage) + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { clineMessages: [currentMessage], clineMessagesSeq: 6 } }, + }), + ) + }) + expect(internSpy).toHaveBeenCalledWith(currentMessage) + internSpy.mockClear() + act(() => { + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { clineMessages: [staleMessage], clineMessagesSeq: 6 } }, + }), + ) + }) + expect(internSpy).not.toHaveBeenCalledWith(staleMessage) + }) + + it("does not intern an update for a message outside the current task", () => { + const currentMessage = { ts: 1, type: "say", say: "text", text: "current" } as ClineMessage + const unknownMessage = { ts: 2, type: "say", say: "text", text: "late" } as ClineMessage + const internSpy = vi.spyOn(StringCache.prototype, "intern") + vi.spyOn(console, "warn").mockImplementation(() => undefined) + render( + + + , + ) + internSpy.mockClear() + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { data: { type: "messageUpdated", clineMessage: unknownMessage } }), + ) + }) + + expect(internSpy).not.toHaveBeenCalledWith(unknownMessage) + }) + it("initializes with empty allowedCommands array", () => { render( diff --git a/webview-ui/src/utils/__tests__/stringCache.spec.ts b/webview-ui/src/utils/__tests__/stringCache.spec.ts new file mode 100644 index 0000000000..7ef523d91d --- /dev/null +++ b/webview-ui/src/utils/__tests__/stringCache.spec.ts @@ -0,0 +1,69 @@ +import { StringCache } from "../stringCache" + +describe("StringCache", () => { + it.each([ + { entries: 2, characters: 1000, values: ["a", "b", "c"], retained: ["b", "c"] }, + { entries: 2, characters: 3, values: ["abc"], retained: ["abc"] }, + { entries: 0, characters: 1000, values: ["a"], retained: [] }, + ])( + "honors independent limits: $entries entries, $characters characters", + ({ entries, characters, values, retained }) => { + const cache = new StringCache(undefined, entries, characters) + cache.intern(values) + expect([...cache["cache"].keys()]).toEqual(retained) + }, + ) + + it("bounds entry count and retained characters, skipping oversized strings", () => { + const cache = new StringCache(undefined, 2, 6) + cache.intern(["aaa", "bbb", "ccc"]) + expect(cache.size).toBe(2) + expect([...cache["cache"].keys()]).toEqual(["bbb", "ccc"]) + cache.intern("oversized") + expect(cache.size).toBe(2) + cache.intern("dddd") + expect([...cache["cache"].keys()]).toEqual(["dddd"]) + cache.clear() + cache.intern(["aaa", "bbb"]) + expect(cache.size).toBe(2) + }) + + it("interns strings recursively without retaining duplicate or empty values", () => { + const cache = new StringCache() + const value = { + label: "shared", + nested: [{ label: "shared" }, { label: "other" }], + empty: "", + count: 1, + } + + expect(cache.intern(value)).toBe(value) + expect(cache.size).toBe(2) + + cache.intern({ label: "shared", another: "other" }) + expect(cache.size).toBe(2) + }) + + it("skips values rejected by its filter while interning accepted siblings", () => { + const cache = new StringCache((value) => !("partial" in value) || value.partial !== true) + const partial = { partial: true, text: "streaming", nested: { text: "transient" } } + const complete = { partial: false, text: "complete", nested: { partial: true, text: "still-complete" } } + + cache.intern(partial) + expect(cache.size).toBe(0) + + cache.intern(complete) + expect(cache.size).toBe(2) + }) + + it("clears interned values and can be reused", () => { + const cache = new StringCache() + cache.intern({ text: "first" }) + + cache.clear() + expect(cache.size).toBe(0) + + cache.intern({ text: "second" }) + expect(cache.size).toBe(1) + }) +}) diff --git a/webview-ui/src/utils/stringCache.ts b/webview-ui/src/utils/stringCache.ts new file mode 100644 index 0000000000..f969525bf1 --- /dev/null +++ b/webview-ui/src/utils/stringCache.ts @@ -0,0 +1,79 @@ +export type StringCacheFilter = (value: object) => boolean + +/** + * Rebinds strings in deserialized values to canonical instances held by this cache. + * FIFO eviction bounds both entry overhead and retained UTF-16 characters. + * Values are mutated in place to avoid allocating another object tree per message. + */ +export class StringCache { + private readonly cache = new Map() + + private retainedCharacters = 0 + + constructor( + private readonly shouldIntern?: StringCacheFilter, + private readonly maxEntries = 10_000, + private readonly maxCharacters = 2_000_000, + ) {} + + intern(value: T): T { + if (value !== null && typeof value === "object" && this.shouldIntern?.(value) === false) { + return value + } + + return this.internValue(value) as T + } + + clear(): void { + this.cache.clear() + this.retainedCharacters = 0 + } + + get size(): number { + return this.cache.size + } + + private internValue(value: unknown): unknown { + if (typeof value === "string") { + if (value.length === 0) { + return value + } + + const existing = this.cache.get(value) + if (existing !== undefined) { + return existing + } + + if (this.maxEntries <= 0 || value.length > this.maxCharacters) { + return value + } + while (this.cache.size >= this.maxEntries || this.retainedCharacters + value.length > this.maxCharacters) { + const oldest = this.cache.keys().next().value + if (oldest === undefined) break + this.cache.delete(oldest) + this.retainedCharacters -= oldest.length + } + this.cache.set(value, value) + this.retainedCharacters += value.length + return value + } + + if (Array.isArray(value)) { + for (let index = 0; index < value.length; index++) { + value[index] = this.internValue(value[index]) + } + return value + } + + if (value === null || typeof value !== "object") { + return value + } + + const record = value as Record + for (const key of Object.keys(record)) { + record[key] = this.internValue(record[key]) + } + + return value + } +}