Merge pull request #2363 from modelcontextprotocol/v2/fix/2028-stdio-delayed-responses

fix: reset relay request-response waits on progress notifications (#2028)

clients/web/src/test/core/mcp/remote/progressToken.test.ts

@@ -0,0 +1,115 @@
+/**
+ * Unit tests for the shared `progressTokenOf` guard (#2028) — the string/safe
+ * integer extractor both relay waits (browser `RemoteClientTransport` and Node
+ * `RemoteSession`) use to correlate a `notifications/progress` with the request
+ * whose timeout it should re-arm. Covers the protocol's accept set and, since
+ * the guard delegates to the SDK's `ProgressTokenSchema`, the fractional and
+ * past-`MAX_SAFE_INTEGER` values a bare `typeof` check would have let through.
+ */
+import { describe, it, expect } from "vitest";
+import { progressTokenOf } from "@inspector/core/mcp/remote/progressToken.js";
+import type { JSONRPCMessage } from "@modelcontextprotocol/client";
+
+// Cast helper: these fixtures are deliberately off-spec to exercise the guard
+// branches, which the typed JSONRPCMessage shape would otherwise forbid.
+const msg = (m: unknown): JSONRPCMessage => m as JSONRPCMessage;
+
+describe("progressTokenOf (#2028)", () => {
+  it("returns a numeric progressToken from a progress notification", () => {
+    expect(
+      progressTokenOf({
+        jsonrpc: "2.0",
+        method: "notifications/progress",
+        params: { progressToken: 42, progress: 1, total: 3 },
+      }),
+    ).toBe(42);
+  });
+
+  it("returns a string progressToken", () => {
+    expect(
+      progressTokenOf({
+        jsonrpc: "2.0",
+        method: "notifications/progress",
+        params: { progressToken: "abc", progress: 1 },
+      }),
+    ).toBe("abc");
+  });
+
+  it("returns undefined for a different notification method", () => {
+    expect(
+      progressTokenOf({
+        jsonrpc: "2.0",
+        method: "notifications/message",
+        params: { level: "info", data: {} },
+      }),
+    ).toBeUndefined();
+  });
+
+  it("returns undefined for a JSON-RPC response (no method)", () => {
+    expect(
+      progressTokenOf({ jsonrpc: "2.0", id: 1, result: {} }),
+    ).toBeUndefined();
+  });
+
+  it("returns undefined when a progress note carries no params", () => {
+    expect(
+      progressTokenOf(
+        msg({ jsonrpc: "2.0", method: "notifications/progress" }),
+      ),
+    ).toBeUndefined();
+  });
+
+  it("returns undefined when params is null", () => {
+    expect(
+      progressTokenOf(
+        msg({ jsonrpc: "2.0", method: "notifications/progress", params: null }),
+      ),
+    ).toBeUndefined();
+  });
+
+  it("returns undefined when the progressToken is neither string nor number", () => {
+    expect(
+      progressTokenOf(
+        msg({
+          jsonrpc: "2.0",
+          method: "notifications/progress",
+          params: { progressToken: { nested: true } },
+        }),
+      ),
+    ).toBeUndefined();
+  });
+
+  it("rejects a fractional numeric token (the protocol allows only integers)", () => {
+    expect(
+      progressTokenOf(
+        msg({
+          jsonrpc: "2.0",
+          method: "notifications/progress",
+          params: { progressToken: 1.5, progress: 1 },
+        }),
+      ),
+    ).toBeUndefined();
+  });
+
+  it("rejects an integer past MAX_SAFE_INTEGER", () => {
+    expect(
+      progressTokenOf(
+        msg({
+          jsonrpc: "2.0",
+          method: "notifications/progress",
+          params: { progressToken: Number.MAX_SAFE_INTEGER + 1, progress: 1 },
+        }),
+      ),
+    ).toBeUndefined();
+  });
+
+  it("accepts a safe integer at the boundary", () => {
+    expect(
+      progressTokenOf({
+        jsonrpc: "2.0",
+        method: "notifications/progress",
+        params: { progressToken: Number.MAX_SAFE_INTEGER, progress: 1 },
+      }),
+    ).toBe(Number.MAX_SAFE_INTEGER);
+  });
+});

clients/web/src/test/core/mcp/remote/remoteClientTransport.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect, vi } from "vitest";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { RemoteClientTransport } from "@inspector/core/mcp/remote/remoteClientTransport.js";
 import type { MCPServerConfig } from "@inspector/core/mcp/types.js";
 
@@ -686,4 +686,127 @@ describe("RemoteClientTransport", () => {
       await transport.close();
     });
   });
