refactor(agent): implement parallel resource cleanup

This change improves the resource cleanup process by handling browser sessions,
shell processes, and sub-agents in parallel rather than sequentially.

The implementation:

1. Refactors the cleanup method into smaller helper methods for each resource type
2. Uses Promise.all to execute cleanup operations concurrently
3. Filters tools by status during the initial grouping to simplify the code

This approach significantly speeds up the cleanup process, especially when
dealing with multiple resources of different types.

Author: bhouston

commit_message.md

@@ -1,16 +1,12 @@
-feat(agent): implement incremental resource cleanup for agent lifecycle
+refactor(agent): implement parallel resource cleanup
 
-This commit adds a new cleanup method to the BackgroundTools class that handles
-cleaning up browser sessions, shell processes, and sub-agents associated with an
-agent when it completes its task, encounters an error, or is terminated.
+This change improves the resource cleanup process by handling browser sessions,
+shell processes, and sub-agents in parallel rather than sequentially.
 
-The changes include:
+The implementation:
+1. Refactors the cleanup method into smaller helper methods for each resource type
+2. Uses Promise.all to execute cleanup operations concurrently
+3. Filters tools by status during the initial grouping to simplify the code
 
-- Adding a cleanup method to BackgroundTools that cleans up resources
-- Calling cleanup when agents complete successfully
-- Calling cleanup when agents encounter errors
-- Calling cleanup when agents are terminated
-- Enhancing global cleanup to first attempt to clean up any still-running agents
-- Adding tests for the new cleanup functionality
-
-Fixes #236
+This approach significantly speeds up the cleanup process, especially when
+dealing with multiple resources of different types.

packages/agent/src/core/backgroundTools.ts

@@ -173,89 +173,115 @@ export class BackgroundTools {
     // Group tools by type for better cleanup organization
     const browserTools = tools.filter(
       (tool): tool is BrowserBackgroundTool =>
-        tool.type === BackgroundToolType.BROWSER,
+        tool.type === BackgroundToolType.BROWSER &&
+        tool.status === BackgroundToolStatus.RUNNING,
     );
 
     const shellTools = tools.filter(
       (tool): tool is ShellBackgroundTool =>
-        tool.type === BackgroundToolType.SHELL,
+        tool.type === BackgroundToolType.SHELL &&
+        tool.status === BackgroundToolStatus.RUNNING,
     );
 
     const agentTools = tools.filter(
       (tool): tool is AgentBackgroundTool =>
-        tool.type === BackgroundToolType.AGENT,
+        tool.type === BackgroundToolType.AGENT &&
+        tool.status === BackgroundToolStatus.RUNNING,
     );
 
-    // Clean up browser sessions
-    for (const tool of browserTools) {
-      if (tool.status === BackgroundToolStatus.RUNNING) {
-        try {
-          const browserManager = (
-            globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager }
-          ).__BROWSER_MANAGER__;
-          if (browserManager) {
-            await browserManager.closeSession(tool.id);
-          }
-          this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
-        } catch (error) {
-          this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
-            error: error instanceof Error ? error.message : String(error),
-          });
-        }
+    // Create cleanup promises for each resource type
+    const browserCleanupPromises = browserTools.map((tool) =>
+      this.cleanupBrowserSession(tool),
+    );
+    const shellCleanupPromises = shellTools.map((tool) =>
+      this.cleanupShellProcess(tool),
+    );
+    const agentCleanupPromises = agentTools.map((tool) =>
+      this.cleanupSubAgent(tool),
+    );
+
+    // Wait for all cleanup operations to complete in parallel
+    await Promise.all([
+      ...browserCleanupPromises,
+      ...shellCleanupPromises,
+      ...agentCleanupPromises,
+    ]);
+  }
+
+  /**
+   * Cleans up a browser session
+   * @param tool The browser tool to clean up
+   */
+  private async cleanupBrowserSession(
+    tool: BrowserBackgroundTool,
+  ): Promise<void> {
+    try {
+      const browserManager = (
+        globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager }
+      ).__BROWSER_MANAGER__;
+      if (browserManager) {
+        await browserManager.closeSession(tool.id);
       }
+      this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
+    } catch (error) {
+      this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
+        error: error instanceof Error ? error.message : String(error),
+      });
     }
+  }
 
-    // Clean up shell processes
-    for (const tool of shellTools) {
-      if (tool.status === BackgroundToolStatus.RUNNING) {
-        try {
-          const processState = processStates.get(tool.id);
-          if (processState && !processState.state.completed) {
-            processState.process.kill('SIGTERM');
+  /**
+   * Cleans up a shell process
+   * @param tool The shell tool to clean up
+   */
+  private async cleanupShellProcess(tool: ShellBackgroundTool): Promise<void> {
+    try {
+      const processState = processStates.get(tool.id);
+      if (processState && !processState.state.completed) {
+        processState.process.kill('SIGTERM');
 
-            // Force kill after a short timeout if still running
-            await new Promise<void>((resolve) => {
-              setTimeout(() => {
-                try {
-                  if (!processState.state.completed) {
-                    processState.process.kill('SIGKILL');
-                  }
-                } catch {
-                  // Ignore errors on forced kill
-                }
-                resolve();
-              }, 500);
-            });
-          }
-          this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
-        } catch (error) {
-          this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
-            error: error instanceof Error ? error.message : String(error),
-          });
-        }
+        // Force kill after a short timeout if still running
+        await new Promise<void>((resolve) => {
+          setTimeout(() => {
+            try {
+              if (!processState.state.completed) {
+                processState.process.kill('SIGKILL');
+              }
+            } catch {
+              // Ignore errors on forced kill
+            }
+            resolve();
+          }, 500);
+        });
       }
+      this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
+    } catch (error) {
+      this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
+        error: error instanceof Error ? error.message : String(error),
+      });
     }
+  }
 
-    // Clean up sub-agents
-    for (const tool of agentTools) {
-      if (tool.status === BackgroundToolStatus.RUNNING) {
-        try {
-          const agentState = agentStates.get(tool.id);
-          if (agentState && !agentState.aborted) {
-            // Set the agent as aborted and completed
-            agentState.aborted = true;
-            agentState.completed = true;
+  /**
+   * Cleans up a sub-agent
+   * @param tool The agent tool to clean up
+   */
+  private async cleanupSubAgent(tool: AgentBackgroundTool): Promise<void> {
+    try {
+      const agentState = agentStates.get(tool.id);
+      if (agentState && !agentState.aborted) {
+        // Set the agent as aborted and completed
+        agentState.aborted = true;
+        agentState.completed = true;
 
-            // Clean up resources owned by this sub-agent
-            await agentState.context.backgroundTools.cleanup();
-          }
-          this.updateToolStatus(tool.id, BackgroundToolStatus.TERMINATED);
-        } catch (error) {
-          this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
-            error: error instanceof Error ? error.message : String(error),
-          });
-        }
+        // Clean up resources owned by this sub-agent
+        await agentState.context.backgroundTools.cleanup();
       }
+      this.updateToolStatus(tool.id, BackgroundToolStatus.TERMINATED);
+    } catch (error) {
+      this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
+        error: error instanceof Error ? error.message : String(error),
+      });
     }
   }
 }