modelgrid/ts/ui/server.ts

/**
 * UI Server
 *
 * Serves the ModelGrid operations console on its own port, separate from
 * the OpenAI-compatible API. Assets come from one of two sources:
 *   - 'disk':   read on demand from `ts_web/` (dev loop, hot edits)
 *   - 'bundle': from the generated `ts_bundled/bundle.ts` module
 *               (default, required for `deno compile` single-binary builds)
 *
 * Plus a single JSON endpoint `/_ui/overview` that the SPA calls to render
 * the Overview view without cross-origin fetches into the API server.
 */

import * as http from 'node:http';
import * as fs from 'node:fs/promises';
import { dirname, extname, join, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';
import type { IUiConfig } from '../interfaces/config.ts';
import type { IHealthResponse } from '../interfaces/api.ts';
import { logger } from '../logger.ts';
import { VERSION } from '../constants.ts';
import type { ContainerManager } from '../containers/container-manager.ts';
import type { ClusterManager } from '../cluster/cluster-manager.ts';
import { GpuDetector } from '../hardware/gpu-detector.ts';
import { buildHealthSnapshot } from '../helpers/health.ts';

interface IBundledFile {
  path: string;
  contentBase64: string;
}

interface IAssetEntry {
  bytes: Uint8Array;
  contentType: string;
}

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
const REPO_ROOT = resolve(__dirname, '..', '..');
const TS_WEB_DIR = join(REPO_ROOT, 'ts_web');

export class UiServer {
  private server?: http.Server;
  private config: IUiConfig;
  private containerManager: ContainerManager;
  private clusterManager: ClusterManager;
  private gpuDetector: GpuDetector;
  private bundleMap: Map<string, IAssetEntry> | null = null;
  private activeAssetSource: 'disk' | 'bundle' = 'bundle';
  private startTime = 0;

  constructor(
    config: IUiConfig,
    containerManager: ContainerManager,
    clusterManager: ClusterManager,
  ) {
    this.config = config;
    this.containerManager = containerManager;
    this.clusterManager = clusterManager;
    this.gpuDetector = new GpuDetector();
  }

  public async start(): Promise<void> {
    if (this.server) {
      logger.warn('UI server is already running');
      return;
    }

    this.activeAssetSource = this.resolveAssetSource();
    if (this.activeAssetSource === 'bundle') {
      this.bundleMap = await this.loadBundleMap();
      if (!this.bundleMap) {
        logger.warn(
          'UI bundle not found (ts_bundled/bundle.ts missing). ' +
            'Falling back to disk mode — run `deno task bundle:ui` before `deno compile`.',
        );
        this.activeAssetSource = 'disk';
      }
    }

    this.startTime = Date.now();

    this.server = http.createServer(async (req, res) => {
      try {
        await this.handleRequest(req, res);
      } catch (err) {
        logger.error(`UI request error: ${err instanceof Error ? err.message : String(err)}`);
        if (!res.headersSent) {
          res.writeHead(500, { 'Content-Type': 'text/plain' });
          res.end('Internal server error');
        }
      }
    });

    await new Promise<void>((resolve, reject) => {
      this.server!.listen(this.config.port, this.config.host, () => {
        logger.success(
          `UI server started on ${this.config.host}:${this.config.port} ` +
            `(asset source: ${this.activeAssetSource})`,
        );
        resolve();
      });
      this.server!.on('error', (error) => {
        logger.error(`UI server error: ${error.message}`);
        reject(error);
      });
    });
  }

  public async stop(): Promise<void> {
    if (!this.server) return;
    await new Promise<void>((resolve) => {
      this.server!.close(() => resolve());
    });
    this.server = undefined;
    logger.log('UI server stopped');
  }

  public getInfo(): { running: boolean; host: string; port: number; assetSource: string } {
    return {
      running: !!this.server,
      host: this.config.host,
      port: this.config.port,
      assetSource: this.activeAssetSource,
    };
  }

  private async handleRequest(
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<void> {
    const url = new URL(req.url || '/', `http://${req.headers.host || 'localhost'}`);
    const path = url.pathname;

    if (req.method !== 'GET' && req.method !== 'HEAD') {
      res.writeHead(405, { 'Content-Type': 'text/plain', 'Allow': 'GET, HEAD' });
      res.end('Method Not Allowed');
      return;
    }

    if (path === '/_ui/overview') {
      await this.handleOverview(res);
      return;
    }

    await this.serveAsset(path, res);
  }

  private async handleOverview(res: http.ServerResponse): Promise<void> {
    const statuses = await this.containerManager.getAllStatus();
    const models = await this.containerManager.getAllAvailableModels();
    const gpus = await this.gpuDetector.detectGpus();

    const health: IHealthResponse = buildHealthSnapshot({
      statuses,
      modelCount: models.size,
      gpus,
      startTime: this.startTime,
      version: VERSION,
    });

    const clusterConfig = this.clusterManager.getConfig();

    const body = {
      health,
      node: {
        name: clusterConfig?.nodeName ?? 'modelgrid-local',
        role: clusterConfig?.role ?? 'standalone',
        version: VERSION,
      },
    };

    res.writeHead(200, {
      'Content-Type': 'application/json; charset=utf-8',
      'Cache-Control': 'no-store',
    });
    res.end(JSON.stringify(body));
  }

  private async serveAsset(path: string, res: http.ServerResponse): Promise<void> {
    const normalized = path === '/' ? '/index.html' : path;

    if (this.activeAssetSource === 'bundle' && this.bundleMap) {
      const hit = this.bundleMap.get(normalized);
      if (hit) {
        this.writeAsset(res, hit);
        return;
      }
      // SPA fallback: any unknown non-asset path gets index.html.
      if (!hasKnownAssetExtension(normalized)) {
        const shell = this.bundleMap.get('/index.html');
        if (shell) {
          this.writeAsset(res, shell);
          return;
        }
      }
      res.writeHead(404, { 'Content-Type': 'text/plain' });
      res.end('Not Found');
      return;
    }

    // Disk mode
    const safe = normalizePath(normalized);
    if (!safe) {
      res.writeHead(400, { 'Content-Type': 'text/plain' });
      res.end('Bad Request');
      return;
    }
    const full = join(TS_WEB_DIR, safe);
    try {
      const bytes = await fs.readFile(full);
      this.writeAsset(res, {
        bytes: new Uint8Array(bytes),
        contentType: contentTypeForPath(safe),
      });
      return;
    } catch (err) {
      if ((err as NodeJS.ErrnoException).code === 'ENOENT') {
        if (!hasKnownAssetExtension(safe)) {
          try {
            const shell = await fs.readFile(join(TS_WEB_DIR, 'index.html'));
            this.writeAsset(res, {
              bytes: new Uint8Array(shell),
              contentType: 'text/html; charset=utf-8',
            });
            return;
          } catch {
            // fall through to 404
          }
        }
        res.writeHead(404, { 'Content-Type': 'text/plain' });
        res.end('Not Found');
        return;
      }
      throw err;
    }
  }

  private writeAsset(res: http.ServerResponse, asset: IAssetEntry): void {
    res.writeHead(200, {
      'Content-Type': asset.contentType,
      'Content-Length': asset.bytes.byteLength,
      'Cache-Control': 'no-cache',
    });
    res.end(asset.bytes);
  }

  private resolveAssetSource(): 'disk' | 'bundle' {
    const envOverride = typeof Deno !== 'undefined' ? Deno.env.get('UI_ASSET_SOURCE') : undefined;
    const picked = (envOverride || this.config.assetSource || 'bundle').toLowerCase();
    if (picked === 'disk' || picked === 'bundle') return picked;
    logger.warn(`Unknown UI_ASSET_SOURCE "${picked}", defaulting to bundle`);
    return 'bundle';
  }

  private async loadBundleMap(): Promise<Map<string, IAssetEntry> | null> {
    try {
      // The bundle module is generated by `deno task bundle:ui`.
      // @ts-ignore — generated file may not exist until the bundle task runs.
      const mod = await import('../../ts_bundled/bundle.ts');
      const files: IBundledFile[] = mod.files ?? [];
      const map = new Map<string, IAssetEntry>();
      for (const file of files) {
        map.set(`/${file.path}`, {
          bytes: decodeBase64(file.contentBase64),
          contentType: contentTypeForPath(file.path),
        });
      }
      return map;
    } catch {
      return null;
    }
  }
}

function decodeBase64(input: string): Uint8Array {
  const binary = atob(input);
  const bytes = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i);
  return bytes;
}

function normalizePath(path: string): string | null {
  // Strip leading slashes, reject traversal.
  const stripped = path.replace(/^\/+/, '');
  if (stripped.includes('..')) return null;
  return stripped;
}

function hasKnownAssetExtension(path: string): boolean {
  return extname(path) !== '';
}

function contentTypeForPath(path: string): string {
  const ext = extname(path).toLowerCase().replace(/^\./, '');
  const types: Record<string, string> = {
    html: 'text/html; charset=utf-8',
    js: 'application/javascript; charset=utf-8',
    mjs: 'application/javascript; charset=utf-8',
    css: 'text/css; charset=utf-8',
    json: 'application/json; charset=utf-8',
    map: 'application/json; charset=utf-8',
    svg: 'image/svg+xml',
    png: 'image/png',
    jpg: 'image/jpeg',
    jpeg: 'image/jpeg',
    gif: 'image/gif',
    ico: 'image/x-icon',
    webp: 'image/webp',
    woff: 'font/woff',
    woff2: 'font/woff2',
    ttf: 'font/ttf',
    otf: 'font/otf',
    txt: 'text/plain; charset=utf-8',
  };
  return types[ext] || 'application/octet-stream';
}