From eda4640c51bb713e260b42847323861808accbe4 Mon Sep 17 00:00:00 2001 From: "Ben Houston (via MyCoder)" Date: Fri, 14 Mar 2025 19:50:57 +0000 Subject: [PATCH] test: fix backgroundTools.cleanup test for AgentTracker integration 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. --- .../src/core/backgroundTools.cleanup.test.ts | 117 +++++++------- packages/agent/src/core/backgroundTools.ts | 13 +- packages/agent/src/tools/getTools.ts | 2 + .../agent/src/tools/interaction/agentStart.ts | 52 +++--- .../src/tools/interaction/agentTracker.ts | 148 ++++++++++++++++++ packages/agent/src/tools/system/listAgents.ts | 115 ++++++++++++++ 6 files changed, 353 insertions(+), 94 deletions(-) create mode 100644 packages/agent/src/tools/interaction/agentTracker.ts create mode 100644 packages/agent/src/tools/system/listAgents.ts diff --git a/packages/agent/src/core/backgroundTools.cleanup.test.ts b/packages/agent/src/core/backgroundTools.cleanup.test.ts index 3adec5d..82e2118 100644 --- 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 }; - state: { completed: boolean }; - command?: string; - stdout?: string[]; - stderr?: string[]; - showStdIn?: boolean; - showStdout?: boolean; -}; - -type MockAgentState = { - aborted: boolean; - completed: boolean; - context: { - backgroundTools: { - cleanup: ReturnType; - }; + process: ChildProcess & { kill: ReturnType }; + 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(), + 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 }; 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 index 45c61c1..098c722 100644 --- 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 { 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 index 79ee272..a2a760f 100644 --- 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 index da25239..f1bd4f5 100644 --- 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 = new Map(); +// For backward compatibility +export const agentStates = new Map(); const parameterSchema = z.object({ description: z @@ -100,11 +86,12 @@ export const agentStartTool: Tool = { 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 = { // Store the agent state const agentState: AgentState = { + id: instanceId, goal, prompt, output: '', @@ -134,6 +122,10 @@ export const agentStartTool: Tool = { 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 = { }); // 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 = { } } 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 new file mode 100644 index 0000000..bb6463d --- /dev/null +++ b/packages/agent/src/tools/interaction/agentTracker.ts @@ -0,0 +1,148 @@ +import { v4 as uuidv4 } from 'uuid'; + +import { ToolAgentResult } from '../../core/toolAgent/types.js'; +import { ToolContext } from '../../core/types.js'; + +export enum AgentStatus { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + TERMINATED = 'terminated', +} + +export interface Agent { + id: string; + status: AgentStatus; + startTime: Date; + endTime?: Date; + goal: string; + result?: string; + error?: string; +} + +// Internal agent state tracking (similar to existing agentStates) +export interface AgentState { + id: string; + goal: string; + prompt: string; + output: string; + completed: boolean; + error?: string; + result?: ToolAgentResult; + context: ToolContext; + workingDirectory: string; + tools: unknown[]; + aborted: boolean; +} + +export class AgentTracker { + private agents: Map = new Map(); + private agentStates: Map = new Map(); + + constructor(readonly ownerName: string) {} + + // Register a new agent + public registerAgent(goal: string): string { + const id = uuidv4(); + + // Create agent tracking entry + const agent: Agent = { + id, + status: AgentStatus.RUNNING, + startTime: new Date(), + goal, + }; + + this.agents.set(id, agent); + return id; + } + + // Register agent state + public registerAgentState(id: string, state: AgentState): void { + this.agentStates.set(id, state); + } + + // Update agent status + public updateAgentStatus( + id: string, + status: AgentStatus, + metadata?: { result?: string; error?: string }, + ): boolean { + const agent = this.agents.get(id); + if (!agent) { + return false; + } + + agent.status = status; + + if ( + status === AgentStatus.COMPLETED || + status === AgentStatus.ERROR || + status === AgentStatus.TERMINATED + ) { + agent.endTime = new Date(); + } + + if (metadata) { + if (metadata.result !== undefined) agent.result = metadata.result; + if (metadata.error !== undefined) agent.error = metadata.error; + } + + return true; + } + + // Get a specific agent state + public getAgentState(id: string): AgentState | undefined { + return this.agentStates.get(id); + } + + // Get a specific agent tracking info + public getAgent(id: string): Agent | undefined { + return this.agents.get(id); + } + + // Get all agents with optional filtering + public getAgents(status?: AgentStatus): Agent[] { + if (!status) { + return Array.from(this.agents.values()); + } + + return Array.from(this.agents.values()).filter( + (agent) => agent.status === status, + ); + } + + // Cleanup and terminate agents + public async cleanup(): Promise { + const runningAgents = this.getAgents(AgentStatus.RUNNING); + + await Promise.all( + runningAgents.map((agent) => this.terminateAgent(agent.id)), + ); + } + + // Terminate a specific agent + public async terminateAgent(id: string): Promise { + try { + const agentState = this.agentStates.get(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 + if (agentState.context.backgroundTools) { + await agentState.context.backgroundTools.cleanup(); + } + } + this.updateAgentStatus(id, AgentStatus.TERMINATED); + } catch (error) { + this.updateAgentStatus(id, AgentStatus.ERROR, { + error: error instanceof Error ? error.message : String(error), + }); + } + } +} + +// Create a singleton instance +export const agentTracker = new AgentTracker('global'); diff --git a/packages/agent/src/tools/system/listAgents.ts b/packages/agent/src/tools/system/listAgents.ts new file mode 100644 index 0000000..e60e1bd --- /dev/null +++ b/packages/agent/src/tools/system/listAgents.ts @@ -0,0 +1,115 @@ +import { z } from 'zod'; +import { zodToJsonSchema } from 'zod-to-json-schema'; + +import { Tool } from '../../core/types.js'; +import { AgentStatus, agentTracker } from '../interaction/agentTracker.js'; + +const parameterSchema = z.object({ + status: z + .enum(['all', 'running', 'completed', 'error', 'terminated']) + .optional() + .describe('Filter agents by status (default: "all")'), + verbose: z + .boolean() + .optional() + .describe('Include detailed information about each agent (default: false)'), +}); + +const returnSchema = z.object({ + agents: z.array( + z.object({ + id: z.string(), + status: z.string(), + goal: z.string(), + startTime: z.string(), + endTime: z.string().optional(), + runtime: z.number().describe('Runtime in seconds'), + result: z.string().optional(), + error: z.string().optional(), + }), + ), + count: z.number(), +}); + +type Parameters = z.infer; +type ReturnType = z.infer; + +export const listAgentsTool: Tool = { + name: 'listAgents', + description: 'Lists all sub-agents 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 agents with status: ${status}, verbose: ${verbose}`, + ); + + // Get all agents + let agents = agentTracker.getAgents(); + + // Filter by status if specified + if (status !== 'all') { + const statusEnum = status.toUpperCase() as keyof typeof AgentStatus; + agents = agents.filter( + (agent) => agent.status === AgentStatus[statusEnum], + ); + } + + // Format the response + const formattedAgents = agents.map((agent) => { + const now = new Date(); + const startTime = agent.startTime; + const endTime = agent.endTime || now; + const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds + + const result: { + id: string; + status: string; + goal: string; + startTime: string; + endTime?: string; + runtime: number; + result?: string; + error?: string; + } = { + id: agent.id, + status: agent.status, + goal: agent.goal, + startTime: startTime.toISOString(), + ...(agent.endTime && { endTime: agent.endTime.toISOString() }), + runtime: parseFloat(runtime.toFixed(2)), + }; + + // Add result/error if verbose or if they exist + if (verbose || agent.result) { + result.result = agent.result; + } + + if (verbose && agent.error) { + result.error = agent.error; + } + + return result; + }); + + return { + agents: formattedAgents, + count: formattedAgents.length, + }; + }, + + logParameters: ({ status = 'all', verbose = false }, { logger }) => { + logger.info(`Listing agents with status: ${status}, verbose: ${verbose}`); + }, + + logReturns: (output, { logger }) => { + logger.info(`Found ${output.count} agents`); + }, +};