feat: implement ShellTracker to decouple from backgroundTools

Author: bhouston

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<string, MockProcessState>(),
+    shellTracker: {
+      processStates: new Map<string, MockProcessState>(),
+      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 () => {

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);

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<string, any>; // 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<void> {
-    try {
-      const processState = processStates.get(tool.id);
-      if (processState && !processState.state.completed) {
-        processState.process.kill('SIGTERM');
-
-        // Force kill after a short timeout if still running
-        await new Promise<void>((resolve) => {
-          setTimeout(() => {
-            try {
-              if (!processState.state.completed) {
-                processState.process.kill('SIGKILL');
-              }
-            } catch {
-              // Ignore errors on forced kill
-            }
-            resolve();
-          }, 500);
-        });
-      }
-      this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
-    } catch (error) {
-      this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
-        error: error instanceof Error ? error.message : String(error),
-      });
-    }
-  }
-
   /**
    * Cleans up a sub-agent
    * @param tool The agent tool to clean up

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