From 42c0eb1df9464f897f6bda6d4b52e0dd1bb56d0b Mon Sep 17 00:00:00 2001 From: "@taltas" <6816042+taltas@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:55:12 +0000 Subject: [PATCH 1/2] perf(webview): reduce retained chat state --- webview-ui/src/App.tsx | 2 + webview-ui/src/__tests__/App.spec.tsx | 35 +++++++- webview-ui/src/components/chat/ChatView.tsx | 5 +- .../chat/__tests__/ChatView.spec.tsx | 5 +- .../src/context/ExtensionStateContext.tsx | 58 ++++++++++++-- .../__tests__/ExtensionStateContext.spec.tsx | 80 +++++++++++++++++++ .../src/utils/__tests__/stringCache.spec.ts | 42 ++++++++++ webview-ui/src/utils/stringCache.ts | 61 ++++++++++++++ 8 files changed, 279 insertions(+), 9 deletions(-) create mode 100644 webview-ui/src/utils/__tests__/stringCache.spec.ts create mode 100644 webview-ui/src/utils/stringCache.ts diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx index b1fbf82999..0c527e4e5b 100644 --- a/webview-ui/src/App.tsx +++ b/webview-ui/src/App.tsx @@ -55,6 +55,7 @@ const App = () => { useThemeFixtureProbe() const { + currentTaskId, didHydrateState, showWelcome, settingsImportedAt, @@ -260,6 +261,7 @@ const App = () => { /> )} ({ vscode: { postMessage: vi.fn(), @@ -28,8 +30,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
) @@ -160,6 +163,7 @@ vi.mock("process.env", () => ({ describe("App", () => { beforeEach(() => { vi.clearAllMocks() + nextChatViewInstance = 0 window.removeEventListener("message", () => {}) // Set up default mock return value @@ -205,6 +209,35 @@ describe("App", () => { expect(chatView.getAttribute("data-hidden")).toBe("false") }, 10000) + it("remounts chat on task transitions but preserves it across 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).not.toBe(firstInstance) + + state.currentTaskId = undefined + rerender() + expect(screen.getByTestId("chat-view").getAttribute("data-instance")).not.toBe(secondInstance) + }) + it("shows welcome view when setup is incomplete", () => { mockUseExtensionState.mockReturnValue({ didHydrateState: true, diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 988c98ba30..8b578aa2fa 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -1510,7 +1510,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction `${messageOrGroup.ts}-${index}`, + (index: number, messageOrGroup: ClineMessage) => + `${messageOrGroup.ts}-${index}-${messageOrGroup.partial ? "partial" : "full"}`, [], ) diff --git a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx index 6b2fa177c9..cd75b21577 100644 --- a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx +++ b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx @@ -708,7 +708,10 @@ 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 })).toBe(`${rowTs}-1-full`) + expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs, partial: true })).toBe( + `${rowTs}-1-partial`, + ) }) }) diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index 377c8eb721..0518995319 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -1,4 +1,4 @@ -import React, { createContext, useCallback, useEffect, useState } from "react" +import React, { createContext, useCallback, useEffect, useRef, useState } from "react" import { type ProviderSettings, @@ -33,6 +33,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 @@ -283,9 +284,22 @@ 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 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) @@ -341,6 +355,37 @@ export const ExtensionStateContextProvider: React.FC<{ switch (message.type) { case "state": { const newState = message.state ?? {} + if (Object.prototype.hasOwnProperty.call(newState, "currentTaskId")) { + if (newState.currentTaskId !== currentTaskIdRef.current) { + messageCache.clear() + currentTaskIdRef.current = newState.currentTaskId + } + } + + const messagesAreStale = + newState.clineMessages !== undefined && + newState.clineMessagesSeq !== undefined && + clineMessagesSeqRef.current !== undefined && + newState.clineMessagesSeq <= clineMessagesSeqRef.current + + if (newState.clineMessages && !messagesAreStale) { + for (const clineMessage of newState.clineMessages) { + messageCache.intern(clineMessage) + } + } + if ( + newState.clineMessagesSeq !== undefined && + (clineMessagesSeqRef.current === undefined || + newState.clineMessagesSeq > clineMessagesSeqRef.current) + ) { + clineMessagesSeqRef.current = newState.clineMessagesSeq + } + if (newState.taskHistory) { + historyCache.intern(newState.taskHistory) + } + if (newState.currentTaskItem) { + historyCache.intern(newState.currentTaskItem) + } setState((prevState) => mergeExtensionState(prevState, newState)) setShowWelcome(!checkExistKey(newState.apiConfiguration, newState.zooCodeIsAuthenticated)) setDidHydrateState(true) @@ -409,6 +454,7 @@ export const ExtensionStateContextProvider: React.FC<{ // 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) if (lastIndex !== -1) { + messageCache.intern(clineMessage) const newClineMessages = [...prevState.clineMessages] newClineMessages[lastIndex] = clineMessage return { ...prevState, clineMessages: newClineMessages } @@ -469,6 +515,7 @@ export const ExtensionStateContextProvider: React.FC<{ case "taskHistoryUpdated": { // Efficiently update just the task history without replacing entire state if (message.taskHistory !== undefined) { + historyCache.intern(message.taskHistory) setState((prevState) => ({ ...prevState, taskHistory: message.taskHistory!, @@ -481,6 +528,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 @@ -503,7 +551,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 23ac911585..44f26ae5e9 100644 --- a/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx +++ b/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx @@ -14,6 +14,7 @@ import { } from "@roo-code/types" import { ExtensionStateContextProvider, useExtensionState, mergeExtensionState } from "../ExtensionStateContext" +import { StringCache } from "@/utils/stringCache" const TestComponent = () => { const { allowedCommands, setAllowedCommands, soundEnabled, showRooIgnoredFiles, setShowRooIgnoredFiles } = @@ -103,6 +104,85 @@ const InitialStateTestComponent = () => { } describe("ExtensionStateContext", () => { + it("clears the task-scoped message cache only on explicit task transitions", () => { + const clearSpy = vi.spyOn(StringCache.prototype, "clear") + render( + + + , + ) + 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: { type: "state", state: { currentTaskId: undefined } } }), + ) + window.dispatchEvent( + new MessageEvent("message", { data: { type: "state", state: { currentTaskId: "task-b" } } }), + ) + }) + expect(clearSpy).toHaveBeenCalledTimes(2) + }) + + 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: 4 } }, + }), + ) + }) + expect(internSpy).not.toHaveBeenCalledWith(staleMessage) + + act(() => { + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "state", state: { clineMessages: [currentMessage], clineMessagesSeq: 6 } }, + }), + ) + }) + expect(internSpy).toHaveBeenCalledWith(currentMessage) + }) + + 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..e699c898c8 --- /dev/null +++ b/webview-ui/src/utils/__tests__/stringCache.spec.ts @@ -0,0 +1,42 @@ +import { StringCache } from "../stringCache" + +describe("StringCache", () => { + 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..1fe34e7cca --- /dev/null +++ b/webview-ui/src/utils/stringCache.ts @@ -0,0 +1,61 @@ +export type StringCacheFilter = (value: object) => boolean + +/** + * Rebinds strings in deserialized values to canonical instances held by this cache. + * Values are mutated in place to avoid allocating another object tree per message. + */ +export class StringCache { + private readonly cache = new Map() + + constructor(private readonly shouldIntern?: StringCacheFilter) {} + + 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() + } + + 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 + } + + this.cache.set(value, value) + 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 + } +} From b9a6ba03cb93c0f92f06cfb00c82ae7d488db7cb Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Tue, 27 Oct 2026 03:50:56 +0000 Subject: [PATCH 2/2] fix(webview): bound string caches and guard task identity updates --- webview-ui/src/App.tsx | 2 - webview-ui/src/__tests__/App.spec.tsx | 6 +- webview-ui/src/components/chat/ChatView.tsx | 5 +- .../chat/__tests__/ChatView.spec.tsx | 27 ++- .../src/context/ExtensionStateContext.tsx | 47 +++- .../__tests__/ExtensionStateContext.spec.tsx | 220 +++++++++++++++++- .../src/utils/__tests__/stringCache.spec.ts | 27 +++ webview-ui/src/utils/stringCache.ts | 20 +- 8 files changed, 330 insertions(+), 24 deletions(-) diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx index e467115837..3e05ce64da 100644 --- a/webview-ui/src/App.tsx +++ b/webview-ui/src/App.tsx @@ -55,7 +55,6 @@ const App = () => { useThemeFixtureProbe() const { - currentTaskId, didHydrateState, showWelcome, settingsImportedAt, @@ -272,7 +271,6 @@ const App = () => { /> )} { expect(chatView.getAttribute("data-hidden")).toBe("false") }, 10000) - it("remounts chat on task transitions but preserves it across tab changes", async () => { + it("preserves chat state across task and tab changes", async () => { const state = { didHydrateState: true, showWelcome: false, @@ -232,11 +232,11 @@ describe("App", () => { state.currentTaskId = "task-b" rerender() const secondInstance = screen.getByTestId("chat-view").getAttribute("data-instance") - expect(secondInstance).not.toBe(firstInstance) + expect(secondInstance).toBe(firstInstance) state.currentTaskId = undefined rerender() - expect(screen.getByTestId("chat-view").getAttribute("data-instance")).not.toBe(secondInstance) + expect(screen.getByTestId("chat-view").getAttribute("data-instance")).toBe(secondInstance) }) it.each(["sidebar", "editor"])( diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index c0a52439aa..bb690c1f03 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -1518,7 +1518,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction - `${messageOrGroup.ts}-${index}-${messageOrGroup.partial ? "partial" : "full"}`, + (index: number, messageOrGroup: ClineMessage) => `${messageOrGroup.ts}-${index}`, [], ) diff --git a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx index f47786a3be..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() @@ -750,9 +771,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-full`) + 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-partial`, + `${rowTs}-1`, ) }) }) diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index 2b2bf67689..7c352a2193 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -12,6 +12,7 @@ import { type OrganizationAllowList, type CloudOrganizationMembership, type ExtensionMessage, + type ClineMessage, type ExtensionState, type MarketplaceInstalledMetadata, type SkillMetadata, @@ -186,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 @@ -291,6 +298,7 @@ export const ExtensionStateContextProvider: React.FC<{ }> = ({ children, 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) @@ -360,12 +368,6 @@ export const ExtensionStateContextProvider: React.FC<{ switch (message.type) { case "state": { const newState = message.state ?? {} - if (Object.prototype.hasOwnProperty.call(newState, "currentTaskId")) { - if (newState.currentTaskId !== currentTaskIdRef.current) { - messageCache.clear() - currentTaskIdRef.current = newState.currentTaskId - } - } const messagesAreStale = newState.clineMessages !== undefined && @@ -373,10 +375,31 @@ export const ExtensionStateContextProvider: React.FC<{ 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) { - messageCache.intern(clineMessage) + 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 && @@ -386,6 +409,7 @@ export const ExtensionStateContextProvider: React.FC<{ clineMessagesSeqRef.current = newState.clineMessagesSeq } if (newState.taskHistory) { + historyCache.clear() historyCache.intern(newState.taskHistory) } if (newState.currentTaskItem) { @@ -455,11 +479,17 @@ 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) if (lastIndex !== -1) { - messageCache.intern(clineMessage) const newClineMessages = [...prevState.clineMessages] newClineMessages[lastIndex] = clineMessage return { ...prevState, clineMessages: newClineMessages } @@ -520,6 +550,7 @@ 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, diff --git a/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx b/webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx index 929702607a..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, @@ -113,13 +114,207 @@ 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(() => { @@ -132,13 +327,21 @@ describe("ExtensionStateContext", () => { act(() => { window.dispatchEvent( - new MessageEvent("message", { data: { type: "state", state: { currentTaskId: undefined } } }), + 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" } } }), + 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", () => { @@ -155,7 +358,7 @@ describe("ExtensionStateContext", () => { act(() => { window.dispatchEvent( new MessageEvent("message", { - data: { type: "state", state: { clineMessages: [staleMessage], clineMessagesSeq: 4 } }, + data: { type: "state", state: { clineMessages: [staleMessage], clineMessagesSeq: 5 } }, }), ) }) @@ -169,6 +372,15 @@ describe("ExtensionStateContext", () => { ) }) 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", () => { diff --git a/webview-ui/src/utils/__tests__/stringCache.spec.ts b/webview-ui/src/utils/__tests__/stringCache.spec.ts index e699c898c8..7ef523d91d 100644 --- a/webview-ui/src/utils/__tests__/stringCache.spec.ts +++ b/webview-ui/src/utils/__tests__/stringCache.spec.ts @@ -1,6 +1,33 @@ 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 = { diff --git a/webview-ui/src/utils/stringCache.ts b/webview-ui/src/utils/stringCache.ts index 1fe34e7cca..f969525bf1 100644 --- a/webview-ui/src/utils/stringCache.ts +++ b/webview-ui/src/utils/stringCache.ts @@ -2,12 +2,19 @@ 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() - constructor(private readonly shouldIntern?: StringCacheFilter) {} + 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) { @@ -19,6 +26,7 @@ export class StringCache { clear(): void { this.cache.clear() + this.retainedCharacters = 0 } get size(): number { @@ -36,7 +44,17 @@ export class StringCache { 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 }