diff --git a/packages/ui/src/stores/instanceScopedStores.test.ts b/packages/ui/src/stores/instanceScopedStores.test.ts new file mode 100644 index 00000000..294cf6ef --- /dev/null +++ b/packages/ui/src/stores/instanceScopedStores.test.ts @@ -0,0 +1,92 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test'; +import type { McpStatus } from '@opencode-ai/sdk/v2'; +import type { McpStatusMap } from './useMcpStore'; + +type Deferred = { promise: Promise; resolve: (value: T) => void }; +const deferred = (): Deferred => { + let resolve!: (value: T) => void; + const promise = new Promise((res) => { resolve = res; }); + return { promise, resolve }; +}; + +type McpStatusResult = Awaited['mcp']['status']>>; +let mcpStatusResponse: Deferred = deferred(); +const opencodeModule = await import('@/lib/opencode/client'); +// Derived from the real client rather than spread from it: the client is a +// class instance, so a spread drops every prototype method the other modules +// loaded in this process call at import time. +// SAFETY: `Object.create` returns `any`; the object delegates to the real +// client for everything the two overrides below do not define. +const opencodeClientStub = Object.create(opencodeModule.opencodeClient) as typeof opencodeModule.opencodeClient; +// The SDK client is derived the same way, so only `mcp.status` is replaced and +// every other endpoint keeps its real implementation and type. +type McpApiClient = ReturnType; +const realApiClient = opencodeModule.opencodeClient.getApiClient(); +const mcpApiStub: McpApiClient = Object.create(realApiClient, { + mcp: { value: { ...realApiClient.mcp, status: () => mcpStatusResponse.promise } }, +}); +opencodeClientStub.getApiClient = () => mcpApiStub; +opencodeClientStub.getScopedApiClient = () => mcpApiStub; +mock.module('@/lib/opencode/client', () => ({ ...opencodeModule, opencodeClient: opencodeClientStub })); + +let skillsResponse: Deferred = deferred(); +const runtimeFetchModule = await import('@/lib/runtime-fetch'); +mock.module('@/lib/runtime-fetch', () => ({ + ...runtimeFetchModule, + runtimeFetch: () => skillsResponse.promise, +})); + +const { useMcpStore } = await import('./useMcpStore'); +const { useSkillsStore } = await import('./useSkillsStore'); + +const mcpStatusResult = (data: McpStatusMap): McpStatusResult => ({ + data, + request: new Request('http://localhost/mcp'), + response: new Response(), +}); + +const connectedServer = (name: string): McpStatusMap => ({ + // SAFETY: the store only reads `status` off each entry; the SDK type carries + // fields no consumer in this test path touches. + [name]: { status: 'connected' } as McpStatus, +}); + +describe('instance-scoped stores reject responses from the previous instance', () => { + beforeEach(() => { + mcpStatusResponse = deferred(); + skillsResponse = deferred(); + useMcpStore.getState().resetForRuntimeSwitch(); + useSkillsStore.getState().resetForRuntimeSwitch(); + }); + + test('an MCP status in flight during a switch does not land in the new instance', async () => { + const refresh = useMcpStore.getState().refresh({ directory: '/repo', silent: true }); + + useMcpStore.getState().resetForRuntimeSwitch(); + mcpStatusResponse.resolve(mcpStatusResult(connectedServer('from-instance-a'))); + await refresh; + + expect(useMcpStore.getState().getStatusForDirectory('/repo')).toEqual({}); + }); + + test('an MCP status that arrives with no switch is stored', async () => { + const refresh = useMcpStore.getState().refresh({ directory: '/repo', silent: true }); + mcpStatusResponse.resolve(mcpStatusResult(connectedServer('server-a'))); + await refresh; + + expect(Object.keys(useMcpStore.getState().getStatusForDirectory('/repo'))).toEqual(['server-a']); + }); + + test('a skills load in flight during a switch does not land in the new instance', async () => { + const load = useSkillsStore.getState().loadSkills('/repo'); + + useSkillsStore.getState().resetForRuntimeSwitch(); + skillsResponse.resolve(new Response( + JSON.stringify({ skills: [{ name: 'from-instance-a', path: '/repo/.agents/skills/a/SKILL.md' }] }), + { status: 200, headers: { 'content-type': 'application/json' } }, + )); + await load; + + expect(useSkillsStore.getState().skillsByDirectory['/repo']).toBe(undefined); + }); +}); diff --git a/packages/ui/src/stores/useMcpStore.ts b/packages/ui/src/stores/useMcpStore.ts index ab4ac80a..2946129e 100644 --- a/packages/ui/src/stores/useMcpStore.ts +++ b/packages/ui/src/stores/useMcpStore.ts @@ -54,6 +54,10 @@ type RefreshOptions = { }; const ensureFreshInFlight = new Map>(); +// Bumped on every runtime switch. Status is keyed by directory alone and two +// instances can hold the same project path, so a request already in flight for +// the previous instance would otherwise write its servers over the new one's. +let mcpGeneration = 0; type TestConnectionResult = { status?: McpStatus; @@ -108,6 +112,7 @@ export const useMcpStore = create()( refreshedAtKeys: {}, resetForRuntimeSwitch: () => { + mcpGeneration += 1; ensureFreshInFlight.clear(); set({ byDirectory: {}, @@ -144,9 +149,11 @@ export const useMcpStore = create()( })); } + const generation = mcpGeneration; try { const api = getMcpApiClient(directory); const result = await api.mcp.status(); + if (generation !== mcpGeneration) return; const data = (result.data ?? {}) as McpStatusMap; set((state) => ({ @@ -162,6 +169,7 @@ export const useMcpStore = create()( refreshedAtKeys: { ...state.refreshedAtKeys, [key]: Date.now() }, })); } catch (error) { + if (generation !== mcpGeneration) return; const message = error instanceof Error ? error.message : 'Failed to load MCP status'; set((state) => ({ loadingKeys: { ...state.loadingKeys, [key]: false }, diff --git a/packages/ui/src/stores/useQuotaStore.test.ts b/packages/ui/src/stores/useQuotaStore.test.ts index 036d3d3f..81d2d4d9 100644 --- a/packages/ui/src/stores/useQuotaStore.test.ts +++ b/packages/ui/src/stores/useQuotaStore.test.ts @@ -6,6 +6,7 @@ let isInitialized = true const fetched: string[] = [] type StubPayload = { usageDropdownProviders: string[] } | ProviderResult +let quotaRequestsFail = false; const json = (body: StubPayload) => new Response( JSON.stringify(body), { status: 200, headers: { "content-type": "application/json" } }, @@ -22,6 +23,7 @@ mock.module("@/lib/runtime-fetch", () => ({ ...runtimeFetchModule, runtimeFetch: async (path: string) => { fetched.push(path) + if (quotaRequestsFail) throw new Error("network down") if (path.startsWith("/api/config/settings")) return json({ usageDropdownProviders: ["claude"] }) return json({ providerId: "claude", providerName: "Claude", ok: true, configured: true, usage: null, fetchedAt: 1 }) }, @@ -40,6 +42,7 @@ describe("Usage quotas are loaded once per ready instance", () => { runtimeKey = "url:https://instance-a" isInitialized = true fetched.length = 0 + quotaRequestsFail = false useQuotaStore.getState().resetForRuntimeSwitch() }) @@ -97,4 +100,40 @@ describe("Usage quotas are loaded once per ready instance", () => { expect(fetched).toHaveLength(0) }) + + test("a failed load is not recorded as loaded, so the next ask retries it", async () => { + quotaRequestsFail = true + await useQuotaStore.getState().ensureLoadedForRuntime() + + expect(useQuotaStore.getState().loadedRuntimeKey).toBeNull() + + quotaRequestsFail = false + fetched.length = 0 + await useQuotaStore.getState().ensureLoadedForRuntime() + + expect(fetched.length).toBeGreaterThan(0) + expect(useQuotaStore.getState().loadedRuntimeKey).toBe("url:https://instance-a") + }) + + test("concurrent asks share one load", async () => { + await Promise.all([ + useQuotaStore.getState().ensureLoadedForRuntime(), + useQuotaStore.getState().ensureLoadedForRuntime(), + ]) + + expect(fetched.filter((path) => path.startsWith("/api/quota/"))).toHaveLength(1) + }) + + test("a switch drops the previous instance's display settings", async () => { + await useQuotaStore.getState().ensureLoadedForRuntime() + expect(useQuotaStore.getState().dropdownProviderIds).toEqual(["claude"]) + useQuotaStore.getState().setDisplayMode("remaining") + + useQuotaStore.getState().resetForRuntimeSwitch() + + // `dropdownProviderIds` decides which providers get queried, so carrying it + // over would ask the new instance through the old one's selection. + expect(useQuotaStore.getState().dropdownProviderIds.length).toBeGreaterThan(1) + expect(useQuotaStore.getState().displayMode).toBe("usage") + }) }) diff --git a/packages/ui/src/stores/useQuotaStore.ts b/packages/ui/src/stores/useQuotaStore.ts index a6f8b396..934231e4 100644 --- a/packages/ui/src/stores/useQuotaStore.ts +++ b/packages/ui/src/stores/useQuotaStore.ts @@ -16,6 +16,7 @@ const QUOTA_REFRESH_INTERVAL_MS = 3 * 60 * 1000; // instance, so both belong to that instance. Bumped on every reset so a // response in flight for the previous instance cannot land in the new one. let quotaGeneration = 0; +let inFlightRuntimeLoad: Promise | null = null; let quotaAutoRefreshConsumers = 0; let quotaAutoRefreshInterval: number | null = null; @@ -38,8 +39,10 @@ interface QuotaStore extends QuotaSettingsState { loadSettings: () => Promise; fetchAllQuotas: () => Promise; - fetchQuotas: (providerIds: QuotaProviderId[]) => Promise; - fetchProviderQuota: (providerId: QuotaProviderId) => Promise; + /** Resolves true when at least one provider answered — see `ensureLoadedForRuntime`. */ + fetchQuotas: (providerIds: QuotaProviderId[]) => Promise; + /** Resolves true when the instance answered, false on a transport failure. */ + fetchProviderQuota: (providerId: QuotaProviderId) => Promise; setSelectedProvider: (providerId: QuotaProviderId | null) => void; setDisplayMode: (mode: 'usage' | 'remaining') => void; setDropdownProviderIds: (providerIds: QuotaProviderId[]) => void; @@ -104,6 +107,13 @@ const parseSettings = (data: Record | null): QuotaSettingsState }; }; +const defaultQuotaSettings = (): QuotaSettingsState => ({ + displayMode: 'usage', + dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id), + selectedModels: {}, + expandedFamilies: {}, +}); + const loadSettingsFromRuntime = async (): Promise => { const runtimeSettings = getRegisteredRuntimeAPIs()?.settings; if (runtimeSettings) { @@ -127,12 +137,7 @@ const loadSettingsFromRuntime = async (): Promise => { } } - return { - displayMode: 'usage', - dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id), - selectedModels: {}, - expandedFamilies: {}, - }; + return defaultQuotaSettings(); }; export const useQuotaStore = create()( @@ -165,18 +170,20 @@ export const useQuotaStore = create()( const generation = quotaGeneration; set({ isLoading: true, error: null }); try { - await Promise.all( + const answered = await Promise.all( providerIds.map((providerId) => get().fetchProviderQuota(providerId)) ); - if (generation !== quotaGeneration) return; + if (generation !== quotaGeneration) return false; set({ isLoading: false, lastUpdated: Date.now() }); + return answered.some(Boolean); } catch (error) { - if (generation !== quotaGeneration) return; + if (generation !== quotaGeneration) return false; const message = error instanceof Error ? error.message : 'Failed to fetch quotas'; set({ isLoading: false, error: message }); + return false; } }, @@ -196,15 +203,16 @@ export const useQuotaStore = create()( throw new Error(payload?.error || 'Failed to fetch quota'); } - if (generation !== quotaGeneration) return; + if (generation !== quotaGeneration) return false; const result = payload as ProviderResult; set((state) => { const next = state.results.filter((entry) => entry.providerId !== providerId); next.push(result); return { results: next, error: null }; }); + return true; } catch (error) { - if (generation !== quotaGeneration) return; + if (generation !== quotaGeneration) return false; const message = error instanceof Error ? error.message : 'Failed to fetch quota'; const fallback: ProviderResult = { providerId, @@ -220,6 +228,7 @@ export const useQuotaStore = create()( next.push(fallback); return { results: next, error: message }; }); + return false; } finally { if (generation === quotaGeneration) { set((state) => ({ @@ -237,18 +246,34 @@ export const useQuotaStore = create()( // then cached as if it were the answer. if (!useConfigStore.getState().isInitialized) return; if (get().loadedRuntimeKey === runtimeKey) return; - set({ loadedRuntimeKey: runtimeKey }); + if (inFlightRuntimeLoad) return inFlightRuntimeLoad; + const generation = quotaGeneration; - await get().loadSettings(); - if (generation !== quotaGeneration) return; - const { dropdownProviderIds, fetchQuotas } = get(); - if (dropdownProviderIds.length === 0) return; - await fetchQuotas(dropdownProviderIds); + inFlightRuntimeLoad = (async () => { + await get().loadSettings(); + if (generation !== quotaGeneration) return; + const { dropdownProviderIds, fetchQuotas } = get(); + if (dropdownProviderIds.length === 0) return; + const answered = await fetchQuotas(dropdownProviderIds); + // Mark the instance loaded only once it actually answered. Claiming it + // up front meant a load that failed on a cold or briefly unreachable + // instance was never attempted again — Usage would stay empty until + // the three-minute refresh, or forever after a switch. + if (answered && generation === quotaGeneration) set({ loadedRuntimeKey: runtimeKey }); + })().finally(() => { inFlightRuntimeLoad = null; }); + + return inFlightRuntimeLoad; }, resetForRuntimeSwitch: () => { quotaGeneration += 1; + inFlightRuntimeLoad = null; set({ + // Display mode, the provider selection and the per-provider model + // picks all come from the instance's own settings, and + // `dropdownProviderIds` decides what gets fetched — carrying them + // over would query the new instance through the old one's choices. + ...defaultQuotaSettings(), results: [], loadedRuntimeKey: null, selectedProviderId: null, diff --git a/packages/ui/src/stores/useSkillsStore.ts b/packages/ui/src/stores/useSkillsStore.ts index 068eae4c..eee64ca5 100644 --- a/packages/ui/src/stores/useSkillsStore.ts +++ b/packages/ui/src/stores/useSkillsStore.ts @@ -198,6 +198,10 @@ const SKILLS_LOAD_CACHE_TTL_MS = 5000; const DEFAULT_SKILLS_CACHE_KEY = '__default__'; const skillsLastLoadedAt = new Map(); const skillsLoadInFlight = new Map>(); +// Bumped on every runtime switch. Skills are discovered on the connected +// instance and cached by directory, which two instances can share, so a load +// already in flight for the previous instance must not write into the new one. +let skillsGeneration = 0; const getSkillsCacheKey = (directory: string | null): string => { return directory?.trim() || DEFAULT_SKILLS_CACHE_KEY; @@ -286,6 +290,7 @@ export const useSkillsStore = create()( skillDraft: null, resetForRuntimeSwitch: () => { + skillsGeneration += 1; skillsLastLoadedAt.clear(); skillsLoadInFlight.clear(); set({ skills: [], skillsByDirectory: {}, isLoading: false }); @@ -316,6 +321,7 @@ export const useSkillsStore = create()( return inFlight; } + const generation = skillsGeneration; const request = (async () => { set({ isLoading: true }); // Failure must never look like an empty project. The mirror is the @@ -361,6 +367,7 @@ export const useSkillsStore = create()( data.externalSkills ?? null, ); + if (generation !== skillsGeneration) return false; set((state) => { const next: Partial = { skillsByDirectory: { ...state.skillsByDirectory, [cacheKey]: visibleSkills }, @@ -379,6 +386,7 @@ export const useSkillsStore = create()( } console.error("Failed to load skills:", lastError); + if (generation !== skillsGeneration) return false; set((state) => { const next: Partial = { skillsByDirectory: { ...state.skillsByDirectory, [cacheKey]: previousSkills },