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 + } +}