diff --git a/packages/web/server/__tests__/mux-websocket.test.ts b/packages/web/server/__tests__/mux-websocket.test.ts index 48f38e96f..6e2d12940 100644 --- a/packages/web/server/__tests__/mux-websocket.test.ts +++ b/packages/web/server/__tests__/mux-websocket.test.ts @@ -227,6 +227,23 @@ describe("SessionBroadcaster", () => { expect(callback).not.toHaveBeenCalled(); }); + it("does not warn for aborted fetches after shutdown starts", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + mockFetch.mockRejectedValueOnce(new Error("This operation was aborted")); + + const callback = vi.fn(); + broadcaster.subscribe(callback); + broadcaster.shutdown(); + await vi.advanceTimersByTimeAsync(10); + + expect(callback).not.toHaveBeenCalled(); + expect(warnSpy).not.toHaveBeenCalledWith( + "[SessionBroadcaster] fetchSnapshot error:", + "This operation was aborted", + ); + warnSpy.mockRestore(); + }); + it("returns null on non-OK response", async () => { mockFetch.mockResolvedValueOnce({ ok: false, status: 500 }); @@ -371,4 +388,28 @@ describe("TerminalManager.open — re-attach skipped when tmux session is gone ( expect(mockPtySpawn).toHaveBeenCalledTimes(2); expect(exitCb).not.toHaveBeenCalled(); }); + + it("detaches PTYs and reports code 0 during manager shutdown", async () => { + const pty = { + onData: vi.fn(), + onExit: vi.fn((cb: (evt: { exitCode: number }) => Promise | void) => { + capturedOnExit = cb; + }), + write: vi.fn(), + resize: vi.fn(), + kill: vi.fn(), + }; + mockPtySpawn.mockImplementationOnce(() => pty); + + const mgr = new TerminalManager("/usr/bin/tmux"); + const exitCb = vi.fn(); + mgr.subscribe("ao-177", undefined, vi.fn(), exitCb); + + mgr.shutdown(); + await capturedOnExit!({ exitCode: 1 }); + + expect(pty.write).toHaveBeenCalledWith("\x02d"); + expect(mockTmuxHasSession).not.toHaveBeenCalled(); + expect(exitCb).toHaveBeenCalledWith(0); + }); }); diff --git a/packages/web/server/__tests__/server-compatibility.test.ts b/packages/web/server/__tests__/server-compatibility.test.ts index efcfc01db..a4b13891f 100644 --- a/packages/web/server/__tests__/server-compatibility.test.ts +++ b/packages/web/server/__tests__/server-compatibility.test.ts @@ -57,6 +57,16 @@ describe("start-all.ts", () => { expect(source).toMatch(/spawnManagedDaemonChild/); expect(source).toMatch(/detached:\s*!\s*isWindows\(\)/); }); + + it("uses a configurable 15s shutdown grace period", () => { + expect(source).toMatch(/DEFAULT_SHUTDOWN_GRACE_MS\s*=\s*15_000/); + expect(source).toMatch(/AO_SHUTDOWN_GRACE_MS/); + }); + + it("does not treat the shutdown safety fallback as command failure", () => { + expect(source).not.toMatch(/process\.exit\(1\)/); + expect(source).toMatch(/safety fallback/); + }); }); describe("OrchestratorConfig compatibility", () => { diff --git a/packages/web/server/direct-terminal-ws.ts b/packages/web/server/direct-terminal-ws.ts index 99f028b1e..561c43a0e 100644 --- a/packages/web/server/direct-terminal-ws.ts +++ b/packages/web/server/direct-terminal-ws.ts @@ -5,14 +5,19 @@ import { createServer, type Server } from "node:http"; import type { WebSocketServer } from "ws"; +import { isWindows } from "@aoagents/ao-core"; import { findTmux } from "./tmux-utils.js"; import { createMuxWebSocket } from "./mux-websocket.js"; export interface DirectTerminalServer { server: Server; - shutdown: () => void; + shutdown: (drainMs?: number) => void; } +type MuxWebSocketServer = WebSocketServer & { + shutdownMuxResources?: () => void; +}; + /** * Create the direct terminal WebSocket server. * Separated from listen() so tests can control lifecycle. @@ -20,7 +25,7 @@ export interface DirectTerminalServer { export function createDirectTerminalServer(tmuxPath?: string | null): DirectTerminalServer { const TMUX = tmuxPath ?? findTmux(); - let muxWss: WebSocketServer | null = null; + let muxWss: MuxWebSocketServer | null = null; const metrics = { totalConnections: 0, @@ -65,22 +70,31 @@ export function createDirectTerminalServer(tmuxPath?: string | null): DirectTerm const pathname = new URL(request.url ?? "/", "ws://localhost").pathname; if (pathname === "/mux" && muxWss) { - muxWss.handleUpgrade(request, socket, head, (ws) => { - muxWss!.emit("connection", ws, request); + const activeMuxWss = muxWss; + activeMuxWss.handleUpgrade(request, socket, head, (ws) => { + activeMuxWss.emit("connection", ws, request); }); } else { socket.destroy(); } }); - function shutdown() { - // Terminate all connected mux clients — this triggers their 'close' events - // which unsubscribe terminal callbacks and kill PTY processes. - if (muxWss) { + function shutdown(drainMs = 0) { + muxWss?.shutdownMuxResources?.(); + + const closeMuxClients = () => { + if (!muxWss) return; for (const client of muxWss.clients) { - client.terminate(); + client.close(1001, "server shutdown"); } muxWss.close(); + }; + + if (drainMs > 0) { + const drainTimer = setTimeout(closeMuxClients, drainMs); + drainTimer.unref(); + } else { + closeMuxClients(); } server.close(); } @@ -102,7 +116,7 @@ if (isMainModule) { const TMUX = findTmux(); if (TMUX) { console.log(`[DirectTerminal] Using tmux: ${TMUX}`); - } else if (process.platform === "win32") { + } else if (isWindows()) { console.log(`[DirectTerminal] Windows mode — using named pipe relay to PTY hosts`); } else { console.log(`[DirectTerminal] No tmux available — terminal relay may be limited`); @@ -114,12 +128,18 @@ if (isMainModule) { console.log(`[DirectTerminal] WebSocket server listening on port ${PORT}`); }); + let shuttingDown = false; + function handleShutdown(signal: string) { + if (shuttingDown) return; + shuttingDown = true; console.log(`[DirectTerminal] Received ${signal}, shutting down...`); - shutdown(); + shutdown(1_000); const forceExitTimer = setTimeout(() => { - console.error("[DirectTerminal] Forced shutdown after timeout"); - process.exit(1); + console.warn( + "[DirectTerminal] Shutdown cleanup still pending after timeout; exiting cleanly", + ); + process.exit(0); }, 5000); forceExitTimer.unref(); } diff --git a/packages/web/server/mux-websocket.ts b/packages/web/server/mux-websocket.ts index 5515574df..a9a8aea72 100644 --- a/packages/web/server/mux-websocket.ts +++ b/packages/web/server/mux-websocket.ts @@ -63,6 +63,7 @@ export class SessionBroadcaster { private errorSubscribers = new Set<(error: string) => void>(); private intervalId: ReturnType | null = null; private polling = false; + private shuttingDown = false; private readonly baseUrl: string; constructor(nextPort: string) { @@ -165,11 +166,19 @@ export class SessionBroadcaster { } catch (err) { clearTimeout(timeoutId); const msg = err instanceof Error ? err.message : String(err); + if (this.shuttingDown && isAbortError(err)) { + return { sessions: null, error: null }; + } console.warn("[SessionBroadcaster] fetchSnapshot error:", msg); return { sessions: null, error: msg }; } } + shutdown(): void { + this.shuttingDown = true; + this.disconnect(); + } + private disconnect(): void { if (this.intervalId !== null) { clearInterval(this.intervalId); @@ -178,6 +187,13 @@ export class SessionBroadcaster { } } +function isAbortError(err: unknown): boolean { + return ( + err instanceof Error && + (err.name === "AbortError" || err.message === "This operation was aborted") + ); +} + // node-pty is an optionalDependency — load dynamically /* eslint-disable @typescript-eslint/consistent-type-imports -- node-pty is optional; static import would crash if missing */ type IPty = import("node-pty").IPty; @@ -232,6 +248,7 @@ const REATTACH_RESET_GRACE_MS = 5_000; export class TerminalManager { private terminals = new Map(); private TMUX: string; + private shuttingDown = false; constructor(tmuxPath?: string) { const resolved = tmuxPath ?? findTmux(); @@ -381,9 +398,25 @@ export class TerminalManager { // request, in-flight terminal) for up to the subprocess timeout when // tmux is slow to respond. pty.onExit(async ({ exitCode }) => { - console.log(`[MuxServer] PTY exited for ${id} with code ${exitCode}`); + const reportedExitCode = this.shuttingDown ? 0 : exitCode; + if (this.shuttingDown) { + console.log(`[MuxServer] PTY detached for ${id} during shutdown`); + } else { + console.log(`[MuxServer] PTY exited for ${id} with code ${exitCode}`); + } terminal.pty = null; + if (this.shuttingDown) { + if (terminal.resetTimer) { + clearTimeout(terminal.resetTimer); + terminal.resetTimer = undefined; + } + for (const cb of terminal.exitCallbacks) { + cb(reportedExitCode); + } + return; + } + // Skip the re-attach loop entirely when the underlying tmux session is // gone (e.g. user pressed Ctrl-C in the pane and the launch command // exited, taking the only window with it). Without this guard we @@ -392,17 +425,14 @@ export class TerminalManager { // clean user-initiated termination — see issue #1756. The // MAX_REATTACH_ATTEMPTS bound from #1640 still covers tmux server // hiccups where the session does still exist. - if ( - terminal.subscribers.size > 0 && - !(await tmuxHasSession(this.TMUX, tmuxSessionId)) - ) { + if (terminal.subscribers.size > 0 && !(await tmuxHasSession(this.TMUX, tmuxSessionId))) { console.log(`[MuxServer] tmux session ${tmuxSessionId} is gone, not re-attaching`); if (terminal.resetTimer) { clearTimeout(terminal.resetTimer); terminal.resetTimer = undefined; } for (const cb of terminal.exitCallbacks) { - cb(exitCode); + cb(reportedExitCode); } return; } @@ -433,7 +463,7 @@ export class TerminalManager { // Notify subscribers that the terminal has exited (re-attach failed or no subscribers) for (const cb of terminal.exitCallbacks) { - cb(exitCode); + cb(reportedExitCode); } }); @@ -495,10 +525,16 @@ export class TerminalManager { terminal.resetTimer = undefined; } if (terminal.pty) { - terminal.pty.kill(); - terminal.pty = null; + if (this.shuttingDown) { + terminal.pty.write("\x02d"); + } else { + terminal.pty.kill(); + terminal.pty = null; + } + } + if (!this.shuttingDown) { + this.terminals.delete(key); } - this.terminals.delete(key); } }; } @@ -511,6 +547,30 @@ export class TerminalManager { if (!terminal) return ""; return terminal.buffer.join(""); } + + shutdown(graceMs = 1_000): void { + if (this.shuttingDown) return; + this.shuttingDown = true; + + for (const terminal of this.terminals.values()) { + if (terminal.resetTimer) { + clearTimeout(terminal.resetTimer); + terminal.resetTimer = undefined; + } + terminal.pty?.write("\x02d"); + } + + const killTimer = setTimeout(() => { + for (const terminal of this.terminals.values()) { + if (terminal.pty) { + terminal.pty.kill(); + terminal.pty = null; + } + } + this.terminals.clear(); + }, graceMs); + killTimer.unref(); + } } // ── Windows Pipe Relay (extracted for testability) ── @@ -637,9 +697,7 @@ export function handleWindowsPipeMessage( try { const status = JSON.parse(payload.toString("utf-8")) as { alive: boolean }; if (!status.alive && ws.readyState === WS_OPEN) { - ws.send( - JSON.stringify({ ch: "terminal", id, type: "exited", code: 0, ...echo }), - ); + ws.send(JSON.stringify({ ch: "terminal", id, type: "exited", code: 0, ...echo })); } } catch { /* ignore parse errors */ @@ -705,6 +763,11 @@ export function createMuxWebSocket(tmuxPath?: string | null): WebSocketServer | const broadcaster = new SessionBroadcaster(nextPort); const wss = new WebSocketServer({ noServer: true }); + const shutdown = (): void => { + broadcaster.shutdown(); + terminalManager?.shutdown(); + }; + (wss as WebSocketServer & { shutdownMuxResources?: () => void }).shutdownMuxResources = shutdown; wss.on("connection", (ws) => { console.log("[MuxServer] New mux connection"); @@ -759,7 +822,14 @@ export function createMuxWebSocket(tmuxPath?: string | null): WebSocketServer | if (type === "open") { if (isWindows()) { handleWindowsPipeMessage( - msg as { id: string; type: string; projectId?: string; data?: string; cols?: number; rows?: number }, + msg as { + id: string; + type: string; + projectId?: string; + data?: string; + cols?: number; + rows?: number; + }, ws, winPipes, winPipeBuffers, @@ -841,7 +911,13 @@ export function createMuxWebSocket(tmuxPath?: string | null): WebSocketServer | } else if (type === "resize" && "cols" in msg && "rows" in msg) { if (isWindows()) { handleWindowsPipeMessage( - msg as { id: string; type: string; projectId?: string; cols: number; rows: number }, + msg as { + id: string; + type: string; + projectId?: string; + cols: number; + rows: number; + }, ws, winPipes, winPipeBuffers, @@ -853,7 +929,14 @@ export function createMuxWebSocket(tmuxPath?: string | null): WebSocketServer | } else if (type === "close") { if (isWindows()) { handleWindowsPipeMessage( - msg as { id: string; type: string; projectId?: string; data?: string; cols?: number; rows?: number }, + msg as { + id: string; + type: string; + projectId?: string; + data?: string; + cols?: number; + rows?: number; + }, ws, winPipes, winPipeBuffers, diff --git a/packages/web/server/start-all.ts b/packages/web/server/start-all.ts index d54b11c33..02c936380 100644 --- a/packages/web/server/start-all.ts +++ b/packages/web/server/start-all.ts @@ -23,12 +23,35 @@ const __dirname = dirname(__filename); const pkgRoot = resolve(__dirname, ".."); const children: ChildProcess[] = []; +const exitedChildren = new WeakSet(); markDaemonShutdownHandlerInstalled(); +const DEFAULT_SHUTDOWN_GRACE_MS = 15_000; + function log(label: string, msg: string): void { process.stdout.write(`[${label}] ${msg}\n`); } +function warn(label: string, msg: string): void { + process.stderr.write(`[${label}] warn: ${msg}\n`); +} + +function getShutdownGraceMs(): number { + const raw = process.env["AO_SHUTDOWN_GRACE_MS"]; + if (!raw) return DEFAULT_SHUTDOWN_GRACE_MS; + + const parsed = Number.parseInt(raw, 10); + if (!Number.isFinite(parsed) || parsed <= 0) { + warn( + "start-all", + `invalid AO_SHUTDOWN_GRACE_MS=${JSON.stringify(raw)}, using ${DEFAULT_SHUTDOWN_GRACE_MS}ms`, + ); + return DEFAULT_SHUTDOWN_GRACE_MS; + } + + return parsed; +} + function spawnProcess( label: string, command: string, @@ -60,6 +83,7 @@ function spawnProcess( }); child.on("exit", (code) => { + exitedChildren.add(child); log(label, `exited with code ${code}`); if (!shuttingDown && opts?.restart && code !== 0 && restarts < maxRestarts) { restarts++; @@ -129,24 +153,44 @@ function cleanup(): void { if (shuttingDown) return; shuttingDown = true; - let alive = children.length; + const activeChildren = children.filter((child) => !exitedChildren.has(child)); + let alive = activeChildren.length; if (alive === 0) { process.exit(0); return; } - // Force exit after 5s if children don't exit cleanly + const shutdownGraceMs = getShutdownGraceMs(); + const startedAt = Date.now(); + + // Safety fallback: after the grace period, force-kill any child process + // trees that are still alive and exit successfully. The operator intent was + // shutdown; this path prevents orphans rather than indicating command failure. const forceTimer = setTimeout(() => { - log("start-all", "Children did not exit in time, forcing shutdown"); - process.exit(1); - }, 5000); + const remaining = activeChildren.filter((child) => !exitedChildren.has(child)); + const labels = remaining.map((child) => `${child.pid ?? "unknown"}`).join(", "); + warn( + "start-all", + `shutdown grace period (${shutdownGraceMs}ms) elapsed; force-killing remaining child process trees (${labels || "unknown"}) as a safety fallback`, + ); + for (const child of remaining) { + const pid = child.pid; + if (pid) { + void killProcessTree(pid, "SIGKILL"); + } else { + child.kill("SIGKILL"); + } + } + process.exit(0); + }, shutdownGraceMs); forceTimer.unref(); - for (const child of children) { + for (const child of activeChildren) { child.on("exit", () => { alive--; if (alive <= 0) { clearTimeout(forceTimer); + log("start-all", `all children shut down cleanly in ${Date.now() - startedAt}ms`); process.exit(0); } });