fix(transport): preserve backpressure across framing and routes

Ordinary daemon writes must remain visible to terminal catch-up and socket admission. Pace chunked writes and explicit completion requests, include waiting logical messages in pressure readings, and propagate origin completion through routed client transports.
This commit is contained in:
Mohamed Boudra
2026-10-30 15:51:50 +01:00
parent 98f667f93e
commit f28f0fee3d
9 changed files with 278 additions and 56 deletions
+37
View File
@@ -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<void>((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);
});
+23 -18
View File
@@ -1884,9 +1884,9 @@ export class DaemonClient {
})
).routeId;
}
sendHostFrame(routeId: string, data: string | Uint8Array): void {
sendHostFrame(routeId: string, data: string | Uint8Array): Promise<void> {
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<void> {
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<void> {
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<void> {
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<void> {
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));
}
+1 -1
View File
@@ -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),
);
@@ -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]);
});
+31 -13
View File
@@ -14,41 +14,58 @@ export class ChunkedTransport {
private closed = false;
private pending: Promise<void> | undefined;
private queuedBytes = 0;
private activeBytes = 0;
constructor(
private readonly io: {
frameBytes: number;
send: (frame: Frame) => void | Promise<void>;
send: (frame: Frame, waitForDrain: boolean) => void | Promise<void>;
deliver: (frame: Frame) => void;
},
) {
if (io.frameBytes <= HEADER) throw new Error("Invalid transport frame budget");
}
send(frame: Frame): Promise<void> {
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<void> {
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<void>;
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);
}
}
@@ -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);
+52 -21
View File
@@ -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<void> {
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, {
@@ -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<typeof WebSocket>];
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<void>((resolve) => server.close(() => resolve()));
}
});