+
+  describe("progress notifications keep the SSE response wait alive (#2028)", () => {
+    // A remote backend holds POST /api/mcp/send open until the JSON-RPC
+    // response, and the browser bounds that wait with `sseResponseTimeoutMs`.
+    // For a long tool that reports progress, that wait must re-arm on each
+    // `notifications/progress` — mirroring the SDK client's
+    // `resetTimeoutOnProgress` — or the web client alone times out a call that
+    // every other transport rides out. Log messages must NOT re-arm it.
+    function backendHoldingSend(): {
+      fetchFn: typeof fetch;
+      getSse: () => ReturnType<typeof createPushableSseStream>;
+      getSentId: () => string | number | undefined;
+    } {
+      let sse!: ReturnType<typeof createPushableSseStream>;
+      let sentId: string | number | undefined;
+      const fetchFn = vi
+        .fn<typeof fetch>()
+        .mockImplementation(async (input, init) => {
+          const url = String(input);
+          if (url.endsWith("/api/mcp/connect") && init?.method === "POST") {
+            return new Response(JSON.stringify({ sessionId: "s-progress" }), {
+              status: 200,
+            });
+          }
+          if (url.includes("/api/mcp/events")) {
+            sse = createPushableSseStream();
+            return sse.response;
+          }
+          if (url.endsWith("/api/mcp/send") && init?.method === "POST") {
+            const body = JSON.parse(String(init.body)) as {
+              message: { id?: string | number };
+            };
+            sentId = body.message.id;
+            // Backend accepted the request; the response arrives later over SSE.
+            return new Response(JSON.stringify({ ok: true }), { status: 200 });
+          }
+          return new Response("not found", { status: 404 });
+        });
+      return { fetchFn, getSse: () => sse, getSentId: () => sentId };
+    }
+
+    // Fake timers so the deadline is driven deterministically rather than
+    // raced against a real `setTimeout` — a real-sleep test can let the timeout
+    // fire before a loaded runner processes the enqueued SSE frame (#1596: a
+    // race is fixed with fake timers, never with headroom). Each push is
+    // followed by a 0ms async advance to flush the SSE consumer's stream read
+    // so the frame is parsed before time moves.
+    const TIMEOUT_MS = 1000;
+    const flushSse = () => vi.advanceTimersByTimeAsync(0);
+
+    beforeEach(() => vi.useFakeTimers());
+    afterEach(() => vi.useRealTimers());
+
+    it("re-arms the wait on each progress note, so a call outliving the timeout still resolves", async () => {
+      const { fetchFn, getSse, getSentId } = backendHoldingSend();
+      const transport = new RemoteClientTransport(
+        { baseUrl, fetchFn, sseResponseTimeoutMs: TIMEOUT_MS },
+        config,
+      );
+      await transport.start();
+
+      const sent = transport.send({
+        jsonrpc: "2.0",
+        id: 1,
+        method: "tools/call",
+      });
+      await flushSse();
+
+      // Three progress notes, each after 900ms (< the 1000ms window): without
+      // the re-arm the wait dies at 1000ms, well before the third lands.
+      for (let i = 0; i < 3; i++) {
+        getSse().pushMessage({
+          jsonrpc: "2.0",
+          method: "notifications/progress",
+          params: { progressToken: getSentId(), progress: i, total: 3 },
+        });
+        await flushSse();
+        await vi.advanceTimersByTimeAsync(900);
+      }
+      getSse().pushMessage({
+        jsonrpc: "2.0",
+        id: getSentId(),
+        result: { ok: true },
+      });
+      await flushSse();
+
+      await expect(sent).resolves.toBeUndefined();
+      await transport.close();
+    });
+
+    it("does not re-arm on notifications/message, so a log-only call still times out", async () => {
+      const { fetchFn, getSse, getSentId } = backendHoldingSend();
+      const transport = new RemoteClientTransport(
+        { baseUrl, fetchFn, sseResponseTimeoutMs: TIMEOUT_MS },
+        config,
+      );
+      await transport.start();
+
+      const sent = transport.send({
+        jsonrpc: "2.0",
+        id: 1,
+        method: "tools/call",
+      });
+      // Hold the expected rejection now (as the RemoteSession timeout tests do)
+      // so it is asserted rather than suppressed by an empty catch.
+      const rejection = expect(sent).rejects.toThrow(
+        /Timed out waiting for MCP response on SSE/,
+      );
+      await flushSse();
+
+      // A log message before the deadline must not extend it.
+      getSse().pushMessage({
+        jsonrpc: "2.0",
+        method: "notifications/message",
+        params: { level: "info", data: { msg: "tick", token: getSentId() } },
+      });
+      await flushSse();
+      await vi.advanceTimersByTimeAsync(TIMEOUT_MS);
+
+      await rejection;
+      await transport.close();
+    });
+  });
 });

