diff --git a/.changeset/ws-register-after-open.md b/.changeset/ws-register-after-open.md new file mode 100644 index 0000000000..c5bdd78acd --- /dev/null +++ b/.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. diff --git a/packages/world-vercel/src/ws-protocol-conformance.test.ts b/packages/world-vercel/src/ws-protocol-conformance.test.ts index deffac7c67..bab9d63db6 100644 --- a/packages/world-vercel/src/ws-protocol-conformance.test.ts +++ b/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 }; } diff --git a/packages/world-vercel/src/ws-transport.test.ts b/packages/world-vercel/src/ws-transport.test.ts index f7b63ea7ad..c9cdb3df0e 100644 --- a/packages/world-vercel/src/ws-transport.test.ts +++ b/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?.(); diff --git a/packages/world-vercel/src/ws-transport.ts b/packages/world-vercel/src/ws-transport.ts index 2b0e7397d3..4245c6ca00 100644 --- a/packages/world-vercel/src/ws-transport.ts +++ b/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 { 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 }; }