Decouple Shell Process Tracking from backgroundTools

[bhouston]

Decouple Shell Process Tracking from backgroundTools

Description

Currently, the backgroundTools tracker is used to track shell processes, browser sessions, and agents in a single implementation. This creates unnecessary coupling between these different types of tools. This issue proposes to refactor the shell process tracking into its own dedicated tracker.

Proposed Changes

1. Create a new ShellTracker class in packages/agent/src/tools/system/shellTracker.ts that will be responsible for tracking shell processes
2. Move shell-specific tracking logic from backgroundTools.ts to this new class
3. Implement a new listShells tool that uses the dedicated ShellTracker
4. Update shellStart.ts and shellMessage.ts to use the new tracker instead of backgroundTools
5. Remove shell-related code from backgroundTools.ts

Technical Details

• The new ShellTracker should maintain a registry of shell processes similar to how processStates works now
• It should provide methods for registering, updating, and retrieving shell processes
• The listShells tool should have similar filtering capabilities as listBackgroundTools but focused only on shell processes
• The implementation should maintain backward compatibility where possible

Benefits

• Reduced coupling between different tool types
• More focused and maintainable code
• Easier to extend or modify shell-specific functionality
• Clearer separation of concerns

Acceptance Criteria

• ShellTracker class is implemented
• listShells tool is implemented
• shellStart and shellMessage are updated to use the new tracker
• Shell-specific code is removed from backgroundTools.ts
• All tests pass
• Documentation is updated

[bhouston]

Technical Implementation Details for Shell Tracker

Current Implementation Analysis

• Currently, shell processes are tracked in two places:

  1. processStates map in shellStart.ts tracks process-specific details like stdout/stderr buffers and process objects
  2. backgroundTools registry provides a higher-level status tracking and metadata storage

• The coupling happens because:

  • shellStart registers processes with backgroundTools
  • shellStart updates backgroundTools when process state changes
  • backgroundTools.cleanup() references processStates directly to kill processes

Proposed ShellTracker Implementation

// shellTracker.ts
import { v4 as uuidv4 } from 'uuid';
import type { ChildProcess } from 'child_process';

export enum ShellProcessStatus {
  RUNNING = 'running',
  COMPLETED = 'completed',
  ERROR = 'error',
  TERMINATED = 'terminated',
}

export interface ShellProcess {
  id: string;
  status: ShellProcessStatus;
  startTime: Date;
  endTime?: Date;
  command: string;
  exitCode?: number | null;
  signaled?: boolean;
  error?: string;
}

// Internal process state tracking
export type ProcessState = {
  process: ChildProcess;
  command: string;
  stdout: string[];
  stderr: string[];
  state: {
    completed: boolean;
    signaled: boolean;
    exitCode: number | null;
  };
  showStdIn: boolean;
  showStdout: boolean;
};

export class ShellTracker {
  private processes: Map<string, ShellProcess> = new Map();
  private processStates: Map<string, ProcessState> = new Map();

  constructor(readonly ownerName: string) {}

  // Register a new shell process
  public registerProcess(command: string, process: ChildProcess, options: { showStdIn: boolean, showStdout: boolean }): string {
    const id = uuidv4();
    
    // Create shell process tracking entry
    const shellProcess: ShellProcess = {
      id,
      status: ShellProcessStatus.RUNNING,
      startTime: new Date(),
      command,
    };
    
    // Create process state entry
    const processState: ProcessState = {
      process,
      command,
      stdout: [],
      stderr: [],
      state: { completed: false, signaled: false, exitCode: null },
      showStdIn: options.showStdIn,
      showStdout: options.showStdout,
    };
    
    this.processes.set(id, shellProcess);
    this.processStates.set(id, processState);
    
    return id;
  }

  // Update process status
  public updateProcessStatus(
    id: string,
    status: ShellProcessStatus,
    metadata?: { exitCode?: number | null; signaled?: boolean; error?: string }
  ): boolean {
    const process = this.processes.get(id);
    if (!process) {
      return false;
    }

    process.status = status;

    if (
      status === ShellProcessStatus.COMPLETED ||
      status === ShellProcessStatus.ERROR ||
      status === ShellProcessStatus.TERMINATED
    ) {
      process.endTime = new Date();
    }

    if (metadata) {
      if (metadata.exitCode !== undefined) process.exitCode = metadata.exitCode;
      if (metadata.signaled !== undefined) process.signaled = metadata.signaled;
      if (metadata.error !== undefined) process.error = metadata.error;
    }

    return true;
  }

  // Get a specific process state
  public getProcessState(id: string): ProcessState | undefined {
    return this.processStates.get(id);
  }

  // Get a specific process tracking info
  public getProcess(id: string): ShellProcess | undefined {
    return this.processes.get(id);
  }

  // Get all processes with optional filtering
  public getProcesses(status?: ShellProcessStatus): ShellProcess[] {
    if (!status) {
      return Array.from(this.processes.values());
    }
    
    return Array.from(this.processes.values()).filter(
      process => process.status === status
    );
  }

  // Add stdout to process buffer
  public addStdout(id: string, data: string): void {
    const processState = this.processStates.get(id);
    if (processState) {
      processState.stdout.push(data);
    }
  }

  // Add stderr to process buffer
  public addStderr(id: string, data: string): void {
    const processState = this.processStates.get(id);
    if (processState) {
      processState.stderr.push(data);
    }
  }

  // Get combined stdout
  public getStdout(id: string): string {
    const processState = this.processStates.get(id);
    return processState ? processState.stdout.join('') : '';
  }

  // Get combined stderr
  public getStderr(id: string): string {
    const processState = this.processStates.get(id);
    return processState ? processState.stderr.join('') : '';
  }

