diff --git a/packages/client/src/daemon-client-transport-types.ts b/packages/client/src/daemon-client-transport-types.ts index 64f84cf8c1..e632e1e3bd 100644 --- a/packages/client/src/daemon-client-transport-types.ts +++ b/packages/client/src/daemon-client-transport-types.ts @@ -4,7 +4,7 @@ export interface DaemonTransport { maxFrameBytes?: number; send: (data: string | Uint8Array | ArrayBuffer) => void | Promise; close: (code?: number, reason?: string) => void; - onMessage: (handler: (data: unknown, isBinary: boolean) => void) => () => void; + onMessage: (handler: (data: unknown, isBinary: boolean) => void | Promise) => () => void; onOpen: (handler: () => void) => () => void; onClose: (handler: (event?: unknown) => void) => () => void; onError: (handler: (event?: unknown) => void) => () => void; diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 9d77a5a91c..a2521a0fb9 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1923,18 +1923,34 @@ export class DaemonClient { }); } } - async openHostLink(channelId: string): Promise { + openHostLink(channelId: string): Promise; + openHostLink( + channelId: string, + options: { flowControl: true }, + ): Promise<{ routeId: string; flowControl: boolean }>; + async openHostLink( + channelId: string, + options?: { flowControl: true }, + ): Promise { this.requireConnectedHosts(); if (this.getLastServerInfoMessage()?.features?.hostMultiplexing !== true) throw Object.assign(new Error("Update the destination daemon"), { code: "unsupported_capability", }); - return ( - await this.sendCorrelatedSessionRequest({ - message: { type: "hosts.link.open.request", channelId }, - responseType: "hosts.link.open.response", - }) - ).routeId; + const response = await this.sendCorrelatedSessionRequest({ + message: { + type: "hosts.link.open.request", + channelId, + ...(options?.flowControl ? { flowControl: true } : {}), + }, + responseType: "hosts.link.open.response", + }); + return options?.flowControl + ? { routeId: response.routeId, flowControl: response.flowControl === true } + : response.routeId; + } + acknowledgeHostLink(routeId: string, sequence: number): Promise { + return this.sendSessionMessageStrict({ type: "hosts.link.ack.request", routeId, sequence }); } sendHostFrame(routeId: string, data: string | Uint8Array): Promise { if (typeof data === "string") { diff --git a/packages/client/src/hosts/index.ts b/packages/client/src/hosts/index.ts index aa5c910158..150cf48071 100644 --- a/packages/client/src/hosts/index.ts +++ b/packages/client/src/hosts/index.ts @@ -26,11 +26,13 @@ export function createHostChannelTransport( return () => { const opens = new Set<() => void>(); const closes = new Set<(event?: unknown) => void>(); - const messages = new Set<(data: unknown, binary: boolean) => void>(); + const messages = new Set<(data: unknown, binary: boolean) => void | Promise>(); const errors = new Set<(error: unknown) => void>(); let routeId: string | null = null; let closed = false; const cleanups: (() => void)[] = []; + let flowControl = false; + let receivedFrames = 0; const close = (error?: unknown) => { if (closed) return; closed = true; @@ -45,10 +47,28 @@ export function createHostChannelTransport( if (error) for (const handler of errors) handler(error); for (const handler of closes) handler(error); }; + const deliver = (data: unknown, binary: boolean) => { + if (closed) return; + const sequence = ++receivedFrames; + try { + const deliveries = [...messages].map(async (handler) => { + await handler(data, binary); + }); + void Promise.all(deliveries) + .then(async () => { + if (flowControl && !closed && routeId) + await origin.acknowledgeHostLink(routeId, sequence); + return undefined; + }) + .catch(close); + } catch (error) { + close(error); + } + }; cleanups.push( origin.subscribeRawMessages((message) => { if (message.type === "hosts.route.frame" && message.payload.subscriptionId === routeId) { - for (const handler of messages) handler(message.payload.data, false); + deliver(message.payload.data, false); } if (message.type === "hosts.route.closed" && message.payload.subscriptionId === routeId) close( @@ -61,7 +81,7 @@ export function createHostChannelTransport( ); cleanups.push( origin.subscribeHostFrames((frame) => { - if (frame.routeId === routeId) for (const handler of messages) handler(frame.payload, true); + if (frame.routeId === routeId) deliver(frame.payload, true); }), ); cleanups.push( @@ -71,7 +91,15 @@ export function createHostChannelTransport( ); let opening: Promise; if (first) opening = Promise.resolve(first); - else opening = link ? origin.openHostLink(scopeId) : origin.openHostRoute(serverId, scopeId); + else if (link) { + opening = origin.openHostLink(scopeId, { flowControl: true }).then((opened) => { + // COMPAT(hostLinkFlowControl): added after v0.15.3 (#9109); remove after + // 2027-05-08 once the daemon floor supports forwarding acknowledgments. + // Older destinations omit flowControl and receive no acknowledgment messages. + flowControl = opened.flowControl; + return opened.routeId; + }); + } else opening = origin.openHostRoute(serverId, scopeId); first = null; void opening .then((id) => { diff --git a/packages/protocol/src/chunked-transport.test.ts b/packages/protocol/src/chunked-transport.test.ts index 3979b48885..caca79c0dc 100644 --- a/packages/protocol/src/chunked-transport.test.ts +++ b/packages/protocol/src/chunked-transport.test.ts @@ -390,3 +390,65 @@ it("unpaced transports retain ordinary outbound admission and accept inbound chu expect(delivered).toEqual(["A"]); await transport.send("small"); }); + +it("fills an ordered writer's window while keeping chunk sequences contiguous", async () => { + const delivered: (string | Uint8Array)[] = []; + const receiver = new ChunkedTransport({ + frameBytes: 8, + send() {}, + deliver: (frame) => { + delivered.push(frame); + }, + }); + receiver.enabled = true; + const drain = Promise.withResolvers(); + const frames: (string | Uint8Array)[] = []; + const sender = new ChunkedTransport({ + frameBytes: 8, + supportsConcurrentWrites: true, + send: (frame) => { + frames.push(frame); + receiver.receive(frame); + return drain.promise; + }, + deliver() {}, + }); + sender.enabled = true; + const first = sender.send("one"); + const second = sender.send("two"); + expect(delivered).toEqual(["one", "two"]); + expect(sender.queuedByteLength).toBe(0); + const large = sender.send("large message"); + const after = sender.send("after"); + expect(frames).toHaveLength(3); + expect(sender.queuedByteLength).toBe(5); + drain.resolve(); + await Promise.all([first, second, large, after]); + expect(delivered).toEqual(["one", "two", "large message", "after"]); + expect(sender.queuedByteLength).toBe(0); +}); + +it("accounts for independently draining writes and a rejected admission", async () => { + const firstDrain = Promise.withResolvers(); + const secondDrain = Promise.withResolvers(); + const sender = new ChunkedTransport({ + frameBytes: 8, + supportsConcurrentWrites: true, + send: (frame) => { + if (frame === "reject") throw new Error("capacity"); + return frame === "one" ? firstDrain.promise : secondDrain.promise; + }, + deliver() {}, + }); + const first = sender.send("one"); + const second = sender.send("two"); + expect(() => sender.send("reject")).toThrow("capacity"); + expect(sender.queuedByteLength).toBe(0); + secondDrain.resolve(); + await second; + expect(sender.queuedByteLength).toBe(0); + firstDrain.resolve(); + await first; + expect(sender.queuedByteLength).toBe(0); + await sender.send("usable"); +}); diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index 8faa22cd96..51d8b888a8 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -23,6 +23,8 @@ export class ChunkedTransport { private readonly io: { frameBytes: number; supportsDrain?: boolean; + /** Writes are admitted synchronously in order; promises track only their drain. */ + supportsConcurrentWrites?: boolean; send: (frame: Frame, waitForDrain: boolean) => void | Promise; deliver: (frame: Frame) => void; onActivity?: () => void; @@ -46,9 +48,11 @@ export class ChunkedTransport { const chunk = this.enabled && this.io.supportsDrain !== false && byteLength > this.io.frameBytes; if (this.closed) throw new Error("Transport closed"); + let started = false; const write = () => { if (this.closed) throw new Error("Transport closed"); - this.activeBytes = byteLength; + started = true; + this.activeBytes += byteLength; // 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 @@ -76,27 +80,31 @@ export class ChunkedTransport { result = this.pending ? this.pending.then(write) : write(); } catch (error) { this.queuedBytes -= byteLength; - this.activeBytes = 0; + if (started) this.activeBytes -= byteLength; throw error; } if (!result) { this.queuedBytes -= byteLength; - this.activeBytes = 0; + if (started) this.activeBytes -= byteLength; return Promise.resolve(); } const completion = Promise.resolve(result).finally(() => { this.queuedBytes -= byteLength; - this.activeBytes = 0; - }); - const tail = completion.then( - () => undefined, - () => undefined, - ); - this.pending = tail; - void tail.then(() => { - if (this.pending === tail) this.pending = undefined; - return undefined; + if (started) this.activeBytes -= byteLength; }); + // Ordinary writes already entered an ordered writer and can use its bounded + // window. Chunk sequences remain a barrier so their frames never interleave. + if (chunk || !this.io.supportsConcurrentWrites) { + const tail = completion.then( + () => undefined, + () => undefined, + ); + this.pending = tail; + void tail.then(() => { + if (this.pending === tail) this.pending = undefined; + return undefined; + }); + } return completion; } diff --git a/packages/protocol/src/hosts.ts b/packages/protocol/src/hosts.ts index 7d36a50236..8311d4268c 100644 --- a/packages/protocol/src/hosts.ts +++ b/packages/protocol/src/hosts.ts @@ -116,8 +116,19 @@ export const HostsLinkOpenRequestSchema = z.object({ type: z.literal("hosts.link.open.request"), requestId: z.string(), channelId: z.string().min(1).max(256), + flowControl: z.boolean().optional(), }); export const HostsLinkOpenResponseSchema = z.object({ type: z.literal("hosts.link.open.response"), - payload: z.object({ requestId: z.string(), routeId: z.string() }), + payload: z.object({ + requestId: z.string(), + routeId: z.string(), + flowControl: z.boolean().optional(), + }), +}); + +export const HostsLinkAckSchema = z.object({ + type: z.literal("hosts.link.ack.request"), + routeId: z.string(), + sequence: z.number().int().positive(), }); diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index c7d14f5ce8..19ded2d075 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -5,6 +5,7 @@ import { HostsUpdateRequestSchema, HostsRemoveRequestSchema, HostsLinkOpenRequestSchema, + HostsLinkAckSchema, HostsRouteOpenRequestSchema, HostsRouteFrameRequestSchema, HostsRouteCloseRequestSchema, @@ -3397,6 +3398,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ HostsUpdateRequestSchema, HostsRemoveRequestSchema, HostsLinkOpenRequestSchema, + HostsLinkAckSchema, HostsRouteOpenRequestSchema, HostsRouteFrameRequestSchema, HostsRouteCloseRequestSchema, diff --git a/packages/server/src/server/authorization/operation-permissions.ts b/packages/server/src/server/authorization/operation-permissions.ts index dfe7393ee7..91573e9cc6 100644 --- a/packages/server/src/server/authorization/operation-permissions.ts +++ b/packages/server/src/server/authorization/operation-permissions.ts @@ -11,6 +11,7 @@ const INBOUND_PERMISSION = { "hosts.update.request": "tunnel.manage", "hosts.remove.request": "tunnel.manage", "hosts.link.open.request": null, + "hosts.link.ack.request": null, "hosts.route.open.request": null, "hosts.route.frame.request": null, "hosts.route.close.request": null, diff --git a/packages/server/src/server/connected-hosts/flow-control.ts b/packages/server/src/server/connected-hosts/flow-control.ts new file mode 100644 index 0000000000..eafc52e9e0 --- /dev/null +++ b/packages/server/src/server/connected-hosts/flow-control.ts @@ -0,0 +1,51 @@ +export class HostChannelOutputLimitError extends Error { + constructor() { + super("Host channel output limit reached"); + } +} + +/** Completion belongs to the receiver's forwarding write, not just our physical socket. */ +export class HostChannelFlowControl { + private sequence = 0; + private bytes = 0; + private closed = false; + private readonly pending = new Map< + number, + { + resolve(): void; + reject(error: Error): void; + } + >(); + + constructor(private readonly reserveShared: (bytes: number) => () => void) {} + + async send(bytes: number, write: () => Promise): Promise { + if (this.closed) throw new Error("Host channel closed"); + // Bound both payload retention and metadata for a producer of tiny frames. + if (this.bytes + bytes > 32 * 1024 * 1024 || this.pending.size >= 4096) + throw new HostChannelOutputLimitError(); + const release = this.reserveShared(bytes); + this.bytes += bytes; + const sequence = ++this.sequence; + const acknowledgment = new Promise((resolve, reject) => { + this.pending.set(sequence, { resolve, reject }); + }); + try { + await Promise.all([acknowledgment, write()]); + } finally { + this.pending.delete(sequence); + this.bytes -= bytes; + release(); + } + } + + acknowledge(sequence: number): void { + this.pending.get(sequence)?.resolve(); + } + + close(): void { + this.closed = true; + for (const pending of this.pending.values()) pending.reject(new Error("Host channel closed")); + this.pending.clear(); + } +} 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 416c6adb31..5b618f0670 100644 --- a/packages/server/src/server/connected-hosts/hosts.e2e.test.ts +++ b/packages/server/src/server/connected-hosts/hosts.e2e.test.ts @@ -1141,7 +1141,7 @@ test("source-key registration remains idempotent at directory capacity", async ( }); /** Real WebSocket proxy records the A–B physical wire without replacing either daemon. */ -async function hostWire(b: TestPaseoDaemon) { +async function hostWire(b: TestPaseoDaemon, options: { legacyFlowControl?: boolean } = {}) { const http = createServer(); const wss = new WebSocketServer({ server: http }); const frames: { binary: boolean; bytes: Buffer; direction: "toB" | "fromB" }[] = []; @@ -1154,7 +1154,14 @@ async function hostWire(b: TestPaseoDaemon) { for (const send of pending.splice(0)) send(); }); downstream.on("message", (data, binary) => { - const bytes = Buffer.from(data as Buffer); + let bytes = Buffer.from(data as Buffer); + if (options.legacyFlowControl && !binary) { + const envelope = JSON.parse(bytes.toString()); + if (envelope.message?.type === "hosts.link.open.request") { + delete envelope.message.flowControl; + bytes = Buffer.from(JSON.stringify(envelope)); + } + } frames.push({ binary, bytes, direction: "toB" }); const send = () => { if (upstream.readyState === 1) upstream.send(bytes, { binary }); @@ -2390,6 +2397,169 @@ test( }, ); +test("a host channel waits for its downstream consumer without blocking sibling requests", async () => { + const b = await daemon(); + const destination = await api(b); + const workspace = await destination.paseo.workspaces.create({ + source: { kind: "directory", path: b.staticDir }, + }); + const agent = await workspace.agents.create({ + config: { provider: "claude/claude-test", modeId: "full-access" }, + title: "Flow-controlled history", + }); + const text = "x".repeat(240 * 1024); + for (let index = 0; index < 50; index++) + await b.daemon.agentManager.appendTimelineItem(agent.id, { type: "user_message", text }); + + const channel = createHostChannelTransport( + destination.client, + b.daemon.getServerId()!, + "held-consumer", + null, + true, + )({ url: "ws://host-link/ws" }); + const received = Promise.withResolvers(); + const forward = Promise.withResolvers(); + let binaryFrames = 0; + channel.onMessage((_data, binary) => { + if (!binary) return; + binaryFrames++; + received.resolve(); + return forward.promise; + }); + const consumer = new DaemonClient({ + url: "ws://host-link/ws", + clientId: "held-consumer", + reconnect: { enabled: false }, + transportFactory: () => channel, + }); + const remote = createPaseoApi(consumer); + cleanup.push(async () => { + forward.resolve(); + await remote.dispose(); + await consumer.close(); + }); + await consumer.connect(); + const history = remote.agents.ref(agent.id).timeline.refetch({ limit: 0 }); + void history.catch(() => {}); // Cleanup can close the request after an earlier assertion fails. + await received.promise; + // A round trip on the same physical link must progress while this channel is held. + await destination.client.ping(); + expect(binaryFrames).toBe(1); + forward.resolve(); + const response = await history; + const messages = response.entries.flatMap((row) => + row.item.type === "user_message" ? [row.item.text] : [], + ); + expect(messages).toHaveLength(50); + expect(messages.every((message) => message === text)).toBe(true); + await consumer.ping(); +}); + +test("small host-channel responses fill the window before earlier forwarding completes", async () => { + const b = await daemon(); + const destination = await api(b); + const channel = createHostChannelTransport( + destination.client, + b.daemon.getServerId()!, + "concurrent-consumer", + null, + true, + )({ url: "ws://host-link/ws" }); + const first = Promise.withResolvers(); + const forward = Promise.withResolvers(); + let responses = 0; + channel.onMessage((data, binary) => { + if (binary || JSON.parse(String(data)).message?.type !== "pong") return; + responses++; + first.resolve(); + return forward.promise; + }); + const consumer = new DaemonClient({ + url: "ws://host-link/ws", + clientId: "concurrent-consumer", + reconnect: { enabled: false }, + transportFactory: () => channel, + }); + cleanup.push(async () => { + forward.resolve(); + await consumer.close(); + }); + await consumer.connect(); + const requests = Promise.all(Array.from({ length: 10 }, () => consumer.ping())); + void requests.catch(() => {}); + await first.promise; + await destination.client.ping(); + await destination.client.ping(); + expect(responses).toBe(10); + forward.resolve(); + await requests; +}); + +test("legacy host links receive no acknowledgment messages", async () => { + const b = await daemon(); + const wire = await hostWire(b, { legacyFlowControl: true }); + const origin = createHostClient(wire.connection, { + clientId: "legacy-flow-origin", + reconnect: { enabled: false }, + }); + cleanup.push(() => origin.close()); + await origin.connect(); + const consumer = new DaemonClient({ + url: "ws://host-link/ws", + clientId: "legacy-flow-consumer", + reconnect: { enabled: false }, + transportFactory: createHostChannelTransport( + origin, + b.daemon.getServerId()!, + "legacy", + null, + true, + ), + }); + cleanup.push(() => consumer.close()); + await consumer.connect(); + await consumer.ping(); + await origin.ping(); + const operations = wire.frames + .filter((frame) => !frame.binary && frame.direction === "toB") + .map((frame) => JSON.parse(frame.bytes.toString()).message?.type); + expect(operations).toContain("hosts.link.open.request"); + expect(operations).not.toContain("hosts.link.ack.request"); +}); + +test("mixed asynchronous and synchronous receiver failures close only their host channel", async () => { + const b = await daemon(); + const destination = await api(b); + const channel = createHostChannelTransport( + destination.client, + b.daemon.getServerId()!, + "failing-consumer", + null, + true, + )({ url: "ws://host-link/ws" }); + const failed = Promise.withResolvers(); + channel.onError(failed.resolve); + channel.onMessage(async () => { + throw new Error("asynchronous receiver failure"); + }); + channel.onMessage(() => { + throw new Error("synchronous receiver failure"); + }); + const consumer = new DaemonClient({ + url: "ws://host-link/ws", + clientId: "failing-consumer", + reconnect: { enabled: false }, + transportFactory: () => channel, + }); + cleanup.push(() => consumer.close()); + const connecting = consumer.connect().catch((error: unknown) => error); + expect(await failed.promise).toBeInstanceOf(Error); + await connecting; + expect(consumer.isConnected).toBe(false); + await destination.client.ping(); +}); + test("a pre-write host budget rejection leaves sibling routes connected", async () => { const { HostTransportBudget } = await import("@getpaseo/client/node"); const b = await daemon(); diff --git a/packages/server/src/server/connected-hosts/index.ts b/packages/server/src/server/connected-hosts/index.ts index bfff76239c..526910f5cc 100644 --- a/packages/server/src/server/connected-hosts/index.ts +++ b/packages/server/src/server/connected-hosts/index.ts @@ -384,7 +384,7 @@ export class ConnectedHosts { originSessionId: string; permissions: readonly DaemonPermission[]; signal: AbortSignal; - frame(data: string | Uint8Array): void; + frame(data: string | Uint8Array): void | Promise; ended(code: string, message: string): void; }): Promise<{ send(data: string | Uint8Array): void; close(): void }> { await this.ready; @@ -499,7 +499,7 @@ export class ConnectedHosts { return; } } - input.frame(binary ? (data as Uint8Array) : String(data)); + return input.frame(binary ? (data as Uint8Array) : String(data)); }); await ready; if (closed) throw new HostError("host_offline", "Host closed during connection"); diff --git a/packages/server/src/server/connected-hosts/link.test.ts b/packages/server/src/server/connected-hosts/link.test.ts new file mode 100644 index 0000000000..04cecaae3e --- /dev/null +++ b/packages/server/src/server/connected-hosts/link.test.ts @@ -0,0 +1,118 @@ +import { afterEach, expect, test } from "vitest"; +import { DaemonLinkEndpoint, type VirtualHostSocket } from "./index.js"; +import { SessionDelivery } from "../session/owned-subscriptions/index.js"; +import type { SessionInboundMessage, SessionOutboundMessage } from "@getpaseo/protocol/messages"; + +const cleanup: (() => Promise)[] = []; +afterEach(async () => { + for (const close of cleanup.splice(0).toReversed()) await close(); +}); + +async function link() { + const source = {}; + const messages: SessionOutboundMessage[] = []; + const sockets: VirtualHostSocket[] = []; + const delivery = new SessionDelivery((_source, message) => { + messages.push(message); + }); + delivery.attach(source, true); + const endpoint = new DaemonLinkEndpoint({ + identity: "physical", + permissions: () => ["daemon.read"], + routed: () => false, + bufferedAmount: () => 0, + delivery, + admit: async (socket) => { + sockets.push(socket); + }, + }); + cleanup.push(async () => { + endpoint.close(); + await delivery.close(); + }); + const dispatch = (message: SessionInboundMessage) => + delivery.request(source, message, () => endpoint.dispatch(message)); + const open = async (channelId: string, flowControl = true) => { + await dispatch({ + type: "hosts.link.open.request", + requestId: channelId, + channelId, + flowControl, + }); + const response = messages.find( + (message) => + message.type === "hosts.link.open.response" && message.payload.requestId === channelId, + ); + expect(response).toMatchObject({ + type: "hosts.link.open.response", + payload: flowControl ? { flowControl: true } : {}, + }); + if (response?.type !== "hosts.link.open.response") throw new Error("Missing opened channel"); + return { socket: sockets.at(-1)!, id: response.payload.routeId }; + }; + const ack = (routeId: string, sequence: number) => + dispatch({ type: "hosts.link.ack.request", routeId, sequence }); + const frames = () => + messages + .filter((message) => message.type === "hosts.route.frame") + .map((message) => message.payload.data); + return { open, ack, frames }; +} + +test("a channel waits for forwarding acknowledgments while its sibling keeps sending", async () => { + const host = await link(); + const slow = await host.open("slow"); + const sibling = await host.open("sibling"); + let completed = false; + const first = slow.socket.send("first").then(() => { + completed = true; + return undefined; + }); + let secondCompleted = false; + const second = slow.socket.send("second").then(() => { + secondCompleted = true; + return undefined; + }); + await expect.poll(host.frames).toEqual(["first", "second"]); + expect(completed).toBe(false); + const independent = sibling.socket.send("independent"); + await expect.poll(host.frames).toEqual(["first", "second", "independent"]); + await host.ack(sibling.id, 1); + await independent; + await host.ack(slow.id, 1); + await first; + await expect.poll(host.frames).toEqual(["first", "second", "independent"]); + await host.ack(slow.id, 1); // A duplicate acknowledgment cannot release the next frame. + expect(completed).toBe(true); + expect(secondCompleted).toBe(false); + await host.ack(slow.id, 2); + await second; +}); + +test("closing a blocked channel rejects its pending writes and leaves a sibling usable", async () => { + const host = await link(); + const slow = await host.open("slow"); + const sibling = await host.open("sibling", false); + const first = slow.socket.send("first"); + const second = slow.socket.send("second"); + const results = Promise.allSettled([first, second]); + await expect.poll(host.frames).toEqual(["first", "second"]); + slow.socket.close(); + expect((await results).map((result) => result.status)).toEqual(["rejected", "rejected"]); + await sibling.socket.send("legacy sibling"); + expect(host.frames()).toEqual(["first", "second", "legacy sibling"]); +}); + +test("an unpaced producer exhausting its acknowledgment window closes only its channel", async () => { + const host = await link(); + const slow = await host.open("slow"); + const sibling = await host.open("sibling", false); + const writes: Promise[] = []; + // A producer ignoring send completion must remain bounded even for tiny messages. + for (let index = 0; index < 5000; index++) + writes.push(slow.socket.send("x").catch((error: unknown) => error)); + await Promise.all(writes); + expect(slow.socket.readyState).toBe(3); + await sibling.socket.send("still connected"); + expect(host.frames().at(-1)).toBe("still connected"); +}); diff --git a/packages/server/src/server/connected-hosts/link.ts b/packages/server/src/server/connected-hosts/link.ts index a18c89aeb0..7f917ba2d6 100644 --- a/packages/server/src/server/connected-hosts/link.ts +++ b/packages/server/src/server/connected-hosts/link.ts @@ -1,3 +1,4 @@ +import { HostChannelFlowControl, HostChannelOutputLimitError } from "./flow-control.js"; import type { SessionDelivery } from "../session/owned-subscriptions/index.js"; import { EventEmitter } from "node:events"; import { createHash } from "node:crypto"; @@ -15,6 +16,7 @@ import type { /** A logical socket has no network resources or authority of its own. */ export class VirtualHostSocket extends EventEmitter { readonly maxFrameBytes = MAX_HOST_CHANNEL_FRAME_BYTES; + readonly supportsConcurrentWrites = true; readyState = 1; constructor( private readonly output: (data: string | Uint8Array) => Promise, @@ -61,7 +63,11 @@ export class VirtualHostSocket extends EventEmitter { /** One authenticated physical source owns all its channels and their admission ceiling. */ export class DaemonLinkEndpoint { - private readonly channels = new Map(); + private readonly channels = new Map< + string, + { scopeId: string; socket: VirtualHostSocket; flowControl?: HostChannelFlowControl } + >(); + private pendingOutputBytes = 0; private closed = false; constructor( private readonly parent: { @@ -76,6 +82,7 @@ export class DaemonLinkEndpoint { handles(message: SessionInboundMessage): boolean { return ( message.type === "hosts.link.open.request" || + message.type === "hosts.link.ack.request" || (message.type === "hosts.route.frame.request" && this.channels.has(message.routeId)) || (message.type === "hosts.route.close.request" && ((!!message.routeId && this.channels.has(message.routeId)) || @@ -83,6 +90,10 @@ export class DaemonLinkEndpoint { ); } async dispatch(message: SessionInboundMessage): Promise { + if (message.type === "hosts.link.ack.request") { + this.channels.get(message.routeId)?.flowControl?.acknowledge(message.sequence); + return; + } if (message.type === "hosts.link.open.request") { const error = (code: string, text: string) => this.parent.delivery.reply({ @@ -103,8 +114,19 @@ export class DaemonLinkEndpoint { return; } let socket: VirtualHostSocket; + const flowControl = message.flowControl + ? new HostChannelFlowControl((bytes) => { + if (this.pendingOutputBytes + bytes > 256 * 1024 * 1024) + throw new HostChannelOutputLimitError(); + this.pendingOutputBytes += bytes; + return () => { + this.pendingOutputBytes -= bytes; + }; + }) + : undefined; const owner = this.parent.delivery.begin("host-link", undefined, () => { this.channels.delete(owner.id); + flowControl?.close(); socket?.close(); }); const id = owner.id; @@ -130,15 +152,25 @@ export class DaemonLinkEndpoint { throw new Error("Host channel output limit reached"); } try { - if (output) await owner.emitAndWait(output); - else await owner.emitBinaryAndWait(frame!); + const write = async () => { + if (output) await owner.emitAndWait(output); + else await owner.emitBinaryAndWait(frame!); + }; + if (flowControl) await flowControl.send(bytes, write); + else await write(); } catch (failure) { - socket.close(); + socket.close( + failure instanceof HostChannelOutputLimitError ? 1009 : 1000, + failure instanceof HostChannelOutputLimitError + ? failure.message + : "Host channel closed", + ); throw failure; } }, (code, reason) => { this.channels.delete(id); + flowControl?.close(); let failure = "host_offline"; if (code === 1009) failure = "resource_limit"; if (code === 1008) failure = "permission_denied"; @@ -167,7 +199,7 @@ export class DaemonLinkEndpoint { }); }, ); - this.channels.set(id, { scopeId: message.channelId, socket }); + this.channels.set(id, { scopeId: message.channelId, socket, flowControl }); try { await this.parent.admit(socket, ceiling); if (this.closed || socket.readyState !== 1) { @@ -176,7 +208,11 @@ export class DaemonLinkEndpoint { } this.parent.delivery.reply({ type: "hosts.link.open.response", - payload: { requestId: message.requestId, routeId: id }, + payload: { + requestId: message.requestId, + routeId: id, + ...(flowControl ? { flowControl: true } : {}), + }, }); } catch { socket.close(); diff --git a/packages/server/src/server/connected-hosts/session.test.ts b/packages/server/src/server/connected-hosts/session.test.ts new file mode 100644 index 0000000000..bbf019f124 --- /dev/null +++ b/packages/server/src/server/connected-hosts/session.test.ts @@ -0,0 +1,70 @@ +import { expect, test } from "vitest"; +import { ConnectedHostSession, type ConnectedHosts } from "./index.js"; +import { SessionDelivery } from "../session/owned-subscriptions/index.js"; + +// The upstream host and final writer are independently controlled transport peers. +test.each(["text", "binary"] as const)( + "routed %s completion waits for the final writer", + async (kind) => { + const source = {}; + const writing = Promise.withResolvers(); + const drain = Promise.withResolvers(); + const write = () => { + writing.resolve(); + return drain.promise; + }; + const delivery = new SessionDelivery((_source, message) => { + if (message.type === "hosts.route.frame") return write(); + }, write); + delivery.attach(source, true); + let receive!: Parameters[0]["frame"]; + const unexpected = (): never => { + throw new Error("Unexpected host operation"); + }; + const session = new ConnectedHostSession( + { + list: unexpected, + observe: unexpected, + add: unexpected, + update: unexpected, + remove: unexpected, + openRoute: async (input) => { + receive = input.frame; + return { send: () => {}, close: () => {} }; + }, + }, + delivery, + { sourceId: "owner", manageSources: false }, + () => [], + () => {}, + () => false, + "origin", + () => 0, + ); + const request = { + type: "hosts.route.open.request" as const, + requestId: "open", + serverId: "remote", + scopeId: "scope", + }; + await delivery.request(source, request, () => session.dispatch(request)!); + let completed = false; + const forwarded = Promise.resolve( + receive(kind === "text" ? "payload" : new Uint8Array([1, 2])), + ).then(() => { + completed = true; + return undefined; + }); + try { + await writing.promise; + expect(completed).toBe(false); + drain.resolve(); + await forwarded; + expect(completed).toBe(true); + } finally { + drain.resolve(); + await forwarded; + await delivery.close(); + } + }, +); diff --git a/packages/server/src/server/connected-hosts/session.ts b/packages/server/src/server/connected-hosts/session.ts index 811daaa723..00d9f10a24 100644 --- a/packages/server/src/server/connected-hosts/session.ts +++ b/packages/server/src/server/connected-hosts/session.ts @@ -23,7 +23,10 @@ interface SessionRoute { export class ConnectedHostSession { private readonly routes = new Map(); constructor( - private readonly hosts: ConnectedHosts, + private readonly hosts: Pick< + ConnectedHosts, + "list" | "observe" | "add" | "update" | "remove" | "openRoute" + >, private readonly delivery: SessionDelivery, private readonly actor: HostActor, private readonly permissions: () => DaemonPermission[], @@ -206,7 +209,7 @@ export class ConnectedHostSession { originSessionId: this.originSessionId, permissions: this.permissions(), signal: owner.signal, - frame: (data) => { + frame: async (data) => { const output: SessionOutboundMessage | null = typeof data === "string" ? { @@ -223,8 +226,8 @@ export class ConnectedHostSession { fail("resource_limit", "The receiving client cannot keep up with host output"); return; } - if (output) owner.emit(output); - else owner.emitBinary(binary!); + if (output) await owner.emitAndWait(output); + else await owner.emitBinaryAndWait(binary!); }, ended: fail, }); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 17d0b6d109..11cf383818 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -451,6 +451,7 @@ function bufferFromWsData(data: Buffer | ArrayBuffer | Buffer[] | string): Buffe export interface WebSocketLike { maxFrameBytes?: number; + supportsConcurrentWrites?: boolean; readyState: number; bufferedAmount?: number; send: ( @@ -1307,6 +1308,7 @@ export class VoiceAssistantWebSocketServer { if (!transport) { transport = new ChunkedTransport({ frameBytes: ws.maxFrameBytes ?? 32 * 1024 * 1024, + supportsConcurrentWrites: ws.supportsConcurrentWrites, onFailure: (err) => this.closePhysicalSocket({ ws,