Skip to content

Commit 14f65ba

Browse files
easonliang28easonliang28
committed
fix(webview): persist view-local state safely
1 parent 2940c14 commit 14f65ba

6 files changed

Lines changed: 306 additions & 38 deletions

File tree

‎src/core/webview/ClineProvider.ts‎

Lines changed: 60 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@ export class ClineProvider
166166
private static activeInstances: Set<ClineProvider> = new Set()
167167
private static nextViewId = 0
168168
private static readonly MAX_PERSISTED_VIEW_STATES = 50
169+
private static persistedViewStateWriteQueue: Promise<void> = Promise.resolve()
169170
private disposables: vscode.Disposable[] = []
170171
private webviewDisposables: vscode.Disposable[] = []
171172
private view?: vscode.WebviewView | vscode.WebviewPanel
@@ -435,51 +436,63 @@ export class ClineProvider
435436
}
436437
}
437438

438-
private getPersistedViewStates(): Record<string, PersistedViewState> {
439-
const viewStates = this.contextProxy.getValue("viewStates")
439+
private getPersistedViewStates(options: { fresh?: boolean } = {}): Record<string, PersistedViewState> {
440+
const viewStates = options.fresh
441+
? this.context.globalState.get<GlobalState["viewStates"]>("viewStates")
442+
: this.contextProxy.getValue("viewStates")
440443

441444
if (!viewStates || typeof viewStates !== "object" || Array.isArray(viewStates)) {
442445
return {}
443446
}
444447

445-
return viewStates
448+
return { ...viewStates }
446449
}
447450

448451
private async savePersistedViewState(values: Partial<PersistedViewState>): Promise<void> {
449-
const states = this.getPersistedViewStates()
450-
const current = states[this.viewStateId] ?? {}
451-
const next: PersistedViewState = { ...current }
452+
const write = ClineProvider.persistedViewStateWriteQueue.then(async () => {
453+
const states = this.getPersistedViewStates({ fresh: true })
454+
const current = states[this.viewStateId] ?? {}
455+
const next: PersistedViewState = { ...current }
456+
457+
if ("mode" in values) {
458+
if (values.mode === undefined || values.mode === null) {
459+
delete next.mode
460+
} else {
461+
next.mode = values.mode
462+
}
463+
}
452464

453-
if ("mode" in values) {
454-
if (values.mode === undefined || values.mode === null) {
455-
delete next.mode
456-
} else {
457-
next.mode = values.mode
465+
if ("currentApiConfigName" in values) {
466+
if (values.currentApiConfigName === undefined || values.currentApiConfigName === null) {
467+
delete next.currentApiConfigName
468+
} else {
469+
next.currentApiConfigName = values.currentApiConfigName
470+
}
458471
}
459-
}
460472

461-
if ("currentApiConfigName" in values) {
462-
if (values.currentApiConfigName === undefined || values.currentApiConfigName === null) {
463-
delete next.currentApiConfigName
473+
if (!next.mode && !next.currentApiConfigName) {
474+
delete states[this.viewStateId]
464475
} else {
465-
next.currentApiConfigName = values.currentApiConfigName
476+
next.updatedAt = values.updatedAt ?? Date.now()
477+
states[this.viewStateId] = next
466478
}
467-
}
468479

469-
if (!next.mode && !next.currentApiConfigName) {
470-
delete states[this.viewStateId]
471-
} else {
472-
next.updatedAt = values.updatedAt ?? Date.now()
473-
states[this.viewStateId] = next
474-
}
480+
await this.contextProxy.setValue("viewStates", this.prunePersistedViewStates(states))
481+
})
475482

476-
await this.contextProxy.setValue("viewStates", this.prunePersistedViewStates(states))
483+
ClineProvider.persistedViewStateWriteQueue = write.catch(() => {})
484+
await write
477485
}
478486

479487
private async clearPersistedViewState(viewStateId = this.viewStateId): Promise<void> {
480-
const states = this.getPersistedViewStates()
481-
delete states[viewStateId]
482-
await this.contextProxy.setValue("viewStates", states)
488+
const write = ClineProvider.persistedViewStateWriteQueue.then(async () => {
489+
const states = this.getPersistedViewStates({ fresh: true })
490+
delete states[viewStateId]
491+
await this.contextProxy.setValue("viewStates", states)
492+
})
493+
494+
ClineProvider.persistedViewStateWriteQueue = write.catch(() => {})
495+
await write
483496
}
484497

485498
private prunePersistedViewStates(states: Record<string, PersistedViewState>): Record<string, PersistedViewState> {
@@ -3041,6 +3054,7 @@ export class ClineProvider
30413054
public async setValue<K extends keyof RooCodeSettings>(key: K, value: RooCodeSettings[K]) {
30423055
await this.contextProxy.setValue(key, value)
30433056
this._updateViewLocalStateFromMutation({ [key]: value })
3057+
await this._persistViewLocalStateFromMutation({ [key]: value })
30443058
}
30453059

30463060
public getValue<K extends keyof RooCodeSettings>(key: K) {
@@ -3054,6 +3068,7 @@ export class ClineProvider
30543068
public async setValues(values: RooCodeSettings) {
30553069
await this.contextProxy.setValues(values)
30563070
this._updateViewLocalStateFromMutation(values)
3071+
await this._persistViewLocalStateFromMutation(values)
30573072
}
30583073

30593074
/**
@@ -3103,6 +3118,24 @@ export class ClineProvider
31033118
}
31043119
}
31053120

3121+
private async _persistViewLocalStateFromMutation(
3122+
values: Partial<RooCodeSettings> & Partial<ExtensionState>,
3123+
): Promise<void> {
3124+
const persistedValues: Partial<PersistedViewState> = {}
3125+
3126+
if ("mode" in values) {
3127+
persistedValues.mode = values.mode as PersistedViewState["mode"]
3128+
}
3129+
3130+
if ("currentApiConfigName" in values) {
3131+
persistedValues.currentApiConfigName = values.currentApiConfigName
3132+
}
3133+
3134+
if ("mode" in persistedValues || "currentApiConfigName" in persistedValues) {
3135+
await this.savePersistedViewState(persistedValues)
3136+
}
3137+
}
3138+
31063139
/**
31073140
* Clear view-local state cache so that getState() falls back to ContextProxy defaults.
31083141
*/

‎src/core/webview/__tests__/ClineProvider.parallelMode.spec.ts‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -691,12 +691,12 @@ describe("ClineProvider - Parallel Mode Support", () => {
691691
)
692692
const provider2 = new ClineProvider(mockContext, mockOutputChannel, "editor", new ContextProxy(mockContext))
693693

694-
// Access viewLocalState via private property for testing
694+
await (provider1 as any).saveViewState("mode", "architect")
695+
695696
const state1 = await provider1.getState()
696697
const state2 = await provider2.getState()
697698

698-
// Both should start with the same default mode from global state
699-
expect(state1.mode).toBe("code")
699+
expect(state1.mode).toBe("architect")
700700
expect(state2.mode).toBe("code")
701701

702702
await provider1.dispose()
@@ -843,6 +843,31 @@ describe("ClineProvider - Parallel Mode Support", () => {
843843

844844
await provider.dispose()
845845
})
846+
it("should merge concurrent persisted updates from separate provider instances without lost viewStates", async () => {
847+
const provider1 = new ClineProvider(
848+
mockContext,
849+
mockOutputChannel,
850+
"sidebar",
851+
new ContextProxy(mockContext),
852+
)
853+
const provider2 = new ClineProvider(mockContext, mockOutputChannel, "editor", new ContextProxy(mockContext))
854+
855+
await (provider1 as any).setViewStateId("stable-sidebar-view")
856+
await (provider2 as any).setViewStateId("stable-editor-view")
857+
858+
await Promise.all([
859+
(provider1 as any).saveViewState("mode", "architect"),
860+
(provider2 as any).saveViewState("currentApiConfigName", "editor-profile"),
861+
])
862+
863+
expect(mockContext.globalState.get("viewStates" as any)).toMatchObject({
864+
"stable-sidebar-view": { mode: "architect" },
865+
"stable-editor-view": { currentApiConfigName: "editor-profile" },
866+
})
867+
868+
await provider1.dispose()
869+
await provider2.dispose()
870+
})
846871
})
847872

848873
describe("loadViewState", () => {
@@ -1042,6 +1067,32 @@ describe("ClineProvider - Parallel Mode Support", () => {
10421067

10431068
await provider.dispose()
10441069
})
1070+
1071+
it("should persist setValue mutations for view-local mode", async () => {
1072+
const provider = new ClineProvider(mockContext, mockOutputChannel, "sidebar", new ContextProxy(mockContext))
1073+
1074+
await (provider as any).setViewStateId("stable-sidebar-view")
1075+
await provider.setValue("mode" as any, "architect" as any)
1076+
1077+
expect(provider.contextProxy.getValue("viewStates" as any)).toMatchObject({
1078+
"stable-sidebar-view": { mode: "architect" },
1079+
})
1080+
1081+
await provider.dispose()
1082+
})
1083+
1084+
it("should persist setValues mutations for view-local API profile", async () => {
1085+
const provider = new ClineProvider(mockContext, mockOutputChannel, "sidebar", new ContextProxy(mockContext))
1086+
1087+
await (provider as any).setViewStateId("stable-sidebar-view")
1088+
await provider.setValues({ currentApiConfigName: "profile-from-set-values" } as any)
1089+
1090+
expect(provider.contextProxy.getValue("viewStates" as any)).toMatchObject({
1091+
"stable-sidebar-view": { currentApiConfigName: "profile-from-set-values" },
1092+
})
1093+
1094+
await provider.dispose()
1095+
})
10451096
})
10461097

10471098
describe("profile mutations", () => {

‎src/core/webview/__tests__/webviewMessageHandler.spec.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,7 @@ describe("webviewMessageHandler - webviewDidLaunch", () => {
234234
await webviewMessageHandler(mockClineProvider, { type: "webviewDidLaunch", viewStateId: "view-1" })
235235
await new Promise((resolve) => setImmediate(resolve))
236236

237+
expect((mockClineProvider as any).setViewStateId).toHaveBeenCalledWith("view-1")
237238
expect((mockClineProvider as any).providerSettingsManager.hasConfig).toHaveBeenCalledWith("view-local-profile")
238239
expect((mockClineProvider as any).providerSettingsManager.hasConfig).not.toHaveBeenCalledWith("shared-profile")
239240
})

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,14 @@ import {
1111
} from "@roo-code/types"
1212

1313
import { ExtensionStateContextProvider, useExtensionState, mergeExtensionState } from "../ExtensionStateContext"
14+
import { vscode } from "@src/utils/vscode"
15+
16+
vi.mock("@src/utils/vscode", () => ({
17+
vscode: {
18+
postMessage: vi.fn(),
19+
getViewStateId: vi.fn(() => "view-a"),
20+
},
21+
}))
1422

1523
const TestComponent = () => {
1624
const { allowedCommands, setAllowedCommands, soundEnabled, showRooIgnoredFiles, setShowRooIgnoredFiles } =
@@ -71,7 +79,80 @@ const ApiConfigTestComponent = () => {
7179
)
7280
}
7381

82+
const ViewLocalStateTestComponent = () => {
83+
const { mode, setMode, currentApiConfigName, setCurrentApiConfigName } = useExtensionState()
84+
85+
return (
86+
<div>
87+
<div data-testid="view-local-mode">{mode}</div>
88+
<div data-testid="view-local-api-config">{currentApiConfigName}</div>
89+
<button data-testid="set-local-mode" onClick={() => setMode("ask" as any)}>
90+
Set Local Mode
91+
</button>
92+
<button data-testid="set-local-api-config" onClick={() => setCurrentApiConfigName("local-profile")}>
93+
Set Local API Config
94+
</button>
95+
</div>
96+
)
97+
}
98+
7499
describe("ExtensionStateContext", () => {
100+
beforeEach(() => {
101+
vi.clearAllMocks()
102+
})
103+
104+
it("posts webviewDidLaunch with the stable viewStateId from vscode API", () => {
105+
render(
106+
<ExtensionStateContextProvider>
107+
<TestComponent />
108+
</ExtensionStateContextProvider>,
109+
)
110+
111+
expect(vscode.getViewStateId).toHaveBeenCalled()
112+
expect(vscode.postMessage).toHaveBeenCalledWith({ type: "webviewDidLaunch", viewStateId: "view-a" })
113+
})
114+
115+
it("reseeds view-local mode and API profile from a new state payload after local edits", () => {
116+
render(
117+
<ExtensionStateContextProvider>
118+
<ViewLocalStateTestComponent />
119+
</ExtensionStateContextProvider>,
120+
)
121+
122+
act(() => {
123+
window.dispatchEvent(
124+
new MessageEvent("message", {
125+
data: {
126+
type: "state",
127+
state: { mode: "code", currentApiConfigName: "profile-a", apiConfiguration: {} },
128+
},
129+
}),
130+
)
131+
})
132+
expect(screen.getByTestId("view-local-mode")).toHaveTextContent("code")
133+
expect(screen.getByTestId("view-local-api-config")).toHaveTextContent("profile-a")
134+
135+
act(() => {
136+
screen.getByTestId("set-local-mode").click()
137+
screen.getByTestId("set-local-api-config").click()
138+
})
139+
expect(screen.getByTestId("view-local-mode")).toHaveTextContent("ask")
140+
expect(screen.getByTestId("view-local-api-config")).toHaveTextContent("local-profile")
141+
142+
act(() => {
143+
window.dispatchEvent(
144+
new MessageEvent("message", {
145+
data: {
146+
type: "state",
147+
state: { mode: "architect", currentApiConfigName: "profile-b", apiConfiguration: {} },
148+
},
149+
}),
150+
)
151+
})
152+
expect(screen.getByTestId("view-local-mode")).toHaveTextContent("architect")
153+
expect(screen.getByTestId("view-local-api-config")).toHaveTextContent("profile-b")
154+
})
155+
75156
it("initializes with empty allowedCommands array", () => {
76157
render(
77158
<ExtensionStateContextProvider>

0 commit comments

Comments
 (0)