Skip to content

Commit 92a5f88

Browse files
bhoustonbhouston
authored
Merge pull request #239 from drivecore/feature/236-incremental-resource-cleanup
Implement Incremental Resource Cleanup for Agent Lifecycle Management
2 parents 4020130 + 123bb22 commit 92a5f88

File tree

7 files changed

+414
-4
lines changed

7 files changed

+414
-4
lines changed

‎commit_message.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
refactor(agent): implement parallel resource cleanup
2+
3+
This change improves the resource cleanup process by handling browser sessions,
4+
shell processes, and sub-agents in parallel rather than sequentially.
5+
6+
The implementation:
7+
1. Refactors the cleanup method into smaller helper methods for each resource type
8+
2. Uses Promise.all to execute cleanup operations concurrently
9+
3. Filters tools by status during the initial grouping to simplify the code
10+
11+
This approach significantly speeds up the cleanup process, especially when
12+
dealing with multiple resources of different types.

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

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
2+
3+
// Import mocked modules
4+
import { BrowserManager } from '../tools/browser/BrowserManager.js';
5+
import { agentStates } from '../tools/interaction/agentStart.js';
6+
import { processStates } from '../tools/system/shellStart.js';
7+
8+
import { BackgroundTools, BackgroundToolStatus } from './backgroundTools';
9+
import { Tool } from './types';
10+
11+
// Define types for our mocks that match the actual types
12+
type MockProcessState = {
13+
process: { kill: ReturnType<typeof vi.fn> };
14+
state: { completed: boolean };
15+
command?: string;
16+
stdout?: string[];
17+
stderr?: string[];
18+
showStdIn?: boolean;
19+
showStdout?: boolean;
20+
};
21+
22+
type MockAgentState = {
23+
aborted: boolean;
24+
completed: boolean;
25+
context: {
26+
backgroundTools: {
27+
cleanup: ReturnType<typeof vi.fn>;
28+
};
29+
};
30+
goal?: string;
31+
prompt?: string;
32+
output?: string;
33+
workingDirectory?: string;
34+
tools?: Tool[];
35+
};
36+
37+
// Mock dependencies
38+
vi.mock('../tools/browser/BrowserManager.js', () => {
39+
return {
40+
BrowserManager: class MockBrowserManager {
41+
closeSession = vi.fn().mockResolvedValue(undefined);
42+
},
43+
};
44+
});
45+
46+
vi.mock('../tools/system/shellStart.js', () => {
47+
return {
48+
processStates: new Map<string, MockProcessState>(),
49+
};
50+
});
51+
52+
vi.mock('../tools/interaction/agentStart.js', () => {
53+
return {
54+
agentStates: new Map<string, MockAgentState>(),
55+
};
56+
});
57+
58+
describe('BackgroundTools cleanup', () => {
59+
let backgroundTools: BackgroundTools;
60+
61+
// Setup mocks for globalThis and process states
62+
beforeEach(() => {
63+
backgroundTools = new BackgroundTools('test-agent');
64+
65+
// Reset mocks
66+
vi.resetAllMocks();
67+
68+
// Setup global browser manager
69+
(
70+
globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager }
71+
).__BROWSER_MANAGER__ = {
72+
closeSession: vi.fn().mockResolvedValue(undefined),
73+
} as unknown as BrowserManager;
74+
75+
// Setup mock process states
76+
const mockProcess = {
77+
kill: vi.fn(),
78+
};
79+
80+
const mockProcessState: MockProcessState = {
81+
process: mockProcess,
82+
state: { completed: false },
83+
command: 'test command',
84+
stdout: [],
85+
stderr: [],
86+
showStdIn: false,
87+
showStdout: false,
88+
};
89+
90+
processStates.clear();
91+
processStates.set('shell-1', mockProcessState as any);
92+
93+
// Setup mock agent states
94+
const mockAgentState: MockAgentState = {
95+
aborted: false,
96+
completed: false,
97+
context: {
98+
backgroundTools: {
99+
cleanup: vi.fn().mockResolvedValue(undefined),
100+
},
101+
},
102+
goal: 'test goal',
103+
prompt: 'test prompt',
104+
output: '',
105+
workingDirectory: '/test',
106+
tools: [],
107+
};
108+
109+
agentStates.clear();
110+
agentStates.set('agent-1', mockAgentState as any);
111+
});
112+
113+
afterEach(() => {
114+
vi.resetAllMocks();
115+
116+
// Clear global browser manager
117+
(
118+
globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager }
119+
).__BROWSER_MANAGER__ = undefined;
120+
121+
// Clear mock states
122+
processStates.clear();
123+
agentStates.clear();
124+
});
125+
126+
it('should clean up browser sessions', async () => {
127+
// Register a browser tool
128+
const browserId = backgroundTools.registerBrowser('https://example.com');
129+
130+
// Run cleanup
131+
await backgroundTools.cleanup();
132+
133+
// Check that closeSession was called
134+
expect(
135+
(globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
136+
.__BROWSER_MANAGER__.closeSession,
137+
).toHaveBeenCalledWith(browserId);
138+
139+
// Check that tool status was updated
140+
const tool = backgroundTools.getToolById(browserId);
141+
expect(tool?.status).toBe(BackgroundToolStatus.COMPLETED);
142+
});
143+
144+
it('should clean up shell processes', async () => {
145+
// Register a shell tool
146+
const shellId = backgroundTools.registerShell('echo "test"');
147+
148+
// Get mock process state
149+
const mockProcessState = processStates.get('shell-1');
150+
151+
// Set the shell ID to match
152+
processStates.set(shellId, processStates.get('shell-1') as any);
153+
154+
// Run cleanup
155+
await backgroundTools.cleanup();
156+
157+
// Check that kill was called
158+
expect(mockProcessState?.process.kill).toHaveBeenCalledWith('SIGTERM');
159+
160+
// Check that tool status was updated
161+
const tool = backgroundTools.getToolById(shellId);
162+
expect(tool?.status).toBe(BackgroundToolStatus.COMPLETED);
163+
});
164+
165+
it('should clean up sub-agents', async () => {
166+
// Register an agent tool
167+
const agentId = backgroundTools.registerAgent('Test goal');
168+
169+
// Get mock agent state
170+
const mockAgentState = agentStates.get('agent-1');
171+
172+
// Set the agent ID to match
173+
agentStates.set(agentId, agentStates.get('agent-1') as any);
174+
175+
// Run cleanup
176+
await backgroundTools.cleanup();
177+
178+
// Check that agent was marked as aborted
179+
expect(mockAgentState?.aborted).toBe(true);
180+
expect(mockAgentState?.completed).toBe(true);
181+
182+
// Check that cleanup was called on the agent's background tools
183+
expect(mockAgentState?.context.backgroundTools.cleanup).toHaveBeenCalled();
184+
185+
// Check that tool status was updated
186+
const tool = backgroundTools.getToolById(agentId);
187+
expect(tool?.status).toBe(BackgroundToolStatus.TERMINATED);
188+
});
189+
190+
it('should handle errors during cleanup', async () => {
191+
// Register a browser tool
192+
const browserId = backgroundTools.registerBrowser('https://example.com');
193+
194+
// Make closeSession throw an error
195+
(
196+
(globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
197+
.__BROWSER_MANAGER__.closeSession as ReturnType<typeof vi.fn>
198+
).mockRejectedValue(new Error('Test error'));
199+
200+
// Run cleanup
201+
await backgroundTools.cleanup();
202+
203+
// Check that tool status was updated to ERROR
204+
const tool = backgroundTools.getToolById(browserId);
205+
expect(tool?.status).toBe(BackgroundToolStatus.ERROR);
206+
expect(tool?.metadata.error).toBe('Test error');
207+
});
208+
209+
it('should only clean up running tools', async () => {
210+
// Register a browser tool and mark it as completed
211+
const browserId = backgroundTools.registerBrowser('https://example.com');
212+
backgroundTools.updateToolStatus(browserId, BackgroundToolStatus.COMPLETED);
213+
214+
// Run cleanup
215+
await backgroundTools.cleanup();
216+
217+
// Check that closeSession was not called
218+
expect(
219+
(globalThis as unknown as { __BROWSER_MANAGER__: BrowserManager })
220+
.__BROWSER_MANAGER__.closeSession,
221+
).not.toHaveBeenCalled();
222+
});
223+
});

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

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
import { v4 as uuidv4 } from 'uuid';
22

