Merge pull request #1235 from samhvw8/feat/context-proxy

Feat ContextProxy to improve state management

Author: mrubens

src/core/__tests__/contextProxy.test.ts

@@ -0,0 +1,243 @@
+import * as vscode from "vscode"
+import { ContextProxy } from "../contextProxy"
+import { logger } from "../../utils/logging"
+import { GLOBAL_STATE_KEYS, SECRET_KEYS } from "../../shared/globalState"
+import { ApiConfiguration } from "../../shared/api"
+
+// Mock shared/globalState
+jest.mock("../../shared/globalState", () => ({
+	GLOBAL_STATE_KEYS: ["apiProvider", "apiModelId", "mode"],
+	SECRET_KEYS: ["apiKey", "openAiApiKey"],
+	GlobalStateKey: {},
+	SecretKey: {},
+}))
+
+// Mock shared/api
+jest.mock("../../shared/api", () => ({
+	API_CONFIG_KEYS: ["apiProvider", "apiModelId"],
+	ApiConfiguration: {},
+}))
+
+// Mock VSCode API
+jest.mock("vscode", () => ({
+	Uri: {
+		file: jest.fn((path) => ({ path })),
+	},
+	ExtensionMode: {
+		Development: 1,
+		Production: 2,
+		Test: 3,
+	},
+}))
+
+describe("ContextProxy", () => {
+	let proxy: ContextProxy
+	let mockContext: any
+	let mockGlobalState: any
+	let mockSecrets: any
+
+	beforeEach(() => {
+		// Reset mocks
+		jest.clearAllMocks()
+
+		// Mock globalState
+		mockGlobalState = {
+			get: jest.fn(),
+			update: jest.fn().mockResolvedValue(undefined),
+		}
+
+		// Mock secrets
+		mockSecrets = {
+			get: jest.fn().mockResolvedValue("test-secret"),
+			store: jest.fn().mockResolvedValue(undefined),
+			delete: jest.fn().mockResolvedValue(undefined),
+		}
+
+		// Mock the extension context
+		mockContext = {
+			globalState: mockGlobalState,
+			secrets: mockSecrets,
+			extensionUri: { path: "/test/extension" },
+			extensionPath: "/test/extension",
+			globalStorageUri: { path: "/test/storage" },
+			logUri: { path: "/test/logs" },
+			extension: { packageJSON: { version: "1.0.0" } },
+			extensionMode: vscode.ExtensionMode.Development,
+		}
+
+		// Create proxy instance
+		proxy = new ContextProxy(mockContext)
+	})
+
+	describe("read-only pass-through properties", () => {
+		it("should return extension properties from the original context", () => {
+			expect(proxy.extensionUri).toBe(mockContext.extensionUri)
+			expect(proxy.extensionPath).toBe(mockContext.extensionPath)
+			expect(proxy.globalStorageUri).toBe(mockContext.globalStorageUri)
+			expect(proxy.logUri).toBe(mockContext.logUri)
+			expect(proxy.extension).toBe(mockContext.extension)
+			expect(proxy.extensionMode).toBe(mockContext.extensionMode)
+		})
+	})
+
+	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 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)
+			}
+		})
+	})
+
+	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 cache
+			const result = proxy.getGlobalState("unknown-key", "default-value")
+			expect(result).toBe("default-value")
+		})
+	})
+
+	describe("updateGlobalState", () => {
+		it("should update state directly in original context", async () => {
+			await proxy.updateGlobalState("test-key", "new-value")
+
+			// Should have called original context
+			expect(mockGlobalState.update).toHaveBeenCalledWith("test-key", "new-value")
+
+			// Should have stored the value in cache
+			const storedValue = await proxy.getGlobalState("test-key")
+			expect(storedValue).toBe("new-value")
+		})
+	})
+
+	describe("getSecret", () => {
+		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 cached value
+			const result = proxy.getSecret("api-key")
+			expect(result).toBe("cached-secret")
+		})
+	})
+
+	describe("storeSecret", () => {
+		it("should store secret directly in original context", async () => {
+			await proxy.storeSecret("api-key", "new-secret")
+
+			// Should have called original context
+			expect(mockSecrets.store).toHaveBeenCalledWith("api-key", "new-secret")
+
+			// Should have stored the value in cache
+			const storedValue = await proxy.getSecret("api-key")
+			expect(storedValue).toBe("new-secret")
+		})
+
+		it("should handle undefined value for secret deletion", async () => {
+			await proxy.storeSecret("api-key", undefined)
+
+			// 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()
+		})
+
+		describe("getApiConfiguration", () => {
+			it("should combine global state and secrets into a single ApiConfiguration object", async () => {
+				// Mock data in state cache
+				await proxy.updateGlobalState("apiProvider", "anthropic")
+				await proxy.updateGlobalState("apiModelId", "test-model")
+				// Mock data in secrets cache
+				await proxy.storeSecret("apiKey", "test-api-key")
+
+				const config = proxy.getApiConfiguration()
+
+				// Should contain values from global state
+				expect(config.apiProvider).toBe("anthropic")
+				expect(config.apiModelId).toBe("test-model")
+				// Should contain values from secrets
+				expect(config.apiKey).toBe("test-api-key")
+			})
+
+			it("should handle special case for apiProvider defaulting", async () => {
+				// Clear apiProvider but set apiKey
+				await proxy.updateGlobalState("apiProvider", undefined)
+				await proxy.storeSecret("apiKey", "test-api-key")
+
+				const config = proxy.getApiConfiguration()
+
+				// Should default to anthropic when apiKey exists
+				expect(config.apiProvider).toBe("anthropic")
+
+				// Clear both apiProvider and apiKey
+				await proxy.updateGlobalState("apiProvider", undefined)
+				await proxy.storeSecret("apiKey", undefined)
+
+				const configWithoutKey = proxy.getApiConfiguration()
+
+				// Should default to openrouter when no apiKey exists
+				expect(configWithoutKey.apiProvider).toBe("openrouter")
+			})
+		})
+
+		describe("updateApiConfiguration", () => {
+			it("should update both global state and secrets", async () => {
+				const apiConfig: ApiConfiguration = {
+					apiProvider: "anthropic",
+					apiModelId: "claude-latest",
+					apiKey: "test-api-key",
+				}
+
+				await proxy.updateApiConfiguration(apiConfig)
+
+				// Should update global state
+				expect(mockGlobalState.update).toHaveBeenCalledWith("apiProvider", "anthropic")
+				expect(mockGlobalState.update).toHaveBeenCalledWith("apiModelId", "claude-latest")
+				// Should update secrets
+				expect(mockSecrets.store).toHaveBeenCalledWith("apiKey", "test-api-key")
+
+				// Check that values are in cache
+				expect(proxy.getGlobalState("apiProvider")).toBe("anthropic")
+				expect(proxy.getGlobalState("apiModelId")).toBe("claude-latest")
+				expect(proxy.getSecret("apiKey")).toBe("test-api-key")
+			})
+
+			it("should ignore keys that aren't in either GLOBAL_STATE_KEYS or SECRET_KEYS", async () => {
+				// Use type assertion to add an invalid key
+				const apiConfig = {
+					apiProvider: "anthropic",
+					invalidKey: "should be ignored",
+				} as ApiConfiguration & { invalidKey: string }
+
+				await proxy.updateApiConfiguration(apiConfig)
+
+				// Should update keys in GLOBAL_STATE_KEYS
+				expect(mockGlobalState.update).toHaveBeenCalledWith("apiProvider", "anthropic")
+				// Should not call update/store for invalid keys
+				expect(mockGlobalState.update).not.toHaveBeenCalledWith("invalidKey", expect.anything())
+				expect(mockSecrets.store).not.toHaveBeenCalledWith("invalidKey", expect.anything())
+			})
+		})
+	})
+})

