refactor(agent): merge Agent and AgentState into AgentInfo type

- Created new AgentInfo type that combines all fields from Agent and AgentState
- Updated AgentTracker to use a single Map for agent data
- Maintained backward compatibility with Agent and AgentState types
- Updated tests to use the new type
- Fixed tests that were broken by the refactoring

Closes #377

Author: bhouston

packages/agent/src/tools/agent/AgentTracker.ts

@@ -10,58 +10,121 @@ export enum AgentStatus {
   TERMINATED = 'terminated',
 }
 
-export interface Agent {
+export interface AgentInfo {
+  // Basic identification and status
   agentId: string;
   status: AgentStatus;
   startTime: Date;
   endTime?: Date;
   goal: string;
+
+  // Result information
   result?: string;
   error?: string;
-}
 
-// Internal agent state tracking (similar to existing agentStates)
-export interface AgentState {
-  agentId: string;
-  goal: string;
+  // Internal state information
   prompt: string;
   output: string;
   capturedLogs: string[]; // Captured log messages from agent and immediate tools
   completed: boolean;
-  error?: string;
-  result?: ToolAgentResult;
+  result_detailed?: ToolAgentResult;
   context: ToolContext;
   workingDirectory: string;
   tools: unknown[];
   aborted: boolean;
   parentMessages: string[]; // Messages from parent agent
 }
 
+// For backward compatibility
+export type Agent = Pick<
+  AgentInfo,
+  'agentId' | 'status' | 'startTime' | 'endTime' | 'goal' | 'result' | 'error'
+>;
+export type AgentState = Pick<
+  AgentInfo,
+  | 'agentId'
+  | 'goal'
+  | 'prompt'
+  | 'output'
+  | 'capturedLogs'
+  | 'completed'
+  | 'error'
+  | 'context'
+  | 'workingDirectory'
+  | 'tools'
+  | 'aborted'
+  | 'parentMessages'
+> & { result?: ToolAgentResult };
+
 export class AgentTracker {
-  private agents: Map<string, Agent> = new Map();
-  private agentStates: Map<string, AgentState> = new Map();
+  private agentInfos: Map<string, AgentInfo> = new Map();
 
   constructor(public ownerAgentId: string | undefined) {}
 
   // Register a new agent
   public registerAgent(goal: string): string {
     const agentId = uuidv4();
 
-    // Create agent tracking entry
-    const agent: Agent = {
+    // Create basic agent info entry
+    const agentInfo: Partial<AgentInfo> = {
       agentId: agentId,
       status: AgentStatus.RUNNING,
       startTime: new Date(),
       goal,
+      // Initialize arrays and default values
+      capturedLogs: [],
+      completed: false,
+      aborted: false,
+      parentMessages: [],
+      output: '',
     };
 
-    this.agents.set(agentId, agent);
+    this.agentInfos.set(agentId, agentInfo as AgentInfo);
     return agentId;
   }
 
-  // Register agent state
+  // Register agent state - for backward compatibility
   public registerAgentState(agentId: string, state: AgentState): void {
-    this.agentStates.set(agentId, state);
+    const agentInfo = this.agentInfos.get(agentId);
+
+    if (!agentInfo) {
+      // If agent doesn't exist yet (shouldn't happen in normal flow), create it
+      const newAgentInfo: AgentInfo = {
+        agentId: state.agentId,
+        status: AgentStatus.RUNNING,
+        startTime: new Date(),
+        goal: state.goal,
+        prompt: state.prompt,
+        output: state.output,
+        capturedLogs: state.capturedLogs,
+        completed: state.completed,
+        error: state.error,
+        result_detailed: state.result,
+        context: state.context,
+        workingDirectory: state.workingDirectory,
+        tools: state.tools,
+        aborted: state.aborted,
+        parentMessages: state.parentMessages,
+      };
+      this.agentInfos.set(agentId, newAgentInfo);
+      return;
+    }
+
+    // Update existing agent info with state data
+    Object.assign(agentInfo, {
+      goal: state.goal,
+      prompt: state.prompt,
+      output: state.output,
+      capturedLogs: state.capturedLogs,
+      completed: state.completed,
+      error: state.error,
+      result_detailed: state.result,
+      context: state.context,
+      workingDirectory: state.workingDirectory,
+      tools: state.tools,
+      aborted: state.aborted,
+      parentMessages: state.parentMessages,
+    });
   }
 
   // Update agent status
@@ -70,48 +133,95 @@ export class AgentTracker {
     status: AgentStatus,
     metadata?: { result?: string; error?: string },
   ): boolean {
-    const agent = this.agents.get(agentId);
-    if (!agent) {
+    const agentInfo = this.agentInfos.get(agentId);
+    if (!agentInfo) {
       return false;
     }
 
-    agent.status = status;
+    agentInfo.status = status;
 
     if (
       status === AgentStatus.COMPLETED ||
       status === AgentStatus.ERROR ||
       status === AgentStatus.TERMINATED
     ) {
-      agent.endTime = new Date();
+      agentInfo.endTime = new Date();
     }
 
     if (metadata) {
-      if (metadata.result !== undefined) agent.result = metadata.result;
-      if (metadata.error !== undefined) agent.error = metadata.error;
+      if (metadata.result !== undefined) agentInfo.result = metadata.result;
+      if (metadata.error !== undefined) agentInfo.error = metadata.error;
     }
 
     return true;
   }
 
-  // Get a specific agent state
+  // Get a specific agent info
+  public getAgentInfo(agentId: string): AgentInfo | undefined {
+    return this.agentInfos.get(agentId);
+  }
+
+  // Get a specific agent state - for backward compatibility
   public getAgentState(agentId: string): AgentState | undefined {
-    return this.agentStates.get(agentId);
+    const agentInfo = this.agentInfos.get(agentId);
+    if (!agentInfo) return undefined;
+
+    // Convert AgentInfo to AgentState
+    const state: AgentState = {
+      agentId: agentInfo.agentId,
+      goal: agentInfo.goal,
+      prompt: agentInfo.prompt,
+      output: agentInfo.output,
+      capturedLogs: agentInfo.capturedLogs,
+      completed: agentInfo.completed,
+      error: agentInfo.error,
+      result: agentInfo.result_detailed,
+      context: agentInfo.context,
+      workingDirectory: agentInfo.workingDirectory,
+      tools: agentInfo.tools,
+      aborted: agentInfo.aborted,
+      parentMessages: agentInfo.parentMessages,
+    };
+
+    return state;
   }
 
-  // Get a specific agent tracking info
+  // Get a specific agent tracking info - for backward compatibility
   public getAgent(agentId: string): Agent | undefined {
-    return this.agents.get(agentId);
+    const agentInfo = this.agentInfos.get(agentId);
+    if (!agentInfo) return undefined;
+
+    // Convert AgentInfo to Agent
+    const agent: Agent = {
+      agentId: agentInfo.agentId,
+      status: agentInfo.status,
+      startTime: agentInfo.startTime,
+      endTime: agentInfo.endTime,
+      goal: agentInfo.goal,
+      result: agentInfo.result,
+      error: agentInfo.error,
+    };
+
+    return agent;
   }
 
   // Get all agents with optional filtering
   public getAgents(status?: AgentStatus): Agent[] {
+    const agents = Array.from(this.agentInfos.values()).map((info) => ({
+      agentId: info.agentId,
+      status: info.status,
+      startTime: info.startTime,
+      endTime: info.endTime,
+      goal: info.goal,
+      result: info.result,
+      error: info.error,
+    }));
+
     if (!status) {
-      return Array.from(this.agents.values());
+      return agents;
     }
 
-    return Array.from(this.agents.values()).filter(
-      (agent) => agent.status === status,
-    );
+    return agents.filter((agent) => agent.status === status);
   }
 
   /**
@@ -122,11 +232,13 @@ export class AgentTracker {
     description: string;
     status: AgentStatus;
   }> {
-    return this.getAgents(AgentStatus.RUNNING).map((agent) => ({
-      agentId: agent.agentId,
-      description: agent.goal,
-      status: agent.status,
-    }));
+    return Array.from(this.agentInfos.values())
+      .filter((info) => info.status === AgentStatus.RUNNING)
+      .map((info) => ({
+        agentId: info.agentId,
+        description: info.goal,
+        status: info.status,
+      }));
   }
 
   // Cleanup and terminate agents
@@ -141,16 +253,16 @@ export class AgentTracker {
   // Terminate a specific agent
   public async terminateAgent(agentId: string): Promise<void> {
     try {
-      const agentState = this.agentStates.get(agentId);
-      if (agentState && !agentState.aborted) {
+      const agentInfo = this.agentInfos.get(agentId);
+      if (agentInfo && !agentInfo.aborted) {
         // Set the agent as aborted and completed
-        agentState.aborted = true;
-        agentState.completed = true;
+        agentInfo.aborted = true;
+        agentInfo.completed = true;
 
         // Clean up resources owned by this sub-agent
-        await agentState.context.agentTracker.cleanup();
-        await agentState.context.shellTracker.cleanup();
-        await agentState.context.browserTracker.cleanup();
+        await agentInfo.context.agentTracker.cleanup();
+        await agentInfo.context.shellTracker.cleanup();
+        await agentInfo.context.browserTracker.cleanup();
       }
       this.updateAgentStatus(agentId, AgentStatus.TERMINATED);
     } catch (error) {

packages/agent/src/tools/agent/__tests__/logCapture.test.ts

@@ -45,15 +45,15 @@ describe('Log Capture in AgentTracker', () => {
       context,
     );
 
-    // Get the agent state
-    const agentState = agentTracker.getAgentState(startResult.agentId);
-    expect(agentState).toBeDefined();
+    // Get the agent info directly
+    const agentInfo = agentTracker.getAgentInfo(startResult.agentId);
+    expect(agentInfo).toBeDefined();
 
-    if (!agentState) return; // TypeScript guard
+    if (!agentInfo) return; // TypeScript guard
 
-    // For testing purposes, manually add logs to the agent state
+    // For testing purposes, manually add logs to the agent info
     // In a real scenario, these would be added by the log listener
-    agentState.capturedLogs = [
+    agentInfo.capturedLogs = [
       'This log message should be captured',
       '[WARN] This warning message should be captured',
       '[ERROR] This error message should be captured',
@@ -62,28 +62,28 @@ describe('Log Capture in AgentTracker', () => {
     ];
 
     // Check that the right messages were captured
-    expect(agentState.capturedLogs.length).toBe(5);
-    expect(agentState.capturedLogs).toContain(
+    expect(agentInfo.capturedLogs.length).toBe(5);
+    expect(agentInfo.capturedLogs).toContain(
       'This log message should be captured',
     );
-    expect(agentState.capturedLogs).toContain(
+    expect(agentInfo.capturedLogs).toContain(
       '[WARN] This warning message should be captured',
     );
-    expect(agentState.capturedLogs).toContain(
+    expect(agentInfo.capturedLogs).toContain(
       '[ERROR] This error message should be captured',
     );
-    expect(agentState.capturedLogs).toContain(
+    expect(agentInfo.capturedLogs).toContain(
       'This tool log message should be captured',
     );
-    expect(agentState.capturedLogs).toContain(
+    expect(agentInfo.capturedLogs).toContain(
       '[WARN] This tool warning message should be captured',
     );
 
     // Make sure deep messages were not captured
-    expect(agentState.capturedLogs).not.toContain(
+    expect(agentInfo.capturedLogs).not.toContain(
       'This deep log message should NOT be captured',
     );
-    expect(agentState.capturedLogs).not.toContain(
+    expect(agentInfo.capturedLogs).not.toContain(
       '[ERROR] This deep error message should NOT be captured',
     );
 
@@ -109,7 +109,7 @@ describe('Log Capture in AgentTracker', () => {
     );
 
     // Check that the logs were cleared after being retrieved
-    expect(agentState.capturedLogs.length).toBe(0);
+    expect(agentInfo.capturedLogs.length).toBe(0);
   });
 
   it('should not include log section if no logs were captured', async () => {