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

packages/client/src/index.ts

@@ -106,5 +106,11 @@ export { fromJsonSchema } from './fromJsonSchema';
 export type { InputRequiredOptions } from '@modelcontextprotocol/core-internal';
 export { withInputRequired } 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/client/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. Node and browser shims stay lazy — there, module evaluation is
+// process/page startup and boot latency is the cost that matters.
+preloadSchemas();
+
 /**
  * Whether `fetch()` may throw `TypeError` due to CORS. CORS is a browser-only concept —
  * in Cloudflare Workers, a `TypeError` from `fetch` is always a real network/configuration

packages/client/test/client/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/client/test/helpers/ensureBuilt.ts

@@ -17,7 +17,20 @@ const execFileAsync = promisify(execFile);
 const DIST_SENTINELS: Record<string, string[]> = {
     // Include late-written artifacts (dts pass, CJS validators) so an
     // interrupted build never satisfies the fast path.
-    client: ['index.mjs', 'stdio.mjs', 'validators/ajv.cjs', 'index.d.mts', 'index.d.cts'],
+    client: [
+        'index.mjs',
+        'stdio.mjs',
+        // Shim entries in both formats — the workerdSchemaPreload suite reads them all.
+        'shimsWorkerd.mjs',
+        'shimsWorkerd.cjs',
+        'shimsNode.mjs',
+        'shimsNode.cjs',
+        'shimsBrowser.mjs',
+        'shimsBrowser.cjs',
+        'validators/ajv.cjs',
+        'index.d.mts',
+        'index.d.cts'
+    ],
     core: ['index.mjs', 'internal.mjs', 'index.d.mts', 'internal.d.mts']
 };