src/core/contextProxy.ts

@@ -0,0 +1,154 @@
+import * as vscode from "vscode"
+import { logger } from "../utils/logging"
+import { ApiConfiguration, API_CONFIG_KEYS } from "../shared/api"
+import { GLOBAL_STATE_KEYS, SECRET_KEYS, GlobalStateKey, SecretKey } from "../shared/globalState"
+
+export class ContextProxy {
+	private readonly originalContext: vscode.ExtensionContext
+	private stateCache: Map<string, any>
+	private secretCache: Map<string, string | undefined>
+
+	constructor(context: vscode.ExtensionContext) {
+		// Initialize properties first
+		this.originalContext = context
+		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")
+	}
+
+	// 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
+	}
+	get extensionPath(): string {
+		return this.originalContext.extensionPath
+	}
+	get globalStorageUri(): vscode.Uri {
+		return this.originalContext.globalStorageUri
+	}
+	get logUri(): vscode.Uri {
+		return this.originalContext.logUri
+	}
+	get extension(): vscode.Extension<any> | undefined {
+		return this.originalContext.extension
+	}
+	get extensionMode(): vscode.ExtensionMode {
+		return this.originalContext.extensionMode
+	}
+
+	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)
+	}
+
+	updateGlobalState<T>(key: string, value: T): Thenable<void> {
+		this.stateCache.set(key, value)
+		return this.originalContext.globalState.update(key, value)
+	}
+
+	getSecret(key: string): string | undefined {
+		return this.secretCache.get(key)
+	}
+	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)
+		}
+	}
+
+	/**
+	 * Gets a complete ApiConfiguration object by fetching values
+	 * from both global state and secrets storage
+	 */
+	getApiConfiguration(): ApiConfiguration {
+		// Create an empty ApiConfiguration object
+		const config: ApiConfiguration = {}
+
+		// Add all API-related keys from global state
+		for (const key of API_CONFIG_KEYS) {
+			const value = this.getGlobalState(key)
+			if (value !== undefined) {
+				// Use type assertion to avoid TypeScript error
+				;(config as any)[key] = value
+			}
+		}
+
+		// Add all secret values
+		for (const key of SECRET_KEYS) {
+			const value = this.getSecret(key)
+			if (value !== undefined) {
+				// Use type assertion to avoid TypeScript error
+				;(config as any)[key] = value
+			}
+		}
+
+		// Handle special case for apiProvider if needed (same logic as current implementation)
+		if (!config.apiProvider) {
+			if (config.apiKey) {
+				config.apiProvider = "anthropic"
+			} else {
+				config.apiProvider = "openrouter"
+			}
+		}
+
+		return config
+	}
+
+	/**
+	 * Updates an ApiConfiguration by persisting each property
+	 * to the appropriate storage (global state or secrets)
+	 */
+	async updateApiConfiguration(apiConfiguration: ApiConfiguration): Promise<void> {
+		const promises: Array<Thenable<void>> = []
+
+		// For each property, update the appropriate storage
+		Object.entries(apiConfiguration).forEach(([key, value]) => {
+			if (SECRET_KEYS.includes(key as SecretKey)) {
+				promises.push(this.storeSecret(key, value))
+			} else if (API_CONFIG_KEYS.includes(key as GlobalStateKey)) {
+				promises.push(this.updateGlobalState(key, value))
+			}
+			// Ignore keys that aren't in either list
+		})
+
+		await Promise.all(promises)
+	}
+}