From fb89dd8503f1bb4c5181b5a9f580e0227a38411f Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Tue, 18 Mar 2025 11:00:51 -0400 Subject: [PATCH] refactor: replace background tools with scoped resource trackers This commit replaces the global background tools approach with individual resource trackers (AgentTracker, ShellTracker, and BrowserTracker) that are scoped to each agent instance. This improves encapsulation and resource management by ensuring that each agent is responsible for its own resources. - Remove backgroundTools.ts and related files - Refactor resource trackers to be scoped to the agent - Add cleanup methods to each tracker - Update tool implementations to use the new trackers Closes #305 --- packages/agent/CHANGELOG.md | 3 +- .../src/core/backgroundTools.cleanup.test.ts | 206 ----------------- .../agent/src/core/backgroundTools.test.ts | 87 -------- packages/agent/src/core/backgroundTools.ts | 207 ------------------ packages/agent/src/core/types.ts | 9 +- packages/agent/src/index.ts | 5 +- .../agent/src/tools/browser/browseMessage.ts | 24 +- .../agent/src/tools/browser/browseStart.ts | 32 +-- .../agent/src/tools/browser/browserTracker.ts | 144 ++++++++++++ .../agent/src/tools/browser/listBrowsers.ts | 102 +++++++++ packages/agent/src/tools/getTools.test.ts | 8 +- packages/agent/src/tools/getTools.ts | 4 +- .../src/tools/interaction/agentMessage.ts | 15 +- .../agent/src/tools/interaction/agentStart.ts | 34 +-- .../src/tools/interaction/agentTools.test.ts | 8 +- .../src/tools/interaction/agentTracker.ts | 11 +- .../src/tools/interaction/subAgent.test.ts | 8 +- .../agent/src/tools/interaction/subAgent.ts | 38 ++-- .../src/tools/system/ShellTracker.test.ts | 4 +- .../agent/src/tools/system/ShellTracker.ts | 17 +- packages/agent/src/tools/system/listAgents.ts | 4 +- .../tools/system/listBackgroundTools.test.ts | 25 --- .../src/tools/system/listBackgroundTools.ts | 114 ---------- .../agent/src/tools/system/listShells.test.ts | 3 +- packages/agent/src/tools/system/listShells.ts | 4 +- .../src/tools/system/shellMessage.test.ts | 3 +- .../agent/src/tools/system/shellMessage.ts | 6 +- .../agent/src/tools/system/shellStart.test.ts | 4 +- packages/agent/src/tools/system/shellStart.ts | 4 +- packages/cli/CHANGELOG.md | 3 +- packages/cli/src/commands/$default.ts | 12 +- packages/cli/src/index.ts | 5 +- packages/cli/src/utils/cleanup.ts | 72 ------ 33 files changed, 361 insertions(+), 864 deletions(-) delete mode 100644 packages/agent/src/core/backgroundTools.cleanup.test.ts delete mode 100644 packages/agent/src/core/backgroundTools.test.ts delete mode 100644 packages/agent/src/core/backgroundTools.ts create mode 100644 packages/agent/src/tools/browser/browserTracker.ts create mode 100644 packages/agent/src/tools/browser/listBrowsers.ts delete mode 100644 packages/agent/src/tools/system/listBackgroundTools.test.ts delete mode 100644 packages/agent/src/tools/system/listBackgroundTools.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 321b230..069f820 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -1,9 +1,8 @@ # [mycoder-agent-v1.4.2](https://github.com/drivecore/mycoder/compare/mycoder-agent-v1.4.1...mycoder-agent-v1.4.2) (2025-03-14) - ### Bug Fixes -* improve profiling ([79a3df2](https://github.com/drivecore/mycoder/commit/79a3df2db13b8372666c6604ebe1666d33663be9)) +- improve profiling ([79a3df2](https://github.com/drivecore/mycoder/commit/79a3df2db13b8372666c6604ebe1666d33663be9)) # [mycoder-agent-v1.4.1](https://github.com/drivecore/mycoder/compare/mycoder-agent-v1.4.0...mycoder-agent-v1.4.1) (2025-03-14) diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/packages/agent/src/core/backgroundTools.cleanup.test.ts deleted file mode 100644 index 28bedef..0000000 --- a/packages/agent/src/core/backgroundTools.cleanup.test.ts +++ /dev/null @@ -1,206 +0,0 @@ -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 { agentTracker } from '../tools/interaction/agentTracker.js'; -import { shellTracker } from '../tools/system/ShellTracker.js'; - -import { BackgroundTools, BackgroundToolStatus } from './backgroundTools'; - -// Import the ChildProcess type for mocking -import type { ChildProcess } from 'child_process'; - -// Define types for our mocks that match the actual types -type MockProcessState = { - process: ChildProcess & { kill: ReturnType }; - state: { - completed: boolean; - signaled: boolean; - exitCode: number | null; - }; - command: string; - stdout: string[]; - stderr: string[]; - showStdIn: boolean; - showStdout: boolean; -}; - -// Mock dependencies -vi.mock('../tools/browser/BrowserManager.js', () => { - return { - BrowserManager: class MockBrowserManager { - closeSession = vi.fn().mockResolvedValue(undefined); - }, - }; -}); - -vi.mock('../tools/system/ShellTracker.js', () => { - return { - shellTracker: { - processStates: new Map(), - cleanupAllShells: vi.fn().mockResolvedValue(undefined), - }, - }; -}); - -vi.mock('../tools/interaction/agentTracker.js', () => { - return { - agentTracker: { - terminateAgent: vi.fn().mockResolvedValue(undefined), - getAgentState: vi.fn().mockImplementation((id: string) => { - return { - id, - aborted: false, - completed: false, - context: { - backgroundTools: { - cleanup: vi.fn().mockResolvedValue(undefined), - }, - }, - goal: 'test goal', - prompt: 'test prompt', - output: '', - workingDirectory: '/test', - tools: [], - }; - }), - }, - }; -}); - -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(), - stdin: null, - stdout: null, - stderr: null, - stdio: null, - } as unknown as ChildProcess & { kill: ReturnType }; - - const mockProcessState: MockProcessState = { - process: mockProcess, - state: { - completed: false, - signaled: false, - exitCode: null, - }, - command: 'test command', - stdout: [], - stderr: [], - showStdIn: false, - showStdout: false, - }; - - shellTracker.processStates.clear(); - shellTracker.processStates.set('shell-1', mockProcessState as any); - - // Reset the agentTracker mock - vi.mocked(agentTracker.terminateAgent).mockClear(); - }); - - afterEach(() => { - vi.resetAllMocks(); - - // Clear global browser manager - ( - globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager } - ).__BROWSER_MANAGER__ = undefined; - - // Clear mock states - shellTracker.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 () => { - // Run cleanup - await backgroundTools.cleanup(); - - // Check that shellTracker.cleanupAllShells was called - expect(shellTracker.cleanupAllShells).toHaveBeenCalled(); - }); - - it('should clean up sub-agents', async () => { - // Register an agent tool - const agentId = backgroundTools.registerAgent('Test goal'); - - // Run cleanup - await backgroundTools.cleanup(); - - // Check that terminateAgent was called with the agent ID - expect(agentTracker.terminateAgent).toHaveBeenCalledWith(agentId); - - // 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.test.ts b/packages/agent/src/core/backgroundTools.test.ts deleted file mode 100644 index bc0a56b..0000000 --- a/packages/agent/src/core/backgroundTools.test.ts +++ /dev/null @@ -1,87 +0,0 @@ -import { describe, expect, it, vi, beforeEach } from 'vitest'; - -import { - BackgroundTools, - BackgroundToolStatus, - BackgroundToolType, -} from './backgroundTools.js'; - -// Mock uuid to return predictable IDs for testing -vi.mock('uuid', () => ({ - v4: vi.fn().mockReturnValue('test-id-1'), // Always return the same ID for simplicity in tests -})); - -describe('BackgroundToolRegistry', () => { - let backgroundTools: BackgroundTools; - beforeEach(() => { - // Clear all registered tools before each test - backgroundTools = new BackgroundTools('test'); - backgroundTools.tools = new Map(); - }); - - it('should register a browser process', () => { - const id = backgroundTools.registerBrowser('https://example.com'); - - expect(id).toBe('test-id-1'); - - const tool = backgroundTools.getToolById(id); - expect(tool).toBeDefined(); - if (tool) { - expect(tool.type).toBe(BackgroundToolType.BROWSER); - expect(tool.status).toBe(BackgroundToolStatus.RUNNING); - if (tool.type === BackgroundToolType.BROWSER) { - expect(tool.metadata.url).toBe('https://example.com'); - } - } - }); - - it('should return false when updating non-existent tool', () => { - const updated = backgroundTools.updateToolStatus( - 'non-existent-id', - BackgroundToolStatus.COMPLETED, - ); - - expect(updated).toBe(false); - }); - - it('should clean up old completed tools', () => { - // Create tools with specific dates - const registry = backgroundTools as any; - - // Add a completed tool from 25 hours ago - const oldTool = { - id: 'old-tool', - type: BackgroundToolType.BROWSER, - status: BackgroundToolStatus.COMPLETED, - startTime: new Date(Date.now() - 25 * 60 * 60 * 1000), - endTime: new Date(Date.now() - 25 * 60 * 60 * 1000), - agentId: 'agent-1', - metadata: { url: 'https://example.com' }, - }; - - // Add a completed tool from 10 hours ago - const recentTool = { - id: 'recent-tool', - type: BackgroundToolType.BROWSER, - status: BackgroundToolStatus.COMPLETED, - startTime: new Date(Date.now() - 10 * 60 * 60 * 1000), - endTime: new Date(Date.now() - 10 * 60 * 60 * 1000), - agentId: 'agent-1', - metadata: { url: 'https://example.com' }, - }; - - // Add a running tool from 25 hours ago - const oldRunningTool = { - id: 'old-running-tool', - type: BackgroundToolType.BROWSER, - status: BackgroundToolStatus.RUNNING, - startTime: new Date(Date.now() - 25 * 60 * 60 * 1000), - agentId: 'agent-1', - metadata: { url: 'https://example.com' }, - }; - - registry.tools.set('old-tool', oldTool); - registry.tools.set('recent-tool', recentTool); - registry.tools.set('old-running-tool', oldRunningTool); - }); -}); diff --git a/packages/agent/src/core/backgroundTools.ts b/packages/agent/src/core/backgroundTools.ts deleted file mode 100644 index c4768e2..0000000 --- a/packages/agent/src/core/backgroundTools.ts +++ /dev/null @@ -1,207 +0,0 @@ -import { v4 as uuidv4 } from 'uuid'; - -// These imports will be used by the cleanup method -import { BrowserManager } from '../tools/browser/BrowserManager.js'; -import { agentTracker } from '../tools/interaction/agentTracker.js'; -import { shellTracker } from '../tools/system/ShellTracker.js'; - -// Types of background processes we can track -export enum BackgroundToolType { - BROWSER = 'browser', - AGENT = 'agent', -} - -// Status of a background process -export enum BackgroundToolStatus { - RUNNING = 'running', - COMPLETED = 'completed', - ERROR = 'error', - TERMINATED = 'terminated', -} - -// Common interface for all background processes -export interface BackgroundTool { - id: string; - type: BackgroundToolType; - status: BackgroundToolStatus; - startTime: Date; - endTime?: Date; - metadata: Record; // Additional tool-specific information -} - -// Browser process specific data -export interface BrowserBackgroundTool extends BackgroundTool { - type: BackgroundToolType.BROWSER; - metadata: { - url?: string; - error?: string; - }; -} - -// Agent process specific data (for future use) -export interface AgentBackgroundTool extends BackgroundTool { - type: BackgroundToolType.AGENT; - metadata: { - goal?: string; - error?: string; - }; -} - -// Utility type for all background tool types -export type AnyBackgroundTool = BrowserBackgroundTool | AgentBackgroundTool; - -/** - * Registry to keep track of all background processes - */ -export class BackgroundTools { - tools: Map = new Map(); - - // Private constructor for singleton pattern - constructor(readonly ownerName: string) {} - - // Register a new browser process - public registerBrowser(url?: string): string { - const id = uuidv4(); - const tool: BrowserBackgroundTool = { - id, - type: BackgroundToolType.BROWSER, - status: BackgroundToolStatus.RUNNING, - startTime: new Date(), - metadata: { - url, - }, - }; - this.tools.set(id, tool); - return id; - } - - // Register a new agent process (for future use) - public registerAgent(goal?: string): string { - const id = uuidv4(); - const tool: AgentBackgroundTool = { - id, - type: BackgroundToolType.AGENT, - status: BackgroundToolStatus.RUNNING, - startTime: new Date(), - metadata: { - goal, - }, - }; - this.tools.set(id, tool); - return id; - } - - // Update the status of a process - public updateToolStatus( - id: string, - status: BackgroundToolStatus, - metadata?: Record, - ): boolean { - const tool = this.tools.get(id); - if (!tool) { - return false; - } - - tool.status = status; - - if ( - status === BackgroundToolStatus.COMPLETED || - status === BackgroundToolStatus.ERROR || - status === BackgroundToolStatus.TERMINATED - ) { - tool.endTime = new Date(); - } - - if (metadata) { - tool.metadata = { ...tool.metadata, ...metadata }; - } - - return true; - } - - public getTools(): AnyBackgroundTool[] { - const result: AnyBackgroundTool[] = []; - for (const tool of this.tools.values()) { - result.push(tool); - } - return result; - } - - // Get a specific process by ID - 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 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 agentCleanupPromises = agentTools.map((tool) => - this.cleanupSubAgent(tool), - ); - - // Clean up shell processes using ShellTracker - await shellTracker.cleanupAllShells(); - - // Wait for all cleanup operations to complete in parallel - await Promise.all([...browserCleanupPromises, ...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 sub-agent - * @param tool The agent tool to clean up - */ - private async cleanupSubAgent(tool: AgentBackgroundTool): Promise { - try { - // Delegate to the agent tracker - await agentTracker.terminateAgent(tool.id); - 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/core/types.ts b/packages/agent/src/core/types.ts index a0a411e..ade4501 100644 --- a/packages/agent/src/core/types.ts +++ b/packages/agent/src/core/types.ts @@ -1,9 +1,11 @@ import { z } from 'zod'; import { JsonSchema7Type } from 'zod-to-json-schema'; +import { BrowserTracker } from '../tools/browser/browserTracker.js'; +import { AgentTracker } from '../tools/interaction/agentTracker.js'; +import { ShellTracker } from '../tools/system/ShellTracker.js'; import { Logger } from '../utils/logger.js'; -import { BackgroundTools } from './backgroundTools.js'; import { TokenTracker } from './tokens.js'; import { ModelProvider } from './toolAgent/config.js'; @@ -23,13 +25,16 @@ export type ToolContext = { tokenCache?: boolean; userPrompt?: boolean; agentId?: string; // Unique identifier for the agent, used for background tool tracking + agentName?: string; // Name of the agent, used for browser tracker provider: ModelProvider; model?: string; baseUrl?: string; apiKey?: string; maxTokens: number; temperature: number; - backgroundTools: BackgroundTools; + agentTracker: AgentTracker; + shellTracker: ShellTracker; + browserTracker: BrowserTracker; }; export type Tool, TReturn = any> = { diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 20c5fa1..5c026ea 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -9,7 +9,6 @@ export * from './tools/system/respawn.js'; export * from './tools/system/sequenceComplete.js'; export * from './tools/system/shellMessage.js'; export * from './tools/system/shellExecute.js'; -export * from './tools/system/listBackgroundTools.js'; export * from './tools/system/listShells.js'; export * from './tools/system/ShellTracker.js'; @@ -20,7 +19,10 @@ export * from './tools/browser/browseMessage.js'; export * from './tools/browser/browseStart.js'; export * from './tools/browser/PageController.js'; export * from './tools/browser/BrowserAutomation.js'; +export * from './tools/browser/listBrowsers.js'; +export * from './tools/browser/browserTracker.js'; +export * from './tools/interaction/agentTracker.js'; // Tools - Interaction export * from './tools/interaction/subAgent.js'; export * from './tools/interaction/userPrompt.js'; @@ -28,7 +30,6 @@ export * from './tools/interaction/userPrompt.js'; // Core export * from './core/executeToolCall.js'; export * from './core/types.js'; -export * from './core/backgroundTools.js'; // Tool Agent Core export { toolAgent } from './core/toolAgent/toolAgentCore.js'; export * from './core/toolAgent/config.js'; diff --git a/packages/agent/src/tools/browser/browseMessage.ts b/packages/agent/src/tools/browser/browseMessage.ts index a6b35d5..7ed3704 100644 --- a/packages/agent/src/tools/browser/browseMessage.ts +++ b/packages/agent/src/tools/browser/browseMessage.ts @@ -1,11 +1,11 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; import { sleep } from '../../utils/sleep.js'; +import { BrowserSessionStatus } from './browserTracker.js'; import { filterPageContent } from './filterPageContent.js'; import { browserSessions, SelectorType } from './types.js'; @@ -72,7 +72,7 @@ export const browseMessageTool: Tool = { execute: async ( { instanceId, actionType, url, selector, selectorType, text }, - { logger, pageFilter, backgroundTools }, + { logger, pageFilter, browserTracker, ..._ }, ): Promise => { // Validate action format @@ -186,10 +186,10 @@ export const browseMessageTool: Tool = { await session.browser.close(); browserSessions.delete(instanceId); - // Update background tool registry when browser is explicitly closed - backgroundTools.updateToolStatus( + // Update browser tracker when browser is explicitly closed + browserTracker.updateSessionStatus( instanceId, - BackgroundToolStatus.COMPLETED, + BrowserSessionStatus.COMPLETED, { closedExplicitly: true, }, @@ -206,11 +206,15 @@ export const browseMessageTool: Tool = { } catch (error) { logger.error('Browser action failed:', { error }); - // Update background tool registry with error status if action fails - backgroundTools.updateToolStatus(instanceId, BackgroundToolStatus.ERROR, { - error: errorToString(error), - actionType, - }); + // Update browser tracker with error status if action fails + browserTracker.updateSessionStatus( + instanceId, + BrowserSessionStatus.ERROR, + { + error: errorToString(error), + actionType, + }, + ); return { status: 'error', diff --git a/packages/agent/src/tools/browser/browseStart.ts b/packages/agent/src/tools/browser/browseStart.ts index ad41298..738b5bf 100644 --- a/packages/agent/src/tools/browser/browseStart.ts +++ b/packages/agent/src/tools/browser/browseStart.ts @@ -1,13 +1,12 @@ import { chromium } from '@playwright/test'; -import { v4 as uuidv4 } from 'uuid'; import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; import { sleep } from '../../utils/sleep.js'; +import { BrowserSessionStatus } from './browserTracker.js'; import { filterPageContent } from './filterPageContent.js'; import { browserSessions } from './types.js'; @@ -43,7 +42,14 @@ export const browseStartTool: Tool = { execute: async ( { url, timeout = 30000 }, - { logger, headless, userSession, pageFilter, backgroundTools }, + { + logger, + headless, + userSession, + pageFilter, + browserTracker, + ..._ // Unused parameters + }, ): Promise => { logger.verbose(`Starting browser session${url ? ` at ${url}` : ''}`); logger.verbose( @@ -52,10 +58,8 @@ export const browseStartTool: Tool = { logger.verbose(`Webpage processing mode: ${pageFilter}`); try { - const instanceId = uuidv4(); - - // Register this browser session with the background tool registry - backgroundTools.registerBrowser(url); + // Register this browser session with the tracker + const instanceId = browserTracker.registerBrowser(url); // Launch browser const launchOptions = { @@ -95,10 +99,10 @@ export const browseStartTool: Tool = { // Setup cleanup handlers browser.on('disconnected', () => { browserSessions.delete(instanceId); - // Update background tool registry when browser disconnects - backgroundTools.updateToolStatus( + // Update browser tracker when browser disconnects + browserTracker.updateSessionStatus( instanceId, - BackgroundToolStatus.TERMINATED, + BrowserSessionStatus.TERMINATED, ); }); @@ -142,10 +146,10 @@ export const browseStartTool: Tool = { logger.verbose('Browser session started successfully'); logger.verbose(`Content length: ${content.length} characters`); - // Update background tool registry with running status - backgroundTools.updateToolStatus( + // Update browser tracker with running status + browserTracker.updateSessionStatus( instanceId, - BackgroundToolStatus.RUNNING, + BrowserSessionStatus.RUNNING, { url: url || 'about:blank', contentLength: content.length, @@ -160,7 +164,7 @@ export const browseStartTool: Tool = { } catch (error) { logger.error(`Failed to start browser: ${errorToString(error)}`); - // No need to update background tool registry here as we don't have a valid instanceId + // No need to update browser tracker here as we don't have a valid instanceId // when an error occurs before the browser is properly initialized return { diff --git a/packages/agent/src/tools/browser/browserTracker.ts b/packages/agent/src/tools/browser/browserTracker.ts new file mode 100644 index 0000000..31c2bc1 --- /dev/null +++ b/packages/agent/src/tools/browser/browserTracker.ts @@ -0,0 +1,144 @@ +import { v4 as uuidv4 } from 'uuid'; + +import { BrowserManager } from './BrowserManager.js'; +import { browserSessions } from './types.js'; + +// Status of a browser session +export enum BrowserSessionStatus { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + TERMINATED = 'terminated', +} + +// Browser session tracking data +export interface BrowserSessionInfo { + id: string; + status: BrowserSessionStatus; + startTime: Date; + endTime?: Date; + metadata: { + url?: string; + contentLength?: number; + closedExplicitly?: boolean; + error?: string; + actionType?: string; + }; +} + +/** + * Registry to keep track of browser sessions + */ +export class BrowserTracker { + private sessions: Map = new Map(); + + constructor(public ownerAgentId: string | undefined) {} + + // Register a new browser session + public registerBrowser(url?: string): string { + const id = uuidv4(); + const session: BrowserSessionInfo = { + id, + status: BrowserSessionStatus.RUNNING, + startTime: new Date(), + metadata: { + url, + }, + }; + this.sessions.set(id, session); + return id; + } + + // Update the status of a browser session + public updateSessionStatus( + id: string, + status: BrowserSessionStatus, + metadata?: Record, + ): boolean { + const session = this.sessions.get(id); + if (!session) { + return false; + } + + session.status = status; + + if ( + status === BrowserSessionStatus.COMPLETED || + status === BrowserSessionStatus.ERROR || + status === BrowserSessionStatus.TERMINATED + ) { + session.endTime = new Date(); + } + + if (metadata) { + session.metadata = { ...session.metadata, ...metadata }; + } + + return true; + } + + // Get all browser sessions + public getSessions(): BrowserSessionInfo[] { + return Array.from(this.sessions.values()); + } + + // Get a specific browser session by ID + public getSessionById(id: string): BrowserSessionInfo | undefined { + return this.sessions.get(id); + } + + // Filter sessions by status + public getSessionsByStatus( + status: BrowserSessionStatus, + ): BrowserSessionInfo[] { + return this.getSessions().filter((session) => session.status === status); + } + + /** + * Cleans up all browser sessions associated with this tracker + * @returns A promise that resolves when cleanup is complete + */ + public async cleanup(): Promise { + const sessions = this.getSessionsByStatus(BrowserSessionStatus.RUNNING); + + // Create cleanup promises for each session + const cleanupPromises = sessions.map((session) => + this.cleanupBrowserSession(session), + ); + + // Wait for all cleanup operations to complete in parallel + await Promise.all(cleanupPromises); + } + + /** + * Cleans up a browser session + * @param session The browser session to clean up + */ + private async cleanupBrowserSession( + session: BrowserSessionInfo, + ): Promise { + try { + const browserManager = ( + globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager } + ).__BROWSER_MANAGER__; + + if (browserManager) { + await browserManager.closeSession(session.id); + } else { + // Fallback to closing via browserSessions if BrowserManager is not available + const browserSession = browserSessions.get(session.id); + if (browserSession) { + await browserSession.page.context().close(); + await browserSession.browser.close(); + browserSessions.delete(session.id); + } + } + + this.updateSessionStatus(session.id, BrowserSessionStatus.COMPLETED); + } catch (error) { + this.updateSessionStatus(session.id, BrowserSessionStatus.ERROR, { + error: error instanceof Error ? error.message : String(error), + }); + } + } +} diff --git a/packages/agent/src/tools/browser/listBrowsers.ts b/packages/agent/src/tools/browser/listBrowsers.ts new file mode 100644 index 0000000..a370af7 --- /dev/null +++ b/packages/agent/src/tools/browser/listBrowsers.ts @@ -0,0 +1,102 @@ +import { z } from 'zod'; +import { zodToJsonSchema } from 'zod-to-json-schema'; + +import { Tool } from '../../core/types.js'; + +import { BrowserSessionStatus } from './browserTracker.js'; + +const parameterSchema = z.object({ + status: z + .enum(['all', 'running', 'completed', 'error', 'terminated']) + .optional() + .describe('Filter browser sessions by status (default: "all")'), + verbose: z + .boolean() + .optional() + .describe( + 'Include detailed metadata about each browser session (default: false)', + ), +}); + +const returnSchema = z.object({ + sessions: z.array( + z.object({ + id: z.string(), + status: z.string(), + startTime: z.string(), + endTime: z.string().optional(), + runtime: z.number().describe('Runtime in seconds'), + url: z.string().optional(), + metadata: z.record(z.any()).optional(), + }), + ), + count: z.number(), +}); + +type Parameters = z.infer; +type ReturnType = z.infer; + +export const listBrowsersTool: Tool = { + name: 'listBrowsers', + description: 'Lists all browser sessions and their status', + logPrefix: '🔍', + parameters: parameterSchema, + returns: returnSchema, + parametersJsonSchema: zodToJsonSchema(parameterSchema), + returnsJsonSchema: zodToJsonSchema(returnSchema), + + execute: async ( + { status = 'all', verbose = false }, + { logger, browserTracker, ..._ }, + ): Promise => { + logger.verbose( + `Listing browser sessions with status: ${status}, verbose: ${verbose}`, + ); + + // Get all browser sessions + const sessions = browserTracker.getSessions(); + + // Filter by status if specified + const filteredSessions = + status === 'all' + ? sessions + : sessions.filter((session) => { + const statusEnum = + status.toUpperCase() as keyof typeof BrowserSessionStatus; + return session.status === BrowserSessionStatus[statusEnum]; + }); + + // Format the response + const formattedSessions = filteredSessions.map((session) => { + const now = new Date(); + const startTime = session.startTime; + const endTime = session.endTime || now; + const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds + + return { + id: session.id, + status: session.status, + startTime: startTime.toISOString(), + ...(session.endTime && { endTime: session.endTime.toISOString() }), + runtime: parseFloat(runtime.toFixed(2)), + url: session.metadata.url, + ...(verbose && { metadata: session.metadata }), + }; + }); + + return { + sessions: formattedSessions, + count: formattedSessions.length, + }; + }, + + logParameters: ({ status = 'all', verbose = false }, { logger }) => { + logger.info( + `Listing browser sessions with status: ${status}, verbose: ${verbose}`, + ); + }, + + logReturns: (output, { logger }) => { + logger.info(`Found ${output.count} browser sessions`); + }, +}; diff --git a/packages/agent/src/tools/getTools.test.ts b/packages/agent/src/tools/getTools.test.ts index 211e116..362a2d8 100644 --- a/packages/agent/src/tools/getTools.test.ts +++ b/packages/agent/src/tools/getTools.test.ts @@ -1,11 +1,13 @@ import { describe, it, expect } from 'vitest'; -import { BackgroundTools } from '../core/backgroundTools.js'; import { TokenTracker } from '../core/tokens.js'; import { ToolContext } from '../core/types.js'; import { MockLogger } from '../utils/mockLogger.js'; +import { BrowserTracker } from './browser/browserTracker.js'; import { getTools } from './getTools.js'; +import { AgentTracker } from './interaction/agentTracker.js'; +import { ShellTracker } from './system/ShellTracker.js'; // Mock context export const getMockToolContext = (): ToolContext => ({ @@ -20,7 +22,9 @@ export const getMockToolContext = (): ToolContext => ({ model: 'claude-3-7-sonnet-20250219', maxTokens: 4096, temperature: 0.7, - backgroundTools: new BackgroundTools('test'), + agentTracker: new AgentTracker('test'), + shellTracker: new ShellTracker('test'), + browserTracker: new BrowserTracker('test'), }); describe('getTools', () => { diff --git a/packages/agent/src/tools/getTools.ts b/packages/agent/src/tools/getTools.ts index 6c8f3b2..598744c 100644 --- a/packages/agent/src/tools/getTools.ts +++ b/packages/agent/src/tools/getTools.ts @@ -4,13 +4,13 @@ import { Tool } from '../core/types.js'; // Import tools import { browseMessageTool } from './browser/browseMessage.js'; import { browseStartTool } from './browser/browseStart.js'; +import { listBrowsersTool } from './browser/listBrowsers.js'; import { subAgentTool } from './interaction/subAgent.js'; import { userPromptTool } from './interaction/userPrompt.js'; import { fetchTool } from './io/fetch.js'; import { textEditorTool } from './io/textEditor.js'; import { createMcpTool } from './mcp.js'; import { listAgentsTool } from './system/listAgents.js'; -import { listBackgroundToolsTool } from './system/listBackgroundTools.js'; import { listShellsTool } from './system/listShells.js'; import { sequenceCompleteTool } from './system/sequenceComplete.js'; import { shellMessageTool } from './system/shellMessage.js'; @@ -32,6 +32,7 @@ export function getTools(options?: GetToolsOptions): Tool[] { const tools: Tool[] = [ textEditorTool as unknown as Tool, subAgentTool as unknown as Tool, + listBrowsersTool as unknown as Tool, /*agentStartTool as unknown as Tool, agentMessageTool as unknown as Tool,*/ sequenceCompleteTool as unknown as Tool, @@ -42,7 +43,6 @@ export function getTools(options?: GetToolsOptions): Tool[] { browseMessageTool as unknown as Tool, //respawnTool as unknown as Tool, this is a confusing tool for now. sleepTool as unknown as Tool, - listBackgroundToolsTool as unknown as Tool, listShellsTool as unknown as Tool, listAgentsTool as unknown as Tool, ]; diff --git a/packages/agent/src/tools/interaction/agentMessage.ts b/packages/agent/src/tools/interaction/agentMessage.ts index fd3e773..d846a53 100644 --- a/packages/agent/src/tools/interaction/agentMessage.ts +++ b/packages/agent/src/tools/interaction/agentMessage.ts @@ -1,7 +1,6 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { agentStates } from './agentStart.js'; @@ -51,7 +50,7 @@ export const agentMessageTool: Tool = { execute: async ( { instanceId, guidance, terminate }, - { logger, backgroundTools }, + { logger, ..._ }, ): Promise => { logger.verbose( `Interacting with sub-agent ${instanceId}${guidance ? ' with guidance' : ''}${terminate ? ' with termination request' : ''}`, @@ -77,18 +76,6 @@ export const agentMessageTool: Tool = { agentState.aborted = true; agentState.completed = true; - // Update background tool registry with terminated status - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.TERMINATED, - { - terminatedByUser: true, - }, - ); - - // 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 f1bd4f5..0a10651 100644 --- a/packages/agent/src/tools/interaction/agentStart.ts +++ b/packages/agent/src/tools/interaction/agentStart.ts @@ -1,7 +1,6 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; import { getDefaultSystemPrompt, AgentConfig, @@ -10,7 +9,7 @@ import { toolAgent } from '../../core/toolAgent/toolAgentCore.js'; import { Tool, ToolContext } from '../../core/types.js'; import { getTools } from '../getTools.js'; -import { AgentStatus, agentTracker, AgentState } from './agentTracker.js'; +import { AgentStatus, AgentState } from './agentTracker.js'; // For backward compatibility export const agentStates = new Map(); @@ -74,7 +73,7 @@ export const agentStartTool: Tool = { returns: returnSchema, returnsJsonSchema: zodToJsonSchema(returnSchema), execute: async (params, context) => { - const { logger, backgroundTools } = context; + const { logger, agentTracker } = context; // Validate parameters const { @@ -89,9 +88,6 @@ export const agentStartTool: Tool = { // Register this agent with the agent tracker const instanceId = agentTracker.registerAgent(goal); - // For backward compatibility, also register with background tools - backgroundTools.registerAgent(goal); - logger.verbose(`Registered agent with ID: ${instanceId}`); // Construct a well-structured prompt @@ -150,20 +146,6 @@ export const agentStartTool: Tool = { result.result.substring(0, 100) + (result.result.length > 100 ? '...' : ''), }); - - // For backward compatibility - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.COMPLETED, - { - result: - result.result.substring(0, 100) + - (result.result.length > 100 ? '...' : ''), - }, - ); - - // Clean up resources when agent completes successfully - await backgroundTools.cleanup(); } } catch (error) { // Update agent state with the error @@ -176,18 +158,6 @@ export const agentStartTool: Tool = { agentTracker.updateAgentStatus(instanceId, AgentStatus.ERROR, { error: error instanceof Error ? error.message : String(error), }); - - // For backward compatibility - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.ERROR, - { - 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/agent/src/tools/interaction/agentTools.test.ts b/packages/agent/src/tools/interaction/agentTools.test.ts index 6e1c26f..032c02c 100644 --- a/packages/agent/src/tools/interaction/agentTools.test.ts +++ b/packages/agent/src/tools/interaction/agentTools.test.ts @@ -1,12 +1,14 @@ import { describe, expect, it, vi } from 'vitest'; -import { BackgroundTools } from '../../core/backgroundTools.js'; import { TokenTracker } from '../../core/tokens.js'; import { ToolContext } from '../../core/types.js'; import { MockLogger } from '../../utils/mockLogger.js'; +import { BrowserTracker } from '../browser/browserTracker.js'; +import { ShellTracker } from '../system/ShellTracker.js'; import { agentMessageTool } from './agentMessage.js'; import { agentStartTool, agentStates } from './agentStart.js'; +import { AgentTracker } from './agentTracker.js'; // Mock the toolAgent function vi.mock('../../core/toolAgent/toolAgentCore.js', () => ({ @@ -29,7 +31,9 @@ const mockContext: ToolContext = { model: 'claude-3-7-sonnet-20250219', maxTokens: 4096, temperature: 0.7, - backgroundTools: new BackgroundTools('test'), + agentTracker: new AgentTracker('test'), + shellTracker: new ShellTracker('test'), + browserTracker: new BrowserTracker('test'), }; describe('Agent Tools', () => { diff --git a/packages/agent/src/tools/interaction/agentTracker.ts b/packages/agent/src/tools/interaction/agentTracker.ts index bb6463d..20b9a42 100644 --- a/packages/agent/src/tools/interaction/agentTracker.ts +++ b/packages/agent/src/tools/interaction/agentTracker.ts @@ -39,7 +39,7 @@ export class AgentTracker { private agents: Map = new Map(); private agentStates: Map = new Map(); - constructor(readonly ownerName: string) {} + constructor(public ownerAgentId: string | undefined) {} // Register a new agent public registerAgent(goal: string): string { @@ -131,9 +131,9 @@ export class AgentTracker { agentState.completed = true; // Clean up resources owned by this sub-agent - if (agentState.context.backgroundTools) { - await agentState.context.backgroundTools.cleanup(); - } + await agentState.context.agentTracker.cleanup(); + await agentState.context.shellTracker.cleanup(); + await agentState.context.browserTracker.cleanup(); } this.updateAgentStatus(id, AgentStatus.TERMINATED); } catch (error) { @@ -143,6 +143,3 @@ export class AgentTracker { } } } - -// Create a singleton instance -export const agentTracker = new AgentTracker('global'); diff --git a/packages/agent/src/tools/interaction/subAgent.test.ts b/packages/agent/src/tools/interaction/subAgent.test.ts index 6b0dff7..11bc9c5 100644 --- a/packages/agent/src/tools/interaction/subAgent.test.ts +++ b/packages/agent/src/tools/interaction/subAgent.test.ts @@ -1,10 +1,12 @@ import { describe, expect, it, vi } from 'vitest'; -import { BackgroundTools } from '../../core/backgroundTools.js'; import { TokenTracker } from '../../core/tokens.js'; import { ToolContext } from '../../core/types.js'; import { MockLogger } from '../../utils/mockLogger.js'; +import { BrowserTracker } from '../browser/browserTracker.js'; +import { ShellTracker } from '../system/ShellTracker.js'; +import { AgentTracker } from './agentTracker.js'; import { subAgentTool } from './subAgent.js'; // Mock the toolAgent function @@ -33,7 +35,9 @@ const mockContext: ToolContext = { model: 'claude-3-7-sonnet-20250219', maxTokens: 4096, temperature: 0.7, - backgroundTools: new BackgroundTools('test'), + agentTracker: new AgentTracker('test'), + shellTracker: new ShellTracker('test'), + browserTracker: new BrowserTracker('test'), }; describe('subAgentTool', () => { diff --git a/packages/agent/src/tools/interaction/subAgent.ts b/packages/agent/src/tools/interaction/subAgent.ts index ac32616..65d3091 100644 --- a/packages/agent/src/tools/interaction/subAgent.ts +++ b/packages/agent/src/tools/interaction/subAgent.ts @@ -1,17 +1,17 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { - BackgroundTools, - BackgroundToolStatus, -} from '../../core/backgroundTools.js'; import { getDefaultSystemPrompt, AgentConfig, } from '../../core/toolAgent/config.js'; import { toolAgent } from '../../core/toolAgent/toolAgentCore.js'; import { Tool, ToolContext } from '../../core/types.js'; +import { BrowserTracker } from '../browser/browserTracker.js'; import { getTools } from '../getTools.js'; +import { ShellTracker } from '../system/ShellTracker.js'; + +import { AgentTracker } from './agentTracker.js'; const parameterSchema = z.object({ description: z @@ -69,7 +69,7 @@ export const subAgentTool: Tool = { returns: returnSchema, returnsJsonSchema: zodToJsonSchema(returnSchema), execute: async (params, context) => { - const { logger, backgroundTools } = context; + const { logger, agentTracker } = context; // Validate parameters const { @@ -81,13 +81,15 @@ export const subAgentTool: Tool = { } = parameterSchema.parse(params); // Register this sub-agent with the background tool registry - const subAgentId = backgroundTools.registerAgent(goal); + const subAgentId = agentTracker.registerAgent(goal); logger.verbose(`Registered sub-agent with ID: ${subAgentId}`); const localContext = { ...context, workingDirectory: workingDirectory ?? context.workingDirectory, - backgroundTools: new BackgroundTools(`subAgent: ${goal}`), + agentTracker: new AgentTracker(subAgentId), + shellTracker: new ShellTracker(subAgentId), + browserTracker: new BrowserTracker(subAgentId), }; // Construct a well-structured prompt @@ -114,24 +116,14 @@ export const subAgentTool: Tool = { const result = await toolAgent(prompt, tools, config, localContext); // Update background tool registry with completed status - backgroundTools.updateToolStatus( - subAgentId, - BackgroundToolStatus.COMPLETED, - { - result: - result.result.substring(0, 100) + - (result.result.length > 100 ? '...' : ''), - }, - ); return { response: result.result }; - } catch (error) { - // Update background tool registry with error status - backgroundTools.updateToolStatus(subAgentId, BackgroundToolStatus.ERROR, { - error: error instanceof Error ? error.message : String(error), - }); - - throw error; + } finally { + await Promise.all([ + localContext.agentTracker.cleanup(), + localContext.shellTracker.cleanup(), + localContext.browserTracker.cleanup(), + ]); } }, logParameters: (input, { logger }) => { diff --git a/packages/agent/src/tools/system/ShellTracker.test.ts b/packages/agent/src/tools/system/ShellTracker.test.ts index 7fd8fdb..2f22be9 100644 --- a/packages/agent/src/tools/system/ShellTracker.test.ts +++ b/packages/agent/src/tools/system/ShellTracker.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { ShellStatus, shellTracker } from './ShellTracker.js'; +import { ShellStatus, ShellTracker } from './ShellTracker.js'; // Mock uuid to return predictable IDs for testing vi.mock('uuid', () => ({ @@ -12,6 +12,8 @@ vi.mock('uuid', () => ({ })); describe('ShellTracker', () => { + const shellTracker = new ShellTracker('test'); + beforeEach(() => { // Clear all registered shells before each test shellTracker['shells'] = new Map(); diff --git a/packages/agent/src/tools/system/ShellTracker.ts b/packages/agent/src/tools/system/ShellTracker.ts index c7dd4bf..d85308c 100644 --- a/packages/agent/src/tools/system/ShellTracker.ts +++ b/packages/agent/src/tools/system/ShellTracker.ts @@ -44,20 +44,10 @@ export interface ShellProcess { * Registry to keep track of shell processes */ export class ShellTracker { - private static instance: ShellTracker; private shells: Map = new Map(); public processStates: Map = new Map(); - // Private constructor for singleton pattern - private constructor() {} - - // Get the singleton instance - public static getInstance(): ShellTracker { - if (!ShellTracker.instance) { - ShellTracker.instance = new ShellTracker(); - } - return ShellTracker.instance; - } + constructor(public ownerAgentId: string | undefined) {} // Register a new shell process public registerShell(command: string): string { @@ -158,7 +148,7 @@ export class ShellTracker { /** * Cleans up all running shell processes */ - public async cleanupAllShells(): Promise { + public async cleanup(): Promise { const runningShells = this.getShells(ShellStatus.RUNNING); const cleanupPromises = runningShells.map((shell) => this.cleanupShellProcess(shell.id), @@ -166,6 +156,3 @@ export class ShellTracker { await Promise.all(cleanupPromises); } } - -// Export a singleton instance -export const shellTracker = ShellTracker.getInstance(); diff --git a/packages/agent/src/tools/system/listAgents.ts b/packages/agent/src/tools/system/listAgents.ts index e60e1bd..ea028fe 100644 --- a/packages/agent/src/tools/system/listAgents.ts +++ b/packages/agent/src/tools/system/listAgents.ts @@ -2,7 +2,7 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; import { Tool } from '../../core/types.js'; -import { AgentStatus, agentTracker } from '../interaction/agentTracker.js'; +import { AgentStatus } from '../interaction/agentTracker.js'; const parameterSchema = z.object({ status: z @@ -45,7 +45,7 @@ export const listAgentsTool: Tool = { execute: async ( { status = 'all', verbose = false }, - { logger }, + { logger, agentTracker }, ): Promise => { logger.verbose( `Listing agents with status: ${status}, verbose: ${verbose}`, diff --git a/packages/agent/src/tools/system/listBackgroundTools.test.ts b/packages/agent/src/tools/system/listBackgroundTools.test.ts deleted file mode 100644 index 3b80dba..0000000 --- a/packages/agent/src/tools/system/listBackgroundTools.test.ts +++ /dev/null @@ -1,25 +0,0 @@ -import { describe, expect, it, vi } from 'vitest'; - -import { BackgroundTools } from '../../core/backgroundTools.js'; - -import { listBackgroundToolsTool } from './listBackgroundTools.js'; - -describe('listBackgroundTools tool', () => { - const mockLogger = { - debug: vi.fn(), - verbose: vi.fn(), - info: vi.fn(), - warn: vi.fn(), - error: vi.fn(), - }; - - it('should list background tools', async () => { - const result = await listBackgroundToolsTool.execute({}, { - logger: mockLogger as any, - backgroundTools: new BackgroundTools('test'), - } as any); - - expect(result.count).toEqual(0); - expect(result.tools).toHaveLength(0); - }); -}); diff --git a/packages/agent/src/tools/system/listBackgroundTools.ts b/packages/agent/src/tools/system/listBackgroundTools.ts deleted file mode 100644 index 41526a7..0000000 --- a/packages/agent/src/tools/system/listBackgroundTools.ts +++ /dev/null @@ -1,114 +0,0 @@ -import { z } from 'zod'; -import { zodToJsonSchema } from 'zod-to-json-schema'; - -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; -import { Tool } from '../../core/types.js'; - -const parameterSchema = z.object({ - status: z - .enum(['all', 'running', 'completed', 'error', 'terminated']) - .optional() - .describe('Filter tools by status (default: "all")'), - type: z - .enum(['all', 'browser', 'agent']) - .optional() - .describe('Filter tools by type (default: "all")'), - verbose: z - .boolean() - .optional() - .describe('Include detailed metadata about each tool (default: false)'), -}); - -const returnSchema = z.object({ - tools: z.array( - z.object({ - id: z.string(), - type: z.string(), - status: z.string(), - startTime: z.string(), - endTime: z.string().optional(), - runtime: z.number().describe('Runtime in seconds'), - metadata: z.record(z.any()).optional(), - }), - ), - count: z.number(), -}); - -type Parameters = z.infer; -type ReturnType = z.infer; - -export const listBackgroundToolsTool: Tool = { - name: 'listBackgroundTools', - description: 'Lists all background tools (browsers, agents) and their status', - logPrefix: '🔍', - parameters: parameterSchema, - returns: returnSchema, - parametersJsonSchema: zodToJsonSchema(parameterSchema), - returnsJsonSchema: zodToJsonSchema(returnSchema), - - execute: async ( - { status = 'all', type = 'all', verbose = false }, - { logger, backgroundTools }, - ): Promise => { - logger.verbose( - `Listing background tools with status: ${status}, type: ${type}, verbose: ${verbose}`, - ); - - // Get all tools for this agent - const tools = backgroundTools.getTools(); - - // Filter by status if specified - const filteredByStatus = - status === 'all' - ? tools - : tools.filter((tool) => { - const statusEnum = - status.toUpperCase() as keyof typeof BackgroundToolStatus; - return tool.status === BackgroundToolStatus[statusEnum]; - }); - - // Filter by type if specified - const filteredTools = - type === 'all' - ? filteredByStatus - : filteredByStatus.filter( - (tool) => tool.type.toLowerCase() === type.toLowerCase(), - ); - - // Format the response - const formattedTools = filteredTools.map((tool) => { - const now = new Date(); - const startTime = tool.startTime; - const endTime = tool.endTime || now; - const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds - - return { - id: tool.id, - type: tool.type, - status: tool.status, - startTime: startTime.toISOString(), - ...(tool.endTime && { endTime: tool.endTime.toISOString() }), - runtime: parseFloat(runtime.toFixed(2)), - ...(verbose && { metadata: tool.metadata }), - }; - }); - - return { - tools: formattedTools, - count: formattedTools.length, - }; - }, - - logParameters: ( - { status = 'all', type = 'all', verbose = false }, - { logger }, - ) => { - logger.info( - `Listing ${type} background tools with status: ${status}, verbose: ${verbose}`, - ); - }, - - logReturns: (output, { logger }) => { - logger.info(`Found ${output.count} background tools`); - }, -}; diff --git a/packages/agent/src/tools/system/listShells.test.ts b/packages/agent/src/tools/system/listShells.test.ts index 10f13a1..94dc984 100644 --- a/packages/agent/src/tools/system/listShells.test.ts +++ b/packages/agent/src/tools/system/listShells.test.ts @@ -4,7 +4,7 @@ import { ToolContext } from '../../core/types.js'; import { getMockToolContext } from '../getTools.test.js'; import { listShellsTool } from './listShells.js'; -import { ShellStatus, shellTracker } from './ShellTracker.js'; +import { ShellStatus, ShellTracker } from './ShellTracker.js'; const toolContext: ToolContext = getMockToolContext(); @@ -15,6 +15,7 @@ vi.spyOn(Date, 'now').mockImplementation(() => mockNow); describe('listShellsTool', () => { beforeEach(() => { // Clear shells before each test + const shellTracker = new ShellTracker('test'); shellTracker['shells'] = new Map(); // Set up some test shells with different statuses diff --git a/packages/agent/src/tools/system/listShells.ts b/packages/agent/src/tools/system/listShells.ts index 0f4639f..7222dbd 100644 --- a/packages/agent/src/tools/system/listShells.ts +++ b/packages/agent/src/tools/system/listShells.ts @@ -3,7 +3,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema'; import { Tool } from '../../core/types.js'; -import { ShellStatus, shellTracker } from './ShellTracker.js'; +import { ShellStatus } from './ShellTracker.js'; const parameterSchema = z.object({ status: z @@ -45,7 +45,7 @@ export const listShellsTool: Tool = { execute: async ( { status = 'all', verbose = false }, - { logger }, + { logger, shellTracker }, ): Promise => { logger.verbose( `Listing shell processes with status: ${status}, verbose: ${verbose}`, diff --git a/packages/agent/src/tools/system/shellMessage.test.ts b/packages/agent/src/tools/system/shellMessage.test.ts index 7b63a5b..5997812 100644 --- a/packages/agent/src/tools/system/shellMessage.test.ts +++ b/packages/agent/src/tools/system/shellMessage.test.ts @@ -6,7 +6,7 @@ import { getMockToolContext } from '../getTools.test.js'; import { shellMessageTool, NodeSignals } from './shellMessage.js'; import { shellStartTool } from './shellStart.js'; -import { shellTracker } from './ShellTracker.js'; +import { ShellTracker } from './ShellTracker.js'; const toolContext: ToolContext = getMockToolContext(); @@ -22,6 +22,7 @@ const getInstanceId = ( describe('shellMessageTool', () => { let testInstanceId = ''; + const shellTracker = new ShellTracker('test'); beforeEach(() => { shellTracker.processStates.clear(); diff --git a/packages/agent/src/tools/system/shellMessage.ts b/packages/agent/src/tools/system/shellMessage.ts index 3dca577..3cf4265 100644 --- a/packages/agent/src/tools/system/shellMessage.ts +++ b/packages/agent/src/tools/system/shellMessage.ts @@ -4,7 +4,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema'; import { Tool } from '../../core/types.js'; import { sleep } from '../../utils/sleep.js'; -import { ShellStatus, shellTracker } from './ShellTracker.js'; +import { ShellStatus } from './ShellTracker.js'; // Define NodeJS signals as an enum export enum NodeSignals { @@ -95,7 +95,7 @@ export const shellMessageTool: Tool = { execute: async ( { instanceId, stdin, signal, showStdIn, showStdout }, - { logger }, + { logger, shellTracker }, ): Promise => { logger.verbose( `Interacting with shell process ${instanceId}${stdin ? ' with input' : ''}${signal ? ` with signal ${signal}` : ''}`, @@ -227,7 +227,7 @@ export const shellMessageTool: Tool = { } }, - logParameters: (input, { logger }) => { + logParameters: (input, { logger, shellTracker }) => { const processState = shellTracker.processStates.get(input.instanceId); const showStdIn = input.showStdIn !== undefined diff --git a/packages/agent/src/tools/system/shellStart.test.ts b/packages/agent/src/tools/system/shellStart.test.ts index 01ee643..30e0e81 100644 --- a/packages/agent/src/tools/system/shellStart.test.ts +++ b/packages/agent/src/tools/system/shellStart.test.ts @@ -5,11 +5,13 @@ import { sleep } from '../../utils/sleep.js'; import { getMockToolContext } from '../getTools.test.js'; import { shellStartTool } from './shellStart.js'; -import { shellTracker } from './ShellTracker.js'; +import { ShellTracker } from './ShellTracker.js'; const toolContext: ToolContext = getMockToolContext(); describe('shellStartTool', () => { + const shellTracker = new ShellTracker('test'); + beforeEach(() => { shellTracker.processStates.clear(); }); diff --git a/packages/agent/src/tools/system/shellStart.ts b/packages/agent/src/tools/system/shellStart.ts index 44f96a5..20ee1cc 100644 --- a/packages/agent/src/tools/system/shellStart.ts +++ b/packages/agent/src/tools/system/shellStart.ts @@ -7,7 +7,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; -import { ShellStatus, shellTracker } from './ShellTracker.js'; +import { ShellStatus } from './ShellTracker.js'; import type { ProcessState } from './ShellTracker.js'; @@ -81,7 +81,7 @@ export const shellStartTool: Tool = { showStdIn = false, showStdout = false, }, - { logger, workingDirectory }, + { logger, workingDirectory, shellTracker }, ): Promise => { if (showStdIn) { logger.info(`Command input: ${command}`); diff --git a/packages/cli/CHANGELOG.md b/packages/cli/CHANGELOG.md index d8663ba..488f37d 100644 --- a/packages/cli/CHANGELOG.md +++ b/packages/cli/CHANGELOG.md @@ -1,9 +1,8 @@ # [mycoder-v1.4.1](https://github.com/drivecore/mycoder/compare/mycoder-v1.4.0...mycoder-v1.4.1) (2025-03-14) - ### Bug Fixes -* improve profiling ([79a3df2](https://github.com/drivecore/mycoder/commit/79a3df2db13b8372666c6604ebe1666d33663be9)) +- improve profiling ([79a3df2](https://github.com/drivecore/mycoder/commit/79a3df2db13b8372666c6604ebe1666d33663be9)) # [mycoder-v1.4.0](https://github.com/drivecore/mycoder/compare/mycoder-v1.3.1...mycoder-v1.4.0) (2025-03-14) diff --git a/packages/cli/src/commands/$default.ts b/packages/cli/src/commands/$default.ts index b359941..e95647e 100644 --- a/packages/cli/src/commands/$default.ts +++ b/packages/cli/src/commands/$default.ts @@ -14,7 +14,9 @@ import { DEFAULT_CONFIG, AgentConfig, ModelProvider, - BackgroundTools, + BrowserTracker, + ShellTracker, + AgentTracker, } from 'mycoder-agent'; import { TokenTracker } from 'mycoder-agent/dist/core/tokens.js'; @@ -101,8 +103,6 @@ export async function executePrompt( // 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 = @@ -183,7 +183,9 @@ export async function executePrompt( model: config.model, maxTokens: config.maxTokens, temperature: config.temperature, - backgroundTools, + shellTracker: new ShellTracker('mainAgent'), + agentTracker: new AgentTracker('mainAgent'), + browserTracker: new BrowserTracker('mainAgent'), apiKey, }); @@ -201,7 +203,7 @@ export async function executePrompt( // Capture the error with Sentry captureException(error); } finally { - await backgroundTools.cleanup(); + // No cleanup needed here as it's handled by the cleanup utility } logger.log( diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts index 3c60fde..a3afbb2 100644 --- a/packages/cli/src/index.ts +++ b/packages/cli/src/index.ts @@ -13,7 +13,7 @@ import { command as toolsCommand } from './commands/tools.js'; import { SharedOptions, sharedOptions } from './options.js'; import { initSentry, captureException } from './sentry/index.js'; import { getConfigFromArgv, loadConfig } from './settings/config.js'; -import { cleanupResources, setupForceExit } from './utils/cleanup.js'; +import { setupForceExit } from './utils/cleanup.js'; import { enableProfiling, mark, reportTimings } from './utils/performance.js'; mark('After imports'); @@ -90,9 +90,6 @@ await main() // Report timings if profiling is enabled await reportTimings(); - // Clean up all resources before exit - await cleanupResources(); - // Setup a force exit as a failsafe // This ensures the process will exit even if there are lingering handles setupForceExit(5000); diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index 4b17828..b971361 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -1,75 +1,3 @@ -import { BrowserManager, shellTracker } from 'mycoder-agent'; -import { agentStates } from 'mycoder-agent/dist/tools/interaction/agentStart.js'; - -/** - * Handles cleanup of resources before application exit - * Ensures all browser sessions and shell processes are terminated - */ -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 unknown as { __BROWSER_MANAGER__?: BrowserManager } - ).__BROWSER_MANAGER__; - if (browserManager) { - console.log('Closing all browser sessions...'); - await browserManager.closeAllSessions(); - } - } catch (error) { - console.error('Error closing browser sessions:', error); - } - - // 2. Clean up shell processes - try { - const runningShells = shellTracker.getShells(); - if (runningShells.length > 0) { - console.log(`Terminating ${runningShells.length} shell processes...`); - await shellTracker.cleanupAllShells(); - } - } catch (error) { - console.error('Error terminating shell processes:', error); - } - - // 3. Give async operations a moment to complete - await new Promise((resolve) => setTimeout(resolve, 1000)); - - console.log('Cleanup completed'); -} - /** * Force exits the process after a timeout * This is a failsafe to ensure the process exits even if there are lingering handles