Implement Incremental Resource Cleanup for Agent Lifecycle Management

commit_message.md

@@ -0,0 +1,12 @@
+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.

packages/agent/src/core/backgroundTools.cleanup.test.ts

@@ -0,0 +1,223 @@
+import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
+
+// Import mocked modules
+import { BrowserManager } from '../tools/browser/BrowserManager.js';
+import { agentStates } from '../tools/interaction/agentStart.js';
+import { processStates } from '../tools/system/shellStart.js';
+
+import { BackgroundTools, BackgroundToolStatus } from './backgroundTools';
+import { Tool } from './types';
+
+// Define types for our mocks that match the actual types
+type MockProcessState = {
+  process: { kill: ReturnType<typeof vi.fn> };
+  state: { completed: boolean };
+  command?: string;
+  stdout?: string[];
+  stderr?: string[];
+  showStdIn?: boolean;
+  showStdout?: boolean;
+};
+
+type MockAgentState = {
+  aborted: boolean;
+  completed: boolean;
+  context: {
+    backgroundTools: {
+      cleanup: ReturnType<typeof vi.fn>;
+    };
+  };
+  goal?: string;
+  prompt?: string;
+  output?: string;
+  workingDirectory?: string;
+  tools?: Tool[];
+};
+
+// Mock dependencies
+vi.mock('../tools/browser/BrowserManager.js', () => {
+  return {
+    BrowserManager: class MockBrowserManager {
+      closeSession = vi.fn().mockResolvedValue(undefined);
+    },
+  };
+});
+
+vi.mock('../tools/system/shellStart.js', () => {
+  return {
+    processStates: new Map<string, MockProcessState>(),
+  };
+});
+
+vi.mock('../tools/interaction/agentStart.js', () => {
+  return {
+    agentStates: new Map<string, MockAgentState>(),
+  };
+});
+
+describe('BackgroundTools cleanup', () => {
+  let backgroundTools: BackgroundTools;
+
+  // Setup mocks for globalThis and process states
+  beforeEach(() => {
+    backgroundTools = new BackgroundTools('test-agent');
+
+    // Reset mocks
+    vi.resetAllMocks();
+
+    // Setup global browser manager
+    (
+      globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager }
+    ).__BROWSER_MANAGER__ = {
+      closeSession: vi.fn().mockResolvedValue(undefined),
+    } as unknown as BrowserManager;
+
+    // Setup mock process states
+    const mockProcess = {
+      kill: vi.fn(),
+    };
+
+    const mockProcessState: MockProcessState = {
+      process: mockProcess,
+      state: { completed: false },
+      command: 'test command',
+      stdout: [],
+      stderr: [],
+      showStdIn: false,
+      showStdout: false,
+    };
+
+    processStates.clear();
+    processStates.set('shell-1', mockProcessState as any);
+
+    // Setup mock agent states
+    const mockAgentState: MockAgentState = {
+      aborted: false,
+      completed: false,
+      context: {
+        backgroundTools: {
+          cleanup: vi.fn().mockResolvedValue(undefined),
+        },
+      },
+      goal: 'test goal',
+      prompt: 'test prompt',
+      output: '',
+      workingDirectory: '/test',
+      tools: [],
+    };
+
+    agentStates.clear();
+    agentStates.set('agent-1', mockAgentState as any);
+  });
+
+  afterEach(() => {
+    vi.resetAllMocks();
+
+    // Clear global browser manager
+    (
+      globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager }
+    ).__BROWSER_MANAGER__ = undefined;
+
+    // Clear mock states
+    processStates.clear();
+    agentStates.clear();
+  });
+
+  it('should clean up browser sessions', async () => {
+    // Register a browser tool
+    const browserId = backgroundTools.registerBrowser('https://example.com');
+
+    // Run cleanup
+    await backgroundTools.cleanup();
+
+    // Check that closeSession was called
+    expect(
+      (globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
+        .__BROWSER_MANAGER__.closeSession,
+    ).toHaveBeenCalledWith(browserId);
+
+    // Check that tool status was updated
+    const tool = backgroundTools.getToolById(browserId);
+    expect(tool?.status).toBe(BackgroundToolStatus.COMPLETED);
+  });
+
+  it('should clean up shell processes', async () => {
+    // Register a shell tool
+    const shellId = backgroundTools.registerShell('echo "test"');
+
+    // Get mock process state
+    const mockProcessState = processStates.get('shell-1');
+
+    // Set the shell ID to match
+    processStates.set(shellId, processStates.get('shell-1') as any);
+
+    // Run cleanup
+    await backgroundTools.cleanup();
+
+    // Check that kill was called
+    expect(mockProcessState?.process.kill).toHaveBeenCalledWith('SIGTERM');
+
+    // Check that tool status was updated
+    const tool = backgroundTools.getToolById(shellId);
+    expect(tool?.status).toBe(BackgroundToolStatus.COMPLETED);
+  });
+
+  it('should clean up sub-agents', async () => {
+    // Register an agent tool
+    const agentId = backgroundTools.registerAgent('Test goal');
+
+    // Get mock agent state
+    const mockAgentState = agentStates.get('agent-1');
+
+    // Set the agent ID to match
+    agentStates.set(agentId, agentStates.get('agent-1') as any);
+
+    // Run cleanup
+    await backgroundTools.cleanup();
+
+    // Check that agent was marked as aborted
+    expect(mockAgentState?.aborted).toBe(true);
+    expect(mockAgentState?.completed).toBe(true);
+
+    // Check that cleanup was called on the agent's background tools
+    expect(mockAgentState?.context.backgroundTools.cleanup).toHaveBeenCalled();
+
+    // Check that tool status was updated
+    const tool = backgroundTools.getToolById(agentId);
+    expect(tool?.status).toBe(BackgroundToolStatus.TERMINATED);
+  });
+
+  it('should handle errors during cleanup', async () => {
+    // Register a browser tool
+    const browserId = backgroundTools.registerBrowser('https://example.com');
+
+    // Make closeSession throw an error
+    (
+      (globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
+        .__BROWSER_MANAGER__.closeSession as ReturnType<typeof vi.fn>
+    ).mockRejectedValue(new Error('Test error'));
+
+    // Run cleanup
+    await backgroundTools.cleanup();
+
+    // Check that tool status was updated to ERROR
+    const tool = backgroundTools.getToolById(browserId);
+    expect(tool?.status).toBe(BackgroundToolStatus.ERROR);
+    expect(tool?.metadata.error).toBe('Test error');
+  });
+
+  it('should only clean up running tools', async () => {
+    // Register a browser tool and mark it as completed
+    const browserId = backgroundTools.registerBrowser('https://example.com');
+    backgroundTools.updateToolStatus(browserId, BackgroundToolStatus.COMPLETED);
+
+    // Run cleanup
+    await backgroundTools.cleanup();
+
+    // Check that closeSession was not called
+    expect(
+      (globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
+        .__BROWSER_MANAGER__.closeSession,
+    ).not.toHaveBeenCalled();
+  });
+});

packages/agent/src/core/backgroundTools.ts

@@ -1,5 +1,10 @@
 import { v4 as uuidv4 } from 'uuid';
 
+// These imports will be used by the cleanup method
+import { BrowserManager } from '../tools/browser/BrowserManager.js';
+import { agentStates } from '../tools/interaction/agentStart.js';
+import { processStates } from '../tools/system/shellStart.js';
+
 // Types of background processes we can track
 export enum BackgroundToolType {
   SHELL = 'shell',
@@ -22,7 +27,7 @@
   status: BackgroundToolStatus;
   startTime: Date;
   endTime?: Date;
   metadata: Record<string, any>; // Additional tool-specific information
 }
 
 // Shell process specific data
@@ -32,6 +37,7 @@
     command: string;
     exitCode?: number | null;
     signaled?: boolean;
+    error?: string;
   };
 }
 
