drivecore · bhouston · Mar 12, 2025 · Mar 12, 2025
diff --git a/packages/agent/src/tools/browser/BrowserManager.ts b/packages/agent/src/tools/browser/BrowserManager.ts
@@ -14,6 +14,15 @@ export class BrowserManager {
     headless: true,
     defaultTimeout: 30000,
   };
+
+  constructor() {
+    // Store a reference to the instance globally for cleanup
+    // This allows the CLI to access the instance for cleanup
+    (globalThis as any).__BROWSER_MANAGER__ = this;
+
+    // Set up cleanup handlers for graceful shutdown
+    this.setupGlobalCleanup();
+  }
 
   async createSession(config?: BrowserConfig): Promise<BrowserSession> {
     try {
@@ -80,14 +89,41 @@ export class BrowserManager {
       this.sessions.delete(session.id);
     });
 
-    // Handle process exit
+    // No need to add individual process handlers for each session
+    // We'll handle all sessions in the global cleanup
+  }
+
+  /**
+   * Sets up global cleanup handlers for all browser sessions
+   */
+  private setupGlobalCleanup(): void {
+    // Use beforeExit for async cleanup
+    process.on('beforeExit', () => {
+      this.closeAllSessions().catch((err) => {
+        console.error('Error closing browser sessions:', err);
+      });
+    });
+
+    // Use exit for synchronous cleanup (as a fallback)
     process.on('exit', () => {
-      this.closeSession(session.id).catch(() => {});
+      // Can only do synchronous operations here
+      for (const session of this.sessions.values()) {
+        try {
+          // Attempt synchronous close - may not fully work
+          session.browser.close();
+        } catch (e) {
+          // Ignore errors during exit
+        }
+      }
     });
-
-    // Handle unexpected errors
-    process.on('uncaughtException', () => {
-      this.closeSession(session.id).catch(() => {});
+
+    // Handle SIGINT (Ctrl+C)
+    process.on('SIGINT', () => {
+      this.closeAllSessions().catch(() => {})
+        .finally(() => {
+          // Give a moment for cleanup to complete
+          setTimeout(() => process.exit(0), 500);
+        });
     });
   }
 

diff --git a/packages/agent/src/tools/system/shellStart.ts b/packages/agent/src/tools/system/shellStart.ts
@@ -25,6 +25,7 @@ type ProcessState = {
 };
 
 // Global map to store process state
+// This is exported so it can be accessed for cleanup
 export const processStates: Map<string, ProcessState> = new Map();
 
 const parameterSchema = z.object({

diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts
@@ -13,6 +13,7 @@ import { command as toolsCommand } from './commands/tools.js';
 import { sharedOptions } from './options.js';
 import { initSentry, captureException } from './sentry/index.js';
 import { getConfig } from './settings/config.js';
+import { cleanupResources, setupForceExit } from './utils/cleanup.js';
 import { enableProfiling, mark, reportTimings } from './utils/performance.js';
 
 mark('After imports');
@@ -82,4 +83,11 @@ await main()
   .finally(async () => {
     // Report timings if profiling is enabled
     await reportTimings();
+
+    // Clean up all resources before exit
+    await cleanupResources();
+
+    // Setup a force exit as a failsafe
+    // This ensures the process will exit even if there are lingering handles
+    setupForceExit(5000);
   });
diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts
@@ -0,0 +1,66 @@
+import { BrowserManager, processStates } from 'mycoder-agent';
+
+/**
+ * Handles cleanup of resources before application exit
+ * Ensures all browser sessions and shell processes are terminated
+ */
+export async function cleanupResources(): Promise<void> {
+  console.log('Cleaning up resources before exit...');
+
+  // 1. Clean up browser sessions
+  try {
+    // Get the BrowserManager instance - this is a singleton
+    const browserManager = (globalThis as any).__BROWSER_MANAGER__ as BrowserManager | undefined;
+    if (browserManager) {
+      console.log('Closing all browser sessions...');
+      await browserManager.closeAllSessions();
+    }
+  } catch (error) {
+    console.error('Error closing browser sessions:', error);
+  }
+
+  // 2. Clean up shell processes
+  try {
+    if (processStates.size > 0) {
+      console.log(`Terminating ${processStates.size} shell processes...`);
+      for (const [id, state] of processStates.entries()) {
+        if (!state.state.completed) {
+          console.log(`Terminating process ${id}...`);
+          try {
+            state.process.kill('SIGTERM');
+            // Force kill after a short timeout if still running
+            setTimeout(() => {
+              try {
+                if (!state.state.completed) {
+                  state.process.kill('SIGKILL');
+                }
+              } catch (e) {
+                // Ignore errors on forced kill
+              }
+            }, 500);
+          } catch (e) {
+            console.error(`Error terminating process ${id}:`, e);
+          }
+        }
+      }
+    }
+  } catch (error) {
+    console.error('Error terminating shell processes:', error);
+  }
+
+  // 3. Give async operations a moment to complete
+  await new Promise((resolve) => setTimeout(resolve, 1000));
+
+  console.log('Cleanup completed');
+}
+
+/**
+ * Force exits the process after a timeout
+ * This is a failsafe to ensure the process exits even if there are lingering handles
+ */
+export function setupForceExit(timeoutMs = 5000): void {
+  setTimeout(() => {
+    console.log(`Forcing exit after ${timeoutMs}ms timeout`);
+    process.exit(0);
+  }, timeoutMs);
+}