chrome-ai-bridge 2.3.9 → 2.4.0

package/build/extension/relay-server.ts

@@ -0,0 +1,539 @@
+/**
+ * RelayServer - WebSocket server for Extension communication
+ */
+
+import WebSocket, { WebSocketServer } from 'ws';
+import { EventEmitter } from 'events';
+import crypto from 'crypto';
+import http from 'http';
+import fs from 'fs';
+
+// デバッグログをファイルに出力
+const DEBUG_LOG_PATH = '/tmp/relay-server-debug.log';
+function debugLog(...args: any[]) {
+  const timestamp = new Date().toISOString();
+  const message = `[${timestamp}] ${args.map(a => typeof a === 'object' ? JSON.stringify(a) : String(a)).join(' ')}\n`;
+  fs.appendFileSync(DEBUG_LOG_PATH, message);
+}
+
+export interface RelayServerOptions {
+  port?: number; // 0 for auto-assign
+  host?: string;
+  token?: string; // Authentication token
+  sessionId?: string;
+}
+
+export interface CDPCommand {
+  id: number;
+  method: string;
+  params?: any;
+}
+
+export interface CDPEvent {
+  method: string;
+  params?: any;
+}
+
+export class RelayServer extends EventEmitter {
+  private wss: WebSocketServer | null = null;
+  private ws: WebSocket | null = null; // Single connection (1 tab per server)
+  private port: number = 0;
+  private host: string;
+  private token: string;
+  private sessionId: string;
+  private instanceId: string;
+  private startedAt: number;
+  private tabId: number | null = null;
+  private ready: boolean = false;
+  private nextId = 1;
+  private pending = new Map<number, {
+    resolve: (value: any) => void;
+    reject: (err: Error) => void;
+    method: string;
+    startedAt: number;
+  }>();
+  private discoveryServer: http.Server | null = null;
+  private discoveryPort: number | null = null;
+  private keepAliveTimer: ReturnType<typeof setInterval> | null = null;
+  private _lastDiscoveryOptions: {
+    tabUrl?: string;
+    tabId?: number;
+    newTab?: boolean;
+    allowTabTakeover?: boolean;
+  } = {};
+
+  constructor(options: RelayServerOptions = {}) {
+    super();
+    this.host = options.host || '127.0.0.1';
+    this.token = options.token || this.generateToken();
+    this.sessionId = options.sessionId || this.generateSessionId();
+    this.instanceId = crypto.randomUUID();
+    this.startedAt = Date.now();
+    this.port = options.port || 0;
+  }
+
+  /**
+   * Start WebSocket server
+   */
+  async start(): Promise<number> {
+    return new Promise((resolve, reject) => {
+      this.wss = new WebSocketServer({
+        host: this.host,
+        port: this.port
+      });
+
+      this.wss.on('listening', () => {
+        const address = this.wss!.address() as WebSocket.AddressInfo;
+        this.port = address.port;
+        debugLog(`[RelayServer] Listening on ws://${this.host}:${this.port}`);
+        resolve(this.port);
+      });
+
+      this.wss.on('error', (error) => {
+        debugLog('[RelayServer] Server error:', error);
+        reject(error);
+      });
+
+      this.wss.on('connection', (ws, req) => {
+        this.handleConnection(ws, req);
+      });
+    });
+  }
+
+  /**
+   * Handle WebSocket connection from Extension
+   */
+  private handleConnection(ws: WebSocket, req: any) {
+    debugLog('[RelayServer] New connection from Extension');
+
+    // Validate token
+    const url = new URL(req.url || '', `ws://${this.host}`);
+    const clientToken = url.searchParams.get('token');
+    const clientSessionId = url.searchParams.get('sid');
+
+    if (clientToken !== this.token) {
+      debugLog('[RelayServer] Invalid token');
+      ws.close(1008, 'Invalid token');
+      return;
+    }
+    if (clientSessionId && clientSessionId !== this.sessionId) {
+      debugLog('[RelayServer] Invalid session id', {expected: this.sessionId, received: clientSessionId});
+      ws.close(1008, 'Invalid session id');
+      return;
+    }
+
+    // Only allow one connection
+    if (this.ws) {
+      debugLog('[RelayServer] Connection already exists');
+      ws.close(1008, 'Connection already exists');
+      return;
+    }
+
+    this.ws = ws;
+    this.startKeepAlive();
+
+    ws.on('message', (data) => {
+      this.handleMessage(data.toString());
+    });
+
+    // Guard: only update state if this socket is still the current one.
+    // Prevents a stale socket's close event from corrupting a newer connection.
+    ws.on('close', () => {
+      if (this.ws !== ws) {
+        debugLog('[RelayServer] Stale socket closed (ignored — already replaced)');
+        return;
+      }
+      debugLog('[RelayServer] Extension disconnected');
+      this.stopKeepAlive();
+      this.rejectPendingRequests(
+        new Error('RELAY_DISCONNECTED: Extension socket closed before request completion'),
+      );
+      this.ws = null;
+      this.ready = false;
+      this.emit('disconnected');
+    });
+
+    ws.on('error', (error) => {
+      debugLog('[RelayServer] WebSocket error:', error);
+    });
+
+    debugLog('[RelayServer] Extension connected');
+  }
+
+  private rejectPendingRequests(error: Error): void {
+    if (this.pending.size === 0) return;
+    const pendingEntries = Array.from(this.pending.entries());
+    this.pending.clear();
+    for (const [id, pending] of pendingEntries) {
+      debugLog('[RelayServer] Rejecting pending request', {
+        id,
+        method: pending.method,
+        startedAt: pending.startedAt,
+        reason: error.message,
+      });
+      pending.reject(error);
+    }
+  }
+
+  /**
+   * Handle message from Extension
+   */
+  private handleMessage(data: string) {
+    try {
+      const message = JSON.parse(data);
+
+      if (typeof message.id === 'number' && (message.result !== undefined || message.error !== undefined)) {
+        const pending = this.pending.get(message.id);
+        if (pending) {
+          this.pending.delete(message.id);
+          if (message.error) {
+            const error =
+              typeof message.error === 'string'
+                ? new Error(message.error)
+                : new Error(message.error.message || 'Unknown error');
+            pending.reject(error);
+          } else {
+            pending.resolve(message.result);
+          }
+          return;
+        }
+
+        if (message.error) {
+          const error =
+            typeof message.error === 'string'
+              ? message.error
+              : message.error.message || 'Unknown error';
+          this.emit('cdp-error', { id: message.id, error });
+        } else {
+          this.emit('cdp-result', { id: message.id, result: message.result });
+        }
+        return;
+      }
+
+      if (message?.method === 'forwardCDPEvent' && message.params) {
+        this.emit('cdp-event', {
+          method: message.params.method,
+          params: message.params.params,
+          sessionId: message.params.sessionId,
+        });
+        return;
+      }
+
+      switch (message.type) {
+        case 'ready':
+          this.tabId = message.tabId;
+          this.ready = true;
+          debugLog(`[RelayServer] Connection ready for tab ${this.tabId}`);
+          this.emit('ready', this.tabId);
+          // Release discovery port after WebSocket is established.
+          // 1-second grace period lets Extension finish processing the response.
+          setTimeout(() => this.stopDiscoveryServer(), 1000);
+          break;
+        case 'pong':
+          debugLog('[RelayServer] Received keep-alive pong');
+          break;
+
+        case 'forwardCDPResult':
+          this.emit('cdp-result', { id: message.id, result: message.result });
+          break;
+
+        case 'forwardCDPError':
+          this.emit('cdp-error', { id: message.id, error: message.error });
+          break;
+
+        case 'forwardCDPEvent':
+          this.emit('cdp-event', {
+            method: message.method,
+            params: message.params
+          });
+          break;
+
+        case 'detached':
+          debugLog(`[RelayServer] Tab ${message.tabId} detached: ${message.reason}`);
+          this.emit('detached', message.reason);
+          break;
+
+        default:
+          debugLog('[RelayServer] Unknown message type:', message.type);
+      }
+    } catch (error) {
+      debugLog('[RelayServer] Failed to parse message:', error);
+    }
+  }
+
+  /**
+   * Send CDP command to Extension
+   */
+  sendCDPCommand(id: number, method: string, params?: any): void {
+    if (!this.ws || !this.ready) {
+      throw new Error('Extension not connected or not ready');
+    }
+
+    this.ws.send(JSON.stringify({
+      type: 'forwardCDPCommand',
+      id,
+      method,
+      params
+    }));
+  }
+
+  sendMessage(message: any): void {
+    if (!this.ws || !this.ready) {
+      throw new Error(
+        `Extension not connected or not ready (connected=${Boolean(this.ws)}, ready=${this.ready})`,
+      );
+    }
+    if (this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error('WebSocket not open');
+    }
+    this.ws.send(JSON.stringify(message));
+  }
+
+  async sendRequest(method: string, params?: any, timeoutMs = 30000): Promise<any> {
+    if (!this.ws || !this.ready) {
+      throw new Error(
+        `Extension not connected or not ready (method=${method}, connected=${Boolean(this.ws)}, ready=${this.ready})`,
+      );
+    }
+    if (this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error('WebSocket not open');
+    }
+    const id = this.nextId++;
+    const payload = {id, method, params};
+    const startedAt = Date.now();
+    const response = new Promise<any>((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        this.pending.delete(id);
+        reject(new Error(`RELAY_REQUEST_TIMEOUT: method=${method} timeoutMs=${timeoutMs}`));
+      }, timeoutMs);
+      timeoutId.unref();
+      this.pending.set(id, {
+        resolve: (value: any) => {
+          clearTimeout(timeoutId);
+          debugLog(`[RelayServer] Request success: ${method}`, {id, elapsedMs: Date.now() - startedAt});
+          resolve(value);
+        },
+        reject: (err: Error) => {
+          clearTimeout(timeoutId);
+          debugLog(`[RelayServer] Request failed: ${method}`, {id, elapsedMs: Date.now() - startedAt, error: err.message});
+          reject(err);
+        },
+        method,
+        startedAt,
+      });
+    });
+    try {
+      this.ws.send(JSON.stringify(payload));
+      debugLog(`[RelayServer] Request sent: ${method}`, {id});
+    } catch (error) {
+      this.pending.delete(id);
+      throw error;
+    }
+    return response;
+  }
+
+  /**
+   * Start simple discovery HTTP server for extension to find relay URL.
+   * Extension polls this endpoint when user clicks the extension icon.
+   */
+  async startDiscoveryServer(options: {
+    tabUrl?: string;
+    tabId?: number;
+    newTab?: boolean;
+    allowTabTakeover?: boolean;
+  } = {}): Promise<number | null> {
+    this._lastDiscoveryOptions = options;
+    const ports = [38765, 38766, 38767, 38768, 38769, 38770, 38771, 38772, 38773, 38774, 38775];
+    const wsUrl = this.getConnectionURL();
+
+    for (const port of ports) {
+      const started = await new Promise<boolean>((resolve) => {
+        const server = http.createServer(async (req, res) => {
+          res.setHeader('Access-Control-Allow-Origin', '*');
+
+          if (req.method === 'GET' && req.url === '/relay-info') {
+            res.setHeader('Content-Type', 'application/json');
+            res.end(JSON.stringify({
+              wsUrl,
+              tabUrl: options.tabUrl || null,
+              tabId: options.tabId ?? null,
+              newTab: Boolean(options.newTab),
+              allowTabTakeover: Boolean(options.allowTabTakeover),
+              sessionId: this.sessionId,
+              startedAt: this.startedAt,
+              instanceId: this.instanceId,
+              expiresAt: Date.now() + 60000,
+            }));
+            return;
+          }
+
+          if (req.method === 'POST' && req.url === '/reload-extension') {
+            res.setHeader('Content-Type', 'application/json');
+            if (!this.ws || !this.ready) {
+              res.statusCode = 503;
+              res.end(JSON.stringify({ error: 'Extension not connected' }));
+              return;
+            }
+            try {
+              await this.sendRequest('reloadExtension');
+              res.end(JSON.stringify({ success: true }));
+            } catch (err: any) {
+              // Extension reloads and drops connection - this is expected
+              res.end(JSON.stringify({ success: true, note: 'Extension reloading' }));
+            }
+            return;
+          }
+
+          res.statusCode = 404;
+          res.end('Not Found');
+        });
+
+        server.on('error', (error: any) => {
+          if (error?.code === 'EADDRINUSE') {
+            resolve(false);
+            return;
+          }
+          debugLog('[RelayServer] Discovery server error:', error);
+          resolve(false);
+        });
+
+        server.listen(port, this.host, () => {
+          this.discoveryServer = server;
+          this.discoveryPort = port;
+          debugLog(`[RelayServer] Discovery available on http://${this.host}:${port}/relay-info`);
+          resolve(true);
+        });
+      });
+
+      if (started) {
+        return port;
+      }
+    }
+
+    debugLog('[RelayServer] Could not start discovery server on any port');
+    return null;
+  }
+
+  /**
+   * Release the discovery HTTP server (port).
+   * Called automatically after the Extension WebSocket connects (ready event).
+   * The port becomes available for other sessions.
+   */
+  stopDiscoveryServer(): void {
+    if (this.discoveryServer) {
+      this.discoveryServer.close();
+      debugLog(`[RelayServer] Discovery server released (port ${this.discoveryPort})`);
+      this.discoveryServer = null;
+      this.discoveryPort = null;
+    }
+  }
+
+  /**
+   * Re-acquire a discovery port (e.g. after WebSocket disconnect for reconnection).
+   * Uses the same options as the last startDiscoveryServer() call.
+   */
+  async restartDiscoveryServer(options?: {
+    tabUrl?: string;
+    tabId?: number;
+    newTab?: boolean;
+    allowTabTakeover?: boolean;
+  }): Promise<number | null> {
+    this.stopDiscoveryServer();
+    return this.startDiscoveryServer(options || this._lastDiscoveryOptions);
+  }
+
+  /**
+   * Stop server
+   */
+  async stop(): Promise<void> {
+    this.stopKeepAlive();
+
+    if (this.ws) {
+      try {
+        this.ws.close();
+      } catch {
+        // ignore close errors
+      }
+      this.ws = null;
+    }
+    this.ready = false;
+    this.tabId = null;
+
+    this.rejectPendingRequests(
+      new Error('RELAY_STOPPED: Relay stopped before request completion'),
+    );
+
+    this.stopDiscoveryServer();
+
+    if (this.wss) {
+      return new Promise((resolve) => {
+        this.wss!.close(() => {
+          this.wss = null;
+          debugLog('[RelayServer] Server stopped');
+          resolve();
+        });
+      });
+    }
+  }
+
+  /**
+   * Start keep-alive ping to prevent Service Worker from sleeping
+   */
+  private startKeepAlive(): void {
+    this.stopKeepAlive();
+    this.keepAliveTimer = setInterval(() => {
+      if (this.ws?.readyState === WebSocket.OPEN) {
+        this.ws.send(JSON.stringify({ type: 'ping' }));
+        debugLog('[RelayServer] Sent keep-alive ping');
+      }
+    }, 30000); // 30 seconds
+    debugLog('[RelayServer] Keep-alive started');
+  }
+
+  /**
+   * Stop keep-alive ping
+   */
+  private stopKeepAlive(): void {
+    if (this.keepAliveTimer) {
+      clearInterval(this.keepAliveTimer);
+      this.keepAliveTimer = null;
+      debugLog('[RelayServer] Keep-alive stopped');
+    }
+  }
+
+  /**
+   * Generate random token
+   */
+  private generateToken(): string {
+    return crypto.randomBytes(32).toString('hex');
+  }
+
+  getPort(): number {
+    return this.port;
+  }
+
+  getToken(): string {
+    return this.token;
+  }
+
+  getTabId(): number | null {
+    return this.tabId;
+  }
+
+  isReady(): boolean {
+    return this.ready;
+  }
+
+
+  getConnectionURL(): string {
+    return `ws://${this.host}:${this.port}?token=${this.token}&sid=${encodeURIComponent(this.sessionId)}`;
+  }
+
+  getSessionId(): string {
+    return this.sessionId;
+  }
+
+  private generateSessionId(): string {
+    return crypto.randomBytes(16).toString('hex');
+  }
+}