diff --git a/packages/ui/src/stores/useAgentsStore.ts b/packages/ui/src/stores/useAgentsStore.ts index aa083528..742d2343 100644 --- a/packages/ui/src/stores/useAgentsStore.ts +++ b/packages/ui/src/stores/useAgentsStore.ts @@ -630,6 +630,7 @@ async function performConfigRefresh(options: { if (refreshProviders) { useConfigStore.getState().invalidateModelMetadataCache(); + useConfigStore.getState().invalidateProviderCache(mode === "active" ? currentDirectory : undefined); } const sdkRefreshTasks: Promise[] = []; diff --git a/packages/ui/src/stores/useConfigStore.test.ts b/packages/ui/src/stores/useConfigStore.test.ts new file mode 100644 index 00000000..3f2f18bf --- /dev/null +++ b/packages/ui/src/stores/useConfigStore.test.ts @@ -0,0 +1,313 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test'; + +const DIRECTORY = '/workspace/project'; +const OTHER_DIRECTORY = '/workspace/other'; +const STORAGE_KEY = 'config-store'; + +let storage = new Map(); +let liveProviderId = 'live'; +let liveProviderIdsByDirectory = new Map(); +let liveProviderVariants: Record> | undefined; +let getProvidersCalls = 0; +let withDirectoryCalls: Array = []; +let currentFetchDirectory: string | null = DIRECTORY; +let configListener: ((event: { scopes: string[]; source?: string; timestamp: number }) => void | Promise) | null = null; + +const makeStorage = (): Storage => ({ + getItem: (key: string) => storage.get(key) ?? null, + setItem: (key: string, value: string) => { + storage.set(key, value); + }, + removeItem: (key: string) => { + storage.delete(key); + }, + clear: () => { + storage.clear(); + }, + key: (index: number) => Array.from(storage.keys())[index] ?? null, + get length() { + return storage.size; + }, +}) as Storage; + +const provider = (id: string, modelId = `${id}-model`, variants?: Record>) => ({ + id, + name: id, + source: 'config' as const, + env: [], + options: {}, + models: [ + { + id: modelId, + name: modelId, + providerID: id, + api: { id: 'chat', url: '', npm: '' }, + capabilities: { + temperature: true, + reasoning: false, + attachment: false, + toolcall: true, + input: { text: true, audio: false, image: false, video: false, pdf: false }, + output: { text: true, audio: false, image: false, video: false, pdf: false }, + interleaved: false, + }, + cost: { input: 0, output: 0, cache: { read: 0, write: 0 } }, + limit: { context: 0, output: 0 }, + options: {}, + release_date: '', + status: 'active' as const, + headers: {}, + attachment: false, + reasoning: false, + temperature: true, + tool_call: true, + ...(variants ? { variants } : {}), + }, + ], +}); + +const providerResponse = (id: string, modelId = `${id}-model`, variants?: Record>) => ({ + id, + name: id, + source: 'config' as const, + env: [], + options: {}, + models: { + [modelId]: { + id: modelId, + name: modelId, + providerID: id, + api: { id: 'chat', url: '', npm: '' }, + capabilities: { + temperature: true, + reasoning: false, + attachment: false, + toolcall: true, + input: { text: true, audio: false, image: false, video: false, pdf: false }, + output: { text: true, audio: false, image: false, video: false, pdf: false }, + interleaved: false, + }, + cost: { input: 0, output: 0, cache: { read: 0, write: 0 } }, + limit: { context: 0, output: 0 }, + options: {}, + release_date: '', + status: 'active' as const, + headers: {}, + attachment: false, + reasoning: false, + temperature: true, + tool_call: true, + ...(variants ? { variants } : {}), + }, + }, +}); + +mock.module('@/stores/utils/safeStorage', () => ({ + getSafeStorage: () => makeStorage(), +})); + +mock.module('@/lib/opencode/client', () => ({ + opencodeClient: { + setDirectory: mock(() => undefined), + getDirectory: mock(() => DIRECTORY), + checkHealth: mock(async () => true), + withDirectory: mock(async (directory: string | null, callback: () => Promise) => { + withDirectoryCalls.push(directory); + const previous = currentFetchDirectory; + currentFetchDirectory = directory; + try { + return await callback(); + } finally { + currentFetchDirectory = previous; + } + }), + getProviders: mock(async () => { + getProvidersCalls += 1; + const id = liveProviderIdsByDirectory.get(currentFetchDirectory ?? '') ?? liveProviderId; + return { providers: [providerResponse(id, `${id}-model`, liveProviderVariants)], default: { default: id } }; + }), + listAgents: mock(async () => []), + }, +})); + +mock.module('@/contexts/runtimeAPIRegistry', () => ({ + getRegisteredRuntimeAPIs: mock(() => null), +})); + +mock.module('@/lib/runtime-fetch', () => ({ + runtimeFetch: mock(async () => new Response(JSON.stringify({}), { + headers: { 'Content-Type': 'application/json' }, + })), +})); + +mock.module('@/lib/persistence', () => ({ + updateDesktopSettings: mock(async () => undefined), +})); + +mock.module('@/lib/startupTrace', () => ({ + markStartupTrace: mock(() => undefined), + measureStartupTrace: mock(async (_name: string, callback: () => Promise) => callback()), +})); + +mock.module('@/lib/configSync', () => ({ + emitConfigChange: mock(() => undefined), + scopeMatches: mock((event: { scopes: string[] }, scope: string) => event.scopes.includes('all') || event.scopes.includes(scope)), + subscribeToConfigChanges: mock((listener: typeof configListener) => { + configListener = listener; + return () => { + if (configListener === listener) { + configListener = null; + } + }; + }), +})); + +const { useConfigStore } = await import('./useConfigStore'); + +describe('useConfigStore provider persistence', () => { + beforeEach(() => { + storage = new Map(); + liveProviderId = 'live'; + liveProviderIdsByDirectory = new Map(); + liveProviderVariants = undefined; + getProvidersCalls = 0; + withDirectoryCalls = []; + currentFetchDirectory = DIRECTORY; + useConfigStore.setState({ + activeDirectoryKey: DIRECTORY, + directoryScoped: {}, + providers: [], + defaultProviders: {}, + currentProviderId: '', + currentModelId: '', + currentVariant: undefined, + selectedProviderId: '', + isConnected: true, + isInitialized: false, + }); + }); + + test('strips persisted provider snapshots while preserving other directory state', async () => { + storage.set(STORAGE_KEY, JSON.stringify({ + state: { + activeDirectoryKey: DIRECTORY, + directoryScoped: { + [DIRECTORY]: { + providers: [provider('stale')], + agents: [{ name: 'build', mode: 'primary' }], + currentProviderId: 'stale', + currentModelId: 'stale-model', + currentAgentName: 'build', + selectedProviderId: 'stale', + agentModelSelections: { build: { providerId: 'stale', modelId: 'stale-model' } }, + defaultProviders: { default: 'stale' }, + }, + [OTHER_DIRECTORY]: { + providers: [provider('other-stale')], + agents: [{ name: 'review', mode: 'primary' }], + currentProviderId: 'other-stale', + currentModelId: 'other-stale-model', + currentAgentName: 'review', + selectedProviderId: 'other-stale', + agentModelSelections: {}, + defaultProviders: { default: 'other-stale' }, + }, + }, + currentProviderId: 'stale', + currentModelId: 'stale-model', + selectedProviderId: 'stale', + defaultProviders: { default: 'stale' }, + }, + version: 0, + })); + + await useConfigStore.persist.rehydrate(); + + const hydrated = useConfigStore.getState(); + expect(hydrated.providers).toEqual([]); + expect(hydrated.defaultProviders).toEqual({}); + expect(hydrated.directoryScoped[DIRECTORY]?.providers).toEqual([]); + expect(hydrated.directoryScoped[DIRECTORY]?.defaultProviders).toEqual({}); + expect(hydrated.directoryScoped[DIRECTORY]?.agents).toEqual([{ name: 'build', mode: 'primary' }]); + expect(hydrated.directoryScoped[DIRECTORY]?.currentAgentName).toBe('build'); + expect(hydrated.directoryScoped[OTHER_DIRECTORY]?.providers).toEqual([]); + + liveProviderId = 'fresh'; + await hydrated.initializeApp(); + + const reloaded = useConfigStore.getState(); + expect(getProvidersCalls).toBe(1); + expect(reloaded.providers.map((entry) => entry.id)).toEqual(['fresh']); + expect(reloaded.directoryScoped[DIRECTORY]?.providers.map((entry) => entry.id)).toEqual(['fresh']); + expect(reloaded.currentProviderId).toBe('fresh'); + expect(reloaded.currentModelId).toBe('fresh-model'); + }); + + test('provider config events refresh all known directory provider caches immediately', async () => { + useConfigStore.setState({ + activeDirectoryKey: DIRECTORY, + providers: [provider('active-stale')], + defaultProviders: { default: 'active-stale' }, + currentProviderId: 'active-stale', + currentModelId: 'active-stale-model', + selectedProviderId: 'active-stale', + directoryScoped: { + [DIRECTORY]: { + providers: [provider('active-stale')], + agents: [], + currentProviderId: 'active-stale', + currentModelId: 'active-stale-model', + currentAgentName: undefined, + selectedProviderId: 'active-stale', + agentModelSelections: {}, + defaultProviders: { default: 'active-stale' }, + }, + [OTHER_DIRECTORY]: { + providers: [provider('inactive-cached')], + agents: [], + currentProviderId: 'inactive-cached', + currentModelId: 'inactive-cached-model', + currentAgentName: undefined, + selectedProviderId: 'inactive-cached', + agentModelSelections: {}, + defaultProviders: { default: 'inactive-cached' }, + }, + }, + }); + + liveProviderIdsByDirectory = new Map([ + [DIRECTORY, 'active-live'], + [OTHER_DIRECTORY, 'inactive-live'], + ]); + expect(configListener).not.toBeNull(); + await configListener?.({ scopes: ['providers'], timestamp: Date.now() }); + + const state = useConfigStore.getState(); + expect(getProvidersCalls).toBe(2); + expect(new Set(withDirectoryCalls)).toEqual(new Set([DIRECTORY, OTHER_DIRECTORY])); + expect(state.directoryScoped[DIRECTORY]?.providers.map((entry) => entry.id)).toEqual(['active-live']); + expect(state.directoryScoped[OTHER_DIRECTORY]?.providers.map((entry) => entry.id)).toEqual(['inactive-live']); + expect(state.directoryScoped[OTHER_DIRECTORY]?.defaultProviders).toEqual({ default: 'inactive-live' }); + }); + + test('provider reload preserves a valid current variant', async () => { + useConfigStore.setState({ + activeDirectoryKey: DIRECTORY, + currentProviderId: 'live', + currentModelId: 'live-model', + currentVariant: 'fast', + selectedProviderId: 'live', + settingsDefaultVariant: 'slow', + directoryScoped: {}, + }); + + liveProviderId = 'live'; + liveProviderVariants = { fast: {}, slow: {} }; + await useConfigStore.getState().loadProviders({ source: 'test:variant' }); + + const state = useConfigStore.getState(); + expect(state.currentProviderId).toBe('live'); + expect(state.currentModelId).toBe('live-model'); + expect(state.currentVariant).toBe('fast'); + }); +}); diff --git a/packages/ui/src/stores/useConfigStore.ts b/packages/ui/src/stores/useConfigStore.ts index 94f01f94..eb4243cf 100644 --- a/packages/ui/src/stores/useConfigStore.ts +++ b/packages/ui/src/stores/useConfigStore.ts @@ -28,6 +28,7 @@ const FALLBACK_PROVIDER_ID = "opencode"; const FALLBACK_MODEL_ID = "big-pickle"; const GIT_UTILITY_PROVIDER_ID = "zen"; const GIT_UTILITY_PREFERRED_MODEL_ID = "big-pickle"; +const PROVIDER_CONFIG_REFRESH_CONCURRENCY = 4; const normalizeSttSilenceThresholdDb = (value: unknown): number | undefined => { if (typeof value !== 'number' || !Number.isFinite(value)) { @@ -181,6 +182,7 @@ type ProviderModel = Provider["models"][string]; type ProviderWithModelList = Omit & { models: ProviderModel[] }; type GitModelSelection = { providerId: string; modelId: string }; +type ProviderModelSelection = { providerId: string; modelId: string; variant?: string } | null; const normalizeOptionalString = (value: unknown): string | undefined => { if (typeof value !== "string") { @@ -202,6 +204,67 @@ const hasProviderModel = ( return provider.models.some((model) => model.id === modelId); }; +const resolveProviderModelSelection = ({ + providers, + currentProviderId, + currentModelId, + currentVariant, + settingsDefaultModel, + settingsDefaultVariant, +}: { + providers: ProviderWithModelList[]; + currentProviderId?: string; + currentModelId?: string; + currentVariant?: string; + settingsDefaultModel?: string; + settingsDefaultVariant?: string; +}): ProviderModelSelection => { + const resolveVariant = (providerId: string, modelId: string, variant?: string): string | undefined => { + if (!variant) { + return undefined; + } + + const model = providers + .find((provider) => provider.id === providerId) + ?.models.find((entry) => entry.id === modelId) as { variants?: Record } | undefined; + + return model?.variants && Object.prototype.hasOwnProperty.call(model.variants, variant) + ? variant + : undefined; + }; + + if (currentProviderId && currentModelId && hasProviderModel(providers, currentProviderId, currentModelId)) { + return { + providerId: currentProviderId, + modelId: currentModelId, + variant: resolveVariant(currentProviderId, currentModelId, currentVariant), + }; + } + + if (settingsDefaultModel) { + const parsed = parseModelString(settingsDefaultModel); + if (parsed && hasProviderModel(providers, parsed.providerId, parsed.modelId)) { + return { + providerId: parsed.providerId, + modelId: parsed.modelId, + variant: resolveVariant(parsed.providerId, parsed.modelId, settingsDefaultVariant), + }; + } + } + + if (hasProviderModel(providers, FALLBACK_PROVIDER_ID, FALLBACK_MODEL_ID)) { + return { providerId: FALLBACK_PROVIDER_ID, modelId: FALLBACK_MODEL_ID }; + } + + const firstProvider = providers[0]; + const firstModel = firstProvider?.models[0]; + if (firstProvider && firstModel) { + return { providerId: firstProvider.id, modelId: firstModel.id }; + } + + return null; +}; + const resolveGitGenerationModelSelection = ({ providers, settingsZenModel, @@ -528,6 +591,43 @@ interface DirectoryScopedConfig { defaultProviders: { [key: string]: string }; } +const clearProviderDataFromDirectoryScoped = ( + directoryScoped: Record, +): Record => { + const next: Record = {}; + + for (const [directoryKey, snapshot] of Object.entries(directoryScoped)) { + next[directoryKey] = { + ...snapshot, + providers: [], + defaultProviders: {}, + }; + } + + return next; +}; + +const stripProviderCacheFromPersistedState = (persistedState: unknown): Partial => { + if (!persistedState || typeof persistedState !== 'object') { + return {}; + } + + const persisted = persistedState as Partial; + const sanitized: Partial = { + ...persisted, + providers: [], + defaultProviders: {}, + }; + + if (persisted.directoryScoped) { + sanitized.directoryScoped = clearProviderDataFromDirectoryScoped( + persisted.directoryScoped as Record, + ); + } + + return sanitized; +}; + interface ConfigStore { activeDirectoryKey: string; @@ -621,6 +721,7 @@ interface ConfigStore { loadProviders: (options?: { directory?: string | null; source?: string }) => Promise; loadAgents: (options?: { directory?: string | null; source?: string }) => Promise; invalidateModelMetadataCache: () => void; + invalidateProviderCache: (directory?: string | null) => void; setProvider: (providerId: string) => void; setModel: (modelId: string) => void; setCurrentVariant: (variant: string | undefined) => void; @@ -974,6 +1075,57 @@ export const useConfigStore = create()( } }, + invalidateProviderCache: (directory) => { + const targetDirectoryKey = directory === undefined ? null : toDirectoryKey(directory); + + set((state) => { + const nextState: Partial = {}; + let scopedChanged = false; + const nextDirectoryScoped: Record = { + ...state.directoryScoped, + }; + + const clearSnapshot = (snapshot: DirectoryScopedConfig): DirectoryScopedConfig => { + if (snapshot.providers.length === 0 && Object.keys(snapshot.defaultProviders).length === 0) { + return snapshot; + } + + scopedChanged = true; + return { + ...snapshot, + providers: [], + defaultProviders: {}, + }; + }; + + if (targetDirectoryKey) { + const snapshot = state.directoryScoped[targetDirectoryKey]; + if (snapshot) { + nextDirectoryScoped[targetDirectoryKey] = clearSnapshot(snapshot); + } + } else { + for (const [directoryKey, snapshot] of Object.entries(state.directoryScoped)) { + nextDirectoryScoped[directoryKey] = clearSnapshot(snapshot); + } + } + + if (scopedChanged) { + nextState.directoryScoped = nextDirectoryScoped; + } + + if (targetDirectoryKey === null || targetDirectoryKey === state.activeDirectoryKey) { + if (state.providers.length > 0) { + nextState.providers = []; + } + if (Object.keys(state.defaultProviders).length > 0) { + nextState.defaultProviders = {}; + } + } + + return Object.keys(nextState).length > 0 ? nextState : state; + }); + }, + loadProviders: async (options) => { const requestedDirectory = options?.directory ?? fromDirectoryKey(get().activeDirectoryKey); const effectiveDirectory = requestedDirectory ?? opencodeClient.getDirectory() ?? null; @@ -1034,10 +1186,38 @@ export const useConfigStore = create()( defaultProviders: {}, }; + const currentProviderId = state.activeDirectoryKey === directoryKey + ? state.currentProviderId + : baseSnapshot.currentProviderId; + const currentModelId = state.activeDirectoryKey === directoryKey + ? state.currentModelId + : baseSnapshot.currentModelId; + const currentVariant = state.activeDirectoryKey === directoryKey + ? state.currentVariant + : baseSnapshot.currentVariant; + const resolvedModel = resolveProviderModelSelection({ + providers: processedProviders, + currentProviderId, + currentModelId, + currentVariant, + settingsDefaultModel: state.settingsDefaultModel, + settingsDefaultVariant: state.settingsDefaultVariant, + }); + const currentSelectedProviderId = state.activeDirectoryKey === directoryKey + ? state.selectedProviderId + : baseSnapshot.selectedProviderId; + const selectedProviderId = processedProviders.some((provider) => provider.id === currentSelectedProviderId) + ? currentSelectedProviderId + : (resolvedModel?.providerId ?? processedProviders[0]?.id ?? ""); + const nextSnapshot: DirectoryScopedConfig = { ...baseSnapshot, providers: processedProviders, defaultProviders: defaults, + currentProviderId: resolvedModel?.providerId ?? "", + currentModelId: resolvedModel?.modelId ?? "", + currentVariant: resolvedModel?.variant, + selectedProviderId, }; const nextState: Partial = { @@ -1050,29 +1230,10 @@ export const useConfigStore = create()( if (state.activeDirectoryKey === directoryKey) { nextState.providers = processedProviders; nextState.defaultProviders = defaults; - - if (!state.currentProviderId && !state.currentModelId && state.settingsDefaultModel) { - const parsed = parseModelString(state.settingsDefaultModel); - if (parsed) { - const settingsProvider = processedProviders.find((p) => p.id === parsed.providerId); - if (settingsProvider?.models.some((m) => m.id === parsed.modelId)) { - const model = settingsProvider.models.find((m) => m.id === parsed.modelId); - const currentVariant = state.settingsDefaultVariant && (model as { variants?: Record } | undefined)?.variants?.[state.settingsDefaultVariant] - ? state.settingsDefaultVariant - : undefined; - - nextState.currentProviderId = parsed.providerId; - nextState.currentModelId = parsed.modelId; - nextState.currentVariant = currentVariant; - nextState.selectedProviderId = parsed.providerId; - - nextSnapshot.currentProviderId = parsed.providerId; - nextSnapshot.currentModelId = parsed.modelId; - nextSnapshot.currentVariant = currentVariant; - nextSnapshot.selectedProviderId = parsed.providerId; - } - } - } + nextState.currentProviderId = nextSnapshot.currentProviderId; + nextState.currentModelId = nextSnapshot.currentModelId; + nextState.currentVariant = nextSnapshot.currentVariant; + nextState.selectedProviderId = selectedProviderId; } return nextState; @@ -2267,6 +2428,8 @@ export const useConfigStore = create()( if (debug) console.log("Initializing app..."); markStartupTrace('initApp:skipped', { reason: 'checkConnection already verified health' }); + get().invalidateProviderCache(); + if (debug) console.log("Loading providers and agents..."); await Promise.all([ get().loadProviders({ source: 'initializeApp' }), @@ -2349,16 +2512,20 @@ export const useConfigStore = create()( { name: "config-store", storage: createJSONStorage(() => getSafeStorage()), + merge: (persistedState, currentState) => ({ + ...currentState, + ...stripProviderCacheFromPersistedState(persistedState), + }), partialize: (state) => ({ activeDirectoryKey: state.activeDirectoryKey, - directoryScoped: state.directoryScoped, + directoryScoped: clearProviderDataFromDirectoryScoped(state.directoryScoped), currentProviderId: state.currentProviderId, currentModelId: state.currentModelId, currentVariant: state.currentVariant, currentAgentName: state.currentAgentName, selectedProviderId: state.selectedProviderId, agentModelSelections: state.agentModelSelections, - defaultProviders: state.defaultProviders, + defaultProviders: {}, settingsDefaultModel: state.settingsDefaultModel, settingsDefaultVariant: state.settingsDefaultVariant, settingsDefaultAgent: state.settingsDefaultAgent, @@ -2380,6 +2547,31 @@ if (typeof window !== "undefined") { window.__zustand_config_store__ = useConfigStore; } +const refreshKnownProviderDirectories = async (source: string): Promise => { + const state = useConfigStore.getState(); + const directoryKeys = Array.from(new Set([ + state.activeDirectoryKey, + ...Object.keys(state.directoryScoped), + ])).filter((key) => key.length > 0); + + state.invalidateProviderCache(); + + let nextIndex = 0; + const workerCount = Math.min(PROVIDER_CONFIG_REFRESH_CONCURRENCY, directoryKeys.length); + const workers = Array.from({ length: workerCount }, async () => { + while (nextIndex < directoryKeys.length) { + const directoryKey = directoryKeys[nextIndex]; + nextIndex += 1; + await useConfigStore.getState().loadProviders({ + directory: fromDirectoryKey(directoryKey), + source, + }); + } + }); + + await Promise.all(workers); +}; + let unsubscribeConfigStoreChanges: (() => void) | null = null; if (!unsubscribeConfigStoreChanges) { @@ -2392,8 +2584,7 @@ if (!unsubscribeConfigStoreChanges) { } if (scopeMatches(event, "providers")) { - const { loadProviders } = useConfigStore.getState(); - tasks.push(loadProviders({ source: 'configChange:providers' })); + tasks.push(refreshKnownProviderDirectories('configChange:providers')); } if (tasks.length > 0) {