fix: reject instance-served responses that outlive their instance

MCP status and skills were cleared on a switch but their in-flight requests were not, so a response for the previous instance could still write itself over the new one's — the same race the quota store already guards. Both now carry a generation. Usage also claimed an instance as loaded before it had answered, so a load that failed on a cold or briefly unreachable instance was never attempted again, and the previous instance's display mode and provider selection survived a switch — which decided what the new instance was even asked for.
2026-09-03 11:49:34 +03:00
parent 17fd46d4e5
commit f0ced9d361
5 changed files with 191 additions and 19 deletions
@@ -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<T> = { promise: Promise<T>; resolve: (value: T) => void };
+const deferred = <T>(): Deferred<T> => {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((res) => { resolve = res; });
+  return { promise, resolve };
+};
+
+type McpStatusResult = Awaited<ReturnType<ReturnType<typeof opencodeModule.opencodeClient.getApiClient>['mcp']['status']>>;
+let mcpStatusResponse: Deferred<McpStatusResult> = 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<typeof opencodeModule.opencodeClient.getApiClient>;
+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<Response> = 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);
+  });
+});
@@ -54,6 +54,10 @@ type RefreshOptions = {
 };

 const ensureFreshInFlight = new Map<string, Promise<void>>();
+// 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<McpStore>()(
    refreshedAtKeys: {},

    resetForRuntimeSwitch: () => {
+      mcpGeneration += 1;
      ensureFreshInFlight.clear();
      set({
        byDirectory: {},
@@ -144,9 +149,11 @@ export const useMcpStore = create<McpStore>()(
        }));
      }

+      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<McpStore>()(
          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 },
@@ -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")
+  })
 })
@@ -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<void> | null = null;
 let quotaAutoRefreshConsumers = 0;
 let quotaAutoRefreshInterval: number | null = null;

@@ -38,8 +39,10 @@ interface QuotaStore extends QuotaSettingsState {

  loadSettings: () => Promise<void>;
  fetchAllQuotas: () => Promise<void>;
-  fetchQuotas: (providerIds: QuotaProviderId[]) => Promise<void>;
-  fetchProviderQuota: (providerId: QuotaProviderId) => Promise<void>;
+  /** Resolves true when at least one provider answered — see `ensureLoadedForRuntime`. */
+  fetchQuotas: (providerIds: QuotaProviderId[]) => Promise<boolean>;
+  /** Resolves true when the instance answered, false on a transport failure. */
+  fetchProviderQuota: (providerId: QuotaProviderId) => Promise<boolean>;
  setSelectedProvider: (providerId: QuotaProviderId | null) => void;
  setDisplayMode: (mode: 'usage' | 'remaining') => void;
  setDropdownProviderIds: (providerIds: QuotaProviderId[]) => void;
@@ -104,6 +107,13 @@ const parseSettings = (data: Record<string, unknown> | null): QuotaSettingsState
  };
 };

+const defaultQuotaSettings = (): QuotaSettingsState => ({
+  displayMode: 'usage',
+  dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id),
+  selectedModels: {},
+  expandedFamilies: {},
+});
+
 const loadSettingsFromRuntime = async (): Promise<QuotaSettingsState> => {
  const runtimeSettings = getRegisteredRuntimeAPIs()?.settings;
  if (runtimeSettings) {
@@ -127,12 +137,7 @@ const loadSettingsFromRuntime = async (): Promise<QuotaSettingsState> => {
    }
  }

-  return {
-    displayMode: 'usage',
-    dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id),
-    selectedModels: {},
-    expandedFamilies: {},
-  };
+  return defaultQuotaSettings();
 };

 export const useQuotaStore = create<QuotaStore>()(
@@ -165,18 +170,20 @@ export const useQuotaStore = create<QuotaStore>()(
        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<QuotaStore>()(
            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<QuotaStore>()(
            next.push(fallback);
            return { results: next, error: message };
          });
+          return false;
        } finally {
          if (generation === quotaGeneration) {
            set((state) => ({
@@ -237,18 +246,34 @@ export const useQuotaStore = create<QuotaStore>()(
        // 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,
@@ -198,6 +198,10 @@ const SKILLS_LOAD_CACHE_TTL_MS = 5000;
 const DEFAULT_SKILLS_CACHE_KEY = '__default__';
 const skillsLastLoadedAt = new Map<string, number>();
 const skillsLoadInFlight = new Map<string, Promise<boolean>>();
+// 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<SkillsStore>()(
        skillDraft: null,

        resetForRuntimeSwitch: () => {
+          skillsGeneration += 1;
          skillsLastLoadedAt.clear();
          skillsLoadInFlight.clear();
          set({ skills: [], skillsByDirectory: {}, isLoading: false });
@@ -316,6 +321,7 @@ export const useSkillsStore = create<SkillsStore>()(
            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<SkillsStore>()(
                  data.externalSkills ?? null,
                );

+                if (generation !== skillsGeneration) return false;
                set((state) => {
                  const next: Partial<SkillsStore> = {
                    skillsByDirectory: { ...state.skillsByDirectory, [cacheKey]: visibleSkills },
@@ -379,6 +386,7 @@ export const useSkillsStore = create<SkillsStore>()(
            }

            console.error("Failed to load skills:", lastError);
+            if (generation !== skillsGeneration) return false;
            set((state) => {
              const next: Partial<SkillsStore> = {
                skillsByDirectory: { ...state.skillsByDirectory, [cacheKey]: previousSkills },