update new way to manage state

src/core/__tests__/contextProxy.test.ts

@@ -1,6 +1,7 @@
 import * as vscode from "vscode"
 import { ContextProxy } from "../contextProxy"
 import { logger } from "../../utils/logging"
+import { GLOBAL_STATE_KEYS, SECRET_KEYS } from "../../shared/globalState"
 
 // Mock the logger
 jest.mock("../../utils/logging", () => ({
@@ -12,6 +13,12 @@ jest.mock("../../utils/logging", () => ({
 	},
 }))
 
+// Mock shared/globalState
+jest.mock("../../shared/globalState", () => ({
+	GLOBAL_STATE_KEYS: ["apiProvider", "apiModelId", "mode"],
+	SECRET_KEYS: ["apiKey", "openAiApiKey"],
+}))
+
 // Mock VSCode API
 jest.mock("vscode", () => ({
 	Uri: {
@@ -42,7 +49,7 @@ describe("ContextProxy", () => {
 
 		// Mock secrets
 		mockSecrets = {
-			get: jest.fn(),
+			get: jest.fn().mockResolvedValue("test-secret"),
 			store: jest.fn().mockResolvedValue(undefined),
 			delete: jest.fn().mockResolvedValue(undefined),
 		}
@@ -74,98 +81,80 @@ describe("ContextProxy", () => {
 		})
 	})
 
-	describe("getGlobalState", () => {
-		it("should return pending change when it exists", async () => {
-			// Set up a pending change
-			await proxy.updateGlobalState("test-key", "new-value")
-
-			// Should return the pending value
-			const result = await proxy.getGlobalState("test-key")
-			expect(result).toBe("new-value")
-
-			// Original context should not be called
-			expect(mockGlobalState.get).not.toHaveBeenCalled()
+	describe("constructor", () => {
+		it("should initialize state cache with all global state keys", () => {
+			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length)
+			for (const key of GLOBAL_STATE_KEYS) {
+				expect(mockGlobalState.get).toHaveBeenCalledWith(key)
+			}
 		})
 
-		it("should fall back to original context when no pending change exists", async () => {
-			// Set up original context value
-			mockGlobalState.get.mockReturnValue("original-value")
+		it("should initialize secret cache with all secret keys", () => {
+			expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_KEYS.length)
+			for (const key of SECRET_KEYS) {
+				expect(mockSecrets.get).toHaveBeenCalledWith(key)
+			}
+		})
+	})
 
-			// Should get from original context
-			const result = await proxy.getGlobalState("test-key")
-			expect(result).toBe("original-value")
-			expect(mockGlobalState.get).toHaveBeenCalledWith("test-key", undefined)
+	describe("getGlobalState", () => {
+		it("should return value from cache when it exists", async () => {
+			// Manually set a value in the cache
+			await proxy.updateGlobalState("test-key", "cached-value")
+
+			// Should return the cached value
+			const result = proxy.getGlobalState("test-key")
+			expect(result).toBe("cached-value")
+
+			// Original context should be called once during updateGlobalState
+			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length) // Only from initialization
 		})
 
 		it("should handle default values correctly", async () => {
-			// No value in either pending or original
-			mockGlobalState.get.mockImplementation((key: string, defaultValue: any) => defaultValue)
-
-			// Should return the default value
-			const result = await proxy.getGlobalState("test-key", "default-value")
+			// No value in cache
+			const result = proxy.getGlobalState("unknown-key", "default-value")
 			expect(result).toBe("default-value")
 		})
 	})
 
 	describe("updateGlobalState", () => {
-		it("should buffer changes without calling original context", async () => {
+		it("should update state directly in original context", async () => {
 			await proxy.updateGlobalState("test-key", "new-value")
 
 			// Should have called logger.debug
-			expect(logger.debug).toHaveBeenCalledWith(expect.stringContaining("buffering state update"))
+			expect(logger.debug).toHaveBeenCalledWith(expect.stringContaining("updating state for key"))
 
-			// Should not have called original context
-			expect(mockGlobalState.update).not.toHaveBeenCalled()
+			// Should have called original context
+			expect(mockGlobalState.update).toHaveBeenCalledWith("test-key", "new-value")
 
-			// Should have stored the value in pendingStateChanges
+			// Should have stored the value in cache
 			const storedValue = await proxy.getGlobalState("test-key")
 			expect(storedValue).toBe("new-value")
 		})
-
-		it("should throw an error when context is disposed", async () => {
-			await proxy.dispose()
-
-			await expect(proxy.updateGlobalState("test-key", "new-value")).rejects.toThrow(
-				"Cannot update state on disposed context",
-			)
-		})
 	})
 
 	describe("getSecret", () => {
-		it("should return pending secret when it exists", async () => {
-			// Set up a pending secret
-			await proxy.storeSecret("api-key", "secret123")
+		it("should return value from cache when it exists", async () => {
+			// Manually set a value in the cache
+			await proxy.storeSecret("api-key", "cached-secret")
 
-			// Should return the pending value
-			const result = await proxy.getSecret("api-key")
-			expect(result).toBe("secret123")
-
-			// Original context should not be called
-			expect(mockSecrets.get).not.toHaveBeenCalled()
-		})
-
-		it("should fall back to original context when no pending secret exists", async () => {
-			// Set up original context value
-			mockSecrets.get.mockResolvedValue("original-secret")
-
-			// Should get from original context
-			const result = await proxy.getSecret("api-key")
-			expect(result).toBe("original-secret")
-			expect(mockSecrets.get).toHaveBeenCalledWith("api-key")
+			// Should return the cached value
+			const result = proxy.getSecret("api-key")
+			expect(result).toBe("cached-secret")
 		})
 	})
 
 	describe("storeSecret", () => {
-		it("should buffer secret changes without calling original context", async () => {
+		it("should store secret directly in original context", async () => {
 			await proxy.storeSecret("api-key", "new-secret")
 
 			// Should have called logger.debug
-			expect(logger.debug).toHaveBeenCalledWith(expect.stringContaining("buffering secret update"))
+			expect(logger.debug).toHaveBeenCalledWith(expect.stringContaining("storing secret for key"))
 
-			// Should not have called original context
-			expect(mockSecrets.store).not.toHaveBeenCalled()
+			// Should have called original context
+			expect(mockSecrets.store).toHaveBeenCalledWith("api-key", "new-secret")
 
-			// Should have stored the value in pendingSecretChanges
+			// Should have stored the value in cache
 			const storedValue = await proxy.getSecret("api-key")
 			expect(storedValue).toBe("new-secret")
 		})
@@ -173,110 +162,12 @@ describe("ContextProxy", () => {
 		it("should handle undefined value for secret deletion", async () => {
 			await proxy.storeSecret("api-key", undefined)
 
-			// Should have stored undefined in pendingSecretChanges
+			// Should have called delete on original context
+			expect(mockSecrets.delete).toHaveBeenCalledWith("api-key")
+
+			// Should have stored undefined in cache
 			const storedValue = await proxy.getSecret("api-key")
 			expect(storedValue).toBeUndefined()
 		})
-
-		it("should throw an error when context is disposed", async () => {
-			await proxy.dispose()
-
-			await expect(proxy.storeSecret("api-key", "new-secret")).rejects.toThrow(
-				"Cannot store secret on disposed context",
-			)
-		})
-	})
-
-	describe("saveChanges", () => {
-		it("should apply state changes to original context", async () => {
-			// Set up pending changes
-			await proxy.updateGlobalState("key1", "value1")
-			await proxy.updateGlobalState("key2", "value2")
-
-			// Save changes
-			await proxy.saveChanges()
-
-			// Should have called update on original context
-			expect(mockGlobalState.update).toHaveBeenCalledTimes(2)
-			expect(mockGlobalState.update).toHaveBeenCalledWith("key1", "value1")
-			expect(mockGlobalState.update).toHaveBeenCalledWith("key2", "value2")
-
-			// Should have cleared pending changes
-			expect(proxy.hasPendingChanges()).toBe(false)
-		})
-
-		it("should apply secret changes to original context", async () => {
-			// Set up pending changes
-			await proxy.storeSecret("secret1", "value1")
-			await proxy.storeSecret("secret2", undefined)
-
-			// Save changes
-			await proxy.saveChanges()
-
-			// Should have called store and delete on original context
-			expect(mockSecrets.store).toHaveBeenCalledTimes(1)
-			expect(mockSecrets.store).toHaveBeenCalledWith("secret1", "value1")
-			expect(mockSecrets.delete).toHaveBeenCalledTimes(1)
-			expect(mockSecrets.delete).toHaveBeenCalledWith("secret2")
-
-			// Should have cleared pending changes
-			expect(proxy.hasPendingChanges()).toBe(false)
-		})
-
-		it("should do nothing when there are no pending changes", async () => {
-			await proxy.saveChanges()
-
-			expect(mockGlobalState.update).not.toHaveBeenCalled()
-			expect(mockSecrets.store).not.toHaveBeenCalled()
-			expect(mockSecrets.delete).not.toHaveBeenCalled()
-		})
-
-		it("should throw an error when context is disposed", async () => {
-			await proxy.dispose()
-
-			await expect(proxy.saveChanges()).rejects.toThrow("Cannot save changes on disposed context")
-		})
-	})
-
-	describe("dispose", () => {
-		it("should save pending changes to original context", async () => {
-			// Set up pending changes
-			await proxy.updateGlobalState("key1", "value1")
-			await proxy.storeSecret("secret1", "value1")
-
-			// Dispose
-			await proxy.dispose()
-
-			// Should have saved changes
-			expect(mockGlobalState.update).toHaveBeenCalledWith("key1", "value1")
-			expect(mockSecrets.store).toHaveBeenCalledWith("secret1", "value1")
-
-			// Should be marked as disposed
-			expect(proxy.hasPendingChanges()).toBe(false)
-		})
-	})
-
-	describe("hasPendingChanges", () => {
-		it("should return false when no changes are pending", () => {
-			expect(proxy.hasPendingChanges()).toBe(false)
-		})
-
-		it("should return true when state changes are pending", async () => {
-			await proxy.updateGlobalState("key", "value")
-			expect(proxy.hasPendingChanges()).toBe(true)
-		})
-
-		it("should return true when secret changes are pending", async () => {
-			await proxy.storeSecret("key", "value")
-			expect(proxy.hasPendingChanges()).toBe(true)
-		})
-
-		it("should return false after changes are saved", async () => {
-			await proxy.updateGlobalState("key", "value")
-			expect(proxy.hasPendingChanges()).toBe(true)
-
-			await proxy.saveChanges()
-			expect(proxy.hasPendingChanges()).toBe(false)
-		})
 	})
 })

src/core/contextProxy.ts

@@ -1,25 +1,53 @@
 import * as vscode from "vscode"
 import { logger } from "../utils/logging"
+import { GLOBAL_STATE_KEYS, SECRET_KEYS } from "../shared/globalState"
 
-/**
- * A proxy class for vscode.ExtensionContext that buffers state changes
- * and only commits them when explicitly requested or during disposal.
- */
 export class ContextProxy {
 	private readonly originalContext: vscode.ExtensionContext
-	private pendingStateChanges: Map<string, any>
-	private pendingSecretChanges: Map<string, string | undefined>
-	private disposed: boolean
+	private stateCache: Map<string, any>
+	private secretCache: Map<string, string | undefined>
 
 	constructor(context: vscode.ExtensionContext) {
+		// Initialize properties first
 		this.originalContext = context
-		this.pendingStateChanges = new Map()
-		this.pendingSecretChanges = new Map()
-		this.disposed = false
+		this.stateCache = new Map()
+		this.secretCache = new Map()
+
+		// Initialize state cache with all defined global state keys
+		this.initializeStateCache()
+
+		// Initialize secret cache with all defined secret keys
+		this.initializeSecretCache()
+
 		logger.debug("ContextProxy created")
 	}
 
-	// Read-only pass-through properties
+	// Helper method to initialize state cache
+	private initializeStateCache(): void {
+		for (const key of GLOBAL_STATE_KEYS) {
+			try {
+				const value = this.originalContext.globalState.get(key)
+				this.stateCache.set(key, value)
+			} catch (error) {
+				logger.error(`Error loading global ${key}: ${error instanceof Error ? error.message : String(error)}`)
+			}
+		}
+	}
+
+	// Helper method to initialize secret cache
+	private initializeSecretCache(): void {
+		for (const key of SECRET_KEYS) {
+			// Get actual value and update cache when promise resolves
+			;(this.originalContext.secrets.get(key) as Promise<string | undefined>)
+				.then((value) => {
+					this.secretCache.set(key, value)
+				})
+				.catch((error: Error) => {
+					logger.error(`Error loading secret ${key}: ${error.message}`)
+				})
+		}
+	}
+
 	get extensionUri(): vscode.Uri {
 		return this.originalContext.extensionUri
 	}
@@ -39,85 +67,30 @@ export class ContextProxy {
 		return this.originalContext.extensionMode
 	}
 
-	// State management methods
-	async getGlobalState<T>(key: string): Promise<T | undefined>
-	async getGlobalState<T>(key: string, defaultValue: T): Promise<T>
-	async getGlobalState<T>(key: string, defaultValue?: T): Promise<T | undefined> {
-		// Check pending changes first
-		if (this.pendingStateChanges.has(key)) {
-			const value = this.pendingStateChanges.get(key) as T | undefined
-			return value !== undefined ? value : (defaultValue as T | undefined)
-		}
-		// Fall back to original context
-		return this.originalContext.globalState.get<T>(key, defaultValue as T)
+	getGlobalState<T>(key: string): T | undefined
+	getGlobalState<T>(key: string, defaultValue: T): T
+	getGlobalState<T>(key: string, defaultValue?: T): T | undefined {
+		const value = this.stateCache.get(key) as T | undefined
+		return value !== undefined ? value : (defaultValue as T | undefined)
 	}
 
-	async updateGlobalState<T>(key: string, value: T): Promise<void> {
-		if (this.disposed) {
-			throw new Error("Cannot update state on disposed context")
-		}
-		logger.debug(`ContextProxy: buffering state update for key "${key}"`)
-		this.pendingStateChanges.set(key, value)
+	updateGlobalState<T>(key: string, value: T): Thenable<void> {
+		this.stateCache.set(key, value)
+		return this.originalContext.globalState.update(key, value)
 	}
 
-	// Secret storage methods
-	async getSecret(key: string): Promise<string | undefined> {
-		// Check pending changes first
-		if (this.pendingSecretChanges.has(key)) {
-			return this.pendingSecretChanges.get(key)
-		}
-		// Fall back to original context
-		return this.originalContext.secrets.get(key)
+	getSecret(key: string): string | undefined {
+		return this.secretCache.get(key)
 	}
 
-	async storeSecret(key: string, value?: string): Promise<void> {
-		if (this.disposed) {
-			throw new Error("Cannot store secret on disposed context")
+	storeSecret(key: string, value?: string): Thenable<void> {
+		// Update cache
+		this.secretCache.set(key, value)
+		// Write directly to context
+		if (value === undefined) {
+			return this.originalContext.secrets.delete(key)
+		} else {
+			return this.originalContext.secrets.store(key, value)
 		}
-		logger.debug(`ContextProxy: buffering secret update for key "${key}"`)
-		this.pendingSecretChanges.set(key, value)
-	}
-
-	// Save pending changes to actual context
-	async saveChanges(): Promise<void> {
-		if (this.disposed) {
-			throw new Error("Cannot save changes on disposed context")
-		}
-
-		// Apply state changes
-		if (this.pendingStateChanges.size > 0) {
-			logger.debug(`ContextProxy: applying ${this.pendingStateChanges.size} buffered state changes`)
-			for (const [key, value] of this.pendingStateChanges.entries()) {
-				await this.originalContext.globalState.update(key, value)
-			}
-			this.pendingStateChanges.clear()
-		}
-
-		// Apply secret changes
-		if (this.pendingSecretChanges.size > 0) {
-			logger.debug(`ContextProxy: applying ${this.pendingSecretChanges.size} buffered secret changes`)
-			for (const [key, value] of this.pendingSecretChanges.entries()) {
-				if (value === undefined) {
-					await this.originalContext.secrets.delete(key)
-				} else {
-					await this.originalContext.secrets.store(key, value)
-				}
-			}
-			this.pendingSecretChanges.clear()
-		}
-	}
-
-	// Called when the provider is disposing
-	async dispose(): Promise<void> {
-		if (!this.disposed) {
-			logger.debug("ContextProxy: disposing and saving pending changes")
-			await this.saveChanges()
-			this.disposed = true
-		}
-	}
-
-	// Method to check if there are pending changes
-	hasPendingChanges(): boolean {
-		return this.pendingStateChanges.size > 0 || this.pendingSecretChanges.size > 0
 	}
 }

src/core/webview/ClineProvider.ts

@@ -16,7 +16,7 @@ import { SecretKey, GlobalStateKey, SECRET_KEYS, GLOBAL_STATE_KEYS } from "../..
 import { HistoryItem } from "../../shared/HistoryItem"
 import { ApiConfigMeta, ExtensionMessage } from "../../shared/ExtensionMessage"
 import { checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, WebviewMessage } from "../../shared/WebviewMessage"
-import { Mode, CustomModePrompts, PromptComponent, defaultModeSlug } from "../../shared/modes"
+import { Mode, CustomModePrompts, PromptComponent, defaultModeSlug, ModeConfig } from "../../shared/modes"
 import { checkExistKey } from "../../shared/checkExistApiConfig"
 import { EXPERIMENT_IDS, experiments as Experiments, experimentDefault, ExperimentId } from "../../shared/experiments"
 import { downloadTask } from "../../integrations/misc/export-markdown"
@@ -119,8 +119,6 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		this.customModesManager?.dispose()
 		this.outputChannel.appendLine("Disposed all disposables")
 		// Dispose the context proxy to commit any pending changes
-		await this.contextProxy.dispose()
-		this.outputChannel.appendLine("Disposed context proxy")
 		ClineProvider.activeInstances.delete(this)
 
 		// Unregister from McpServerManager
@@ -2082,22 +2080,26 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		// Add promise for custom modes which is handled separately
 		const customModesPromise = this.customModesManager.getCustomModes()
 
-		// Wait for all promises to resolve
-		const [stateResults, secretResults, customModes] = await Promise.all([
-			Promise.all(statePromises),
-			Promise.all(secretPromises),
+		let idx = 0
+		const secretValuesArray = await Promise.all([
+			...statePromises,
+			...secretPromises,
 			customModesPromise,
 		])
 
 		// Populate stateValues and secretValues
-		GLOBAL_STATE_KEYS.forEach((key, index) => {
-			stateValues[key] = stateResults[index]
+		GLOBAL_STATE_KEYS.forEach((key, _) => {
+			stateValues[key] = secretValuesArray[idx]
+			idx = idx + 1
 		})
 
 		SECRET_KEYS.forEach((key, index) => {
-			secretValues[key] = secretResults[index]
+			secretValues[key] = secretValuesArray[idx]
+			idx = idx + 1
 		})
 
+		let customModes = secretValuesArray[idx] as ModeConfig[] | undefined
+
 		// Determine apiProvider with the same logic as before
 		let apiProvider: ApiProvider
 		if (stateValues.apiProvider) {
@@ -2219,12 +2221,6 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 	async updateGlobalState(key: GlobalStateKey, value: any) {
 		this.outputChannel.appendLine(`Updating global state: ${key}`)
 		await this.contextProxy.updateGlobalState(key, value)
-
-		// // If we have a lot of pending changes, consider saving them periodically
-		// if (this.contextProxy.hasPendingChanges() && Math.random() < 0.1) { // 10% chance to save changes
-		// 	this.outputChannel.appendLine("Periodically flushing context state changes")
-		// 	await this.contextProxy.saveChanges()
-		// }
 	}
 
 	async getGlobalState(key: GlobalStateKey) {
@@ -2256,13 +2252,13 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		}
 
 		for (const key of this.context.globalState.keys()) {
-			// Still using original context for listing keys
 			await this.contextProxy.updateGlobalState(key, undefined)
 		}
 
 		for (const key of SECRET_KEYS) {
 			await this.storeSecret(key, undefined)
 		}
+
 		await this.configManager.resetAllConfigs()
 		await this.customModesManager.resetCustomModes()
 		if (this.cline) {

src/core/webview/__tests__/ClineProvider.test.ts

@@ -1303,17 +1303,6 @@ describe("ClineProvider", () => {
 			// Verify state was posted to webview
 			expect(mockPostMessage).toHaveBeenCalledWith(expect.objectContaining({ type: "state" }))
 		})
-
-		test("disposes the contextProxy when provider is disposed", async () => {
-			// Setup mock Cline instance
-			const mockCline = {
-				abortTask: jest.fn(),
-			}
-			// @ts-ignore - accessing private property for testing
-			provider.cline = mockCline
-			await provider.dispose()
-			expect(mockContextProxy.dispose).toHaveBeenCalled()
-		})
 	})
 
 	describe("updateCustomMode", () => {
@@ -1602,15 +1591,4 @@ describe("ContextProxy integration", () => {
 		expect(mockContextProxy.updateGlobalState).toBeDefined()
 		expect(mockContextProxy.storeSecret).toBeDefined()
 	})
-
-	test("contextProxy is properly disposed", async () => {
-		// Setup mock Cline instance
-		const mockCline = {
-			abortTask: jest.fn(),
-		}
-		// @ts-ignore - accessing private property for testing
-		provider.cline = mockCline
-		await provider.dispose()
-		expect(mockContextProxy.dispose).toHaveBeenCalled()
-	})
 })

src/shared/api.ts

@@ -108,7 +108,7 @@ export const API_CONFIG_KEYS: GlobalStateKey[] = [
 	"lmStudioBaseUrl",
 	"lmStudioDraftModelId",
 	"lmStudioSpeculativeDecodingEnabled",
-	"mistralCodestralUrl", // New option for Codestral URL
+	"mistralCodestralUrl",
 	"azureApiVersion",
 	"openRouterUseMiddleOutTransform",
 	"openAiStreamingEnabled",