Decouple Browser Session Tracking from backgroundTools

[bhouston]

Decouple Browser Session 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 browser session tracking into its own dedicated tracker.

Proposed Changes

1. Create a new BrowserTracker class in packages/agent/src/tools/browser/browserTracker.ts that will be responsible for tracking browser sessions
2. Move browser-specific tracking logic from backgroundTools.ts to this new class
3. Implement a new listBrowsers tool that uses the dedicated BrowserTracker
4. Update browseStart.ts and browseMessage.ts to use the new tracker instead of backgroundTools
5. Remove browser-related code from backgroundTools.ts

Technical Details

• The new BrowserTracker should work alongside the existing browserSessions map to provide status tracking
• It should provide methods for registering, updating, and retrieving browser sessions
• The listBrowsers tool should have similar filtering capabilities as listBackgroundTools but focused only on browser sessions
• 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 browser-specific functionality
• Clearer separation of concerns

Acceptance Criteria

• BrowserTracker class is implemented
• listBrowsers tool is implemented
• browseStart and browseMessage are updated to use the new tracker
• Browser-specific code is removed from backgroundTools.ts
• All tests pass
• Documentation is updated

[bhouston]

Technical Implementation Details for Browser Tracker

Current Implementation Analysis

• Currently, browser sessions are tracked in two places:

  1. browserSessions map in types.ts tracks browser-specific objects like browser instances and pages
  2. backgroundTools registry provides a higher-level status tracking and metadata storage

• The coupling happens because:

  • browseStart registers browser sessions with backgroundTools
  • browseStart updates backgroundTools when browser state changes
  • backgroundTools.cleanup() references BrowserManager directly to close sessions

Proposed BrowserTracker Implementation

// browserTracker.ts
import { v4 as uuidv4 } from 'uuid';
import { Browser, Page } from '@playwright/test';

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

export interface BrowserSession {
  id: string;
  status: BrowserSessionStatus;
  startTime: Date;
  endTime?: Date;
  url?: string;
  contentLength?: number;
  error?: string;
}

// Internal browser session state tracking (similar to existing browserSessions)
export interface BrowserSessionState {
  id: string;
  browser: Browser;
  page: Page;
}

export class BrowserTracker {
  private sessions: Map<string, BrowserSession> = new Map();
  private sessionStates: Map<string, BrowserSessionState> = new Map();

  constructor(readonly ownerName: string) {}

  // Register a new browser session
  public registerSession(url?: string): string {
    const id = uuidv4();
    
    // Create browser session tracking entry
    const session: BrowserSession = {
      id,
      status: BrowserSessionStatus.RUNNING,
      startTime: new Date(),
      url,
    };
    
    this.sessions.set(id, session);
    return id;
  }

  // Register browser state objects
  public registerSessionState(id: string, browser: Browser, page: Page): void {
    this.sessionStates.set(id, { id, browser, page });
  }

  // Update session status
  public updateSessionStatus(
    id: string,
    status: BrowserSessionStatus,
    metadata?: { url?: string; contentLength?: number; error?: string }
  ): boolean {
    const session = this.sessions.get(id);
    if (!session) {
      return false;
    }

    session.status = status;

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

    if (metadata) {
      if (metadata.url !== undefined) session.url = metadata.url;
      if (metadata.contentLength !== undefined) session.contentLength = metadata.contentLength;
      if (metadata.error !== undefined) session.error = metadata.error;
    }

    return true;
  }

  // Get a specific session state
  public getSessionState(id: string): BrowserSessionState | undefined {
    return this.sessionStates.get(id);
  }

  // Get a specific session tracking info
  public getSession(id: string): BrowserSession | undefined {
    return this.sessions.get(id);
  }

  // Get all sessions with optional filtering
  public getSessions(status?: BrowserSessionStatus): BrowserSession[] {
    if (!status) {
      return Array.from(this.sessions.values());
    }
    
    return Array.from(this.sessions.values()).filter(
      session => session.status === status
    );
  }

  // Cleanup and close sessions
  public async cleanup(): Promise<void> {
    const runningSessions = this.getSessions(BrowserSessionStatus.RUNNING);
    
    await Promise.all(
      runningSessions.map(session => this.closeSession(session.id))
    );
  }

