@aiworker/web 1.22.2 → 1.23.1

package/server/lib/relay/downstream-scheduler.js

@@ -0,0 +1,154 @@
+// Downstream credit belongs to the end client, not the relay's TCP connection.
+// Producers await send(), so HTTP readers stop before building a plaintext queue.
+import { decodeTunnelFrame } from './tunnel-codec.js';
+import { MAX_PLAINTEXT_FRAME_BYTES } from './e2ee.js';
+
+export const DOWNSTREAM_CHUNK_BYTES = 8 * 1024;
+const MIN_WINDOW_BYTES = 64 * 1024;
+const MAX_WINDOW_BYTES = 1024 * 1024;
+const MAX_QUEUED_BYTES = 4 * 1024 * 1024;
+const MAX_QUEUED_FRAMES = 4096;
+const MAX_OUTSTANDING_FRAMES = 4096;
+const QUEUE_DELAY_TARGET_MS = 100;
+
+export const createDownstreamScheduler = ({ sendBatch, onError, maxBatchFrames = 32, now = () => performance.now() }) => {
+  const queues = new Map();
+  const outstanding = [];
+  let queuedBytes = 0;
+  let queuedFrames = 0;
+  let sentBytes = 0;
+  let acknowledgedBytes = 0;
+  let windowBytes = MIN_WINDOW_BYTES;
+  let minRtt = Infinity;
+  let adjustmentBoundary = MIN_WINDOW_BYTES;
+  let windowLimited = false;
+  let draining = false;
+  let closed = false;
+  let sending = [];
+
+  const close = () => {
+    closed = true;
+    for (const queue of queues.values()) for (const entry of queue) entry.resolve();
+    for (const entry of sending) entry.resolve();
+    sending = [];
+    queues.clear();
+    outstanding.length = 0;
+    queuedBytes = 0;
+    queuedFrames = 0;
+  };
+
+  const fail = error => {
+    close();
+    onError(error);
+  };
+
+  const drain = async () => {
+    if (draining || closed) return;
+    draining = true;
+    try {
+      while (!closed) {
+        const selected = [];
+        let batchBytes = 1;
+        while (selected.length < maxBatchFrames) {
+          // Stream zero carries only tiny keepalives; it must not wait for credit.
+          const next = queues.has(0) ? [0, queues.get(0)] : queues.entries().next().value;
+          if (!next) break;
+          const [streamId, queue] = next;
+          const entry = queue[0];
+          // Do not skip a larger head indefinitely in favour of smaller frames.
+          if (streamId !== 0 && sentBytes - acknowledgedBytes + entry.frame.length > windowBytes) {
+            windowLimited = true;
+            break;
+          }
+          if (streamId !== 0 && outstanding.length >= MAX_OUTSTANDING_FRAMES) break;
+          if (batchBytes + 4 + entry.frame.length > MAX_PLAINTEXT_FRAME_BYTES) break;
+          queue.shift();
+          queues.delete(streamId);
+          if (queue.length) queues.set(streamId, queue);
+          queuedBytes -= entry.frame.length;
+          queuedFrames -= 1;
+          if (streamId !== 0) {
+            sentBytes += entry.frame.length;
+            outstanding.push({ end: sentBytes, at: now() });
+          }
+          selected.push(entry);
+          batchBytes += 4 + entry.frame.length;
+        }
+        if (!selected.length) break;
+        sending = selected;
+        try {
+          // Selection precedes encryption. Never reorder encrypted IV counters.
+          await sendBatch(selected.map(entry => entry.frame));
+        } finally {
+          for (const entry of selected) entry.resolve();
+          sending = [];
+        }
+      }
+    } catch (error) {
+      fail(error);
+    } finally {
+      draining = false;
+    }
+  };
+
+  return {
+    send(frame) {
+      if (closed) return Promise.resolve();
+      const { streamId } = decodeTunnelFrame(frame);
+      if (queuedBytes + frame.length > MAX_QUEUED_BYTES || queuedFrames >= MAX_QUEUED_FRAMES) {
+        fail(new Error('relay downstream queue limit exceeded'));
+        return Promise.resolve();
+      }
+      return new Promise(resolve => {
+        const queue = queues.get(streamId) ?? [];
+        queue.push({ frame, resolve, streamId });
+        queues.set(streamId, queue);
+        queuedBytes += frame.length;
+        queuedFrames += 1;
+        // Let concurrent streams join this round before selecting the next frame.
+        queueMicrotask(() => { void drain(); });
+      });
+    },
+    acknowledge(bytes) {
+      if (closed) return;
+      if (!Number.isSafeInteger(bytes) || bytes <= acknowledgedBytes || bytes > sentBytes) {
+        throw new Error('invalid relay delivery acknowledgement');
+      }
+      // An ACK may cover several complete frames, never a partial frame.
+      const endIndex = outstanding.findIndex(entry => entry.end === bytes);
+      if (endIndex < 0) throw new Error('relay acknowledgement is not a frame boundary');
+      const rtt = Math.max(1, now() - outstanding[endIndex].at);
+      outstanding.splice(0, endIndex + 1);
+      acknowledgedBytes = bytes;
+      minRtt = Math.min(minRtt, rtt);
+      // Grow on a clear path, reduce when the client's delivery delay grows.
+      // Adjust once per window, not once per frame in a burst of ACKs.
+      if (bytes >= adjustmentBoundary) {
+        if (rtt > minRtt + QUEUE_DELAY_TARGET_MS) {
+          windowBytes = Math.max(MIN_WINDOW_BYTES, Math.floor(windowBytes / 2));
+        } else if (windowLimited) {
+          // Sparse token traffic must not inflate a future bash burst's window.
+          windowBytes = Math.min(MAX_WINDOW_BYTES, windowBytes * 2);
+        }
+        windowLimited = false;
+        adjustmentBoundary = bytes + windowBytes;
+      }
+      void drain();
+    },
+    cancel(streamId) {
+      // Already-selected bytes may still arrive and must still be ACKed, but
+      // cancelled producers need not wait for an in-progress encryption call.
+      for (const entry of sending) if (entry.streamId === streamId) entry.resolve();
+      const queue = queues.get(streamId);
+      if (!queue) return;
+      queues.delete(streamId);
+      for (const entry of queue) {
+        queuedBytes -= entry.frame.length;
+        queuedFrames -= 1;
+        entry.resolve();
+      }
+      void drain();
+    },
+    close,
+  };
+};

