diff --git a/src/core/__tests__/contextProxy.test.ts b/src/core/__tests__/contextProxy.test.ts index 794cd91497..e6f1bfc9ca 100644 --- a/src/core/__tests__/contextProxy.test.ts +++ b/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) - }) }) }) diff --git a/src/core/contextProxy.ts b/src/core/contextProxy.ts index e4672ae225..7c429c86cf 100644 --- a/src/core/contextProxy.ts +++ b/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 - private pendingSecretChanges: Map - private disposed: boolean + private stateCache: Map + private secretCache: Map 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) + .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(key: string): Promise - async getGlobalState(key: string, defaultValue: T): Promise - async getGlobalState(key: string, defaultValue?: T): Promise { - // 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(key, defaultValue as T) + getGlobalState(key: string): T | undefined + getGlobalState(key: string, defaultValue: T): T + getGlobalState(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(key: string, value: T): Promise { - 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(key: string, value: T): Thenable { + this.stateCache.set(key, value) + return this.originalContext.globalState.update(key, value) } - // Secret storage methods - async getSecret(key: string): Promise { - // 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 { - if (this.disposed) { - throw new Error("Cannot store secret on disposed context") + storeSecret(key: string, value?: string): Thenable { + // 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 { - 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 { - 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 } } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index ae53e58df9..748ba2525b 100644 --- a/src/core/webview/ClineProvider.ts +++ b/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) { diff --git a/src/core/webview/__tests__/ClineProvider.test.ts b/src/core/webview/__tests__/ClineProvider.test.ts index 20778b8802..9463be25b7 100644 --- a/src/core/webview/__tests__/ClineProvider.test.ts +++ b/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() - }) }) diff --git a/src/shared/api.ts b/src/shared/api.ts index 7c5c65fe90..981fcf8d76 100644 --- a/src/shared/api.ts +++ b/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",