[backport] [world-vercel] Use stream control frame for transparent reconnection (#1766)

Author: VaguelySerious

.changeset/stream-reconnect-stable.md

@@ -0,0 +1,5 @@
+---
+'@workflow/world-vercel': patch
+---
+
+Use stream control frame to transparently reconnect on server timeout

packages/world-vercel/src/streamer.test.ts

@@ -1,5 +1,10 @@
 import { afterEach, describe, expect, it, vi } from 'vitest';
-import { encodeMultiChunks, MAX_CHUNKS_PER_REQUEST } from './streamer.js';
+import {
+  encodeMultiChunks,
+  MAX_CHUNKS_PER_REQUEST,
+  parseStreamControlFrame,
+  STREAM_CONTROL_FRAME_SIZE,
+} from './streamer.js';
 
 describe('encodeMultiChunks', () => {
   /**
@@ -269,3 +274,198 @@ describe('writeToStreamMulti pagination', () => {
     ]);
   });
 });
+
+/**
+ * Build a control frame matching the workflow-server format.
+ */
+function buildControlFrame(done: boolean, nextIndex: number): Uint8Array {
+  const frame = new Uint8Array(STREAM_CONTROL_FRAME_SIZE);
+  // Bytes 0-3: zero-frame marker (already 0x00)
+  frame[4] = done ? 1 : 0;
+  new DataView(frame.buffer).setUint32(5, nextIndex, false);
+  // Magic footer "WFCT"
+  frame.set(new Uint8Array([0x57, 0x46, 0x43, 0x54]), 9);
+  return frame;
+}
+
+describe('parseStreamControlFrame', () => {
+  it('parses a valid done=true control frame', () => {
+    const frame = buildControlFrame(true, 42);
+    const result = parseStreamControlFrame(frame);
+    expect(result).toEqual({
+      done: true,
+      nextIndex: 42,
+      totalLength: STREAM_CONTROL_FRAME_SIZE,
+    });
+  });
+
+  it('parses a valid done=false (timeout) control frame', () => {
+    const frame = buildControlFrame(false, 100);
+    const result = parseStreamControlFrame(frame);
+    expect(result).toEqual({
+      done: false,
+      nextIndex: 100,
+      totalLength: STREAM_CONTROL_FRAME_SIZE,
+    });
+  });
+
+  it('parses control frame appended after data bytes', () => {
+    const data = new Uint8Array([1, 2, 3, 4, 5]);
+    const frame = buildControlFrame(false, 7);
+    const combined = new Uint8Array(data.length + frame.length);
+    combined.set(data, 0);
+    combined.set(frame, data.length);
+
+    const result = parseStreamControlFrame(combined);
+    expect(result).toEqual({
+      done: false,
+      nextIndex: 7,
+      totalLength: STREAM_CONTROL_FRAME_SIZE,
+    });
+  });
+
+  it('returns null for buffer shorter than control frame size', () => {
+    expect(parseStreamControlFrame(new Uint8Array(12))).toBeNull();
+    expect(parseStreamControlFrame(new Uint8Array(0))).toBeNull();
+  });
+
+  it('returns null when magic footer does not match', () => {
+    const frame = buildControlFrame(true, 0);
+    frame[12] = 0xff; // corrupt magic footer
+    expect(parseStreamControlFrame(frame)).toBeNull();
+  });
+
+  it('returns null when zero-frame marker is not all zeros', () => {
+    const frame = buildControlFrame(true, 0);
+    frame[0] = 1; // corrupt zero-frame marker
+    expect(parseStreamControlFrame(frame)).toBeNull();
+  });
+
+  it('handles nextIndex=0', () => {
+    const frame = buildControlFrame(false, 0);
+    const result = parseStreamControlFrame(frame);
+    expect(result).toEqual({
+      done: false,
+      nextIndex: 0,
+      totalLength: STREAM_CONTROL_FRAME_SIZE,
+    });
+  });
+
+  it('handles large nextIndex values', () => {
+    const frame = buildControlFrame(true, 0xffffffff);
+    const result = parseStreamControlFrame(frame);
+    expect(result?.nextIndex).toBe(0xffffffff);
+  });
+});
+
+describe('readFromStream reconnection', () => {
+  /** Collect every byte from a ReadableStream into one Uint8Array. */
+  async function drain(
+    stream: ReadableStream<Uint8Array>
+  ): Promise<Uint8Array> {
+    const reader = stream.getReader();
+    const parts: Uint8Array[] = [];
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      parts.push(value);
+    }
+    const len = parts.reduce((s, p) => s + p.length, 0);
+    const out = new Uint8Array(len);
+    let off = 0;
+    for (const p of parts) {
+      out.set(p, off);
+      off += p.length;
+    }
+    return out;
+  }
+
+  function chunkedStream(chunks: Uint8Array[]): ReadableStream<Uint8Array> {
+    let i = 0;
+    return new ReadableStream({
+      pull(controller) {
+        if (i < chunks.length) {
+          controller.enqueue(chunks[i++]);
+        } else {
+          controller.close();
+        }
+      },
+    });
+  }
+
+  function streamResponse(...chunks: Uint8Array[]): Response {
+    return new Response(chunkedStream(chunks), {
+      status: 200,
+      headers: { 'Content-Type': 'application/octet-stream' },
+    });
+  }
+
+  async function getStreamer() {
+    const { createStreamer } = await import('./streamer.js');
+    return createStreamer();
+  }
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('reconnects when server sends done=false and resumes from nextIndex', async () => {
+    const chunk1 = new TextEncoder().encode('aaa');
+    const chunk2 = new TextEncoder().encode('bbb');
+    const timeout = buildControlFrame(false, 3);
+    const done = buildControlFrame(true, 6);
+
+    const fetchSpy = vi
+      .spyOn(globalThis, 'fetch')
+      .mockResolvedValueOnce(streamResponse(chunk1, timeout))
+      .mockResolvedValueOnce(streamResponse(chunk2, done));
+
+    const streamer = await getStreamer();
+    const result = await drain(await streamer.readFromStream('strm_test'));
+
+    const expected = new Uint8Array([...chunk1, ...chunk2]);
+    expect(result).toEqual(expected);
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+
+    const secondUrl = new URL(fetchSpy.mock.calls[1][0] as string);
+    expect(secondUrl.searchParams.get('startIndex')).toBe('3');
+  });
+
+  it('falls through when no control frame is present (backward compat)', async () => {
+    const data = new TextEncoder().encode('legacy server');
+
+    vi.spyOn(globalThis, 'fetch').mockResolvedValueOnce(streamResponse(data));
+
+    const streamer = await getStreamer();
+    const result = await drain(await streamer.readFromStream('strm_test'));
+
+    expect(result).toEqual(data);
+  });
+
+  it('propagates network error to consumer without retrying', async () => {
+    const data = new TextEncoder().encode('partial');
+
+    let callCount = 0;
+    const errorStream = new ReadableStream<Uint8Array>({
+      pull(controller) {
+        if (callCount === 0) {
+          callCount++;
+          controller.enqueue(data);
+        } else {
+          controller.error(new Error('connection reset'));
+        }
+      },
+    });
+
+    vi.spyOn(globalThis, 'fetch').mockResolvedValueOnce(
+      new Response(errorStream, { status: 200 })
+    );
+
+    const streamer = await getStreamer();
+    // readFromStream reads the full response via arrayBuffer(), so
+    // a mid-stream error rejects the readFromStream promise itself.
+    await expect(streamer.readFromStream('strm_test')).rejects.toThrow(
+      'connection reset'
+    );
+  });
+});

packages/world-vercel/src/streamer.ts

@@ -23,6 +23,66 @@ export const MAX_CHUNKS_PER_REQUEST = 1000;
 // (partial writes, long-lived reads), and duplex streams are incompatible
 // with undici's experimental H2 support.
 
+/**
+ * Stream control frame constants, mirroring workflow-server's format.
+ *
+ * Control frame (13 bytes):
+ *   [0-3]  Zero-frame marker (0x00 0x00 0x00 0x00)
+ *   [4]    Flags — bit 0: done (1 = complete, 0 = timeout/reconnect)
+ *   [5-8]  nextIndex — big-endian uint32, chunk index to resume from
+ *   [9-12] Magic footer — "WFCT" (0x57 0x46 0x43 0x54)
+ */
+export const STREAM_CONTROL_FRAME_SIZE = 13;
+const STREAM_CONTROL_MAGIC = new Uint8Array([0x57, 0x46, 0x43, 0x54]);
+
+export interface StreamControlFrame {
+  done: boolean;
+  nextIndex: number;
+}
+
+/**
+ * Try to parse a stream control frame from the tail of a buffer.
+ * Returns the parsed frame and the byte length of the control data,
+ * or null if no valid control frame is present.
+ */
+export function parseStreamControlFrame(
+  buffer: Uint8Array
+): (StreamControlFrame & { totalLength: number }) | null {
+  if (buffer.length < STREAM_CONTROL_FRAME_SIZE) return null;
+
+  const offset = buffer.length - STREAM_CONTROL_FRAME_SIZE;
+
+  // Check zero-frame marker (bytes 0-3 must be 0x00)
+  if (
+    buffer[offset] !== 0 ||
+    buffer[offset + 1] !== 0 ||
+    buffer[offset + 2] !== 0 ||
+    buffer[offset + 3] !== 0
+  ) {
+    return null;
+  }
+
+  // Check magic footer at bytes 9-12
+  if (
+    buffer[offset + 9] !== STREAM_CONTROL_MAGIC[0] ||
+    buffer[offset + 10] !== STREAM_CONTROL_MAGIC[1] ||
+    buffer[offset + 11] !== STREAM_CONTROL_MAGIC[2] ||
+    buffer[offset + 12] !== STREAM_CONTROL_MAGIC[3]
+  ) {
+    return null;
+  }
+
+  const flags = buffer[offset + 4];
+  const view = new DataView(buffer.buffer, buffer.byteOffset + offset + 5, 4);
+  const nextIndex = view.getUint32(0, false);
+
+  return {
+    done: (flags & 1) === 1,
+    nextIndex,
+    totalLength: STREAM_CONTROL_FRAME_SIZE,
+  };
+}
+
 function getStreamUrl(
   name: string,
   runId: string | undefined,
@@ -188,21 +248,90 @@ export function createStreamer(config?: APIConfig): Streamer {
     },
 
     async readFromStream(name: string, startIndex?: number) {
-      const httpConfig = await getHttpConfig(config);
-      const url = getStreamUrl(name, undefined, httpConfig);
-      if (typeof startIndex === 'number') {
-        url.searchParams.set('startIndex', String(startIndex));
+      let currentStartIndex = startIndex ?? 0;
+
+      // Cap reconnections to prevent infinite loops if the server
+      // never completes the stream. 50 reconnects at 2-min server
+      // timeout ≈ 100 minutes of streaming.
+      const MAX_RECONNECTS = 50;
+      let reconnectCount = 0;
+
+      const connect = async (): Promise<Response> => {
+        const httpConfig = await getHttpConfig(config);
+        // Use v3 to receive the stream control frame for reconnection.
+        const url = new URL(
+          `${httpConfig.baseUrl}/v3/stream/${encodeURIComponent(name)}`
+        );
+        url.searchParams.set('startIndex', String(currentStartIndex));
+        const response = await fetch(url, {
+          headers: httpConfig.headers,
+        });
+        if (!response.ok) {
+          throw new Error(`Failed to fetch stream: ${response.status}`);
+        }
+        if (!response.body) {
+          throw new Error('No response body for stream');
+        }
+        return response;
+      };
+
+      // Read the entire response as bytes, then strip any trailing
+      // control frame. This avoids the complexity of hold-back buffering
+      // inside a ReadableStream pull loop, which interacts poorly with
+      // the proxy and byte-stream wrappers in production.
+      const readFull = async (): Promise<{
+        data: Uint8Array;
+        control: (StreamControlFrame & { totalLength: number }) | null;
+      }> => {
+        const response = await connect();
+        const buffer = new Uint8Array(await response.arrayBuffer());
+        const control = parseStreamControlFrame(buffer);
+        if (control) {
+          const dataLen = buffer.length - control.totalLength;
+          return {
+            data: dataLen > 0 ? buffer.subarray(0, dataLen) : new Uint8Array(0),
+            control,
+          };
+        }
+        return { data: buffer, control: null };
+      };
+
+      // Collect all data, transparently reconnecting on server timeout.
+      const parts: Uint8Array[] = [];
+      for (;;) {
+        const { data, control } = await readFull();
+        if (data.length > 0) {
+          parts.push(data);
+        }
+
+        if (!control || control.done) {
+          // Stream complete or no control frame (older server).
+          break;
+        }
+
+        // Timeout — reconnect from the next chunk index.
+        reconnectCount++;
+        if (reconnectCount > MAX_RECONNECTS) {
+          throw new Error(
+            `Stream exceeded maximum reconnection attempts (${MAX_RECONNECTS})`
+          );
+        }
+        currentStartIndex = control.nextIndex;
       }
-      const response = await fetch(url, {
-        headers: httpConfig.headers,
+
+      // Return a ReadableStream that emits the collected data.
+      let emitted = false;
+      return new ReadableStream<Uint8Array>({
+        pull(controller) {
+          if (!emitted) {
+            emitted = true;
+            for (const part of parts) {
+              controller.enqueue(part);
+            }
+          }
+          controller.close();
+        },
       });
-      if (!response.ok) {
-        throw new Error(`Failed to fetch stream: ${response.status}`);
-      }
-      if (!response.body) {
-        throw new Error('No response body for stream');
-      }
-      return response.body as ReadableStream<Uint8Array>;
     },
 
     async getStreamChunks(