package/server/lib/relay/downstream-scheduler.test.js

@@ -0,0 +1,162 @@
+import { describe, expect, test } from 'bun:test';
+import { createDownstreamScheduler, DOWNSTREAM_CHUNK_BYTES } from './downstream-scheduler.js';
+import { encodeTunnelFrame, decodeTunnelFrame, TunnelFrameType } from './tunnel-codec.js';
+
+const tick = () => new Promise(resolve => setTimeout(resolve, 0));
+const body = (streamId, value = 0, size = DOWNSTREAM_CHUNK_BYTES) =>
+  encodeTunnelFrame(TunnelFrameType.HttpBody, streamId, new Uint8Array(size).fill(value));
+
+function fixture() {
+  const sent = [];
+  const errors = [];
+  let time = 0;
+  const scheduler = createDownstreamScheduler({
+    sendBatch: async frames => { sent.push(...frames); },
+    onError: error => errors.push(error.message),
+    now: () => time,
+  });
+  return { scheduler, sent, errors, advance: ms => { time += ms; } };
+}
+
+describe('relay downstream scheduling', () => {
+  test('stops without end-client ACKs even when the relay socket accepts everything', async () => {
+    const { scheduler, sent } = fixture();
+    const pending = Array.from({ length: 100 }, () => scheduler.send(body(1)));
+    await tick();
+    const bytes = sent.reduce((sum, frame) => sum + frame.length, 0);
+    expect(bytes).toBeLessThanOrEqual(64 * 1024);
+    expect(bytes).toBeGreaterThan(48 * 1024);
+    await tick();
+    expect(sent.reduce((sum, frame) => sum + frame.length, 0)).toBe(bytes);
+    scheduler.acknowledge(bytes);
+    await tick();
+    expect(sent.reduce((sum, frame) => sum + frame.length, 0)).toBeGreaterThan(bytes);
+    scheduler.close();
+    await Promise.all(pending);
+  });
+
+  test('rotates streams while preserving every delta and each stream end', async () => {
+    const { scheduler, sent } = fixture();
+    const pending = [];
+    for (let value = 0; value < 3; value++) pending.push(scheduler.send(body(1, value)));
+    pending.push(scheduler.send(encodeTunnelFrame(TunnelFrameType.StreamEnd, 1, new Uint8Array())));
+    pending.push(scheduler.send(body(3, 99, 10)));
+    await Promise.all(pending);
+    expect(sent.map(frame => decodeTunnelFrame(frame).streamId)).toEqual([1, 3, 1, 1, 1]);
+    expect(sent.filter(frame => decodeTunnelFrame(frame).streamId === 1).map(frame => decodeTunnelFrame(frame).payload[0])).toEqual([0, 1, 2, undefined]);
+    scheduler.close();
+  });
+
+  test('keepalives bypass credit; cancellation and close release blocked producers', async () => {
+    const { scheduler, sent } = fixture();
+    const pending = Array.from({ length: 20 }, () => scheduler.send(body(1)));
+    await tick();
+    const before = sent.length;
+    await scheduler.send(encodeTunnelFrame(TunnelFrameType.Pong, 0, new Uint8Array()));
+    expect(sent.length).toBe(before + 1);
+    scheduler.cancel(1);
+    await Promise.all(pending);
+    const blocked = scheduler.send(body(3));
+    scheduler.close();
+    await blocked;
+    // Teardown makes late ACKs harmless.
+    scheduler.acknowledge(123);
+  });
+
+  test('rejects forged, replayed, regressing and partial-frame ACKs', async () => {
+    const { scheduler, sent } = fixture();
+    await scheduler.send(body(1));
+    const bytes = sent[0].length;
+    for (const value of [0, -1, 1, bytes + 1, NaN, Infinity]) {
+      expect(() => scheduler.acknowledge(value)).toThrow();
+    }
+    scheduler.acknowledge(bytes);
+    expect(() => scheduler.acknowledge(bytes)).toThrow();
+    scheduler.close();
+  });
+
+  test('bounds pending memory and reports overflow instead of silently dropping deltas', async () => {
+    const { scheduler, errors } = fixture();
+    const pending = Array.from({ length: 600 }, () => scheduler.send(body(1)));
+    await Promise.all(pending);
+    expect(errors).toEqual(['relay downstream queue limit exceeded']);
+  });
+
+  test('send failure settles selected and queued producers', async () => {
+    const errors = [];
+    const scheduler = createDownstreamScheduler({
+      sendBatch: async () => { throw new Error('wire failed'); },
+      onError: error => errors.push(error.message),
+    });
+    await Promise.all([scheduler.send(body(1)), scheduler.send(body(3))]);
+    expect(errors).toEqual(['wire failed']);
+  });
+
+  test('teardown releases a producer even while encryption is still pending', async () => {
+    let release;
+    let started = false;
+    const scheduler = createDownstreamScheduler({
+      sendBatch: () => new Promise(resolve => { started = true; release = resolve; }),
+      onError: error => { throw error; },
+    });
+    const pending = scheduler.send(body(1));
+    await tick();
+    expect(started).toBe(true);
+    scheduler.close();
+    await pending;
+    release();
+    await tick();
+  });
+
+  test('sparse traffic cannot inflate the window for a later bulk burst', async () => {
+    const { scheduler, sent, advance } = fixture();
+    let acknowledged = 0;
+    for (let i = 0; i < 100; i++) {
+      const frame = body(1);
+      await scheduler.send(frame);
+      advance(50);
+      acknowledged += frame.length;
+      scheduler.acknowledge(acknowledged);
+    }
+    const before = sent.length;
+    const pending = Array.from({ length: 100 }, () => scheduler.send(body(1)));
+    await tick();
+    expect(sent.slice(before).reduce((sum, frame) => sum + frame.length, 0)).toBeLessThanOrEqual(64 * 1024);
+    scheduler.close();
+    await Promise.all(pending);
+  });
+
+  test('adapts to a clear high-RTT path, then contracts when delivery backs up', async () => {
+    const { scheduler, sent, advance, errors } = fixture();
+    let acknowledged = 0;
+    let maxFlight = 0;
+    const pending = Array.from({ length: 400 }, () => scheduler.send(body(1)));
+    for (let round = 0; round < 12; round++) {
+      // Keep enough queued output to exercise growth all the way to the cap.
+      while (pending.length < sent.length + 400) pending.push(scheduler.send(body(1)));
+      await tick();
+      const delivered = sent.reduce((sum, frame) => sum + frame.length, 0);
+      maxFlight = Math.max(maxFlight, delivered - acknowledged);
+      expect(delivered - acknowledged).toBeLessThanOrEqual(1024 * 1024);
+      advance(200);
+      scheduler.acknowledge(delivered);
+      acknowledged = delivered;
+    }
+    // Reach the cap with a backlogged sender, allowing for whole-frame sizing.
+    expect(maxFlight).toBeGreaterThan(1024 * 1024 - body(1).length);
+    // Keep the sender backlogged while ACK latency jumps to a second.
+    for (let round = 0; round < 12; round++) {
+      for (let i = 0; i < 30; i++) pending.push(scheduler.send(body(1)));
+      await tick();
+      const delivered = sent.reduce((sum, frame) => sum + frame.length, 0);
+      advance(1000);
+      scheduler.acknowledge(delivered);
+      acknowledged = delivered;
+    }
+    await tick();
+    expect(sent.reduce((sum, frame) => sum + frame.length, 0) - acknowledged).toBeLessThanOrEqual(64 * 1024);
+    expect(errors).toEqual([]);
+    scheduler.close();
+    await Promise.all(pending);
+  });
+});

