From b5bb48981791acda74ee46b93d2d85e27e93a538 Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Tue, 11 Mar 2025 22:28:57 -0400 Subject: [PATCH 1/2] feat: implement background tool tracking (issue #112) --- .../agent/src/core/backgroundTools.test.ts | 185 +++++++++++++++++ packages/agent/src/core/backgroundTools.ts | 194 ++++++++++++++++++ packages/agent/src/core/types.ts | 1 + packages/agent/src/index.ts | 2 + .../agent/src/tools/browser/browseMessage.ts | 25 +++ .../agent/src/tools/browser/browseStart.ts | 28 ++- packages/agent/src/tools/getTools.ts | 4 +- .../tools/system/listBackgroundTools.test.ts | 52 +++++ .../src/tools/system/listBackgroundTools.ts | 118 +++++++++++ .../agent/src/tools/system/shellMessage.ts | 40 ++++ packages/agent/src/tools/system/shellStart.ts | 31 ++- 11 files changed, 676 insertions(+), 4 deletions(-) create mode 100644 packages/agent/src/core/backgroundTools.test.ts create mode 100644 packages/agent/src/core/backgroundTools.ts create mode 100644 packages/agent/src/tools/system/listBackgroundTools.test.ts create mode 100644 packages/agent/src/tools/system/listBackgroundTools.ts diff --git a/packages/agent/src/core/backgroundTools.test.ts b/packages/agent/src/core/backgroundTools.test.ts new file mode 100644 index 0000000..ec75544 --- /dev/null +++ b/packages/agent/src/core/backgroundTools.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it, vi, beforeEach } from 'vitest'; + +import { + backgroundToolRegistry, + BackgroundToolStatus, + BackgroundToolType, +} from './backgroundTools.js'; + +// Mock uuid to return predictable IDs for testing +vi.mock('uuid', () => ({ + v4: vi.fn().mockReturnValue('test-id-1'), // Always return the same ID for simplicity in tests +})); + +describe('BackgroundToolRegistry', () => { + beforeEach(() => { + // Clear all registered tools before each test + const registry = backgroundToolRegistry as any; + registry.tools = new Map(); + }); + + it('should register a shell process', () => { + const id = backgroundToolRegistry.registerShell('agent-1', 'ls -la'); + + expect(id).toBe('test-id-1'); + + const tool = backgroundToolRegistry.getToolById(id); + expect(tool).toBeDefined(); + if (tool) { + expect(tool.type).toBe(BackgroundToolType.SHELL); + expect(tool.status).toBe(BackgroundToolStatus.RUNNING); + expect(tool.agentId).toBe('agent-1'); + if (tool.type === BackgroundToolType.SHELL) { + expect(tool.metadata.command).toBe('ls -la'); + } + } + }); + + it('should register a browser process', () => { + const id = backgroundToolRegistry.registerBrowser( + 'agent-1', + 'https://example.com', + ); + + expect(id).toBe('test-id-1'); + + const tool = backgroundToolRegistry.getToolById(id); + expect(tool).toBeDefined(); + if (tool) { + expect(tool.type).toBe(BackgroundToolType.BROWSER); + expect(tool.status).toBe(BackgroundToolStatus.RUNNING); + expect(tool.agentId).toBe('agent-1'); + if (tool.type === BackgroundToolType.BROWSER) { + expect(tool.metadata.url).toBe('https://example.com'); + } + } + }); + + it('should update tool status', () => { + const id = backgroundToolRegistry.registerShell('agent-1', 'sleep 10'); + + const updated = backgroundToolRegistry.updateToolStatus( + id, + BackgroundToolStatus.COMPLETED, + { + exitCode: 0, + }, + ); + + expect(updated).toBe(true); + + const tool = backgroundToolRegistry.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 = backgroundToolRegistry.updateToolStatus( + 'non-existent-id', + BackgroundToolStatus.COMPLETED, + ); + + expect(updated).toBe(false); + }); + + it('should get tools by agent ID', () => { + // For this test, we'll directly manipulate the tools map + const registry = backgroundToolRegistry as any; + registry.tools = new Map(); + + // Add tools directly to the map with different agent IDs + registry.tools.set('id1', { + id: 'id1', + type: BackgroundToolType.SHELL, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId: 'agent-1', + metadata: { command: 'ls -la' }, + }); + + registry.tools.set('id2', { + id: 'id2', + type: BackgroundToolType.BROWSER, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId: 'agent-1', + metadata: { url: 'https://example.com' }, + }); + + registry.tools.set('id3', { + id: 'id3', + type: BackgroundToolType.SHELL, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId: 'agent-2', + metadata: { command: 'echo hello' }, + }); + + const agent1Tools = backgroundToolRegistry.getToolsByAgent('agent-1'); + const agent2Tools = backgroundToolRegistry.getToolsByAgent('agent-2'); + + expect(agent1Tools.length).toBe(2); + expect(agent2Tools.length).toBe(1); + }); + + it('should clean up old completed tools', () => { + // Create tools with specific dates + const registry = backgroundToolRegistry as any; + + // Add a completed tool from 25 hours ago + const oldTool = { + id: 'old-tool', + type: BackgroundToolType.SHELL, + 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' }, + }; + + // Add a completed tool from 10 hours ago + const recentTool = { + id: 'recent-tool', + type: BackgroundToolType.SHELL, + 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' }, + }; + + // Add a running tool from 25 hours ago + const oldRunningTool = { + id: 'old-running-tool', + type: BackgroundToolType.SHELL, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(Date.now() - 25 * 60 * 60 * 1000), + agentId: 'agent-1', + metadata: { command: 'sleep 100' }, + }; + + registry.tools.set('old-tool', oldTool); + registry.tools.set('recent-tool', recentTool); + registry.tools.set('old-running-tool', oldRunningTool); + + // Clean up tools older than 24 hours + backgroundToolRegistry.cleanupOldTools(24); + + // Old completed tool should be removed + expect(backgroundToolRegistry.getToolById('old-tool')).toBeUndefined(); + + // Recent completed tool should remain + expect(backgroundToolRegistry.getToolById('recent-tool')).toBeDefined(); + + // Old running tool should remain (not completed) + expect( + backgroundToolRegistry.getToolById('old-running-tool'), + ).toBeDefined(); + }); +}); diff --git a/packages/agent/src/core/backgroundTools.ts b/packages/agent/src/core/backgroundTools.ts new file mode 100644 index 0000000..ebe850b --- /dev/null +++ b/packages/agent/src/core/backgroundTools.ts @@ -0,0 +1,194 @@ +import { v4 as uuidv4 } from 'uuid'; + +// Types of background processes we can track +export enum BackgroundToolType { + SHELL = 'shell', + BROWSER = 'browser', + AGENT = 'agent', +} + +// Status of a background process +export enum BackgroundToolStatus { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + TERMINATED = 'terminated', +} + +// Common interface for all background processes +export interface BackgroundTool { + id: string; + type: BackgroundToolType; + status: BackgroundToolStatus; + startTime: Date; + endTime?: Date; + agentId: string; // To track which agent created this process + 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; + }; +} + +// Browser process specific data +export interface BrowserBackgroundTool extends BackgroundTool { + type: BackgroundToolType.BROWSER; + metadata: { + url?: string; + }; +} + +// Agent process specific data (for future use) +export interface AgentBackgroundTool extends BackgroundTool { + type: BackgroundToolType.AGENT; + metadata: { + goal?: string; + }; +} + +// Utility type for all background tool types +export type AnyBackgroundTool = + | ShellBackgroundTool + | BrowserBackgroundTool + | AgentBackgroundTool; + +/** + * Registry to keep track of all background processes + */ +export class BackgroundToolRegistry { + private static instance: BackgroundToolRegistry; + private tools: Map = new Map(); + + // Private constructor for singleton pattern + private constructor() {} + + // Get the singleton instance + public static getInstance(): BackgroundToolRegistry { + if (!BackgroundToolRegistry.instance) { + BackgroundToolRegistry.instance = new BackgroundToolRegistry(); + } + return BackgroundToolRegistry.instance; + } + + // Register a new shell process + public registerShell(agentId: string, command: string): string { + const id = uuidv4(); + const tool: ShellBackgroundTool = { + id, + type: BackgroundToolType.SHELL, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId, + metadata: { + command, + }, + }; + this.tools.set(id, tool); + return id; + } + + // Register a new browser process + public registerBrowser(agentId: string, url?: string): string { + const id = uuidv4(); + const tool: BrowserBackgroundTool = { + id, + type: BackgroundToolType.BROWSER, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId, + metadata: { + url, + }, + }; + this.tools.set(id, tool); + return id; + } + + // Register a new agent process (for future use) + public registerAgent(agentId: string, goal?: string): string { + const id = uuidv4(); + const tool: AgentBackgroundTool = { + id, + type: BackgroundToolType.AGENT, + status: BackgroundToolStatus.RUNNING, + startTime: new Date(), + agentId, + metadata: { + goal, + }, + }; + this.tools.set(id, tool); + return id; + } + + // Update the status of a process + public updateToolStatus( + id: string, + status: BackgroundToolStatus, + metadata?: Record, + ): boolean { + const tool = this.tools.get(id); + if (!tool) { + return false; + } + + tool.status = status; + + if ( + status === BackgroundToolStatus.COMPLETED || + status === BackgroundToolStatus.ERROR || + status === BackgroundToolStatus.TERMINATED + ) { + tool.endTime = new Date(); + } + + if (metadata) { + tool.metadata = { ...tool.metadata, ...metadata }; + } + + return true; + } + + // Get all processes for a specific agent + public getToolsByAgent(agentId: string): AnyBackgroundTool[] { + const result: AnyBackgroundTool[] = []; + for (const tool of this.tools.values()) { + if (tool.agentId === agentId) { + result.push(tool); + } + } + return result; + } + + // Get a specific process by ID + public getToolById(id: string): AnyBackgroundTool | undefined { + return this.tools.get(id); + } + + // Clean up completed processes (optional, for maintenance) + public cleanupOldTools(olderThanHours: number = 24): void { + const cutoffTime = new Date(Date.now() - olderThanHours * 60 * 60 * 1000); + + for (const [id, tool] of this.tools.entries()) { + // Remove if it's completed/error/terminated AND older than cutoff + if ( + tool.endTime && + tool.endTime < cutoffTime && + (tool.status === BackgroundToolStatus.COMPLETED || + tool.status === BackgroundToolStatus.ERROR || + tool.status === BackgroundToolStatus.TERMINATED) + ) { + this.tools.delete(id); + } + } + } +} + +// Export singleton instance +export const backgroundToolRegistry = BackgroundToolRegistry.getInstance(); diff --git a/packages/agent/src/core/types.ts b/packages/agent/src/core/types.ts index 738a79b..f2f13b9 100644 --- a/packages/agent/src/core/types.ts +++ b/packages/agent/src/core/types.ts @@ -20,6 +20,7 @@ export type ToolContext = { customPrompt?: string; tokenCache?: boolean; enableUserPrompt?: boolean; + agentId?: string; // Unique identifier for the agent, used for background tool tracking }; export type Tool, TReturn = any> = { diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index a816669..401de56 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -9,6 +9,7 @@ export * from './tools/system/respawn.js'; export * from './tools/system/sequenceComplete.js'; export * from './tools/system/shellMessage.js'; export * from './tools/system/shellExecute.js'; +export * from './tools/system/listBackgroundTools.js'; // Tools - Browser export * from './tools/browser/BrowserManager.js'; @@ -25,6 +26,7 @@ export * from './tools/interaction/userPrompt.js'; // Core export * from './core/executeToolCall.js'; export * from './core/types.js'; +export * from './core/backgroundTools.js'; // Tool Agent Core export { toolAgent } from './core/toolAgent/toolAgentCore.js'; export * from './core/toolAgent/config.js'; diff --git a/packages/agent/src/tools/browser/browseMessage.ts b/packages/agent/src/tools/browser/browseMessage.ts index 5643321..abe07c3 100644 --- a/packages/agent/src/tools/browser/browseMessage.ts +++ b/packages/agent/src/tools/browser/browseMessage.ts @@ -1,6 +1,10 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { + backgroundToolRegistry, + BackgroundToolStatus, +} from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; import { sleep } from '../../utils/sleep.js'; @@ -184,6 +188,16 @@ export const browseMessageTool: Tool = { await session.page.context().close(); await session.browser.close(); browserSessions.delete(instanceId); + + // Update background tool registry when browser is explicitly closed + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.COMPLETED, + { + closedExplicitly: true, + }, + ); + logger.verbose('Browser session closed successfully'); return { status: 'closed' }; } @@ -194,6 +208,17 @@ export const browseMessageTool: Tool = { } } catch (error) { logger.error('Browser action failed:', { error }); + + // Update background tool registry with error status if action fails + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.ERROR, + { + error: errorToString(error), + actionType, + }, + ); + return { status: 'error', error: errorToString(error), diff --git a/packages/agent/src/tools/browser/browseStart.ts b/packages/agent/src/tools/browser/browseStart.ts index 8ba5a65..8d95000 100644 --- a/packages/agent/src/tools/browser/browseStart.ts +++ b/packages/agent/src/tools/browser/browseStart.ts @@ -3,6 +3,10 @@ import { v4 as uuidv4 } from 'uuid'; import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { + backgroundToolRegistry, + BackgroundToolStatus, +} from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; import { sleep } from '../../utils/sleep.js'; @@ -42,7 +46,7 @@ export const browseStartTool: Tool = { execute: async ( { url, timeout = 30000 }, - { logger, headless, userSession, pageFilter }, + { logger, headless, userSession, pageFilter, agentId }, ): Promise => { logger.verbose(`Starting browser session${url ? ` at ${url}` : ''}`); logger.verbose( @@ -53,6 +57,9 @@ export const browseStartTool: Tool = { try { const instanceId = uuidv4(); + // Register this browser session with the background tool registry + backgroundToolRegistry.registerBrowser(agentId || 'unknown', url); + // Launch browser const launchOptions = { headless, @@ -91,6 +98,11 @@ export const browseStartTool: Tool = { // Setup cleanup handlers browser.on('disconnected', () => { browserSessions.delete(instanceId); + // Update background tool registry when browser disconnects + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.TERMINATED, + ); }); // Navigate to URL if provided @@ -133,6 +145,16 @@ export const browseStartTool: Tool = { logger.verbose('Browser session started successfully'); logger.verbose(`Content length: ${content.length} characters`); + // Update background tool registry with running status + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.RUNNING, + { + url: url || 'about:blank', + contentLength: content.length, + }, + ); + return { instanceId, status: 'initialized', @@ -140,6 +162,10 @@ export const browseStartTool: Tool = { }; } catch (error) { logger.error(`Failed to start browser: ${errorToString(error)}`); + + // No need to update background tool registry here as we don't have a valid instanceId + // when an error occurs before the browser is properly initialized + return { instanceId: '', status: 'error', diff --git a/packages/agent/src/tools/getTools.ts b/packages/agent/src/tools/getTools.ts index 39033d8..0314056 100644 --- a/packages/agent/src/tools/getTools.ts +++ b/packages/agent/src/tools/getTools.ts @@ -5,10 +5,10 @@ import { browseMessageTool } from './browser/browseMessage.js'; import { browseStartTool } from './browser/browseStart.js'; import { agentMessageTool } from './interaction/agentMessage.js'; import { agentStartTool } from './interaction/agentStart.js'; -import { subAgentTool } from './interaction/subAgent.js'; import { userPromptTool } from './interaction/userPrompt.js'; import { fetchTool } from './io/fetch.js'; import { textEditorTool } from './io/textEditor.js'; +import { listBackgroundToolsTool } from './system/listBackgroundTools.js'; import { respawnTool } from './system/respawn.js'; import { sequenceCompleteTool } from './system/sequenceComplete.js'; import { shellMessageTool } from './system/shellMessage.js'; @@ -27,7 +27,6 @@ export function getTools(options?: GetToolsOptions): Tool[] { // Force cast to Tool type to avoid TypeScript issues const tools: Tool[] = [ textEditorTool as unknown as Tool, - subAgentTool as unknown as Tool, agentStartTool as unknown as Tool, agentMessageTool as unknown as Tool, sequenceCompleteTool as unknown as Tool, @@ -38,6 +37,7 @@ export function getTools(options?: GetToolsOptions): Tool[] { browseMessageTool as unknown as Tool, respawnTool as unknown as Tool, sleepTool as unknown as Tool, + listBackgroundToolsTool as unknown as Tool, ]; // Only include userPrompt tool if enabled diff --git a/packages/agent/src/tools/system/listBackgroundTools.test.ts b/packages/agent/src/tools/system/listBackgroundTools.test.ts new file mode 100644 index 0000000..5f9fddf --- /dev/null +++ b/packages/agent/src/tools/system/listBackgroundTools.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { listBackgroundToolsTool } from './listBackgroundTools.js'; + +// Mock the entire background tools module +vi.mock('../../core/backgroundTools.js', () => { + return { + backgroundToolRegistry: { + getToolsByAgent: vi.fn().mockReturnValue([ + { + id: 'shell-1', + type: 'shell', + status: 'running', + startTime: new Date(Date.now() - 10000), + agentId: 'agent-1', + metadata: { command: 'ls -la' }, + }, + ]), + }, + BackgroundToolStatus: { + RUNNING: 'running', + COMPLETED: 'completed', + ERROR: 'error', + TERMINATED: 'terminated', + }, + BackgroundToolType: { + SHELL: 'shell', + BROWSER: 'browser', + AGENT: 'agent', + }, + }; +}); + +describe('listBackgroundTools tool', () => { + const mockLogger = { + debug: vi.fn(), + verbose: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + + it('should list background tools', async () => { + const result = await listBackgroundToolsTool.execute({}, { + logger: mockLogger as any, + agentId: 'agent-1', + } as any); + + expect(result.count).toEqual(1); + expect(result.tools).toHaveLength(1); + }); +}); diff --git a/packages/agent/src/tools/system/listBackgroundTools.ts b/packages/agent/src/tools/system/listBackgroundTools.ts new file mode 100644 index 0000000..83eff8f --- /dev/null +++ b/packages/agent/src/tools/system/listBackgroundTools.ts @@ -0,0 +1,118 @@ +import { z } from 'zod'; +import { zodToJsonSchema } from 'zod-to-json-schema'; + +import { + backgroundToolRegistry, + BackgroundToolStatus, +} from '../../core/backgroundTools.js'; +import { Tool } from '../../core/types.js'; + +const parameterSchema = z.object({ + status: z + .enum(['all', 'running', 'completed', 'error', 'terminated']) + .optional() + .describe('Filter tools by status (default: "all")'), + type: z + .enum(['all', 'shell', 'browser', 'agent']) + .optional() + .describe('Filter tools by type (default: "all")'), + verbose: z + .boolean() + .optional() + .describe('Include detailed metadata about each tool (default: false)'), +}); + +const returnSchema = z.object({ + tools: z.array( + z.object({ + id: z.string(), + type: z.string(), + status: z.string(), + startTime: z.string(), + endTime: z.string().optional(), + runtime: z.number().describe('Runtime in seconds'), + metadata: z.record(z.any()).optional(), + }), + ), + count: z.number(), +}); + +type Parameters = z.infer; +type ReturnType = z.infer; + +export const listBackgroundToolsTool: Tool = { + name: 'listBackgroundTools', + description: + 'Lists all background tools (shells, browsers, agents) and their status', + logPrefix: '🔍', + parameters: parameterSchema, + returns: returnSchema, + parametersJsonSchema: zodToJsonSchema(parameterSchema), + returnsJsonSchema: zodToJsonSchema(returnSchema), + + execute: async ( + { status = 'all', type = 'all', verbose = false }, + { logger, agentId }, + ): Promise => { + logger.verbose( + `Listing background tools with status: ${status}, type: ${type}, verbose: ${verbose}`, + ); + + // Get all tools for this agent + const tools = backgroundToolRegistry.getToolsByAgent(agentId || 'unknown'); + + // Filter by status if specified + const filteredByStatus = + status === 'all' + ? tools + : tools.filter((tool) => { + const statusEnum = + status.toUpperCase() as keyof typeof BackgroundToolStatus; + return tool.status === BackgroundToolStatus[statusEnum]; + }); + + // Filter by type if specified + const filteredTools = + type === 'all' + ? filteredByStatus + : filteredByStatus.filter( + (tool) => tool.type.toLowerCase() === type.toLowerCase(), + ); + + // Format the response + const formattedTools = filteredTools.map((tool) => { + const now = new Date(); + const startTime = tool.startTime; + const endTime = tool.endTime || now; + const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds + + return { + id: tool.id, + type: tool.type, + status: tool.status, + startTime: startTime.toISOString(), + ...(tool.endTime && { endTime: tool.endTime.toISOString() }), + runtime: parseFloat(runtime.toFixed(2)), + ...(verbose && { metadata: tool.metadata }), + }; + }); + + return { + tools: formattedTools, + count: formattedTools.length, + }; + }, + + logParameters: ( + { status = 'all', type = 'all', verbose = false }, + { logger }, + ) => { + logger.info( + `Listing ${type} background tools with status: ${status}, verbose: ${verbose}`, + ); + }, + + logReturns: (output, { logger }) => { + logger.info(`Found ${output.count} background tools`); + }, +}; diff --git a/packages/agent/src/tools/system/shellMessage.ts b/packages/agent/src/tools/system/shellMessage.ts index 49f3f7c..487ac8c 100644 --- a/packages/agent/src/tools/system/shellMessage.ts +++ b/packages/agent/src/tools/system/shellMessage.ts @@ -1,6 +1,10 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { + backgroundToolRegistry, + BackgroundToolStatus, +} from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { sleep } from '../../utils/sleep.js'; @@ -111,6 +115,16 @@ export const shellMessageTool: Tool = { if (signal) { const wasKilled = processState.process.kill(signal); if (!wasKilled) { + // Update background tool registry if signal failed + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.ERROR, + { + error: `Failed to send signal ${signal}`, + signalAttempted: signal, + }, + ); + return { stdout: '', stderr: '', @@ -119,7 +133,33 @@ export const shellMessageTool: Tool = { error: `Failed to send signal ${signal} to process (process may have already terminated)`, }; } + processState.state.signaled = true; + + // Update background tool registry with signal information + if ( + signal === 'SIGTERM' || + signal === 'SIGKILL' || + signal === 'SIGINT' + ) { + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.TERMINATED, + { + signal, + terminatedByUser: true, + }, + ); + } else { + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.RUNNING, + { + signal, + signaled: true, + }, + ); + } } // Send input if provided diff --git a/packages/agent/src/tools/system/shellStart.ts b/packages/agent/src/tools/system/shellStart.ts index 170aa3f..86fbbba 100644 --- a/packages/agent/src/tools/system/shellStart.ts +++ b/packages/agent/src/tools/system/shellStart.ts @@ -4,6 +4,10 @@ import { v4 as uuidv4 } from 'uuid'; import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { + backgroundToolRegistry, + BackgroundToolStatus, +} from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { errorToString } from '../../utils/errorToString.js'; @@ -98,7 +102,7 @@ export const shellStartTool: Tool = { showStdIn = false, showStdout = false, }, - { logger, workingDirectory }, + { logger, workingDirectory, agentId }, ): Promise => { if (showStdIn) { logger.info(`Command input: ${command}`); @@ -107,7 +111,12 @@ export const shellStartTool: Tool = { return new Promise((resolve) => { try { + // Generate a unique ID for this process const instanceId = uuidv4(); + + // Register this shell process with the background tool registry + backgroundToolRegistry.registerShell(agentId || 'unknown', command); + let hasResolved = false; // Split command into command and args @@ -154,6 +163,16 @@ export const shellStartTool: Tool = { process.on('error', (error) => { logger.error(`[${instanceId}] Process error: ${error.message}`); processState.state.completed = true; + + // Update background tool registry with error status + backgroundToolRegistry.updateToolStatus( + instanceId, + BackgroundToolStatus.ERROR, + { + error: error.message, + }, + ); + if (!hasResolved) { hasResolved = true; resolve({ @@ -175,6 +194,16 @@ 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; + backgroundToolRegistry.updateToolStatus(instanceId, status, { + exitCode: code, + signaled: signal !== null, + }); + if (!hasResolved) { hasResolved = true; // If we haven't resolved yet, this happened within the timeout From 4a3bcc72f27af5fdbeeb407a748d5ecf3b7faed5 Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 12 Mar 2025 06:41:27 -0400 Subject: [PATCH 2/2] feat: add agent tracking to background tools --- .../src/tools/interaction/agentMessage.ts | 6 ++++ .../agent/src/tools/interaction/agentStart.ts | 23 +++++++++++-- .../agent/src/tools/interaction/subAgent.ts | 32 ++++++++++++++++--- 3 files changed, 53 insertions(+), 8 deletions(-) diff --git a/packages/agent/src/tools/interaction/agentMessage.ts b/packages/agent/src/tools/interaction/agentMessage.ts index 2fc8abc..89f4dce 100644 --- a/packages/agent/src/tools/interaction/agentMessage.ts +++ b/packages/agent/src/tools/interaction/agentMessage.ts @@ -1,6 +1,7 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { backgroundToolRegistry, BackgroundToolStatus } from '../../core/backgroundTools.js'; import { Tool } from '../../core/types.js'; import { agentStates } from './agentStart.js'; @@ -75,6 +76,11 @@ export const agentMessageTool: Tool = { if (terminate) { agentState.aborted = true; agentState.completed = true; + + // Update background tool registry with terminated status + backgroundToolRegistry.updateToolStatus(instanceId, BackgroundToolStatus.TERMINATED, { + terminatedByUser: true + }); return { output: agentState.output || 'Sub-agent terminated before completion', diff --git a/packages/agent/src/tools/interaction/agentStart.ts b/packages/agent/src/tools/interaction/agentStart.ts index 12deff4..3e80509 100644 --- a/packages/agent/src/tools/interaction/agentStart.ts +++ b/packages/agent/src/tools/interaction/agentStart.ts @@ -2,6 +2,7 @@ import { v4 as uuidv4 } from 'uuid'; import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { backgroundToolRegistry, BackgroundToolStatus, BackgroundToolType } from '../../core/backgroundTools.js'; import { getDefaultSystemPrompt, getModel, @@ -90,6 +91,8 @@ export const agentStartTool: Tool = { returns: returnSchema, returnsJsonSchema: zodToJsonSchema(returnSchema), execute: async (params, context) => { + const { logger, agentId } = context; + // Validate parameters const { description, @@ -99,6 +102,13 @@ export const agentStartTool: Tool = { relevantFilesDirectories, enableUserPrompt = false, } = parameterSchema.parse(params); + + // Create an instance ID + const instanceId = uuidv4(); + + // Register this agent with the background tool registry + backgroundToolRegistry.registerAgent(agentId || 'unknown', goal); + logger.verbose(`Registered agent with ID: ${instanceId}`); // Construct a well-structured prompt const prompt = [ @@ -115,9 +125,6 @@ export const agentStartTool: Tool = { const tools = getTools({ enableUserPrompt }); - // Create an instance ID - const instanceId = uuidv4(); - // Store the agent state const agentState: AgentState = { goal, @@ -147,6 +154,11 @@ export const agentStartTool: Tool = { state.completed = true; state.result = result; state.output = result.result; + + // Update background tool registry with completed status + backgroundToolRegistry.updateToolStatus(instanceId, BackgroundToolStatus.COMPLETED, { + result: result.result.substring(0, 100) + (result.result.length > 100 ? '...' : '') + }); } } catch (error) { // Update agent state with the error @@ -154,6 +166,11 @@ export const agentStartTool: Tool = { if (state && !state.aborted) { state.completed = true; state.error = error instanceof Error ? error.message : String(error); + + // Update background tool registry with error status + backgroundToolRegistry.updateToolStatus(instanceId, BackgroundToolStatus.ERROR, { + error: error instanceof Error ? error.message : String(error) + }); } } return true; diff --git a/packages/agent/src/tools/interaction/subAgent.ts b/packages/agent/src/tools/interaction/subAgent.ts index fc2da6e..c36b933 100644 --- a/packages/agent/src/tools/interaction/subAgent.ts +++ b/packages/agent/src/tools/interaction/subAgent.ts @@ -1,6 +1,7 @@ import { z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; +import { backgroundToolRegistry, BackgroundToolStatus, BackgroundToolType } from '../../core/backgroundTools.js'; import { getDefaultSystemPrompt, getModel, @@ -68,6 +69,8 @@ export const subAgentTool: Tool = { returns: returnSchema, returnsJsonSchema: zodToJsonSchema(returnSchema), execute: async (params, context) => { + const { logger, agentId } = context; + // Validate parameters const { description, @@ -76,6 +79,10 @@ export const subAgentTool: Tool = { workingDirectory, relevantFilesDirectories, } = parameterSchema.parse(params); + + // Register this sub-agent with the background tool registry + const subAgentId = backgroundToolRegistry.registerAgent(agentId || 'unknown', goal); + logger.verbose(`Registered sub-agent with ID: ${subAgentId}`); // Construct a well-structured prompt const prompt = [ @@ -97,11 +104,26 @@ export const subAgentTool: Tool = { ...subAgentConfig, }; - const result = await toolAgent(prompt, tools, config, { - ...context, - workingDirectory: workingDirectory ?? context.workingDirectory, - }); - return { response: result.result }; + try { + const result = await toolAgent(prompt, tools, config, { + ...context, + workingDirectory: workingDirectory ?? context.workingDirectory, + }); + + // Update background tool registry with completed status + backgroundToolRegistry.updateToolStatus(subAgentId, BackgroundToolStatus.COMPLETED, { + result: result.result.substring(0, 100) + (result.result.length > 100 ? '...' : '') + }); + + return { response: result.result }; + } catch (error) { + // Update background tool registry with error status + backgroundToolRegistry.updateToolStatus(subAgentId, BackgroundToolStatus.ERROR, { + error: error instanceof Error ? error.message : String(error) + }); + + throw error; + } }, logParameters: (input, { logger }) => { logger.info(`Delegating task "${input.description}"`);