trace.ai-cli/services/domeHandler.js

A powerful AI-powered CLI tool

const WebSocket = require('ws');
const chalk = require('chalk');
const { processWithAI } = require('./aiService');
const { markdownToAnsi } = require('../utils/markdown');

class DomeHandler {
  constructor(cli) {
    this.cli = cli;
    this._connections = new Map();
  }

  query(url, query) {
    return new Promise((resolve, reject) => {
      const conn = this._getOrCreate(url);

      const timeout = setTimeout(() => {
        conn.pending = null;
        reject(new Error('Query timed out'));
      }, 30000);

      conn.pending = { resolve, reject, timeout, query };

      const send = () => {
        conn.ws.send(JSON.stringify({ type: 'dome', query }));
      };

      if (conn.ws.readyState === WebSocket.OPEN) {
        send();
      } else if (conn.ws.readyState === WebSocket.CONNECTING) {
        conn.ws.once('open', send);
      } else {
        clearTimeout(timeout);
        conn.pending = null;
        this._connections.delete(url);
        this.query(url, query).then(resolve).catch(reject);
      }
    });
  }

  _getOrCreate(url) {
    if (this._connections.has(url)) {
      const conn = this._connections.get(url);
      if (conn.ws.readyState === WebSocket.OPEN || conn.ws.readyState === WebSocket.CONNECTING) {
        return conn;
      }
      this._connections.delete(url);
    }

    const ws = new WebSocket(url);
    const conn = { ws, pending: null };

    ws.on('message', (data) => {
      try {
        const parsed = JSON.parse(data.toString());
        if (parsed.type === 'keepalive') return;
        if (!conn.pending) return;

        if (parsed.type === 'dome_response') {
          const { resolve, timeout, query: origQuery } = conn.pending;
          clearTimeout(timeout);
          conn.pending = null;

          let contextText = '';
          if (parsed.context) contextText += parsed.context + '\n\n';
          if (parsed.results && Array.isArray(parsed.results) && parsed.results.length > 0) {
            contextText += 'Search Results:\n';
            parsed.results.forEach((result, idx) => {
              contextText += `\n[${idx + 1}] ${result.metadata?.tabTitle || 'Result'}\n`;
              if (result.metadata?.url) contextText += `URL: ${result.metadata.url}\n`;
              if (result.text) contextText += `${result.text}\n`;
            });
          }

          const q = parsed.query || origQuery;
          processWithAI(q, contextText, this.cli.aiMode)
            .then(r => resolve(markdownToAnsi(r.text || r)))
            .catch(err => resolve(`Error processing: ${err.message}`));
        }
      } catch (e) {
        // ignore parse errors
      }
    });

    ws.on('close', () => {
      if (conn.pending) {
        clearTimeout(conn.pending.timeout);
        conn.pending.reject(new Error('Connection closed'));
        conn.pending = null;
      }
      this._connections.delete(url);
    });

    ws.on('error', (err) => {
      if (conn.pending) {
        clearTimeout(conn.pending.timeout);
        conn.pending.reject(new Error(`Connection error: ${err.message}`));
        conn.pending = null;
      }
      this._connections.delete(url);
    });

    this._connections.set(url, conn);
    return conn;
  }

  disconnect(url) {
    if (url) {
      const conn = this._connections.get(url);
      if (conn) {
        conn.ws.close();
        this._connections.delete(url);
      }
    } else {
      for (const [, conn] of this._connections) {
        conn.ws.close();
      }
      this._connections.clear();
    }
  }
}

module.exports = DomeHandler;