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`); + }, +};