perf(streams): connect WebSocket after first write (#4104)

## Summary & Motivation

Keeps WebSocket setup off the first stream group, then starts the background upgrade when the second HTTP request is dispatched. Later groups continue over HTTP without waiting until the socket is OPEN, when the serialized writer switches transports at a confirmed request boundary. One-group streams never create a socket.

An ambiguous HTTP outcome poisons the writer and retires any provisional socket before it can carry a frame. Initial upgrades have a dedicated 10-second background timeout; the existing 250ms bound remains scoped to reconnects after an established socket closes. WS write spans include chunk sequence and count for direct takeover analysis.

## Test Plan

Tests cover dispatch overlap, continued HTTP writes while connecting, one-group streams, close during background connection, ambiguous HTTP outcomes, timeout and late OPEN behavior, and trace propagation. The world-vercel suite and typecheck pass locally.

Author: alangenfeld

.changeset/quiet-streams-connect.md

@@ -0,0 +1,5 @@
+---
+'@workflow/world-vercel': patch
+---
+
+Start stream WebSocket connections after the first HTTP group without delaying later HTTP writes.

packages/world-vercel/src/http-core.ts

@@ -549,7 +549,13 @@ export interface InstrumentedFetchOptions extends HttpClientSpanOptions {
    * Non-2xx bodies consumed by `buildError` are still reported here.
    */
   deferTransportSuccessUntilBody?: boolean;
-  /** Error code used when the request itself fails before a response arrives. */
+  /**
+   * Called synchronously after the request promise is created, before awaiting
+   * its response. This observes local dispatch only; it does not imply that any
+   * bytes reached the origin. Must not throw.
+   */
+  onRequestDispatched?: () => void;
+  /** Error code used when the request itself fails before a response arrived. */
   transportErrorCode?: 'TRANSPORT' | 'STREAM_ERROR';
 }
 
@@ -585,6 +591,7 @@ export async function instrumentedFetch(
     durationAttribute,
     onTransportOutcome,
     deferTransportSuccessUntilBody = false,
+    onRequestDispatched,
     transportErrorCode = 'TRANSPORT',
   } = opts;
   const label = logLabel ?? url;
@@ -623,8 +630,8 @@ export async function instrumentedFetch(
         span?.setAttributes({
           ...WorkflowHttpTransport(nodeAgents ? 'node-http' : 'undici'),
         });
-        response = nodeAgents
-          ? await nodeHttpFetch(url, {
+        const request = nodeAgents
+          ? nodeHttpFetch(url, {
               method,
               headers,
               body,
@@ -637,14 +644,16 @@ export async function instrumentedFetch(
               headersTimeoutMs: NODE_HTTP_HEADERS_TIMEOUT_MS,
               bodyTimeoutMs: NODE_HTTP_BODY_TIMEOUT_MS,
             })
-          : await fetch(url, {
+          : fetch(url, {
               method,
               headers,
               body,
               signal,
               // eslint-disable-next-line @typescript-eslint/no-explicit-any -- undici dispatcher type doesn't match @types/node's RequestInit
               dispatcher,
             } as any);
+        onRequestDispatched?.();
+        response = await request;
       } catch (error) {
         const elapsed = Date.now() - start;
         // Report the raw error, before the timeout mapping below rewraps it: the

packages/world-vercel/src/streamer.ts

@@ -211,7 +211,8 @@ export async function writeStreamSessionOverHttp(
   name: string,
   chunks: (string | Uint8Array)[],
   config?: APIConfig,
-  attributes?: Attributes
+  attributes?: Attributes,
+  onRequestDispatched?: () => void
 ): Promise<void> {
   const httpConfig = await getHttpConfig(config);
   httpConfig.headers.set('X-Stream-Multi', 'true');
@@ -231,6 +232,7 @@ export async function writeStreamSessionOverHttp(
       logLabel: url.pathname,
       spanName: 'workflow.stream.write',
       durationAttribute: 'workflow.stream.write.chunk_rtt',
+      onRequestDispatched: offset === 0 ? onRequestDispatched : undefined,
       attributes: {
         ...streamSpanAttributes({
           runId,
@@ -282,13 +284,14 @@ export function createStreamer(config?: APIConfig): Streamer {
                 name,
                 writerId,
                 config,
-                (chunks, attributes) =>
+                (chunks, attributes, onRequestDispatched) =>
                   writeStreamSessionOverHttp(
                     runId,
                     name,
                     chunks,
                     config,
-                    attributes
+                    attributes,
+                    onRequestDispatched
                   ),
                 () => closeStreamSessionOverHttp(runId, name, config)
               );

packages/world-vercel/src/trace-propagation.test.ts

@@ -345,14 +345,16 @@ describe('ws stream transport upgrade trace propagation', () => {
     await tracer.startActiveSpan('flow-invocation', async (span) => {
       traceId = span.spanContext().traceId;
       spanId = span.spanContext().spanId;
-      createStreamWriteSession(
+      const session = createStreamWriteSession(
         'wrun_1',
         'user',
         'wrtr_01ARZ3NDEKTSV4RRFFQ69G5FAV',
         { token: 'test-token' },
-        async () => {},
+        async (_chunks, _attributes, dispatched) => dispatched?.(),
         async () => {}
       );
+      await session.write(0, ['first']);
+      await session.write(1, ['second']);
       await vi.waitFor(() => expect(wsUpgrades).toHaveLength(1));
       span.end();
     });

packages/world-vercel/src/ws-stream-connect.ts

@@ -1,11 +1,10 @@
 import type { WebSocket } from 'ws';
 import { trace } from './telemetry.js';
 
-/**
- * Initial implementation-level wait for an opted-in stream socket to open.
- * Tune from deployment measurements; this is not protocol semantics.
- */
-export const STREAM_WS_CONNECT_BUDGET_MS = 250;
+/** Background initial upgrade bound; writes continue over HTTP meanwhile. */
+export const STREAM_WS_INITIAL_CONNECT_TIMEOUT_MS = 10_000;
+/** Bounded wait while reconnecting after an established socket drains/closes. */
+export const STREAM_WS_RECONNECT_BUDGET_MS = 250;
 /** Cleanup-only wait after semantic close success; not protocol semantics. */
 export const STREAM_WS_CLOSE_BUDGET_MS = 250;
 

packages/world-vercel/src/ws-stream-session.test.ts

@@ -122,6 +122,7 @@ beforeEach(() => {
   injectTraceContextIntoHeaders.mockClear();
   writeSpans.length = 0;
   delete process.env.WORKFLOW_STREAMS_TRANSPORT;
+  delete process.env.WORKFLOW_REQUEST_TIMEOUT_MS;
 });
 
 afterEach(() => {
@@ -131,17 +132,27 @@ afterEach(() => {
 });
 
 function makeSession(
-  config: { token?: string } | undefined = { token: 'token' }
+  config: { token?: string } | undefined = { token: 'token' },
+  connectAfterFirstWrite = false
 ) {
-  const writeHttp = vi.fn().mockResolvedValue(undefined);
+  const writeHttp = vi.fn(
+    async (
+      _chunks: (string | Uint8Array)[],
+      _attributes?: Record<string, unknown>,
+      onRequestDispatched?: () => void
+    ) => {
+      onRequestDispatched?.();
+    }
+  );
   const closeHttp = vi.fn().mockResolvedValue(undefined);
   const session = createStreamWriteSession(
     'wrun_1',
     'stream/1',
     writerId,
     config,
     writeHttp,
-    closeHttp
+    closeHttp,
+    connectAfterFirstWrite
   );
   activeSessions.push(session);
   return { session, writeHttp, closeHttp };
@@ -158,6 +169,124 @@ describe('v1 stream WebSocket writer lifecycle', () => {
     expect(closeHttp).toHaveBeenCalledTimes(1);
   });
 
+  it('starts WS during the second HTTP group and keeps writing HTTP until OPEN', async () => {
+    process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
+    let releaseSecond: (() => void) | undefined;
+    const secondPending = new Promise<void>((resolve) => {
+      releaseSecond = resolve;
+    });
+    const { session, writeHttp } = makeSession({ token: 'token' }, true);
+
+    await session.write(0, ['one']);
+    expect(sockets).toHaveLength(0);
+    expect(writeHttp).toHaveBeenNthCalledWith(
+      1,
+      ['one'],
+      expect.objectContaining({
+        'workflow.stream.ws.session_first_write': true,
+        'workflow.stream.ws.http_group_ordinal': 1,
+      })
+    );
+
+    writeHttp.mockImplementationOnce(
+      async (_chunks, _attributes, dispatched) => {
+        dispatched?.();
+        await secondPending;
+      }
+    );
+    const second = session.write(1, ['two']);
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+    expect(writeHttp).toHaveBeenNthCalledWith(
+      2,
+      ['two'],
+      expect.objectContaining({
+        'workflow.stream.ws.connect_after_http_group': 2,
+      }),
+      expect.any(Function)
+    );
+    expect(sockets[0].sent).toHaveLength(0);
+    releaseSecond?.();
+    await second;
+
+    const third = session.write(2, ['three']);
+    await third;
+    expect(writeHttp.mock.calls[2]?.[0]).toEqual(['three']);
+    sockets[0].open();
+
+    const fourth = session.write(3, ['four']);
+    await vi.waitFor(() => expect(sockets[0].sent).toHaveLength(1));
+    expect((await decodeOne(sockets[0].sent[0])).meta).toMatchObject({
+      type: 'write',
+      chunkSeq: 3,
+      numChunks: 1,
+    });
+    sockets[0].reply(
+      encodeFrame({ type: 'write_ack', reqId: 1 }, new Uint8Array())
+    );
+    await fourth;
+  });
+
+  it('closes over HTTP without waiting for the background socket', async () => {
+    process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
+    const { session, closeHttp } = makeSession({ token: 'token' }, true);
+    await session.write(0, ['one']);
+    await session.write(1, ['two']);
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+
+    await session.close();
+    expect(closeHttp).toHaveBeenCalledTimes(1);
+    expect(sockets[0].closed).toContainEqual([1000, 'stream closed over HTTP']);
+  });
+
+  it('retires a provisional socket after an ambiguous second HTTP outcome', async () => {
+    process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
+    const error = new Error('second HTTP outcome unknown');
+    const { session, writeHttp } = makeSession({ token: 'token' }, true);
+    await session.write(0, ['one']);
+    let rejectSecond: ((error: Error) => void) | undefined;
+    const secondPending = new Promise<void>((_resolve, reject) => {
+      rejectSecond = reject;
+    });
+    writeHttp.mockImplementationOnce(
+      async (_chunks, _attributes, dispatched) => {
+        dispatched?.();
+        await secondPending;
+      }
+    );
+
+    const second = session.write(1, ['two']);
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+    rejectSecond?.(error);
+    await expect(second).rejects.toBe(error);
+    await expect(session.write(2, ['three'])).rejects.toBe(error);
+    expect(sockets[0].sent).toHaveLength(0);
+    expect(sockets[0].closed).toContainEqual([
+      1011,
+      'unknown stream write outcome',
+    ]);
+  });
+
+  it('does not start a socket after an ambiguous first HTTP outcome', async () => {
+    process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
+    const error = new Error('HTTP outcome unknown');
+    const { session, writeHttp } = makeSession({ token: 'token' }, true);
+    writeHttp.mockRejectedValueOnce(error);
+
+    await expect(session.write(0, ['one'])).rejects.toBe(error);
+    await expect(session.write(1, ['two'])).rejects.toBe(error);
+    expect(sockets).toHaveLength(0);
+  });
+
+  it('closes a one-group stream without starting a socket', async () => {
+    process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
+    const { session, closeHttp } = makeSession({ token: 'token' }, true);
+
+    await session.write(0, ['only']);
+    await session.close();
+    expect(closeHttp).toHaveBeenCalledTimes(1);
+    expect(sockets).toHaveLength(0);
+  });
+
   it('sends immediately over HTTP while the initial socket connects, then switches to WS', async () => {
     process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
     let releaseHttp: (() => void) | undefined;
@@ -239,27 +368,33 @@ describe('v1 stream WebSocket writer lifecycle', () => {
     });
   });
 
-  it('tombstones to HTTP when the background connect budget expires', async () => {
+  it('bounds the background connect attempt without delaying HTTP writes', async () => {
     vi.useFakeTimers();
     process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
-    const { session, writeHttp } = makeSession();
+    process.env.WORKFLOW_REQUEST_TIMEOUT_MS = '10000';
+    const { session, writeHttp } = makeSession({ token: 'token' }, true);
     const first = session.write(0, ['one']);
-    await vi.advanceTimersByTimeAsync(0);
-    expect(writeHttp).toHaveBeenCalledWith(
-      ['one'],
-      expect.objectContaining({
-        'workflow.stream.ws.session_first_write': true,
-        'workflow.stream.ws.connecting_at_write': true,
-      })
+    await vi.waitFor(() =>
+      expect(writeHttp).toHaveBeenCalledWith(
+        ['one'],
+        expect.objectContaining({
+          'workflow.stream.ws.session_first_write': true,
+          'workflow.stream.ws.connecting_at_write': false,
+          'workflow.stream.ws.connect_deferred_at_write': true,
+        })
+      )
     );
     await first;
+    expect(sockets).toHaveLength(0);
+    await session.write(1, ['two']);
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
 
-    await vi.advanceTimersByTimeAsync(250);
-    const second = session.write(1, ['two']);
-    await second;
-    expect(writeHttp).toHaveBeenCalledTimes(2);
+    await vi.advanceTimersByTimeAsync(10_000);
+    await session.write(2, ['three']);
+    expect(writeHttp).toHaveBeenCalledTimes(3);
     expect(writeHttp.mock.calls[0]?.[0]).toEqual(['one']);
-    expect(writeHttp.mock.calls[1]).toEqual([['two']]);
+    expect(writeHttp.mock.calls[1]?.[0]).toEqual(['two']);
+    expect(writeHttp.mock.calls[2]).toEqual([['three']]);
     expect(sockets[0].sent).toHaveLength(0);
     expect(sockets[0].closed).toContainEqual([1000, 'connect budget expired']);
     sockets[0].open();
@@ -766,15 +901,16 @@ describe('v1 stream WebSocket writer lifecycle', () => {
     expect(sockets).toHaveLength(0);
   });
 
-  it('forwards streamer-wrapper disposal after session materialization', async () => {
+  it('keeps streamer-wrapper disposal socket-free before its first write', async () => {
     process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws';
     const { createStreamer } = await import('./streamer.js');
     const session = createStreamer({
       token: 'token',
     }).streams.createWriteSession?.('wrun_1', 'stream/1', { writerId });
-    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(sockets).toHaveLength(0);
     await session?.dispose?.();
-    expect(sockets[0].closed).toContainEqual([1000, 'stream writer disposed']);
+    expect(sockets).toHaveLength(0);
   });
 
   it('disposes transport without sending protocol close', async () => {