From e1c8f7192425c23e66ed386cedec95f16e21e0b8 Mon Sep 17 00:00:00 2001 From: Cliff Hall Date: Tue, 2 Jun 2026 11:55:24 -0400 Subject: [PATCH] fix(storage): await store-file writes instead of polling (#1412) (#1413) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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) * 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) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../integration/auth/node/storage.test.ts | 30 +++--- .../mcp/inspectorClient-oauth-e2e.test.ts | 27 +++--- .../test/integration/storage/adapters.test.ts | 92 +++++++++++++++---- core/storage/adapters/file-storage.ts | 5 + core/storage/store-io.ts | 58 ++++++++++-- test-servers/src/test-helpers.ts | 84 ----------------- 6 files changed, 153 insertions(+), 143 deletions(-) diff --git a/clients/web/src/test/integration/auth/node/storage.test.ts b/clients/web/src/test/integration/auth/node/storage.test.ts index 1cf79d4b..12e317f3 100644 --- a/clients/web/src/test/integration/auth/node/storage.test.ts +++ b/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( - 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; }; }; - const parsed = await waitForStateFile( - 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, diff --git a/clients/web/src/test/integration/mcp/inspectorClient-oauth-e2e.test.ts b/clients/web/src/test/integration/mcp/inspectorClient-oauth-e2e.test.ts index e87e2381..14d177cb 100644 --- a/clients/web/src/test/integration/mcp/inspectorClient-oauth-e2e.test.ts +++ b/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; }; }; - const parsed = await waitForStateFile( - 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); diff --git a/clients/web/src/test/integration/storage/adapters.test.ts b/clients/web/src/test/integration/storage/adapters.test.ts index fc6d00b6..8c248a3f 100644 --- a/clients/web/src/test/integration/storage/adapters.test.ts +++ b/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; diff --git a/core/storage/adapters/file-storage.ts b/core/storage/adapters/file-storage.ts index 5d45c7bd..079b2548 100644 --- a/core/storage/adapters/file-storage.ts +++ b/core/storage/adapters/file-storage.ts @@ -20,6 +20,11 @@ export function createFileStorageAdapter( ): ReturnType { 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), diff --git a/core/storage/store-io.ts b/core/storage/store-io.ts index 79c53fcf..d42129ef 100644 --- a/core/storage/store-io.ts +++ b/core/storage/store-io.ts @@ -57,20 +57,66 @@ export async function readStoreFile(filePath: string): Promise { } } +/** + * 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>(); + /** * 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 { - 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 => { + 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 { + if (filePath !== undefined) { + await pendingWrites.get(path.resolve(filePath)); + return; + } + await Promise.all(pendingWrites.values()); } /** diff --git a/test-servers/src/test-helpers.ts b/test-servers/src/test-helpers.ts index 081d8a26..244d4fe2 100644 --- a/test-servers/src/test-helpers.ts +++ b/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( - filePath: string, - predicate: (parsed: unknown) => boolean, - options?: WaitForStateFileOptions, -): Promise { - 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;