diff --git a/packages/app/e2e/browser/composer-diff-stat.spec.ts b/packages/app/e2e/browser/composer-diff-stat.spec.ts index 9b80222432..959567f58b 100644 --- a/packages/app/e2e/browser/composer-diff-stat.spec.ts +++ b/packages/app/e2e/browser/composer-diff-stat.spec.ts @@ -18,7 +18,7 @@ function visibleMainPane(page: Page) { } function composerChangesPill(page: Page) { - return page.getByTestId("composer-diff-stat-pill"); + return page.getByRole("button", { name: "Open Changes tab", exact: true }); } async function revealComposerChangesInExplorer(page: Page) { @@ -498,6 +498,7 @@ test.describe("send review feedback", () => { test("only open tabs in this workspace appear in compact and wide recipient menus", async ({ page, }) => { + test.setTimeout(120000); const workspace = await seedChangedAgent("review-feedback-multiple-"); const other = await seedChangedAgent("review-feedback-other-"); try { diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 7348455407..90ea2602df 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1567,6 +1567,7 @@ export class DaemonClient { return transport.maxFrameBytes ?? 32 * 1024 * 1024; }, send: (frame) => transport.send(frame), + onFailure: () => transport.close(4003, "Transport send failed"), deliver: (frame) => this.handleTransportMessage(frame), }); this.lastServerInfoMessage = null; @@ -2106,10 +2107,7 @@ export class DaemonClient { try { submitted?.(); const payload = typeof frame === "string" ? frame : asUint8Array(frame)!; - const transport = this.transport; - const write = this.chunkedTransport!.send(payload); - void write.catch(() => transport.close(4003, "Transport send failed")); - return write; + return this.chunkedTransport!.send(payload); } finally { this.endTraceSection(isOpen); } diff --git a/packages/protocol/src/chunked-transport.test.ts b/packages/protocol/src/chunked-transport.test.ts index 6aa8a19ce7..4e758b05ee 100644 --- a/packages/protocol/src/chunked-transport.test.ts +++ b/packages/protocol/src/chunked-transport.test.ts @@ -164,3 +164,139 @@ it("requests drain for chunks while ordinary frames keep the immediate write pat await transport.send("abcdefghij"); expect(waits).toEqual([false, true, true, true]); }); + +it("charges queued ASCII messages by their UTF-8 bytes", async () => { + let drain!: () => void; + let blocked = true; + const transport = new ChunkedTransport({ + frameBytes: 32 * 1024 * 1024, + send: () => { + if (blocked) { + blocked = false; + return new Promise((resolve) => { + drain = resolve; + }); + } + }, + deliver: () => {}, + }); + transport.enabled = true; + const first = transport.send(new Uint8Array([1])); + const text = "x".repeat(10 * 1024 * 1024); + const queued = Array.from({ length: 4 }, () => transport.send(text)); + const settled = Promise.allSettled(queued); + expect(transport.queuedByteLength).toBe(40 * 1024 * 1024); + drain(); + await first; + expect((await settled).map((result) => result.status)).toEqual([ + "fulfilled", + "fulfilled", + "fulfilled", + "fulfilled", + ]); +}); + +it("notifies the socket owner when an awaited send overflows the queue", async () => { + let drain!: () => void; + const failures: unknown[] = []; + const transport = new ChunkedTransport({ + frameBytes: 32 * 1024 * 1024, + send: () => + new Promise((resolve) => { + drain = resolve; + }), + deliver: () => {}, + onFailure: (error) => { + failures.push(error); + }, + }); + const first = transport.send(new Uint8Array(100 * 1024 * 1024)); + const overflow = transport.send("extra"); + await expect(overflow).rejects.toThrow("queue"); + drain(); + await first; + expect(failures).toHaveLength(1); +}); + +it("bounds retained reassembly memory for many one-byte chunks", async () => { + const { spawnSync } = await import("node:child_process"); + // Exercise the built package as a Node consumer; CI builds it before protocol tests. + const entry = new URL("../dist/chunked-transport.js", import.meta.url).href; + const result = spawnSync( + process.execPath, + [ + "--expose-gc", + "--input-type=module", + "-e", + ` + import { ChunkedTransport } from ${JSON.stringify(entry)}; + let delivered = 0; + const receiver = new ChunkedTransport({ frameBytes: 8, send() {}, deliver(frame) { delivered = frame.byteLength; } }); + receiver.enabled = true; + global.gc(); + const before = process.memoryUsage().heapUsed; + receiver.receive(new Uint8Array([0x31, 1, 1, 1, 65])); + for (let index = 0; index < 300000; index++) receiver.receive(new Uint8Array([0x31, 1, 1, 0, 65])); + global.gc(); + const growth = process.memoryUsage().heapUsed - before; + receiver.receive(new Uint8Array([0x31, 1, 1, 2, 66])); + console.log(JSON.stringify({ growth, delivered })); + `, + ], + { encoding: "utf8", timeout: 15000 }, + ); + expect(result.status, result.stderr).toBe(0); + const observed = JSON.parse(result.stdout); + expect(observed.delivered).toBe(300002); + expect(observed.growth).toBeLessThan(8 * 1024 * 1024); +}, 20000); + +it.each(["ASCII", "中", "🌍", "\ud800", "\udc00", "\ud800x"])( + "accounts for queued UTF-8 text %j exactly", + async (text) => { + let drain!: () => void; + let blocked = true; + const transport = new ChunkedTransport({ + frameBytes: 1024, + send: () => { + if (blocked) { + blocked = false; + return new Promise((resolve) => { + drain = resolve; + }); + } + }, + deliver: () => {}, + }); + const first = transport.send(new Uint8Array([1])); + const second = transport.send(text); + expect(transport.queuedByteLength).toBe(new TextEncoder().encode(text).byteLength); + drain(); + await Promise.all([first, second]); + }, +); + +it.each([false, true])( + "notifies the owner once when a writer fails (async=%s)", + async (asynchronous) => { + const failures: unknown[] = []; + const error = new Error("write failed"); + const transport = new ChunkedTransport({ + frameBytes: 1024, + deliver: () => {}, + onFailure: (failure) => { + failures.push(failure); + }, + send: () => { + if (asynchronous) return Promise.reject(error); + throw error; + }, + }); + await expect(async () => transport.send("hello", { waitForDrain: true })).rejects.toThrow( + "write failed", + ); + expect(failures).toEqual([error]); + expect(() => transport.send("again")).toThrow("Transport closed"); + expect(failures).toEqual([error]); + }, +); diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index e32f0e3d52..83087d0a75 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -8,7 +8,7 @@ const HEADER = 4; /** Connection-scoped ordered framing, below serialization and routing, above encryption. */ export class ChunkedTransport { enabled = false; - private parts: Uint8Array[] | undefined; + private assembly: Uint8Array | undefined; private size = 0; private kind = 0; private closed = false; @@ -21,6 +21,7 @@ export class ChunkedTransport { frameBytes: number; send: (frame: Frame, waitForDrain: boolean) => void | Promise; deliver: (frame: Frame) => void; + onFailure?: (error: unknown) => void; }, ) { if (io.frameBytes <= HEADER) throw new Error("Invalid transport frame budget"); @@ -34,12 +35,8 @@ export class ChunkedTransport { 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; + send(frame: Frame, options: { waitForDrain?: boolean; byteLength?: number } = {}): Promise { + const byteLength = options.byteLength ?? frameByteLength(frame); 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"); @@ -47,12 +44,18 @@ export class ChunkedTransport { if (this.closed) throw new Error("Transport closed"); this.activeBytes = byteLength; return chunk - ? this.sendChunks(bytes!, typeof frame === "string") - : this.io.send(frame, waitForDrain); + ? this.sendChunks( + typeof frame === "string" ? new TextEncoder().encode(frame) : frame, + typeof frame === "string", + ) + : this.io.send(frame, options.waitForDrain === true); }; if (this.queuedBytes + byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) { - this.close(); - return Promise.reject(new Error("Logical output queue exceeded its byte limit")); + const error = new Error("Logical output queue exceeded its byte limit"); + this.fail(error); + const rejected = Promise.reject(error); + void rejected.catch(() => {}); + return rejected; } this.queuedBytes += byteLength; let result: void | Promise; @@ -61,6 +64,7 @@ export class ChunkedTransport { } catch (error) { this.queuedBytes -= byteLength; this.activeBytes = 0; + this.fail(error); throw error; } if (!result) { @@ -70,7 +74,7 @@ export class ChunkedTransport { } const pending = Promise.resolve(result) .catch((error) => { - this.close(); + this.fail(error); throw error; }) .finally(() => { @@ -109,7 +113,7 @@ export class ChunkedTransport { if (this.closed) return; try { if (typeof frame === "string" || frame[0] !== OPCODE) { - if (this.parts) throw new Error("Interrupted chunked message"); + if (this.assembly) throw new Error("Interrupted chunked message"); this.io.deliver(frame); return; } @@ -123,27 +127,32 @@ export class ChunkedTransport { throw new Error("Invalid chunk header"); const start = (frame[3]! & 1) !== 0; const end = (frame[3]! & 2) !== 0; - if (start === !!this.parts) throw new Error("Invalid chunk sequence"); + if (start === !!this.assembly) throw new Error("Invalid chunk sequence"); if (start) { - this.parts = []; + this.assembly = new Uint8Array(0); this.kind = frame[2]!; this.size = 0; } if (frame[2] !== this.kind) throw new Error("Chunk kind changed"); this.size += frame.length - HEADER; if (this.size > MAX_WEBSOCKET_MESSAGE_BYTES) throw new Error("Logical message too large"); - this.parts!.push(frame.slice(HEADER)); - if (!end) return; - const data = new Uint8Array(this.size); - let offset = 0; - for (const part of this.parts!) { - data.set(part, offset); - offset += part.length; + const payload = frame.subarray(HEADER); + if (this.assembly!.byteLength < this.size) { + const capacity = Math.min( + MAX_WEBSOCKET_MESSAGE_BYTES, + Math.max(this.size, this.assembly!.byteLength * 2, 16384), + ); + const grown = new Uint8Array(capacity); + grown.set(this.assembly!.subarray(0, this.size - payload.byteLength)); + this.assembly = grown; } - this.parts = undefined; + this.assembly!.set(payload, this.size - payload.byteLength); + if (!end) return; + const data = this.assembly!.subarray(0, this.size); + this.assembly = undefined; this.size = 0; this.io.deliver( - this.kind === 0 ? new TextDecoder("utf-8", { fatal: true }).decode(data) : data, + this.kind === 0 ? new TextDecoder("utf-8", { fatal: true }).decode(data) : data.slice(), ); } catch (error) { this.close(); @@ -151,10 +160,36 @@ export class ChunkedTransport { } } + private fail(error: unknown): void { + if (this.closed) return; + this.close(); + this.io.onFailure?.(error); + } + close(): void { this.closed = true; this.enabled = false; - this.parts = undefined; + this.assembly = undefined; this.size = 0; } } + +function frameByteLength(frame: Frame): number { + if (typeof frame !== "string") return frame.byteLength; + let bytes = 0; + for (let index = 0; index < frame.length; index++) { + const code = frame.charCodeAt(index); + if (code <= 0x7f) bytes++; + else if (code <= 0x7ff) bytes += 2; + else if ( + code >= 0xd800 && + code <= 0xdbff && + frame.charCodeAt(index + 1) >= 0xdc00 && + frame.charCodeAt(index + 1) <= 0xdfff + ) { + bytes += 4; + index++; + } else bytes += 3; + } + return bytes; +} diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 0b9cecd9fa..7c246a4ba8 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -1275,7 +1275,9 @@ export class VoiceAssistantWebSocketServer { return { sockets: writableSockets, payload, bytes: outboundFrameByteLength(payload) }; } - private negotiateChunking(ws: WebSocketLike, hello: WSHelloMessage): void { + private sendHelloResponse(ws: WebSocketLike, session: Session, hello: WSHelloMessage): void { + // The peer learns our support from this ordinary response before receiving chunks. + this.sendToClient(ws, this.createServerInfoMessage(session)); this.chunkedSocket(ws).enabled = hello.capabilities?.[CLIENT_CAPS.binaryChunkingV1] === true; } @@ -1286,6 +1288,12 @@ export class VoiceAssistantWebSocketServer { if (!transport) { transport = new ChunkedTransport({ frameBytes: ws.maxFrameBytes ?? 32 * 1024 * 1024, + onFailure: (err) => + this.closePhysicalSocket({ + ws, + logMessage: "Closing physical WebSocket after outbound framing failure", + logFields: { err }, + }), send: (frame, waitForDrain) => { const params = { socket: ws, @@ -1325,7 +1333,7 @@ export class VoiceAssistantWebSocketServer { : 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); + return transport.send(frame, { waitForDrain, byteLength: bytes }); } private async sendMessageToClientAndWait( @@ -1379,7 +1387,6 @@ export class VoiceAssistantWebSocketServer { try { 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"); @@ -1733,7 +1740,6 @@ export class VoiceAssistantWebSocketServer { pending: PendingConnection; }): Promise { const { ws, message, pending } = params; - this.negotiateChunking(ws, message); if (message.protocolVersion < 1) { this.clearPendingConnection(ws); @@ -1816,7 +1822,7 @@ export class VoiceAssistantWebSocketServer { this.externalSessionsByKey.set(sessionKey, connection); } pending.identity.sessionId = connection.session.getSessionId(); - this.sendToClient(ws, this.createServerInfoMessage(connection.session)); + this.sendHelloResponse(ws, connection.session, message); connection.connectionLogger.info( { ...toConnectionLogFields(pending.identity), @@ -1952,7 +1958,7 @@ export class VoiceAssistantWebSocketServer { existing.sockets.add(ws); this.sessions.set(ws, existing); pending.identity.sessionId = existing.session.getSessionId(); - this.sendToClient(ws, this.createServerInfoMessage(existing.session)); + this.sendHelloResponse(ws, existing.session, message); pending.connectionLogger.info( { ...toConnectionLogFields(pending.identity), diff --git a/packages/server/src/server/websocket/physical-socket.test.ts b/packages/server/src/server/websocket/physical-socket.test.ts index 315f6bd2eb..6dce3e151d 100644 --- a/packages/server/src/server/websocket/physical-socket.test.ts +++ b/packages/server/src/server/websocket/physical-socket.test.ts @@ -170,3 +170,46 @@ test("ordinary frames expose a paused real socket's backlog and terminate at the await new Promise((resolve) => server.close(() => resolve())); } }); + +test("an awaited framing failure terminates its real socket immediately", 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; + const transport = new ChunkedTransport({ + frameBytes: 32 * 1024 * 1024, + deliver: () => {}, + onFailure: () => socket.terminate(), + send: (frame) => + sendBoundedPhysicalFrameAndWait({ + socket, + frame, + onHighWater: () => socket.terminate(), + }).then(() => undefined), + }); + transport.enabled = true; + try { + client.pause(); + tcp.cork(); + const closed = once(socket, "close"); + const first = transport.send(new Uint8Array(100 * 1024 * 1024), { waitForDrain: true }); + const overflow = transport.send(new Uint8Array([1]), { waitForDrain: true }); + const settled = Promise.allSettled([first, overflow]); + await closed; + expect(socket.readyState).toBe(WebSocket.CLOSED); + expect((await settled).map((result) => result.status)).toEqual(["rejected", "rejected"]); + } finally { + transport.close(); + client.terminate(); + socket.terminate(); + await new Promise((resolve) => server.close(() => resolve())); + } +});