diff --git a/docs/protocol-compatibility.md b/docs/protocol-compatibility.md index 46b9eacdff..4d719bc15f 100644 --- a/docs/protocol-compatibility.md +++ b/docs/protocol-compatibility.md @@ -67,9 +67,17 @@ consume slices incrementally. Today the receiver buffers up to the existing Disconnect clears partial state and negotiation. The lower transport supplies its plaintext frame budget, including encryption -overhead. The sender waits for each physical write before sending the next chunk. +overhead. The sender waits for drain between chunks while ordinary messages keep +the existing delivery path. Slow-client protections count messages queued behind +an active transfer as well as bytes already in the socket. The framing module never selects an encryption construction. +Native uploads retain the platform's existing queue limits. React Native's +WebSocket does not report its queued bytes, so the client cannot wait for native +drain. Android's installed OkHttp transport closes when a send would put more +than 16 MiB in its queue. Negotiating chunking does not establish large native +upload support; daemon-to-native responses still use the daemon's paced writes. + ## Owned observations `owned_subscriptions` and `server_info.features.ownedSubscriptions` negotiate the source-owned diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts index 058c9ba77b..ad388eab0f 100644 --- a/packages/client/src/daemon-client.test.ts +++ b/packages/client/src/daemon-client.test.ts @@ -8504,3 +8504,40 @@ test("uses the frame budget selected by an asynchronously opened transport", asy true, ); }); + +test.each([false, true])("routed sends wait for origin drain (binary=%s)", async (binary) => { + const mock = createMockTransport(); + let drain!: () => void; + let blocked = false; + const send = mock.transport.send; + mock.transport.send = (frame) => { + send(frame); + if (blocked) + return new Promise((resolve) => { + drain = resolve; + }); + }; + const client = new DaemonClient({ + url: "ws://test", + clientId: "routed-drain", + transportFactory: () => mock.transport, + reconnect: { enabled: false }, + }); + clients.push(client); + const connecting = client.connect(); + mock.triggerOpen({ features: { binaryChunkingV1: true } }); + await connecting; + blocked = true; + let finished = false; + const write = Promise.resolve( + client.sendHostFrame("route", binary ? new Uint8Array([1]) : "hello"), + ).then(() => { + finished = true; + return undefined; + }); + await Promise.resolve(); + expect(finished).toBe(false); + drain(); + await write; + expect(finished).toBe(true); +}); diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 8f8ec74ed4..7348455407 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1884,9 +1884,9 @@ export class DaemonClient { }) ).routeId; } - sendHostFrame(routeId: string, data: string | Uint8Array): void { + sendHostFrame(routeId: string, data: string | Uint8Array): Promise { if (typeof data === "string") { - this.sendSessionMessageStrict({ + return this.sendSessionMessageStrict({ type: "hosts.route.frame.request", routeId, data, @@ -1894,7 +1894,7 @@ export class DaemonClient { }); } else { if (!this.isConnected || !this.transport) throw this.unavailableError(); - this.sendTransportFrame(encodeHostRouteFrame(routeId, data)); + return this.sendTransportFrame(encodeHostRouteFrame(routeId, data)); } } private readonly hostFrameListeners = new Set<(frame: HostRouteFrame) => void>(); @@ -2060,7 +2060,7 @@ export class DaemonClient { messageType: string, message: unknown, submitted?: () => void, - ): void { + ): Promise { const isOpen = this.config.trace?.isEnabled() === true && this.beginTraceSection("paseo.ws.message.outbound", { @@ -2069,25 +2069,30 @@ export class DaemonClient { ...describeRequestIdentity(message), }); try { - this.sendTransportFrame(JSON.stringify(message), submitted); + return this.sendTransportFrame(JSON.stringify(message), submitted); } finally { this.endTraceSection(isOpen); } } - private sendSessionPayload(payload: SessionInboundMessage): void { - this.sendJsonMessage("session", payload.type, { type: "session", message: payload }, () => { - if (!("requestId" in payload) || typeof payload.requestId !== "string") return; - for (const waiter of this.waiters) { - if (waiter.requestId === payload.requestId) waiter.submitted = true; - } - }); + private sendSessionPayload(payload: SessionInboundMessage): Promise { + return this.sendJsonMessage( + "session", + payload.type, + { type: "session", message: payload }, + () => { + if (!("requestId" in payload) || typeof payload.requestId !== "string") return; + for (const waiter of this.waiters) { + if (waiter.requestId === payload.requestId) waiter.submitted = true; + } + }, + ); } private sendTransportFrame( frame: string | Uint8Array | ArrayBuffer, submitted?: () => void, - ): void { + ): Promise { if (!this.transport) { throw this.unavailableError(); } @@ -2102,9 +2107,9 @@ export class DaemonClient { submitted?.(); const payload = typeof frame === "string" ? frame : asUint8Array(frame)!; const transport = this.transport; - void this.chunkedTransport!.send(payload).catch(() => - transport.close(4003, "Transport send failed"), - ); + const write = this.chunkedTransport!.send(payload); + void write.catch(() => transport.close(4003, "Transport send failed")); + return write; } finally { this.endTraceSection(isOpen); } @@ -2376,13 +2381,13 @@ export class DaemonClient { }); } - private sendSessionMessageStrict(message: SessionInboundMessage): void { + private sendSessionMessageStrict(message: SessionInboundMessage): Promise { if (!this.transport || this.connectionState.status !== "connected") { throw this.unavailableError(); } const payload = SessionInboundMessageSchema.parse(message); try { - this.sendSessionPayload(payload); + return this.sendSessionPayload(payload); } catch (error) { throw error instanceof Error ? error : new Error(String(error)); } diff --git a/packages/client/src/hosts/index.ts b/packages/client/src/hosts/index.ts index 97b25f06da..b1c6451365 100644 --- a/packages/client/src/hosts/index.ts +++ b/packages/client/src/hosts/index.ts @@ -92,7 +92,7 @@ export function createHostChannelTransport( code: "host_offline", serverId, }); - origin.sendHostFrame( + return origin.sendHostFrame( routeId, typeof data === "string" || data instanceof Uint8Array ? data : new Uint8Array(data), ); diff --git a/packages/protocol/src/chunked-transport.test.ts b/packages/protocol/src/chunked-transport.test.ts index 0c2db38d73..6aa8a19ce7 100644 --- a/packages/protocol/src/chunked-transport.test.ts +++ b/packages/protocol/src/chunked-transport.test.ts @@ -84,13 +84,16 @@ it("waits for drain before the next chunk and preserves logical ordering", async }); sender.enabled = true; const first = sender.send("abcdefghij"); - const second = sender.send("next"); + const next = new Uint8Array([1, 2, 3, 4]); + const second = sender.send(next); await Promise.resolve(); expect(frames).toHaveLength(1); + expect(sender.queuedByteLength).toBe(4); drain(); await Promise.all([first, second]); expect(frames).toHaveLength(4); - expect(frames[3]).toBe("next"); + expect(frames[3]).toEqual(next); + expect(sender.queuedByteLength).toBe(0); }); it.each([ @@ -146,3 +149,18 @@ it("bounds logical messages waiting behind a blocked writer", async () => { await first; } }); + +it("requests drain for chunks while ordinary frames keep the immediate write path", async () => { + const waits: boolean[] = []; + const transport = new ChunkedTransport({ + frameBytes: 8, + send: (_frame, wait) => { + waits.push(wait); + }, + deliver: () => {}, + }); + transport.enabled = true; + await transport.send("small"); + await transport.send("abcdefghij"); + expect(waits).toEqual([false, true, true, true]); +}); diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index 2b9eb40b91..e32f0e3d52 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -14,41 +14,58 @@ export class ChunkedTransport { private closed = false; private pending: Promise | undefined; private queuedBytes = 0; + private activeBytes = 0; constructor( private readonly io: { frameBytes: number; - send: (frame: Frame) => void | Promise; + send: (frame: Frame, waitForDrain: boolean) => void | Promise; deliver: (frame: Frame) => void; }, ) { if (io.frameBytes <= HEADER) throw new Error("Invalid transport frame budget"); } - send(frame: Frame): Promise { - const bytes = typeof frame === "string" ? new TextEncoder().encode(frame) : frame; - if (bytes.byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) - throw new Error("Logical message too large"); - const chunk = this.enabled && bytes.byteLength > this.io.frameBytes; + /** Bytes waiting behind the active write; physical buffering owns the active frame. */ + get queuedByteLength(): number { + return this.queuedBytes - this.activeBytes; + } + get isWriting(): boolean { + return this.pending !== undefined; + } + + send(frame: Frame, waitForDrain = false): Promise { + const budget = this.enabled ? this.io.frameBytes : MAX_WEBSOCKET_MESSAGE_BYTES; + let bytes: Uint8Array | undefined; + if (typeof frame !== "string") bytes = frame; + else if (frame.length * 3 > budget) bytes = new TextEncoder().encode(frame); + const byteLength = bytes?.byteLength ?? (frame as string).length * 3; + if (byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) throw new Error("Logical message too large"); + const chunk = this.enabled && byteLength > this.io.frameBytes; if (this.closed) throw new Error("Transport closed"); const write = () => { if (this.closed) throw new Error("Transport closed"); - return chunk ? this.sendChunks(bytes, typeof frame === "string") : this.io.send(frame); + this.activeBytes = byteLength; + return chunk + ? this.sendChunks(bytes!, typeof frame === "string") + : this.io.send(frame, waitForDrain); }; - if (this.queuedBytes + bytes.byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) { + if (this.queuedBytes + byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) { this.close(); return Promise.reject(new Error("Logical output queue exceeded its byte limit")); } - this.queuedBytes += bytes.byteLength; + this.queuedBytes += byteLength; let result: void | Promise; try { result = this.pending ? this.pending.then(write) : write(); } catch (error) { - this.queuedBytes -= bytes.byteLength; + this.queuedBytes -= byteLength; + this.activeBytes = 0; throw error; } if (!result) { - this.queuedBytes -= bytes.byteLength; + this.queuedBytes -= byteLength; + this.activeBytes = 0; return Promise.resolve(); } const pending = Promise.resolve(result) @@ -57,7 +74,8 @@ export class ChunkedTransport { throw error; }) .finally(() => { - this.queuedBytes -= bytes.byteLength; + this.queuedBytes -= byteLength; + this.activeBytes = 0; }); this.pending = pending; void pending.then( @@ -83,7 +101,7 @@ export class ChunkedTransport { (offset === 0 ? 1 : 0) | (end === bytes.byteLength ? 2 : 0), ]); wire.set(bytes.subarray(offset, end), HEADER); - await this.io.send(wire); + await this.io.send(wire, true); } } diff --git a/packages/server/src/server/daemon-e2e/terminal.e2e.test.ts b/packages/server/src/server/daemon-e2e/terminal.e2e.test.ts index 633bd3bc1c..3c131ecab3 100644 --- a/packages/server/src/server/daemon-e2e/terminal.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/terminal.e2e.test.ts @@ -1836,3 +1836,58 @@ describe("list terminals across directories", () => { ]); }); }); + +test("a paused terminal consumer catches up through a snapshot instead of retaining every output frame", async () => { + const { existsSync } = await import("node:fs"); + const cwd = tmpCwd(); + tempDirs.push(cwd); + const marker = path.join(cwd, "output-complete"); + const created = await createTerminalInWorkspace(ctx.client, { + cwd, + options: { + command: process.execPath, + args: [ + "-e", + ` + const fs = require('node:fs'); + process.stdin.setRawMode?.(true); + process.stdin.resume(); + process.stdin.once('data', () => { + const chunk = Buffer.alloc(256 * 1024, 120); + let remaining = 128; + const write = () => { + if (remaining-- === 0) { fs.writeFileSync(${JSON.stringify(marker)}, 'done'); return; } + process.stdout.write(chunk, write); + }; + write(); + }); + `, + ], + }, + }); + const ws = await connectRawWebSocket(ctx.daemon.port); + try { + const slot = await subscribeRawTerminal(ws, created.terminal!.id, "slow-terminal"); + await waitForRawBinaryFrame( + ws, + (frame) => frame.slot === slot && frame.opcode === TerminalStreamOpcode.Snapshot, + ); + const capture = startTerminalOutputCapture(ws, slot); + try { + ws.pause(); + sendRawTerminalInput(ws, slot, "go\n"); + await waitForCondition(() => existsSync(marker), 20000); + ws.resume(); + await waitForCondition(() => capture.snapshotCount > 0, 10000); + expect(ws.readyState).toBe(WebSocket.OPEN); + expect(capture.frames.reduce((total, frame) => total + frame.payloadBytes, 0)).toBeLessThan( + 32 * 1024 * 1024, + ); + } finally { + ws.resume(); + capture.stop(); + } + } finally { + await closeWebSocket(ws); + } +}, 40000); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index bc86fa02a2..0b9cecd9fa 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -126,6 +126,7 @@ import { outboundFrameByteLength, physicalSocketHasCapacity, sendBoundedPhysicalFrameAndWait, + sendBoundedPhysicalFrame, } from "./websocket/physical-socket.js"; const WS_CLOSE_DAEMON_AUTH_FAILED = 4401; @@ -1285,13 +1286,18 @@ export class VoiceAssistantWebSocketServer { if (!transport) { transport = new ChunkedTransport({ frameBytes: ws.maxFrameBytes ?? 32 * 1024 * 1024, - send: async (frame) => { - const sent = await sendBoundedPhysicalFrameAndWait({ + send: (frame, waitForDrain) => { + const params = { socket: ws, frame, onHighWater: () => this.closeAtOutboundHighWater(ws), - }); - if (!sent) throw new Error("Source is closed"); + }; + if (waitForDrain) + return sendBoundedPhysicalFrameAndWait(params).then((sent) => { + if (!sent) throw new Error("Source is closed"); + return undefined; + }); + if (!sendBoundedPhysicalFrame(params)) throw new Error("Source is closed"); }, deliver: (frame) => this.handleRawMessage(ws, typeof frame === "string" ? frame : Buffer.from(frame)), @@ -1301,6 +1307,27 @@ export class VoiceAssistantWebSocketServer { return transport; } + private socketBufferedAmount(ws: WebSocketLike): number | null { + const queued = this.chunkedSockets.get(ws)?.queuedByteLength ?? 0; + if (typeof ws.bufferedAmount === "number") return ws.bufferedAmount + queued; + return queued > 0 ? queued : null; + } + + private sendSocketFrame( + ws: WebSocketLike, + frame: string | Uint8Array, + bytes: number, + waitForDrain = false, + ): Promise { + const transport = this.chunkedSocket(ws); + const admittedBytes = transport.isWriting + ? bytes + : Math.min(bytes, ws.maxFrameBytes ?? 32 * 1024 * 1024); + if (!this.ensureOutboundCapacity(ws, admittedBytes)) + return Promise.reject(new Error("Source is closed")); + return transport.send(frame, waitForDrain); + } + private async sendMessageToClientAndWait( ws: WebSocketLike, message: WSOutboundMessage, @@ -1308,7 +1335,7 @@ export class VoiceAssistantWebSocketServer { const prepared = this.prepareJsonDelivery([ws], message, true); if (!prepared) throw new Error("Source message delivery unavailable"); try { - await this.chunkedSocket(ws).send(prepared.payload); + await this.sendSocketFrame(ws, prepared.payload, prepared.bytes, true); this.runtimeMetrics.recordOutboundMessage(message, ws.bufferedAmount); } catch (err) { this.logger.warn({ err }, "ws_send_failed"); @@ -1335,7 +1362,7 @@ export class VoiceAssistantWebSocketServer { return; } try { - await this.chunkedSocket(ws).send(frame); + await this.sendSocketFrame(ws, frame, frame.byteLength, true); this.runtimeMetrics.recordOutboundBinaryFrame(ws.bufferedAmount); } catch (err) { this.logger.warn({ err }, "ws_send_failed"); @@ -1346,16 +1373,14 @@ export class VoiceAssistantWebSocketServer { private sendFrameToClient( ws: WebSocketLike, frame: string | Uint8Array, - _frameBytes: number, + frameBytes: number, recordSent: () => void, ): void { try { - void this.chunkedSocket(ws) - .send(frame) - .then(recordSent, (err) => { - this.logger.warn({ err }, "ws_send_failed"); - ws.close(1011, "Transport send failed"); - }); + void this.sendSocketFrame(ws, frame, frameBytes).then(recordSent, (err) => { + this.logger.warn({ err }, "ws_send_failed"); + ws.close(1011, "Transport send failed"); + }); } catch (err) { this.logger.warn({ err }, "ws_send_failed"); } @@ -1363,7 +1388,13 @@ export class VoiceAssistantWebSocketServer { private ensureOutboundCapacity(ws: WebSocketLike, frameBytes: number): boolean { if (ws.readyState !== 1) return false; - if (physicalSocketHasCapacity(ws, frameBytes)) return true; + if ( + physicalSocketHasCapacity( + { bufferedAmount: this.socketBufferedAmount(ws) ?? undefined }, + frameBytes, + ) + ) + return true; this.closeAtOutboundHighWater(ws); return false; @@ -1536,16 +1567,15 @@ export class VoiceAssistantWebSocketServer { if (!connection) { return null; } - if (source) return (source as WebSocketLike).bufferedAmount ?? null; + if (source) return this.socketBufferedAmount(source as WebSocketLike); // Relay-attached sockets are a WebSocketLike that doesn't expose // bufferedAmount. Return null when no socket gives a signal so the // terminal fallback can't mistake "no signal" for "client keeping up"; // a direct ws reports its real buffered bytes (0 when drained). let maxBuffered: number | null = null; for (const socket of connection.sockets) { - if (typeof socket.bufferedAmount === "number") { - maxBuffered = Math.max(maxBuffered ?? 0, socket.bufferedAmount); - } + const buffered = this.socketBufferedAmount(socket); + if (buffered !== null) maxBuffered = Math.max(maxBuffered ?? 0, buffered); } return maxBuffered; }, @@ -1582,8 +1612,9 @@ export class VoiceAssistantWebSocketServer { let total = 0; for (const socket of this.sessions.keys()) { if (socket instanceof VirtualHostSocket) continue; - if (typeof socket.bufferedAmount !== "number") return null; - total += socket.bufferedAmount; + const buffered = this.socketBufferedAmount(socket); + if (buffered === null) return null; + total += buffered; } return total; }, @@ -2589,7 +2620,7 @@ export class VoiceAssistantWebSocketServer { identity: JSON.stringify([connection.principalId, connection.clientId]), permissions: () => connection.session.getPermissions(), routed: () => connection.session.isRouted(), - bufferedAmount: () => ws.bufferedAmount ?? null, + bufferedAmount: () => this.socketBufferedAmount(ws), delivery: connection.session.delivery, admit: (socket, permissions) => this.attachSocket(socket, undefined, undefined, false, { diff --git a/packages/server/src/server/websocket/physical-socket.test.ts b/packages/server/src/server/websocket/physical-socket.test.ts index c81b1410d6..315f6bd2eb 100644 --- a/packages/server/src/server/websocket/physical-socket.test.ts +++ b/packages/server/src/server/websocket/physical-socket.test.ts @@ -120,3 +120,53 @@ test("the awaitable physical send rejects callback errors", async () => { }), ).rejects.toThrow("send failed"); }); + +test("ordinary frames expose a paused real socket's backlog and terminate at the physical backstop", async () => { + const { once } = await import("node:events"); + const { WebSocket, WebSocketServer } = await import("ws"); + const { ChunkedTransport } = await import("@getpaseo/protocol/chunked-transport"); + const server = new WebSocketServer({ host: "127.0.0.1", port: 0 }); + await once(server, "listening"); + const connected = once(server, "connection"); + const address = server.address(); + if (typeof address === "string" || !address) throw new Error("Missing server address"); + const client = new WebSocket(`ws://127.0.0.1:${address.port}`); + await once(client, "open"); + const [socket] = (await connected) as [InstanceType]; + const tcp = (socket as unknown as { _socket: import("node:net").Socket })._socket; + let terminated = false; + const transport = new ChunkedTransport({ + frameBytes: 32 * 1024 * 1024, + deliver: () => {}, + send: (frame, wait) => { + const options = { + socket, + frame, + onHighWater: () => { + terminated = true; + socket.terminate(); + }, + }; + if (wait) return sendBoundedPhysicalFrameAndWait(options).then(() => undefined); + sendBoundedPhysicalFrame(options); + }, + }); + try { + client.pause(); + tcp.cork(); + const frame = new Uint8Array(256 * 1024); + for (let count = 0; count < 200; count++) await transport.send(frame); + expect(socket.bufferedAmount).toBeGreaterThan(4 * 1024 * 1024); + expect(transport.queuedByteLength).toBe(0); + for (let count = 0; count < 100; count++) { + if (terminated) break; + await transport.send(frame); + } + expect(terminated).toBe(true); + } finally { + transport.close(); + client.terminate(); + socket.terminate(); + await new Promise((resolve) => server.close(() => resolve())); + } +});