  // Cleanup and terminate processes
  public async cleanup(): Promise<void> {
    const runningProcesses = this.getProcesses(ShellProcessStatus.RUNNING);
    
    await Promise.all(
      runningProcesses.map(process => this.terminateProcess(process.id))
    );
  }

  // Terminate a specific process
  public async terminateProcess(id: string): Promise<void> {
    try {
      const processState = this.processStates.get(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.updateProcessStatus(id, ShellProcessStatus.TERMINATED);
    } catch (error) {
      this.updateProcessStatus(id, ShellProcessStatus.ERROR, {
        error: error instanceof Error ? error.message : String(error),
      });
    }
  }
}

// Create a singleton instance
export const shellTracker = new ShellTracker('global');

Implementation of listShells Tool

// listShells.ts
import { z } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';

import { ShellProcessStatus, shellTracker } from './shellTracker.js';
import { Tool } from '../../core/types.js';

const parameterSchema = z.object({
  status: z
    .enum(['all', 'running', 'completed', 'error', 'terminated'])
    .optional()
    .describe('Filter shell processes by status (default: "all")'),
  verbose: z
    .boolean()
    .optional()
    .describe('Include detailed information about each process (default: false)'),
});

const returnSchema = z.object({
  processes: z.array(
    z.object({
      id: z.string(),
      status: z.string(),
      command: z.string(),
      startTime: z.string(),
      endTime: z.string().optional(),
      runtime: z.number().describe('Runtime in seconds'),
      exitCode: z.number().nullable().optional(),
      signaled: z.boolean().optional(),
      error: z.string().optional(),
      stdout: z.string().optional(),
      stderr: z.string().optional(),
    }),
  ),
  count: z.number(),
});

type Parameters = z.infer<typeof parameterSchema>;
type ReturnType = z.infer<typeof returnSchema>;

export const listShellsTool: Tool<Parameters, ReturnType> = {
  name: 'listShells',
  description: 'Lists all shell processes and their status',
  logPrefix: '💻',
  parameters: parameterSchema,
  returns: returnSchema,
  parametersJsonSchema: zodToJsonSchema(parameterSchema),
  returnsJsonSchema: zodToJsonSchema(returnSchema),

  execute: async (
    { status = 'all', verbose = false },
    { logger },
  ): Promise<ReturnType> => {
    logger.verbose(
      `Listing shell processes with status: ${status}, verbose: ${verbose}`,
    );

    // Get all processes
    let processes = shellTracker.getProcesses();
    
    // Filter by status if specified
    if (status !== 'all') {
      const statusEnum = status.toUpperCase() as keyof typeof ShellProcessStatus;
      processes = processes.filter(
        process => process.status === ShellProcessStatus[statusEnum]
      );
    }

    // Format the response
    const formattedProcesses = processes.map(process => {
      const now = new Date();
      const startTime = process.startTime;
      const endTime = process.endTime || now;
      const runtime = (endTime.getTime() - startTime.getTime()) / 1000; // in seconds

      const result: any = {
        id: process.id,
        status: process.status,
        command: process.command,
        startTime: startTime.toISOString(),
        ...(process.endTime && { endTime: process.endTime.toISOString() }),
        runtime: parseFloat(runtime.toFixed(2)),
        ...(process.exitCode !== undefined && { exitCode: process.exitCode }),
        ...(process.signaled !== undefined && { signaled: process.signaled }),
        ...(process.error && { error: process.error }),
      };

      // Add stdout/stderr if verbose
      if (verbose) {
        result.stdout = shellTracker.getStdout(process.id);
        result.stderr = shellTracker.getStderr(process.id);
      }

      return result;
    });

    return {
      processes: formattedProcesses,
      count: formattedProcesses.length,
    };
  },

  logParameters: ({ status = 'all', verbose = false }, { logger }) => {
    logger.info(
      `Listing shell processes with status: ${status}, verbose: ${verbose}`,
    );
  },

  logReturns: (output, { logger }) => {
    logger.info(`Found ${output.count} shell processes`);
  },
};

Changes Required in shellStart.ts

1. Replace the global processStates map with the shellTracker singleton
2. Update process registration to use shellTracker.registerProcess()
3. Update status handling to use shellTracker.updateProcessStatus()
4. Update stdout/stderr handling to use shellTracker.addStdout() and shellTracker.addStderr()

Migration Strategy

1. Implement the ShellTracker class first
2. Add the listShells tool
3. Modify shellStart to use both the old backgroundTools and new shellTracker temporarily to ensure compatibility
4. Once all tests pass, remove the backgroundTools usage from shellStart
5. Update backgroundTools.cleanup() to delegate shell process cleanup to shellTracker
6. Finally, remove shell-specific code from backgroundTools

[bhouston]

/mycoder can you create a PR That realized this issue?

[bhouston]

Hi @bhouston,

I've implemented the ShellTracker class and decoupled shell process tracking from backgroundTools as requested.

A PR has been created at #296 that addresses all the requirements in this issue:

1. ✅ Created a new ShellTracker class in packages/agent/src/tools/system/ShellTracker.ts
2. ✅ Moved shell-specific tracking logic from backgroundTools.ts to this new class
3. ✅ Implemented a new listShells tool that uses the dedicated ShellTracker
4. ✅ Updated shellStart.ts and shellMessage.ts to use the new tracker instead of backgroundTools
5. ✅ Removed shell-related code from backgroundTools.ts

All tests are passing, and the implementation maintains backward compatibility while providing a cleaner separation of concerns.

Please review the PR when you have a chance. Let me know if you need any changes or have any questions!

[bhouston]

This is completed.