diff --git a/src/core/__tests__/contextProxy.test.ts b/src/core/__tests__/contextProxy.test.ts index 9f0c20b0c4..0ac98bbc8c 100644 --- a/src/core/__tests__/contextProxy.test.ts +++ b/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") + }) + }) }) diff --git a/src/core/contextProxy.ts b/src/core/contextProxy.ts index 7c429c86cf..f27cc7f65c 100644 --- a/src/core/contextProxy.ts +++ b/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 { + 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): Promise { + const promises: Thenable[] = [] + + for (const [key, value] of Object.entries(values)) { + promises.push(this.setValue(key, value)) + } + + return Promise.all(promises) + } } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 8d8732bdd7..017371e558 100644 --- a/src/core/webview/ClineProvider.ts +++ b/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[] = [] - - // 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({