Skip to content

Commit b9a6ba0

Browse files
edelaunaedelauna
committed
fix(webview): bound string caches and guard task identity updates
1 parent de8291c commit b9a6ba0

8 files changed

Lines changed: 330 additions & 24 deletions

File tree

‎webview-ui/src/App.tsx‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@ const App = () => {
5555
useThemeFixtureProbe()
5656

5757
const {
58-
currentTaskId,
5958
didHydrateState,
6059
showWelcome,
6160
settingsImportedAt,
@@ -272,7 +271,6 @@ const App = () => {
272271
/>
273272
)}
274273
<ChatView
275-
key={currentTaskId ?? "no-task"}
276274
ref={chatViewRef}
277275
isHidden={tab !== "chat"}
278276
showAnnouncement={showAnnouncement}

‎webview-ui/src/__tests__/App.spec.tsx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ describe("App", () => {
210210
expect(chatView.getAttribute("data-hidden")).toBe("false")
211211
}, 10000)
212212

213-
it("remounts chat on task transitions but preserves it across tab changes", async () => {
213+
it("preserves chat state across task and tab changes", async () => {
214214
const state = {
215215
didHydrateState: true,
216216
showWelcome: false,
@@ -232,11 +232,11 @@ describe("App", () => {
232232
state.currentTaskId = "task-b"
233233
rerender(<AppWithProviders />)
234234
const secondInstance = screen.getByTestId("chat-view").getAttribute("data-instance")
235-
expect(secondInstance).not.toBe(firstInstance)
235+
expect(secondInstance).toBe(firstInstance)
236236

237237
state.currentTaskId = undefined
238238
rerender(<AppWithProviders />)
239-
expect(screen.getByTestId("chat-view").getAttribute("data-instance")).not.toBe(secondInstance)
239+
expect(screen.getByTestId("chat-view").getAttribute("data-instance")).toBe(secondInstance)
240240
})
241241

242242
it.each(["sidebar", "editor"])(

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1518,7 +1518,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
15181518
// regular message
15191519
return (
15201520
<ChatRow
1521-
key={`${messageOrGroup.ts}-${messageOrGroup.partial ? "partial" : "full"}`}
1521+
key={messageOrGroup.ts}
15221522
message={messageOrGroup}
15231523
isExpanded={expandedRows[messageOrGroup.ts] || false}
15241524
onToggleExpand={toggleRowExpansion} // This was already stabilized
@@ -1573,8 +1573,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
15731573
)
15741574

15751575
const computeMessageKey = useCallback(
1576-
(index: number, messageOrGroup: ClineMessage) =>
1577-
`${messageOrGroup.ts}-${index}-${messageOrGroup.partial ? "partial" : "full"}`,
1576+
(index: number, messageOrGroup: ClineMessage) => `${messageOrGroup.ts}-${index}`,
15781577
[],
15791578
)
15801579

‎webview-ui/src/components/chat/__tests__/ChatView.spec.tsx‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
mockVscodePostMessage,
77
renderWithExtensionState,
88
waitFor,
9+
screen,
910
act,
1011
fireEvent,
1112
} from "@/utils/test-utils"
@@ -49,6 +50,8 @@ vi.mock("use-sound", () => ({
4950
}),
5051
}))
5152

53+
let nextRowInstance = 0
54+
5255
// Mock components that use ESM dependencies
5356
vi.mock("../ChatRow", () => ({
5457
default: function MockChatRow({
@@ -60,6 +63,7 @@ vi.mock("../ChatRow", () => ({
6063
onSuggestionClick?: (suggestion: SuggestionItem, event?: React.MouseEvent) => void
6164
isFollowUpAnswered?: boolean
6265
}) {
66+
const [instance] = React.useState(() => ++nextRowInstance)
6367
if (message.type === "ask" && message.ask === "followup" && message.text) {
6468
try {
6569
const followUp = JSON.parse(message.text) as { suggest?: SuggestionItem[] }
@@ -81,7 +85,11 @@ vi.mock("../ChatRow", () => ({
8185
}
8286
}
8387

84-
return <div data-testid="chat-row">{JSON.stringify(message)}</div>
88+
return (
89+
<div data-testid="chat-row" data-instance={instance}>
90+
{JSON.stringify(message)}
91+
</div>
92+
)
8593
},
8694
}))
8795

@@ -721,6 +729,19 @@ describe("ChatView - Virtualization Configuration", () => {
721729
mockVirtuosoState.lastConfig = null
722730
})
723731

732+
it("preserves the row instance when a partial message completes", async () => {
733+
renderChatView()
734+
const task: ClineMessage = { type: "say", say: "task", ts: 1, text: "Task" }
735+
const partial: ClineMessage = { type: "say", say: "text", ts: 2, text: "Streaming", partial: true }
736+
mockPostMessage({ clineMessages: [task, partial] })
737+
const row = await screen.findByTestId("chat-row")
738+
const instance = row.getAttribute("data-instance")
739+
expect(instance).not.toBeNull()
740+
mockPostMessage({ clineMessages: [task, { ...partial, text: "Completed", partial: false }] })
741+
await waitFor(() => expect(screen.getByTestId("chat-row")).toHaveTextContent("Completed"))
742+
expect(screen.getByTestId("chat-row").getAttribute("data-instance")).toBe(instance)
743+
})
744+
724745
it("keeps the off-screen render buffer tight for chat rows", async () => {
725746
renderChatView()
726747

@@ -750,9 +771,9 @@ describe("ChatView - Virtualization Configuration", () => {
750771

751772
expect(mockVirtuosoState.lastConfig?.defaultItemHeight).toBe(180)
752773
expect(mockVirtuosoState.lastConfig?.increaseViewportBy).toEqual({ top: 600, bottom: 800 })
753-
expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs })).toBe(`${rowTs}-1-full`)
774+
expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs })).toBe(`${rowTs}-1`)
754775
expect(mockVirtuosoState.lastConfig?.computeItemKey?.(1, { type: "say", ts: rowTs, partial: true })).toBe(
755-
`${rowTs}-1-partial`,
776+
`${rowTs}-1`,
756777
)
757778
})
758779
})

‎webview-ui/src/context/ExtensionStateContext.tsx‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
type OrganizationAllowList,
1313
type CloudOrganizationMembership,
1414
type ExtensionMessage,
15+
type ClineMessage,
1516
type ExtensionState,
1617
type MarketplaceInstalledMetadata,
1718
type SkillMetadata,
@@ -186,6 +187,12 @@ export const mergeExtensionState = (prevState: ExtensionState, newState: Partial
186187
) {
187188
rest.clineMessages = prevState.clineMessages
188189
rest.clineMessagesSeq = prevState.clineMessagesSeq
190+
rest.currentTaskId = prevState.currentTaskId
191+
} else if (newState.clineMessages !== undefined) {
192+
rest.currentTaskId = newState.currentTaskId
193+
} else {
194+
// Pushes without messages cannot establish task identity.
195+
rest.currentTaskId = prevState.currentTaskId
189196
}
190197

191198
// Note that we completely replace the previous apiConfiguration and customSupportPrompts objects
@@ -291,6 +298,7 @@ export const ExtensionStateContextProvider: React.FC<{
291298
}> = ({ children, initialState }) => {
292299
const [historyCache] = useState(() => new StringCache())
293300
const [messageCache] = useState(() => new StringCache((value) => !("partial" in value) || value.partial !== true))
301+
const messagesRef = useRef(initialState?.clineMessages ?? [])
294302
const currentTaskIdRef = useRef(initialState?.currentTaskId)
295303
const clineMessagesSeqRef = useRef(initialState?.clineMessagesSeq)
296304

@@ -360,23 +368,38 @@ export const ExtensionStateContextProvider: React.FC<{
360368
switch (message.type) {
361369
case "state": {
362370
const newState = message.state ?? {}
363-
if (Object.prototype.hasOwnProperty.call(newState, "currentTaskId")) {
364-
if (newState.currentTaskId !== currentTaskIdRef.current) {
365-
messageCache.clear()
366-
currentTaskIdRef.current = newState.currentTaskId
367-
}
368-
}
369371

370372
const messagesAreStale =
371373
newState.clineMessages !== undefined &&
372374
newState.clineMessagesSeq !== undefined &&
373375
clineMessagesSeqRef.current !== undefined &&
374376
newState.clineMessagesSeq <= clineMessagesSeqRef.current
375377

378+
if (!messagesAreStale && newState.clineMessages !== undefined) {
379+
if (newState.currentTaskId !== currentTaskIdRef.current) {
380+
messageCache.clear()
381+
messagesRef.current = []
382+
currentTaskIdRef.current = newState.currentTaskId
383+
}
384+
}
385+
376386
if (newState.clineMessages && !messagesAreStale) {
387+
const previousMessages = new Map(messagesRef.current.map((item) => [item.ts, item]))
377388
for (const clineMessage of newState.clineMessages) {
378-
messageCache.intern(clineMessage)
389+
const previous = previousMessages.get(clineMessage.ts)
390+
if (
391+
previous &&
392+
Object.keys(clineMessage).length === Object.keys(previous).length &&
393+
Object.entries(clineMessage).every(
394+
([key, value]) => previous[key as keyof ClineMessage] === value,
395+
)
396+
) {
397+
Object.assign(clineMessage, previous)
398+
} else {
399+
messageCache.intern(clineMessage)
400+
}
379401
}
402+
messagesRef.current = newState.clineMessages
380403
}
381404
if (
382405
newState.clineMessagesSeq !== undefined &&
@@ -386,6 +409,7 @@ export const ExtensionStateContextProvider: React.FC<{
386409
clineMessagesSeqRef.current = newState.clineMessagesSeq
387410
}
388411
if (newState.taskHistory) {
412+
historyCache.clear()
389413
historyCache.intern(newState.taskHistory)
390414
}
391415
if (newState.currentTaskItem) {
@@ -455,11 +479,17 @@ export const ExtensionStateContextProvider: React.FC<{
455479
}
456480
case "messageUpdated": {
457481
const clineMessage = message.clineMessage!
482+
const messageIndex = messagesRef.current.findIndex((item) => item.ts === clineMessage.ts)
483+
if (messageIndex !== -1) {
484+
messageCache.intern(clineMessage)
485+
messagesRef.current = messagesRef.current.map((item, index) =>
486+
index === messageIndex ? clineMessage : item,
487+
)
488+
}
458489
setState((prevState) => {
459490
// 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
460491
const lastIndex = findLastIndex(prevState.clineMessages, (msg) => msg.ts === clineMessage.ts)
461492
if (lastIndex !== -1) {
462-
messageCache.intern(clineMessage)
463493
const newClineMessages = [...prevState.clineMessages]
464494
newClineMessages[lastIndex] = clineMessage
465495
return { ...prevState, clineMessages: newClineMessages }
@@ -520,6 +550,7 @@ export const ExtensionStateContextProvider: React.FC<{
520550
case "taskHistoryUpdated": {
521551
// Efficiently update just the task history without replacing entire state
522552
if (message.taskHistory !== undefined) {
553+
historyCache.clear()
523554
historyCache.intern(message.taskHistory)
524555
setState((prevState) => ({
525556
...prevState,

0 commit comments

Comments
 (0)