Refactor server and session management

- Remove metrics and system info modules
- Simplify browser session handling
- Update Supabase authentication to return user ID
- Modify Chrome arguments and environment configuration
- Update package scripts to use environment file

Author: k9p5

.env.example

@@ -2,11 +2,11 @@
 CHROME_PATH=/usr/bin/google-chrome
 
 # Defines Chrome arguments for security and privacy
-CHROME_ARGS="--no-sandbox --incognito --disable-sync --disable-logging --disable-ftp --disable-password-generation --disable-save-password-bubble --disable-autofill --disable-translate"
+CHROME_ARGS="--no-sandbox --enable-zero-copy"
 
 # Defines the database access credentials
 SUPABASE_SERVICE_ROLE_KEY=xxx
 SUPABASE_URL=xxx
 
-# Defines the GCP project ID
+# Defines the GCP project ID (optional)
 GCP_PROJECT_ID=xxx

capture-gpu.ts

@@ -3,7 +3,7 @@ import puppeteer from 'puppeteer-core';
 (async () => {
   // Connect to the existing Chrome instance via WebSocket
   const browser = await puppeteer.connect({
-    browserWSEndpoint: 'ws://localhost:3000?token=run-npm-run-keygen-to-generate-api-key',
+    browserWSEndpoint: `ws://localhost:3000?token=${process.env.API_KEY}`,
   })
 
   // Create a new page

package-lock.json

@@ -8,12 +8,12 @@
     "dev": "tsx --watch --env-file=.env src/app.ts",
     "serve": "tsx src/app.ts",
     "keygen": "tsx keygen.ts",
-    "gpuinfo": "tsx capture-gpu.ts"
+    "gpuinfo": "tsx --env-file=.env capture-gpu.ts"
   },
   "devDependencies": {
     "@types/express": "^5.0.0",
     "@types/node": "^22.2.0",
-    "@types/ws": "^8.5.13",
+    "@types/ws": "^8.5.14",
     "esbuild": "^0.24.2",
     "tsx": "^4.19.2",
     "typescript": "^5.5.3"

package.json

@@ -5,6 +5,7 @@ import wsserver from "./server";
 import * as sessions from "./routes/sessions";
 import * as health from "./routes/health";
 import * as env from "./environment";
+import './utilization';
 
 export const app = express();
 export const server = http.createServer(app);
@@ -21,5 +22,3 @@ server.on('upgrade', wsserver);
 server.listen(env.PORT, async () => {
   console.log(`Server running on port ${env.PORT}`);
 });
-
-import('./metrics').then(module => module.default());

src/app.ts

@@ -0,0 +1,141 @@
+import puppeteer, { Browser as PuppeteerBrowser, Target } from 'puppeteer-core';
+import { MAX_BROWSER_INSTANCES, CHROME_PATH, CHROME_ARGS } from './environment';
+import { CHROME_SECURITY_FLAGS } from './constants';
+import { report } from './monitoring';
+import { v4 as uuid } from 'uuid';
+import { supabase } from './lib/supabase';
+
+export const browsers = new Map<string, Browser>();
+
+export class Browser {
+  public readonly user?: string | null;
+  public readonly id: string = uuid();
+  public instance?: PuppeteerBrowser;
+  public pages: string[] = [];
+
+  public constructor(user?: string | null) {
+    this.user = user;
+  }
+
+  public async init() {
+    try {
+      const { error } = await supabase
+        .from('browsers')
+        .insert([{ id: this.id, user: this.user }])
+        .single();
+
+      if (error) {
+        console.error('Error creating browser database entry:', error);
+        return;
+      }
+
+      this.instance = await puppeteer.launch({
+        executablePath: CHROME_PATH,
+        headless: true,
+        args: [...CHROME_SECURITY_FLAGS, ...CHROME_ARGS],
+      });
+
+      this.instance.on('targetcreated', this.handlePageCreated.bind(this));
+      console.log('Initialized browser:', this.id);
+    } catch (e) {
+      console.error('Error initializing browser:', e);
+    } finally {
+      browsers.set(this.id, this);
+      this.handleUtilizationReport();
+    }
+  }
+
+  public async close() {
+    if (!this.instance) return;
+
+    try {
+      await this.instance.close();
+      console.log('Closed browser:', this.id);
+    } catch (e) {
+      console.error('Error closing browser:', e);
+    } finally {
+      let error = await supabase
+        .from('browsers')
+        .update({ closed_at: new Date().toISOString() })
+        .is('closed_at', null)
+        .eq('id', this.id)
+        .single()
+        .then(res => res.error);
+
+      if (error) {
+        console.error('Error updating browser closed_at:', error);
+      }
+
+      error = await supabase
+        .from('browser_tabs')
+        .update([{ browser: this.id, closed_at: new Date().toISOString() }])
+        .eq('browser', this.id)
+        .then(res => res.error);
+
+      if (error) {
+        console.error('Error updating browser tabs:', error);
+      }
+
+      browsers.delete(this.id);
+      this.instance = undefined;
+      this.handleUtilizationReport();
+    }
+  }
+
+  private async handlePageCreated(target: Target) {
+    if (target.type() !== 'page') return;
+
+    const page = await target.page();
+
+    const { data, error } = await supabase
+      .from('browser_tabs')
+      .insert([{ browser: this.id }])
+      .select()
+      .single();
+
+    if (error || !data) {
+      console.error('Error creating tab database entry:', error);
+    } else {
+      this.pages.push(data.id);
+    }
+
+    page?.on('close', () => this.handlePageClosed(data.id)());
+  }
+
+  private handlePageClosed(tabId: string) {
+    return async () => {
+      const { error } = await supabase
+        .from('browser_tabs')
+        .update({ closed_at: new Date().toISOString() })
+        .eq('id', tabId)
+        .single();
+
+      if (error) {
+        console.error('Error updating tab closed_at:', error);
+      }
+
+      this.pages = this.pages.filter(id => id !== tabId);
+    }
+  }
+
+  private async handleUtilizationReport() {
+    return await report([{
+      type: 'custom.googleapis.com/browser_utilization',
+      value: browsers.size / MAX_BROWSER_INSTANCES,
+    }]);
+  }
+}
+
+/**
+ * Handle SIGINT, SIGTERM, and SIGQUIT signals
+ */
+['SIGINT', 'SIGTERM', 'SIGQUIT'].forEach(signal => {
+  process.on(signal, () => {
+    console.log(`${signal} received`);
+    for (const browser of browsers.values()) {
+      browser.close();
+    }
+    browsers.clear();
+    process.exit(0);
+  });
+});

src/browser.ts

@@ -0,0 +1,16 @@
+export const CHROME_SECURITY_FLAGS = [
+  '--incognito',                     // Ensures fresh session
+  '--disable-dev-shm-usage',         // Prevents excessive memory use in /dev/shm
+  '--renderer-process-limit=3',      // Limits renderer processes (prevents high RAM usage)
+  '--js-flags="--max_old_space_size=3072"', // Limits JavaScript heap memory to 3GB
+  '--disable-background-networking', // Reduces background memory usage
+  '--disable-sync',                  // Prevents syncing (reduces memory footprint)
+  '--disable-extensions',            // Prevents high-memory extension usage
+  '--disable-password-generation',   // Prevents password generation
+  '--disable-save-password-bubble',  // Prevents saving password bubble
+  '--disable-autofill',              // Prevents autofill
+  '--disable-translate',             // Prevents translation
+  '--disable-ftp',                   // Prevents FTP
+];
+export const METADATA_URL = 'http://metadata.google.internal/computeMetadata/v1';
+export const METADATA_HEADERS = { 'Metadata-Flavor': 'Google' };

src/constants.ts

@@ -4,7 +4,6 @@ export const PORT = process.env.PORT || 3000;
 export const SUPABASE_URL = process.env.SUPABASE_URL;
 export const SUPABASE_SERVICE_ROLE_KEY = process.env.SUPABASE_SERVICE_ROLE_KEY;
 export const CHROME_ARGS = process.env.CHROME_ARGS?.split(' ') || [];
-export const API_KEY = process.env.API_KEY;
 export const CHROME_PATH = process.env.CHROME_PATH;
 export const MAX_BROWSER_INSTANCES = Math.max(os.cpus().length - 1, 1, parseInt(process.env.MAX_BROWSER_INSTANCES || '0'))
 export const GCP_PROJECT_ID = process.env.GCP_PROJECT_ID;
@@ -21,7 +20,3 @@ if (SUPABASE_URL && SUPABASE_SERVICE_ROLE_KEY) {
   console.log('SUPABASE_URL:', SUPABASE_URL);
   console.log('SUPABASE_SERVICE_ROLE_KEY:', SUPABASE_SERVICE_ROLE_KEY.slice(0, 4) + '...');
 }
-
-if (API_KEY) {
-  console.log('API_KEY:', API_KEY.slice(0, 4) + '...');
-}

src/environment.ts

@@ -1,35 +1,32 @@
-import { createClient, SupabaseClient } from '@supabase/supabase-js';
+import { createClient } from '@supabase/supabase-js';
 import * as env from '../environment';
 
-let supabase: SupabaseClient<any, "public", any> | undefined = undefined;
-
-if (env.SUPABASE_URL && env.SUPABASE_SERVICE_ROLE_KEY) {
-  supabase = createClient(
-    env.SUPABASE_URL,
-    env.SUPABASE_SERVICE_ROLE_KEY,
-    {
-      auth: {
-        autoRefreshToken: false,
-        persistSession: false,
-      },
-    }
-  );
-} else if (!env.API_KEY) {
-  console.error('Missing environment variables. API_KEY or SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY are required.');
-}
+if (!env.SUPABASE_URL || !env.SUPABASE_SERVICE_ROLE_KEY) {
+  throw new Error('Missing environment variables. SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY are required.');
+} 
+
+export const supabase = createClient(
+  env.SUPABASE_URL,
+  env.SUPABASE_SERVICE_ROLE_KEY,
+  {
+    auth: {
+      autoRefreshToken: false,
+      persistSession: false,
+    },
+  }
+)
 
-export async function validateApiKey(key: string): Promise<boolean> {
-  if (key === env.API_KEY) return true;
-  if (!supabase) return false;
+export async function getUser(token?: string | null): Promise<string | null> {
+  if (!token) return null;
 
   const { data, error } = await supabase
     .from('api_keys')
     .select('*')
-    .eq('key', key)
+    .eq('key', token)
     .single();
 
   if (error || !data) {
-    return false;
+    return null;
   }
 
   // Update last used timestamp
@@ -38,7 +35,7 @@ export async function validateApiKey(key: string): Promise<boolean> {
     .update({ last_used_at: new Date().toISOString() })
     .eq('id', data.id);
 
-  return true;
+  return data.user;
 }
 
 // export async function generateApiKey(userId: string): Promise<string | null> {