clients/web/src/test/integration/mcp/remote/remote-session.test.ts

@@ -7,11 +7,16 @@
  */
 
 import { describe, it, expect, vi } from "vitest";
-import type { Transport } from "@modelcontextprotocol/client";
+import type { Transport, JSONRPCMessage } from "@modelcontextprotocol/client";
 import { RemoteSession } from "@inspector/core/mcp/remote/node/remote-session.js";
 import type { FetchRequestEntryBase } from "@inspector/core/mcp/types.js";
 import { AuthChallengeError } from "@inspector/core/auth/challenge.js";
 
+// Stage a deliberately off-spec message as `unknown` and narrow with a single
+// cast — the guard branches these fixtures exercise are ones the typed
+// `JSONRPCMessage` shape forbids, and `unknown` is the safe top type to pass in.
+const asMsg = (m: unknown): JSONRPCMessage => m as JSONRPCMessage;
+
 function makeFetchEntry(
   overrides: Partial<FetchRequestEntryBase> = {},
 ): FetchRequestEntryBase {
@@ -231,6 +236,76 @@ describe("RemoteSession", () => {
     await expect(wait).resolves.toBeUndefined();
   });
 
+  it("a matching progress notification re-arms the wait past its original deadline (#2028)", async () => {
+    vi.useFakeTimers();
+    const session = new RemoteSession("s-progress-reset");
+    const wait = session.waitForRequestResponse(5, 1000);
+    // Almost at the deadline, then a progress note for this request arrives.
+    await vi.advanceTimersByTimeAsync(900);
+    session.onMessage({
+      jsonrpc: "2.0",
+      method: "notifications/progress",
+      params: { progressToken: 5, progress: 1, total: 10 },
+    });
+    // Past the *original* 1000ms deadline: without the re-arm this would have
+    // already rejected, and the resolve below would be a no-op on a dead wait.
+    await vi.advanceTimersByTimeAsync(900);
+    session.onMessage({ jsonrpc: "2.0", id: 5, result: {} });
+    await expect(wait).resolves.toBeUndefined();
+    vi.useRealTimers();
+  });
+
+  it("a notifications/message does NOT re-arm the wait (log messages don't extend the deadline, #2028)", async () => {
+    vi.useFakeTimers();
+    const session = new RemoteSession("s-message-no-reset");
+    const wait = session.waitForRequestResponse(6, 1000);
+    const rejection = expect(wait).rejects.toThrow(/timed out/);
+    await vi.advanceTimersByTimeAsync(500);
+    session.onMessage({
+      jsonrpc: "2.0",
+      method: "notifications/message",
+      params: { level: "info", data: { msg: "tick" } },
+    });
+    await vi.advanceTimersByTimeAsync(500);
+    await rejection;
+    vi.useRealTimers();
+  });
+
+  it("noteRequestProgress is a no-op for an id with no pending wait", () => {
+    const session = new RemoteSession("s-progress-noop");
+    expect(() => session.noteRequestProgress(999)).not.toThrow();
+  });
+
+  it("a progress notification on a timeoutMs=0 wait leaves it timerless and it still resolves", async () => {
+    const session = new RemoteSession("s-progress-notimer");
+    const wait = session.waitForRequestResponse(7, 0);
+    session.onMessage({
+      jsonrpc: "2.0",
+      method: "notifications/progress",
+      params: { progressToken: 7, progress: 1 },
+    });
+    session.onMessage({ jsonrpc: "2.0", id: 7, result: {} });
+    await expect(wait).resolves.toBeUndefined();
+  });
+
+  it("a progress notification with a non-token progressToken re-arms nothing", async () => {
+    vi.useFakeTimers();
+    const session = new RemoteSession("s-progress-bad-token");
+    const wait = session.waitForRequestResponse(8, 1000);
+    const rejection = expect(wait).rejects.toThrow(/timed out/);
+    // A malformed token (object) must be ignored, not coerced to a key.
+    session.onMessage(
+      asMsg({
+        jsonrpc: "2.0",
+        method: "notifications/progress",
+        params: { progressToken: { bad: true }, progress: 1 },
+      }),
+    );
+    await vi.advanceTimersByTimeAsync(1000);
+    await rejection;
+    vi.useRealTimers();
+  });
+
   it("endSend is a no-op when no send is active", () => {
     const session = new RemoteSession("s-endsend-noop");
     expect(session.hasActiveSend()).toBe(false);

core/mcp/remote/node/remote-session.ts

@@ -10,6 +10,7 @@ import type { AuthChallenge } from "../../../auth/challenge.js";
 import { AuthChallengeError } from "../../../auth/challenge.js";
 import type { RemoteAuthProviderHandle } from "./tokenAuthProvider.js";
 import type { RemoteAuthState } from "../types.js";
+import { progressTokenOf } from "../progressToken.js";
 
 export interface SessionEvent {
   type: RemoteEvent["type"];
@@ -19,6 +20,14 @@ export interface SessionEvent {
 type RequestWait = {
   resolve: () => void;
   reject: (error: Error) => void;
+  /**
+   * Re-arm this wait's timeout for another full window. Called when a
+   * `notifications/progress` for this request arrives, mirroring the SDK
+   * client's `resetTimeoutOnProgress` so the relay's own deadline never fires
+   * before the client's progress-aware one (#2028). A no-op when the wait was
+   * created without a timeout.
+   */
+  resetTimeout: () => void;
 };
 
 export class RemoteSession {
@@ -172,35 +181,57 @@ export class RemoteSession {
     timeoutMs = 60_000,
   ): Promise<void> {
     return new Promise((resolve, reject) => {
-      const timer =
-        timeoutMs > 0
-          ? setTimeout(() => {
-              this.requestWaits.delete(requestId);
-              reject(
-                new Error(
-                  `MCP request ${String(requestId)} timed out after ${timeoutMs}ms`,
-                ),
-              );
-            }, timeoutMs)
-          : undefined;
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const arm = () => {
+        if (timeoutMs <= 0) return;
+        timer = setTimeout(() => {
+          this.requestWaits.delete(requestId);
+          reject(
+            new Error(
+              `MCP request ${String(requestId)} timed out after ${timeoutMs}ms`,
+            ),
+          );
+        }, timeoutMs);
+      };
+      const clear = () => {
+        if (timer !== undefined) {
+          clearTimeout(timer);
+          timer = undefined;
+        }
+      };
+      arm();
 
       this.requestWaits.set(requestId, {
         resolve: () => {
-          if (timer !== undefined) {
-            clearTimeout(timer);
-          }
+          clear();
           resolve();
         },
         reject: (error) => {
-          if (timer !== undefined) {
-            clearTimeout(timer);
-          }
+          clear();
           reject(error);
         },
+        // A progress notification for this request re-arms the full window, so
+        // a tool that keeps reporting progress never trips the relay's own
+        // deadline — the same guarantee the SDK client's `resetTimeoutOnProgress`
+        // gives on a direct transport (#2028).
+        resetTimeout: () => {
+          if (timer === undefined) return;
+          clear();
+          arm();
+        },
       });
     });
   }
 
+  /**
+   * Re-arm a pending request's wait timeout. Invoked from {@link onMessage} when
+   * a `notifications/progress` carrying this request's `progressToken` arrives.
+   * A no-op when no wait is pending for the id.
+   */
+  noteRequestProgress(requestId: string | number): void {
+    this.requestWaits.get(requestId)?.resetTimeout();
+  }
+
   cancelRequestWait(requestId: string | number): void {
     const wait = this.requestWaits.get(requestId);
     if (!wait) {
@@ -265,6 +296,15 @@ export class RemoteSession {
     ) {
       this.settleRequestWait(message.id);
     }
+    // A `notifications/progress` carries the originating request's id in
+    // `params.progressToken` (the SDK stamps `progressToken: messageId`). Re-arm
+    // that request's wait so the relay's deadline tracks the client's
+    // progress-aware one instead of firing at a flat 60s (#2028). Log messages
+    // (`notifications/message`) deliberately do NOT reset it, matching the SDK.
+    const progressToken = progressTokenOf(message);
+    if (progressToken !== undefined) {
+      this.noteRequestProgress(progressToken);
+    }
     this.pushEvent({ type: "message", data: message });
   }
 

core/mcp/remote/progressToken.ts

@@ -0,0 +1,40 @@
+/**
+ * Shared helper for the remote relay's progress-aware request-response waits
+ * (#2028). Pure and browser-safe, so both the browser `RemoteClientTransport`
+ * and the Node `RemoteSession` import it rather than each carrying a copy.
+ */
+
+import { ProgressTokenSchema } from "@modelcontextprotocol/core";
+import type { JSONRPCMessage } from "@modelcontextprotocol/client";
+
+/**
+ * The token identifying the request a `notifications/progress` belongs to, or
+ * `undefined` for any other message. The SDK stamps `progressToken: messageId`
+ * on the outbound request, so a progress notification's `progressToken` is the
+ * id of the request whose wait should be re-armed.
+ *
+ * Validation is delegated to the SDK's own `ProgressTokenSchema` rather than a
+ * hand-rolled `typeof` check — the protocol allows only a string or a *safe*
+ * integer, and zod's `.int()` rejects a fractional or past-`MAX_SAFE_INTEGER`
+ * value that a bare `typeof token === "number"` would wave through (the same
+ * drift `isProgressToken` in `core/mcp/inspectorClient.ts` exists to avoid).
+ */
+export function progressTokenOf(
+  message: JSONRPCMessage,
+): string | number | undefined {
+  if (
+    "method" in message &&
+    message.method === "notifications/progress" &&
+    "params" in message &&
+    typeof message.params === "object" &&
+    message.params !== null
+  ) {
+    const parsed = ProgressTokenSchema.safeParse(
+      (message.params as { progressToken?: unknown }).progressToken,
+    );
+    if (parsed.success) {
+      return parsed.data;
+    }
+  }
+  return undefined;
+}

core/mcp/remote/remoteClientTransport.ts

@@ -31,6 +31,7 @@ import type {
   RemoteSendResponse,
 } from "./types.js";
 import { oauthTokensToRemoteAuthState } from "./types.js";
+import { progressTokenOf } from "./progressToken.js";
 
 export interface AuthRecoveryHandlers {
   handleAuthChallenge(
@@ -85,6 +86,12 @@ const DEFAULT_SSE_RESPONSE_TIMEOUT_MS = 60_000;
 type SseResponseWait = {
   resolve: () => void;
   reject: (error: Error) => void;
+  /**
+   * Re-arm this wait's timeout for another full window when a
+   * `notifications/progress` for the request arrives, mirroring the SDK
+   * client's `resetTimeoutOnProgress` (#2028).
+   */
+  resetTimeout: () => void;
 };
 
 function requestIdForMessage(
@@ -524,6 +531,13 @@ export class RemoteClientTransport implements Transport {
           if (parsed.type === "message") {
             const msg = parsed.data as JSONRPCMessage;
             this.settleSseResponseWait(msg);
+            // Keep the SSE wait alive while progress flows for a long call
+            // (#2028) — the settle above never matches a notification, so this
+            // is the only place a progress note touches the wait.
+            const progressToken = progressTokenOf(msg);
+            if (progressToken !== undefined) {
+              this.resetSseResponseWait(progressToken);
+            }
             this.onmessage?.(msg, undefined);
           } else if (
             parsed.type === "fetch_request" &&
@@ -675,14 +689,18 @@ export class RemoteClientTransport implements Transport {
 
   private waitForSseResponse(requestId: string | number): Promise<void> {
     return new Promise((resolve, reject) => {
-      const timer = setTimeout(() => {
-        this.sseResponseWaits.delete(requestId);
-        reject(
-          new Error(
-            `Timed out waiting for MCP response on SSE (${this.sseResponseTimeoutMs}ms)`,
-          ),
-        );
-      }, this.sseResponseTimeoutMs);
+      let timer: ReturnType<typeof setTimeout>;
+      const arm = () => {
+        timer = setTimeout(() => {
+          this.sseResponseWaits.delete(requestId);
+          reject(
+            new Error(
+              `Timed out waiting for MCP response on SSE (${this.sseResponseTimeoutMs}ms)`,
+            ),
+          );
+        }, this.sseResponseTimeoutMs);
+      };
+      arm();
       this.sseResponseWaits.set(requestId, {
         resolve: () => {
           clearTimeout(timer);
@@ -692,10 +710,22 @@ export class RemoteClientTransport implements Transport {
           clearTimeout(timer);
           reject(error);
         },
+        // A progress notification for this request re-arms the full window so
+        // the browser-side deadline tracks the SDK client's progress-aware one
+        // rather than firing at a flat 60s (#2028).
+        resetTimeout: () => {
+          clearTimeout(timer);
+          arm();
+        },
       });
     });
   }
 
+  /** Re-arm a pending request's SSE wait timeout on a matching progress note. */
+  private resetSseResponseWait(requestId: string | number): void {
+    this.sseResponseWaits.get(requestId)?.resetTimeout();
+  }
+
   private async postSend(
     message: JSONRPCMessage,
     options: TransportSendOptions | undefined,