WebSocket-dispatched Remote invocations carry no request context

[yuzisama224]

The WebSocket upgrade is authenticated but carries nothing about the request forward, so every streamed Remote invocation is dispatched with no request in scope — even where the unary /api path has one.

Where the request is lost

The upgrade route checks trust and then hands the socket to the mux (packages/api/gateway/src/index.ts:205-229):

handler: (req, socket, head) => {
  const rejection = webCtx.connection.requestRejection(req)
  if (rejection !== undefined) {
    rejectRemoteStreamUpgrade(socket, rejection)
    return
  }
  mux.handleUpgrade(req, socket, head)
},

req is an ordinary HTTP request that has passed through whatever sits in front of the process. It is used for the trust check and then dropped. Everything that streams — session/follow, session/control, workspace/follow — is dispatched from frames on that socket, and none of it can see it.

This is the same gap as the companion post about client-connection (#5790), one transport over. That one covers unary /api calls; this one covers everything that streams, and the two together are what an endpoint owner actually needs, because a session renders over the stream and is cancelled over /api.

All references are to dsh-v0.1.2-rc.1 (a66e470).

The part that is not obvious

Entering a scope around handleUpgrade, or around the accept callback, does not work. We measured it: the request is readable inside the accept callback and undefined in every message listener, because ws delivers socket events on the socket's own async context rather than the upgrade's. By the time the first frame arrives the scope has unwound.

Scoping only the mux's open closure is not enough either. That covers prepareInvocation, and therefore the lookup resolvers, but a Remote method declared as a generator has not run its body when open returns, so anything the method itself reads falls outside.

The scope has to be re-entered per socket event. The smallest place that does that is the accepted socket's own emit, wrapped once where the per-socket connection is built (packages/api/gateway/src/stream-server.ts:48-57) — the first point at which anything is per-socket at all.

The change we would propose

RemoteStreamMuxServer takes one more constructor argument, and it is optional:

export type RemoteStreamDispatchScope = (
  req: IncomingMessage,
  dispatch: () => boolean,
) => boolean

handleUpgrade wraps the accepted socket's emit with it only when one was supplied; a three-argument caller gets a socket whose emit is the one ws gave it, which is exactly today's behaviour. index.ts supplies (req, dispatch) => webCtx.connection.runWithRequest(req, dispatch).

Nothing new is injected for that: connection is already injected on that exact context (packages/api/gateway/src/index.ts:205). What the change does need is that runWithRequest be declared on HostConnectionHandle, the interface that types ctx.connection — the companion change declares it there. This is why this change does not compile on its own even though it touches only gateway files. There is no second storage, and one request context serves both transports.

The diff is two files and about 25 lines. It is attached or linked below.

Wrapping emit is the blunt part of it, and we would rather be told a better one. ws gives no per-socket async context and no dispatch hook, so the alternatives we found were worse: patching on/once per event name misses listeners registered later, and moving the scope into RemoteStreamMuxConnection would have to re-enter it at each of several read sites instead of one.

What we are asking

1. Is per-request context on the WebSocket path something you would want, given the companion post's scope exists?
2. Is wrapping emit acceptable, or is there a hook we missed?
3. Should the argument be optional, as proposed, or should the mux always take a scope and every caller pass one?

We run this as a local patch over the published 0.1.2-rc.1 package today, verified two ways: interleaved frames from two live sockets, each dispatched as the request that opened its own socket and never as the other, with an unstamped socket yielding no request at all rather than its neighbour's; and the same claim through a real browser. Those tests live in our repository, not here — upstream homes would be packages/api/gateway/tests/stream-server.host.spec.ts and gateway-stream.host.spec.ts, and we would write them if there is appetite.

diff --git a/packages/api/gateway/src/index.ts b/packages/api/gateway/src/index.ts
index 4ec7c3f9fd..335ff3ff7b 100644
--- a/packages/api/gateway/src/index.ts
+++ b/packages/api/gateway/src/index.ts
@@ -207,6 +207,7 @@ export class TypertGatewayService extends Service implements TypertGateway {
         (endpoint, payload, signal) => this.openWireStream(endpoint, payload, signal),
         this.wireStream.failure,
         resolved.websocketHeartbeatIntervalMs,
+        (req, dispatch) => webCtx.connection.runWithRequest(req, dispatch),
       )
       webCtx.effect(() => {
         const route: WebUpgradeRoute = {
diff --git a/packages/api/gateway/src/stream-server.ts b/packages/api/gateway/src/stream-server.ts
index 3a36e97ecd..1ec92ba7b0 100644
--- a/packages/api/gateway/src/stream-server.ts
+++ b/packages/api/gateway/src/stream-server.ts
@@ -19,6 +19,17 @@ export type RemoteStreamOpener = (
 /** Convert an invocation or carrier failure to a stable wire value. */
 export type RemoteStreamFailureMapper = (error: unknown) => RemoteStreamFailure
 
+/**
+ * Re-enter the upgrade request's dispatch scope for one socket event.
+ * @param req - upgrade request that opened the socket delivering this event.
+ * @param dispatch - the socket's own event dispatch.
+ * @returns whatever `dispatch` returned.
+ */
+export type RemoteStreamDispatchScope = (
+  req: IncomingMessage,
+  dispatch: () => boolean,
+) => boolean
+
 const MAX_MISSED_HEARTBEATS = 2
 
 /** Own the no-server WebSocket acceptor and every active logical stream. */
@@ -32,11 +43,14 @@ export class RemoteStreamMuxServer {
    * @param open - Gateway stream dispatcher.
    * @param failure - Gateway error-to-wire mapper.
    * @param heartbeatIntervalMs - interval between WebSocket Ping control frames.
+   * @param scope - re-enter the upgrade request's dispatch scope per socket
+   * event; omitted, the accepted socket's own dispatch is left untouched.
    */
   constructor(
     private readonly open: RemoteStreamOpener,
     private readonly failure: RemoteStreamFailureMapper,
     private readonly heartbeatIntervalMs: number,
+    private readonly scope?: RemoteStreamDispatchScope,
   ) {}
 
   /**
@@ -50,6 +64,16 @@ export class RemoteStreamMuxServer {
       this.missedHeartbeats.set(websocket, 0)
       websocket.on('pong', () => { this.missedHeartbeats.set(websocket, 0) })
       this.startHeartbeat()
+      const scope = this.scope
+      if (scope !== undefined) {
+        // `ws` delivers every event on the socket's own async context rather
+        // than the upgrade's, so a scope entered around this callback has
+        // already unwound when the first frame arrives. Re-enter it per event
+        // instead, which also covers a Remote generator whose body runs after
+        // `open` returns.
+        const emit = websocket.emit.bind(websocket)
+        websocket.emit = (event, ...args) => scope(req, () => emit(event, ...args))
+      }
       const connection = new RemoteStreamMuxConnection(websocket, this.open, this.failure)
       const done = connection.run()
       this.connections.add(done)