From cc9b321f802e94d169bb74786195139137d4d232 Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Fri, 17 Jul 2026 14:19:32 -0700 Subject: [PATCH] fix(pty-daemon): pause flooding PTYs on backpressure instead of destroying the shared socket --- packages/pty-daemon/package.json | 2 +- packages/pty-daemon/src/Pty/Pty.ts | 27 +++++ packages/pty-daemon/src/Server/Server.ts | 52 +++++++++ .../src/SessionStore/SessionStore.test.ts | 2 + .../src/SessionStore/snapshot.test.ts | 2 + .../pty-daemon/src/handlers/handlers.test.ts | 2 + .../pty-daemon/test/byte-fidelity.test.ts | 2 + packages/pty-daemon/test/flow-control.test.ts | 107 ++++++++++++++++++ 8 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 packages/pty-daemon/test/flow-control.test.ts diff --git a/packages/pty-daemon/package.json b/packages/pty-daemon/package.json index e65436ba84c..8cae8efafe2 100644 --- a/packages/pty-daemon/package.json +++ b/packages/pty-daemon/package.json @@ -30,7 +30,7 @@ "build:daemon": "bun run build.ts", "typecheck": "tsc --noEmit --emitDeclarationOnly false", "test": "bun test src/protocol src/SessionStore src/handlers src/Pty/Pty.test.ts test/no-encoding-hops.test.ts", - "test:integration": "node --experimental-strip-types --test test/integration.test.ts test/control-plane.test.ts test/signal-recovery.test.ts test/byte-fidelity.test.ts test/handoff.test.ts test/foreground-process.test.ts", + "test:integration": "node --experimental-strip-types --test test/integration.test.ts test/control-plane.test.ts test/signal-recovery.test.ts test/byte-fidelity.test.ts test/handoff.test.ts test/foreground-process.test.ts test/flow-control.test.ts", "build:types": "tsc -p tsconfig.types.json" }, "dependencies": { diff --git a/packages/pty-daemon/src/Pty/Pty.ts b/packages/pty-daemon/src/Pty/Pty.ts index 6bf8a83d4a7..18b4b9dc222 100644 --- a/packages/pty-daemon/src/Pty/Pty.ts +++ b/packages/pty-daemon/src/Pty/Pty.ts @@ -26,6 +26,13 @@ export interface Pty { kill(signal?: NodeJS.Signals): void; onData(cb: PtyOnData): void; onExit(cb: PtyOnExit): void; + /** + * Flow control: stop reading from the PTY master. The kernel PTY buffer + * (~64KB) fills and the foreground process blocks on write, throttling + * itself — same mechanism as VS Code's ptyHost pause/resume. + */ + pause(): void; + resume(): void; /** * The kernel master fd backing this PTY. Required for daemon-upgrade * fd-handoff (Phase 2): the successor daemon process inherits this fd @@ -114,6 +121,16 @@ class NodePtyAdapter implements Pty { this.exitCallbacks.push(cb); } + pause(): void { + if (this.exited) return; + this.term.pause(); + } + + resume(): void { + if (this.exited) return; + this.term.resume(); + } + private scheduleKillEscalation( signal: NodeJS.Signals, targets: ProcessSignalTarget[], @@ -331,6 +348,16 @@ class AdoptedPty implements Pty { this.exitCallbacks.push(cb); } + pause(): void { + if (this.exitFired) return; + this.reader.pause(); + } + + resume(): void { + if (this.exitFired) return; + this.reader.resume(); + } + private scheduleKillEscalation( signal: NodeJS.Signals, targets: ProcessSignalTarget[], diff --git a/packages/pty-daemon/src/Server/Server.ts b/packages/pty-daemon/src/Server/Server.ts index 5b2db938575..ec82057e950 100644 --- a/packages/pty-daemon/src/Server/Server.ts +++ b/packages/pty-daemon/src/Server/Server.ts @@ -35,6 +35,8 @@ export interface ServerOptions { daemonVersion: string; bufferCap?: number; outboundBufferCap?: number; + /** Pause producing PTYs when a subscriber buffers past this (flow control). */ + outboundPauseThreshold?: number; /** * Override for the PTY-spawn factory. Production leaves this unset; * `defaultSpawn` (real node-pty) is used. Tests inject a fake here so @@ -45,10 +47,20 @@ export interface ServerOptions { const DEFAULT_OUTBOUND_BUFFER_CAP_BYTES = 8 * 1024 * 1024; +// Flow control: when a subscriber's socket buffers more than this, pause the +// producing PTYs instead of letting the buffer grow toward the destroy cap. +// The kernel PTY buffer then fills and the foreground process blocks on +// write — the flood throttles itself at the source (VS Code ptyHost does the +// same via pause/resume; our congestion signal is local writableLength +// instead of renderer ACKs). Resume on socket 'drain'. +const DEFAULT_OUTBOUND_PAUSE_THRESHOLD_BYTES = 1 * 1024 * 1024; + interface ConnState extends Conn { socket: net.Socket; decoder: FrameDecoder; negotiated: number | null; + /** Session ids whose PTYs we paused because this conn was congested. */ + pausedSessions: Set; } export class Server { @@ -317,11 +329,17 @@ export class Server { decoder: new FrameDecoder(), negotiated: null, subscriptions: new Set(), + pausedSessions: new Set(), send: (msg, payload) => writeMessage(socket, msg, payload, outboundBufferCap), }; this.conns.add(conn); + // Flow control: the buffer only reaches the pause threshold after + // write() has returned false, so a 'drain' (buffer fully flushed) is + // guaranteed to follow and is our resume signal. + socket.on("drain", () => this.resumePausedSessions(conn)); + socket.on("data", (chunk) => { try { conn.decoder.push(chunk); @@ -460,13 +478,24 @@ export class Server { * subscribed to this session id receives the output / exit frames. */ private wireSession(session: Session): void { + const pauseThreshold = + this.opts.outboundPauseThreshold ?? + DEFAULT_OUTBOUND_PAUSE_THRESHOLD_BYTES; session.pty.onData((chunk) => { this.store.appendOutput(session, chunk); const out: ServerMessage = { type: "output", id: session.id }; + let congested = false; for (const c of this.conns) { if (!c.subscriptions.has(session.id)) continue; c.send(out, chunk); + if (!c.socket.destroyed && c.socket.writableLength > pauseThreshold) { + c.pausedSessions.add(session.id); + congested = true; + } } + // Paused PTYs emit no further onData, so this fires at most once + // per congestion episode; the conn's 'drain' resumes us. + if (congested) session.pty.pause(); }); session.pty.onExit((info) => { session.exited = true; @@ -483,6 +512,7 @@ export class Server { c.send(ev); c.subscriptions.delete(session.id); } + c.pausedSessions.delete(session.id); } // Delete the session immediately. Without this, every closed // terminal pane left a row in the store forever — list-reply @@ -500,6 +530,28 @@ export class Server { private dropConn(conn: ConnState): void { this.conns.delete(conn); + // A destroyed socket never drains — release its paused producers here. + this.resumePausedSessions(conn); + } + + /** + * Resume every PTY paused on behalf of `conn`, unless another congested + * conn still holds it paused. Called on socket 'drain' and on conn drop. + */ + private resumePausedSessions(conn: ConnState): void { + for (const id of conn.pausedSessions) { + conn.pausedSessions.delete(id); + let heldElsewhere = false; + for (const other of this.conns) { + if (other !== conn && other.pausedSessions.has(id)) { + heldElsewhere = true; + break; + } + } + if (heldElsewhere) continue; + const session = this.store.get(id); + if (session && !session.exited) session.pty.resume(); + } } } diff --git a/packages/pty-daemon/src/SessionStore/SessionStore.test.ts b/packages/pty-daemon/src/SessionStore/SessionStore.test.ts index f2ab3992997..5153768e9ce 100644 --- a/packages/pty-daemon/src/SessionStore/SessionStore.test.ts +++ b/packages/pty-daemon/src/SessionStore/SessionStore.test.ts @@ -17,6 +17,8 @@ function fakePty(meta: { cols: number; rows: number }): Pty { onData: () => {}, onExit: () => {}, getMasterFd: () => -1, + pause: () => {}, + resume: () => {}, }; } diff --git a/packages/pty-daemon/src/SessionStore/snapshot.test.ts b/packages/pty-daemon/src/SessionStore/snapshot.test.ts index 79ee26ce301..67209904b4c 100644 --- a/packages/pty-daemon/src/SessionStore/snapshot.test.ts +++ b/packages/pty-daemon/src/SessionStore/snapshot.test.ts @@ -22,6 +22,8 @@ function fakePty(pid: number, meta: { cols: number; rows: number }): Pty { onData: () => {}, onExit: () => {}, getMasterFd: () => -1, + pause: () => {}, + resume: () => {}, }; } diff --git a/packages/pty-daemon/src/handlers/handlers.test.ts b/packages/pty-daemon/src/handlers/handlers.test.ts index 8facfa1f3f9..6d0c073eb61 100644 --- a/packages/pty-daemon/src/handlers/handlers.test.ts +++ b/packages/pty-daemon/src/handlers/handlers.test.ts @@ -38,6 +38,8 @@ function makeFakePty(state: FakePtyState, meta: SpawnOptions["meta"]): Pty { onData: () => {}, onExit: () => {}, getMasterFd: () => -1, + pause: () => {}, + resume: () => {}, }; } diff --git a/packages/pty-daemon/test/byte-fidelity.test.ts b/packages/pty-daemon/test/byte-fidelity.test.ts index f324e48ac7d..a1cfc91d0c3 100644 --- a/packages/pty-daemon/test/byte-fidelity.test.ts +++ b/packages/pty-daemon/test/byte-fidelity.test.ts @@ -64,6 +64,8 @@ function makeDriveablePty(meta: SpawnOptions["meta"]): DriveablePty { resize: () => {}, kill: () => {}, getMasterFd: () => -1, + pause: () => {}, + resume: () => {}, onData: (cb: PtyOnData) => { onDataCbs.push(cb); }, diff --git a/packages/pty-daemon/test/flow-control.test.ts b/packages/pty-daemon/test/flow-control.test.ts new file mode 100644 index 00000000000..17913d55ece --- /dev/null +++ b/packages/pty-daemon/test/flow-control.test.ts @@ -0,0 +1,107 @@ +// Flow control: a flooding PTY with a slow consumer must pause the producer +// (kernel backpressure) instead of destroying the shared daemon socket. +// +// Regression: writeMessage() used to socket.destroy() when writableLength +// exceeded the outbound cap — one flooding terminal severed the connection +// for every session in the org. +// +// Runs under Node (`node --experimental-strip-types --test`). + +import { strict as assert } from "node:assert"; +import * as os from "node:os"; +import * as path from "node:path"; +import { after, before, test } from "node:test"; +import { Server } from "../src/Server/index.ts"; +import { connectAndHello, payloadAsString } from "./helpers/client.ts"; + +const sockPath = path.join(os.tmpdir(), `pty-daemon-flow-${process.pid}.sock`); +let server: Server; + +// Small thresholds so the flood trips them in milliseconds. The old destroy +// behavior fires at outboundBufferCap; with flow control the buffer must stop +// growing at outboundPauseThreshold, far below the cap. +const PAUSE_THRESHOLD = 64 * 1024; +const DESTROY_CAP = 4 * 1024 * 1024; + +const FLOOD_META = { + shell: "/bin/sh", + argv: ["-c", 'yes "flood-line-0123456789-abcdefghijklmnopqrstuvwxyz"'], + cols: 80, + rows: 24, +}; + +before(async () => { + server = new Server({ + socketPath: sockPath, + daemonVersion: "0.0.0-test", + outboundPauseThreshold: PAUSE_THRESHOLD, + outboundBufferCap: DESTROY_CAP, + }); + await server.listen(); +}); + +after(async () => { + await server.close(); +}); + +test("slow consumer pauses the flood instead of destroying the connection", async () => { + const c = await connectAndHello(sockPath); + c.send({ type: "open", id: "flow-0", meta: FLOOD_META }); + await c.waitFor((m) => m.type === "open-ok" && m.id === "flow-0"); + c.send({ type: "subscribe", id: "flow-0", replay: false }); + await c.waitFor( + (m) => m.type === "output" && payloadAsString(m).includes("flood-line"), + 3000, + ); + + // Stop reading. At full `yes` rate the old behavior blows through any + // buffer cap within this window and destroys the socket. + c.socket.pause(); + await new Promise((r) => setTimeout(r, 1500)); + assert.equal( + c.closed(), + false, + "daemon destroyed the shared connection under flood (flow control missing)", + ); + + // Resume reading: buffered output drains, PTY resumes, data flows again. + c.socket.resume(); + await c.waitForNext( + (m) => m.type === "output" && payloadAsString(m).includes("flood-line"), + 3000, + ); + assert.equal(c.closed(), false); + + c.send({ type: "close", id: "flow-0", signal: "SIGKILL" }); + await c.waitFor((m) => m.type === "closed" && m.id === "flow-0", 3000); + await c.close(); +}); + +test("consumer disconnect while paused resumes the PTY for the next subscriber", async () => { + const c1 = await connectAndHello(sockPath); + c1.send({ type: "open", id: "flow-1", meta: FLOOD_META }); + await c1.waitFor((m) => m.type === "open-ok" && m.id === "flow-1"); + c1.send({ type: "subscribe", id: "flow-1", replay: false }); + await c1.waitFor((m) => m.type === "output" && m.id === "flow-1", 3000); + + // Congest c1 until the PTY is paused, then drop the socket hard. The + // daemon must resume the PTY via the conn-drop path (a destroyed socket + // never emits 'drain'). + c1.socket.pause(); + await new Promise((r) => setTimeout(r, 800)); + c1.socket.destroy(); + + const c2 = await connectAndHello(sockPath); + c2.send({ type: "subscribe", id: "flow-1", replay: false }); + await c2.waitFor( + (m) => + m.type === "output" && + m.id === "flow-1" && + payloadAsString(m).includes("flood-line"), + 3000, + ); + + c2.send({ type: "close", id: "flow-1", signal: "SIGKILL" }); + await c2.waitFor((m) => m.type === "closed" && m.id === "flow-1", 3000); + await c2.close(); +});