@runuai/host 0.9.147 → 0.9.149

package/lib/mcp-gateway.ts

@@ -26,12 +26,17 @@ import {
   writeFileSync,
 } from "node:fs";
 import {
-  createServer,
+  createServer as createHttpServer,
   type IncomingMessage,
-  type Server,
+  type Server as HttpServer,
   type ServerResponse,
 } from "node:http";
 import { randomBytes, timingSafeEqual } from "node:crypto";
+import {
+  createServer as createNetServer,
+  type Server as NetServer,
+  type Socket,
+} from "node:net";
 import { Readable, pipeline } from "node:stream";
 import { resolve } from "node:path";
 
@@ -56,24 +61,31 @@ import { agentClisReady } from "./standard-image";
 import { isSafeHostTaskId, assertSafeHostTaskId } from "./task-identity";
 import type { McpGatewayCapability } from "../src/protocol";
 import { isEnvironmentTransportProvider } from "./task-environment/types";
+import { sidecarTaskId } from "./sidecar";
+import {
+  TASK_FORWARDER_ACCEPT_MAGIC,
+  TASK_FORWARDER_MAX_PREAMBLE_BYTES,
+  TASK_FORWARDER_PREAMBLE_HEADER_BYTES,
+  TASK_FORWARDER_RESPONSE_BYTES,
+  decodeTaskForwarderPreamble,
+  decodeTaskForwarderResponse,
+  encodeTaskForwarderChallenge,
+  taskForwarderPreambleLength,
+  type TaskForwarderIdentity,
+} from "./task-gateway-protocol";
+import {
+  currentTaskGatewayRegistration,
+  isTaskGatewayGenerationCurrent,
+  resolveTaskGatewayEnv,
+  subscribeTaskGatewayRevocations,
+  verifyTaskGatewayAuthentication,
+} from "./task-gateway-state";
+import { MCP_GATEWAY_PORT } from "./task-gateway-port";
 
 export type { McpGatewayCapability } from "../src/protocol";
 
-const DEFAULT_MCP_GATEWAY_PORT = 5877;
-
-function gatewayPort(value: string | undefined): number {
-  if (value === undefined || !/^\d+$/.test(value)) {
-    return DEFAULT_MCP_GATEWAY_PORT;
-  }
-  const port = Number(value);
-  return Number.isSafeInteger(port) && port >= 1 && port <= 65_535
-    ? port
-    : DEFAULT_MCP_GATEWAY_PORT;
-}
+export { MCP_GATEWAY_PORT };
 
-export const MCP_GATEWAY_PORT = gatewayPort(
-  process.env.UAI_MCP_GATEWAY_PORT,
-);
 /** Loopback works only where the container runtime forwards
  * host.docker.internal INTO the host loopback — macOS Docker
  * (Desktop/OrbStack). Native Linux docker resolves host.docker.internal to
@@ -102,7 +114,7 @@ function gatewayBindAddress(): string {
  * throws so the config write fails visibly and is retried on the next session
  * ensure, rather than baking a dead URL into the task. */
 let cachedAppleGatewayHost: string | null = null;
