Take the ws events transport only once its handshake is done

.changeset/ws-register-after-open.md

@@ -0,0 +1,7 @@
+---
+'@workflow/world-vercel': patch
+---
+
+Only take the WebSocket events transport once its handshake has completed. A
+write issued mid-connect now goes over HTTP instead of waiting for the socket,
+which measured faster for the first write of a run.

packages/world-vercel/src/ws-protocol-conformance.test.ts

@@ -277,13 +277,17 @@ async function callThroughFixture(
   // Stand in for the flow route: without an open channel the write resolves
   // none and goes over HTTP, which is the whole point of the explicit pair.
   openWsChannel(input.runId, { token: 'test-token' });
-  const pending = createWorkflowRunEventV4(input, { token: 'test-token' });
-  void pending.catch(() => {});
-  // Let the transport construct its socket (it awaits the header thunk).
+  // Let the transport construct its socket (it awaits the header thunk), and
+  // finish the handshake *before* issuing the write. The write path only takes
+  // the socket once it is genuinely open — a frame issued mid-connect resolves
+  // no transport and goes over HTTP, which would leave this fixture asserting
+  // on a socket nothing was ever sent through.
   await vi.waitFor(() => expect(sockets.length).toBeGreaterThan(0));
   const socket = sockets[0];
   const seen = attachFixtureServer(socket, handler);
   socket.open();
+  const pending = createWorkflowRunEventV4(input, { token: 'test-token' });
+  void pending.catch(() => {});
   return { pending, socket, seen };
 }
 

packages/world-vercel/src/ws-transport.test.ts

@@ -1006,10 +1006,13 @@ describe('transport selection', () => {
       expect(sockets).toHaveLength(0);
     });
 
-    it('scopes the channel to one run and memoizes it per URL', () => {
+    it('scopes the channel to one run and memoizes it per URL', async () => {
       process.env.WORKFLOW_EVENTS_TRANSPORT = 'ws';
       openWsChannel('wrun_1', directConfig);
       openWsChannel('wrun_2', directConfig);
+      // Resolution now requires an open socket, not merely a claimed channel.
+      await tick();
+      for (const socket of sockets) socket.open();
 
       const first = resolveWsTransport('wrun_1', directConfig);
       const second = resolveWsTransport('wrun_1', directConfig);
@@ -1021,11 +1024,33 @@ describe('transport selection', () => {
       );
     });
 
-    it('returns null once the channel is closed', () => {
+    it('withholds the channel until the handshake completes', async () => {
+      // The contract this change exists for. Measured on a WS-enabled
+      // deployment, `run_started` — the write that routinely lands mid-connect
+      // — cost p50 269ms / p95 3.77s when it waited for the socket, against
+      // p50 79ms / p95 134ms when it fell back to HTTP, while every write
+      // issued after the socket was up cost ~65ms. Waiting made the first
+      // write of a run slower than not using the socket at all.
+      process.env.WORKFLOW_EVENTS_TRANSPORT = 'ws';
+      openWsChannel('wrun_1', directConfig);
+      await tick();
+
+      // The socket exists and is mid-handshake: claimed, but not writable.
+      expect(sockets).toHaveLength(1);
+      expect(sockets[0].readyState).toBe(0);
+      expect(resolveWsTransport('wrun_1', directConfig)).toBeNull();
+
+      sockets[0].open();
+      expect(resolveWsTransport('wrun_1', directConfig)).not.toBeNull();
+    });
+
+    it('returns null once the channel is closed', async () => {
       // What makes a late write — one the runtime issues after the invocation
       // that opened the channel has returned — fall back instead of failing.
       process.env.WORKFLOW_EVENTS_TRANSPORT = 'ws';
       const release = openWsChannel('wrun_1', directConfig);
+      await tick();
+      sockets[0]?.open();
       expect(resolveWsTransport('wrun_1', directConfig)).not.toBeNull();
 
       release?.();

packages/world-vercel/src/ws-transport.ts

@@ -301,6 +301,25 @@ class WsEventsTransport {
     conn?.ws.close(1000, reason);
   }
 
+  /**
+   * Whether a write can go out *now*, without waiting on a handshake.
+   *
+   * This is deliberately narrower than "does this run have a channel". The
+   * write path uses it to decide between the socket and HTTP, and the honest
+   * answer during a connect is HTTP: measured on a WS-enabled deployment,
+   * `run_started` — the one write that routinely lands mid-handshake — cost
+   * p50 269ms / p95 3.77s over the socket against p50 79ms / p95 134ms when it
+   * fell back, while every write issued after the socket was up cost ~65ms.
+   * Waiting for the handshake made the first write of a run slower than not
+   * using the socket at all.
+   *
+   * A reconnect reads as not-ready for the same reason: the frames that would
+   * queue behind it are better served by the transport that needs no setup.
+   */
+  get isReadyForWrites(): boolean {
+    return this.connection?.ws.readyState === WebSocket.OPEN;
+  }
+
   private ensureConnected(): Promise<Connection> {
     const conn = this.connection;
     if (conn && conn.ws.readyState === WebSocket.OPEN) {
@@ -919,5 +938,9 @@ export function resolveWsTransport(
   const wsUrl = resolveChannelUrl(runId, config);
   if (!wsUrl) return null;
   const transport = wsState.transports.get(wsUrl);
-  return transport ? { transport, wsUrl } : null;
+  // Membership answers "does this run have a channel"; `isReadyForWrites` adds
+  // "and can it take a frame without a handshake first". Both have to hold:
+  // see the getter for why a mid-connect write is better off on HTTP.
+  if (!transport || !transport.isReadyForWrites) return null;
+  return { transport, wsUrl };
 }