diff --git a/src/core/__tests__/contextProxy.test.ts b/src/core/__tests__/contextProxy.test.ts index ea9e79736d..e2d6c4ad12 100644 --- a/src/core/__tests__/contextProxy.test.ts +++ b/src/core/__tests__/contextProxy.test.ts @@ -102,6 +102,29 @@ describe("ContextProxy", () => { const result = proxy.getGlobalState("apiProvider", "default-value") expect(result).toBe("default-value") }) + + it("should bypass cache for pass-through state keys", async () => { + // Setup mock return value + mockGlobalState.get.mockReturnValue("pass-through-value") + + // Use a pass-through key (taskHistory) + const result = proxy.getGlobalState("taskHistory" as GlobalStateKey) + + // Should get value directly from original context + expect(result).toBe("pass-through-value") + expect(mockGlobalState.get).toHaveBeenCalledWith("taskHistory") + }) + + it("should respect default values for pass-through state keys", async () => { + // Setup mock to return undefined + mockGlobalState.get.mockReturnValue(undefined) + + // Use a pass-through key with default value + const result = proxy.getGlobalState("taskHistory" as GlobalStateKey, "default-value") + + // Should return default value when original context returns undefined + expect(result).toBe("default-value") + }) }) describe("updateGlobalState", () => { @@ -115,6 +138,21 @@ describe("ContextProxy", () => { const storedValue = await proxy.getGlobalState("apiProvider") expect(storedValue).toBe("new-value") }) + + it("should bypass cache for pass-through state keys", async () => { + await proxy.updateGlobalState("taskHistory" as GlobalStateKey, "new-value") + + // Should update original context + expect(mockGlobalState.update).toHaveBeenCalledWith("taskHistory", "new-value") + + // Setup mock for subsequent get + mockGlobalState.get.mockReturnValue("new-value") + + // Should get fresh value from original context + const storedValue = proxy.getGlobalState("taskHistory" as GlobalStateKey) + expect(storedValue).toBe("new-value") + expect(mockGlobalState.get).toHaveBeenCalledWith("taskHistory") + }) }) describe("getSecret", () => { diff --git a/src/core/contextProxy.ts b/src/core/contextProxy.ts index 40d8b475bc..698713179d 100644 --- a/src/core/contextProxy.ts +++ b/src/core/contextProxy.ts @@ -10,6 +10,7 @@ import { ConfigurationValues, isSecretKey, isGlobalStateKey, + isPassThroughStateKey, } from "../shared/globalState" import { API_CONFIG_KEYS, ApiConfiguration } from "../shared/api" @@ -80,11 +81,18 @@ export class ContextProxy { getGlobalState(key: GlobalStateKey): T | undefined getGlobalState(key: GlobalStateKey, defaultValue: T): T getGlobalState(key: GlobalStateKey, defaultValue?: T): T | undefined { + if (isPassThroughStateKey(key)) { + const value = this.originalContext.globalState.get(key) + return value === undefined || value === null ? defaultValue : (value as T) + } const value = this.stateCache.get(key) as T | undefined return value !== undefined ? value : (defaultValue as T | undefined) } updateGlobalState(key: GlobalStateKey, value: T) { + if (isPassThroughStateKey(key)) { + return this.originalContext.globalState.update(key, value) + } this.stateCache.set(key, value) return this.originalContext.globalState.update(key, value) } @@ -114,12 +122,14 @@ export class ContextProxy { setValue(key: ConfigurationKey, value: any) { if (isSecretKey(key)) { return this.storeSecret(key, value) - } else if (isGlobalStateKey(key)) { - return this.updateGlobalState(key, value) - } else { - logger.warn(`Unknown key: ${key}. Storing as global state.`) + } + + if (isGlobalStateKey(key)) { return this.updateGlobalState(key, value) } + + logger.warn(`Unknown key: ${key}. Storing as global state.`) + return this.updateGlobalState(key, value) } /** diff --git a/src/shared/globalState.ts b/src/shared/globalState.ts index 4c0adb0fcf..4331effce9 100644 --- a/src/shared/globalState.ts +++ b/src/shared/globalState.ts @@ -124,6 +124,8 @@ export const GLOBAL_STATE_KEYS = [ "maxWorkspaceFiles", ] as const +export const PASS_THROUGH_STATE_KEYS = ["taskHistory"] as const + type CheckGlobalStateKeysExhaustiveness = Exclude extends never ? true : false @@ -133,3 +135,6 @@ export const isSecretKey = (key: string): key is SecretKey => SECRET_KEYS.includ export const isGlobalStateKey = (key: string): key is GlobalStateKey => GLOBAL_STATE_KEYS.includes(key as GlobalStateKey) + +export const isPassThroughStateKey = (key: string): key is (typeof PASS_THROUGH_STATE_KEYS)[number] => + PASS_THROUGH_STATE_KEYS.includes(key as (typeof PASS_THROUGH_STATE_KEYS)[number])