diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts index ad388eab0f..dc264c4f46 100644 --- a/packages/client/src/daemon-client.test.ts +++ b/packages/client/src/daemon-client.test.ts @@ -8541,3 +8541,69 @@ test.each([false, true])("routed sends wait for origin drain (binary=%s)", async await write; expect(finished).toBe(true); }); + +test("a queued admission failure rejects its RPC without closing the client", async () => { + const mock = createMockTransport(); + const send = mock.transport.send; + let block = false; + let drain!: () => void; + mock.transport.send = (frame) => { + if (!block) return send(frame); + if (typeof frame !== "string") + return new Promise((resolve) => { + drain = resolve; + }); + throw new Error("capacity"); + }; + const client = new DaemonClient({ + url: "ws://test", + clientId: "queued-admission", + transportFactory: () => mock.transport, + reconnect: { enabled: false }, + }); + clients.push(client); + const connecting = client.connect(); + mock.triggerOpen(); + await connecting; + block = true; + const first = client.sendHostFrame("route", new Uint8Array([1])); + const ping = client.ping({ timeoutMs: 1000 }); + void ping.catch(() => {}); + drain(); + await first; + await expect(ping).rejects.toThrow("capacity"); + expect(client.isConnected).toBe(true); +}); + +test("RPC deadlines remain active while a preceding frame waits for drain", async () => { + const mock = createMockTransport(); + const send = mock.transport.send; + let blocked = false; + let drain!: () => void; + mock.transport.send = (frame) => { + if (!blocked) return send(frame); + if (typeof frame !== "string") + return new Promise((resolve) => { + drain = resolve; + }); + }; + const client = new DaemonClient({ + url: "ws://test", + clientId: "drain-deadline", + transportFactory: () => mock.transport, + reconnect: { enabled: false }, + }); + clients.push(client); + const connecting = client.connect(); + mock.triggerOpen(); + await connecting; + blocked = true; + const first = client.sendHostFrame("route", new Uint8Array([1])); + try { + await expect(client.ping({ timeoutMs: 20 })).rejects.toThrow("Timeout waiting for message"); + expect(client.isConnected).toBe(true); + } finally { + drain(); + await first; + } +}); diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 90ea2602df..9a2965130b 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1267,7 +1267,7 @@ function toReasonCode(reason: string | null | undefined): string | null { } interface PendingSend { - send: () => void; + send: () => void | Promise; resolve: () => void; reject: (error: Error) => void; timeoutHandle: ReturnType; @@ -2166,7 +2166,7 @@ export class DaemonClient { private sendSessionMessageOrThrow(message: SessionInboundMessage): Promise { return this.sendWhenConnected(() => { const payload = SessionInboundMessageSchema.parse(message); - this.sendSessionPayload(payload); + return this.sendSessionPayload(payload); }); } @@ -2175,13 +2175,12 @@ export class DaemonClient { return this.sendWhenConnected(() => undefined); } - private sendWhenConnected(send: () => void): Promise { + private sendWhenConnected(send: () => void | Promise): Promise { const status = this.connectionState.status; // If connected, send immediately if (this.transport && status === "connected") { - send(); - return Promise.resolve(); + return Promise.resolve(send()); } // If connecting, queue the message to be sent once connected @@ -2220,8 +2219,7 @@ export class DaemonClient { clearTimeout(pending.timeoutHandle); try { if (this.transport && this.connectionState.status === "connected") { - pending.send(); - pending.resolve(); + void Promise.resolve(pending.send()).then(pending.resolve, pending.reject); } else { pending.reject(new DaemonConnectionError("Connection lost before message could be sent")); } @@ -2277,19 +2275,17 @@ export class DaemonClient { { ...params.options, requestId: params.requestId }, ); + // The response deadline still applies while the transport waits for drain. + void promise.catch(() => undefined); try { - await this.sendSessionMessageOrThrow(wire.message); + const sent = this.sendSessionMessageOrThrow(wire.message); + const result = await Promise.race([promise, sent.then(() => promise)]); + if (result.kind === "error") throw result.error; + return result.value; } catch (error) { const err = error instanceof Error ? error : new Error(String(error)); cancel(err); - void promise.catch(() => undefined); throw err; - } - - try { - const result = await promise; - if (result.kind === "error") throw result.error; - return result.value; } finally { await wire.finish(); } diff --git a/packages/client/src/node/connector.test.ts b/packages/client/src/node/connector.test.ts index a1eba1ef15..882a2f5706 100644 --- a/packages/client/src/node/connector.test.ts +++ b/packages/client/src/node/connector.test.ts @@ -67,12 +67,13 @@ test("direct connector preserves 25 MiB binary frames in both directions", async await ready; const bytes = new Uint8Array(25 * 1024 * 1024).fill(173); const echoed = once(socket, "message"); - transport.send(bytes); + const sending = transport.send(bytes); const [sent, binary] = await echoed; expect(binary).toBe(true); expect(Buffer.from(sent as Uint8Array).equals(bytes)).toBe(true); socket.send(bytes); expect(Buffer.from(await received).equals(bytes)).toBe(true); + await sending; }); import { @@ -145,9 +146,10 @@ test("encrypted connector rejects overload synchronously and preserves the physi await new Promise((resolve) => setImmediate(resolve)); expect(errors).toEqual([]); expect(closed).toBe(false); - transport.send("sibling control"); + const sibling = transport.send("sibling control"); peer.resume(); await drained; + await sibling; await expect.poll(() => received.includes("sibling control")).toBe(true); expect(received.filter((data) => data instanceof ArrayBuffer).length).toBeGreaterThan(0); expect( diff --git a/packages/protocol/src/chunked-transport.test.ts b/packages/protocol/src/chunked-transport.test.ts index 4e758b05ee..18ea469521 100644 --- a/packages/protocol/src/chunked-transport.test.ts +++ b/packages/protocol/src/chunked-transport.test.ts @@ -276,27 +276,75 @@ it.each(["ASCII", "中", "🌍", "\ud800", "\udc00", "\ud800x"])( }, ); +it("notifies the owner once when an accepted writer fails", async () => { + const failures: unknown[] = []; + const error = new Error("write failed"); + const transport = new ChunkedTransport({ + frameBytes: 1024, + deliver: () => {}, + onFailure: (failure) => { + failures.push(failure); + }, + send: () => Promise.reject(error), + }); + await expect(transport.send("hello", { waitForDrain: true })).rejects.toThrow("write failed"); + expect(failures).toEqual([error]); + expect(() => transport.send("again")).toThrow("Transport closed"); + expect(failures).toEqual([error]); +}); + it.each([false, true])( - "notifies the owner once when a writer fails (async=%s)", - async (asynchronous) => { + "a pre-write rejection preserves the sender (chunked=%s)", + async (chunked) => { const failures: unknown[] = []; - const error = new Error("write failed"); + let reject = true; const transport = new ChunkedTransport({ - frameBytes: 1024, + frameBytes: 8, deliver: () => {}, - onFailure: (failure) => { - failures.push(failure); + onFailure: (error) => { + failures.push(error); }, send: () => { - if (asynchronous) return Promise.reject(error); - throw error; + if (reject) throw new Error("capacity"); }, }); - 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]); + transport.enabled = chunked; + expect(() => transport.send("abcdefghij")).toThrow("capacity"); + expect(failures).toEqual([]); + reject = false; + await transport.send("sibling"); }, ); + +it("a queued pre-write rejection does not poison the next logical send", async () => { + let drain!: () => void; + const failures: unknown[] = []; + const delivered: string[] = []; + const transport = new ChunkedTransport({ + frameBytes: 1024, + deliver: () => {}, + onFailure: (error) => { + failures.push(error); + }, + send: (frame) => { + if (frame === "blocked") + return new Promise((resolve) => { + drain = resolve; + }); + if (frame === "reject") throw new Error("capacity"); + delivered.push(frame as string); + }, + }); + const first = transport.send("blocked"); + const rejected = transport.send("reject"); + const sibling = transport.send("sibling"); + const results = Promise.allSettled([first, rejected, sibling]); + drain(); + expect((await results).map((result) => result.status)).toEqual([ + "fulfilled", + "rejected", + "fulfilled", + ]); + expect(delivered).toEqual(["sibling"]); + expect(failures).toEqual([]); +}); diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index 83087d0a75..536c7f1f49 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -43,12 +43,19 @@ export class ChunkedTransport { const write = () => { if (this.closed) throw new Error("Transport closed"); this.activeBytes = byteLength; - return chunk + // A synchronous rejection did not admit a frame. Only accepted writes can + // fail the connection; queued admission failures still belong to their sender. + const result = chunk ? this.sendChunks( typeof frame === "string" ? new TextEncoder().encode(frame) : frame, typeof frame === "string", ) : this.io.send(frame, options.waitForDrain === true); + if (!result) return; + return result.catch((error) => { + this.fail(error); + throw error; + }); }; if (this.queuedBytes + byteLength > MAX_WEBSOCKET_MESSAGE_BYTES) { const error = new Error("Logical output queue exceeded its byte limit"); @@ -64,7 +71,6 @@ export class ChunkedTransport { } catch (error) { this.queuedBytes -= byteLength; this.activeBytes = 0; - this.fail(error); throw error; } if (!result) { @@ -72,29 +78,25 @@ export class ChunkedTransport { this.activeBytes = 0; return Promise.resolve(); } - const pending = Promise.resolve(result) - .catch((error) => { - this.fail(error); - throw error; - }) - .finally(() => { - this.queuedBytes -= byteLength; - this.activeBytes = 0; - }); - this.pending = pending; - void pending.then( - () => { - if (this.pending === pending) this.pending = undefined; - return undefined; - }, - () => {}, + const completion = Promise.resolve(result).finally(() => { + this.queuedBytes -= byteLength; + this.activeBytes = 0; + }); + const tail = completion.then( + () => undefined, + () => undefined, ); - return pending; + this.pending = tail; + void tail.then(() => { + if (this.pending === tail) this.pending = undefined; + return undefined; + }); + return completion; } - private async sendChunks(bytes: Uint8Array, text: boolean): Promise { + private sendChunks(bytes: Uint8Array, text: boolean): Promise { const capacity = this.io.frameBytes - HEADER; - for (let offset = 0; offset < bytes.byteLength; offset += capacity) { + const writeChunk = (offset: number) => { if (this.closed) throw new Error("Transport closed"); const end = Math.min(bytes.byteLength, offset + capacity); const wire = new Uint8Array(HEADER + end - offset); @@ -105,8 +107,14 @@ export class ChunkedTransport { (offset === 0 ? 1 : 0) | (end === bytes.byteLength ? 2 : 0), ]); wire.set(bytes.subarray(offset, end), HEADER); - await this.io.send(wire, true); - } + return this.io.send(wire, true); + }; + const first = writeChunk(0); + return (async () => { + await first; + for (let offset = capacity; offset < bytes.byteLength; offset += capacity) + await writeChunk(offset); + })(); } receive(frame: Frame): void { diff --git a/packages/server/src/server/connected-hosts/hosts.e2e.test.ts b/packages/server/src/server/connected-hosts/hosts.e2e.test.ts index a587d8c04f..4c6d79aced 100644 --- a/packages/server/src/server/connected-hosts/hosts.e2e.test.ts +++ b/packages/server/src/server/connected-hosts/hosts.e2e.test.ts @@ -2129,3 +2129,42 @@ test( await remote.dispose(); }, ); + +test("a pre-write host budget rejection leaves sibling routes connected", async () => { + const { HostTransportBudget } = await import("@getpaseo/client/node"); + const b = await daemon(); + const origin = createHostClient( + `127.0.0.1:${b.port}`, + { clientId: "budget-origin", reconnect: { enabled: false } }, + new HostTransportBudget(1024 * 1024), + ); + cleanup.push(() => origin.close()); + await origin.connect(); + const serverId = origin.getLastServerInfoMessage()!.serverId; + const routeIds = [ + await origin.openHostLink("budget-first"), + await origin.openHostLink("budget-sibling"), + ]; + const peers = routeIds.map( + (id, index) => + new DaemonClient({ + url: "ws://host-link/ws", + clientId: `budget-child-${index}`, + reconnect: { enabled: false }, + transportFactory: createHostChannelTransport(origin, serverId, `budget-${index}`, id), + }), + ); + for (const peer of peers) { + cleanup.push(() => peer.close()); + await peer.connect(); + } + await expect(async () => + origin.sendHostFrame(routeIds[0]!, new Uint8Array(2 * 1024 * 1024)), + ).rejects.toMatchObject({ code: "resource_limit" }); + await peers[1]!.ping({ timeoutMs: 2000 }); + expect([origin.isConnected, ...peers.map((peer) => peer.isConnected)]).toEqual([ + true, + true, + true, + ]); +}, 30000); diff --git a/packages/server/src/server/connected-hosts/index.ts b/packages/server/src/server/connected-hosts/index.ts index 96d6ac35d2..15b1e16b97 100644 --- a/packages/server/src/server/connected-hosts/index.ts +++ b/packages/server/src/server/connected-hosts/index.ts @@ -510,10 +510,13 @@ export class ConnectedHosts { return; } if (!verified) throw new HostError("host_offline", "Host identity has not been verified"); - try { - transport.send(data); - } catch { + const rejectWrite = () => route.close("resource_limit", "Host buffer limit reached or transport unavailable"); + try { + const write = transport.send(data); + if (write) void write.catch(rejectWrite); + } catch { + rejectWrite(); } }, }; diff --git a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts index 32d3a1f36c..0670bd2898 100644 --- a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts @@ -593,7 +593,7 @@ async function waitForCapturedLog( } }, 90000); - test("oversized attachments fail without disconnecting or resubmitting over the relay", async () => { + test("large messages cross the relay in both directions and logical oversize rejects without disconnecting", async () => { await startRelay({ useLocalRelay: true }); const daemon = await createTestPaseoDaemon({ relayEnabled: true, @@ -623,42 +623,72 @@ async function waitForCapturedLog( const statuses: string[] = []; client.subscribeConnectionStatus((state) => statuses.push(state.status)); for (const imageMiB of [30, 61]) { - // Attachments are base64 in the JSON, then E2EE base64 expands the frame again. const images = [ { data: "A".repeat(4 * Math.ceil((imageMiB * 1024 * 1024) / 3)), mimeType: "image/png" }, ]; - await expect( - withTimeout( - client.createWorkspace({ - source: { kind: "directory", path: daemon.staticDir }, - agent: { - provider: "claude", - cwd: daemon.staticDir, - initialPrompt: "Describe these screens.", - images, - }, - }), - 10000, - "Workspace creation did not settle", - ), - ).rejects.toThrow(/too large.*32\.0 MiB/); - expect((await client.fetchAgents()).entries).toEqual([]); + await withTimeout( + client.createWorkspace({ + source: { kind: "directory", path: daemon.staticDir }, + agent: { + provider: "claude", + cwd: daemon.staticDir, + initialPrompt: "Describe these screens.", + images, + }, + }), + 30000, + "Chunked workspace creation did not settle", + ); } + expect((await client.fetchAgents()).entries).toHaveLength(2); + await expect( + client.createWorkspace({ + source: { kind: "directory", path: daemon.staticDir }, + agent: { + provider: "claude", + cwd: daemon.staticDir, + initialPrompt: "Describe these screens.", + images: [{ data: "A".repeat(101 * 1024 * 1024), mimeType: "image/png" }], + }, + }), + ).rejects.toThrow(/too large.*100\.0 MiB/); + expect((await client.fetchAgents()).entries).toHaveLength(2); const agent = await client.createAgent({ provider: "claude", cwd: daemon.staticDir }); + await client.sendMessage(agent.id, "Describe these screens.", { + images: [{ data: "A".repeat(40 * 1024 * 1024), mimeType: "image/png" }], + }); await expect( client.sendMessage(agent.id, "Describe these screens.", { - images: [{ data: "A".repeat(40 * 1024 * 1024), mimeType: "image/png" }], + images: [{ data: "A".repeat(101 * 1024 * 1024), mimeType: "image/png" }], }), - ).rejects.toThrow(/too large.*32\.0 MiB/); + ).rejects.toThrow(/too large.*100\.0 MiB/); await client.sendMessage(agent.id, "A small follow-up still works."); - expect((await client.fetchAgents()).entries).toHaveLength(1); + expect((await client.fetchAgents()).entries).toHaveLength(3); + // Respect the per-item content budget while exceeding the relay frame limit. + const text = "x".repeat(240 * 1024); + for (let row = 0; row < 220; row++) + await daemon.daemon.agentManager.appendTimelineItem(agent.id, { + type: "user_message", + text, + }); + const timeline = await client.fetchAgentTimeline(agent.id, { + projection: "canonical", + limit: 0, + timeout: 30000, + }); + const delivered = timeline.entries.filter( + (row) => row.item.type === "user_message" && row.item.text === text, + ); + expect(delivered).toHaveLength(220); + expect(delivered.length * Buffer.byteLength(text)).toBeGreaterThan(50 * 1024 * 1024); + await client.ping(); expect(statuses).toEqual(["connected"]); } finally { await client?.close(); await daemon.close(); await stopRelay(); } - }, 60000); + }, 120000); test("daemon connects to relay and client ping/pong works through relay", async () => { process.env.PASEO_PRIMARY_LAN_IP = "192.168.1.12"; diff --git a/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts index 5131456ee9..b708ebb149 100644 --- a/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts +++ b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts @@ -90,7 +90,11 @@ test( TEST_TIMEOUT_MS, ); -async function connectSocket(port: number, clientId: string): Promise { +async function connectSocket( + port: number, + clientId: string, + capabilities?: Record, +): Promise { const socket = new WebSocket(`ws://127.0.0.1:${port}/ws`); await new Promise((resolve, reject) => { socket.once("open", resolve); @@ -98,7 +102,7 @@ async function connectSocket(port: number, clientId: string): Promise }); await sendAndWait( socket, - { type: "hello", clientId, clientType: "browser", protocolVersion: 1 }, + { type: "hello", clientId, clientType: "browser", protocolVersion: 1, capabilities }, (message) => message.type === "session" && message.message.type === "status" && @@ -166,3 +170,88 @@ function sendAndWait( socket.send(JSON.stringify(message)); return response; } + +test("an awaited timeline framing overflow closes the actual daemon socket", async () => { + const { createHostClient } = await import("@getpaseo/client/node"); + const { createPaseoApi } = await import("@getpaseo/client"); + const { default: pino } = await import("pino"); + const { Writable } = await import("node:stream"); + const messages: string[] = []; + const logger = pino( + { level: "warn" }, + new Writable({ + write(chunk, _encoding, done) { + messages.push(JSON.parse(chunk.toString()).msg); + done(); + }, + }), + ); + daemon = await createTestPaseoDaemon({ logger }); + const client = createHostClient(`127.0.0.1:${daemon.port}`, { + clientId: "framing-seed", + reconnect: { enabled: false }, + }); + await client.connect(); + const api = createPaseoApi(client); + try { + const workspace = await api.workspaces.create({ + source: { kind: "directory", path: daemon.staticDir }, + }); + const agents = []; + const text = "x".repeat(240 * 1024); + for (const [index, rows] of [350, 90].entries()) { + const agent = await workspace.agents.create({ + config: { provider: "claude/claude-test", modeId: "full-access" }, + title: `Framing overflow ${index}`, + }); + agents.push(agent.id); + for (let row = 0; row < rows; row++) + await daemon.daemon.agentManager.appendTimelineItem(agent.id, { + type: "user_message", + text, + }); + } + const source = await connectSocket(daemon.port, "framing-overflow", { + binary_chunking_v1: true, + }); + sockets.push(source); + const closed = new Promise((resolve) => source.once("close", () => resolve())); + const firstChunk = new Promise((resolve) => { + const listener = (data: RawData, binary: boolean) => { + if (!binary || (data as Buffer)[0] !== 0x31) return; + source.off("message", listener); + source.pause(); + resolve(); + }; + source.on("message", listener); + }); + const fetch = (agentId: string, requestId: string) => + source.send( + JSON.stringify({ + type: "session", + message: { + type: "fetch_agent_timeline_request", + requestId, + agentId, + direction: "tail", + projection: "canonical", + limit: 0, + }, + }), + ); + fetch(agents[0]!, "first"); + await firstChunk; + fetch(agents[1]!, "overflow"); + await expect + .poll(() => messages.includes("Closing physical WebSocket after outbound framing failure"), { + timeout: 15000, + }) + .toBe(true); + expect(messages).not.toContain("Closing physical WebSocket at outbound high-water mark"); + source.resume(); + await closed; + } finally { + await api.dispose(); + await client.close(); + } +}, 60000);