test: fix backgroundTools.cleanup test for AgentTracker integration

bhouston · bhouston · commit eda4640c51bb · 2025-03-14T19:50:57.000Z
Update the test to mock the new agentTracker module instead of agentStates.
This ensures the test correctly verifies that the backgroundTools.cleanupSubAgent
method properly delegates to agentTracker.terminateAgent.
diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/packages/agent/src/core/backgroundTools.cleanup.test.ts
@@ -2,36 +2,27 @@ 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 { processStates } from '../tools/system/shellStart.js';
 
 import { BackgroundTools, BackgroundToolStatus } from './backgroundTools';
-import { Tool } from './types';
+
+// 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: { kill: ReturnType<typeof vi.fn> };
-  state: { completed: boolean };
-  command?: string;
-  stdout?: string[];
-  stderr?: string[];
-  showStdIn?: boolean;
-  showStdout?: boolean;
-};
-
-type MockAgentState = {
-  aborted: boolean;
-  completed: boolean;
-  context: {
-    backgroundTools: {
-      cleanup: ReturnType<typeof vi.fn>;
-    };
+  process: ChildProcess & { kill: ReturnType<typeof vi.fn> };
+  state: {
+    completed: boolean;
+    signaled: boolean;
+    exitCode: number | null;
   };
-  goal?: string;
-  prompt?: string;
-  output?: string;
-  workingDirectory?: string;
-  tools?: Tool[];
+  command: string;
+  stdout: string[];
+  stderr: string[];
+  showStdIn: boolean;
+  showStdout: boolean;
 };
 
 // Mock dependencies
@@ -49,9 +40,28 @@ vi.mock('../tools/system/shellStart.js', () => {
   };
 });
 
-vi.mock('../tools/interaction/agentStart.js', () => {
+vi.mock('../tools/interaction/agentTracker.js', () => {
   return {
-    agentStates: new Map<string, MockAgentState>(),
+    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: [],
+        };
+      }),
+    },
   };
 });
 
