Merge pull request #1396 from samhvw8/refactor-context-proxy-more

feat(contextProxy): add setValue and setValues methods to simplify

Author: cte

src/core/__tests__/contextProxy.test.ts

@@ -154,4 +154,105 @@ describe("ContextProxy", () => {
 			expect(storedValue).toBeUndefined()
 		})
 	})
+
+	describe("setValue", () => {
+		it("should route secret keys to storeSecret", async () => {
+			// Spy on storeSecret
+			const storeSecretSpy = jest.spyOn(proxy, "storeSecret")
+
+			// Test with a known secret key
+			await proxy.setValue("openAiApiKey", "test-api-key")
+
+			// Should have called storeSecret
+			expect(storeSecretSpy).toHaveBeenCalledWith("openAiApiKey", "test-api-key")
+
+			// Should have stored the value in secret cache
+			const storedValue = proxy.getSecret("openAiApiKey")
+			expect(storedValue).toBe("test-api-key")
+		})
+
+		it("should route global state keys to updateGlobalState", async () => {
+			// Spy on updateGlobalState
+			const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
+
+			// Test with a known global state key
+			await proxy.setValue("apiModelId", "gpt-4")
+
+			// Should have called updateGlobalState
+			expect(updateGlobalStateSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
+
+			// Should have stored the value in state cache
+			const storedValue = proxy.getGlobalState("apiModelId")
+			expect(storedValue).toBe("gpt-4")
+		})
+
+		it("should handle unknown keys as global state with warning", async () => {
+			// Spy on the logger
+			const warnSpy = jest.spyOn(logger, "warn")
+
+			// Spy on updateGlobalState
+			const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
+
+			// Test with an unknown key
+			await proxy.setValue("unknownKey", "some-value")
+
+			// Should have logged a warning
+			expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("Unknown key: unknownKey"))
+
+			// Should have called updateGlobalState
+			expect(updateGlobalStateSpy).toHaveBeenCalledWith("unknownKey", "some-value")
+
+			// Should have stored the value in state cache
+			const storedValue = proxy.getGlobalState("unknownKey")
+			expect(storedValue).toBe("some-value")
+		})
+	})
+
+	describe("setValues", () => {
+		it("should process multiple values correctly", async () => {
+			// Spy on setValue
+			const setValueSpy = jest.spyOn(proxy, "setValue")
+
+			// Test with multiple values
+			await proxy.setValues({
+				apiModelId: "gpt-4",
+				apiProvider: "openai",
+				mode: "test-mode",
+			})
+
+			// Should have called setValue for each key
+			expect(setValueSpy).toHaveBeenCalledTimes(3)
+			expect(setValueSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
+			expect(setValueSpy).toHaveBeenCalledWith("apiProvider", "openai")
+			expect(setValueSpy).toHaveBeenCalledWith("mode", "test-mode")
+
+			// Should have stored all values in state cache
+			expect(proxy.getGlobalState("apiModelId")).toBe("gpt-4")
+			expect(proxy.getGlobalState("apiProvider")).toBe("openai")
+			expect(proxy.getGlobalState("mode")).toBe("test-mode")
+		})
+
+		it("should handle both secret and global state keys", async () => {
+			// Spy on storeSecret and updateGlobalState
+			const storeSecretSpy = jest.spyOn(proxy, "storeSecret")
+			const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
+
+			// Test with mixed keys
+			await proxy.setValues({
+				apiModelId: "gpt-4", // global state
+				openAiApiKey: "test-api-key", // secret
+				unknownKey: "some-value", // unknown
+			})
+
+			// Should have called appropriate methods
+			expect(storeSecretSpy).toHaveBeenCalledWith("openAiApiKey", "test-api-key")
+			expect(updateGlobalStateSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
+			expect(updateGlobalStateSpy).toHaveBeenCalledWith("unknownKey", "some-value")
+
+			// Should have stored values in appropriate caches
+			expect(proxy.getSecret("openAiApiKey")).toBe("test-api-key")
+			expect(proxy.getGlobalState("apiModelId")).toBe("gpt-4")
+			expect(proxy.getGlobalState("unknownKey")).toBe("some-value")
+		})
+	})
 })

src/core/contextProxy.ts

@@ -93,4 +93,40 @@ export class ContextProxy {
 			return this.originalContext.secrets.store(key, value)
 		}
 	}
+	/**
+	 * Set a value in either secrets or global state based on key type.
+	 * If the key is in SECRET_KEYS, it will be stored as a secret.
+	 * If the key is in GLOBAL_STATE_KEYS or unknown, it will be stored in global state.
+	 * @param key The key to set
+	 * @param value The value to set
+	 * @returns A promise that resolves when the operation completes
+	 */
+	setValue(key: string, value: any): Thenable<void> {
+		if (SECRET_KEYS.includes(key as any)) {
+			return this.storeSecret(key, value)
+		}
+
+		if (GLOBAL_STATE_KEYS.includes(key as any)) {
+			return this.updateGlobalState(key, value)
+		}
+
+		logger.warn(`Unknown key: ${key}. Storing as global state.`)
+		return this.updateGlobalState(key, value)
+	}
+
+	/**
+	 * Set multiple values at once. Each key will be routed to either
+	 * secrets or global state based on its type.
+	 * @param values An object containing key-value pairs to set
+	 * @returns A promise that resolves when all operations complete
+	 */
+	async setValues(values: Record<string, any>): Promise<void[]> {
+		const promises: Thenable<void>[] = []
+
+		for (const [key, value] of Object.entries(values)) {
+			promises.push(this.setValue(key, value))
+		}
+
+		return Promise.all(promises)
+	}
 }

src/core/webview/ClineProvider.ts

@@ -1688,20 +1688,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 			}
 		}
 
-		// Create an array of promises to update state
-		const promises: Promise<any>[] = []
-
-		// For each property in apiConfiguration, update the appropriate state
-		Object.entries(apiConfiguration).forEach(([key, value]) => {
-			// Check if this key is a secret
-			if (SECRET_KEYS.includes(key as SecretKey)) {
-				promises.push(this.storeSecret(key as SecretKey, value))
-			} else {
-				promises.push(this.updateGlobalState(key as GlobalStateKey, value))
-			}
-		})
-
-		await Promise.all(promises)
+		// Use the new setValues method to handle routing values to secrets or global state
+		await this.contextProxy.setValues(apiConfiguration)
 
 		if (this.cline) {
 			this.cline.api = buildApiHandler(apiConfiguration)
@@ -1805,8 +1793,11 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		}
 
 		const openrouter: ApiProvider = "openrouter"
-		await this.updateGlobalState("apiProvider", openrouter)
-		await this.storeSecret("openRouterApiKey", apiKey)
+		await this.contextProxy.setValues({
+			apiProvider: openrouter,
+			openRouterApiKey: apiKey,
+		})
+
 		await this.postStateToWebview()
 		if (this.cline) {
 			this.cline.api = buildApiHandler({ apiProvider: openrouter, openRouterApiKey: apiKey })
@@ -1833,8 +1824,10 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		}
 
 		const glama: ApiProvider = "glama"
-		await this.updateGlobalState("apiProvider", glama)
-		await this.storeSecret("glamaApiKey", apiKey)
+		await this.contextProxy.setValues({
+			apiProvider: glama,
+			glamaApiKey: apiKey,
+		})
 		await this.postStateToWebview()
 		if (this.cline) {
 			this.cline.api = buildApiHandler({