@@ -40,6 +46,7 @@
   type: BackgroundToolType.BROWSER;
   metadata: {
     url?: string;
+    error?: string;
   };
 }
 
@@ -48,6 +55,7 @@
   type: BackgroundToolType.AGENT;
   metadata: {
     goal?: string;
+    error?: string;
   };
 }
 
@@ -118,7 +126,7 @@
   public updateToolStatus(
     id: string,
     status: BackgroundToolStatus,
     metadata?: Record<string, any>,
   ): boolean {
     const tool = this.tools.get(id);
     if (!tool) {
@@ -154,4 +162,126 @@
   public getToolById(id: string): AnyBackgroundTool | undefined {
     return this.tools.get(id);
   }
+
+  /**
+   * Cleans up all resources associated with this agent instance
+   * @returns A promise that resolves when cleanup is complete
+   */
+  public async cleanup(): Promise<void> {
+    const tools = this.getTools();
+
+    // Group tools by type for better cleanup organization
+    const browserTools = tools.filter(
+      (tool): tool is BrowserBackgroundTool =>
+        tool.type === BackgroundToolType.BROWSER &&
+        tool.status === BackgroundToolStatus.RUNNING,
+    );
+
+    const shellTools = tools.filter(
+      (tool): tool is ShellBackgroundTool =>
+        tool.type === BackgroundToolType.SHELL &&
+        tool.status === BackgroundToolStatus.RUNNING,
+    );
+
+    const agentTools = tools.filter(
+      (tool): tool is AgentBackgroundTool =>
+        tool.type === BackgroundToolType.AGENT &&
+        tool.status === BackgroundToolStatus.RUNNING,
+    );
+
+    // 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),
+      });
+    }
+  }
+
+  /**
+   * 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),
+      });
+    }
+  }
+
+  /**
+   * 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),
+      });
+    }
+  }
 }