-async function taskGatewayHost(): Promise<string> {
+export async function taskGatewayHost(): Promise<string> {
   if (pinnedContainerRuntimeProvider() !== "apple-container") {
     return "host.docker.internal";
   }
@@ -145,6 +157,21 @@ const GATEWAY_RETRY_INITIAL_DELAY_MS = 1_000;
 const GATEWAY_RETRY_MAX_DELAY_MS = 15_000;
 const GATEWAY_RETRY_MIN_JITTER = 0.8;
 const GATEWAY_RETRY_JITTER_SPAN = 0.4;
+const TASK_GATEWAY_PREAUTH_TIMEOUT_MS = 3_000;
+const TASK_GATEWAY_MAX_PENDING_PREAUTH = 128;
+const TASK_GATEWAY_MAX_PENDING_PREAUTH_PER_PEER = 16;
+const TASK_GATEWAY_MAX_PENDING_PREAUTH_PER_TASK = 16;
+// MCP streamable-HTTP/SSE and proxied WebSockets can legitimately hold a
+// connection for the life of a session, so active limits are deliberately
+// separate from the tight three-second handshake limits above. They are still
+// finite: once a socket has supplied its first application byte, HTTP's header
+// and request deadlines bound time, while these counters bound cardinality.
+const TASK_GATEWAY_MAX_ACTIVE_CONNECTIONS = 2_048;
+// Direct sockets have not proved which task they belong to yet. Treat this as
+// a small, evictable staging allowance; tokenized MCP moves to the task quota
+// as soon as Node parses and validates the request path.
+const TASK_GATEWAY_MAX_DIRECT_STAGING_CONNECTIONS_PER_PEER = 16;
+const TASK_GATEWAY_MAX_ACTIVE_CONNECTIONS_PER_TASK = 256;
 
 export interface TaskMcpConnection {
   id: string;
@@ -276,7 +303,82 @@ const FORWARD_REQ_HEADERS = [
 ];
 const FORWARD_RES_HEADERS = ["content-type", "mcp-session-id"];
 
-async function handle(req: IncomingMessage, res: ServerResponse): Promise<void> {
+export interface AuthenticatedTaskGatewayConnection
+  extends TaskForwarderIdentity {
+  readonly socket: Socket;
+  /** Re-check immediately before using task authority. */
+  readonly isCurrent: () => boolean;
+}
+
+export type TaskGatewayTlsHandler = (
+  connection: AuthenticatedTaskGatewayConnection,
+) => void;
+
+let taskGatewayTlsHandler: TaskGatewayTlsHandler | null = null;
+
+/**
+ * Register the TLS consumer used by the ADR-153 secrets proxy. There is one
+ * gateway and therefore one consumer; a second registration is a wiring bug.
+ */
+export function registerTaskGatewayTlsHandler(
+  handler: TaskGatewayTlsHandler,
+): () => void {
+  if (taskGatewayTlsHandler !== null) {
+    throw new Error("task gateway TLS handler is already registered");
+  }
+  taskGatewayTlsHandler = handler;
+  return () => {
+    if (taskGatewayTlsHandler === handler) taskGatewayTlsHandler = null;
+  };
+}
+
+function loopbackPeer(socket: Socket): boolean {
+  const address = socket.remoteAddress ?? "";
+  return (
+    address === "::1" ||
+    /^127(?:\.\d{1,3}){3}$/.test(address) ||
+    /^::ffff:127(?:\.\d{1,3}){3}$/i.test(address)
+  );
+}
+
+async function handleEnvRequest(
+  req: IncomingMessage,
+  res: ServerResponse,
+  identity: TaskForwarderIdentity | null,
+): Promise<void> {
+  // Keep the task hook's deliberately tiny response grammar stable. Node adds
+  // Date by default, but the response is local, non-cacheable, and has no use
+  // for a wall-clock header.
+  res.sendDate = false;
+  if (identity === null) return deny(res, 404, "not found");
+  if (req.method !== "GET") {
+    res.setHeader("allow", "GET");
+    return deny(res, 405, "method not allowed");
+  }
+  const snapshot = resolveTaskGatewayEnv(identity);
+  if (snapshot === null) return deny(res, 503, "task env unavailable");
+  const headers = {
+    etag: snapshot.etag,
+    "cache-control": "no-store",
+  };
+  if (req.headers["if-none-match"] === snapshot.etag) {
+    res.writeHead(304, { ...headers, "content-length": "0" });
+    res.end();
+    return;
+  }
+  res.writeHead(200, {
+    ...headers,
+    "content-type": "text/plain; charset=utf-8",
+    "content-length": String(snapshot.body.length),
+  });
+  res.end(snapshot.body);
+}
+
+async function handleMcpRequest(
+  req: IncomingMessage,
+  res: ServerResponse,
+  claimTaskQuota: (socket: Socket, taskId: string) => boolean,
+): Promise<void> {
   const match = /^\/t\/([^/]+)\/([^/?]+)\/?(?:\?.*)?$/.exec(req.url ?? "");
   if (!match) return deny(res, 404, "not found");
   const [, token, slug] = match as unknown as [string, string, string];
@@ -293,6 +395,16 @@ async function handle(req: IncomingMessage, res: ServerResponse): Promise<void>
   ) {
     return deny(res, 401, "unknown task token");
   }
+  // A direct connection is admitted only to an evictable peer staging pool.
+  // Once its unguessable ACL token proves the task, move the close-scoped
+  // reservation into that task's quota before using any task authority. An
+  // authenticated forwarder is already in this quota, and must name the same
+  // task if it sends an MCP request over its connection.
+  if (!claimTaskQuota(req.socket, taskId)) {
+    res.shouldKeepAlive = false;
+    res.setHeader("connection", "close");
+    return deny(res, 429, "task connection limit reached");
+  }
   const entry = acl.connections.find((c) => c.slug === slug);
   if (!entry) return deny(res, 404, `no connection "${slug}" for this task`);
   const conn = getConnection(entry.id);
@@ -383,9 +495,482 @@ async function handle(req: IncomingMessage, res: ServerResponse): Promise<void>
   }
 }
 
+async function handleGatewayRequest(
+  req: IncomingMessage,
+  res: ServerResponse,
+  authenticatedSockets: WeakMap<Socket, TaskForwarderIdentity>,
+  claimTaskQuota: (socket: Socket, taskId: string) => boolean,
+): Promise<void> {
+  const pathname = (() => {
+    try {
+      return new URL(req.url ?? "", "http://uai.invalid").pathname;
+    } catch {
+      return "";
+    }
+  })();
+  if (pathname === "/env" && req.url === "/env") {
+    const socket = req.socket;
+    let identity = authenticatedSockets.get(socket) ?? null;
+    if (identity === null) {
+      const fixedTaskId = sidecarTaskId();
+      if (fixedTaskId !== null && loopbackPeer(socket)) {
+        identity = currentTaskGatewayRegistration(fixedTaskId);
+      }
+    }
+    await handleEnvRequest(req, res, identity);
+    return;
+  }
+  await handleMcpRequest(req, res, claimTaskQuota);
+}
+
+type GatewayTransportServer = NetServer & {
+  closeAllConnections?: () => void;
+};
+
+function readSocketBytes(
+  socket: Socket,
+  length: number,
+  deadline: number,
+): Promise<Buffer> {
+  return new Promise((resolveRead, rejectRead) => {
+    let settled = false;
+    const finish = (error: Error | null, bytes?: Buffer): void => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timer);
+      socket.off("readable", onReadable);
+      socket.off("error", onError);
+      socket.off("end", onEnd);
+      socket.off("close", onClose);
+      if (error) rejectRead(error);
+      else if (bytes) resolveRead(bytes);
+      else rejectRead(new Error("task gateway stream ended"));
+    };
+    const tryRead = (): void => {
+      const bytes = socket.read(length) as Buffer | null;
+      if (bytes === null) return;
+      if (bytes.length !== length) {
+        finish(new Error("truncated task gateway frame"));
+        return;
+      }
+      finish(null, bytes);
+    };
+    const onReadable = (): void => tryRead();
+    const onError = (): void => finish(new Error("task gateway socket error"));
+    const onEnd = (): void => {
+      tryRead();
+      if (!settled) finish(new Error("truncated task gateway frame"));
+    };
+    const onClose = (): void => finish(new Error("task gateway socket closed"));
+    const remainingMs = Math.max(1, deadline - Date.now());
+    const timer = setTimeout(
+      () => finish(new Error("task gateway pre-auth timed out")),
+      remainingMs,
+    );
+    timer.unref();
+    socket.on("readable", onReadable);
+    socket.once("error", onError);
+    socket.once("end", onEnd);
+    socket.once("close", onClose);
+    tryRead();
+  });
+}
+
+interface TaskGatewayDemuxOptions {
+  readonly httpServer: HttpServer;
+  readonly authenticatedSockets: WeakMap<Socket, TaskForwarderIdentity>;
+  readonly authenticatedConnections: Map<Socket, TaskForwarderIdentity>;
+  readonly preauthTimeoutMs: number;
+  readonly maxPendingPreauth: number;
+  readonly maxPendingPreauthPerPeer: number;
+  readonly maxPendingPreauthPerTask: number;
+  readonly maxActiveConnections: number;
+  readonly maxDirectStagingConnectionsPerPeer: number;
+  readonly maxActiveConnectionsPerTask: number;
+  readonly nonce: () => Buffer;
+  readonly tlsHandler: () => TaskGatewayTlsHandler | null;
+}
+
+interface AuthenticatedTaskForwarder {
+  readonly identity: TaskForwarderIdentity;
+  /** Release the task-scoped pre-auth reservation exactly once. */
+  readonly release: () => void;
+}
+
+function createTaskGatewayDemux(options: TaskGatewayDemuxOptions): {
+  accept: (socket: Socket) => void;
+  claimTaskQuota: (socket: Socket, taskId: string) => boolean;
+} {
+  let pendingPreauth = 0;
+  const pendingPeerBySocket = new Map<Socket, string>();
+  const pendingSocketsByPeer = new Map<string, Set<Socket>>();
+  const pendingByTask = new Map<string, number>();
+  let activeConnections = 0;
+  const activeScopeBySocket = new Map<
+    Socket,
+    { kind: "peer" | "task"; key: string }
+  >();
+  const activeSocketsByPeer = new Map<string, Set<Socket>>();
+  const activeSocketsByTask = new Map<string, Set<Socket>>();
+  const activeUnauthenticatedSockets = new Set<Socket>();
+
+  const peerKey = (socket: Socket): string => {
+    if (loopbackPeer(socket)) return "loopback";
+    const address = (socket.remoteAddress ?? "unknown").toLowerCase();
+    return address.startsWith("::ffff:") ? address.slice(7) : address;
+  };
+
+  const releasePreauth = (socket: Socket): void => {
+    const peer = pendingPeerBySocket.get(socket);
+    if (peer === undefined) return;
+    pendingPeerBySocket.delete(socket);
+    pendingPreauth -= 1;
+    const peerSockets = pendingSocketsByPeer.get(peer);
+    peerSockets?.delete(socket);
+    if (peerSockets?.size === 0) pendingSocketsByPeer.delete(peer);
+  };
+
+  const reservePreauth = (socket: Socket, peer: string): void => {
+    pendingPreauth += 1;
+    pendingPeerBySocket.set(socket, peer);
+    const peerSockets = pendingSocketsByPeer.get(peer) ?? new Set<Socket>();
+    peerSockets.add(socket);
+    pendingSocketsByPeer.set(peer, peerSockets);
+  };
+
+  const releaseActive = (socket: Socket): void => {
+    const scope = activeScopeBySocket.get(socket);
+    if (scope === undefined) return;
+    activeScopeBySocket.delete(socket);
+    activeConnections -= 1;
+    const scopedSockets =
+      scope.kind === "task"
+        ? activeSocketsByTask.get(scope.key)
+        : activeSocketsByPeer.get(scope.key);
+    scopedSockets?.delete(socket);
+    if (scope.kind === "peer") activeUnauthenticatedSockets.delete(socket);
+    if (scopedSockets?.size === 0) {
+      if (scope.kind === "task") activeSocketsByTask.delete(scope.key);
+      else activeSocketsByPeer.delete(scope.key);
+    }
+  };
+
+  const evictUnauthenticated = (socket: Socket): void => {
+    // Release synchronously: Socket's close event is asynchronous and the
+    // replacement must be able to reserve the slot in this turn.
+    releaseActive(socket);
+    socket.destroy();
+  };
+
+  /** Transfer one fully-demultiplexed socket into a close-scoped quota.
+   * Forwarder-authenticated traffic and fixed-sidecar traffic are charged to
+   * their proved task. Direct HTTP starts in an evictable peer staging pool;
+   * a valid MCP token moves it to the non-evictable task pool below. */
+  const reserveActive = (
+    socket: Socket,
+    scope: { kind: "peer" | "task"; key: string },
+  ): boolean => {
+    if (activeScopeBySocket.has(socket)) return true;
+    const scopedSockets =
+      scope.kind === "task"
+        ? activeSocketsByTask.get(scope.key)
+        : activeSocketsByPeer.get(scope.key);
+    const scopedLimit =
+      scope.kind === "task"
+        ? options.maxActiveConnectionsPerTask
+        : options.maxDirectStagingConnectionsPerPeer;
+    if ((scopedSockets?.size ?? 0) >= scopedLimit) {
+      if (scope.kind === "task") return false;
+      const oldest = scopedSockets?.values().next().value as Socket | undefined;
+      if (oldest) evictUnauthenticated(oldest);
+    }
+    if (activeConnections >= options.maxActiveConnections) {
+      // Never evict a proved task socket: it may be a long-lived MCP SSE stream
+      // or secrets-proxy tunnel. Unproved direct traffic is replaceable, which
+      // keeps a handful of peers from filling the global pool permanently.
+      const oldest = activeUnauthenticatedSockets.values().next().value as
+        | Socket
+        | undefined;
+      if (!oldest) return false;
+      evictUnauthenticated(oldest);
+    }
+    const currentScopedSockets =
+      scope.kind === "task"
+        ? activeSocketsByTask.get(scope.key)
+        : activeSocketsByPeer.get(scope.key);
+    const nextScopedSockets = currentScopedSockets ?? new Set<Socket>();
+    nextScopedSockets.add(socket);
+    if (scope.kind === "task") {
+      activeSocketsByTask.set(scope.key, nextScopedSockets);
+    } else {
+      activeSocketsByPeer.set(scope.key, nextScopedSockets);
+      activeUnauthenticatedSockets.add(socket);
+    }
+    activeScopeBySocket.set(socket, scope);
+    activeConnections += 1;
+    socket.once("close", () => releaseActive(socket));
+    return true;
+  };
+
+  const claimTaskQuota = (socket: Socket, taskId: string): boolean => {
+    const current = activeScopeBySocket.get(socket);
+    if (current === undefined) return false;
+    if (current.kind === "task") return current.key === taskId;
+
+    const taskSockets = activeSocketsByTask.get(taskId);
+    if (
+      (taskSockets?.size ?? 0) >= options.maxActiveConnectionsPerTask
+    ) {
+      return false;
+    }
+
+    const peerSockets = activeSocketsByPeer.get(current.key);
+    peerSockets?.delete(socket);
+    if (peerSockets?.size === 0) activeSocketsByPeer.delete(current.key);
+    activeUnauthenticatedSockets.delete(socket);
+
+    const nextTaskSockets = taskSockets ?? new Set<Socket>();
+    nextTaskSockets.add(socket);
+    activeSocketsByTask.set(taskId, nextTaskSockets);
+    activeScopeBySocket.set(socket, { kind: "task", key: taskId });
+    return true;
+  };
+
+  const rememberAuthenticated = (
+    socket: Socket,
+    identity: TaskForwarderIdentity,
+  ): void => {
+    if (options.authenticatedConnections.has(socket)) return;
+    options.authenticatedSockets.set(socket, identity);
+    options.authenticatedConnections.set(socket, identity);
+    socket.once("close", () => {
+      options.authenticatedConnections.delete(socket);
+    });
+  };
+
+  const dispatchHttp = (
+    socket: Socket,
+    firstByte: Buffer,
+    identity: TaskForwarderIdentity | null,
+    activeScope: { kind: "peer" | "task"; key: string },
+  ): void => {
+    if (!reserveActive(socket, activeScope)) {
+      socket.destroy();
+      return;
+    }
+    if (identity) rememberAuthenticated(socket, identity);
+    socket.unshift(firstByte);
+    options.httpServer.emit("connection", socket);
+  };
+
+  const dispatchTls = (
+    socket: Socket,
+    firstByte: Buffer,
+    identity: TaskForwarderIdentity,
+  ): void => {
+    const handler = options.tlsHandler();
+    if (
+      !handler ||
+      !isTaskGatewayGenerationCurrent(identity) ||
+      !reserveActive(socket, { kind: "task", key: identity.taskId })
+    ) {
+      socket.destroy();
+      return;
+    }
+    rememberAuthenticated(socket, identity);
+    socket.unshift(firstByte);
+    try {
+      handler({
+        ...identity,
+        socket,
+        isCurrent: () => isTaskGatewayGenerationCurrent(identity),
+      });
+    } catch {
+      socket.destroy();
+    }
+  };
+
+  const authenticate = async (
+    socket: Socket,
+    firstByte: Buffer,
+    deadline: number,
+  ): Promise<AuthenticatedTaskForwarder | null> => {
+    const headerTail = await readSocketBytes(
+      socket,
+      TASK_FORWARDER_PREAMBLE_HEADER_BYTES - 1,
+      deadline,
+    );
+    const header = Buffer.concat([firstByte, headerTail]);
+    const preambleLength = taskForwarderPreambleLength(header);
+    if (
+      preambleLength === null ||
+      preambleLength > TASK_FORWARDER_MAX_PREAMBLE_BYTES
+    ) {
+      return null;
+    }
+    const payload = await readSocketBytes(
+      socket,
+      preambleLength - header.length,
+      deadline,
+    );
+    const identity = decodeTaskForwarderPreamble(
+      Buffer.concat([header, payload], preambleLength),
+    );
+    if (identity === null || !isSafeHostTaskId(identity.taskId)) return null;
+
+    const nonce = options.nonce();
+    socket.write(encodeTaskForwarderChallenge(nonce));
+    const response = await readSocketBytes(
+      socket,
+      TASK_FORWARDER_RESPONSE_BYTES,
+      deadline,
+    );
+    const mac = decodeTaskForwarderResponse(response);
+    if (
+      mac === null ||
+      !verifyTaskGatewayAuthentication(identity, nonce, mac)
+    ) {
+      return null;
+    }
+
+    // A claimed task id is untrusted until its MAC verifies. Charging its
+    // reservation earlier lets a peer lock out a real task with only the
+    // public task id.
+    const taskPending = pendingByTask.get(identity.taskId) ?? 0;
+    if (taskPending >= options.maxPendingPreauthPerTask) return null;
+    pendingByTask.set(identity.taskId, taskPending + 1);
+    let released = false;
+    const release = (): void => {
+      if (released) return;
+      released = true;
+      const next = (pendingByTask.get(identity.taskId) ?? 1) - 1;
+      if (next > 0) pendingByTask.set(identity.taskId, next);
+      else pendingByTask.delete(identity.taskId);
+    };
+    if (!reserveActive(socket, { kind: "task", key: identity.taskId })) {
+      release();
+      return null;
+    }
+    let retained = false;
+    try {
+      socket.write(TASK_FORWARDER_ACCEPT_MAGIC);
+      retained = true;
+      return { identity, release };
+    } finally {
+      // A synchronous write failure must not permanently consume this task's
+      // post-MAC quotas. Once returned, the caller owns the pending release;
+      // the active reservation remains close-scoped.
+      if (!retained) {
+        release();
+        releaseActive(socket);
+      }
+    }
+  };
+
+  return {
+    claimTaskQuota,
+    accept: (socket) => {
+      socket.pause();
+      socket.setNoDelay(true);
+      const peer = peerKey(socket);
+      const peerSockets = pendingSocketsByPeer.get(peer);
+      if (
+        peerSockets &&
+        peerSockets.size >= options.maxPendingPreauthPerPeer
+      ) {
+        // Prefer the newest connection instead of letting a peer reserve its
+        // own entire quota with idle sockets. A real hook/forwarder sends its
+        // first frame immediately, so replacing the oldest pending socket
+        // preserves progress even when several task runtimes share one NATed
+        // source address (notably desktop container VMs).
+        const oldest = peerSockets.values().next().value as Socket | undefined;
+        if (oldest) {
+          releasePreauth(oldest);
+          oldest.destroy();
+        }
+      }
+      if (pendingPreauth >= options.maxPendingPreauth) {
+        // Pending sockets carry no authority and have a short deadline. Drop
+        // the oldest globally so several peers cannot permanently consume all
+        // admission slots before a real forwarder gets a turn.
+        const oldest = pendingPeerBySocket.keys().next().value as
+          | Socket
+          | undefined;
+        if (oldest) {
+          releasePreauth(oldest);
+          oldest.destroy();
+        }
+      }
+      reservePreauth(socket, peer);
+      const deadline = Date.now() + options.preauthTimeoutMs;
+      void (async () => {
+        try {
+          const firstByte = await readSocketBytes(socket, 1, deadline);
+          if (firstByte[0] === 0x00) {
+            const authenticated = await authenticate(socket, firstByte, deadline);
+            if (authenticated === null) {
+              socket.destroy();
+              return;
+            }
+            const { identity } = authenticated;
+            try {
+              rememberAuthenticated(socket, identity);
+              if (!isTaskGatewayGenerationCurrent(identity)) {
+                socket.destroy();
+                return;
+              }
+              // The forwarder does not relay task bytes until it receives the
+              // accept frame, so this byte cannot be unauthenticated pipelining.
+              const applicationByte = await readSocketBytes(socket, 1, deadline);
+              if (applicationByte[0] === 0x16) {
+                dispatchTls(socket, applicationByte, identity);
+              } else {
+                dispatchHttp(socket, applicationByte, identity, {
+                  kind: "task",
+                  key: identity.taskId,
+                });
+              }
+            } finally {
+              authenticated.release();
+            }
+            return;
+          }
+
+          const fixedTaskId = sidecarTaskId();
+          const sidecarIdentity =
+            fixedTaskId !== null && loopbackPeer(socket)
+              ? currentTaskGatewayRegistration(fixedTaskId)
+              : null;
+          if (firstByte[0] === 0x16) {
+            // Direct TLS is valid only for the fixed one-task sidecar. Docker
+            // and Apple must prove identity through the forwarder first.
+            if (sidecarIdentity) dispatchTls(socket, firstByte, sidecarIdentity);
+            else socket.destroy();
+          } else {
+            dispatchHttp(
+              socket,
+              firstByte,
+              sidecarIdentity,
+              sidecarIdentity
+                ? { kind: "task", key: sidecarIdentity.taskId }
+                : { kind: "peer", key: peer },
+            );
+          }
+        } catch {
+          socket.destroy();
+        } finally {
+          releasePreauth(socket);
+        }
+      })();
+    },
+  };
+}
+
 export interface McpGatewayController {
   /** Current bounded capability snapshot. Never contains a raw listener error. */
   state(): McpGatewayCapability;
+  /** Resolve only after this listener has positively bound its configured port. */
+  waitUntilReady(timeoutMs?: number): Promise<void>;
   /** Subscribe to semantic state changes. The current snapshot is read via state(). */
   subscribe(
     listener: (state: McpGatewayCapability) => void,
@@ -405,8 +990,20 @@ export interface McpGatewayListenerOptions {
   retryInitialDelayMs?: number;
   retryMaxDelayMs?: number;
   serverFactory?: (
+    connectionListener: (socket: Socket) => void,
+  ) => GatewayTransportServer;
+  httpServerFactory?: (
     requestListener: (req: IncomingMessage, res: ServerResponse) => void,
-  ) => Server;
+  ) => HttpServer;
+  tlsHandler?: TaskGatewayTlsHandler;
+  preauthTimeoutMs?: number;
+  maxPendingPreauth?: number;
+  maxPendingPreauthPerPeer?: number;
+  maxPendingPreauthPerTask?: number;
+  maxActiveConnections?: number;
+  maxDirectStagingConnectionsPerPeer?: number;
+  maxActiveConnectionsPerTask?: number;
+  nonce?: () => Buffer;
   setRetryTimeout?: (
     callback: () => void,
     delayMs: number,
@@ -468,7 +1065,40 @@ export function createMcpGatewayListener(
       GATEWAY_RETRY_MAX_DELAY_MS,
     ),
   );
-  const serverFactory = options.serverFactory ?? createServer;
+  const serverFactory =
+    options.serverFactory ??
+    ((connectionListener: (socket: Socket) => void) =>
+      createNetServer({ pauseOnConnect: true }, connectionListener));
+  const httpServerFactory = options.httpServerFactory ?? createHttpServer;
+  const preauthTimeoutMs = positiveInteger(
+    options.preauthTimeoutMs,
+    TASK_GATEWAY_PREAUTH_TIMEOUT_MS,
+  );
+  const maxPendingPreauth = positiveInteger(
+    options.maxPendingPreauth,
+    TASK_GATEWAY_MAX_PENDING_PREAUTH,
+  );
+  const maxPendingPreauthPerPeer = positiveInteger(
+    options.maxPendingPreauthPerPeer,
+    TASK_GATEWAY_MAX_PENDING_PREAUTH_PER_PEER,
+  );
+  const maxPendingPreauthPerTask = positiveInteger(
+    options.maxPendingPreauthPerTask,
+    TASK_GATEWAY_MAX_PENDING_PREAUTH_PER_TASK,
+  );
+  const maxActiveConnections = positiveInteger(
+    options.maxActiveConnections,
+    TASK_GATEWAY_MAX_ACTIVE_CONNECTIONS,
+  );
+  const maxDirectStagingConnectionsPerPeer = positiveInteger(
+    options.maxDirectStagingConnectionsPerPeer,
+    TASK_GATEWAY_MAX_DIRECT_STAGING_CONNECTIONS_PER_PEER,
+  );
+  const maxActiveConnectionsPerTask = positiveInteger(
+    options.maxActiveConnectionsPerTask,
+    TASK_GATEWAY_MAX_ACTIVE_CONNECTIONS_PER_TASK,
+  );
+  const nonce = options.nonce ?? (() => randomBytes(32));
   const setRetryTimeout =
     options.setRetryTimeout ??
     ((callback: () => void, delayMs: number) =>
@@ -479,12 +1109,55 @@ export function createMcpGatewayListener(
   const subscribers = new Set<
     (state: McpGatewayCapability) => void
   >();
+  const authenticatedSockets = new WeakMap<Socket, TaskForwarderIdentity>();
+  const authenticatedConnections = new Map<Socket, TaskForwarderIdentity>();
+  let claimTaskQuota = (_socket: Socket, _taskId: string): boolean => false;
+  const httpServer = httpServerFactory((req, res) => {
+    void handleGatewayRequest(
+      req,
+      res,
+      authenticatedSockets,
+      claimTaskQuota,
+    ).catch(() => deny(res, 500, "gateway error"));
+  });
+  httpServer.on("clientError", (_error, socket) => socket.destroy());
+  // This HTTP server never owns the TCP listener: authenticated sockets are
+  // demultiplexed into it. Node starts its headers/request timeout checker on
+  // the `listening` event, so emit that lifecycle event for this hidden
+  // server too or a peer can drip one byte forever after pre-authentication.
+  httpServer.emit("listening");
+  const demux = createTaskGatewayDemux({
+    httpServer,
+    authenticatedSockets,
+    authenticatedConnections,
+    preauthTimeoutMs,
+    maxPendingPreauth,
+    maxPendingPreauthPerPeer,
+    maxPendingPreauthPerTask,
+    maxActiveConnections,
+    maxDirectStagingConnectionsPerPeer,
+    maxActiveConnectionsPerTask,
+    nonce,
+    tlsHandler: () => options.tlsHandler ?? taskGatewayTlsHandler,
+  });
+  claimTaskQuota = demux.claimTaskQuota;
+  const socketsByServer = new Map<GatewayTransportServer, Set<Socket>>();
+  const unsubscribeRevocations = subscribeTaskGatewayRevocations((revoked) => {
+    for (const [socket, identity] of authenticatedConnections) {
+      if (
+        identity.taskId === revoked.taskId &&
+        identity.generation === revoked.generation
+      ) {
+        socket.destroy();
+      }
+    }
+  });
 
   let capability: McpGatewayCapability = Object.freeze({
     status: "starting",
     port,
   });
-  let server: Server | null = null;
+  let server: GatewayTransportServer | null = null;
   let retryTimer: NodeJS.Timeout | null = null;
   let consecutiveFailures = 0;
   let started = false;
@@ -504,7 +1177,13 @@ export function createMcpGatewayListener(
     }
   };
 
-  const closeServer = (target: Server): void => {
+  const closeServer = (target: GatewayTransportServer): void => {
+    const sockets = socketsByServer.get(target);
+    socketsByServer.delete(target);
+    if (sockets) {
+      for (const socket of sockets) socket.destroy();
+      sockets.clear();
+    }
     try {
       target.close(() => undefined);
       target.closeAllConnections?.();
@@ -555,7 +1234,10 @@ export function createMcpGatewayListener(
     retryTimer.unref();
   };
 
-  const listenerFailed = (target: Server, error: unknown): void => {
+  const listenerFailed = (
+    target: GatewayTransportServer,
+    error: unknown,
+  ): void => {
     if (stopped || server !== target) return;
     server = null;
     closeServer(target);
@@ -565,15 +1247,19 @@ export function createMcpGatewayListener(
   beginAttempt = (): void => {
     if (stopped || server !== null || retryTimer !== null) return;
 
-    let nextServer: Server;
+    let nextServer: GatewayTransportServer;
+    const attemptSockets = new Set<Socket>();
     try {
-      nextServer = serverFactory((req, res) => {
-        void handle(req, res).catch(() => deny(res, 500, "gateway error"));
+      nextServer = serverFactory((socket) => {
+        attemptSockets.add(socket);
+        socket.once("close", () => attemptSockets.delete(socket));
+        demux.accept(socket);
       });
     } catch {
       scheduleRetry("bind_failed");
       return;
     }
+    socketsByServer.set(nextServer, attemptSockets);
     server = nextServer;
 
     nextServer.on("error", (error) => {
@@ -589,6 +1275,7 @@ export function createMcpGatewayListener(
       logger.log(`[mcp-gateway] listening on ${bind}:${port}`);
     });
     nextServer.on("close", () => {
+      socketsByServer.delete(nextServer);
       if (stopped || server !== nextServer) return;
       server = null;
       scheduleRetry("bind_failed");
@@ -603,6 +1290,52 @@ export function createMcpGatewayListener(
 
   const controller: McpGatewayListener = {
     state: () => capability,
+    waitUntilReady: (timeoutMs = 5_000) => {
+      if (capability.status === "ready") return Promise.resolve();
+      if (capability.status === "down") {
+        return Promise.reject(
+          new Error(`MCP gateway listener is down (${capability.reason})`),
+        );
+      }
+      if (stopped) {
+        return Promise.reject(new Error("MCP gateway listener is stopped"));
+      }
+      const boundedTimeoutMs =
+        Number.isFinite(timeoutMs) && timeoutMs >= 1
+          ? Math.floor(timeoutMs)
+          : 5_000;
+      return new Promise<void>((resolveReady, rejectReady) => {
+        let settled = false;
+        let unsubscribe = (): void => undefined;
+        const finish = (error?: Error): void => {
+          if (settled) return;
+          settled = true;
+          clearTimeout(timer);
+          unsubscribe();
+          if (error) rejectReady(error);
+          else resolveReady();
+        };
+        const timer = setTimeout(
+          () =>
+            finish(
+              new Error(
+                `MCP gateway listener did not become ready within ${boundedTimeoutMs}ms`,
+              ),
+            ),
+          boundedTimeoutMs,
+        );
+        timer.unref();
+        const observe = (state: McpGatewayCapability): void => {
+          if (state.status === "ready") finish();
+          else if (state.status === "down") {
+            finish(new Error(`MCP gateway listener is down (${state.reason})`));
+          }
+        };
+        unsubscribe = controller.subscribe(observe);
+        // Close the state-change race between the initial check and subscribe.
+        observe(capability);
+      });
+    },
     subscribe: (listener) => {
       subscribers.add(listener);
       return () => subscribers.delete(listener);
@@ -639,6 +1372,11 @@ export function createMcpGatewayListener(
         };
         activeServer.once("close", finish);
         try {
+          const sockets = socketsByServer.get(activeServer);
+          socketsByServer.delete(activeServer);
+          if (sockets) {
+            for (const socket of sockets) socket.destroy();
+          }
           activeServer.close(finish);
           activeServer.closeAllConnections?.();
           if (!activeServer.listening) queueMicrotask(finish);
@@ -646,6 +1384,13 @@ export function createMcpGatewayListener(
           finish();
         }
       }).finally(() => {
+        unsubscribeRevocations();
+        for (const socket of authenticatedConnections.keys()) socket.destroy();
+        authenticatedConnections.clear();
+        httpServer.closeAllConnections();
+        // Clears the hidden server's timeout-check interval. It has no own
+        // listening handle, so close is otherwise a no-op.
+        httpServer.close();
         subscribers.clear();
       });
       return stopPromise;
@@ -662,6 +1407,11 @@ export function startMcpGateway(): McpGatewayController {
   return gatewayController.start();
 }
 
+/** Task admission gate: no runtime starts until the shared listener is bound. */
+export function waitForMcpGatewayReady(timeoutMs = 5_000): Promise<void> {
+  return startMcpGateway().waitUntilReady(timeoutMs);
+}
+
 /** Stop and reset the process-wide listener (also the deterministic test seam). */
 export async function stopMcpGateway(): Promise<void> {
   const active = gatewayController;
@@ -1190,6 +1940,7 @@ export async function setupMcpTaskConfig(
       ? await environment.exec({
           argv: ["sh", "-lc", steps],
           env: RUNTIME_AUTHORITY_ENV,
+          projectEnvProfile: "maintenance",
           timeoutMs: MCP_CONFIG_HOST_TIMEOUT_MS,
           maxOutputBytes: 256 * 1024,
         })