diff --git a/commit_message.md b/commit_message.md new file mode 100644 index 0000000..1121a3c --- /dev/null +++ b/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. \ No newline at end of file diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/packages/agent/src/core/backgroundTools.cleanup.test.ts new file mode 100644 index 0000000..3adec5d --- /dev/null +++ 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 }; + state: { completed: boolean }; + command?: string; + stdout?: string[]; + stderr?: string[]; + showStdIn?: boolean; + showStdout?: boolean; +}; + +type MockAgentState = { + aborted: boolean; + completed: boolean; + context: { + backgroundTools: { + cleanup: ReturnType; + }; + }; + 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(), + }; +}); + +vi.mock('../tools/interaction/agentStart.js', () => { + return { + agentStates: new Map(), + }; +}); + +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 + ).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 index 8e034c3..45c61c1 100644 --- 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,126 @@ 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 { + 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 { + 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 { + 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((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 { + 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 index 00a5ff4..fd3e773 100644 --- a/packages/agent/src/tools/interaction/agentMessage.ts +++ b/packages/agent/src/tools/interaction/agentMessage.ts @@ -86,6 +86,9 @@ export const agentMessageTool: Tool = { }, ); + // 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 index 8b06295..da25239 100644 --- a/packages/agent/src/tools/interaction/agentStart.ts +++ b/packages/agent/src/tools/interaction/agentStart.ts @@ -162,6 +162,9 @@ export const agentStartTool: Tool = { (result.result.length > 100 ? '...' : ''), }, ); + + // Clean up resources when agent completes successfully + await backgroundTools.cleanup(); } } catch (error) { // Update agent state with the error @@ -178,6 +181,9 @@ export const agentStartTool: Tool = { error: error instanceof Error ? error.message : String(error), }, ); + + // Clean up resources when agent encounters an error + await backgroundTools.cleanup(); } } return true; diff --git a/packages/cli/src/commands/$default.ts b/packages/cli/src/commands/$default.ts index 3680d8c..a56fa47 100644 --- a/packages/cli/src/commands/$default.ts +++ b/packages/cli/src/commands/$default.ts @@ -109,6 +109,8 @@ export const command: CommandModule = { // Use command line option if provided, otherwise use config value tokenTracker.tokenCache = config.tokenCache; + const backgroundTools = new BackgroundTools('mainAgent'); + try { // Early API key check based on model provider const providerSettings = @@ -209,7 +211,7 @@ export const command: CommandModule = { model: config.model, maxTokens: config.maxTokens, temperature: config.temperature, - backgroundTools: new BackgroundTools('mainAgent'), + backgroundTools, }); const output = @@ -225,6 +227,8 @@ export const command: CommandModule = { ); // Capture the error with Sentry captureException(error); + } finally { + await backgroundTools.cleanup(); } logger.log( diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index b3fa6e8..44f3ef8 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -1,4 +1,5 @@ import { BrowserManager, processStates } from 'mycoder-agent'; +import { agentStates } from 'mycoder-agent/dist/tools/interaction/agentStart.js'; /** * Handles cleanup of resources before application exit @@ -7,12 +8,43 @@ import { BrowserManager, processStates } from 'mycoder-agent'; export async function cleanupResources(): Promise { console.log('Cleaning up resources before exit...'); + // First attempt to clean up any still-running agents + // This will cascade to their browser sessions and shell processes + try { + // Find all active agent instances + const activeAgents = Array.from(agentStates.entries()).filter( + ([_, state]) => !state.completed && !state.aborted, + ); + + if (activeAgents.length > 0) { + console.log(`Cleaning up ${activeAgents.length} active agents...`); + + for (const [id, state] of activeAgents) { + try { + // Mark the agent as aborted + state.aborted = true; + state.completed = true; + + // Clean up its resources + await state.context.backgroundTools.cleanup(); + } catch (error) { + console.error(`Error cleaning up agent ${id}:`, error); + } + } + } + } catch (error) { + console.error('Error cleaning up agents:', error); + } + + // As a fallback, still clean up any browser sessions and shell processes + // that might not have been caught by the agent cleanup + // 1. Clean up browser sessions try { // Get the BrowserManager instance - this is a singleton - const browserManager = (globalThis as any).__BROWSER_MANAGER__ as - | BrowserManager - | undefined; + const browserManager = ( + globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager } + ).__BROWSER_MANAGER__; if (browserManager) { console.log('Closing all browser sessions...'); await browserManager.closeAllSessions();