ai-terminal-assistant/packages/cli/src/commands/attach.ts

/**
 * Attach Command
 *
 * 连接到远程 Server
 */

import type { Command } from 'commander';
import chalk from 'chalk';
import ora from 'ora';
import { createClient } from '../client/index.js';

export function registerAttachCommand(program: Command): void {
  program
    .command('attach')
    .description('Connect to a remote AI Assistant server')
    .argument('<url>', 'Server URL (e.g., http://192.168.1.100:3000)')
    .option('-t, --token <token>', 'Authentication token')
    .option('-s, --session <id>', 'Session ID to connect to')
    .option('--tui', 'Use TUI mode (default)', true)
    .option('--no-tui', 'Use simple CLI mode')
    .action(async (url: string, options) => {
      const { token, session: sessionId, tui } = options;

      const spinner = ora('Connecting to server...').start();

      try {
        // 创建客户端
        const client = createClient({
          baseUrl: url,
          token,
        });

        // 检查服务器健康状态
        const health = await client.health();
        spinner.succeed(`Connected to server at ${url}`);

        if (tui) {
          // TUI 模式
          const { TUIApp } = await import('../tui/index.js');
          const app = new TUIApp({
            client,
            sessionId,
          });
          await app.start();
        } else {
          // 简单 CLI 模式
          console.log(chalk.gray('─'.repeat(50)));
          console.log(chalk.bold('Server Status:'));
          console.log(`  Status: ${chalk.green(health.status)}`);
          console.log(`  Agent: ${health.agent.coreAvailable ? chalk.green('Available') : chalk.yellow('Not available')}`);
          console.log(`  Auth: ${health.auth.enabled ? chalk.yellow('Enabled') : chalk.gray('Disabled')}`);
          console.log(`  Sessions: ${health.stats.sessions}`);
          console.log(`  WebSocket: ${health.stats.websocket.connections} connections`);
          console.log(chalk.gray('─'.repeat(50)));

          // 如果指定了 session，连接到该 session
          if (sessionId) {
            await connectToSession(client, sessionId);
          } else {
            // 显示可用 sessions 或创建新的
            await showSessionMenu(client);
          }
        }
      } catch (error) {
        spinner.fail('Failed to connect to server');
        if (error instanceof Error) {
          console.error(chalk.red(`Error: ${error.message}`));
        }
        process.exit(1);
      }
    });
}

async function connectToSession(client: ReturnType<typeof createClient>, sessionId: string): Promise<void> {
  const spinner = ora(`Connecting to session ${sessionId}...`).start();

  try {
    // 获取 session 信息
    const { data: session } = await client.getSession(sessionId);
    spinner.succeed(`Connected to session: ${session.name || session.id}`);

    // 获取历史消息
    const { data: messages } = await client.getMessages(sessionId);

    if (messages.length > 0) {
      console.log(chalk.gray('\n─── Recent Messages ───'));
      const recentMessages = messages.slice(-5);
      for (const msg of recentMessages) {
        const prefix = msg.role === 'user' ? chalk.blue('You: ') : chalk.green('AI: ');
        const content = msg.content.length > 100
          ? msg.content.slice(0, 100) + '...'
          : msg.content;
        console.log(prefix + content);
      }
      console.log(chalk.gray('───────────────────────\n'));
    }

    // 启动交互式会话
    await startInteractiveSession(client, sessionId);
  } catch (error) {
    spinner.fail('Failed to connect to session');
    throw error;
  }
}

async function showSessionMenu(client: ReturnType<typeof createClient>): Promise<void> {
  const { default: inquirer } = await import('inquirer');

  // 获取可用 sessions
  const { data: sessions } = await client.listSessions();

  const choices = [
    { name: chalk.green('+ Create new session'), value: 'new' },
    ...sessions.map((s) => ({
      name: `${s.name || s.id} (${s.messageCount} messages)`,
      value: s.id,
    })),
  ];

  const { selection } = await inquirer.prompt([
    {
      type: 'list',
      name: 'selection',
      message: 'Select a session:',
      choices,
    },
  ]);

  if (selection === 'new') {
    const { name } = await inquirer.prompt([
      {
        type: 'input',
        name: 'name',
        message: 'Session name (optional):',
      },
    ]);

    const { data: newSession } = await client.createSession(name || undefined);
    console.log(chalk.green(`Created new session: ${newSession.id}`));
    await startInteractiveSession(client, newSession.id);
  } else {
    await connectToSession(client, selection);
  }
}

async function startInteractiveSession(
  client: ReturnType<typeof createClient>,
  sessionId: string
): Promise<void> {
  const { createInterface } = await import('readline');

  console.log(chalk.gray('Type your message and press Enter. Type /quit to exit.\n'));

  // 连接 WebSocket
  const ws = client.connectWebSocket(sessionId);

  let currentResponse = '';

  ws.onopen = () => {
    console.log(chalk.gray('[Connected to WebSocket]'));
  };

  ws.onmessage = (event) => {
    try {
      const message = JSON.parse(event.data as string);

      switch (message.type) {
        case 'chunk':
          process.stdout.write(message.payload?.content || '');
          currentResponse += message.payload?.content || '';
          break;
        case 'done':
          if (currentResponse) {
            console.log('\n');
            currentResponse = '';
          }
          break;
        case 'error':
          console.error(chalk.red(`\nError: ${message.payload?.message}`));
          break;
      }
    } catch {
      // 忽略解析错误
    }
  };

  ws.onerror = (error) => {
    console.error(chalk.red('WebSocket error:', error));
  };

  ws.onclose = () => {
    console.log(chalk.gray('[Disconnected from WebSocket]'));
    process.exit(0);
  };

  // 创建交互式 readline
  const rl = createInterface({
    input: process.stdin,
    output: process.stdout,
  });

  const prompt = () => {
    rl.question(chalk.blue('You: '), async (input) => {
      const trimmed = input.trim();

      if (trimmed === '/quit' || trimmed === '/exit') {
        console.log(chalk.gray('Goodbye!'));
        ws.close();
        rl.close();
        return;
      }

      if (!trimmed) {
        prompt();
        return;
      }

      // 通过 WebSocket 发送消息
      ws.send(JSON.stringify({
        type: 'message',
        sessionId,
        payload: { content: trimmed },
      }));

      process.stdout.write(chalk.green('AI: '));

      // 等待响应完成后再提示
      const waitForDone = () => {
        if (currentResponse === '') {
          prompt();
        } else {
          setTimeout(waitForDone, 100);
        }
      };
      setTimeout(waitForDone, 500);
    });
  };

  // 等待 WebSocket 连接
  await new Promise<void>((resolve) => {
    if (ws.readyState === WebSocket.OPEN) {
      resolve();
    } else {
      ws.onopen = () => resolve();
    }
  });

  prompt();
}