flynn/src/gateway/server.ts

import { WebSocketServer, WebSocket } from 'ws';
import { randomUUID } from 'crypto';
import { createServer, type Server as HttpServer, type IncomingMessage, type ServerResponse } from 'http';
import { Router } from './router.js';
import { serveStatic } from './static.js';
import { SessionBridge } from './session-bridge.js';
import type { SessionBridgeConfig } from './session-bridge.js';
import { authenticateRequest } from './auth.js';
import type { AuthConfig } from './auth.js';
import {
  parseMessage,
  makeError,
  ErrorCode,
  type OutboundMessage,
} from './protocol.js';
import {
  createSystemHandlers,
  createSessionHandlers,
  createToolHandlers,
  createAgentHandlers,
  createConfigHandlers,
} from './handlers/index.js';
import type { SessionManager } from '../session/manager.js';
import type { Config } from '../config/index.js';
import type { ToolRegistry } from '../tools/registry.js';
import type { ToolExecutor } from '../tools/executor.js';

export interface GatewayServerConfig {
  port: number;
  host?: string;
  sessionManager: SessionManager;
  modelClient: SessionBridgeConfig['modelClient'];
  systemPrompt: string;
  toolRegistry: ToolRegistry;
  toolExecutor: ToolExecutor;
  version?: string;
  auth?: AuthConfig;
  /** Whether to apply token auth to HTTP requests too (default: true when token is set). */
  authHttp?: boolean;
  uiDir?: string;
  config?: Config;
  /** Optional callback for system.restart. Should trigger graceful shutdown + process restart. */
  restart?: () => Promise<void>;
  channelRegistry?: { list(): Array<{ readonly name: string; readonly status: string }> };
}

export class GatewayServer {
  private wss: WebSocketServer | null = null;
  private httpServer: HttpServer | null = null;
  private router: Router;
  private sessionBridge: SessionBridge;
  private connectionMap: Map<WebSocket, string> = new Map();
  private config: GatewayServerConfig;
  private startTime: number = Date.now();

  constructor(config: GatewayServerConfig) {
    this.config = config;

    this.sessionBridge = new SessionBridge({
      sessionManager: config.sessionManager,
      modelClient: config.modelClient,
      systemPrompt: config.systemPrompt,
      toolRegistry: config.toolRegistry,
      toolExecutor: config.toolExecutor,
    });

    this.router = new Router();
    this.registerHandlers();
  }

  private registerHandlers(): void {
    const systemHandlers = createSystemHandlers({
      startTime: this.startTime,
      version: this.config.version ?? '0.1.0',
      getSessionCount: () => this.sessionBridge.listSessions().length,
      getToolCount: () => this.config.toolRegistry.list().length,
      getConnectionCount: () => this.sessionBridge.connectionCount,
      restart: this.config.restart,
      getChannels: this.config.channelRegistry
        ? () => this.config.channelRegistry!.list().map(a => ({ name: a.name, status: a.status }))
        : undefined,
      getUsage: () => ({
        totalSessions: this.config.sessionManager.listSessions().length,
        activeConnections: this.sessionBridge.connectionCount,
      }),
    });

    const sessionHandlers = createSessionHandlers({
      sessionManager: this.config.sessionManager,
      sessionBridge: this.sessionBridge,
    });

    const toolHandlers = createToolHandlers({
      toolRegistry: this.config.toolRegistry,
      toolExecutor: this.config.toolExecutor,
    });

    const agentHandlers = createAgentHandlers({
      sessionBridge: this.sessionBridge,
    });

    // Config handlers (only if config object is provided)
    if (this.config.config) {
      const configHandlers = createConfigHandlers({ config: this.config.config });
      for (const [method, handler] of Object.entries(configHandlers)) {
        this.router.register(method, handler);
      }
    }

    // Register all methods
    for (const [method, handler] of Object.entries(systemHandlers)) {
      this.router.register(method, handler);
    }
    for (const [method, handler] of Object.entries(sessionHandlers)) {
      this.router.register(method, handler);
    }
    for (const [method, handler] of Object.entries(toolHandlers)) {
      this.router.register(method, handler);
    }
    for (const [method, handler] of Object.entries(agentHandlers)) {
      this.router.register(method, handler);
    }
  }

