drivecore · bhouston · Mar 12, 2025 · Mar 12, 2025 · Mar 12, 2025
diff --git a/packages/agent/src/core/backgroundTools.test.ts 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
@@ -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<string, any>; // 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<string, AnyBackgroundTool> = 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<string, any>,
+  ): 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
@@ -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<TParams = Record<string, any>, TReturn = any> = {