From 50c8f3ff7ab0a08c7d8f4e8f55b87d20e752e2a4 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Fri, 4 Jul 2025 18:36:01 +0000 Subject: [PATCH] fix: prevent Node.js process leaks during debug sessions (#5397) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add centralized ProcessRegistry for tracking spawned processes - Implement debug session lifecycle management with automatic cleanup - Enhance extension deactivation to kill all tracked processes - Integrate process tracking in ExecaTerminalProcess and MCP servers - Add comprehensive process tree cleanup with SIGTERM → SIGKILL escalation - Fix test compatibility issues with VSCode mocks Fixes #5397 --- src/core/process/ProcessRegistry.ts | 240 ++++++++++++++++++ src/core/task/__tests__/Task.spec.ts | 17 ++ src/extension.ts | 12 + .../terminal/ExecaTerminalProcess.ts | 34 ++- src/integrations/terminal/TerminalRegistry.ts | 17 +- .../__tests__/ExecaTerminalProcess.spec.ts | 46 +++- src/services/mcp/McpHub.ts | 26 ++ src/services/mcp/__tests__/McpHub.spec.ts | 17 +- 8 files changed, 396 insertions(+), 13 deletions(-) create mode 100644 src/core/process/ProcessRegistry.ts diff --git a/src/core/process/ProcessRegistry.ts b/src/core/process/ProcessRegistry.ts new file mode 100644 index 0000000000..84ad53fc66 --- /dev/null +++ b/src/core/process/ProcessRegistry.ts @@ -0,0 +1,240 @@ +import * as vscode from "vscode" +import { ChildProcess } from "child_process" +import psTree from "ps-tree" + +/** + * Interface for tracking spawned processes + */ +export interface TrackedProcess { + /** The child process instance */ + process: ChildProcess + /** Unique identifier for the process */ + id: string + /** Optional description of what this process is for */ + description?: string + /** Debug session ID if this process is associated with a debug session */ + debugSessionId?: string + /** Timestamp when the process was registered */ + registeredAt: number +} + +/** + * Central registry for tracking all spawned processes to ensure proper cleanup + * during debug session termination and extension deactivation. + */ +export class ProcessRegistry implements vscode.Disposable { + private processes = new Map() + private debugSessionProcesses = new Map>() + private disposables: vscode.Disposable[] = [] + + constructor() { + // Register for debug session events to track process lifecycle + // Only register if vscode.debug is available (not in test environment) + try { + if (vscode.debug && vscode.debug.onDidStartDebugSession) { + this.disposables.push( + vscode.debug.onDidStartDebugSession(this.onDebugSessionStart.bind(this)), + vscode.debug.onDidTerminateDebugSession(this.onDebugSessionTerminate.bind(this)), + ) + } + } catch (error) { + // Ignore errors in test environment where vscode.debug might not be available + } + } + + /** + * Register a process for tracking + */ + register( + process: ChildProcess, + id: string, + options?: { + description?: string + debugSessionId?: string + }, + ): void { + const trackedProcess: TrackedProcess = { + process, + id, + description: options?.description, + debugSessionId: options?.debugSessionId, + registeredAt: Date.now(), + } + + this.processes.set(id, trackedProcess) + + // Track debug session association + if (options?.debugSessionId) { + if (!this.debugSessionProcesses.has(options.debugSessionId)) { + this.debugSessionProcesses.set(options.debugSessionId, new Set()) + } + this.debugSessionProcesses.get(options.debugSessionId)!.add(id) + } + + // Clean up when process exits naturally + process.on("exit", () => { + this.unregister(id) + }) + } + + /** + * Unregister a process from tracking + */ + unregister(id: string): void { + const trackedProcess = this.processes.get(id) + if (trackedProcess) { + // Remove from debug session tracking + if (trackedProcess.debugSessionId) { + const sessionProcesses = this.debugSessionProcesses.get(trackedProcess.debugSessionId) + if (sessionProcesses) { + sessionProcesses.delete(id) + if (sessionProcesses.size === 0) { + this.debugSessionProcesses.delete(trackedProcess.debugSessionId) + } + } + } + this.processes.delete(id) + } + } + + /** + * Kill a specific process and its children + */ + async killProcess(id: string, signal: NodeJS.Signals = "SIGTERM"): Promise { + const trackedProcess = this.processes.get(id) + if (!trackedProcess || !trackedProcess.process.pid) { + return + } + + try { + await this.killProcessTree(trackedProcess.process.pid, signal) + } catch (error) { + console.warn(`Failed to kill process ${id}:`, error) + } finally { + this.unregister(id) + } + } + + /** + * Kill all processes associated with a debug session + */ + async killDebugSessionProcesses(debugSessionId: string): Promise { + const processIds = this.debugSessionProcesses.get(debugSessionId) + if (!processIds) { + return + } + + const killPromises = Array.from(processIds).map((id) => this.killProcess(id)) + await Promise.allSettled(killPromises) + } + + /** + * Kill all tracked processes + */ + async killAllProcesses(): Promise { + const killPromises = Array.from(this.processes.keys()).map((id) => this.killProcess(id)) + await Promise.allSettled(killPromises) + } + + /** + * Get information about all tracked processes + */ + getTrackedProcesses(): TrackedProcess[] { + return Array.from(this.processes.values()) + } + + /** + * Get processes associated with a specific debug session + */ + getDebugSessionProcesses(debugSessionId: string): TrackedProcess[] { + const processIds = this.debugSessionProcesses.get(debugSessionId) + if (!processIds) { + return [] + } + + return Array.from(processIds) + .map((id) => this.processes.get(id)) + .filter((process): process is TrackedProcess => process !== undefined) + } + + private onDebugSessionStart(session: vscode.DebugSession): void { + // Debug session started - we'll track processes as they're spawned + console.log(`Debug session started: ${session.id}`) + } + + private async onDebugSessionTerminate(session: vscode.DebugSession): Promise { + // Debug session terminated - clean up all associated processes + console.log(`Debug session terminated: ${session.id}, cleaning up processes`) + await this.killDebugSessionProcesses(session.id) + } + + /** + * Kill a process tree using ps-tree with timeout-based escalation + */ + private async killProcessTree(pid: number, signal: NodeJS.Signals = "SIGTERM"): Promise { + return new Promise((resolve) => { + // First, try to get the process tree + psTree(pid, (err, children) => { + if (err) { + // Process might already be dead + resolve() + return + } + + // Kill all child processes first + const childPids = children.map((p) => parseInt(p.PID)) + childPids.forEach((childPid) => { + try { + process.kill(childPid, signal) + } catch (error) { + // Process might already be dead + } + }) + + // Kill the main process + try { + process.kill(pid, signal) + } catch (error) { + // Process might already be dead + } + + // If using SIGTERM, set up escalation to SIGKILL after timeout + if (signal === "SIGTERM") { + setTimeout(() => { + // Escalate to SIGKILL if process is still alive + try { + process.kill(pid, "SIGKILL") + childPids.forEach((childPid) => { + try { + process.kill(childPid, "SIGKILL") + } catch (error) { + // Process might already be dead + } + }) + } catch (error) { + // Process is already dead + } + resolve() + }, 5000) // 5 second timeout before escalating to SIGKILL + } else { + resolve() + } + }) + }) + } + + dispose(): void { + // Clean up all processes and event listeners + this.killAllProcesses().catch((error) => { + console.warn("Error during ProcessRegistry disposal:", error) + }) + + this.disposables.forEach((disposable) => disposable.dispose()) + this.disposables = [] + this.processes.clear() + this.debugSessionProcesses.clear() + } +} + +// Global instance for the extension +export const processRegistry = new ProcessRegistry() diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 693f72d1c7..624538649f 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -116,8 +116,25 @@ vi.mock("vscode", () => { stat: vi.fn().mockResolvedValue({ type: 1 }), // FileType.File = 1 }, onDidSaveTextDocument: vi.fn(() => mockDisposable), + onDidChangeWorkspaceFolders: vi.fn(() => mockDisposable), getConfiguration: vi.fn(() => ({ get: (key: string, defaultValue: any) => defaultValue })), }, + Uri: { + file: vi.fn().mockImplementation((path: string) => ({ + scheme: "file", + authority: "", + path: path, + query: "", + fragment: "", + fsPath: path, + with: vi.fn(), + toJSON: vi.fn(), + })), + }, + RelativePattern: vi.fn().mockImplementation((base: string, pattern: string) => ({ + base, + pattern, + })), env: { uriScheme: "vscode", language: "en", diff --git a/src/extension.ts b/src/extension.ts index bd43bcbf8a..60f52599bf 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -30,6 +30,7 @@ import { MdmService } from "./services/mdm/MdmService" import { migrateSettings } from "./utils/migrateSettings" import { autoImportSettings } from "./utils/autoImportSettings" import { API } from "./extension/api" +import { processRegistry } from "./core/process/ProcessRegistry" import { handleUri, @@ -214,6 +215,17 @@ export async function activate(context: vscode.ExtensionContext) { // This method is called when your extension is deactivated. export async function deactivate() { outputChannel.appendLine(`${Package.name} extension deactivated`) + + // Clean up all tracked processes first to prevent leaks + try { + outputChannel.appendLine("Cleaning up tracked processes...") + await processRegistry.killAllProcesses() + processRegistry.dispose() + outputChannel.appendLine("Process cleanup completed") + } catch (error) { + outputChannel.appendLine(`Error during process cleanup: ${error}`) + } + await McpServerManager.cleanup(extensionContext) TelemetryService.instance.shutdown() TerminalRegistry.cleanup() diff --git a/src/integrations/terminal/ExecaTerminalProcess.ts b/src/integrations/terminal/ExecaTerminalProcess.ts index 2b9b97d10a..7ead1304b4 100644 --- a/src/integrations/terminal/ExecaTerminalProcess.ts +++ b/src/integrations/terminal/ExecaTerminalProcess.ts @@ -1,14 +1,17 @@ import { execa, ExecaError } from "execa" import psTree from "ps-tree" import process from "process" +import * as vscode from "vscode" import type { RooTerminal } from "./types" import { BaseTerminalProcess } from "./BaseTerminalProcess" +import { processRegistry } from "../../core/process/ProcessRegistry" export class ExecaTerminalProcess extends BaseTerminalProcess { private terminalRef: WeakRef private aborted = false private pid?: number + private processId?: string constructor(terminal: RooTerminal) { super() @@ -49,6 +52,22 @@ export class ExecaTerminalProcess extends BaseTerminalProcess { })`${command}` this.pid = subprocess.pid + + // Register the process with the ProcessRegistry for cleanup tracking + if (subprocess.pid) { + this.processId = `execa-${subprocess.pid}-${Date.now()}` + try { + const currentDebugSession = vscode.debug.activeDebugSession + processRegistry.register(subprocess, this.processId, { + description: `Terminal command: ${command}`, + debugSessionId: currentDebugSession?.id, + }) + } catch (error) { + // In test environment, vscode.debug might not be available + console.warn(`[ExecaTerminalProcess] Failed to register process: ${error}`) + } + } + const stream = subprocess.iterable({ from: "all", preserveNewlines: true }) this.terminal.setActiveStream(stream, subprocess.pid) @@ -111,6 +130,13 @@ export class ExecaTerminalProcess extends BaseTerminalProcess { this.terminal.setActiveStream(undefined) this.emitRemainingBufferIfListening() this.stopHotTimer() + + // Unregister the process from the ProcessRegistry + if (this.processId) { + processRegistry.unregister(this.processId) + this.processId = undefined + } + this.emit("completed", this.fullOutput) this.emit("continue") } @@ -124,7 +150,13 @@ export class ExecaTerminalProcess extends BaseTerminalProcess { public override abort() { this.aborted = true - if (this.pid) { + // Use ProcessRegistry for cleanup if available, otherwise fall back to manual cleanup + if (this.processId) { + processRegistry.killProcess(this.processId, "SIGINT").catch((error) => { + console.warn(`[ExecaTerminalProcess] Failed to kill process via registry: ${error}`) + }) + } else if (this.pid) { + // Fallback to manual cleanup for processes not registered psTree(this.pid, async (err, children) => { if (!err) { const pids = children.map((p) => parseInt(p.PID)) diff --git a/src/integrations/terminal/TerminalRegistry.ts b/src/integrations/terminal/TerminalRegistry.ts index af334611c3..d9351dbebd 100644 --- a/src/integrations/terminal/TerminalRegistry.ts +++ b/src/integrations/terminal/TerminalRegistry.ts @@ -1,6 +1,7 @@ import * as vscode from "vscode" import { arePathsEqual } from "../../utils/path" +import { processRegistry } from "../../core/process/ProcessRegistry" import { RooTerminal, RooTerminalProvider } from "./types" import { TerminalProcess } from "./TerminalProcess" @@ -34,12 +35,18 @@ export class TerminalRegistry { // should probably live elsewhere. // Register handler for terminal close events to clean up temporary - // directories. + // directories and processes. const closeDisposable = vscode.window.onDidCloseTerminal((vsceTerminal) => { const terminal = this.getTerminalByVSCETerminal(vsceTerminal) if (terminal) { ShellIntegrationManager.zshCleanupTmpDir(terminal.id) + + // Clean up any processes associated with this terminal + // For ExecaTerminal processes, we need to abort them which will handle cleanup + if (terminal.process) { + terminal.process.abort() + } } }) @@ -280,6 +287,14 @@ export class TerminalRegistry { public static cleanup() { // Clean up all temporary directories. ShellIntegrationManager.clear() + + // Clean up all terminal processes + this.terminals.forEach((terminal) => { + if (terminal.process) { + terminal.process.abort() + } + }) + this.disposables.forEach((disposable) => disposable.dispose()) this.disposables = [] } diff --git a/src/integrations/terminal/__tests__/ExecaTerminalProcess.spec.ts b/src/integrations/terminal/__tests__/ExecaTerminalProcess.spec.ts index 873b8f85ab..0e36507e00 100644 --- a/src/integrations/terminal/__tests__/ExecaTerminalProcess.spec.ts +++ b/src/integrations/terminal/__tests__/ExecaTerminalProcess.spec.ts @@ -1,18 +1,29 @@ // npx vitest run integrations/terminal/__tests__/ExecaTerminalProcess.spec.ts -const mockPid = 12345 - vitest.mock("execa", () => { + const mockPid = 12345 const mockKill = vitest.fn() + const execa = vitest.fn((options: any) => { - return (_template: TemplateStringsArray, ...args: any[]) => ({ - pid: mockPid, - iterable: (_opts: any) => - (async function* () { - yield "test output\n" - })(), - kill: mockKill, - }) + return (_template: TemplateStringsArray, ...args: any[]) => { + const mockSubprocess = { + pid: mockPid, + iterable: vitest.fn((_opts: any) => + (async function* () { + yield "test output\n" + })(), + ), + kill: mockKill, + // Add Promise-like behavior to simulate successful completion + then: vitest.fn((resolve) => { + // Simulate successful completion + resolve({ exitCode: 0 }) + return Promise.resolve({ exitCode: 0 }) + }), + catch: vitest.fn(), + } + return mockSubprocess + } }) return { execa, ExecaError: class extends Error {} } }) @@ -21,11 +32,26 @@ vitest.mock("ps-tree", () => ({ default: vitest.fn((_: number, cb: any) => cb(null, [])), })) +vitest.mock("vscode", () => ({ + debug: { + activeDebugSession: undefined, + }, +})) + +vitest.mock("../../core/process/ProcessRegistry", () => ({ + processRegistry: { + register: vitest.fn(), + unregister: vitest.fn(), + killProcess: vitest.fn().mockResolvedValue(undefined), + }, +})) + import { execa } from "execa" import { ExecaTerminalProcess } from "../ExecaTerminalProcess" import type { RooTerminal } from "../types" describe("ExecaTerminalProcess", () => { + const mockPid = 12345 let mockTerminal: RooTerminal let terminalProcess: ExecaTerminalProcess let originalEnv: NodeJS.ProcessEnv diff --git a/src/services/mcp/McpHub.ts b/src/services/mcp/McpHub.ts index 10a74712ef..ab973601ad 100644 --- a/src/services/mcp/McpHub.ts +++ b/src/services/mcp/McpHub.ts @@ -32,11 +32,13 @@ import { import { fileExistsAtPath } from "../../utils/fs" import { arePathsEqual } from "../../utils/path" import { injectVariables } from "../../utils/config" +import { processRegistry } from "../../core/process/ProcessRegistry" export type McpConnection = { server: McpServer client: Client transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport + processId?: string // Track process ID for cleanup } // Base configuration schema for common settings @@ -748,6 +750,21 @@ export class McpHub { } this.connections.push(connection) + // Register the MCP server process with ProcessRegistry for cleanup tracking (stdio only) + if (configInjected.type === "stdio") { + const childProcess = (transport as any).process + if (childProcess && childProcess.pid) { + const processId = `mcp-${name}-${childProcess.pid}-${Date.now()}` + const currentDebugSession = vscode.debug.activeDebugSession + processRegistry.register(childProcess, processId, { + description: `MCP Server: ${name}`, + debugSessionId: currentDebugSession?.id, + }) + // Store the process ID in the connection for cleanup + connection.processId = processId + } + } + // Connect (this will automatically start the transport) await client.connect(transport) connection.server.status = "connected" @@ -920,6 +937,11 @@ export class McpHub { for (const connection of connections) { try { + // Unregister the process from ProcessRegistry if it was registered + if (connection.processId) { + processRegistry.unregister(connection.processId) + } + await connection.transport.close() await connection.client.close() } catch (error) { @@ -1627,6 +1649,10 @@ export class McpHub { this.removeAllFileWatchers() for (const connection of this.connections) { try { + // Ensure process cleanup happens during disposal + if (connection.processId) { + await processRegistry.killProcess(connection.processId) + } await this.deleteConnection(connection.server.name, connection.server.source) } catch (error) { console.error(`Failed to close connection for ${connection.server.name}:`, error) diff --git a/src/services/mcp/__tests__/McpHub.spec.ts b/src/services/mcp/__tests__/McpHub.spec.ts index 98ef4514c2..c72c706388 100644 --- a/src/services/mcp/__tests__/McpHub.spec.ts +++ b/src/services/mcp/__tests__/McpHub.spec.ts @@ -62,6 +62,22 @@ vi.mock("vscode", () => ({ dispose: vi.fn(), }), }, + Uri: { + file: vi.fn().mockImplementation((path: string) => ({ + scheme: "file", + authority: "", + path: path, + query: "", + fragment: "", + fsPath: path, + with: vi.fn(), + toJSON: vi.fn(), + })), + }, + RelativePattern: vi.fn().mockImplementation((base: string, pattern: string) => ({ + base, + pattern, + })), Disposable: { from: vi.fn(), }, @@ -93,7 +109,6 @@ describe("McpHub", () => { // Mock console.error to suppress error messages during tests console.error = vi.fn() - const mockUri: Uri = { scheme: "file", authority: "",