ai-terminal-assistant/packages/server/src/index.ts

/**
 * AI Assistant Server
 *
 * HTTP Server 入口，提供 REST API、WebSocket 和 SSE 支持
 */

import { Hono } from 'hono';
import { cors } from 'hono/cors';
import { logger } from 'hono/logger';
import { createBunWebSocket } from 'hono/bun';

import { sessionsRouter, toolsRouter, configRouter, filesRouter, commandsRouter } from './routes/index.js';
import {
  handleWebSocket,
  handleWebSocketMessage,
  handleWebSocketClose,
  getConnectionStats,
} from './ws.js';
import { handleSSE, getSSEStats } from './sse.js';
import { getSessionManager } from './session/manager.js';
import { initCore, isCoreAvailable, getAgentStats } from './agent/index.js';
import {
  authMiddleware,
  initAuth,
  getAuthConfig,
  generateToken,
  addToken,
  setAuthEnabled,
  maskToken,
} from './auth/index.js';

// 创建 Hono 应用
const app = new Hono();

// WebSocket 升级 (Bun 环境)
const { upgradeWebSocket, websocket } = createBunWebSocket();

// 中间件
app.use('*', logger());
app.use(
  '*',
  cors({
    origin: '*', // 生产环境应该限制
    allowMethods: ['GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS'],
    allowHeaders: ['Content-Type', 'Authorization'],
  })
);

// 认证中间件 (在 CORS 之后)
app.use('*', authMiddleware);

// 健康检查 (跳过认证)
app.get('/health', (c) => {
  const sessionManager = getSessionManager();
  const wsStats = getConnectionStats();
  const sseStats = getSSEStats();
  const authConfig = getAuthConfig();

  return c.json({
    status: 'ok',
    timestamp: new Date().toISOString(),
    agent: {
      coreAvailable: isCoreAvailable(),
    },
    auth: {
      enabled: authConfig.enabled,
      tokenCount: authConfig.tokens.length,
    },
    stats: {
      sessions: sessionManager.count(),
      websocket: wsStats,
      sse: sseStats,
    },
  });
});

// API 版本前缀
const api = new Hono();

// 挂载路由
api.route('/sessions', sessionsRouter);
api.route('/tools', toolsRouter);
api.route('/config', configRouter);
api.route('/files', filesRouter);
api.route('/commands', commandsRouter);

// SSE 事件流
api.get('/sessions/:id/events', handleSSE);

// WebSocket 端点
api.get(
  '/ws/:sessionId',
  upgradeWebSocket((c) => {
    const sessionId = c.req.param('sessionId');

    return {
      onOpen(_event, ws) {
        handleWebSocket(ws, sessionId);
      },
      onMessage(event, ws) {
        handleWebSocketMessage(ws, sessionId, event.data);
      },
      onClose(_event, ws) {
        handleWebSocketClose(ws, sessionId);
      },
      onError(event, ws) {
        console.error('[WS] Error:', event);
        handleWebSocketClose(ws, sessionId);
      },
    };
  })
);

// 挂载 API 到 /api
app.route('/api', api);

// 404 处理
app.notFound((c) => {
  return c.json(
    {
      success: false,
      error: 'Not found',
    },
    404
  );
});

// 错误处理
app.onError((err, c) => {
  console.error('[Server Error]', err);

  return c.json(
    {
      success: false,
      error: err.message || 'Internal server error',
    },
    500
  );
});

// 服务器配置
export interface ServerOptions {
  port?: number;
  host?: string;
  /** 是否启用认证 (远程模式自动启用) */
  auth?: boolean;
  /** 预设的 token */
  token?: string;
}

/**
 * 创建服务器实例
 */
export function createServer(options: ServerOptions = {}) {
  const { port = 3000, host = '127.0.0.1' } = options;

  return {
    app,
    websocket,
    port,
    host,
  };
}

/**
 * 初始化服务器（加载 core 模块等）
 */
export async function initServer(options: ServerOptions = {}): Promise<void> {
  // 初始化 SessionManager（加载持久化的 sessions）
  const sessionManager = getSessionManager();
  await sessionManager.init();

  // 尝试加载 core 模块
  const coreLoaded = await initCore();
  if (coreLoaded) {
    console.log('[Server] Core module initialized');
  } else {
    console.warn('[Server] Core module not available, running in limited mode');
  }

  // 初始化认证
  const { host = '127.0.0.1', auth, token } = options;
  const isRemote = host !== '127.0.0.1' && host !== 'localhost';
  const authEnabled = auth !== undefined ? auth : isRemote;

  initAuth({
    enabled: authEnabled,
    tokens: [],
    skipPaths: ['/health', '/api/health'],
  });

  // 如果启用认证，生成或使用提供的 token
  if (authEnabled) {
    const serverToken = token || generateToken();
    addToken(serverToken);
    console.log(`[Auth] Authentication enabled`);
    console.log(`[Auth] Token: ${serverToken}`);
  }
}

/**
 * 启动服务器 (Bun 环境)
 */
export async function startServer(options: ServerOptions = {}): Promise<void> {
  const { port = 3000, host = '127.0.0.1' } = options;

  // 初始化
  await initServer(options);

  const coreStatus = isCoreAvailable() ? '✅ Core loaded' : '⚠️  Core not available';
  const authConfig = getAuthConfig();
  const authStatus = authConfig.enabled ? '🔐 Enabled' : '🔓 Disabled';

  console.log(`
╔════════════════════════════════════════════╗
║        AI Assistant Server                 ║
╠════════════════════════════════════════════╣
║  REST API:    http://${host}:${port}/api
║  WebSocket:   ws://${host}:${port}/api/ws/:sessionId
║  SSE:         http://${host}:${port}/api/sessions/:id/events
║  Health:      http://${host}:${port}/health
║  Agent:       ${coreStatus}
║  Auth:        ${authStatus}
╚════════════════════════════════════════════╝
  `);

  // Bun.serve 需要在 CLI 包中调用
  // 这里只导出配置
}

// 导出
export { app, websocket };
export { getSessionManager } from './session/manager.js';
export { registerTool, getRegisteredTools } from './routes/tools.js';
export { getConfig, setConfig } from './routes/config.js';
export {
  emitEvent,
  broadcastEvent,
  emitStatusEvent,
  emitLogEvent,
  emitProgressEvent,
  emitFileChangeEvent,
} from './sse.js';
export { broadcastToSession } from './ws.js';
export {
  initCore,
  isCoreAvailable,
  getAgentStats,
  processMessage,
  cancelProcessing,
} from './agent/index.js';
export {
  initAuth,
  getAuthConfig,
  generateToken,
  addToken,
  removeToken,
  setAuthEnabled,
  validateToken,
  maskToken,
} from './auth/index.js';
export * from './types.js';