Skip to content

Commit 1de07ed

Browse files
bhoustonbhouston
committed
Merge remote-tracking branch 'origin/shell-tracker-implementation'
2 parents e8cc18f + 3dca767 commit 1de07ed

15 files changed

+595
-262
lines changed

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

Lines changed: 11 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
33
// Import mocked modules
44
import { BrowserManager } from '../tools/browser/BrowserManager.js';
55
import { agentStates } from '../tools/interaction/agentStart.js';
6-
import { processStates } from '../tools/system/shellStart.js';
6+
import { shellTracker } from '../tools/system/ShellTracker.js';
77

88
import { BackgroundTools, BackgroundToolStatus } from './backgroundTools';
99
import { Tool } from './types';
@@ -43,9 +43,12 @@ vi.mock('../tools/browser/BrowserManager.js', () => {
4343
};
4444
});
4545

46-
vi.mock('../tools/system/shellStart.js', () => {
46+
vi.mock('../tools/system/ShellTracker.js', () => {
4747
return {
48-
processStates: new Map<string, MockProcessState>(),
48+
shellTracker: {
49+
processStates: new Map<string, MockProcessState>(),
50+
cleanupAllShells: vi.fn().mockResolvedValue(undefined),
51+
},
4952
};
5053
});
5154

@@ -87,8 +90,8 @@ describe('BackgroundTools cleanup', () => {
8790
showStdout: false,
8891
};
8992

90-
processStates.clear();
91-
processStates.set('shell-1', mockProcessState as any);
93+
shellTracker.processStates.clear();
94+
shellTracker.processStates.set('shell-1', mockProcessState as any);
9295

9396
// Setup mock agent states
9497
const mockAgentState: MockAgentState = {
@@ -119,7 +122,7 @@ describe('BackgroundTools cleanup', () => {
119122
).__BROWSER_MANAGER__ = undefined;
120123

121124
// Clear mock states
122-
processStates.clear();
125+
shellTracker.processStates.clear();
123126
agentStates.clear();
124127
});
125128

@@ -142,24 +145,11 @@ describe('BackgroundTools cleanup', () => {
142145
});
143146

144147
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-
154148
// Run cleanup
155149
await backgroundTools.cleanup();
156150

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);
151+
// Check that shellTracker.cleanupAllShells was called
152+
expect(shellTracker.cleanupAllShells).toHaveBeenCalled();
163153
});
164154

