Skip to content

Commit 9c03b23

Browse files
steipetesteipete
authored
fix(terminal): keep slow clients connected under heavy output (#107348)
* fix(terminal): coalesce output and throttle slow clients * fix(terminal): satisfy lint and deadcode gates
1 parent c1b55ff commit 9c03b23

13 files changed

Lines changed: 509 additions & 85 deletions

‎src/gateway/gateway-misc.test.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,18 @@ function broadcastChatClassEvents(
382382
}
383383

384384
describe("gateway broadcaster", () => {
385+
it("exposes current buffered bytes for a targeted connection", () => {
386+
const socket = makeRecordingSocket();
387+
socket.bufferedAmount = 1234;
388+
const client = makeOperatorWsClient("c-admin", socket, ["operator.admin"]);
389+
const clients = new Set<GatewayWsClient>([client]);
390+
const { getBufferedAmount } = createGatewayBroadcaster({ clients });
391+
392+
expect(getBufferedAmount("c-admin")).toBe(1234);
393+
clients.delete(client);
394+
expect(getBufferedAmount("c-admin")).toBeUndefined();
395+
});
396+
385397
it("keeps workers outside all generic and targeted gateway broadcasts", () => {
386398
const workerSocket = makeRecordingSocket();
387399
const worker = makeGatewayWsClient("c-worker", workerSocket, {

‎src/gateway/server-broadcast-types.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,6 @@ export type GatewayBroadcastToConnIdsFn = (
2525
connIds: ReadonlySet<string>,
2626
opts?: GatewayBroadcastOpts,
2727
) => void;
28+
29+
/** Current queued outbound bytes for one live gateway connection. */
30+
export type GatewayBufferedAmountFn = (connId: string) => number | undefined;

‎src/gateway/server-broadcast.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import type {
1212
GatewayBroadcastFn,
1313
GatewayBroadcastOpts,
1414
GatewayBroadcastToConnIdsFn,
15+
GatewayBufferedAmountFn,
1516
} from "./server-broadcast-types.js";
1617
import { MAX_BUFFERED_BYTES } from "./server-constants.js";
1718
import type { GatewayWsClient } from "./server/ws-types.js";
@@ -216,5 +217,14 @@ export function createGatewayBroadcaster(params: { clients: Set<GatewayWsClient>
216217
broadcastInternal(event, payload, opts, connIds);
217218
};
218219

219-
return { broadcast, broadcastToConnIds };
220+
const getBufferedAmount: GatewayBufferedAmountFn = (connId) => {
221+
for (const client of params.clients) {
222+
if (client.connId === connId) {
223+
return client.socket.bufferedAmount;
224+
}
225+
}
226+
return undefined;
227+
};
228+
229+
return { broadcast, broadcastToConnIds, getBufferedAmount };
220230
}

‎src/gateway/server-runtime-state.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,11 @@ import type { HooksConfigResolved } from "./hooks.js";
2929
import type { AuthorizedGatewayHttpRequest } from "./http-auth-utils.js";
3030
import { createMcpAppSandboxHttpServer } from "./mcp-app-sandbox-http.js";
3131
import { isLoopbackHost, resolveGatewayListenHosts } from "./net.js";
32-
import type { GatewayBroadcastFn, GatewayBroadcastToConnIdsFn } from "./server-broadcast-types.js";
32+
import type {
33+
GatewayBroadcastFn,
34+
GatewayBroadcastToConnIdsFn,
35+
GatewayBufferedAmountFn,
36+
} from "./server-broadcast-types.js";
3337
import { createGatewayBroadcaster } from "./server-broadcast.js";
3438
import {
3539
type ChatRunEntry,
@@ -127,6 +131,7 @@ export async function createGatewayRuntimeState(params: {
127131
clients: Set<GatewayWsClient>;
128132
broadcast: GatewayBroadcastFn;
129133
broadcastToConnIds: GatewayBroadcastToConnIdsFn;
134+
getBufferedAmount: GatewayBufferedAmountFn;
130135
agentRunSeq: Map<string, number>;
131136
dedupe: Map<string, DedupeEntry>;
132137
chatRunState: ReturnType<typeof createChatRunState>;
@@ -156,7 +161,9 @@ export async function createGatewayRuntimeState(params: {
156161
const resolvePluginRouteRegistry = () =>
157162
params.getPluginRouteRegistry?.() ?? params.pluginRegistry;
158163
const clients = new Set<GatewayWsClient>();
159-
const { broadcast, broadcastToConnIds } = createGatewayBroadcaster({ clients });
164+
const { broadcast, broadcastToConnIds, getBufferedAmount } = createGatewayBroadcaster({
165+
clients,
166+
});
160167

161168
let loadedHooksRequestHandler: HooksRequestHandler | null = null;
162169
const handleHooksRequest: HooksRequestHandler = async (req, res) => {
@@ -461,6 +468,7 @@ export async function createGatewayRuntimeState(params: {
461468
clients,
462469
broadcast,
463470
broadcastToConnIds,
471+
getBufferedAmount,
464472
agentRunSeq,
465473
dedupe,
466474
chatRunState,

‎src/gateway/server.impl.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1132,6 +1132,7 @@ export async function startGatewayServer(
11321132
clients,
11331133
broadcast,
11341134
broadcastToConnIds,
1135+
getBufferedAmount,
11351136
agentRunSeq,
11361137
dedupe,
11371138
chatRunState,
@@ -1245,12 +1246,9 @@ export async function startGatewayServer(
12451246
watchNodeRequestHandler.current = watchNodeHttpRuntime.handleRequest;
12461247
const { TerminalSessionManager, DEFAULT_TERMINAL_DETACH_SECONDS } =
12471248
await import("./terminal/session-manager.js");
1248-
// One PTY store per gateway. Emits each session's bytes only to the owning
1249-
// connection so terminals stay private to the operator that opened them.
1250-
// Startup config is enough here: gateway.terminal.* changes restart the
1251-
// gateway (config-reload-plan), so the grace period never drifts at runtime.
1249+
const { createTerminalSessionTransport } = await import("./terminal/gateway-transport.js");
12521250
const terminalSessions = new TerminalSessionManager({
1253-
emit: (connId, event, payload) => broadcastToConnIds(event, payload, new Set([connId])),
1251+
...createTerminalSessionTransport(broadcastToConnIds, getBufferedAmount),
12541252
detachGraceMs:
12551253
(cfgAtStart.gateway?.terminal?.detachedSessionTimeoutSeconds ??
12561254
DEFAULT_TERMINAL_DETACH_SECONDS) * 1000,

‎src/gateway/terminal/backend.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ export type TerminalBackendExit = {
99
export interface TerminalBackend {
1010
write(data: string): void;
1111
resize(cols: number, rows: number): void;
12+
pause(): void;
13+
resume(): void;
1214
kill(): void;
1315
onData(callback: (data: string) => void): void;
1416
onExit(callback: (exit: TerminalBackendExit) => void): void;
@@ -24,6 +26,8 @@ export async function createLocalTerminalBackend(
2426
return {
2527
write: (data) => pty.write(data),
2628
resize: (cols, rows) => pty.resize(cols, rows),
29+
pause: () => pty.pause(),
30+
resume: () => pty.resume(),
2731
kill: () => pty.kill(),
2832
onData: (callback) => pty.onData(callback),
2933
onExit: (callback) => pty.onExit(callback),

‎src/gateway/terminal/gateway-transport.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import type {
2+
GatewayBroadcastToConnIdsFn,
3+
GatewayBufferedAmountFn,
4+
} from "../server-broadcast-types.js";
5+
6+
export const TERMINAL_EVENT_DATA = "terminal.data" as const;
7+
export const TERMINAL_EVENT_EXIT = "terminal.exit" as const;
8+
9+
/** Adapts terminal ownership to targeted gateway delivery and pressure state. */
10+
export function createTerminalSessionTransport(
11+
broadcastToConnIds: GatewayBroadcastToConnIdsFn,
12+
getBufferedAmount: GatewayBufferedAmountFn,
13+
) {
14+
return {
15+
emit: (connId: string, event: string, payload: unknown) =>
16+
broadcastToConnIds(event, payload, new Set([connId]), {
17+
// PTY flow control is primary; dropping is the last-resort socket cap guard.
18+
dropIfSlow: event === TERMINAL_EVENT_DATA,
19+
}),
20+
getBufferedAmount,
21+
};
22+
}

‎src/gateway/terminal/node-relay.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,10 @@ export async function createNodeRelayBackend(params: {
134134
resize(cols, rows) {
135135
send({ kind: "resize", cols, rows });
136136
},
137+
// Node host pauses its PTY around each awaited progress write; this relay
138+
// has no local PTY or transport-level pause operation to control.
139+
pause() {},
140+
resume() {},
137141
kill() {
138142
abort.abort();
139143
},

‎src/gateway/terminal/output-coalescer.ts‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { truncateUtf8Prefix } from "../../utils/utf8-truncate.js";
2+
3+
const TERMINAL_OUTPUT_COALESCE_WINDOW_MS = 4;
4+
const TERMINAL_OUTPUT_FRAME_BYTES = 64 * 1024;
5+
6+
/** Batches adjacent PTY chunks while keeping each emitted frame UTF-8 bounded. */
7+
export class TerminalOutputCoalescer {
8+
private readonly emit: (data: string) => void;
9+
private chunks: string[] = [];
10+
private bufferedBytes = 0;
11+
private timer: ReturnType<typeof setTimeout> | null = null;
12+
13+
constructor(emit: (data: string) => void) {
14+
this.emit = emit;
15+
}
16+
17+
get isEmpty(): boolean {
18+
return this.chunks.length === 0;
19+
}
20+
21+
push(data: string, opts?: { flushNow?: boolean }): void {
22+
let remaining = data;
23+
while (remaining) {
24+
const available = TERMINAL_OUTPUT_FRAME_BYTES - this.bufferedBytes;
25+
const part = truncateUtf8Prefix(remaining, available);
26+
if (!part) {
27+
this.flush();
28+
continue;
29+
}
30+
this.chunks.push(part);
31+
this.bufferedBytes += Buffer.byteLength(part, "utf8");
32+
remaining = remaining.slice(part.length);
33+
if (this.bufferedBytes >= TERMINAL_OUTPUT_FRAME_BYTES) {
34+
this.flush();
35+
}
36+
}
37+
if (opts?.flushNow) {
38+
this.flush();
39+
} else {
40+
this.schedule();
41+
}
42+
}
43+
44+
flush(): void {
45+
if (this.timer) {
46+
clearTimeout(this.timer);
47+
this.timer = null;
48+
}
49+
if (this.chunks.length === 0) {
50+
return;
51+
}
52+
const data = this.chunks.join("");
53+
this.chunks = [];
54+
this.bufferedBytes = 0;
55+
this.emit(data);
56+
}
57+
58+
clear(): void {
59+
if (this.timer) {
60+
clearTimeout(this.timer);
61+
this.timer = null;
62+
}
63+
this.chunks = [];
64+
this.bufferedBytes = 0;
65+
}
66+
67+
dispose(opts?: { flush?: boolean }): void {
68+
if (opts?.flush) {
69+
this.flush();
70+
} else {
71+
this.clear();
72+
}
73+
}
74+
75+
private schedule(): void {
76+
if (this.timer || this.chunks.length === 0) {
77+
return;
78+
}
79+
this.timer = setTimeout(() => {
80+
this.timer = null;
81+
this.flush();
82+
}, TERMINAL_OUTPUT_COALESCE_WINDOW_MS);
83+
this.timer.unref?.();
84+
}
85+
}

0 commit comments

Comments
 (0)