trace.ai-cli/services/agentHandler.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 AgentHandler {
  constructor(cli) {
    this.cli = cli;
    this.ws = null;
    this._manualDisconnect = false;
  }

  get connected() {
    return this.ws && this.ws.readyState === WebSocket.OPEN;
  }

  async connect(url) {
    if (this.ws) {
      console.log(chalk.yellow('Already connected! Disconnect first using /exit.'));
      return;
    }

    try {
      console.log(chalk.blueBright(`Connecting to ${url}...`));
      this.ws = new WebSocket(url);

      this.ws.on('open', () => {
        this.cli.safeLog(
          chalk.green(' Connected to Keylink agent successfully!') + '\n' +
          chalk.gray('Type your task or @trained agent. Type /exit to disconnect.')
        );
      });

      this.ws.on('message', (data) => this._handleMessage(data));

      this.ws.on('close', () => {
        if (!this._manualDisconnect) {
          this.cli.safeLog(chalk.yellow('\nAgent connection closed.'));
        }
        this.ws = null;
        this._manualDisconnect = false;
      });

      this.ws.on('error', (err) => {
        this.cli.safeLog(chalk.red(`\nKeylink Connection Error: ${err.message}`));
        this.ws = null;
      });
    } catch (error) {
      this.cli.displayError(`Failed to connect: ${error.message}`);
      this.ws = null;
    }
  }

  send(query) {
    if (!this.connected) return;
    this.ws.send(JSON.stringify({ type: 'agent', query }));
  }

  disconnect() {
    if (this.ws) {
      this._manualDisconnect = true;
      this.ws.close();
      this.ws = null;
      this.cli._agentCache = [];
      this.cli.safeLog(chalk.green(' Disconnected from agent.'));
    }
  }

  _handleMessage(data) {
    try {
      const parsed = JSON.parse(data.toString());

      if (parsed.type === 'keepalive') return;

      if (parsed.type === 'agent_response') {
        this._handleAgentResponse(parsed);
      } else if (parsed.type === 'AI_AGENT_DONE' && parsed.summary) {
        const formatted = markdownToAnsi(parsed.summary);
        this.cli.displayResult('Agent Complete', formatted, null);
      } else if (parsed.type === 'AI_AGENT_STEP' && parsed.action) {
        const stepInfo = chalk.cyan(`[Step ${parsed.step}/${parsed.maxSteps}]`);
        const action = chalk.yellow(parsed.action.action);
        this.cli.safeLog(`\n${stepInfo} ${action}: ${chalk.gray(JSON.stringify(parsed.action, null, 2))}`);
      } else if (parsed.type === 'AI_AGENT_STATUS') {
        const status = parsed.status === 'started' ? chalk.green(' Started') : chalk.yellow(parsed.status);
        this.cli.safeLog(`\n${chalk.cyan('[Status]')} ${status}${parsed.query ? ': ' + chalk.white(parsed.query) : ''}`);
      } else if (parsed.type === 'AI_AGENT_VALIDATION') {
        const success = parsed.validation?.success ? chalk.green('') : chalk.red('');
        this.cli.safeLog(`\n${chalk.cyan('[Validation]')} ${success} ${chalk.gray(parsed.validation?.reason || '')}`);
      } else {
        const marker = chalk.yellow('[Agent Response]:');
        this.cli.safeLog(`\n${marker} ${chalk.white(data.toString())}`);
      }
    } catch (e) {
      const marker = chalk.yellow('[Agent Response]:');
      this.cli.safeLog(`\n${marker} ${chalk.white(data.toString())}`);
    }
  }

  _handleAgentResponse(parsed) {
    if (parsed.answer) {
      const formatted = markdownToAnsi(parsed.answer);
      this.cli.displayResult('Trace.Ai Response', formatted, null);
      return;
    }

    if (parsed.context || parsed.results) {
      const frames = ['', '', '', '', '', '', '', '', '', ''];
      let frameIndex = 0;

      const spinner = setInterval(() => {
        process.stdout.write(`\r${chalk.blueBright(frames[frameIndex])} ${chalk.white('Processing your query')}`);
        frameIndex = (frameIndex + 1) % frames.length;
      }, 100);

      let contextText = '';
      if (parsed.context) {
        contextText += parsed.context + '\n\n';
      }
      if (parsed.results && Array.isArray(parsed.results)) {
        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`;
        });
      }

      (async () => {
        try {
          const query = parsed.query || 'Analyze this information';
          const aiResponse = await processWithAI(query, contextText, this.cli.aiMode);

          clearInterval(spinner);
          process.stdout.write(`\r${chalk.green('')} ${chalk.white('Processing your query')} ${chalk.gray('- Complete!')}\n`);

          const formatted = markdownToAnsi(aiResponse.text || aiResponse);
          this.cli.displayResult('Trace.Ai Response', formatted, null);
        } catch (error) {
          clearInterval(spinner);
          process.stdout.write(`\r${chalk.red('')} ${chalk.white('Processing your query')} ${chalk.gray('- Failed!')}\n`);
          this.cli.displayError(error.message);
        }
      })();
      return;
    }

    if (parsed.message) {
      if (!parsed.message.includes('Found') && !parsed.message.includes('context')) {
        this.cli.safeLog('\n' + chalk.gray(parsed.message));
      }
    }
  }
}

module.exports = AgentHandler;