All of the mid-session and step-up auth work

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

@@ -6,12 +6,21 @@ import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
 import type { FetchRequestEntryBase } from "../../types.js";
 import type { RemoteEvent } from "../types.js";
+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";
 
 export interface SessionEvent {
   type: RemoteEvent["type"];
   data: unknown;
 }
 
+type RequestWait = {
+  resolve: () => void;
+  reject: (error: Error) => void;
+};
+
 export class RemoteSession {
   public readonly sessionId: string;
   public transport!: Transport;
@@ -19,11 +28,31 @@ export class RemoteSession {
   private eventConsumer: ((event: SessionEvent) => void) | null = null;
   private transportDead: boolean = false;
   private transportError: string | null = null;
+  private activeSendCount = 0;
+  /**
+   * Suppress duplicate ambient SSE auth echoes after HTTP command-scoped delivery.
+   * Stale transport `onerror` may arrive after the send completes.
+   */
+  private authHttpEchoSuppressUntilMs = 0;
+  static readonly AUTH_HTTP_ECHO_SUPPRESS_MS = 30_000;
+  private readonly requestWaits = new Map<string | number, RequestWait>();
+  private authProviderHandle: RemoteAuthProviderHandle | null = null;
 
   constructor(sessionId: string) {
     this.sessionId = sessionId;
   }
 
+  setAuthProviderHandle(handle: RemoteAuthProviderHandle | null): void {
+    this.authProviderHandle = handle;
+  }
+
+  setAuthState(authState: RemoteAuthState): void {
+    if (!this.authProviderHandle) {
+      throw new Error("Session has no OAuth auth provider");
+    }
+    this.authProviderHandle.setAuthState(authState);
+  }
+
   setTransport(transport: Transport): void {
     this.transport = transport;
   }
@@ -72,6 +101,119 @@ export class RemoteSession {
     return this.eventConsumer !== null;
   }
 
+  beginSend(): void {
+    this.activeSendCount++;
+  }
+
+  endSend(): void {
+    if (this.activeSendCount > 0) {
+      this.activeSendCount--;
+    }
+  }
+
+  /** Command-scoped auth was returned on HTTP; suppress async onerror echo on SSE. */
+  noteAuthChallengeDeliveredViaHttp(): void {
+    this.extendAuthHttpEchoSuppression();
+  }
+
+  private extendAuthHttpEchoSuppression(): void {
+    this.authHttpEchoSuppressUntilMs =
+      Date.now() + RemoteSession.AUTH_HTTP_ECHO_SUPPRESS_MS;
+  }
+
+  private shouldSuppressAuthEcho(): boolean {
+    return Date.now() < this.authHttpEchoSuppressUntilMs;
+  }
+
+  hasActiveSend(): boolean {
+    return this.activeSendCount > 0;
+  }
+
+  pushAuthChallenge(challenge: AuthChallenge): void {
+    this.pushEvent({ type: "auth_challenge", data: challenge });
+  }
+
+  /**
+   * Streamable HTTP `transport.send()` can return before the JSON-RPC response
+   * arrives. The remote send handler awaits this so `/api/mcp/send` returns
+   * `auth_challenge` or `ok: true` only after the MCP round-trip (or auth error).
+   */
+  waitForRequestResponse(
+    requestId: string | number,
+    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;
+
+      this.requestWaits.set(requestId, {
+        resolve: () => {
+          if (timer !== undefined) {
+            clearTimeout(timer);
+          }
+          resolve();
+        },
+        reject: (error) => {
+          if (timer !== undefined) {
+            clearTimeout(timer);
+          }
+          reject(error);
+        },
+      });
+    });
+  }
+
+  cancelRequestWait(requestId: string | number): void {
+    this.requestWaits.delete(requestId);
+  }
+
+  rejectActiveRequestWaits(error: Error): void {
+    for (const wait of this.requestWaits.values()) {
+      wait.reject(error);
+    }
+    this.requestWaits.clear();
+  }
+
+  /**
+   * Auth errors from the MCP transport.
+   * - During an active send: command path owns delivery (HTTP); never SSE.
+   * - After HTTP already returned auth_challenge: swallow async onerror echo.
+   * - Otherwise: ambient SSE (idle session).
+   */
+  handleTransportAuthError(error: unknown): error is AuthChallengeError {
+    if (!(error instanceof AuthChallengeError)) {
+      return false;
+    }
+    this.rejectActiveRequestWaits(error);
+    if (this.hasActiveSend()) {
+      this.extendAuthHttpEchoSuppression();
+      return true;
+    }
+    if (this.shouldSuppressAuthEcho()) {
+      return true;
+    }
+    this.pushAuthChallenge(error.authChallenge);
+    return true;
+  }
+
+  private settleRequestWait(requestId: string | number): void {
+    const wait = this.requestWaits.get(requestId);
+    if (!wait) {
+      return;
+    }
+    this.requestWaits.delete(requestId);
+    wait.resolve();
+  }
+
   pushEvent(event: SessionEvent): void {
     if (this.eventConsumer) {
       this.eventConsumer(event);
@@ -81,6 +223,14 @@ export class RemoteSession {
   }
 
   onMessage(message: JSONRPCMessage): void {
+    if (
+      "id" in message &&
+      message.id !== null &&
+      message.id !== undefined &&
+      ("result" in message || "error" in message)
+    ) {
+      this.settleRequestWait(message.id);
+    }
     this.pushEvent({ type: "message", data: message });
   }