fix(storage): await store-file writes instead of polling (#1412) (#1413)

* fix(storage): await store-file writes instead of polling (#1412)

OAuth persistence is Zustand `persist` → file-storage adapter →
writeStoreFile, which the persist middleware invokes fire-and-forget:
saveTokens updates the in-memory store synchronously while the file
write lags. Tests verified persistence by polling the state file with a
2s/50ms `waitForStateFile()` budget, which flaked under full-suite load
(the file was written, the poll just didn't observe it in time).

Make the write completion awaitable instead:
- Track in-flight writeStoreFile() promises per resolved path in
  store-io and add flushStoreFileWrites(filePath?) to await them (also
  useful for graceful shutdown). Writes to the same path are chained so a
  flush awaits every queued write.
- Convert the OAuth E2E, node storage, and file-storage-adapter tests
  from poll loops (waitForStateFile / vi.waitFor on existsSync /
  setTimeout) to `await flushStoreFileWrites(path)` + a single read.
- Remove the now-unused waitForStateFile() test helper.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* chore(storage): doc load-bearing sync registration; rename stale debug env

Address @claude review on #1413:
- Document that writeStoreFile()'s synchronous pendingWrites registration is
  load-bearing for flushStoreFileWrites(), in both the storage adapter's
  setItem callsite and the pendingWrites docblock, so future refactors don't
  introduce an await before registration and silently re-flake the tests.
- Rename the stale DEBUG_WAIT_FOR_STATE_FILE test env var to DEBUG_PERSIST_PATH
  (waitForStateFile was removed in this PR).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

clients/web/src/test/integration/auth/node/storage.test.ts

@@ -12,7 +12,7 @@ import type {
 import * as fs from "node:fs/promises";
 import * as path from "node:path";
 import * as os from "node:os";
-import { waitForStateFile } from "@modelcontextprotocol/inspector-test-server";
+import { flushStoreFileWrites } from "@inspector/core/storage/store-io.js";
 
 // Unique path per process so parallel test files don't share the same state file
 const testStatePath = path.join(
@@ -457,7 +457,7 @@ describe("OAuth Store (Zustand)", () => {
       `mcp-inspector-oauth-persist-${Date.now()}-${Math.random().toString(36).slice(2)}.json`,
     );
     try {
-      if (process.env.DEBUG_WAIT_FOR_STATE_FILE === "1") {
+      if (process.env.DEBUG_PERSIST_PATH === "1") {
         console.error("[storage-node.test] state file path:", persistTestPath);
       }
       const store = getOAuthStore(persistTestPath);
@@ -478,14 +478,11 @@ describe("OAuth Store (Zustand)", () => {
           >;
         };
       };
-      const parsed = await waitForStateFile<StateShape>(
-        persistTestPath,
-        (p) => {
-          const s = (p as StateShape)?.state?.servers?.[serverUrl];
-          return !!s?.clientInformation;
-        },
-        { timeout: 2000, interval: 50 },
-      );
+      // Persistence is fire-and-forget; await the write rather than polling.
+      await flushStoreFileWrites(persistTestPath);
+      const parsed = JSON.parse(
+        await fs.readFile(persistTestPath, "utf-8"),
+      ) as StateShape;
       expect(parsed.state.servers[serverUrl]?.clientInformation).toEqual(
         clientInfo,
       );
@@ -522,14 +519,11 @@ describe("NodeOAuthStorage with custom storagePath", () => {
           servers: Record<string, { tokens?: { access_token?: string } }>;
         };
       };
-      const parsed = await waitForStateFile<StateShape>(
-        customPath,
-        (p) => {
-          const t = (p as StateShape)?.state?.servers?.[testServerUrl]?.tokens;
-          return t?.access_token === tokens.access_token;
-        },
-        { timeout: 2000, interval: 50 },
-      );
+      // Persistence is fire-and-forget; await the write rather than polling.
+      await flushStoreFileWrites(customPath);
+      const parsed = JSON.parse(
+        await fs.readFile(customPath, "utf-8"),
+      ) as StateShape;
 
       expect(parsed.state.servers[testServerUrl]?.tokens?.access_token).toBe(
         tokens.access_token,

clients/web/src/test/integration/mcp/inspectorClient-oauth-e2e.test.ts

@@ -21,7 +21,6 @@ import { FetchRequestLogState } from "@inspector/core/mcp/state/index.js";
 import { createTransportNode } from "@inspector/core/mcp/node/transport.js";
 import {
   TestServerHttp,
-  waitForStateFile,
   waitForOAuthWellKnown,
   getDefaultServerConfig,
   createOAuthTestServerConfig,
@@ -29,6 +28,7 @@ import {
   getDCRRequests,
   invalidateAccessToken,
 } from "@modelcontextprotocol/inspector-test-server";
+import { flushStoreFileWrites } from "@inspector/core/storage/store-io.js";
 import {
   createOAuthClientConfig,
   completeOAuthAuthorization,
@@ -1752,21 +1752,16 @@ describe("InspectorClient OAuth E2E", () => {
             servers?: Record<string, { tokens?: { access_token?: string } }>;
           };
         };
-        const parsed = await waitForStateFile<StateShape>(
-          customPath,
-          (p) => {
-            const servers = (p as StateShape)?.state?.servers ?? {};
-            return Object.values(servers).some(
-              (s) =>
-                !!(s as { tokens?: { access_token?: string } })?.tokens
-                  ?.access_token,
-            );
-          },
-          { timeout: 2000, interval: 50 },
-        );
-        expect(Object.keys(parsed.state?.servers ?? {}).length).toBeGreaterThan(
-          0,
-        );
+        // Persistence is fire-and-forget; await the write rather than polling.
+        await flushStoreFileWrites(customPath);
+        const parsed = JSON.parse(
+          await fs.readFile(customPath, "utf-8"),
+        ) as StateShape;
+        const servers = parsed.state?.servers ?? {};
+        expect(Object.keys(servers).length).toBeGreaterThan(0);
+        expect(
+          Object.values(servers).some((s) => !!s?.tokens?.access_token),
+        ).toBe(true);
       } finally {
         try {
           await fs.unlink(customPath);

clients/web/src/test/integration/storage/adapters.test.ts

@@ -13,6 +13,10 @@ import { createFileStorageAdapter } from "@inspector/core/storage/adapters/file-
 import { createRemoteStorageAdapter } from "@inspector/core/storage/adapters/remote-storage.js";
 import { createOAuthStore } from "@inspector/core/auth/store.js";
 import { createRemoteApp } from "@inspector/core/mcp/remote/node/server.js";
+import {
+  writeStoreFile,
+  flushStoreFileWrites,
+} from "@inspector/core/storage/store-io.js";
 
 interface StartRemoteServerOptions {
   storageDir?: string;
@@ -147,13 +151,8 @@ describe("Storage adapters", () => {
         tokens: { access_token: "test-token", token_type: "Bearer" },
       });
 
-      // Wait for persistence (Zustand persist is async; poll for file so we don't race with cleanup)
-      await vi.waitFor(
-        () => {
-          expect(existsSync(filePath)).toBe(true);
-        },
-        { timeout: 2000, interval: 20 },
-      );
+      // Persistence is fire-and-forget; await the write rather than polling.
+      await flushStoreFileWrites(filePath);
       const fileContent = readFileSync(filePath, "utf-8");
       const parsed = JSON.parse(fileContent);
       expect(parsed.state.servers["https://example.com"].tokens).toEqual({
@@ -172,12 +171,7 @@ describe("Storage adapters", () => {
       store1.getState().setServerState("https://example.com", {
         tokens: { access_token: "initial-token", token_type: "Bearer" },
       });
-      await vi.waitFor(
-        () => {
-          expect(existsSync(filePath)).toBe(true);
-        },
-        { timeout: 2000, interval: 20 },
-      );
+      await flushStoreFileWrites(filePath);
 
       // Create new store instance (should load persisted state)
       const storage2 = createFileStorageAdapter({ filePath });
@@ -201,7 +195,7 @@ describe("Storage adapters", () => {
       store.getState().setServerState("https://example.com", {
         tokens: { access_token: "test-token", token_type: "Bearer" },
       });
-      await new Promise((resolve) => setTimeout(resolve, 100));
+      await flushStoreFileWrites(filePath);
       expect(existsSync(filePath)).toBe(true);
 
       // Clear all servers (this will persist empty state)
@@ -210,7 +204,7 @@ describe("Storage adapters", () => {
       for (const url of urls) {
         state.clearServerState(url);
       }
-      await new Promise((resolve) => setTimeout(resolve, 100));
+      await flushStoreFileWrites(filePath);
 
       // Verify file still exists but with empty servers
       expect(existsSync(filePath)).toBe(true);
@@ -229,15 +223,75 @@ describe("Storage adapters", () => {
       store.getState().setServerState("https://example.com", {
         tokens: { access_token: "t", token_type: "Bearer" },
       });
-      await vi.waitFor(() => expect(existsSync(filePath)).toBe(true), {
-        timeout: 2000,
-        interval: 20,
-      });
+      await flushStoreFileWrites(filePath);
+      expect(existsSync(filePath)).toBe(true);
       await storage!.removeItem("inspector-oauth-store");
       expect(existsSync(filePath)).toBe(false);
     });
   });
 
+  describe("flushStoreFileWrites", () => {
+    let tempDir: string | null = null;
+
+    afterEach(() => {
+      if (tempDir) {
+        try {
+          rmSync(tempDir, { recursive: true });
+        } catch {
+          // Ignore cleanup errors
+        }
+        tempDir = null;
+      }
+    });
+
+    it("resolves immediately when nothing is in flight", async () => {
+      await expect(flushStoreFileWrites()).resolves.toBeUndefined();
+      await expect(
+        flushStoreFileWrites("/no/such/path.json"),
+      ).resolves.toBeUndefined();
+    });
+
+    it("awaits the pending write for a specific path", async () => {
+      tempDir = mkdtempSync(join(tmpdir(), "inspector-flush-test-"));
+      const filePath = join(tempDir, "store.json");
+
+      // Kick off the write without awaiting it, then flush.
+      const write = writeStoreFile(filePath, '{"hello":"world"}');
+      await flushStoreFileWrites(filePath);
+      expect(existsSync(filePath)).toBe(true);
+      expect(JSON.parse(readFileSync(filePath, "utf-8"))).toEqual({
+        hello: "world",
+      });
+      await write;
+    });
+
+    it("awaits all pending writes when no path is given", async () => {
+      tempDir = mkdtempSync(join(tmpdir(), "inspector-flush-test-"));
+      const a = join(tempDir, "a.json");
+      const b = join(tempDir, "b.json");
+
+      const writes = Promise.all([
+        writeStoreFile(a, '{"n":1}'),
+        writeStoreFile(b, '{"n":2}'),
+      ]);
+      await flushStoreFileWrites();
+      expect(existsSync(a)).toBe(true);
+      expect(existsSync(b)).toBe(true);
+      await writes;
+    });
+
+    it("serializes overlapping writes to the same path (last write wins)", async () => {
+      tempDir = mkdtempSync(join(tmpdir(), "inspector-flush-test-"));
+      const filePath = join(tempDir, "store.json");
+
+      const first = writeStoreFile(filePath, '{"v":1}');
+      const second = writeStoreFile(filePath, '{"v":2}');
+      await Promise.all([first, second]);
+      await flushStoreFileWrites(filePath);
+      expect(JSON.parse(readFileSync(filePath, "utf-8"))).toEqual({ v: 2 });
+    });
+  });
+
   describe("RemoteStorageAdapter", () => {
     let remoteServer: ServerType | null = null;
     let tempDir: string | null = null;

core/storage/adapters/file-storage.ts

@@ -20,6 +20,11 @@ export function createFileStorageAdapter(
 ): ReturnType<typeof createJSONStorage> {
   return createJSONStorage(() => ({
     getItem: async () => readStoreFile(options.filePath),
+    // Do not introduce an `await` before writeStoreFile() here: it registers
+    // the write in pendingWrites synchronously, which is load-bearing for
+    // flushStoreFileWrites() (a microtask hop before registration would make a
+    // flush called right after setItem find an empty map and return early,
+    // silently regressing tests to non-deterministic).
     setItem: async (_name: string, value: string) =>
       writeStoreFile(options.filePath, value),
     removeItem: async () => deleteStoreFile(options.filePath),

core/storage/store-io.ts

@@ -57,20 +57,66 @@ export async function readStoreFile(filePath: string): Promise<string | null> {
   }
 }
 
+/**
+ * In-flight writeStoreFile() promises, keyed by resolved path. Lets callers
+ * await persistence completion via flushStoreFileWrites() instead of polling
+ * the file — Zustand's persist middleware invokes writeStoreFile() fire-and-
+ * forget, so the in-memory store updates synchronously while the file write
+ * lags. Entries are removed once their write settles.
+ *
+ * Load-bearing: pendingWrites.set() below runs synchronously before the first
+ * await in writeStoreFile(), so a flushStoreFileWrites() called right after a
+ * persist sees the in-flight entry. Callers (e.g. the storage adapter's
+ * setItem) must not introduce an await before writeStoreFile() — doing so would
+ * let a flush run before registration and return early.
+ */
+const pendingWrites = new Map<string, Promise<void>>();
+
 /**
  * Write store file atomically (temp file + rename). Ensures parent directory exists.
  * Uses mode 0o600 for the file.
+ *
+ * The returned promise is also tracked per path so flushStoreFileWrites() can
+ * await it. Writes to the same path are chained so a flush awaits every queued
+ * write (and a later write's completion implies all earlier ones settled).
  */
 export async function writeStoreFile(
   filePath: string,
   data: string,
 ): Promise<void> {
-  const dir = path.dirname(filePath);
-  await fs.mkdir(dir, { recursive: true });
-  await writeFile(filePath, data, {
-    encoding: "utf-8",
-    mode: 0o600,
-  });
+  const key = path.resolve(filePath);
+  const run = async (): Promise<void> => {
+    const dir = path.dirname(filePath);
+    await fs.mkdir(dir, { recursive: true });
+    await writeFile(filePath, data, {
+      encoding: "utf-8",
+      mode: 0o600,
+    });
+  };
+  const prior = pendingWrites.get(key);
+  const tracked = (prior ?? Promise.resolve()).catch(() => {}).then(run);
+  pendingWrites.set(key, tracked);
+  try {
+    await tracked;
+  } finally {
+    if (pendingWrites.get(key) === tracked) {
+      pendingWrites.delete(key);
+    }
+  }
+}
+
+/**
+ * Await pending writeStoreFile() writes — those for `filePath` if given, else
+ * all of them. Use in tests after triggering persistence (Zustand persist
+ * writes fire-and-forget) instead of polling the file, and for graceful
+ * shutdown. Resolves immediately when nothing is in flight.
+ */
+export async function flushStoreFileWrites(filePath?: string): Promise<void> {
+  if (filePath !== undefined) {
+    await pendingWrites.get(path.resolve(filePath));
+    return;
+  }
+  await Promise.all(pendingWrites.values());
 }
 
 /**

test-servers/src/test-helpers.ts

@@ -4,7 +4,6 @@
  */
 
 import { vi } from "vitest";
-import * as fs from "node:fs/promises";
 
 export interface WaitForEventOptions {
   timeout?: number;
@@ -75,89 +74,6 @@ export function waitForProgressCount(
   });
 }
 
-export interface WaitForStateFileOptions {
-  timeout?: number;
-  interval?: number;
-}
-
-const DEBUG_WAIT_FOR_STATE_FILE = process.env.DEBUG_WAIT_FOR_STATE_FILE === "1";
-
-function truncate(s: string, maxLen: number): string {
-  if (s.length <= maxLen) return s;
-  return s.slice(0, maxLen) + `... (${s.length} chars total)`;
-}
-
-/**
- * Poll state file until `predicate(parsed)` returns true, then return the parsed value.
- * Uses vi.waitFor under the hood. For use with Zustand persist state.json files.
- *
- * On failure, the thrown error includes:
- * - Whether the failure was a JSON parse error or predicate returned false.
- * - A truncated snippet of what was read (to distinguish partial write vs wrong content).
- * - Attempt count (to see if we timed out early or after many retries).
- *
- * Run with DEBUG_WAIT_FOR_STATE_FILE=1 to log every attempt (parse ok/fail, predicate result).
- */
-export async function waitForStateFile<T = unknown>(
-  filePath: string,
-  predicate: (parsed: unknown) => boolean,
-  options?: WaitForStateFileOptions,
-): Promise<T> {
-  const { timeout = 2000, interval = 50 } = options ?? {};
-  let result: T | undefined;
-  let attemptCount = 0;
-
-  await vi.waitFor(
-    async () => {
-      attemptCount += 1;
-      let raw: string;
-      try {
-        raw = await fs.readFile(filePath, "utf-8");
-      } catch (readErr) {
-        const msg = (readErr as NodeJS.ErrnoException).code ?? String(readErr);
-        if (DEBUG_WAIT_FOR_STATE_FILE) {
-          console.error(
-            `[waitForStateFile] attempt ${attemptCount} read failed:`,
-            msg,
-          );
-        }
-        throw new Error(
-          `waitForStateFile failed: file read error (${msg}). File: ${filePath}. Attempts: ${attemptCount}. Run with DEBUG_WAIT_FOR_STATE_FILE=1 for per-attempt logs.`,
-        );
-      }
-      let parsed: unknown;
-      try {
-        parsed = JSON.parse(raw) as unknown;
-      } catch {
-        if (DEBUG_WAIT_FOR_STATE_FILE) {
-          console.error(
-            `[waitForStateFile] attempt ${attemptCount} JSON parse failed. Raw (first 300):`,
-            truncate(raw, 300),
-          );
-        }
-        throw new Error(
-          `waitForStateFile failed: JSON parse error (file may be mid-write or corrupt). File: ${filePath}. Attempts: ${attemptCount}. Raw snippet: ${truncate(raw, 200)}. Run with DEBUG_WAIT_FOR_STATE_FILE=1 for per-attempt logs.`,
-        );
-      }
-      const predOk = predicate(parsed);
-      if (DEBUG_WAIT_FOR_STATE_FILE) {
-        console.error(
-          `[waitForStateFile] attempt ${attemptCount} parse ok, predicate: ${predOk}`,
-        );
-      }
-      if (!predOk) {
-        throw new Error(
-          `waitForStateFile failed: predicate returned false. File: ${filePath}. Attempts: ${attemptCount}. Parsed snippet: ${truncate(JSON.stringify(parsed), 200)}. Run with DEBUG_WAIT_FOR_STATE_FILE=1 for per-attempt logs.`,
-        );
-      }
-      result = parsed as T;
-      return true;
-    },
-    { timeout, interval },
-  );
-  return result!;
-}
-
 export interface WaitForOAuthWellKnownOptions {
   timeout?: number;
   interval?: number;