@@ -75,11 +85,19 @@ describe('BackgroundTools cleanup', () => {
     // Setup mock process states
     const mockProcess = {
       kill: vi.fn(),
-    };
+      stdin: null,
+      stdout: null,
+      stderr: null,
+      stdio: null,
+    } as unknown as ChildProcess & { kill: ReturnType<typeof vi.fn> };
 
     const mockProcessState: MockProcessState = {
       process: mockProcess,
-      state: { completed: false },
+      state: {
+        completed: false,
+        signaled: false,
+        exitCode: null,
+      },
       command: 'test command',
       stdout: [],
       stderr: [],
@@ -88,26 +106,13 @@ describe('BackgroundTools cleanup', () => {
     };
 
     processStates.clear();
-    processStates.set('shell-1', mockProcessState as any);
-
-    // Setup mock agent states
-    const mockAgentState: MockAgentState = {
-      aborted: false,
-      completed: false,
-      context: {
-        backgroundTools: {
-          cleanup: vi.fn().mockResolvedValue(undefined),
-        },
-      },
-      goal: 'test goal',
-      prompt: 'test prompt',
-      output: '',
-      workingDirectory: '/test',
-      tools: [],
-    };
+    processStates.set(
+      'shell-1',
+      mockProcessState as unknown as MockProcessState,
+    );
 
-    agentStates.clear();
-    agentStates.set('agent-1', mockAgentState as any);
+    // Reset the agentTracker mock
+    vi.mocked(agentTracker.terminateAgent).mockClear();
   });
 
   afterEach(() => {
@@ -120,7 +125,6 @@ describe('BackgroundTools cleanup', () => {
 
     // Clear mock states
     processStates.clear();
-    agentStates.clear();
   });
 
   it('should clean up browser sessions', async () => {
@@ -149,7 +153,10 @@ describe('BackgroundTools cleanup', () => {
     const mockProcessState = processStates.get('shell-1');
 
     // Set the shell ID to match
-    processStates.set(shellId, processStates.get('shell-1') as any);
+    processStates.set(
+      shellId,
+      processStates.get('shell-1') as unknown as MockProcessState,
+    );
 
     // Run cleanup
     await backgroundTools.cleanup();
@@ -166,21 +173,11 @@ describe('BackgroundTools cleanup', () => {
     // Register an agent tool
     const agentId = backgroundTools.registerAgent('Test goal');
 
-    // Get mock agent state
-    const mockAgentState = agentStates.get('agent-1');
-
-    // Set the agent ID to match
-    agentStates.set(agentId, agentStates.get('agent-1') as any);
-
     // Run cleanup
     await backgroundTools.cleanup();
 
-    // Check that agent was marked as aborted
-    expect(mockAgentState?.aborted).toBe(true);
-    expect(mockAgentState?.completed).toBe(true);
-
-    // Check that cleanup was called on the agent's background tools
-    expect(mockAgentState?.context.backgroundTools.cleanup).toHaveBeenCalled();
+    // Check that terminateAgent was called with the agent ID
+    expect(agentTracker.terminateAgent).toHaveBeenCalledWith(agentId);
 
     // Check that tool status was updated
     const tool = backgroundTools.getToolById(agentId);
diff --git a/packages/agent/src/core/backgroundTools.ts b/packages/agent/src/core/backgroundTools.ts
@@ -2,7 +2,7 @@ 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 { agentTracker } from '../tools/interaction/agentTracker.js';
 import { processStates } from '../tools/system/shellStart.js';
 
 // Types of background processes we can track
@@ -268,15 +268,8 @@ export class BackgroundTools {
    */
   private async cleanupSubAgent(tool: AgentBackgroundTool): Promise<void> {
     try {
-      const agentState = agentStates.get(tool.id);
-      if (agentState && !agentState.aborted) {
-        // Set the agent as aborted and completed
-        agentState.aborted = true;
-        agentState.completed = true;
-
-        // Clean up resources owned by this sub-agent
-        await agentState.context.backgroundTools.cleanup();
-      }
+      // Delegate to the agent tracker
+      await agentTracker.terminateAgent(tool.id);
       this.updateToolStatus(tool.id, BackgroundToolStatus.TERMINATED);
     } catch (error) {
       this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
diff --git a/packages/agent/src/tools/getTools.ts b/packages/agent/src/tools/getTools.ts
@@ -9,6 +9,7 @@ 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 { sequenceCompleteTool } from './system/sequenceComplete.js';
 import { shellMessageTool } from './system/shellMessage.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,
+    listAgentsTool as unknown as Tool,
   ];
 
   // Only include userPrompt tool if enabled
diff --git a/packages/agent/src/tools/interaction/agentStart.ts b/packages/agent/src/tools/interaction/agentStart.ts
@@ -1,4 +1,3 @@
-import { v4 as uuidv4 } from 'uuid';
 import { z } from 'zod';
 import { zodToJsonSchema } from 'zod-to-json-schema';
 
@@ -8,26 +7,13 @@ import {
   AgentConfig,
 } from '../../core/toolAgent/config.js';
 import { toolAgent } from '../../core/toolAgent/toolAgentCore.js';
-import { ToolAgentResult } from '../../core/toolAgent/types.js';
 import { Tool, ToolContext } from '../../core/types.js';
 import { getTools } from '../getTools.js';
 
-// Define AgentState type
-type AgentState = {
-  goal: string;
-  prompt: string;
-  output: string;
-  completed: boolean;
-  error?: string;
-  result?: ToolAgentResult;
-  context: ToolContext;
-  workingDirectory: string;
-  tools: Tool[];
-  aborted: boolean;
-};
+import { AgentStatus, agentTracker, AgentState } from './agentTracker.js';
 
-// Global map to store agent state
-export const agentStates: Map<string, AgentState> = new Map();
+// For backward compatibility
+export const agentStates = new Map<string, AgentState>();
 
 const parameterSchema = z.object({
   description: z
@@ -100,11 +86,12 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
       userPrompt = false,
     } = parameterSchema.parse(params);
 
-    // Create an instance ID
-    const instanceId = uuidv4();
+    // Register this agent with the agent tracker
+    const instanceId = agentTracker.registerAgent(goal);
 
-    // Register this agent with the background tool registry
+    // For backward compatibility, also register with background tools
     backgroundTools.registerAgent(goal);
+
     logger.verbose(`Registered agent with ID: ${instanceId}`);
 
     // Construct a well-structured prompt
@@ -124,6 +111,7 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
 
     // Store the agent state
     const agentState: AgentState = {
+      id: instanceId,
       goal,
       prompt,
       output: '',
@@ -134,6 +122,10 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
       aborted: false,
     };
 
+    // Register agent state with the tracker
+    agentTracker.registerAgentState(instanceId, agentState);
+
+    // For backward compatibility
     agentStates.set(instanceId, agentState);
 
     // Start the agent in a separate promise that we don't await
@@ -146,13 +138,20 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
         });
 
         // Update agent state with the result
-        const state = agentStates.get(instanceId);
+        const state = agentTracker.getAgentState(instanceId);
         if (state && !state.aborted) {
           state.completed = true;
           state.result = result;
           state.output = result.result;
 
-          // Update background tool registry with completed status
+          // Update agent tracker with completed status
+          agentTracker.updateAgentStatus(instanceId, AgentStatus.COMPLETED, {
+            result:
+              result.result.substring(0, 100) +
+              (result.result.length > 100 ? '...' : ''),
+          });
+
+          // For backward compatibility
           backgroundTools.updateToolStatus(
             instanceId,
             BackgroundToolStatus.COMPLETED,
@@ -168,12 +167,17 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
         }
       } catch (error) {
         // Update agent state with the error
-        const state = agentStates.get(instanceId);
+        const state = agentTracker.getAgentState(instanceId);
         if (state && !state.aborted) {
           state.completed = true;
           state.error = error instanceof Error ? error.message : String(error);
 
-          // Update background tool registry with error status
+          // Update agent tracker with error status
+          agentTracker.updateAgentStatus(instanceId, AgentStatus.ERROR, {
+            error: error instanceof Error ? error.message : String(error),
+          });
+
+          // For backward compatibility
           backgroundTools.updateToolStatus(
             instanceId,
             BackgroundToolStatus.ERROR,
diff --git a/packages/agent/src/tools/interaction/agentTracker.ts b/packages/agent/src/tools/interaction/agentTracker.ts
diff --git a/packages/agent/src/tools/system/listAgents.ts b/packages/agent/src/tools/system/listAgents.ts
