flynn/src/gateway/server.ts

import { WebSocketServer, WebSocket } from 'ws';
import { randomUUID } from 'crypto';
import type { IncomingMessage } from 'http';
import { Router } from './router.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,
} from './handlers/index.js';
import type { SessionManager } from '../session/manager.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;
}

export class GatewayServer {
  private wss: WebSocketServer | 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,
    });

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

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

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

    // 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> {
    return new Promise((resolve) => {
      this.wss = new WebSocketServer({
        port: this.config.port,
        host: this.config.host ?? '127.0.0.1',
      });

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

      this.wss.on('listening', () => {
        const addr = this.wss?.address();
        const portStr = typeof addr === 'object' && addr ? `:${addr.port}` : '';
        console.log(`Gateway WebSocket server listening on ${this.config.host ?? '127.0.0.1'}${portStr}`);
        resolve();
      });
    });
  }

  async stop(): Promise<void> {
    return new Promise((resolve) => {
      if (!this.wss) {
        resolve();
        return;
      }

      // Close all connections
      for (const [ws, connectionId] of this.connectionMap) {
        this.sessionBridge.disconnect(connectionId);
        ws.close(1001, 'Server shutting down');
      }
      this.connectionMap.clear();

      this.wss.close(() => {
        this.wss = 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);
    });
  }

  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 session bridge (for testing/debugging). */
  getSessionBridge(): SessionBridge {
    return this.sessionBridge;
  }

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