perf: add preloadSchemas() and warm wire schemas eagerly on workerd (#2483)

packages/server/package.json

@@ -73,12 +73,12 @@
             },
             "browser": {
                 "import": {
-                    "types": "./dist/shimsWorkerd.d.mts",
-                    "default": "./dist/shimsWorkerd.mjs"
+                    "types": "./dist/shimsBrowser.d.mts",
+                    "default": "./dist/shimsBrowser.mjs"
                 },
                 "require": {
-                    "types": "./dist/shimsWorkerd.d.cts",
-                    "default": "./dist/shimsWorkerd.cjs"
+                    "types": "./dist/shimsBrowser.d.cts",
+                    "default": "./dist/shimsBrowser.cjs"
                 }
             },
             "node": {

packages/server/src/index.ts

@@ -99,5 +99,11 @@ export type { CacheHint, CacheScope } from '@modelcontextprotocol/core-internal'
 export type { ElicitInputParams, InputRequiredSpec, InputResponseView } from '@modelcontextprotocol/core-internal';
 export { acceptedContent, inputRequired, inputResponse } from '@modelcontextprotocol/core-internal';
 
+// Explicit opt-in to eager wire-schema construction, for platforms that bill
+// request CPU but not module evaluation (isolate-based edge/serverless
+// runtimes). The package's workerd build calls it automatically at module
+// scope; other builds stay lazy unless the application calls it itself.
+export { preloadSchemas } from '@modelcontextprotocol/core-internal';
+
 // re-export curated public API from core
 export * from '@modelcontextprotocol/core-internal/public';

packages/server/src/shimsBrowser.ts

@@ -0,0 +1,27 @@
+/**
+ * Browser runtime shims for server package
+ *
+ * This file is selected via package.json export conditions when bundling for
+ * browsers. It binds the same platform choices as the workerd shim (the
+ * cfWorker validator, the process stub) WITHOUT the module-scope
+ * `preloadSchemas()` call: in a browser, module evaluation is page load —
+ * boot latency — so schema construction stays lazy, exactly like Node.
+ */
+export { CfWorkerJsonSchemaValidator as DefaultJsonSchemaValidator } from '@modelcontextprotocol/core-internal/validators/cfWorker';
+
+/**
+ * Stub process object for non-Node.js environments.
+ * StdioServerTransport is not supported in Cloudflare Workers/browser environments.
+ */
+function notSupported(): never {
+    throw new Error('StdioServerTransport is not supported in this environment. Use StreamableHTTPServerTransport instead.');
+}
+
+export const process = {
+    get stdin(): never {
+        return notSupported();
+    },
+    get stdout(): never {
+        return notSupported();
+    }
+};

packages/server/src/shimsWorkerd.ts

@@ -3,8 +3,17 @@
  *
  * This file is selected via package.json export conditions when running in workerd.
  */
+import { preloadSchemas } from '@modelcontextprotocol/core-internal';
+
 export { CfWorkerJsonSchemaValidator as DefaultJsonSchemaValidator } from '@modelcontextprotocol/core-internal/validators/cfWorker';
 
+// Platform asymmetry: isolate platforms like workerd evaluate module scope
+// during deployment/isolate warm-up, outside any request's billed CPU, while
+// lazy construction would land inside the first request each fresh isolate
+// serves. The Node and browser shims stay lazy — there, module evaluation is
+// process/page startup and boot latency is the cost that matters.
+preloadSchemas();
+
 /**
  * Stub process object for non-Node.js environments.
  * StdioServerTransport is not supported in Cloudflare Workers/browser environments.

packages/server/test/helpers/ensureBuilt.ts

@@ -0,0 +1,105 @@
+/**
+ * Server-package copy of packages/client/test/helpers/ensureBuilt.ts (per-package
+ * test dirs stay self-contained). Keep the locking logic in sync with that file.
+ */
+import { execFile } from 'node:child_process';
+import { existsSync, mkdirSync, rmSync, statSync } from 'node:fs';
+import { basename, join } from 'node:path';
+import { setTimeout as sleep } from 'node:timers/promises';
+import { promisify } from 'node:util';
+
+const execFileAsync = promisify(execFile);
+
+/**
+ * Canonical dist sentinels per package. Every caller waits on the same
+ * complete set: if callers checked their own subsets, a partial dist (from an
+ * interrupted build) could satisfy one caller's fast path while another still
+ * rebuilds — and the rebuild's clean step would wipe files under the first
+ * caller's running tests, which is exactly the race this helper exists to
+ * prevent.
+ */
+const DIST_SENTINELS: Record<string, string[]> = {
+    server: [
+        'index.mjs',
+        'stdio.mjs',
+        'shimsWorkerd.mjs',
+        'shimsWorkerd.cjs',
+        'shimsNode.mjs',
+        'shimsNode.cjs',
+        'shimsBrowser.mjs',
+        'shimsBrowser.cjs'
+    ]
+};
+
+/** The build is killed after this long; execFile's kill still runs our finally. */
+const BUILD_TIMEOUT_MS = 60_000;
+/** How long a waiter polls for another worker's build — outlasts BUILD_TIMEOUT_MS. */
+const WAIT_DEADLINE_MS = 90_000;
+/** A lock older than this belongs to a dead worker (finally never ran) — steal it. */
+const STALE_LOCK_MS = 120_000;
+
+/**
+ * Build a package's dist on demand, safely across parallel vitest workers.
+ *
+ * Vitest runs test files in separate worker processes, so two files that each
+ * "build if dist is missing" can race on a cold checkout: both see no dist,
+ * both spawn `pnpm build`, and tsdown's clean step deletes files out from
+ * under whichever worker is already reading them. This helper makes the build
+ * single-flight with an atomic mkdir lock: the first worker builds, everyone
+ * else waits for the sentinel files to appear and the lock to clear, and once
+ * the dist exists nobody ever rebuilds (so no later clean can wipe it
+ * mid-read).
+ */
+export async function ensureBuilt(pkgDir: string): Promise<void> {
+    const sentinels = DIST_SENTINELS[basename(pkgDir)];
+    if (!sentinels) {
+        throw new Error(`No dist sentinels registered for ${pkgDir} — add the package to DIST_SENTINELS`);
+    }
+    const lockDir = join(pkgDir, '.dist-build-lock');
+    const haveAll = () => sentinels.every(s => existsSync(join(pkgDir, 'dist', s)));
+    const deadline = Date.now() + WAIT_DEADLINE_MS;
+    for (;;) {
+        if (haveAll() && !existsSync(lockDir)) return;
+        try {
+            mkdirSync(lockDir);
+        } catch (err) {
+            if ((err as NodeJS.ErrnoException).code !== 'EEXIST') throw err;
+            // Another worker holds the lock. If the holder died without its
+            // finally running (SIGKILL/OOM), the lock never clears — steal it
+            // once it is older than any live build could be.
+            try {
+                if (Date.now() - statSync(lockDir).mtimeMs > STALE_LOCK_MS) {
+                    rmSync(lockDir, { recursive: true, force: true });
+                    continue;
+                }
+            } catch {
+                continue; // lock vanished between mkdir and stat — re-check now
+            }
+            if (Date.now() > deadline) {
+                throw new Error(
+                    `Timed out waiting for ${pkgDir}/dist (${sentinels.join(', ')}) ` +
+                        `while another worker held the build lock at ${lockDir}`
+                );
+            }
+            await sleep(250);
+            continue;
+        }
+        try {
+            if (!haveAll()) {
+                try {
+                    await execFileAsync('pnpm', ['build'], {
+                        cwd: pkgDir,
+                        timeout: BUILD_TIMEOUT_MS,
+                        maxBuffer: 16 * 1024 * 1024
+                    });
+                } catch (err) {
+                    const stderr = (err as { stderr?: string }).stderr ?? '';
+                    throw new Error(`pnpm build failed in ${pkgDir}: ${(err as Error).message}\n${stderr.slice(-2000)}`);
+                }
+            }
+            return;
+        } finally {
+            rmSync(lockDir, { recursive: true, force: true });
+        }
+    }
+}

packages/server/test/server/barrelClean.test.ts

@@ -1,11 +1,12 @@
-import { execFileSync } from 'node:child_process';
-import { existsSync, readFileSync } from 'node:fs';
+import { readFileSync } from 'node:fs';
 import { createRequire } from 'node:module';
 import { dirname, join } from 'node:path';
 import { fileURLToPath } from 'node:url';
 
 import { beforeAll, describe, expect, test } from 'vitest';
 
+import { ensureBuilt } from '../helpers/ensureBuilt';
+
 const pkgDir = join(dirname(fileURLToPath(import.meta.url)), '../..');
 const distDir = join(pkgDir, 'dist');
 const requireDist = createRequire(join(pkgDir, 'package.json'));
@@ -37,11 +38,12 @@ function rootExportBlockOf(content: string): string {
 }
 
 describe('@modelcontextprotocol/server root entry is browser-safe', () => {
-    beforeAll(() => {
-        if (!existsSync(join(distDir, 'index.mjs')) || !existsSync(join(distDir, 'stdio.mjs'))) {
-            execFileSync('pnpm', ['build'], { cwd: pkgDir, stdio: 'inherit' });
-        }
-    }, 60_000);
+    beforeAll(async () => {
+        // Single-flight across parallel vitest workers: this file and
+        // workerdSchemaPreload.test.ts both read the built dist, and two
+        // unlocked `pnpm build`s would race tsdown's clean step.
+        await ensureBuilt(pkgDir);
+    }, 180_000);
 
     test('dist/index.mjs does not export StdioServerTransport and has no process-stdio runtime imports', () => {
         const entry = readFileSync(join(distDir, 'index.mjs'), 'utf8');
@@ -66,7 +68,7 @@ describe('@modelcontextprotocol/server root entry is browser-safe', () => {
     });
 
     test('runtime shims vendor default validator backends instead of requiring consumers to install them', () => {
-        for (const shim of ['shimsNode.mjs', 'shimsWorkerd.mjs']) {
+        for (const shim of ['shimsNode.mjs', 'shimsWorkerd.mjs', 'shimsBrowser.mjs']) {
             const entry = join(distDir, shim);
             expect(readFileSync(entry, 'utf8')).not.toMatch(VALIDATOR_BACKEND_IMPORT);
 

packages/server/test/server/workerdSchemaPreload.test.ts

@@ -0,0 +1,90 @@
+/**
+ * Platform-conditional schema warm-up pins, asserted against the real build
+ * outputs.
+ *
+ * The wire schemas are built lazily by default — the right trade on
+ * process-per-invocation runtimes, where module evaluation is boot latency.
+ * On isolate platforms (workerd), module scope evaluates during isolate
+ * warm-up outside any request's billed CPU, so the workerd shim calls
+ * `preloadSchemas()` at module scope to keep construction out of the first
+ * request. Two regressions would be silent without these pins:
+ *
+ * 1. The workerd condition loses its module-scope call (a shim refactor drops
+ *    it) — fresh isolates go back to paying schema construction inside the
+ *    first request's CPU.
+ * 2. The node or browser condition gains one (an import shuffle re-eagerizes
+ *    it) — every process boot / page load pays construction for validations
+ *    that may never happen.
+ *
+ * Identity also matters: the warm-up only helps if the shim forces the SAME
+ * memo module the package entry validates through. The shim entry must
+ * import `preloadSchemas` from the shared chunk — a duplicated definition in
+ * the shim chunk would warm a twin module and leave the real one cold.
+ */
+import { readFileSync } from 'node:fs';
+import { dirname, join } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+import { beforeAll, describe, expect, test } from 'vitest';
+
+import { ensureBuilt } from '../helpers/ensureBuilt';
+
+const pkgDir = join(dirname(fileURLToPath(import.meta.url)), '../..');
+const distDir = join(pkgDir, 'dist');
+
+/**
+ * A module-scope `preloadSchemas();` call: statement at column 0 — entry-level
+ * statements are unindented in the unminified build output, while statements
+ * inside function bodies are indented, so a call that merely sits in some
+ * bundled helper body cannot satisfy (or trip) this pin. The CJS build calls
+ * through the required chunk's namespace (`require_src.preloadSchemas();`),
+ * so an optional receiver is allowed.
+ */
+const MODULE_SCOPE_PRELOAD_CALL = /^(?:[\w$]+\.)?preloadSchemas\(\);/m;
+
+/**
+ * An entry-level `import { …, X as preloadSchemas, … } from "./chunk"` (or
+ * the un-aliased `{ preloadSchemas }` form), capturing the chunk specifier.
+ */
+const PRELOAD_IMPORT = /import\s*\{[^}]*\bpreloadSchemas\b[^}]*\}\s*from\s*"(\.\/[^"]+)"/;
+
+function dist(file: string): string {
+    return readFileSync(join(distDir, file), 'utf8');
+}
+
+describe('workerd schema warm-up (built dist)', () => {
+    beforeAll(async () => {
+        await ensureBuilt(pkgDir);
+    }, 180_000);
+
+    test('shimsWorkerd calls preloadSchemas() at module scope (ESM and CJS)', () => {
+        expect(dist('shimsWorkerd.mjs')).toMatch(MODULE_SCOPE_PRELOAD_CALL);
+        expect(dist('shimsWorkerd.cjs')).toMatch(MODULE_SCOPE_PRELOAD_CALL);
+    });
+
+    test('node and browser shims stay lazy — no preloadSchemas reference at all', () => {
+        for (const shim of ['shimsNode.mjs', 'shimsNode.cjs', 'shimsBrowser.mjs', 'shimsBrowser.cjs']) {
+            expect(dist(shim), `${shim} must not re-eagerize schema construction`).not.toMatch(/\bpreloadSchemas\b/);
+        }
+    });
+
+    test('the root entry exports preloadSchemas but never calls it at module scope', () => {
+        const index = dist('index.mjs');
+        expect(index).toMatch(/\bpreloadSchemas\b/);
+        expect(index).not.toMatch(MODULE_SCOPE_PRELOAD_CALL);
+    });
+
+    test('the workerd shim warms the same chunk the root entry uses (no duplicated schema graph)', () => {
+        const entries = ['index.mjs', 'shimsWorkerd.mjs', 'shimsNode.mjs', 'shimsBrowser.mjs', 'stdio.mjs'];
+        for (const entry of entries) {
+            expect(dist(entry), `${entry} must not carry its own preloadSchemas definition`).not.toMatch(/function preloadSchemas\b/);
+        }
+
+        const shimSource = dist('shimsWorkerd.mjs').match(PRELOAD_IMPORT)?.[1];
+        const indexSource = dist('index.mjs').match(PRELOAD_IMPORT)?.[1];
+        expect(shimSource).toBeDefined();
+        expect(indexSource).toBeDefined();
+        expect(shimSource).toBe(indexSource);
+        expect(dist(shimSource!.replace('./', ''))).toMatch(/function preloadSchemas\b/);
+    });
+});

packages/server/tsdown.config.ts

@@ -7,6 +7,7 @@ export default defineConfig({
         'src/stdio.ts',
         'src/shimsNode.ts',
         'src/shimsWorkerd.ts',
+        'src/shimsBrowser.ts',
         'src/validators/ajv.ts',
         'src/validators/cfWorker.ts'
     ],