package/server/lib/relay/e2ee.js

@@ -252,11 +252,12 @@ const parseHandshakeMessage = (raw) => {
   if (parsed.v !== RELAY_PROTOCOL_VERSION) return null;
   // Unknown/missing capability flag = false = legacy behavior.
   const batch = parsed.batch === true;
+  const flowControl = parsed.flowControl === true;
   if (parsed.t === 'ready') {
-    return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch };
+    return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch, flowControl };
   }
   if (parsed.t === 'hello' && typeof parsed.nonce === 'string' && typeof parsed.clientPubJwk === 'object' && parsed.clientPubJwk !== null) {
-    return { t: 'hello', v: RELAY_PROTOCOL_VERSION, clientPubJwk: parsed.clientPubJwk, nonce: parsed.nonce, batch };
+    return { t: 'hello', v: RELAY_PROTOCOL_VERSION, clientPubJwk: parsed.clientPubJwk, nonce: parsed.nonce, batch, flowControl };
   }
   return null;
 };
@@ -275,7 +276,7 @@ const failClosed = (reason) => ({
  *   { type: 'ignore' }                                — drop the frame
  *   { type: 'fail', closeCode, reason }               — close the socket with closeCode
  * @param {CryptoKey} hostEncPrivateKey long-lived ECDH private key
- * @param {{ batch?: boolean }} [options] `batch` defaults true; set false to force legacy behavior
+ * @param {{ batch?: boolean, flowControl?: boolean }} [options] Capabilities default true.
  */
 export const createHostHandshake = (hostEncPrivateKey, options = {}) => {
   const localBatch = options.batch !== false;
@@ -321,15 +322,16 @@ export const createHostHandshake = (hostEncPrivateKey, options = {}) => {
       acceptedClientKeyFingerprint = fingerprint;
       // Batching runs only if both peers advertised it.
       negotiatedBatch = localBatch && message.batch === true;
-      readyText = JSON.stringify(
-        negotiatedBatch
-          ? { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch: true }
-          : { t: 'ready', v: RELAY_PROTOCOL_VERSION },
-      );
+      const flowControl = options.flowControl !== false && message.flowControl === true;
+      const ready = { t: 'ready', v: RELAY_PROTOCOL_VERSION };
+      if (negotiatedBatch) ready.batch = true;
+      if (flowControl) ready.flowControl = true;
+      readyText = JSON.stringify(ready);
       established = true;
       return {
         type: 'established',
         batch: negotiatedBatch,
+        flowControl,
         replyText: readyText,
         channel: {
           encryptor: createFrameEncryptor(keys.hostToClient),

package/server/lib/relay/flow-control.test.js

@@ -0,0 +1,177 @@
+import { describe, expect, test } from 'bun:test';
+import http from 'node:http';
+import { WebSocketServer } from 'ws';
+import { startRelayHost } from './host-client.js';
+import { generateEcdhKeyPair, exportPublicKeyJwk } from './e2ee.js';
+import { createRelayTunnelClient } from '../../../../ui/src/lib/relay/tunnel-client.ts';
+
+const waitFor = async predicate => {
+  for (let attempt = 0; attempt < 1000; attempt++) {
+    if (predicate()) return;
+    await new Promise(resolve => setTimeout(resolve, 5));
+  }
+  throw new Error('test condition timed out');
+};
+
+async function exercise({ clientFlow = true, hostFlow = true, batch = true, cancel = false, websocket = false } = {}) {
+  const total = 512 * 1024;
+  const content = Buffer.alloc(total, 'a');
+  content.fill('b', total / 2);
+  const upstream = http.createServer((req, res) => {
+    // HTTP bearer ownership is unchanged by flow control.
+    if (req.headers.authorization !== 'Bearer fixture-token') { res.writeHead(401); res.end(); return; }
+    res.end(req.url === '/health' ? 'ok' : content);
+  });
+  const upstreamWs = new WebSocketServer({ noServer: true });
+  upstream.on('upgrade', (req, socket, head) => {
+    // Mirror the real URL-token and loopback-origin gates for tunneled WS.
+    const query = new URL(req.url, 'http://localhost');
+    if (query.searchParams.get('oc_url_token') !== 'fixture-url-token' || req.headers.origin !== `http://127.0.0.1:${upstream.address().port}`) {
+      socket.end('HTTP/1.1 403 Forbidden\r\n\r\n');
+      return;
+    }
+    upstreamWs.handleUpgrade(req, socket, head, ws => {
+      ws.send(content.subarray(0, total / 2));
+      ws.send(content.subarray(total / 2));
+      // Let Bun finish the local write turn, then close while the throttled
+      // tunnel still has most of the two messages queued.
+      setTimeout(() => ws.close(1000, 'complete'), 20);
+    });
+  });
+  await new Promise(resolve => upstream.listen(0, '127.0.0.1', resolve));
+  const relay = new WebSocketServer({ host: '127.0.0.1', port: 0 });
+  await new Promise(resolve => relay.on('listening', resolve));
+  let control;
+  let clientSocket;
+  let hostData;
+  const waiting = [];
+  const down = [];
+  // This fixture brokers an isolated room only; it does not replace production
+  // relay authentication. Neither WebSocket leg fails or drops a frame.
+  relay.on('connection', (socket, req) => {
+    const role = new URL(req.url, 'http://localhost').searchParams.get('role');
+    if (role === 'host-control') {
+      control = socket;
+      socket.send(JSON.stringify({ type: 'sync', connectionIds: clientSocket ? ['fixture'] : [] }));
+    } else if (role === 'client') {
+      clientSocket = socket;
+      control?.send(JSON.stringify({ type: 'connected', connectionId: 'fixture' }));
+      socket.on('message', (data, binary) => {
+        if (hostData) hostData.send(data, { binary });
+        else waiting.push({ data, binary });
+      });
+    } else {
+      hostData = socket;
+      for (const frame of waiting.splice(0)) socket.send(frame.data, { binary: frame.binary });
+      socket.on('message', (data, binary) => {
+        down.push({ data, binary });
+      });
+    }
+  });
+  // Drain at a byte budget, not a frame count: frame-size changes cannot make
+  // the test's effective bandwidth change. Preserve whole WS messages in order.
+  let credit = 0;
+  const drain = setInterval(() => {
+    credit = Math.min(128 * 1024, credit + 4096);
+    while (down.length && down[0].data.length <= credit) {
+      const frame = down.shift();
+      credit -= frame.data.length;
+      clientSocket.send(frame.data, { binary: frame.binary });
+    }
+  }, 5);
+  const keys = await generateEcdhKeyPair();
+  const relayUrl = `ws://127.0.0.1:${relay.address().port}/ws`;
+  const host = startRelayHost({
+    relayUrl, localPort: upstream.address().port, batch, flowControl: hostFlow,
+    identity: { serverId: 'fixture', hostEncPrivateKey: keys.privateKey, signRelayAuth: () => ({ ts: 0, sig: '', pk: '' }) },
+  });
+  const client = createRelayTunnelClient({ relayUrl, serverId: 'fixture', hostEncPubJwk: await exportPublicKeyJwk(keys.publicKey), batch, flowControl: clientFlow });
+  const states = [];
+  client.subscribeStatus(status => states.push(status.state));
+  const headers = { authorization: 'Bearer fixture-token' };
+  let received = 0;
+  let finished = false;
+  let wsMessages = 0;
+  let streamFailure;
+  const chunks = [];
+  const abort = new AbortController();
+  try {
+    await waitFor(() => client.getStatus().state === 'connected');
+    let stream;
+    if (websocket) {
+      const socket = client.openWebSocket('/api/terminal/ws?oc_url_token=fixture-url-token');
+      stream = new Promise((resolve, reject) => {
+        socket.onmessage = event => {
+          const bytes = new Uint8Array(event.data);
+          chunks.push(bytes);
+          received += bytes.length;
+          wsMessages++;
+        };
+        socket.onclose = event => {
+          finished = true;
+          if (event.code !== 1000) reject(new Error(event.reason));
+          else resolve();
+        };
+      });
+    } else {
+      const response = await client.fetch('/api/global/event', { headers, signal: abort.signal });
+      stream = (async () => {
+        try {
+          for await (const chunk of response.body) { chunks.push(chunk); received += chunk.length; }
+        } catch (error) {
+          if (!cancel || error.name !== 'AbortError') throw error;
+        } finally { finished = true; }
+      })();
+    }
+    stream = stream.catch(error => { streamFailure = error; });
+    await waitFor(() => received > 0 || streamFailure);
+    if (streamFailure) throw streamFailure;
+    if (cancel) abort.abort();
+    const response = await client.fetch('/health', { headers, signal: AbortSignal.timeout(5000) });
+    expect(await response.text()).toBe('ok');
+    const bytesAtProbe = received;
+    await stream;
+    expect(streamFailure).toBeUndefined();
+    expect(finished).toBe(true);
+    if (!cancel) expect(Buffer.concat(chunks).equals(content)).toBe(true);
+    if (websocket) expect(wsMessages).toBe(2);
+    expect(states).toEqual(['connected']);
+    return { bytesAtProbe, total };
+  } finally {
+    client.close();
+    host.stop();
+    clearInterval(drain);
+    for (const ws of relay.clients) ws.terminate();
+    relay.close();
+    for (const ws of upstreamWs.clients) ws.terminate();
+    upstreamWs.close();
+    upstream.closeAllConnections();
+    upstream.close();
+  }
+}
+
+describe('end-to-end downstream credit', () => {
+  test('small HTTP response overtakes bulk output without losing bytes or reconnecting', async () => {
+    const result = await exercise();
+    expect(result.bytesAtProbe).toBeLessThan(result.total / 2);
+    // Queue peaks depend on real ACK timing and include encryption overhead.
+    // downstream-scheduler.test.js checks credit bounds with a controlled clock.
+  });
+  test('works without batching', async () => {
+    const result = await exercise({ batch: false });
+    expect(result.bytesAtProbe).toBeLessThan(result.total / 2);
+  });
+  test('cancelling a blocked stream leaves unrelated requests usable', async () => {
+    const result = await exercise({ cancel: true });
+    expect(result.bytesAtProbe).toBeLessThan(result.total);
+  });
+  test('fragmented WS messages remain complete and precede the close', async () => {
+    await exercise({ websocket: true });
+  });
+  for (const legacy of [{ clientFlow: false }, { hostFlow: false }]) {
+    test(`legacy fallback ${JSON.stringify(legacy)} preserves delivery`, async () => {
+      const result = await exercise(legacy);
+      expect(result.bytesAtProbe).toBe(result.total);
+    });
+  }
+});

package/server/lib/relay/host-client.js

@@ -4,10 +4,12 @@
 // dispatcher. Spec: .opencode/plans/private-relay/01-protocol-spec.md (Layer 1).
 
 import { WebSocket } from 'ws';
+import { createRequire } from 'node:module';
 
 import { RELAY_PROTOCOL_VERSION, RelayCloseCode, createHostHandshake } from './e2ee.js';
-import { createOutboundFrameBatcher, decodeFrameBatch } from './tunnel-codec.js';
+import { createOutboundFrameBatcher, decodeFrameBatch, decodeTunnelFrame, decodeDeliveryAck, encodeFrameBatch, TunnelFrameType } from './tunnel-codec.js';
 import { createTunnelHost } from './tunnel-host.js';
+import { createDownstreamScheduler, DOWNSTREAM_CHUNK_BYTES } from './downstream-scheduler.js';
 
 const BACKOFF_BASE_MS = 1000;
 const BACKOFF_CAP_MS = 30000;
@@ -48,7 +50,9 @@ const resolveBatchWindowMs = (option) => {
  *   logger?: Pick<Console, 'warn'>,
  * }} options
  */
-export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, onStatus, logger = console, batchWindowMs, batch }) => {
+export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, onStatus, logger = console, batchWindowMs, batch, flowControl }) => {
+  const { version } = createRequire(import.meta.url)('../../../package.json');
+  const platform = process.env.AIWORKER_RUNTIME || 'web';
   const resolveLocalPort = typeof getLocalPort === 'function' ? getLocalPort : () => localPort;
   const localBatch = batch !== false;
   const resolvedBatchWindowMs = resolveBatchWindowMs(batchWindowMs);
@@ -81,6 +85,10 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
     url.searchParams.set('v', String(RELAY_PROTOCOL_VERSION));
     url.searchParams.set('role', role);
     url.searchParams.set('serverId', identity.serverId);
+    // Self-reported diagnostics, not part of relay authentication.
+    url.searchParams.set('appId', 'aiworker');
+    url.searchParams.set('appVersion', version);
+    url.searchParams.set('platform', platform);
     if (connectionId) url.searchParams.set('connectionId', connectionId);
     const auth = identity.signRelayAuth(role, connectionId ?? null);
     url.searchParams.set('ts', String(auth.ts));
@@ -95,6 +103,7 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
     dataSockets.delete(connectionId);
     if (entry.openTimer) clearTimeout(entry.openTimer);
     entry.batcher?.dispose();
+    entry.scheduler?.close();
     entry.tunnel?.close();
     try {
       if (entry.socket.readyState === WebSocket.OPEN || entry.socket.readyState === WebSocket.CONNECTING) {
@@ -118,14 +127,14 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
       return;
     }
 
-    const entry = { socket, tunnel: null, openTimer: null, batcher: null, lastActivityAt: Date.now() };
+    const entry = { socket, tunnel: null, openTimer: null, batcher: null, scheduler: null, lastActivityAt: Date.now() };
     dataSockets.set(connectionId, entry);
     entry.openTimer = setTimeout(() => {
       logger.warn('[Relay] host-data socket open timeout');
       teardownDataSocket(connectionId);
     }, DATA_SOCKET_OPEN_TIMEOUT_MS);
 
-    const handshake = createHostHandshake(identity.hostEncPrivateKey, { batch: localBatch });
+    const handshake = createHostHandshake(identity.hostEncPrivateKey, { batch: localBatch, flowControl });
     let channel = null;
     let batchNegotiated = false;
     // Serialize async message handling so encrypted frame order (and the
@@ -140,11 +149,14 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
         .then(async () => {
           if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN || !channel) return;
           const encrypted = await channel.encryptor.encrypt(plaintext);
+          if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN) return;
           socket.send(encrypted, { binary: true });
         })
         .catch((error) => {
           logger.warn(`[Relay] host-data send failed: ${error?.message ?? error}`);
+          failChannel(RelayCloseCode.ChannelFailure, 'send failed');
         });
+      return sendChain;
     };
 
     const failChannel = (closeCode, reason) => {
@@ -167,17 +179,25 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
         } else if (action.type === 'established') {
           channel = action.channel;
           batchNegotiated = action.batch === true;
-          entry.batcher = batchNegotiated
+          entry.scheduler = action.flowControl ? createDownstreamScheduler({
+            sendBatch: frames => sendEncryptedPlaintext(batchNegotiated ? encodeFrameBatch(frames) : frames[0]),
+            maxBatchFrames: batchNegotiated ? 32 : 1,
+            onError: () => failChannel(RelayCloseCode.ChannelFailure, 'downstream queue failed'),
+          }) : null;
+          entry.batcher = batchNegotiated && !entry.scheduler
             ? createOutboundFrameBatcher({ windowMs: resolvedBatchWindowMs, sendBatch: sendEncryptedPlaintext })
             : null;
           entry.tunnel = createTunnelHost({
             connectionId,
             getLocalPort: resolveLocalPort,
             getBufferedAmount: () => socket.bufferedAmount,
+            responseChunkBytes: entry.scheduler ? DOWNSTREAM_CHUNK_BYTES : undefined,
+            cancelPendingFrames: streamId => entry.scheduler?.cancel(streamId),
             sendFrame: (plaintextFrame) => {
               if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN) return;
+              if (entry.scheduler) return entry.scheduler.send(plaintextFrame);
               if (entry.batcher) entry.batcher.enqueue(plaintextFrame);
-              else sendEncryptedPlaintext(plaintextFrame);
+              else return sendEncryptedPlaintext(plaintextFrame);
             },
           });
           if (action.replyText) socket.send(action.replyText);
@@ -205,16 +225,32 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on
           // in order through the same per-frame handling as legacy.
           for (const frame of decodeFrameBatch(plaintext)) {
             if (dataSockets.get(connectionId) !== entry) return;
-            await entry.tunnel.handleFrame(frame);
+            await dispatchFrame(frame);
           }
         } else {
-          await entry.tunnel.handleFrame(plaintext);
+          await dispatchFrame(plaintext);
         }
       } catch (error) {
         logger.warn(`[Relay] tunnel frame handling failed: ${error?.message ?? error}`);
+        failChannel(RelayCloseCode.ChannelFailure, 'invalid tunnel frame');
       }
     };
 
+    const dispatchFrame = (plaintext) => {
+      const frame = decodeTunnelFrame(plaintext);
+      if (frame.frameType === TunnelFrameType.DeliveryAck) {
+        if (!entry.scheduler || frame.streamId !== 0 || frame.hasMoreFragments) {
+          throw new Error('unexpected delivery acknowledgement');
+        }
+        entry.scheduler.acknowledge(decodeDeliveryAck(frame.payload));
+        return;
+      }
+      // Never await outbound credit on the receive chain: ACKs use this chain too.
+      void entry.tunnel.handleFrame(plaintext).catch(() => {
+        failChannel(RelayCloseCode.ChannelFailure, 'invalid tunnel frame');
+      });
+    };
+
     socket.on('open', () => {
       if (entry.openTimer) {
         clearTimeout(entry.openTimer);