3+
// These imports will be used by the cleanup method
4+
import { BrowserManager } from '../tools/browser/BrowserManager.js';
5+
import { agentStates } from '../tools/interaction/agentStart.js';
6+
import { processStates } from '../tools/system/shellStart.js';
7+
38
// Types of background processes we can track
49
export enum BackgroundToolType {
510
SHELL = 'shell',
@@ -32,6 +37,7 @@ export interface ShellBackgroundTool extends BackgroundTool {
3237
command: string;
3338
exitCode?: number | null;
3439
signaled?: boolean;
40+
error?: string;
3541
};
3642
}
3743

@@ -40,6 +46,7 @@ export interface BrowserBackgroundTool extends BackgroundTool {
4046
type: BackgroundToolType.BROWSER;
4147
metadata: {
4248
url?: string;
49+
error?: string;
4350
};
4451
}
4552

@@ -48,6 +55,7 @@ export interface AgentBackgroundTool extends BackgroundTool {
4855
type: BackgroundToolType.AGENT;
4956
metadata: {
5057
goal?: string;
58+
error?: string;
5159
};
5260
}
5361

@@ -154,4 +162,126 @@ export class BackgroundTools {
154162
public getToolById(id: string): AnyBackgroundTool | undefined {
155163
return this.tools.get(id);
156164
}
165+
166+
/**
167+
* Cleans up all resources associated with this agent instance
168+
* @returns A promise that resolves when cleanup is complete
169+
*/
170+
public async cleanup(): Promise<void> {
171+
const tools = this.getTools();
172+
173+
// Group tools by type for better cleanup organization
174+
const browserTools = tools.filter(
175+
(tool): tool is BrowserBackgroundTool =>
176+
tool.type === BackgroundToolType.BROWSER &&
177+
tool.status === BackgroundToolStatus.RUNNING,
178+
);
179+
180+
const shellTools = tools.filter(
181+
(tool): tool is ShellBackgroundTool =>
182+
tool.type === BackgroundToolType.SHELL &&
183+
tool.status === BackgroundToolStatus.RUNNING,
184+
);
185+
186+
const agentTools = tools.filter(
187+
(tool): tool is AgentBackgroundTool =>
188+
tool.type === BackgroundToolType.AGENT &&
189+
tool.status === BackgroundToolStatus.RUNNING,
190+
);
191+
192+
// Create cleanup promises for each resource type
193+
const browserCleanupPromises = browserTools.map((tool) =>
194+
this.cleanupBrowserSession(tool),
195+
);
196+
const shellCleanupPromises = shellTools.map((tool) =>
197+
this.cleanupShellProcess(tool),
198+
);
199+
const agentCleanupPromises = agentTools.map((tool) =>
200+
this.cleanupSubAgent(tool),
201+
);
202+
203+
// Wait for all cleanup operations to complete in parallel
204+
await Promise.all([
205+
...browserCleanupPromises,
206+
...shellCleanupPromises,
207+
...agentCleanupPromises,
208+
]);
209+
}
210+
211+
/**
212+
* Cleans up a browser session
213+
* @param tool The browser tool to clean up
214+
*/
215+
private async cleanupBrowserSession(
216+
tool: BrowserBackgroundTool,
217+
): Promise<void> {
218+
try {
219+
const browserManager = (
220+
globalThis as unknown as { __BROWSER_MANAGER__?: BrowserManager }
221+
).__BROWSER_MANAGER__;
222+
if (browserManager) {
223+
await browserManager.closeSession(tool.id);
224+
}
225+
this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
226+
} catch (error) {
227+
this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
228+
error: error instanceof Error ? error.message : String(error),
229+
});
230+
}
231+
}
232+
233+
/**
234+
* Cleans up a shell process
235+
* @param tool The shell tool to clean up
236+
*/
237+
private async cleanupShellProcess(tool: ShellBackgroundTool): Promise<void> {
238+
try {
239+
const processState = processStates.get(tool.id);
240+
if (processState && !processState.state.completed) {
241+
processState.process.kill('SIGTERM');
242+
243+
// Force kill after a short timeout if still running
244+
await new Promise<void>((resolve) => {
245+
setTimeout(() => {
246+
try {
247+
if (!processState.state.completed) {
248+
processState.process.kill('SIGKILL');
249+
}
250+
} catch {
251+
// Ignore errors on forced kill
252+
}
253+
resolve();
254+
}, 500);
255+
});
256+
}
257+
this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
258+
} catch (error) {
259+
this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
260+
error: error instanceof Error ? error.message : String(error),
261+
});
262+
}
263+
}
264+
265+
/**
266+
* Cleans up a sub-agent
267+
* @param tool The agent tool to clean up
268+
*/
269+
private async cleanupSubAgent(tool: AgentBackgroundTool): Promise<void> {
270+
try {
271+
const agentState = agentStates.get(tool.id);
272+
if (agentState && !agentState.aborted) {
273+
// Set the agent as aborted and completed
274+
agentState.aborted = true;
275+
agentState.completed = true;
276+
277+
// Clean up resources owned by this sub-agent
278+
await agentState.context.backgroundTools.cleanup();
279+
}
280+
this.updateToolStatus(tool.id, BackgroundToolStatus.TERMINATED);
281+
} catch (error) {
282+
this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
283+
error: error instanceof Error ? error.message : String(error),
284+
});
285+
}
286+
}
157287
}

0 commit comments

Comments
 (0)