feat(agent): implement incremental resource cleanup for agent lifecycle

bhouston · bhouston · commit 576436ef2c7c · 2025-03-12T15:22:37.000-04:00
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. The changes include: - 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
diff --git a/commit_message.md b/commit_message.md
@@ -0,0 +1,16 @@
+feat(agent): implement incremental resource cleanup for agent lifecycle
+
+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.
+
+The changes include:
+
+- 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
diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/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();
+  });
+});
diff --git a/packages/agent/src/core/backgroundTools.ts b/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',
@@ -32,6 +37,7 @@ export interface ShellBackgroundTool extends BackgroundTool {
     command: string;
     exitCode?: number | null;
     signaled?: boolean;
+    error?: string;
   };
 }
 
@@ -40,6 +46,7 @@ export interface BrowserBackgroundTool extends BackgroundTool {
   type: BackgroundToolType.BROWSER;
   metadata: {
     url?: string;
+    error?: string;
   };
 }
 
@@ -48,6 +55,7 @@ export interface AgentBackgroundTool extends BackgroundTool {
   type: BackgroundToolType.AGENT;
   metadata: {
     goal?: string;
+    error?: string;
   };
 }
 
@@ -154,4 +162,100 @@ export class BackgroundTools {
   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,
+    );
+
+    const shellTools = tools.filter(
+      (tool): tool is ShellBackgroundTool =>
+        tool.type === BackgroundToolType.SHELL,
+    );
+
+    const agentTools = tools.filter(
+      (tool): tool is AgentBackgroundTool =>
+        tool.type === BackgroundToolType.AGENT,
+    );
+
+    // 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),
+          });
+        }
+      }
+    }
+
+    // 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');
+
+            // 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;
+
+            // 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),
+          });
+        }
+      }
+    }
+  }
 }
diff --git a/packages/agent/src/tools/interaction/agentMessage.ts b/packages/agent/src/tools/interaction/agentMessage.ts
@@ -86,6 +86,9 @@ export const agentMessageTool: Tool<Parameters, ReturnType> = {
           },
         );
 
+        // Clean up resources when agent is terminated
+        await backgroundTools.cleanup();
+
         return {
           output: agentState.output || 'Sub-agent terminated before completion',
           completed: true,
diff --git a/packages/agent/src/tools/interaction/agentStart.ts b/packages/agent/src/tools/interaction/agentStart.ts
diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts
