feat: add agent tracking to background tools

bhouston · bhouston · commit 4a3bcc72f27a · 2025-03-12T06:41:27.000-04:00
diff --git 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<Parameters, ReturnType> = {
       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
@@ -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<Parameters, ReturnType> = {
   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<Parameters, ReturnType> = {
       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<Parameters, ReturnType> = {
 
     const tools = getTools({ enableUserPrompt });
 
-    // Create an instance ID
-    const instanceId = uuidv4();
-
     // Store the agent state
     const agentState: AgentState = {
       goal,
@@ -147,13 +154,23 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
           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
         const state = agentStates.get(instanceId);
         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
@@ -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<Parameters, ReturnType> = {
   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<Parameters, ReturnType> = {
       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<Parameters, ReturnType> = {
       ...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}"`);
