From 65378e34b035699f61b701679742ba9a7e667215 Mon Sep 17 00:00:00 2001 From: "Ben Houston (via MyCoder)" Date: Fri, 14 Mar 2025 19:52:38 +0000 Subject: [PATCH 1/2] feat: implement ShellTracker to decouple from backgroundTools --- .../src/core/backgroundTools.cleanup.test.ts | 32 ++-- .../agent/src/core/backgroundTools.test.ts | 52 +----- packages/agent/src/core/backgroundTools.ts | 85 +-------- packages/agent/src/tools/getTools.ts | 2 + .../src/tools/system/ShellTracker.test.ts | 119 ++++++++++++ .../agent/src/tools/system/ShellTracker.ts | 171 ++++++++++++++++++ .../src/tools/system/listBackgroundTools.ts | 5 +- .../agent/src/tools/system/listShells.test.ts | 113 ++++++++++++ packages/agent/src/tools/system/listShells.ts | 98 ++++++++++ .../src/tools/system/shellMessage.test.ts | 23 +-- .../agent/src/tools/system/shellMessage.ts | 49 ++--- .../agent/src/tools/system/shellStart.test.ts | 19 +- packages/agent/src/tools/system/shellStart.ts | 53 ++---- 13 files changed, 583 insertions(+), 238 deletions(-) create mode 100644 packages/agent/src/tools/system/ShellTracker.test.ts create mode 100644 packages/agent/src/tools/system/ShellTracker.ts create mode 100644 packages/agent/src/tools/system/listShells.test.ts create mode 100644 packages/agent/src/tools/system/listShells.ts diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/packages/agent/src/core/backgroundTools.cleanup.test.ts index 3adec5d..c04c1f1 100644 --- a/packages/agent/src/core/backgroundTools.cleanup.test.ts +++ b/packages/agent/src/core/backgroundTools.cleanup.test.ts @@ -3,7 +3,7 @@ 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 { shellTracker } from '../tools/system/ShellTracker.js'; import { BackgroundTools, BackgroundToolStatus } from './backgroundTools'; import { Tool } from './types'; @@ -43,9 +43,12 @@ vi.mock('../tools/browser/BrowserManager.js', () => { }; }); -vi.mock('../tools/system/shellStart.js', () => { +vi.mock('../tools/system/ShellTracker.js', () => { return { - processStates: new Map(), + shellTracker: { + processStates: new Map(), + cleanupAllShells: vi.fn().mockResolvedValue(undefined), + }, }; }); @@ -87,8 +90,8 @@ describe('BackgroundTools cleanup', () => { showStdout: false, }; - processStates.clear(); - processStates.set('shell-1', mockProcessState as any); + shellTracker.processStates.clear(); + shellTracker.processStates.set('shell-1', mockProcessState as any); // Setup mock agent states const mockAgentState: MockAgentState = { @@ -119,7 +122,7 @@ describe('BackgroundTools cleanup', () => { ).__BROWSER_MANAGER__ = undefined; // Clear mock states - processStates.clear(); + shellTracker.processStates.clear(); agentStates.clear(); }); @@ -142,24 +145,11 @@ describe('BackgroundTools cleanup', () => { }); 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); + // Check that shellTracker.cleanupAllShells was called + expect(shellTracker.cleanupAllShells).toHaveBeenCalled(); }); it('should clean up sub-agents', async () => { diff --git a/packages/agent/src/core/backgroundTools.test.ts b/packages/agent/src/core/backgroundTools.test.ts index 4b0e5c3..bc0a56b 100644 --- a/packages/agent/src/core/backgroundTools.test.ts +++ b/packages/agent/src/core/backgroundTools.test.ts @@ -19,22 +19,6 @@ describe('BackgroundToolRegistry', () => { backgroundTools.tools = new Map(); }); - it('should register a shell process', () => { - const id = backgroundTools.registerShell('ls -la'); - - expect(id).toBe('test-id-1'); - - const tool = backgroundTools.getToolById(id); - expect(tool).toBeDefined(); - if (tool) { - expect(tool.type).toBe(BackgroundToolType.SHELL); - expect(tool.status).toBe(BackgroundToolStatus.RUNNING); - if (tool.type === BackgroundToolType.SHELL) { - expect(tool.metadata.command).toBe('ls -la'); - } - } - }); - it('should register a browser process', () => { const id = backgroundTools.registerBrowser('https://example.com'); @@ -51,30 +35,6 @@ describe('BackgroundToolRegistry', () => { } }); - it('should update tool status', () => { - const id = backgroundTools.registerShell('sleep 10'); - - const updated = backgroundTools.updateToolStatus( - id, - BackgroundToolStatus.COMPLETED, - { - exitCode: 0, - }, - ); - - expect(updated).toBe(true); - - const tool = backgroundTools.getToolById(id); - expect(tool).toBeDefined(); - if (tool) { - expect(tool.status).toBe(BackgroundToolStatus.COMPLETED); - expect(tool.endTime).toBeDefined(); - if (tool.type === BackgroundToolType.SHELL) { - expect(tool.metadata.exitCode).toBe(0); - } - } - }); - it('should return false when updating non-existent tool', () => { const updated = backgroundTools.updateToolStatus( 'non-existent-id', @@ -91,33 +51,33 @@ describe('BackgroundToolRegistry', () => { // Add a completed tool from 25 hours ago const oldTool = { id: 'old-tool', - type: BackgroundToolType.SHELL, + 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: { command: 'echo old' }, + metadata: { url: 'https://example.com' }, }; // Add a completed tool from 10 hours ago const recentTool = { id: 'recent-tool', - type: BackgroundToolType.SHELL, + 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: { command: 'echo recent' }, + metadata: { url: 'https://example.com' }, }; // Add a running tool from 25 hours ago const oldRunningTool = { id: 'old-running-tool', - type: BackgroundToolType.SHELL, + type: BackgroundToolType.BROWSER, status: BackgroundToolStatus.RUNNING, startTime: new Date(Date.now() - 25 * 60 * 60 * 1000), agentId: 'agent-1', - metadata: { command: 'sleep 100' }, + metadata: { url: 'https://example.com' }, }; registry.tools.set('old-tool', oldTool); diff --git a/packages/agent/src/core/backgroundTools.ts b/packages/agent/src/core/backgroundTools.ts index 45c61c1..b76abd3 100644 --- a/packages/agent/src/core/backgroundTools.ts +++ b/packages/agent/src/core/backgroundTools.ts @@ -3,11 +3,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'; +import { shellTracker } from '../tools/system/ShellTracker.js'; // Types of background processes we can track export enum BackgroundToolType { - SHELL = 'shell', BROWSER = 'browser', AGENT = 'agent', } @@ -30,17 +29,6 @@ export interface BackgroundTool { metadata: Record; // Additional tool-specific information } -// Shell process specific data -export interface ShellBackgroundTool extends BackgroundTool { - type: BackgroundToolType.SHELL; - metadata: { - command: string; - exitCode?: number | null; - signaled?: boolean; - error?: string; - }; -} - // Browser process specific data export interface BrowserBackgroundTool extends BackgroundTool { type: BackgroundToolType.BROWSER; @@ -60,10 +48,7 @@ export interface AgentBackgroundTool extends BackgroundTool { } // Utility type for all background tool types -export type AnyBackgroundTool = - | ShellBackgroundTool - | BrowserBackgroundTool - | AgentBackgroundTool; +export type AnyBackgroundTool = BrowserBackgroundTool | AgentBackgroundTool; /** * Registry to keep track of all background processes @@ -74,22 +59,6 @@ export class BackgroundTools { // Private constructor for singleton pattern constructor(readonly ownerName: string) {} - // Register a new shell process - public registerShell(command: string): string { - const id = uuidv4(); - const tool: ShellBackgroundTool = { - id, - type: BackgroundToolType.SHELL, - status: BackgroundToolStatus.RUNNING, - startTime: new Date(), - metadata: { - command, - }, - }; - this.tools.set(id, tool); - return id; - } - // Register a new browser process public registerBrowser(url?: string): string { const id = uuidv4(); @@ -177,12 +146,6 @@ export class BackgroundTools { 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 && @@ -193,19 +156,15 @@ export class BackgroundTools { const browserCleanupPromises = browserTools.map((tool) => this.cleanupBrowserSession(tool), ); - const shellCleanupPromises = shellTools.map((tool) => - this.cleanupShellProcess(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, - ...shellCleanupPromises, - ...agentCleanupPromises, - ]); + await Promise.all([...browserCleanupPromises, ...agentCleanupPromises]); } /** @@ -230,38 +189,6 @@ export class BackgroundTools { } } - /** - * 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 diff --git a/packages/agent/src/tools/getTools.ts b/packages/agent/src/tools/getTools.ts index 79ee272..39e84d6 100644 --- a/packages/agent/src/tools/getTools.ts +++ b/packages/agent/src/tools/getTools.ts @@ -10,6 +10,7 @@ import { fetchTool } from './io/fetch.js'; import { textEditorTool } from './io/textEditor.js'; import { createMcpTool } from './mcp.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'; import { shellStartTool } from './system/shellStart.js'; @@ -41,6 +42,7 @@ export function getTools(options?: GetToolsOptions): 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, ]; // Only include userPrompt tool if enabled diff --git a/packages/agent/src/tools/system/ShellTracker.test.ts b/packages/agent/src/tools/system/ShellTracker.test.ts new file mode 100644 index 0000000..0d44cdc --- /dev/null +++ b/packages/agent/src/tools/system/ShellTracker.test.ts @@ -0,0 +1,119 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +import { ShellStatus, shellTracker } from './ShellTracker.js'; + +// Mock uuid to return predictable IDs for testing +vi.mock('uuid', () => ({ + v4: vi + .fn() + .mockReturnValueOnce('test-id-1') + .mockReturnValueOnce('test-id-2') + .mockReturnValueOnce('test-id-3'), +})); + +describe('ShellTracker', () => { + beforeEach(() => { + // Clear all registered shells before each test + shellTracker['shells'] = new Map(); + shellTracker.processStates.clear(); + }); + + it('should register a shell process', () => { + const id = shellTracker.registerShell('ls -la'); + + expect(id).toBe('test-id-1'); + + const shell = shellTracker.getShellById(id); + expect(shell).toBeDefined(); + if (shell) { + expect(shell.status).toBe(ShellStatus.RUNNING); + expect(shell.metadata.command).toBe('ls -la'); + } + }); + + it('should update shell status', () => { + const id = shellTracker.registerShell('sleep 10'); + + const updated = shellTracker.updateShellStatus(id, ShellStatus.COMPLETED, { + exitCode: 0, + }); + + expect(updated).toBe(true); + + const shell = shellTracker.getShellById(id); + expect(shell).toBeDefined(); + if (shell) { + expect(shell.status).toBe(ShellStatus.COMPLETED); + expect(shell.endTime).toBeDefined(); + expect(shell.metadata.exitCode).toBe(0); + } + }); + + it('should return false when updating non-existent shell', () => { + const updated = shellTracker.updateShellStatus( + 'non-existent-id', + ShellStatus.COMPLETED, + ); + + expect(updated).toBe(false); + }); + + it('should filter shells by status', () => { + // Create shells with different statuses + const shell1 = { + id: 'shell-1', + status: ShellStatus.RUNNING, + startTime: new Date(), + metadata: { + command: 'command1', + }, + }; + + const shell2 = { + id: 'shell-2', + status: ShellStatus.COMPLETED, + startTime: new Date(), + endTime: new Date(), + metadata: { + command: 'command2', + exitCode: 0, + }, + }; + + const shell3 = { + id: 'shell-3', + status: ShellStatus.ERROR, + startTime: new Date(), + endTime: new Date(), + metadata: { + command: 'command3', + exitCode: 1, + error: 'Error message', + }, + }; + + // Add the shells directly to the map + shellTracker['shells'].set('shell-1', shell1); + shellTracker['shells'].set('shell-2', shell2); + shellTracker['shells'].set('shell-3', shell3); + + // Get all shells + const allShells = shellTracker.getShells(); + expect(allShells.length).toBe(3); + + // Get running shells + const runningShells = shellTracker.getShells(ShellStatus.RUNNING); + expect(runningShells.length).toBe(1); + expect(runningShells[0].id).toBe('shell-1'); + + // Get completed shells + const completedShells = shellTracker.getShells(ShellStatus.COMPLETED); + expect(completedShells.length).toBe(1); + expect(completedShells[0].id).toBe('shell-2'); + + // Get error shells + const errorShells = shellTracker.getShells(ShellStatus.ERROR); + expect(errorShells.length).toBe(1); + expect(errorShells[0].id).toBe('shell-3'); + }); +}); diff --git a/packages/agent/src/tools/system/ShellTracker.ts b/packages/agent/src/tools/system/ShellTracker.ts new file mode 100644 index 0000000..c7dd4bf --- /dev/null +++ b/packages/agent/src/tools/system/ShellTracker.ts @@ -0,0 +1,171 @@ +import { v4 as uuidv4 } from 'uuid'; + +import type { ChildProcess } from 'child_process'; + +// Status of a shell process +export enum ShellStatus { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + TERMINATED = 'terminated', +} + +// Define ProcessState type +export type ProcessState = { + process: ChildProcess; + command: string; + stdout: string[]; + stderr: string[]; + state: { + completed: boolean; + signaled: boolean; + exitCode: number | null; + }; + showStdIn: boolean; + showStdout: boolean; +}; + +// Shell process specific data +export interface ShellProcess { + id: string; + status: ShellStatus; + startTime: Date; + endTime?: Date; + metadata: { + command: string; + exitCode?: number | null; + signaled?: boolean; + error?: string; + [key: string]: any; // Additional shell-specific information + }; +} + +/** + * 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; + } + + // Register a new shell process + public registerShell(command: string): string { + const id = uuidv4(); + const shell: ShellProcess = { + id, + status: ShellStatus.RUNNING, + startTime: new Date(), + metadata: { + command, + }, + }; + this.shells.set(id, shell); + return id; + } + + // Update the status of a shell process + public updateShellStatus( + id: string, + status: ShellStatus, + metadata?: Record, + ): boolean { + const shell = this.shells.get(id); + if (!shell) { + return false; + } + + shell.status = status; + + if ( + status === ShellStatus.COMPLETED || + status === ShellStatus.ERROR || + status === ShellStatus.TERMINATED + ) { + shell.endTime = new Date(); + } + + if (metadata) { + shell.metadata = { ...shell.metadata, ...metadata }; + } + + return true; + } + + // Get all shell processes + public getShells(status?: ShellStatus): ShellProcess[] { + const result: ShellProcess[] = []; + for (const shell of this.shells.values()) { + if (!status || shell.status === status) { + result.push(shell); + } + } + return result; + } + + // Get a specific shell process by ID + public getShellById(id: string): ShellProcess | undefined { + return this.shells.get(id); + } + + /** + * Cleans up a shell process + * @param id The ID of the shell process to clean up + */ + public async cleanupShellProcess(id: string): Promise { + try { + const shell = this.shells.get(id); + if (!shell) { + return; + } + + const processState = this.processStates.get(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.updateShellStatus(id, ShellStatus.TERMINATED); + } catch (error) { + this.updateShellStatus(id, ShellStatus.ERROR, { + error: error instanceof Error ? error.message : String(error), + }); + } + } + + /** + * Cleans up all running shell processes + */ + public async cleanupAllShells(): Promise { + const runningShells = this.getShells(ShellStatus.RUNNING); + const cleanupPromises = runningShells.map((shell) => + this.cleanupShellProcess(shell.id), + ); + await Promise.all(cleanupPromises); + } +} + +// Export a singleton instance +export const shellTracker = ShellTracker.getInstance(); diff --git a/packages/agent/src/tools/system/listBackgroundTools.ts b/packages/agent/src/tools/system/listBackgroundTools.ts index bc7608e..41526a7 100644 --- a/packages/agent/src/tools/system/listBackgroundTools.ts +++ b/packages/agent/src/tools/system/listBackgroundTools.ts @@ -10,7 +10,7 @@ const parameterSchema = z.object({ .optional() .describe('Filter tools by status (default: "all")'), type: z - .enum(['all', 'shell', 'browser', 'agent']) + .enum(['all', 'browser', 'agent']) .optional() .describe('Filter tools by type (default: "all")'), verbose: z @@ -39,8 +39,7 @@ type ReturnType = z.infer; export const listBackgroundToolsTool: Tool = { name: 'listBackgroundTools', - description: - 'Lists all background tools (shells, browsers, agents) and their status', + description: 'Lists all background tools (browsers, agents) and their status', logPrefix: '🔍', parameters: parameterSchema, returns: returnSchema, diff --git a/packages/agent/src/tools/system/listShells.test.ts b/packages/agent/src/tools/system/listShells.test.ts new file mode 100644 index 0000000..95b6647 --- /dev/null +++ b/packages/agent/src/tools/system/listShells.test.ts @@ -0,0 +1,113 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +import { ToolContext } from '../../core/types.js'; +import { getMockToolContext } from '../getTools.test.js'; + +import { listShellsTool } from './listShells.js'; +import { ShellStatus, shellTracker } from './ShellTracker.js'; + +const toolContext: ToolContext = getMockToolContext(); + +// Mock Date.now to return a consistent timestamp +const mockNow = new Date('2023-01-01T00:00:00Z').getTime(); +vi.spyOn(Date, 'now').mockImplementation(() => mockNow); + +describe('listShellsTool', () => { + beforeEach(() => { + // Clear shells before each test + shellTracker['shells'] = new Map(); + + // Set up some test shells with different statuses + const shell1 = { + id: 'shell-1', + status: ShellStatus.RUNNING, + startTime: new Date(mockNow - 1000 * 60 * 5), // 5 minutes ago + metadata: { + command: 'sleep 100', + }, + }; + + const shell2 = { + id: 'shell-2', + status: ShellStatus.COMPLETED, + startTime: new Date(mockNow - 1000 * 60 * 10), // 10 minutes ago + endTime: new Date(mockNow - 1000 * 60 * 9), // 9 minutes ago + metadata: { + command: 'echo "test"', + exitCode: 0, + }, + }; + + const shell3 = { + id: 'shell-3', + status: ShellStatus.ERROR, + startTime: new Date(mockNow - 1000 * 60 * 15), // 15 minutes ago + endTime: new Date(mockNow - 1000 * 60 * 14), // 14 minutes ago + metadata: { + command: 'nonexistentcommand', + exitCode: 127, + error: 'Command not found', + }, + }; + + // Add the shells to the tracker + shellTracker['shells'].set('shell-1', shell1); + shellTracker['shells'].set('shell-2', shell2); + shellTracker['shells'].set('shell-3', shell3); + }); + + it('should list all shells by default', async () => { + const result = await listShellsTool.execute({}, toolContext); + + expect(result.shells.length).toBe(3); + expect(result.count).toBe(3); + + // Check that shells are properly formatted + const shell1 = result.shells.find((s) => s.id === 'shell-1'); + expect(shell1).toBeDefined(); + expect(shell1?.status).toBe(ShellStatus.RUNNING); + expect(shell1?.command).toBe('sleep 100'); + expect(shell1?.runtime).toBeGreaterThan(0); + + // Metadata should not be included by default + expect(shell1?.metadata).toBeUndefined(); + }); + + it('should filter shells by status', async () => { + const result = await listShellsTool.execute( + { status: 'running' }, + toolContext, + ); + + expect(result.shells.length).toBe(1); + expect(result.count).toBe(1); + expect(result.shells[0].id).toBe('shell-1'); + expect(result.shells[0].status).toBe(ShellStatus.RUNNING); + }); + + it('should include metadata when verbose is true', async () => { + const result = await listShellsTool.execute({ verbose: true }, toolContext); + + expect(result.shells.length).toBe(3); + + // Check that metadata is included + const shell3 = result.shells.find((s) => s.id === 'shell-3'); + expect(shell3).toBeDefined(); + expect(shell3?.metadata).toBeDefined(); + expect(shell3?.metadata?.exitCode).toBe(127); + expect(shell3?.metadata?.error).toBe('Command not found'); + }); + + it('should combine status filter with verbose option', async () => { + const result = await listShellsTool.execute( + { status: 'error', verbose: true }, + toolContext, + ); + + expect(result.shells.length).toBe(1); + expect(result.shells[0].id).toBe('shell-3'); + expect(result.shells[0].status).toBe(ShellStatus.ERROR); + expect(result.shells[0].metadata).toBeDefined(); + expect(result.shells[0].metadata?.error).toBe('Command not found'); + }); +}); diff --git a/packages/agent/src/tools/system/listShells.ts b/packages/agent/src/tools/system/listShells.ts new file mode 100644 index 0000000..0f4639f --- /dev/null +++ b/packages/agent/src/tools/system/listShells.ts @@ -0,0 +1,98 @@ +import { z } from 'zod'; +import { zodToJsonSchema } from 'zod-to-json-schema'; + +import { Tool } from '../../core/types.js'; + +import { ShellStatus, shellTracker } from './ShellTracker.js'; + +const parameterSchema = z.object({ + status: z + .enum(['all', 'running', 'completed', 'error', 'terminated']) + .optional() + .describe('Filter shells by status (default: "all")'), + verbose: z + .boolean() + .optional() + .describe('Include detailed metadata about each shell (default: false)'), +}); + +const returnSchema = z.object({ + shells: z.array( + z.object({ + id: z.string(), + status: z.string(), + startTime: z.string(), + endTime: z.string().optional(), + runtime: z.number().describe('Runtime in seconds'), + command: z.string(), + metadata: z.record(z.any()).optional(), + }), + ), + count: z.number(), +}); + +type Parameters = z.infer; +type ReturnType = z.infer; + +export const listShellsTool: Tool = { + name: 'listShells', + description: 'Lists all shell processes and their status', + logPrefix: '🔍', + parameters: parameterSchema, + returns: returnSchema, + parametersJsonSchema: zodToJsonSchema(parameterSchema), + returnsJsonSchema: zodToJsonSchema(returnSchema), + + execute: async ( + { status = 'all', verbose = false }, + { logger }, + ): Promise => { + logger.verbose( + `Listing shell processes with status: ${status}, verbose: ${verbose}`, + ); + + // Get all shells + let shells = shellTracker.getShells(); + + // Filter by status if specified + if (status !== 'all') { + const statusEnum = status.toUpperCase() as keyof typeof ShellStatus; + shells = shells.filter( + (shell) => shell.status === ShellStatus[statusEnum], + ); + } + + // Format the response + const formattedShells = shells.map((shell) => { + const now = new Date(); + const startTime = shell.startTime; + const endTime = shell.endTime || now; + const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds + + return { + id: shell.id, + status: shell.status, + startTime: startTime.toISOString(), + ...(shell.endTime && { endTime: shell.endTime.toISOString() }), + runtime: parseFloat(runtime.toFixed(2)), + command: shell.metadata.command, + ...(verbose && { metadata: shell.metadata }), + }; + }); + + return { + shells: formattedShells, + count: formattedShells.length, + }; + }, + + logParameters: ({ status = 'all', verbose = false }, { logger }) => { + logger.info( + `Listing shell processes with status: ${status}, verbose: ${verbose}`, + ); + }, + + logReturns: (output, { logger }) => { + logger.info(`Found ${output.count} shell processes`); + }, +}; diff --git a/packages/agent/src/tools/system/shellMessage.test.ts b/packages/agent/src/tools/system/shellMessage.test.ts index b78da0e..7b63a5b 100644 --- a/packages/agent/src/tools/system/shellMessage.test.ts +++ b/packages/agent/src/tools/system/shellMessage.test.ts @@ -5,7 +5,8 @@ import { sleep } from '../../utils/sleep.js'; import { getMockToolContext } from '../getTools.test.js'; import { shellMessageTool, NodeSignals } from './shellMessage.js'; -import { processStates, shellStartTool } from './shellStart.js'; +import { shellStartTool } from './shellStart.js'; +import { shellTracker } from './ShellTracker.js'; const toolContext: ToolContext = getMockToolContext(); @@ -23,14 +24,14 @@ describe('shellMessageTool', () => { let testInstanceId = ''; beforeEach(() => { - processStates.clear(); + shellTracker.processStates.clear(); }); afterEach(() => { - for (const processState of processStates.values()) { + for (const processState of shellTracker.processStates.values()) { processState.process.kill(); } - processStates.clear(); + shellTracker.processStates.clear(); }); it('should interact with a running process', async () => { @@ -62,7 +63,7 @@ describe('shellMessageTool', () => { expect(result.completed).toBe(false); // Verify the instance ID is valid - expect(processStates.has(testInstanceId)).toBe(true); + expect(shellTracker.processStates.has(testInstanceId)).toBe(true); }); it('should handle nonexistent process', async () => { @@ -104,7 +105,7 @@ describe('shellMessageTool', () => { expect(result.completed).toBe(true); // Process should still be in processStates even after completion - expect(processStates.has(instanceId)).toBe(true); + expect(shellTracker.processStates.has(instanceId)).toBe(true); }); it('should handle SIGTERM signal correctly', async () => { @@ -207,7 +208,7 @@ describe('shellMessageTool', () => { expect(checkResult.signaled).toBe(true); expect(checkResult.completed).toBe(true); - expect(processStates.has(instanceId)).toBe(true); + expect(shellTracker.processStates.has(instanceId)).toBe(true); }); it('should respect showStdIn and showStdout parameters', async () => { @@ -224,7 +225,7 @@ describe('shellMessageTool', () => { const instanceId = getInstanceId(startResult); // Verify process state has default visibility settings - const processState = processStates.get(instanceId); + const processState = shellTracker.processStates.get(instanceId); expect(processState?.showStdIn).toBe(false); expect(processState?.showStdout).toBe(false); @@ -241,7 +242,7 @@ describe('shellMessageTool', () => { ); // Verify process state still exists - expect(processStates.has(instanceId)).toBe(true); + expect(shellTracker.processStates.has(instanceId)).toBe(true); }); it('should inherit visibility settings from process state', async () => { @@ -260,7 +261,7 @@ describe('shellMessageTool', () => { const instanceId = getInstanceId(startResult); // Verify process state has the specified visibility settings - const processState = processStates.get(instanceId); + const processState = shellTracker.processStates.get(instanceId); expect(processState?.showStdIn).toBe(true); expect(processState?.showStdout).toBe(true); @@ -275,6 +276,6 @@ describe('shellMessageTool', () => { ); // Verify process state still exists - expect(processStates.has(instanceId)).toBe(true); + expect(shellTracker.processStates.has(instanceId)).toBe(true); }); }); diff --git a/packages/agent/src/tools/system/shellMessage.ts b/packages/agent/src/tools/system/shellMessage.ts index 17655d9..3dca577 100644 --- a/packages/agent/src/tools/system/shellMessage.ts +++ b/packages/agent/src/tools/system/shellMessage.ts @@ -1,11 +1,10 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; -import { BackgroundToolStatus } from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { sleep } from '../../utils/sleep.js'; -import { processStates } from './shellStart.js'; +import { ShellStatus, shellTracker } from './ShellTracker.js'; // Define NodeJS signals as an enum export enum NodeSignals { @@ -96,14 +95,14 @@ export const shellMessageTool: Tool = { execute: async ( { instanceId, stdin, signal, showStdIn, showStdout }, - { logger, backgroundTools }, + { logger }, ): Promise => { logger.verbose( `Interacting with shell process ${instanceId}${stdin ? ' with input' : ''}${signal ? ` with signal ${signal}` : ''}`, ); try { - const processState = processStates.get(instanceId); + const processState = shellTracker.processStates.get(instanceId); if (!processState) { throw new Error(`No process found with ID ${instanceId}`); } @@ -118,44 +117,32 @@ export const shellMessageTool: Tool = { // If the process is already terminated, we'll just mark it as signaled anyway processState.state.signaled = true; - // Update background tool registry if signal failed - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.ERROR, - { - error: `Failed to send signal ${signal}: ${String(error)}`, - signalAttempted: signal, - }, - ); + // Update shell tracker if signal failed + shellTracker.updateShellStatus(instanceId, ShellStatus.ERROR, { + error: `Failed to send signal ${signal}: ${String(error)}`, + signalAttempted: signal, + }); logger.verbose( `Failed to send signal ${signal}: ${String(error)}, but marking as signaled anyway`, ); } - // Update background tool registry with signal information + // Update shell tracker with signal information if ( signal === 'SIGTERM' || signal === 'SIGKILL' || signal === 'SIGINT' ) { - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.TERMINATED, - { - signal, - terminatedByUser: true, - }, - ); + shellTracker.updateShellStatus(instanceId, ShellStatus.TERMINATED, { + signal, + terminatedByUser: true, + }); } else { - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.RUNNING, - { - signal, - signaled: true, - }, - ); + shellTracker.updateShellStatus(instanceId, ShellStatus.RUNNING, { + signal, + signaled: true, + }); } } @@ -241,7 +228,7 @@ export const shellMessageTool: Tool = { }, logParameters: (input, { logger }) => { - const processState = processStates.get(input.instanceId); + const processState = shellTracker.processStates.get(input.instanceId); const showStdIn = input.showStdIn !== undefined ? input.showStdIn diff --git a/packages/agent/src/tools/system/shellStart.test.ts b/packages/agent/src/tools/system/shellStart.test.ts index 223560c..01ee643 100644 --- a/packages/agent/src/tools/system/shellStart.test.ts +++ b/packages/agent/src/tools/system/shellStart.test.ts @@ -4,20 +4,21 @@ import { ToolContext } from '../../core/types.js'; import { sleep } from '../../utils/sleep.js'; import { getMockToolContext } from '../getTools.test.js'; -import { processStates, shellStartTool } from './shellStart.js'; +import { shellStartTool } from './shellStart.js'; +import { shellTracker } from './ShellTracker.js'; const toolContext: ToolContext = getMockToolContext(); describe('shellStartTool', () => { beforeEach(() => { - processStates.clear(); + shellTracker.processStates.clear(); }); afterEach(() => { - for (const processState of processStates.values()) { + for (const processState of shellTracker.processStates.values()) { processState.process.kill(); } - processStates.clear(); + shellTracker.processStates.clear(); }); it('should handle fast commands in sync mode', async () => { @@ -83,7 +84,7 @@ describe('shellStartTool', () => { ); // Even sync results should be in processStates - expect(processStates.size).toBeGreaterThan(0); + expect(shellTracker.processStates.size).toBeGreaterThan(0); expect(syncResult.mode).toBe('sync'); expect(syncResult.error).toBeUndefined(); if (syncResult.mode === 'sync') { @@ -101,7 +102,7 @@ describe('shellStartTool', () => { ); if (asyncResult.mode === 'async') { - expect(processStates.has(asyncResult.instanceId)).toBe(true); + expect(shellTracker.processStates.has(asyncResult.instanceId)).toBe(true); } }); @@ -120,7 +121,7 @@ describe('shellStartTool', () => { expect(result.instanceId).toBeDefined(); expect(result.error).toBeUndefined(); - const processState = processStates.get(result.instanceId); + const processState = shellTracker.processStates.get(result.instanceId); expect(processState).toBeDefined(); if (processState?.process.stdin) { @@ -177,7 +178,9 @@ describe('shellStartTool', () => { ); if (asyncResult.mode === 'async') { - const processState = processStates.get(asyncResult.instanceId); + const processState = shellTracker.processStates.get( + asyncResult.instanceId, + ); expect(processState).toBeDefined(); expect(processState?.showStdIn).toBe(true); expect(processState?.showStdout).toBe(true); diff --git a/packages/agent/src/tools/system/shellStart.ts b/packages/agent/src/tools/system/shellStart.ts index c98c7e7..44f96a5 100644 --- a/packages/agent/src/tools/system/shellStart.ts +++ b/packages/agent/src/tools/system/shellStart.ts @@ -4,30 +4,12 @@ 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 type { ChildProcess } from 'child_process'; +import { ShellStatus, shellTracker } from './ShellTracker.js'; -// Define ProcessState type -type ProcessState = { - process: ChildProcess; - command: string; - stdout: string[]; - stderr: string[]; - state: { - completed: boolean; - signaled: boolean; - exitCode: number | null; - }; - showStdIn: boolean; - showStdout: boolean; -}; - -// Global map to store process state -// This is exported so it can be accessed for cleanup -export const processStates: Map = new Map(); +import type { ProcessState } from './ShellTracker.js'; const parameterSchema = z.object({ command: z.string().describe('The shell command to execute'), @@ -99,7 +81,7 @@ export const shellStartTool: Tool = { showStdIn = false, showStdout = false, }, - { logger, workingDirectory, backgroundTools }, + { logger, workingDirectory }, ): Promise => { if (showStdIn) { logger.info(`Command input: ${command}`); @@ -111,8 +93,8 @@ export const shellStartTool: Tool = { // Generate a unique ID for this process const instanceId = uuidv4(); - // Register this shell process with the background tool registry - backgroundTools.registerShell(command); + // Register this shell process with the shell tracker + shellTracker.registerShell(command); let hasResolved = false; @@ -134,8 +116,8 @@ export const shellStartTool: Tool = { showStdout, }; - // Initialize combined process state - processStates.set(instanceId, processState); + // Initialize process state + shellTracker.processStates.set(instanceId, processState); // Handle process events if (process.stdout) @@ -160,14 +142,10 @@ export const shellStartTool: Tool = { logger.error(`[${instanceId}] Process error: ${error.message}`); processState.state.completed = true; - // Update background tool registry with error status - backgroundTools.updateToolStatus( - instanceId, - BackgroundToolStatus.ERROR, - { - error: error.message, - }, - ); + // Update shell tracker with error status + shellTracker.updateShellStatus(instanceId, ShellStatus.ERROR, { + error: error.message, + }); if (!hasResolved) { hasResolved = true; @@ -190,12 +168,9 @@ export const shellStartTool: Tool = { processState.state.signaled = signal !== null; processState.state.exitCode = code; - // Update background tool registry with completed status - const status = - code === 0 - ? BackgroundToolStatus.COMPLETED - : BackgroundToolStatus.ERROR; - backgroundTools.updateToolStatus(instanceId, status, { + // Update shell tracker with completed status + const status = code === 0 ? ShellStatus.COMPLETED : ShellStatus.ERROR; + shellTracker.updateShellStatus(instanceId, status, { exitCode: code, signaled: signal !== null, }); From 3dca7670bed4884650b43d431c09a14d2673eb58 Mon Sep 17 00:00:00 2001 From: "Ben Houston (via MyCoder)" Date: Fri, 14 Mar 2025 19:56:17 +0000 Subject: [PATCH 2/2] fix: update CLI cleanup to use ShellTracker instead of processStates --- packages/agent/src/index.ts | 2 ++ .../src/tools/system/ShellTracker.test.ts | 9 ++++-- .../agent/src/tools/system/listShells.test.ts | 14 +++++---- packages/cli/src/utils/cleanup.ts | 29 ++++--------------- 4 files changed, 21 insertions(+), 33 deletions(-) diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index e0a9567..20c5fa1 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -10,6 +10,8 @@ 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'; // Tools - Browser export * from './tools/browser/BrowserManager.js'; diff --git a/packages/agent/src/tools/system/ShellTracker.test.ts b/packages/agent/src/tools/system/ShellTracker.test.ts index 0d44cdc..7fd8fdb 100644 --- a/packages/agent/src/tools/system/ShellTracker.test.ts +++ b/packages/agent/src/tools/system/ShellTracker.test.ts @@ -104,16 +104,19 @@ describe('ShellTracker', () => { // Get running shells const runningShells = shellTracker.getShells(ShellStatus.RUNNING); expect(runningShells.length).toBe(1); - expect(runningShells[0].id).toBe('shell-1'); + expect(runningShells.length).toBe(1); + expect(runningShells[0]!.id).toBe('shell-1'); // Get completed shells const completedShells = shellTracker.getShells(ShellStatus.COMPLETED); expect(completedShells.length).toBe(1); - expect(completedShells[0].id).toBe('shell-2'); + expect(completedShells.length).toBe(1); + expect(completedShells[0]!.id).toBe('shell-2'); // Get error shells const errorShells = shellTracker.getShells(ShellStatus.ERROR); expect(errorShells.length).toBe(1); - expect(errorShells[0].id).toBe('shell-3'); + expect(errorShells.length).toBe(1); + expect(errorShells[0]!.id).toBe('shell-3'); }); }); diff --git a/packages/agent/src/tools/system/listShells.test.ts b/packages/agent/src/tools/system/listShells.test.ts index 95b6647..10f13a1 100644 --- a/packages/agent/src/tools/system/listShells.test.ts +++ b/packages/agent/src/tools/system/listShells.test.ts @@ -81,8 +81,9 @@ describe('listShellsTool', () => { expect(result.shells.length).toBe(1); expect(result.count).toBe(1); - expect(result.shells[0].id).toBe('shell-1'); - expect(result.shells[0].status).toBe(ShellStatus.RUNNING); + expect(result.shells.length).toBe(1); + expect(result.shells[0]!.id).toBe('shell-1'); + expect(result.shells[0]!.status).toBe(ShellStatus.RUNNING); }); it('should include metadata when verbose is true', async () => { @@ -105,9 +106,10 @@ describe('listShellsTool', () => { ); expect(result.shells.length).toBe(1); - expect(result.shells[0].id).toBe('shell-3'); - expect(result.shells[0].status).toBe(ShellStatus.ERROR); - expect(result.shells[0].metadata).toBeDefined(); - expect(result.shells[0].metadata?.error).toBe('Command not found'); + expect(result.shells.length).toBe(1); + expect(result.shells[0]!.id).toBe('shell-3'); + expect(result.shells[0]!.status).toBe(ShellStatus.ERROR); + expect(result.shells[0]!.metadata).toBeDefined(); + expect(result.shells[0]!.metadata?.error).toBe('Command not found'); }); }); diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index 44f3ef8..4b17828 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -1,4 +1,4 @@ -import { BrowserManager, processStates } from 'mycoder-agent'; +import { BrowserManager, shellTracker } from 'mycoder-agent'; import { agentStates } from 'mycoder-agent/dist/tools/interaction/agentStart.js'; /** @@ -55,29 +55,10 @@ export async function cleanupResources(): Promise { // 2. Clean up shell processes try { - if (processStates.size > 0) { - console.log(`Terminating ${processStates.size} shell processes...`); - for (const [id, state] of processStates.entries()) { - if (!state.state.completed) { - console.log(`Terminating process ${id}...`); - try { - state.process.kill('SIGTERM'); - // Force kill after a short timeout if still running - setTimeout(() => { - try { - if (!state.state.completed) { - state.process.kill('SIGKILL'); - } - // eslint-disable-next-line unused-imports/no-unused-vars - } catch (e) { - // Ignore errors on forced kill - } - }, 500); - } catch (e) { - console.error(`Error terminating process ${id}:`, e); - } - } - } + 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);