  async start(): Promise<void> {
    const host = this.config.host ?? '127.0.0.1';
    const { port } = this.config;

    return new Promise((resolve) => {
      // Create HTTP server first — handles static file requests
      this.httpServer = createServer((req: IncomingMessage, res: ServerResponse) => {
        this.handleHttpRequest(req, res);
      });

      // Attach WebSocket server to the shared HTTP server (no separate port)
      this.wss = new WebSocketServer({ server: this.httpServer });

      this.wss.on('connection', (ws: WebSocket, req: IncomingMessage) => {
        // Auth check on upgrade — only WS connections require auth
        const authResult = authenticateRequest(req, this.config.auth ?? {});
        if (!authResult.authenticated) {
          ws.close(4001, authResult.error ?? 'Authentication failed');
          return;
        }
        this.handleConnection(ws, authResult.identity);
      });

      this.httpServer.listen(port, host, () => {
        console.log(`Gateway server listening on ${host}:${port}`);
        resolve();
      });
    });
  }

  async stop(): Promise<void> {
    // Close all WebSocket connections first
    for (const [ws, connectionId] of this.connectionMap) {
      this.sessionBridge.disconnect(connectionId);
      ws.close(1001, 'Server shutting down');
    }
    this.connectionMap.clear();

    // Close WSS first, then the underlying HTTP server
    await new Promise<void>((resolve) => {
      if (!this.wss) {
        resolve();
        return;
      }
      this.wss.close(() => {
        this.wss = null;
        resolve();
      });
    });

    await new Promise<void>((resolve) => {
      if (!this.httpServer) {
        resolve();
        return;
      }
      this.httpServer.close(() => {
        this.httpServer = null;
        resolve();
      });
    });
  }

  private handleConnection(ws: WebSocket, identity?: string): void {
    const connectionId = randomUUID();
    this.sessionBridge.connect(connectionId);
    this.connectionMap.set(ws, connectionId);

    ws.on('message', async (data) => {
      const raw = data.toString();
      await this.handleMessage(ws, connectionId, raw);
    });

    ws.on('close', () => {
      this.sessionBridge.disconnect(connectionId);
      this.connectionMap.delete(ws);
    });

    ws.on('error', (err) => {
      console.error(`WebSocket error (${connectionId}):`, err.message);
    });
  }

  /**
   * Handle incoming HTTP requests.
   * Optionally applies auth (when authHttp is enabled and a token is configured).
   * Delegates to serveStatic for UI files; returns 404 if no UI dir or file not found.
   */
  private async handleHttpRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
    // Apply auth to HTTP requests when configured
    const authConfig = this.config.auth ?? {};
    if (this.config.authHttp !== false && authConfig.token) {
      const authResult = authenticateRequest(req, authConfig);
      if (!authResult.authenticated) {
        res.writeHead(401, {
          'Content-Type': 'text/plain',
          'WWW-Authenticate': 'Bearer',
        });
        res.end(authResult.error ?? 'Unauthorized');
        return;
      }
    }

    const uiDir = this.config.uiDir;

    if (uiDir) {
      const served = await serveStatic(req, res, uiDir);
      if (served) return;
    }

    // No UI directory configured, or file not found
    res.writeHead(404, { 'Content-Type': 'text/plain' });
    res.end('Not Found');
  }

  private async handleMessage(ws: WebSocket, connectionId: string, raw: string): Promise<void> {
    const request = parseMessage(raw);

    if (!request) {
      this.send(ws, makeError(0, ErrorCode.ParseError, 'Invalid JSON or missing required fields'));
      return;
    }

    // Inject connectionId into params so handlers can identify the client
    if (!request.params) request.params = {};
    request.params.connectionId = connectionId;

    const send = (msg: OutboundMessage) => this.send(ws, msg);
    const response = await this.router.dispatch(request, send);

    if (response) {
      this.send(ws, response);
    }
  }

  private send(ws: WebSocket, msg: OutboundMessage): void {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(msg));
    }
  }

  /** Get the underlying WebSocketServer (for testing). */
  getWss(): WebSocketServer | null {
    return this.wss;
  }

  /** Get the underlying HTTP server (for testing). */
  getHttpServer(): HttpServer | null {
    return this.httpServer;
  }

  /** Get the session bridge (for testing/debugging). */
  getSessionBridge(): SessionBridge {
    return this.sessionBridge;
  }

  /** Get list of registered methods. */
  getMethods(): string[] {
    return this.router.listMethods();
  }
}