Skip to content

feat: implement ShellTracker to decouple from backgroundTools #296

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 18, 2025
Merged

feat: implement ShellTracker to decouple from backgroundTools #296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
65378e3
feat: implement ShellTracker to decouple from backgroundTools
bhouston Mar 14, 2025
3dca767
fix: update CLI cleanup to use ShellTracker instead of processStates
bhouston Mar 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 11 additions & 21 deletions packages/agent/src/core/backgroundTools.cleanup.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
// Import mocked modules
import { BrowserManager } from '../tools/browser/BrowserManager.js';
import { agentStates } from '../tools/interaction/agentStart.js';
import { processStates } from '../tools/system/shellStart.js';
import { shellTracker } from '../tools/system/ShellTracker.js';

import { BackgroundTools, BackgroundToolStatus } from './backgroundTools';
import { Tool } from './types';
Expand Down Expand Up @@ -43,9 +43,12 @@ vi.mock('../tools/browser/BrowserManager.js', () => {
};
});

vi.mock('../tools/system/shellStart.js', () => {
vi.mock('../tools/system/ShellTracker.js', () => {
return {
processStates: new Map<string, MockProcessState>(),
shellTracker: {
processStates: new Map<string, MockProcessState>(),
cleanupAllShells: vi.fn().mockResolvedValue(undefined),
},
};
});

Expand Down Expand Up @@ -87,8 +90,8 @@ describe('BackgroundTools cleanup', () => {
showStdout: false,
};

processStates.clear();
processStates.set('shell-1', mockProcessState as any);
shellTracker.processStates.clear();
shellTracker.processStates.set('shell-1', mockProcessState as any);

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

// Clear mock states
processStates.clear();
shellTracker.processStates.clear();
agentStates.clear();
});

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

it('should clean up shell processes', async () => {
// Register a shell tool
const shellId = backgroundTools.registerShell('echo "test"');

// Get mock process state
const mockProcessState = processStates.get('shell-1');

// Set the shell ID to match
processStates.set(shellId, processStates.get('shell-1') as any);

// Run cleanup
await backgroundTools.cleanup();

// Check that kill was called
expect(mockProcessState?.process.kill).toHaveBeenCalledWith('SIGTERM');

// Check that tool status was updated
const tool = backgroundTools.getToolById(shellId);
expect(tool?.status).toBe(BackgroundToolStatus.COMPLETED);
// Check that shellTracker.cleanupAllShells was called
expect(shellTracker.cleanupAllShells).toHaveBeenCalled();
});

it('should clean up sub-agents', async () => {
Expand Down
52 changes: 6 additions & 46 deletions packages/agent/src/core/backgroundTools.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,22 +19,6 @@ describe('BackgroundToolRegistry', () => {
backgroundTools.tools = new Map();
});

it('should register a shell process', () => {
const id = backgroundTools.registerShell('ls -la');

expect(id).toBe('test-id-1');

const tool = backgroundTools.getToolById(id);
expect(tool).toBeDefined();
if (tool) {
expect(tool.type).toBe(BackgroundToolType.SHELL);
expect(tool.status).toBe(BackgroundToolStatus.RUNNING);
if (tool.type === BackgroundToolType.SHELL) {
expect(tool.metadata.command).toBe('ls -la');
}
}
});

it('should register a browser process', () => {
const id = backgroundTools.registerBrowser('https://example.com');

Expand All @@ -51,30 +35,6 @@ describe('BackgroundToolRegistry', () => {
}
});

it('should update tool status', () => {
const id = backgroundTools.registerShell('sleep 10');

const updated = backgroundTools.updateToolStatus(
id,
BackgroundToolStatus.COMPLETED,
{
exitCode: 0,
},
);

expect(updated).toBe(true);

const tool = backgroundTools.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 = backgroundTools.updateToolStatus(
'non-existent-id',
Expand All @@ -91,33 +51,33 @@ describe('BackgroundToolRegistry', () => {
// Add a completed tool from 25 hours ago
const oldTool = {
id: 'old-tool',
type: BackgroundToolType.SHELL,
type: BackgroundToolType.BROWSER,
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' },
metadata: { url: 'https://example.com' },
};

// Add a completed tool from 10 hours ago
const recentTool = {
id: 'recent-tool',
type: BackgroundToolType.SHELL,
type: BackgroundToolType.BROWSER,
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' },
metadata: { url: 'https://example.com' },
};

// Add a running tool from 25 hours ago
const oldRunningTool = {
id: 'old-running-tool',
type: BackgroundToolType.SHELL,
type: BackgroundToolType.BROWSER,
status: BackgroundToolStatus.RUNNING,
startTime: new Date(Date.now() - 25 * 60 * 60 * 1000),
agentId: 'agent-1',
metadata: { command: 'sleep 100' },
metadata: { url: 'https://example.com' },
};

registry.tools.set('old-tool', oldTool);
Expand Down
85 changes: 6 additions & 79 deletions packages/agent/src/core/backgroundTools.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,10 @@ import { v4 as uuidv4 } from 'uuid';
// These imports will be used by the cleanup method
import { BrowserManager } from '../tools/browser/BrowserManager.js';
import { agentStates } from '../tools/interaction/agentStart.js';
import { processStates } from '../tools/system/shellStart.js';
import { shellTracker } from '../tools/system/ShellTracker.js';

// Types of background processes we can track
export enum BackgroundToolType {
SHELL = 'shell',
BROWSER = 'browser',
AGENT = 'agent',
}
Expand All @@ -30,17 +29,6 @@ export interface BackgroundTool {
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;
error?: string;
};
}

// Browser process specific data
export interface BrowserBackgroundTool extends BackgroundTool {
type: BackgroundToolType.BROWSER;
Expand All @@ -60,10 +48,7 @@ export interface AgentBackgroundTool extends BackgroundTool {
}

// Utility type for all background tool types
export type AnyBackgroundTool =
| ShellBackgroundTool
| BrowserBackgroundTool
| AgentBackgroundTool;
export type AnyBackgroundTool = BrowserBackgroundTool | AgentBackgroundTool;

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

// Register a new shell process
public registerShell(command: string): string {
const id = uuidv4();
const tool: ShellBackgroundTool = {
id,
type: BackgroundToolType.SHELL,
status: BackgroundToolStatus.RUNNING,
startTime: new Date(),
metadata: {
command,
},
};
this.tools.set(id, tool);
return id;
}

// Register a new browser process
public registerBrowser(url?: string): string {
const id = uuidv4();
Expand Down Expand Up @@ -177,12 +146,6 @@ export class BackgroundTools {
tool.status === BackgroundToolStatus.RUNNING,
);

const shellTools = tools.filter(
(tool): tool is ShellBackgroundTool =>
tool.type === BackgroundToolType.SHELL &&
tool.status === BackgroundToolStatus.RUNNING,
);

const agentTools = tools.filter(
(tool): tool is AgentBackgroundTool =>
tool.type === BackgroundToolType.AGENT &&
Expand All @@ -193,19 +156,15 @@ export class BackgroundTools {
const browserCleanupPromises = browserTools.map((tool) =>
this.cleanupBrowserSession(tool),
);
const shellCleanupPromises = shellTools.map((tool) =>
this.cleanupShellProcess(tool),
);
const agentCleanupPromises = agentTools.map((tool) =>
this.cleanupSubAgent(tool),
);

// Clean up shell processes using ShellTracker
await shellTracker.cleanupAllShells();

// Wait for all cleanup operations to complete in parallel
await Promise.all([
...browserCleanupPromises,
...shellCleanupPromises,
...agentCleanupPromises,
]);
await Promise.all([...browserCleanupPromises, ...agentCleanupPromises]);
}

/**
Expand All @@ -230,38 +189,6 @@ export class BackgroundTools {
}
}

/**
* Cleans up a shell process
* @param tool The shell tool to clean up
*/
private async cleanupShellProcess(tool: ShellBackgroundTool): Promise<void> {
try {
const processState = processStates.get(tool.id);
if (processState && !processState.state.completed) {
processState.process.kill('SIGTERM');

// Force kill after a short timeout if still running
await new Promise<void>((resolve) => {
setTimeout(() => {
try {
if (!processState.state.completed) {
processState.process.kill('SIGKILL');
}
} catch {
// Ignore errors on forced kill
}
resolve();
}, 500);
});
}
this.updateToolStatus(tool.id, BackgroundToolStatus.COMPLETED);
} catch (error) {
this.updateToolStatus(tool.id, BackgroundToolStatus.ERROR, {
error: error instanceof Error ? error.message : String(error),
});
}
}

/**
* Cleans up a sub-agent
* @param tool The agent tool to clean up
Expand Down
2 changes: 2 additions & 0 deletions packages/agent/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ export * from './tools/system/sequenceComplete.js';
export * from './tools/system/shellMessage.js';
export * from './tools/system/shellExecute.js';
export * from './tools/system/listBackgroundTools.js';
export * from './tools/system/listShells.js';
export * from './tools/system/ShellTracker.js';

// Tools - Browser
export * from './tools/browser/BrowserManager.js';
Expand Down
2 changes: 2 additions & 0 deletions packages/agent/src/tools/getTools.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { fetchTool } from './io/fetch.js';
import { textEditorTool } from './io/textEditor.js';
import { createMcpTool } from './mcp.js';
import { listBackgroundToolsTool } from './system/listBackgroundTools.js';
import { listShellsTool } from './system/listShells.js';
import { sequenceCompleteTool } from './system/sequenceComplete.js';
import { shellMessageTool } from './system/shellMessage.js';
import { shellStartTool } from './system/shellStart.js';
Expand Down Expand Up @@ -41,6 +42,7 @@ export function getTools(options?: GetToolsOptions): Tool[] {
//respawnTool as unknown as Tool, this is a confusing tool for now.
sleepTool as unknown as Tool,
listBackgroundToolsTool as unknown as Tool,
listShellsTool as unknown as Tool,
];

// Only include userPrompt tool if enabled
Expand Down
Loading