mock-mcp 0.3.0 → 0.3.1

package/README.md

@@ -70,9 +70,10 @@ it("example", async () => {
     instructions: "Respond with a single user described by the schema.",
   };
 
-  fetchMock.get("/user", () =>
-    mockClient.requestMock("/user", "GET", { metadata }) // add mock via mock-mcp
-  );
+  fetchMock.get("/user", async () => {
+    const response = await mockClient.requestMock("/user", "GET", { metadata }) // add mock via mock-mcp
+    return response.data
+  });
 
   const result = await fetch("/user");
   const data = await result.json();
@@ -183,7 +184,7 @@ const mockClient = await connect({
 
 await page.route("**/api/users", async (route) => {
   const url = new URL(route.request().url());
-  const data = await mockClient.requestMock(
+  const { data } = await mockClient.requestMock(
     url.pathname,
     route.request().method()
   );

package/dist/client/batch-mock-collector.d.ts

@@ -1,3 +1,4 @@
+import { type ResolvedMock } from "../types.js";
 type Logger = Pick<Console, "log" | "warn" | "error"> & {
     debug?: (...args: unknown[]) => void;
 };
@@ -33,6 +34,18 @@ export interface BatchMockCollectorOptions {
      * Optional custom logger. Defaults to `console`.
      */
     logger?: Logger;
+    /**
+     * Interval for WebSocket heartbeats in milliseconds. Set to 0 to disable.
+     *
+     * @default 15000
+     */
+    heartbeatIntervalMs?: number;
+    /**
+     * Automatically attempt to reconnect when the WebSocket closes unexpectedly.
+     *
+     * @default true
+     */
+    enableReconnect?: boolean;
 }
 export interface RequestMockOptions {
     body?: unknown;
@@ -44,19 +57,24 @@ export interface RequestMockOptions {
  * the MCP server as a batch for AI-assisted mock generation.
  */
 export declare class BatchMockCollector {
-    private readonly ws;
+    private ws;
     private readonly pendingRequests;
     private readonly queuedRequestIds;
     private readonly timeout;
     private readonly batchDebounceMs;
     private readonly maxBatchSize;
     private readonly logger;
+    private readonly heartbeatIntervalMs;
+    private readonly enableReconnect;
+    private readonly port;
     private batchTimer;
+    private heartbeatTimer;
+    private reconnectTimer;
     private requestIdCounter;
     private closed;
     private readyResolve?;
     private readyReject?;
-    private readonly readyPromise;
+    private readyPromise;
     constructor(options?: BatchMockCollectorOptions);
     /**
      * Ensures the underlying WebSocket connection is ready for use.
@@ -65,7 +83,7 @@ export declare class BatchMockCollector {
     /**
      * Request mock data for a specific endpoint/method pair.
      */
-    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<T>;
+    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<ResolvedMock<T>>;
     /**
      * Wait for all requests that are currently pending to settle. Requests
      * created after this method is called are not included.
@@ -76,11 +94,17 @@ export declare class BatchMockCollector {
      */
     close(code?: number): Promise<void>;
     private setupWebSocket;
+    private createWebSocket;
+    private resetReadyPromise;
+    private startHeartbeat;
+    private stopHeartbeat;
+    private scheduleReconnect;
     private handleMessage;
     private resolveRequest;
     private enqueueRequest;
     private flushQueue;
     private sendBatch;
+    private buildResolvedMock;
     private rejectRequest;
     private failAllPending;
 }

package/dist/client/batch-mock-collector.js

@@ -5,6 +5,7 @@ const DEFAULT_TIMEOUT = 60_000;
 const DEFAULT_BATCH_DEBOUNCE_MS = 0;
 const DEFAULT_MAX_BATCH_SIZE = 50;
 const DEFAULT_PORT = 3002;
+const DEFAULT_HEARTBEAT_INTERVAL_MS = 15_000;
 /**
  * Collects HTTP requests issued during a single macrotask and forwards them to
  * the MCP server as a batch for AI-assisted mock generation.
@@ -17,24 +18,28 @@ export class BatchMockCollector {
     batchDebounceMs;
     maxBatchSize;
     logger;
+    heartbeatIntervalMs;
+    enableReconnect;
+    port;
     batchTimer = null;
+    heartbeatTimer = null;
+    reconnectTimer = null;
     requestIdCounter = 0;
     closed = false;
     readyResolve;
     readyReject;
-    readyPromise;
+    readyPromise = Promise.resolve();
     constructor(options = {}) {
         this.timeout = options.timeout ?? DEFAULT_TIMEOUT;
         this.batchDebounceMs = options.batchDebounceMs ?? DEFAULT_BATCH_DEBOUNCE_MS;
         this.maxBatchSize = options.maxBatchSize ?? DEFAULT_MAX_BATCH_SIZE;
         this.logger = options.logger ?? console;
-        const port = options.port ?? DEFAULT_PORT;
-        this.readyPromise = new Promise((resolve, reject) => {
-            this.readyResolve = resolve;
-            this.readyReject = reject;
-        });
-        const wsUrl = `ws://localhost:${port}`;
-        this.ws = new WebSocket(wsUrl);
+        this.heartbeatIntervalMs =
+            options.heartbeatIntervalMs ?? DEFAULT_HEARTBEAT_INTERVAL_MS;
+        this.enableReconnect = options.enableReconnect ?? true;
+        this.port = options.port ?? DEFAULT_PORT;
+        this.resetReadyPromise();
+        this.ws = this.createWebSocket();
         this.setupWebSocket();
     }
     /**
@@ -70,9 +75,9 @@ export class BatchMockCollector {
             }, this.timeout);
             this.pendingRequests.set(requestId, {
                 request,
-                resolve: (data) => {
+                resolve: (mock) => {
                     settleCompletion({ status: "fulfilled", value: undefined });
-                    resolve(data);
+                    resolve(this.buildResolvedMock(mock));
                 },
                 reject: (error) => {
                     settleCompletion({ status: "rejected", reason: error });
@@ -111,6 +116,14 @@ export class BatchMockCollector {
             clearTimeout(this.batchTimer);
             this.batchTimer = null;
         }
+        if (this.heartbeatTimer) {
+            clearInterval(this.heartbeatTimer);
+            this.heartbeatTimer = null;
+        }
+        if (this.reconnectTimer) {
+            clearTimeout(this.reconnectTimer);
+            this.reconnectTimer = null;
+        }
         this.queuedRequestIds.clear();
         const closePromise = new Promise((resolve) => {
             this.ws.once("close", () => resolve());
@@ -123,6 +136,7 @@ export class BatchMockCollector {
         this.ws.on("open", () => {
             this.logger.log("🔌 Connected to mock MCP WebSocket endpoint");
             this.readyResolve?.();
+            this.startHeartbeat();
         });
         this.ws.on("message", (data) => this.handleMessage(data));
         this.ws.on("error", (error) => {
@@ -132,8 +146,64 @@ export class BatchMockCollector {
         });
         this.ws.on("close", () => {
             this.logger.warn("🔌 WebSocket connection closed");
+            this.stopHeartbeat();
             this.failAllPending(new Error("WebSocket connection closed"));
+            if (!this.closed && this.enableReconnect) {
+                this.scheduleReconnect();
+            }
+        });
+    }
+    createWebSocket() {
+        const wsUrl = `ws://localhost:${this.port}`;
+        return new WebSocket(wsUrl);
+    }
+    resetReadyPromise() {
+        this.readyPromise = new Promise((resolve, reject) => {
+            this.readyResolve = resolve;
+            this.readyReject = reject;
+        });
+    }
+    startHeartbeat() {
+        if (this.heartbeatIntervalMs <= 0 || this.heartbeatTimer) {
+            return;
+        }
+        let lastPong = Date.now();
+        this.ws.on("pong", () => {
+            lastPong = Date.now();
         });
+        this.heartbeatTimer = setInterval(() => {
+            if (this.ws.readyState !== WebSocket.OPEN) {
+                return;
+            }
+            const now = Date.now();
+            if (now - lastPong > this.heartbeatIntervalMs * 2) {
+                this.logger.warn("Heartbeat missed; closing socket to trigger reconnect...");
+                this.ws.close();
+                return;
+            }
+            this.ws.ping();
+        }, this.heartbeatIntervalMs);
+        this.heartbeatTimer.unref?.();
+    }
+    stopHeartbeat() {
+        if (this.heartbeatTimer) {
+            clearInterval(this.heartbeatTimer);
+            this.heartbeatTimer = null;
+        }
+    }
+    scheduleReconnect() {
+        if (this.reconnectTimer || this.closed) {
+            return;
+        }
+        this.reconnectTimer = setTimeout(() => {
+            this.reconnectTimer = null;
+            this.logger.warn("🔄 Reconnecting to mock MCP WebSocket endpoint...");
+            this.stopHeartbeat();
+            this.resetReadyPromise();
+            this.ws = this.createWebSocket();
+            this.setupWebSocket();
+        }, 1_000);
+        this.reconnectTimer.unref?.();
     }
     handleMessage(data) {
         let parsed;
@@ -161,7 +231,13 @@ export class BatchMockCollector {
         }
         clearTimeout(pending.timeoutId);
         this.pendingRequests.delete(mock.requestId);
-        pending.resolve(mock.data);
+        const resolve = () => pending.resolve(mock);
+        if (mock.delayMs && mock.delayMs > 0) {
+            setTimeout(resolve, mock.delayMs);
+        }
+        else {
+            resolve();
+        }
     }
     enqueueRequest(requestId) {
         this.queuedRequestIds.add(requestId);
@@ -206,6 +282,15 @@ export class BatchMockCollector {
         this.logger.debug?.(`📤 Sending batch with ${requests.length} request(s) to MCP server`);
         this.ws.send(JSON.stringify(payload));
     }
+    buildResolvedMock(mock) {
+        return {
+            requestId: mock.requestId,
+            data: mock.data,
+            status: mock.status,
+            headers: mock.headers,
+            delayMs: mock.delayMs,
+        };
+    }
     rejectRequest(requestId, error) {
         const pending = this.pendingRequests.get(requestId);
         if (!pending) {

package/dist/client/connect.d.ts

@@ -1,8 +1,14 @@
-import { BatchMockCollector } from "./batch-mock-collector.js";
-import type { BatchMockCollectorOptions } from "./batch-mock-collector.js";
+import type { BatchMockCollectorOptions, RequestMockOptions } from "./batch-mock-collector.js";
+import type { ResolvedMock } from "../types.js";
 export type ConnectOptions = number | BatchMockCollectorOptions | undefined;
+export interface MockClient {
+    waitUntilReady(): Promise<void>;
+    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<ResolvedMock<T>>;
+    waitForPendingRequests(): Promise<void>;
+    close(code?: number): Promise<void>;
+}
 /**
  * Convenience helper that creates a {@link BatchMockCollector} and waits for the
  * underlying WebSocket connection to become ready before resolving.
  */
-export declare const connect: (options?: ConnectOptions) => Promise<BatchMockCollector | void>;
+export declare const connect: (options?: ConnectOptions) => Promise<MockClient>;

package/dist/client/connect.js

@@ -1,15 +1,29 @@
 import { BatchMockCollector } from "./batch-mock-collector.js";
 import { isEnabled } from "./util.js";
+class DisabledMockClient {
+    async waitUntilReady() {
+        return;
+    }
+    async requestMock() {
+        throw new Error("[mock-mcp] MOCK_MCP is not enabled. Set MOCK_MCP=1 to enable mock generation.");
+    }
+    async waitForPendingRequests() {
+        return;
+    }
+    async close() {
+        return;
+    }
+}
 /**
  * Convenience helper that creates a {@link BatchMockCollector} and waits for the
  * underlying WebSocket connection to become ready before resolving.
  */
 export const connect = async (options) => {
+    const resolvedOptions = typeof options === "number" ? { port: options } : options ?? {};
     if (!isEnabled()) {
         console.log("[mock-mcp] Skipping (set MOCK_MCP=1 to enable)");
-        return;
+        return new DisabledMockClient();
     }
-    const resolvedOptions = typeof options === "number" ? { port: options } : options ?? {};
     const collector = new BatchMockCollector(resolvedOptions);
     await collector.waitUntilReady();
     return collector;

package/dist/client/index.d.ts

@@ -1,2 +1,2 @@
 export { BatchMockCollector, type BatchMockCollectorOptions, type RequestMockOptions, } from "./batch-mock-collector.js";
-export { connect, type ConnectOptions } from "./connect.js";
+export { connect, type ConnectOptions, type MockClient } from "./connect.js";

package/dist/connect.cjs

@@ -51,6 +51,7 @@ var DEFAULT_TIMEOUT = 6e4;
 var DEFAULT_BATCH_DEBOUNCE_MS = 0;
 var DEFAULT_MAX_BATCH_SIZE = 50;
 var DEFAULT_PORT = 3002;
+var DEFAULT_HEARTBEAT_INTERVAL_MS = 15e3;
 var BatchMockCollector = class {
   ws;
   pendingRequests = /* @__PURE__ */ new Map();
@@ -59,24 +60,27 @@ var BatchMockCollector = class {
   batchDebounceMs;
   maxBatchSize;
   logger;
+  heartbeatIntervalMs;
+  enableReconnect;
+  port;
   batchTimer = null;
+  heartbeatTimer = null;
+  reconnectTimer = null;
   requestIdCounter = 0;
   closed = false;
   readyResolve;
   readyReject;
-  readyPromise;
+  readyPromise = Promise.resolve();
   constructor(options = {}) {
     this.timeout = options.timeout ?? DEFAULT_TIMEOUT;
     this.batchDebounceMs = options.batchDebounceMs ?? DEFAULT_BATCH_DEBOUNCE_MS;
     this.maxBatchSize = options.maxBatchSize ?? DEFAULT_MAX_BATCH_SIZE;
     this.logger = options.logger ?? console;
-    const port = options.port ?? DEFAULT_PORT;
-    this.readyPromise = new Promise((resolve, reject) => {
-      this.readyResolve = resolve;
-      this.readyReject = reject;
-    });
-    const wsUrl = `ws://localhost:${port}`;
-    this.ws = new import_ws.default(wsUrl);
+    this.heartbeatIntervalMs = options.heartbeatIntervalMs ?? DEFAULT_HEARTBEAT_INTERVAL_MS;
+    this.enableReconnect = options.enableReconnect ?? true;
+    this.port = options.port ?? DEFAULT_PORT;
+    this.resetReadyPromise();
+    this.ws = this.createWebSocket();
     this.setupWebSocket();
   }
   /**
@@ -117,9 +121,9 @@ var BatchMockCollector = class {
       }, this.timeout);
       this.pendingRequests.set(requestId, {
         request,
-        resolve: (data) => {
+        resolve: (mock) => {
           settleCompletion({ status: "fulfilled", value: void 0 });
-          resolve(data);
+          resolve(this.buildResolvedMock(mock));
         },
         reject: (error) => {
           settleCompletion({ status: "rejected", reason: error });
@@ -162,6 +166,14 @@ var BatchMockCollector = class {
       clearTimeout(this.batchTimer);
       this.batchTimer = null;
     }
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
     this.queuedRequestIds.clear();
     const closePromise = new Promise((resolve) => {
       this.ws.once("close", () => resolve());
@@ -174,6 +186,7 @@ var BatchMockCollector = class {
     this.ws.on("open", () => {
       this.logger.log("\u{1F50C} Connected to mock MCP WebSocket endpoint");
       this.readyResolve?.();
+      this.startHeartbeat();
     });
     this.ws.on("message", (data) => this.handleMessage(data));
     this.ws.on("error", (error) => {
@@ -187,9 +200,67 @@ var BatchMockCollector = class {
     });
     this.ws.on("close", () => {
       this.logger.warn("\u{1F50C} WebSocket connection closed");
+      this.stopHeartbeat();
       this.failAllPending(new Error("WebSocket connection closed"));
+      if (!this.closed && this.enableReconnect) {
+        this.scheduleReconnect();
+      }
     });
   }
+  createWebSocket() {
+    const wsUrl = `ws://localhost:${this.port}`;
+    return new import_ws.default(wsUrl);
+  }
+  resetReadyPromise() {
+    this.readyPromise = new Promise((resolve, reject) => {
+      this.readyResolve = resolve;
+      this.readyReject = reject;
+    });
+  }
+  startHeartbeat() {
+    if (this.heartbeatIntervalMs <= 0 || this.heartbeatTimer) {
+      return;
+    }
+    let lastPong = Date.now();
+    this.ws.on("pong", () => {
+      lastPong = Date.now();
+    });
+    this.heartbeatTimer = setInterval(() => {
+      if (this.ws.readyState !== import_ws.default.OPEN) {
+        return;
+      }
+      const now = Date.now();
+      if (now - lastPong > this.heartbeatIntervalMs * 2) {
+        this.logger.warn(
+          "Heartbeat missed; closing socket to trigger reconnect..."
+        );
+        this.ws.close();
+        return;
+      }
+      this.ws.ping();
+    }, this.heartbeatIntervalMs);
+    this.heartbeatTimer.unref?.();
+  }
+  stopHeartbeat() {
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+  }
+  scheduleReconnect() {
+    if (this.reconnectTimer || this.closed) {
+      return;
+    }
+    this.reconnectTimer = setTimeout(() => {
+      this.reconnectTimer = null;
+      this.logger.warn("\u{1F504} Reconnecting to mock MCP WebSocket endpoint...");
+      this.stopHeartbeat();
+      this.resetReadyPromise();
+      this.ws = this.createWebSocket();
+      this.setupWebSocket();
+    }, 1e3);
+    this.reconnectTimer.unref?.();
+  }
   handleMessage(data) {
     let parsed;
     try {
@@ -217,7 +288,12 @@ var BatchMockCollector = class {
     }
     clearTimeout(pending.timeoutId);
     this.pendingRequests.delete(mock.requestId);
-    pending.resolve(mock.data);
+    const resolve = () => pending.resolve(mock);
+    if (mock.delayMs && mock.delayMs > 0) {
+      setTimeout(resolve, mock.delayMs);
+    } else {
+      resolve();
+    }
   }
   enqueueRequest(requestId) {
     this.queuedRequestIds.add(requestId);
@@ -266,6 +342,15 @@ var BatchMockCollector = class {
     );
     this.ws.send(JSON.stringify(payload));
   }
+  buildResolvedMock(mock) {
+    return {
+      requestId: mock.requestId,
+      data: mock.data,
+      status: mock.status,
+      headers: mock.headers,
+      delayMs: mock.delayMs
+    };
+  }
   rejectRequest(requestId, error) {
     const pending = this.pendingRequests.get(requestId);
     if (!pending) {
@@ -283,12 +368,28 @@ var BatchMockCollector = class {
 };
 
 // src/client/connect.ts
+var DisabledMockClient = class {
+  async waitUntilReady() {
+    return;
+  }
+  async requestMock() {
+    throw new Error(
+      "[mock-mcp] MOCK_MCP is not enabled. Set MOCK_MCP=1 to enable mock generation."
+    );
+  }
+  async waitForPendingRequests() {
+    return;
+  }
+  async close() {
+    return;
+  }
+};
 var connect = async (options) => {
+  const resolvedOptions = typeof options === "number" ? { port: options } : options ?? {};
   if (!isEnabled()) {
     console.log("[mock-mcp] Skipping (set MOCK_MCP=1 to enable)");
-    return;
+    return new DisabledMockClient();
   }
-  const resolvedOptions = typeof options === "number" ? { port: options } : options ?? {};
   const collector = new BatchMockCollector(resolvedOptions);
   await collector.waitUntilReady();
   return collector;

package/dist/connect.d.cts

@@ -1,3 +1,17 @@
+/**
+ * Shape of the mock data that needs to be returned for a request.
+ */
+interface MockResponseDescriptor {
+    requestId: string;
+    data: unknown;
+    status?: number;
+    headers?: Record<string, string>;
+    delayMs?: number;
+}
+interface ResolvedMock<T = unknown> extends Omit<MockResponseDescriptor, "data"> {
+    data: T;
+}
+
 type Logger = Pick<Console, "log" | "warn" | "error"> & {
     debug?: (...args: unknown[]) => void;
 };
@@ -33,63 +47,36 @@ interface BatchMockCollectorOptions {
      * Optional custom logger. Defaults to `console`.
      */
     logger?: Logger;
+    /**
+     * Interval for WebSocket heartbeats in milliseconds. Set to 0 to disable.
+     *
+     * @default 15000
+     */
+    heartbeatIntervalMs?: number;
+    /**
+     * Automatically attempt to reconnect when the WebSocket closes unexpectedly.
+     *
+     * @default true
+     */
+    enableReconnect?: boolean;
 }
 interface RequestMockOptions {
     body?: unknown;
     headers?: Record<string, string>;
     metadata?: Record<string, unknown>;
 }
-/**
- * Collects HTTP requests issued during a single macrotask and forwards them to
- * the MCP server as a batch for AI-assisted mock generation.
- */
-declare class BatchMockCollector {
-    private readonly ws;
-    private readonly pendingRequests;
-    private readonly queuedRequestIds;
-    private readonly timeout;
-    private readonly batchDebounceMs;
-    private readonly maxBatchSize;
-    private readonly logger;
-    private batchTimer;
-    private requestIdCounter;
-    private closed;
-    private readyResolve?;
-    private readyReject?;
-    private readonly readyPromise;
-    constructor(options?: BatchMockCollectorOptions);
-    /**
-     * Ensures the underlying WebSocket connection is ready for use.
-     */
+
+type ConnectOptions = number | BatchMockCollectorOptions | undefined;
+interface MockClient {
     waitUntilReady(): Promise<void>;
-    /**
-     * Request mock data for a specific endpoint/method pair.
-     */
-    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<T>;
-    /**
-     * Wait for all requests that are currently pending to settle. Requests
-     * created after this method is called are not included.
-     */
+    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<ResolvedMock<T>>;
     waitForPendingRequests(): Promise<void>;
-    /**
-     * Close the underlying connection and fail all pending requests.
-     */
     close(code?: number): Promise<void>;
-    private setupWebSocket;
-    private handleMessage;
-    private resolveRequest;
-    private enqueueRequest;
-    private flushQueue;
-    private sendBatch;
-    private rejectRequest;
-    private failAllPending;
 }
-
-type ConnectOptions = number | BatchMockCollectorOptions | undefined;
 /**
  * Convenience helper that creates a {@link BatchMockCollector} and waits for the
  * underlying WebSocket connection to become ready before resolving.
  */
-declare const connect: (options?: ConnectOptions) => Promise<BatchMockCollector | void>;
+declare const connect: (options?: ConnectOptions) => Promise<MockClient>;
 
-export { type ConnectOptions, connect };
+export { type ConnectOptions, type MockClient, connect };

package/dist/index.d.ts

@@ -1,10 +1,11 @@
 #!/usr/bin/env node
 import { TestMockMCPServer, type TestMockMCPServerOptions } from "./server/test-mock-mcp-server.js";
 import { BatchMockCollector, type BatchMockCollectorOptions, type RequestMockOptions } from "./client/batch-mock-collector.js";
-import { connect, type ConnectOptions } from "./client/connect.js";
+import { connect, type ConnectOptions, type MockClient } from "./client/connect.js";
+import type { ResolvedMock } from "./types.js";
 export { TestMockMCPServer };
 export type { TestMockMCPServerOptions };
 export { BatchMockCollector };
 export type { BatchMockCollectorOptions, RequestMockOptions };
 export { connect };
-export type { ConnectOptions };
+export type { ConnectOptions, MockClient, ResolvedMock };

package/dist/server/test-mock-mcp-server.js

@@ -109,9 +109,21 @@ export class TestMockMCPServer {
         if (!batch) {
             throw new Error(`Batch not found: ${batchId}`);
         }
-        const missing = mocks.find((mock) => !batch.requests.some((request) => request.requestId === mock.requestId));
-        if (missing) {
-            throw new Error(`Mock data references unknown requestId: ${missing.requestId}`);
+        const expectedIds = new Set(batch.requests.map((request) => request.requestId));
+        const providedIds = new Set();
+        const unknownMock = mocks.find((mock) => !expectedIds.has(mock.requestId));
+        if (unknownMock) {
+            throw new Error(`Mock data references unknown requestId: ${unknownMock.requestId}`);
+        }
+        for (const mock of mocks) {
+            if (providedIds.has(mock.requestId)) {
+                throw new Error(`Duplicate mock data provided for requestId: ${mock.requestId}`);
+            }
+            providedIds.add(mock.requestId);
+        }
+        const missingIds = Array.from(expectedIds).filter((requestId) => !providedIds.has(requestId));
+        if (missingIds.length > 0) {
+            throw new Error(`Missing mock data for requestId(s): ${missingIds.join(", ")}`);
         }
         if (batch.ws.readyState !== WebSocket.OPEN) {
             this.pendingBatches.delete(batchId);
@@ -204,7 +216,14 @@ export class TestMockMCPServer {
                                     properties: {
                                         requestId: { type: "string" },
                                         data: {
-                                            type: "object",
+                                            anyOf: [
+                                                { type: "object" },
+                                                { type: "array" },
+                                                { type: "string" },
+                                                { type: "number" },
+                                                { type: "boolean" },
+                                                { type: "null" },
+                                            ],
                                         },
                                         status: {
                                             type: "number",
@@ -249,6 +268,14 @@ export class TestMockMCPServer {
         this.logger.error("🔌 Test process connected");
         this.clients.add(ws);
         ws.on("message", (data) => this.handleClientMessage(ws, data));
+        ws.on("ping", () => {
+            try {
+                ws.pong();
+            }
+            catch (error) {
+                this.logger.warn("Failed to respond to ping:", error);
+            }
+        });
         ws.on("close", () => {
             this.logger.error("🔌 Test process disconnected");
             this.clients.delete(ws);

package/dist/types.d.ts

@@ -21,6 +21,9 @@ export interface MockResponseDescriptor {
     headers?: Record<string, string>;
     delayMs?: number;
 }
+export interface ResolvedMock<T = unknown> extends Omit<MockResponseDescriptor, "data"> {
+    data: T;
+}
 export interface BatchMockRequestMessage {
     type: typeof BATCH_MOCK_REQUEST;
     requests: MockRequestDescriptor[];

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mock-mcp",
-  "version": "0.3.0",
+  "version": "0.3.1",
   "description": "An MCP server enabling LLMs to write integration tests through live test environment interaction",
   "main": "./dist/connect.cjs",
   "type": "module",