165155
it('should clean up sub-agents', async () => {

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

Lines changed: 6 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -19,22 +19,6 @@ describe('BackgroundToolRegistry', () => {
1919
backgroundTools.tools = new Map();
2020
});
2121

22-
it('should register a shell process', () => {
23-
const id = backgroundTools.registerShell('ls -la');
24-
25-
expect(id).toBe('test-id-1');
26-
27-
const tool = backgroundTools.getToolById(id);
28-
expect(tool).toBeDefined();
29-
if (tool) {
30-
expect(tool.type).toBe(BackgroundToolType.SHELL);
31-
expect(tool.status).toBe(BackgroundToolStatus.RUNNING);
32-
if (tool.type === BackgroundToolType.SHELL) {
33-
expect(tool.metadata.command).toBe('ls -la');
34-
}
35-
}
36-
});
37-
3822
it('should register a browser process', () => {
3923
const id = backgroundTools.registerBrowser('https://example.com');
4024

@@ -51,30 +35,6 @@ describe('BackgroundToolRegistry', () => {
5135
}
5236
});
5337

54-
it('should update tool status', () => {
55-
const id = backgroundTools.registerShell('sleep 10');
56-
57-
const updated = backgroundTools.updateToolStatus(
58-
id,
59-
BackgroundToolStatus.COMPLETED,
60-
{
61-
exitCode: 0,
62-
},
63-
);
64-
65-
expect(updated).toBe(true);
66-
67-
const tool = backgroundTools.getToolById(id);
68-
expect(tool).toBeDefined();
69-
if (tool) {
70-
expect(tool.status).toBe(BackgroundToolStatus.COMPLETED);
71-
expect(tool.endTime).toBeDefined();
72-
if (tool.type === BackgroundToolType.SHELL) {
73-
expect(tool.metadata.exitCode).toBe(0);
74-
}
75-
}
76-
});
77-
7838
it('should return false when updating non-existent tool', () => {
7939
const updated = backgroundTools.updateToolStatus(
8040
'non-existent-id',
@@ -91,33 +51,33 @@ describe('BackgroundToolRegistry', () => {
9151
// Add a completed tool from 25 hours ago
9252
const oldTool = {
9353
id: 'old-tool',
94-
type: BackgroundToolType.SHELL,
54+
type: BackgroundToolType.BROWSER,
9555
status: BackgroundToolStatus.COMPLETED,
9656
startTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
9757
endTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
9858
agentId: 'agent-1',
99-
metadata: { command: 'echo old' },
59+
metadata: { url: 'https://example.com' },
10060
};
10161

10262
// Add a completed tool from 10 hours ago
10363
const recentTool = {
10464
id: 'recent-tool',
105-
type: BackgroundToolType.SHELL,
65+
type: BackgroundToolType.BROWSER,
10666
status: BackgroundToolStatus.COMPLETED,
10767
startTime: new Date(Date.now() - 10 * 60 * 60 * 1000),
10868
endTime: new Date(Date.now() - 10 * 60 * 60 * 1000),
10969
agentId: 'agent-1',
110-
metadata: { command: 'echo recent' },
70+
metadata: { url: 'https://example.com' },
11171
};
11272

11373
// Add a running tool from 25 hours ago
11474
const oldRunningTool = {
11575
id: 'old-running-tool',
116-
type: BackgroundToolType.SHELL,
76+
type: BackgroundToolType.BROWSER,
11777
status: BackgroundToolStatus.RUNNING,
11878
startTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
11979
agentId: 'agent-1',
120-
metadata: { command: 'sleep 100' },
80+
metadata: { url: 'https://example.com' },
12181
};
12282

12383
registry.tools.set('old-tool', oldTool);

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

Lines changed: 6 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@ import { v4 as uuidv4 } from 'uuid';
33
// These imports will be used by the cleanup method
44
import { BrowserManager } from '../tools/browser/BrowserManager.js';
55
import { agentStates } from '../tools/interaction/agentStart.js';
6-
import { processStates } from '../tools/system/shellStart.js';
6+
import { shellTracker } from '../tools/system/ShellTracker.js';
77

88
// Types of background processes we can track
99
export enum BackgroundToolType {
10-
SHELL = 'shell',
1110
BROWSER = 'browser',
1211
AGENT = 'agent',
1312
}
@@ -30,17 +29,6 @@ export interface BackgroundTool {
3029
metadata: Record<string, any>; // Additional tool-specific information
3130
}
3231

33-
// Shell process specific data
34-
export interface ShellBackgroundTool extends BackgroundTool {
35-
type: BackgroundToolType.SHELL;
36-
metadata: {
37-
command: string;
38-
exitCode?: number | null;
39-
signaled?: boolean;
40-
error?: string;
41-
};
42-
}
43-
4432
// Browser process specific data
4533
export interface BrowserBackgroundTool extends BackgroundTool {
4634
type: BackgroundToolType.BROWSER;
@@ -60,10 +48,7 @@ export interface AgentBackgroundTool extends BackgroundTool {
6048
}
6149

6250
// Utility type for all background tool types
63-
export type AnyBackgroundTool =
64-
| ShellBackgroundTool
65-
| BrowserBackgroundTool
66-
| AgentBackgroundTool;
51+
export type AnyBackgroundTool = BrowserBackgroundTool | AgentBackgroundTool;
6752

6853
/**
6954
* Registry to keep track of all background processes
@@ -74,22 +59,6 @@ export class BackgroundTools {
7459
// Private constructor for singleton pattern
7560
constructor(readonly ownerName: string) {}
7661

77-
// Register a new shell process
78-
public registerShell(command: string): string {
79-
const id = uuidv4();
80-
const tool: ShellBackgroundTool = {
81-
id,
82-
type: BackgroundToolType.SHELL,
83-
status: BackgroundToolStatus.RUNNING,
84-
startTime: new Date(),
85-
metadata: {
86-
command,
87-
},
88-
};
89-
this.tools.set(id, tool);
90-
return id;
91-
}
92-
9362
// Register a new browser process
9463
public registerBrowser(url?: string): string {
9564
const id = uuidv4();
@@ -177,12 +146,6 @@ export class BackgroundTools {
177146
tool.status === BackgroundToolStatus.RUNNING,
178147
);
179148

180-
const shellTools = tools.filter(
181-
(tool): tool is ShellBackgroundTool =>
182-
tool.type === BackgroundToolType.SHELL &&
183-
tool.status === BackgroundToolStatus.RUNNING,
184-
);
185-
186149
const agentTools = tools.filter(
187150
(tool): tool is AgentBackgroundTool =>
188151
tool.type === BackgroundToolType.AGENT &&
@@ -193,19 +156,15 @@ export class BackgroundTools {
193156
const browserCleanupPromises = browserTools.map((tool) =>
194157
this.cleanupBrowserSession(tool),
195158
);
196-
const shellCleanupPromises = shellTools.map((tool) =>
197-
this.cleanupShellProcess(tool),
198-
);
199159
const agentCleanupPromises = agentTools.map((tool) =>
200160
this.cleanupSubAgent(tool),
201161
);
202162

163+
// Clean up shell processes using ShellTracker
164+
await shellTracker.cleanupAllShells();
165+
203166
// Wait for all cleanup operations to complete in parallel
204-
await Promise.all([
205-
...browserCleanupPromises,
206-
...shellCleanupPromises,
207-
...agentCleanupPromises,
208-
]);
167+
await Promise.all([...browserCleanupPromises, ...agentCleanupPromises]);
209168
}
210169

211170
/**
@@ -230,38 +189,6 @@ export class BackgroundTools {
230189
}
231190
}
232191

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-
265192
/**
266193
* Cleans up a sub-agent
267194
* @param tool The agent tool to clean up

‎packages/agent/src/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ export * from './tools/system/sequenceComplete.js';
1010
export * from './tools/system/shellMessage.js';
1111
export * from './tools/system/shellExecute.js';
1212
export * from './tools/system/listBackgroundTools.js';
13+
export * from './tools/system/listShells.js';
14+
export * from './tools/system/ShellTracker.js';
1315

1416
// Tools - Browser
1517
export * from './tools/browser/BrowserManager.js';

‎packages/agent/src/tools/getTools.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { fetchTool } from './io/fetch.js';
1010
import { textEditorTool } from './io/textEditor.js';
1111
import { createMcpTool } from './mcp.js';
1212
import { listBackgroundToolsTool } from './system/listBackgroundTools.js';
13+
import { listShellsTool } from './system/listShells.js';
1314
import { sequenceCompleteTool } from './system/sequenceComplete.js';
1415
import { shellMessageTool } from './system/shellMessage.js';
1516
import { shellStartTool } from './system/shellStart.js';
@@ -41,6 +42,7 @@ export function getTools(options?: GetToolsOptions): Tool[] {
4142
//respawnTool as unknown as Tool, this is a confusing tool for now.
4243
sleepTool as unknown as Tool,
4344
listBackgroundToolsTool as unknown as Tool,
45+
listShellsTool as unknown as Tool,
4446
];
4547

4648
// Only include userPrompt tool if enabled

0 commit comments

Comments
 (0)