mock-mcp 0.2.3 → 0.3.1

package/README.md

@@ -4,7 +4,7 @@
 [![npm](https://img.shields.io/npm/v/mock-mcp.svg)](https://www.npmjs.com/package/mock-mcp)
 ![license](https://img.shields.io/npm/l/mock-mcp)
 
-Mock MCP Server - AI-generated mock data. The project pairs a WebSocket batch bridge with MCP tooling so Cursor, Claude Desktop, or any compatible client can fulfill intercepted requests in real time.
+Mock MCP Server - AI-generated mock data based on your **OpenAPI JSON Schema** definitions. The project pairs a WebSocket batch bridge with MCP tooling so Cursor, Claude Desktop, or any compatible client can fulfill intercepted requests in real time, ensuring strict contract compliance.
 
 ## Table of Contents
 
@@ -25,62 +25,67 @@ Mock MCP Server - AI-generated mock data. The project pairs a WebSocket batch br
 
 1. **Install the package.** Add mock-mcp as a dev dependency inside your project.
 
-   ```bash
-   npm install -D mock-mcp
-   ```
+```bash
+npm install -D mock-mcp
+# or
+yarn add -D mock-mcp
+# or
+pnpm add -D mock-mcp
+```
 
 2. **Configure the Model Context Protocol server.** For example, Claude Desktop can launch the binary through npx:
 
-   ```json
-   {
-     "mock-mcp": {
-       "command": "npx",
-       "args": ["-y", "mock-mcp@latest"]
-     }
-   }
-   ```
+```json
+{
+  "mock-mcp": {
+    "command": "npx",
+    "args": ["-y", "mock-mcp@latest"]
+  }
+}
+```
 
 3. **Connect from your tests.** Use `connect` to retrieve a mock client and request data for intercepted calls.
 
-   ```ts
-   import { render, screen, fireEvent } from "@testing-library/react";
-   import { connect } from "mock-mcp";
-
-   const userSchema = {
-     summary: "Fetch the current user",
-     response: {
-       type: "object",
-       required: ["id", "name"],
-       properties: {
-         id: { type: "number" },
-         name: { type: "string" },
-       },
-     },
-   };
-
-   it("example", async () => {
-     const mockClient = await connect();
-     const metadata = {
-       schemaUrl: "https://example.com/openapi.json#/paths/~1user/get",
-       schema: userSchema,
-       instructions: "Respond with a single user described by the schema.",
-     };
-
-     fetchMock.get("/user", () =>
-       mockClient.requestMock("/user", "GET", { metadata })
-     );
-
-     const result = await fetch("/user");
-     const data = await result.json();
-     expect(data).toEqual({ id: 1, name: "Jane" });
-   }); // 10 minute timeout for AI interaction
-   ```
+```ts
+import { render, screen, fireEvent } from "@testing-library/react";
+import { connect } from "mock-mcp";
+
+const userSchema = {
+  summary: "Fetch the current user",
+  response: {
+    type: "object",
+    required: ["id", "name"],
+    properties: {
+      id: { type: "number" },
+      name: { type: "string" },
+    },
+  },
+};
+
+it("example", async () => {
+  const mockClient = await connect();
+  const metadata = {
+    schemaUrl: "https://example.com/openapi.json#/paths/~1user/get",
+    schema: userSchema,
+    instructions: "Respond with a single user described by the schema.",
+  };
+
+  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();
+  expect(data).toEqual({ id: 1, name: "Jane" });
+}, 10 * 60 * 1000); // 10 minute timeout for AI interaction
+```
 
 4. **Run with MCP enabled.** Prompt your AI client to run the persistent test command and provide mocks through the tools.
 
-   ```
-   Please run the persistent test: `MOCK_MCP=true npm test test/example.test.tsx` and mock fetch data with mock-mcp
-   ```
+```
+Please run the persistent test: `MOCK_MCP=true npm test test/example.test.tsx` and mock fetch data with mock-mcp
+```
 
 ## Why Mock MCP
 
@@ -116,15 +121,20 @@ Traditional:  Write Test → Create Fixtures → Run Test → Maintain Fixtures
                                 ↑                          ↓
                                 └──────── Pain Loop ───────┘
 
-Mock MCP:     Write Test → AI Generates Data → Run Test → Solidify Code
-                             ↑                           ↓
-                             └─────── Evolution ─────────┘
+Mock MCP:     Write Test → AI Generates Data (Schema-Compliant) → Run Test → Solidify Code
+                             ↑                                    ↓
+                             └───────────── Evolution ────────────┘
 ```
 
+### Schema-Driven Accuracy
+
+Unlike "hallucinated" mocks, Mock MCP uses your actual **OpenAPI JSON Schema** definitions to ground the AI. This ensures that generated data not only looks real but strictly adheres to your API contracts, catching integration issues early.
+
 ## What Mock MCP Does
 
 Mock MCP pairs a WebSocket batch bridge with MCP tooling to move intercepted requests from tests to AI helpers and back again.
 
+- **Schema-aware generation** uses your provided metadata (OpenAPI JSON Schema) to ensure mocks match production behavior.
 - **Batch-aware test client** collects every network interception inside a single macrotask and waits for the full response set.
 - **MCP tooling** exposes `get_pending_batches` and `provide_batch_mock_data` so AI agents understand the waiting requests and push data back.
 - **WebSocket bridge** connects the test runner to the MCP server while hiding transport details from both sides.
@@ -174,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()
   );
@@ -189,9 +199,17 @@ await page.route("**/api/users", async (route) => {
 
 Batch behaviour stays automatic: additional `requestMock` calls issued in the same macrotask are grouped, forwarded, and resolved together.
 
+Need to pause the test until everything in-flight resolves? Call `waitForPendingRequests` to block on the current set of pending requests (anything started after the call is not included):
+
+```ts
+// After routing a few requests
+await mockClient.waitForPendingRequests();
+// Safe to assert on the results produced by the mocked responses
+```
+
 ## Describe Requests with Metadata
 
-`requestMock` accepts an optional third argument (`RequestMockOptions`) that is forwarded without modification to the MCP server. The most important field in that object is `metadata`, which lets the test process describe each request with the exact OpenAPI/JSON Schema fragment, sample payloads, or test context that the AI client needs to build a response.
+`requestMock` accepts an optional third argument (`RequestMockOptions`) that is forwarded without modification to the MCP server. The most important field in that object is `metadata`, which lets the test process describe each request with the exact OpenAPI JSON Schema fragment, sample payloads, or test context that the AI client needs to build a response.
 
 When an MCP client calls `get_pending_batches`, every `requests[].metadata` entry from the test run is included in the response. That is the channel the LLM uses to understand the requested endpoint before supplying data through `provide_batch_mock_data`. Metadata is also persisted when batch logging is enabled, so you can audit what was sent to the model.
 
@@ -261,6 +279,7 @@ The library exports primitives so you can embed the workflow inside bespoke runn
 
 - `TestMockMCPServer` starts and stops the WebSocket plus MCP tooling bridge programmatically.
 - `BatchMockCollector` provides a low-level batching client used directly inside test environments.
+- `BatchMockCollector.waitForPendingRequests()` waits for the currently pending mock requests to settle (resolves when all finish, rejects if any fail).
 - `connect(options)` instantiates `BatchMockCollector` and waits for the WebSocket connection to open.
 
 Each class accepts logger overrides, timeout tweaks, and other ergonomics surfaced in the technical design.

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;
 };
@@ -5,7 +6,7 @@ export interface BatchMockCollectorOptions {
     /**
      * TCP port exposed by {@link TestMockMCPServer}.
      *
-     * @default 8080
+     * @default 3002
      */
     port?: number;
     /**
@@ -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,17 +83,28 @@ 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.
+     */
+    waitForPendingRequests(): Promise<void>;
     /**
      * Close the underlying connection and fail all pending requests.
      */
     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

@@ -1,9 +1,11 @@
 import WebSocket from "ws";
 import { BATCH_MOCK_REQUEST, BATCH_MOCK_RESPONSE, } from "../types.js";
+import { isEnabled } from "./util.js";
 const DEFAULT_TIMEOUT = 60_000;
 const DEFAULT_BATCH_DEBOUNCE_MS = 0;
 const DEFAULT_MAX_BATCH_SIZE = 50;
-const DEFAULT_PORT = 8080;
+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.
@@ -16,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();
     }
     /**
@@ -59,24 +65,45 @@ export class BatchMockCollector {
             headers: options.headers,
             metadata: options.metadata,
         };
+        let settleCompletion;
+        const completion = new Promise((resolve) => {
+            settleCompletion = resolve;
+        });
         return new Promise((resolve, reject) => {
             const timeoutId = setTimeout(() => {
-                this.pendingRequests.delete(requestId);
-                reject(new Error(`Mock request timed out after ${this.timeout}ms: ${method} ${endpoint}`));
+                this.rejectRequest(requestId, new Error(`Mock request timed out after ${this.timeout}ms: ${method} ${endpoint}`));
             }, this.timeout);
             this.pendingRequests.set(requestId, {
                 request,
-                resolve: (data) => {
-                    resolve(data);
+                resolve: (mock) => {
+                    settleCompletion({ status: "fulfilled", value: undefined });
+                    resolve(this.buildResolvedMock(mock));
                 },
                 reject: (error) => {
+                    settleCompletion({ status: "rejected", reason: error });
                     reject(error);
                 },
                 timeoutId,
+                completion,
             });
             this.enqueueRequest(requestId);
         });
     }
+    /**
+     * Wait for all requests that are currently pending to settle. Requests
+     * created after this method is called are not included.
+     */
+    async waitForPendingRequests() {
+        if (!isEnabled()) {
+            return;
+        }
+        const pendingCompletions = Array.from(this.pendingRequests.values()).map((pending) => pending.completion);
+        const results = await Promise.all(pendingCompletions);
+        const rejected = results.find((result) => result.status === "rejected");
+        if (rejected) {
+            throw rejected.reason;
+        }
+    }
     /**
      * Close the underlying connection and fail all pending requests.
      */
@@ -89,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());
@@ -101,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) => {
@@ -110,9 +146,65 @@ 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;
         try {
@@ -139,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);
@@ -184,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 isEnabled = process.env.MOCK_MCP !== undefined && process.env.MOCK_MCP !== "0";
-    if (!isEnabled) {
+    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/client/util.d.ts

@@ -0,0 +1 @@
+export declare const isEnabled: () => boolean;

package/dist/client/util.js

@@ -0,0 +1,3 @@
+export const isEnabled = () => {
+    return process.env.MOCK_MCP !== undefined && process.env.MOCK_MCP !== "0";
+};

package/dist/connect.cjs

@@ -41,11 +41,17 @@ var import_ws = __toESM(require("ws"), 1);
 var BATCH_MOCK_REQUEST = "BATCH_MOCK_REQUEST";
 var BATCH_MOCK_RESPONSE = "BATCH_MOCK_RESPONSE";
 
+// src/client/util.ts
+var isEnabled = () => {
+  return process.env.MOCK_MCP !== void 0 && process.env.MOCK_MCP !== "0";
+};
+
 // src/client/batch-mock-collector.ts
 var DEFAULT_TIMEOUT = 6e4;
 var DEFAULT_BATCH_DEBOUNCE_MS = 0;
 var DEFAULT_MAX_BATCH_SIZE = 50;
-var DEFAULT_PORT = 8080;
+var DEFAULT_PORT = 3002;
+var DEFAULT_HEARTBEAT_INTERVAL_MS = 15e3;
 var BatchMockCollector = class {
   ws;
   pendingRequests = /* @__PURE__ */ new Map();
@@ -54,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();
   }
   /**
@@ -97,10 +106,14 @@ var BatchMockCollector = class {
       headers: options.headers,
       metadata: options.metadata
     };
+    let settleCompletion;
+    const completion = new Promise((resolve) => {
+      settleCompletion = resolve;
+    });
     return new Promise((resolve, reject) => {
       const timeoutId = setTimeout(() => {
-        this.pendingRequests.delete(requestId);
-        reject(
+        this.rejectRequest(
+          requestId,
           new Error(
             `Mock request timed out after ${this.timeout}ms: ${method} ${endpoint}`
           )
@@ -108,17 +121,39 @@ var BatchMockCollector = class {
       }, this.timeout);
       this.pendingRequests.set(requestId, {
         request,
-        resolve: (data) => {
-          resolve(data);
+        resolve: (mock) => {
+          settleCompletion({ status: "fulfilled", value: void 0 });
+          resolve(this.buildResolvedMock(mock));
         },
         reject: (error) => {
+          settleCompletion({ status: "rejected", reason: error });
           reject(error);
         },
-        timeoutId
+        timeoutId,
+        completion
       });
       this.enqueueRequest(requestId);
     });
   }
+  /**
+   * Wait for all requests that are currently pending to settle. Requests
+   * created after this method is called are not included.
+   */
+  async waitForPendingRequests() {
+    if (!isEnabled()) {
+      return;
+    }
+    const pendingCompletions = Array.from(this.pendingRequests.values()).map(
+      (pending) => pending.completion
+    );
+    const results = await Promise.all(pendingCompletions);
+    const rejected = results.find(
+      (result) => result.status === "rejected"
+    );
+    if (rejected) {
+      throw rejected.reason;
+    }
+  }
   /**
    * Close the underlying connection and fail all pending requests.
    */
@@ -131,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());
@@ -143,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) => {
@@ -156,8 +200,66 @@ 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;
@@ -186,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);
@@ -235,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) {
@@ -252,13 +368,28 @@ var BatchMockCollector = class {
 };
 
 // src/client/connect.ts
-var connect = async (options) => {
-  const isEnabled = process.env.MOCK_MCP !== void 0 && process.env.MOCK_MCP !== "0";
-  if (!isEnabled) {
-    console.log("[mock-mcp] Skipping (set MOCK_MCP=1 to enable)");
+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 new DisabledMockClient();
+  }
   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;
 };
@@ -5,7 +19,7 @@ interface BatchMockCollectorOptions {
     /**
      * TCP port exposed by {@link TestMockMCPServer}.
      *
-     * @default 8080
+     * @default 3002
      */
     port?: number;
     /**
@@ -33,58 +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>;
-    /**
-     * Close the underlying connection and fail all pending requests.
-     */
+    requestMock<T = unknown>(endpoint: string, method: string, options?: RequestMockOptions): Promise<ResolvedMock<T>>;
+    waitForPendingRequests(): Promise<void>;
     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.2.3",
+  "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",
@@ -70,4 +70,4 @@
       "path": "cz-conventional-changelog"
     }
   }
-}
+}