  // Close a specific session
  public async closeSession(id: string): Promise<void> {
    try {
      const sessionState = this.sessionStates.get(id);
      if (sessionState) {
        await sessionState.page.close().catch(() => {});
        await sessionState.browser.close().catch(() => {});
        this.sessionStates.delete(id);
      }
      this.updateSessionStatus(id, BrowserSessionStatus.COMPLETED);
    } catch (error) {
      this.updateSessionStatus(id, BrowserSessionStatus.ERROR, {
        error: error instanceof Error ? error.message : String(error),
      });
    }
  }
}

// Create a singleton instance
export const browserTracker = new BrowserTracker('global');

Implementation of listBrowsers Tool

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

import { BrowserSessionStatus, browserTracker } from './browserTracker.js';
import { Tool } from '../../core/types.js';

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

const returnSchema = z.object({
  sessions: z.array(
    z.object({
      id: z.string(),
      status: z.string(),
      url: z.string().optional(),
      startTime: z.string(),
      endTime: z.string().optional(),
      runtime: z.number().describe('Runtime in seconds'),
      contentLength: z.number().optional(),
      error: z.string().optional(),
    }),
  ),
  count: z.number(),
});

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

export const listBrowsersTool: Tool<Parameters, ReturnType> = {
  name: 'listBrowsers',
  description: 'Lists all browser sessions 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 browser sessions with status: ${status}, verbose: ${verbose}`,
    );

    // Get all sessions
    let sessions = browserTracker.getSessions();
    
    // Filter by status if specified
    if (status !== 'all') {
      const statusEnum = status.toUpperCase() as keyof typeof BrowserSessionStatus;
      sessions = sessions.filter(
        session => session.status === BrowserSessionStatus[statusEnum]
      );
    }

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

      return {
        id: session.id,
        status: session.status,
        ...(session.url && { url: session.url }),
        startTime: startTime.toISOString(),
        ...(session.endTime && { endTime: session.endTime.toISOString() }),
        runtime: parseFloat(runtime.toFixed(2)),
        ...(session.contentLength !== undefined && { contentLength: session.contentLength }),
        ...(verbose && session.error && { error: session.error }),
      };
    });

    return {
      sessions: formattedSessions,
      count: formattedSessions.length,
    };
  },

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

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

Integration with BrowserManager

The existing BrowserManager class should be updated to use the new BrowserTracker. Here's how to integrate them:

// BrowserManager.js
import { browserTracker, BrowserSessionStatus } from './browserTracker.js';

export class BrowserManager {
  // ... existing code ...

  // Update closeSession to use browserTracker
  async closeSession(sessionId: string): Promise<void> {
    try {
      await browserTracker.closeSession(sessionId);
    } catch (error) {
      this.logger.error(`Error closing session ${sessionId}: ${error}`);
    }
  }

  // Update cleanup to use browserTracker
  async cleanup(): Promise<void> {
    await browserTracker.cleanup();
  }
}

Changes Required in browseStart.ts

1. Update the session registration to use browserTracker.registerSession()
2. Update the browser session state registration to use browserTracker.registerSessionState()
3. Update status handling to use browserTracker.updateSessionStatus()
4. Make sure the existing browserSessions map is still maintained for backward compatibility

Migration Strategy

1. Implement the BrowserTracker class first
2. Add the listBrowsers tool
3. Modify browseStart to use both the old backgroundTools and new browserTracker temporarily to ensure compatibility
4. Update BrowserManager to use browserTracker for session management
5. Once all tests pass, remove the backgroundTools usage from browseStart
6. Update backgroundTools.cleanup() to delegate browser session cleanup to browserTracker
7. Finally, remove browser-specific code from backgroundTools

Handling the Relationship with browserSessions

The existing browserSessions map is used directly in various places. We have two options:

1. Keep both: Maintain both browserSessions and browserTracker.sessionStates in sync
2. Refactor browserSessions: Make browserSessions a getter that accesses browserTracker.sessionStates

The first approach is safer for backward compatibility but requires more maintenance. The second approach is cleaner but requires more upfront changes.

[bhouston]

/mycoder can you create a PR that realizes this issue?

[bhouston]

I've created a PR that implements the decoupling of browser session tracking from backgroundTools as requested:

PR: #297

The implementation follows the approach outlined in your comment, with a dedicated BrowserTracker class and a new listBrowsers tool. The changes maintain backward compatibility while providing a cleaner separation of concerns between different types of background tools.

All tests are passing, indicating that the refactoring maintains the expected functionality.

Let me know if you'd like any adjustments to the implementation!

[bhouston]

This is completed.