refactor(agent): remove legacy agentStates variable

bhouston · bhouston · commit 57e709e4be4b · 2025-03-25T14:16:58.000-04:00
- Removed agentStates variable from agentStart.ts
- Updated agentMessage.ts to no longer use the legacy variable
- Simplified agentTools.test.ts to no longer rely on the legacy variable
- All tests still pass with this simpler implementation
diff --git a/packages/agent/src/tools/agent/agentMessage.ts b/packages/agent/src/tools/agent/agentMessage.ts
@@ -3,8 +3,6 @@ import { zodToJsonSchema } from 'zod-to-json-schema';
 
 import { Tool } from '../../core/types.js';
 
-import { agentStates } from './agentStart.js';
-
 const parameterSchema = z.object({
   agentId: z.string().describe('The ID returned by agentStart'),
   guidance: z
@@ -65,18 +63,8 @@ export const agentMessageTool: Tool<Parameters, ReturnType> = {
     );
 
     try {
-      // First try to get the agent from the tracker
-      let agent = agentTracker.getAgent(agentId);
-
-      // Fall back to legacy agentStates for backward compatibility
-      if (!agent && agentStates.has(agentId)) {
-        // If found in legacy store, register it with the tracker for future use
-        const legacyState = agentStates.get(agentId)!;
-        agentTracker.registerAgent(legacyState);
-
-        // Try again with the newly registered agent
-        agent = agentTracker.getAgent(agentId);
-      }
+      // Get the agent from the tracker
+      const agent = agentTracker.getAgent(agentId);
 
       if (!agent) {
         throw new Error(`No sub-agent found with ID ${agentId}`);
diff --git a/packages/agent/src/tools/agent/agentStart.ts b/packages/agent/src/tools/agent/agentStart.ts
@@ -11,10 +11,7 @@ import { Tool, ToolContext } from '../../core/types.js';
 import { LogLevel, Logger, LoggerListener } from '../../utils/logger.js';
 import { getTools } from '../getTools.js';
 
-import { AgentStatus, AgentInfo } from './AgentTracker.js';
-
-// For backward compatibility
-export const agentStates = new Map<string, AgentInfo>();
+import { AgentStatus } from './AgentTracker.js';
 
 // Generate a random color for an agent
 // Avoid colors that are too light or too similar to error/warning colors
@@ -189,11 +186,8 @@ export const agentStartTool: Tool<Parameters, ReturnType> = {
 
     logger.debug(`Registered agent with ID: ${agentId}`);
 
-    // For backward compatibility
-    const agent = agentTracker.getAgent(agentId);
-    if (agent) {
-      agentStates.set(agentId, agent);
-    }
+    // Get the agent for verification (not used but useful for debugging)
+    const _agent = agentTracker.getAgent(agentId);
 
     // Start the agent in a separate promise that we don't await
     // eslint-disable-next-line promise/catch-or-return
diff --git a/packages/agent/src/tools/agent/agentTools.test.ts b/packages/agent/src/tools/agent/agentTools.test.ts
@@ -7,7 +7,7 @@ import { SessionTracker } from '../session/SessionTracker.js';
 import { ShellTracker } from '../shell/ShellTracker.js';
 
 import { agentMessageTool } from './agentMessage.js';
-import { agentStartTool, agentStates } from './agentStart.js';
+import { agentStartTool } from './agentStart.js';
 import { AgentTracker } from './AgentTracker.js';
 
 // Mock the toolAgent function
@@ -57,9 +57,6 @@ describe('Agent Tools', () => {
       expect(agent).toHaveProperty('goal', 'Test the agent tools');
       expect(agent).toHaveProperty('completed', false);
       expect(agent).toHaveProperty('aborted', false);
-
-      // Verify it was also added to legacy agentStates for backward compatibility
-      expect(agentStates.has(result.agentId)).toBe(true);
     });
   });
 
