Skip to content

Commit 792c58f

Browse files
bhoustonbhouston
authored
Merge pull request #216 from drivecore/feature/issue-112-background-tool-tracking
feat: implement background tool tracking (issue #112)
2 parents cf5ffc1 + 4a3bcc7 commit 792c58f

14 files changed

+729
-12
lines changed

‎packages/agent/src/core/backgroundTools.test.ts

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
import { describe, expect, it, vi, beforeEach } from 'vitest';
2+
3+
import {
4+
backgroundToolRegistry,
5+
BackgroundToolStatus,
6+
BackgroundToolType,
7+
} from './backgroundTools.js';
8+
9+
// Mock uuid to return predictable IDs for testing
10+
vi.mock('uuid', () => ({
11+
v4: vi.fn().mockReturnValue('test-id-1'), // Always return the same ID for simplicity in tests
12+
}));
13+
14+
describe('BackgroundToolRegistry', () => {
15+
beforeEach(() => {
16+
// Clear all registered tools before each test
17+
const registry = backgroundToolRegistry as any;
18+
registry.tools = new Map();
19+
});
20+
21+
it('should register a shell process', () => {
22+
const id = backgroundToolRegistry.registerShell('agent-1', 'ls -la');
23+
24+
expect(id).toBe('test-id-1');
25+
26+
const tool = backgroundToolRegistry.getToolById(id);
27+
expect(tool).toBeDefined();
28+
if (tool) {
29+
expect(tool.type).toBe(BackgroundToolType.SHELL);
30+
expect(tool.status).toBe(BackgroundToolStatus.RUNNING);
31+
expect(tool.agentId).toBe('agent-1');
32+
if (tool.type === BackgroundToolType.SHELL) {
33+
expect(tool.metadata.command).toBe('ls -la');
34+
}
35+
}
36+
});
37+
38+
it('should register a browser process', () => {
39+
const id = backgroundToolRegistry.registerBrowser(
40+
'agent-1',
41+
'https://example.com',
42+
);
43+
44+
expect(id).toBe('test-id-1');
45+
46+
const tool = backgroundToolRegistry.getToolById(id);
47+
expect(tool).toBeDefined();
48+
if (tool) {
49+
expect(tool.type).toBe(BackgroundToolType.BROWSER);
50+
expect(tool.status).toBe(BackgroundToolStatus.RUNNING);
51+
expect(tool.agentId).toBe('agent-1');
52+
if (tool.type === BackgroundToolType.BROWSER) {
53+
expect(tool.metadata.url).toBe('https://example.com');
54+
}
55+
}
56+
});
57+
58+
it('should update tool status', () => {
59+
const id = backgroundToolRegistry.registerShell('agent-1', 'sleep 10');
60+
61+
const updated = backgroundToolRegistry.updateToolStatus(
62+
id,
63+
BackgroundToolStatus.COMPLETED,
64+
{
65+
exitCode: 0,
66+
},
67+
);
68+
69+
expect(updated).toBe(true);
70+
71+
const tool = backgroundToolRegistry.getToolById(id);
72+
expect(tool).toBeDefined();
73+
if (tool) {
74+
expect(tool.status).toBe(BackgroundToolStatus.COMPLETED);
75+
expect(tool.endTime).toBeDefined();
76+
if (tool.type === BackgroundToolType.SHELL) {
77+
expect(tool.metadata.exitCode).toBe(0);
78+
}
79+
}
80+
});
81+
82+
it('should return false when updating non-existent tool', () => {
83+
const updated = backgroundToolRegistry.updateToolStatus(
84+
'non-existent-id',
85+
BackgroundToolStatus.COMPLETED,
86+
);
87+
88+
expect(updated).toBe(false);
89+
});
90+
91+
it('should get tools by agent ID', () => {
92+
// For this test, we'll directly manipulate the tools map
93+
const registry = backgroundToolRegistry as any;
94+
registry.tools = new Map();
95+
96+
// Add tools directly to the map with different agent IDs
97+
registry.tools.set('id1', {
98+
id: 'id1',
99+
type: BackgroundToolType.SHELL,
100+
status: BackgroundToolStatus.RUNNING,
101+
startTime: new Date(),
102+
agentId: 'agent-1',
103+
metadata: { command: 'ls -la' },
104+
});
105+
106+
registry.tools.set('id2', {
107+
id: 'id2',
108+
type: BackgroundToolType.BROWSER,
109+
status: BackgroundToolStatus.RUNNING,
110+
startTime: new Date(),
111+
agentId: 'agent-1',
112+
metadata: { url: 'https://example.com' },
113+
});
114+
115+
registry.tools.set('id3', {
116+
id: 'id3',
117+
type: BackgroundToolType.SHELL,
118+
status: BackgroundToolStatus.RUNNING,
119+
startTime: new Date(),
120+
agentId: 'agent-2',
121+
metadata: { command: 'echo hello' },
122+
});
123+
124+
const agent1Tools = backgroundToolRegistry.getToolsByAgent('agent-1');
125+
const agent2Tools = backgroundToolRegistry.getToolsByAgent('agent-2');
126+
127+
expect(agent1Tools.length).toBe(2);
128+
expect(agent2Tools.length).toBe(1);
129+
});
130+
131+
it('should clean up old completed tools', () => {
132+
// Create tools with specific dates
133+
const registry = backgroundToolRegistry as any;
134+
135+
// Add a completed tool from 25 hours ago
136+
const oldTool = {
137+
id: 'old-tool',
138+
type: BackgroundToolType.SHELL,
139+
status: BackgroundToolStatus.COMPLETED,
140+
startTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
141+
endTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
142+
agentId: 'agent-1',
143+
metadata: { command: 'echo old' },
144+
};
145+
146+
// Add a completed tool from 10 hours ago
147+
const recentTool = {
148+
id: 'recent-tool',
149+
type: BackgroundToolType.SHELL,
150+
status: BackgroundToolStatus.COMPLETED,
151+
startTime: new Date(Date.now() - 10 * 60 * 60 * 1000),
152+
endTime: new Date(Date.now() - 10 * 60 * 60 * 1000),
153+
agentId: 'agent-1',
154+
metadata: { command: 'echo recent' },
155+
};
156+
157+
// Add a running tool from 25 hours ago
158+
const oldRunningTool = {
159+
id: 'old-running-tool',
160+
type: BackgroundToolType.SHELL,
161+
status: BackgroundToolStatus.RUNNING,
162+
startTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
163+
agentId: 'agent-1',
164+
metadata: { command: 'sleep 100' },
165+
};
166+
167+
registry.tools.set('old-tool', oldTool);
168+
registry.tools.set('recent-tool', recentTool);
169+
registry.tools.set('old-running-tool', oldRunningTool);
170+
171+
// Clean up tools older than 24 hours
172+
backgroundToolRegistry.cleanupOldTools(24);
173+
174+
// Old completed tool should be removed
175+
expect(backgroundToolRegistry.getToolById('old-tool')).toBeUndefined();
176+
177+
// Recent completed tool should remain
178+
expect(backgroundToolRegistry.getToolById('recent-tool')).toBeDefined();
179+
180+
// Old running tool should remain (not completed)
181+
expect(
182+
backgroundToolRegistry.getToolById('old-running-tool'),
183+
).toBeDefined();
184+
});
185+
});

‎packages/agent/src/core/backgroundTools.ts

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
import { v4 as uuidv4 } from 'uuid';
2+
3+
// Types of background processes we can track
4+
export enum BackgroundToolType {
5+
SHELL = 'shell',
6+
BROWSER = 'browser',
7+
AGENT = 'agent',
8+
}
9+
10+
// Status of a background process
11+
export enum BackgroundToolStatus {
12+
RUNNING = 'running',
13+
COMPLETED = 'completed',
14+
ERROR = 'error',
15+
TERMINATED = 'terminated',
16+
}
17+
18+
// Common interface for all background processes
19+
export interface BackgroundTool {
20+
id: string;
21+
type: BackgroundToolType;
22+
status: BackgroundToolStatus;
23+
startTime: Date;
24+
endTime?: Date;
25+
agentId: string; // To track which agent created this process
26+
metadata: Record<string, any>; // Additional tool-specific information
27+
}
28+
29+
// Shell process specific data
30+
export interface ShellBackgroundTool extends BackgroundTool {
31+
type: BackgroundToolType.SHELL;
32+
metadata: {
33+
command: string;
34+
exitCode?: number | null;
35+
signaled?: boolean;
36+
};
37+
}
38+
39+
// Browser process specific data
40+
export interface BrowserBackgroundTool extends BackgroundTool {
41+
type: BackgroundToolType.BROWSER;
42+
metadata: {
43+
url?: string;
44+
};
45+
}
46+
47+
// Agent process specific data (for future use)
48+
export interface AgentBackgroundTool extends BackgroundTool {
49+
type: BackgroundToolType.AGENT;
50+
metadata: {
51+
goal?: string;
52+
};
53+
}
54+
55+
// Utility type for all background tool types
56+
export type AnyBackgroundTool =
57+
| ShellBackgroundTool
58+
| BrowserBackgroundTool
59+
| AgentBackgroundTool;
60+
61+
/**
62+
* Registry to keep track of all background processes
63+
*/
64+
export class BackgroundToolRegistry {
65+
private static instance: BackgroundToolRegistry;
66+
private tools: Map<string, AnyBackgroundTool> = new Map();
67+
68+
// Private constructor for singleton pattern
69+
private constructor() {}
70+
71+
// Get the singleton instance
72+
public static getInstance(): BackgroundToolRegistry {
73+
if (!BackgroundToolRegistry.instance) {
74+
BackgroundToolRegistry.instance = new BackgroundToolRegistry();
75+
}
76+
return BackgroundToolRegistry.instance;
77+
}
78+
79+
// Register a new shell process
80+
public registerShell(agentId: string, command: string): string {
81+
const id = uuidv4();
82+
const tool: ShellBackgroundTool = {
83+
id,
84+
type: BackgroundToolType.SHELL,
85+
status: BackgroundToolStatus.RUNNING,
86+
startTime: new Date(),
87+
agentId,
88+
metadata: {
89+
command,
90+
},
91+
};
92+
this.tools.set(id, tool);
93+
return id;
94+
}
95+
96+
// Register a new browser process
97+
public registerBrowser(agentId: string, url?: string): string {
98+
const id = uuidv4();
99+
const tool: BrowserBackgroundTool = {
100+
id,
101+
type: BackgroundToolType.BROWSER,
102+
status: BackgroundToolStatus.RUNNING,
103+
startTime: new Date(),
104+
agentId,
105+
metadata: {
106+
url,
107+
},
108+
};
109+
this.tools.set(id, tool);
110+
return id;
111+
}
112+
113+
// Register a new agent process (for future use)
114+
public registerAgent(agentId: string, goal?: string): string {
115+
const id = uuidv4();
116+
const tool: AgentBackgroundTool = {
117+
id,
118+
type: BackgroundToolType.AGENT,
119+
status: BackgroundToolStatus.RUNNING,
120+
startTime: new Date(),
121+
agentId,
122+
metadata: {
123+
goal,
124+
},
125+
};
126+
this.tools.set(id, tool);
127+
return id;
128+
}
129+
130+
// Update the status of a process
131+
public updateToolStatus(
132+
id: string,
133+
status: BackgroundToolStatus,
134+
metadata?: Record<string, any>,
135+
): boolean {
136+
const tool = this.tools.get(id);
137+
if (!tool) {
138+
return false;
139+
}
140+
141+
tool.status = status;
142+
143+
if (
144+
status === BackgroundToolStatus.COMPLETED ||
145+
status === BackgroundToolStatus.ERROR ||
146+
status === BackgroundToolStatus.TERMINATED
147+
) {
148+
tool.endTime = new Date();
149+
}
150+
151+
if (metadata) {
152+
tool.metadata = { ...tool.metadata, ...metadata };
153+
}
154+
155+
return true;
156+
}
157+
158+
// Get all processes for a specific agent
159+
public getToolsByAgent(agentId: string): AnyBackgroundTool[] {
160+
const result: AnyBackgroundTool[] = [];
161+
for (const tool of this.tools.values()) {
162+
if (tool.agentId === agentId) {
163+
result.push(tool);
164+
}
165+
}
166+
return result;
167+
}
168+
169+
// Get a specific process by ID
170+
public getToolById(id: string): AnyBackgroundTool | undefined {
171+
return this.tools.get(id);
172+
}
173+
174+
// Clean up completed processes (optional, for maintenance)
175+
public cleanupOldTools(olderThanHours: number = 24): void {
176+
const cutoffTime = new Date(Date.now() - olderThanHours * 60 * 60 * 1000);
177+
178+
for (const [id, tool] of this.tools.entries()) {
179+
// Remove if it's completed/error/terminated AND older than cutoff
180+
if (
181+
tool.endTime &&
182+
tool.endTime < cutoffTime &&
183+
(tool.status === BackgroundToolStatus.COMPLETED ||
184+
tool.status === BackgroundToolStatus.ERROR ||
185+
tool.status === BackgroundToolStatus.TERMINATED)
186+
) {
187+
this.tools.delete(id);
188+
}
189+
}
190+
}
191+
}
192+
193+
// Export singleton instance
194+
export const backgroundToolRegistry = BackgroundToolRegistry.getInstance();

‎packages/agent/src/core/types.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ export type ToolContext = {
2020
customPrompt?: string;
2121
tokenCache?: boolean;
2222
enableUserPrompt?: boolean;
23+
agentId?: string; // Unique identifier for the agent, used for background tool tracking
2324
};
2425

2526
export type Tool<TParams = Record<string, any>, TReturn = any> = {

0 commit comments

Comments
 (0)