Skip to content

Commit 82f9f23

Browse files
easonliang28easonliang28
committed
fix(webview): persist view-local state safely
1 parent 8885969 commit 82f9f23

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
@@ -428,51 +429,63 @@ export class ClineProvider
428429
}
429430
}
430431

431-
private getPersistedViewStates(): Record<string, PersistedViewState> {
432-
const viewStates = this.contextProxy.getValue("viewStates")
432+
private getPersistedViewStates(options: { fresh?: boolean } = {}): Record<string, PersistedViewState> {
433+
const viewStates = options.fresh
434+
? this.context.globalState.get<GlobalState["viewStates"]>("viewStates")
435+
: this.contextProxy.getValue("viewStates")
433436

434437
if (!viewStates || typeof viewStates !== "object" || Array.isArray(viewStates)) {
435438
return {}
436439
}
437440

438-
return viewStates
441+
return { ...viewStates }
439442
}
440443

441444
private async savePersistedViewState(values: Partial<PersistedViewState>): Promise<void> {
442-
const states = this.getPersistedViewStates()
443-
const current = states[this.viewStateId] ?? {}
444-
const next: PersistedViewState = { ...current }
445+
const write = ClineProvider.persistedViewStateWriteQueue.then(async () => {
446+
const states = this.getPersistedViewStates({ fresh: true })
447+
const current = states[this.viewStateId] ?? {}
448+
const next: PersistedViewState = { ...current }
449+
450+
if ("mode" in values) {
451+
if (values.mode === undefined || values.mode === null) {
452+
delete next.mode
453+
} else {
454+
next.mode = values.mode
455+
}
456+
}
445457

446-
if ("mode" in values) {
447-
if (values.mode === undefined || values.mode === null) {
448-
delete next.mode
449-
} else {
450-
next.mode = values.mode
458+
if ("currentApiConfigName" in values) {
459+
if (values.currentApiConfigName === undefined || values.currentApiConfigName === null) {
460+
delete next.currentApiConfigName
461+
} else {
462+
next.currentApiConfigName = values.currentApiConfigName
463+
}
451464
}
452-
}
453465

454-
if ("currentApiConfigName" in values) {
455-
if (values.currentApiConfigName === undefined || values.currentApiConfigName === null) {
456-
delete next.currentApiConfigName
466+
if (!next.mode && !next.currentApiConfigName) {
467+
delete states[this.viewStateId]
457468
} else {
458-
next.currentApiConfigName = values.currentApiConfigName
469+
next.updatedAt = values.updatedAt ?? Date.now()
470+
states[this.viewStateId] = next
459471
}
460-
}
461472

462-
if (!next.mode && !next.currentApiConfigName) {
463-
delete states[this.viewStateId]
464-
} else {
465-
next.updatedAt = values.updatedAt ?? Date.now()
466-
states[this.viewStateId] = next
467-
}
473+
await this.contextProxy.setValue("viewStates", this.prunePersistedViewStates(states))
474+
})
468475

469-
await this.contextProxy.setValue("viewStates", this.prunePersistedViewStates(states))
476+
ClineProvider.persistedViewStateWriteQueue = write.catch(() => {})
477+
await write
470478
}
471479

472480
private async clearPersistedViewState(viewStateId = this.viewStateId): Promise<void> {
473-
const states = this.getPersistedViewStates()
474-
delete states[viewStateId]
475-
await this.contextProxy.setValue("viewStates", states)
481+
const write = ClineProvider.persistedViewStateWriteQueue.then(async () => {
482+
const states = this.getPersistedViewStates({ fresh: true })
483+
delete states[viewStateId]
484+
await this.contextProxy.setValue("viewStates", states)
485+
})
486+
487+
ClineProvider.persistedViewStateWriteQueue = write.catch(() => {})
488+
await write
476489
}
477490

478491
private prunePersistedViewStates(states: Record<string, PersistedViewState>): Record<string, PersistedViewState> {
@@ -3008,6 +3021,7 @@ export class ClineProvider
30083021
public async setValue<K extends keyof RooCodeSettings>(key: K, value: RooCodeSettings[K]) {
30093022
await this.contextProxy.setValue(key, value)
30103023
this._updateViewLocalStateFromMutation({ [key]: value })
3024+
await this._persistViewLocalStateFromMutation({ [key]: value })
30113025
}
30123026

30133027
public getValue<K extends keyof RooCodeSettings>(key: K) {
@@ -3021,6 +3035,7 @@ export class ClineProvider
30213035
public async setValues(values: RooCodeSettings) {
30223036
await this.contextProxy.setValues(values)
30233037
this._updateViewLocalStateFromMutation(values)
3038+
await this._persistViewLocalStateFromMutation(values)
30243039
}
30253040

30263041
/**
@@ -3070,6 +3085,24 @@ export class ClineProvider
30703085
}
30713086
}
30723087

3088+
private async _persistViewLocalStateFromMutation(
3089+
values: Partial<RooCodeSettings> & Partial<ExtensionState>,
3090+
): Promise<void> {
3091+
const persistedValues: Partial<PersistedViewState> = {}
3092+
3093+
if ("mode" in values) {
3094+
persistedValues.mode = values.mode as PersistedViewState["mode"]
3095+
}
3096+
3097+
if ("currentApiConfigName" in values) {
3098+
persistedValues.currentApiConfigName = values.currentApiConfigName
3099+
}
3100+
3101+
if ("mode" in persistedValues || "currentApiConfigName" in persistedValues) {
3102+
await this.savePersistedViewState(persistedValues)
3103+
}
3104+
}
3105+
30733106
/**
30743107
* Clear view-local state cache so that getState() falls back to ContextProxy defaults.
30753108
*/

‎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)