From 576436ef2c7c5f234f088b7dba2e7fd65590738f Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 12 Mar 2025 15:22:37 -0400 Subject: [PATCH 1/3] 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 --- commit_message.md | 16 ++ .../src/core/backgroundTools.cleanup.test.ts | 223 ++++++++++++++++++ packages/agent/src/core/backgroundTools.ts | 104 ++++++++ .../src/tools/interaction/agentMessage.ts | 3 + .../agent/src/tools/interaction/agentStart.ts | 6 + packages/cli/src/utils/cleanup.ts | 38 ++- 6 files changed, 387 insertions(+), 3 deletions(-) create mode 100644 commit_message.md create mode 100644 packages/agent/src/core/backgroundTools.cleanup.test.ts diff --git a/commit_message.md b/commit_message.md new file mode 100644 index 0000000..46b0288 --- /dev/null +++ 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 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..b8f9614 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,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 { + 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((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 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/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(); From 098bc288d974b34f775e97f2a5ce8f75059bb0b0 Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 12 Mar 2025 15:28:55 -0400 Subject: [PATCH 2/3] 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. --- commit_message.md | 22 ++- packages/agent/src/core/backgroundTools.ts | 154 ++++++++++++--------- 2 files changed, 99 insertions(+), 77 deletions(-) diff --git a/commit_message.md b/commit_message.md index 46b0288..1121a3c 100644 --- a/commit_message.md +++ b/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. \ No newline at end of file diff --git a/packages/agent/src/core/backgroundTools.ts b/packages/agent/src/core/backgroundTools.ts index b8f9614..45c61c1 100644 --- a/packages/agent/src/core/backgroundTools.ts +++ b/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 { + 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 { + 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), - }); - } + // 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), + }); } + } - // 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 { + 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), + }); } } } From 123bb227457c935e0087377f00b890a924f786e4 Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 12 Mar 2025 15:34:49 -0400 Subject: [PATCH 3/3] add cleanup in default command. --- packages/cli/src/commands/$default.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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(