Decide WebSocket upgrades in the session, complete them in the host (COL-84)

packages/control-plane/src/cloudflare/websocket-upgrade.ts

@@ -0,0 +1,38 @@
+import type { Logger } from "../logger";
+import type { SessionUpgradeAdmission } from "../session/connection-authenticator";
+
+/**
+ * Complete a WebSocket upgrade the way Workers do it: the session decides,
+ * then the server half of a `WebSocketPair` is attached to the runtime and
+ * the client half rides back on a 101 response. `log` is the session logger;
+ * the decision carries its own request-scoped one.
+ */
+export async function upgradeWebSocket(
+  upgrades: SessionUpgradeAdmission,
+  request: Request,
+  log: Logger
+): Promise<Response> {
+  const decision = await upgrades.authorize(request);
+  if (decision.kind === "reject") return decision.response;
+
+  const pair = new WebSocketPair();
+  const [client, server] = Object.values(pair);
+  try {
+    await decision.attach(server);
+  } catch (error) {
+    log.error("WebSocket upgrade failed", {
+      ws_type: decision.role,
+      error: error instanceof Error ? error : String(error),
+    });
+    // Attachment commits nothing before its one await, so the server half is
+    // either unaccepted or fully attached; closing covers both without
+    // leaving the runtime a socket the client never received.
+    try {
+      server.close(1011, "WebSocket upgrade failed");
+    } catch {
+      // Never accepted: there is nothing to close.
+    }
+    return new Response("WebSocket upgrade failed", { status: 500 });
+  }
+  return new Response(null, { status: 101, webSocket: client });
+}

packages/control-plane/src/session/components.ts

@@ -117,7 +117,10 @@ import { SessionHttpDispatcher } from "./http/dispatcher";
 import { SessionMessageRouter } from "./message-router";
 import { SessionDisconnectHandler } from "./disconnect-handler";
 import type { Clock, SandboxDisconnectMonitor, SessionBroadcaster, SocketRegistry } from "./ports";
-import { SessionConnectionAuthenticator } from "./connection-authenticator";
+import {
+  SessionConnectionAuthenticator,
+  type SessionUpgradeAdmission,
+} from "./connection-authenticator";
 import { SessionSnapshotReader } from "./snapshot-reader";
 import { SessionAccessReader } from "./sandbox-access-reader";
 import { createSessionScopedLogger } from "./session-logger";
@@ -147,6 +150,8 @@ const WS_AUTH_TIMEOUT_MS = 30000; // 30 seconds
 export interface SessionRuntime {
   readonly log: Logger;
   readonly server: SessionServer<WebSocket, ClientInfo>;
+  /** Admission of WebSocket upgrades; the host completes the handshake and attaches its socket. */
+  readonly upgrades: SessionUpgradeAdmission;
   readonly alarms: {
     /** Expire stale authorization leases and re-arm persisted deadlines after a cold start. */
     rehydrate(): void;
@@ -790,13 +795,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi
   };
 
   const server = new SessionServer<WebSocket, ClientInfo>({
-    http: new SessionHttpDispatcher({
-      log,
-      routes,
-      handleWebSocketUpgrade: (request, url, requestLog) =>
-        connectionAuthenticator.handleWebSocketUpgrade(request, url, requestLog),
-      clock,
-    }),
+    http: new SessionHttpDispatcher({ log, routes, clock }),
     messages: new SessionMessageRouter({
       log,
       sockets,
@@ -843,6 +842,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi
   return {
     log,
     server,
+    upgrades: connectionAuthenticator,
     alarms: {
       rehydrate: () =>
         backgroundTasks.submit(