proxy.js

#!/usr/bin/env node
/**
 * Local HTTPS-to-HTTP Proxy for LLM endpoints (e.g. Ollama)
 *
 * Accepts HTTPS requests and forwards them to a local HTTP endpoint.
 * Solves the mixed-content problem when your app runs on HTTPS but
 * your local LLM only speaks HTTP.
 *
 * Configuration priority (highest wins):
 *   1. Environment variables
 *   2. config.json
 *   3. Built-in defaults
 *
 * Usage:
 *   node proxy.js                    # uses config.json or defaults
 *   node proxy.js --config my.json   # uses custom config file
 *   TARGET_HOST=10.0.0.5 node proxy.js  # env var overrides config
 */

const https = require("https");
const http = require("http");
const fs = require("fs");
const path = require("path");

// --- Load config.json ---
const CONFIG_PATH = (() => {
  const argIdx = process.argv.indexOf("--config");
  if (argIdx !== -1 && process.argv[argIdx + 1]) {
    return path.resolve(process.argv[argIdx + 1]);
  }
  return path.join(__dirname, "config.json");
})();

let fileConfig = {};
try {
  fileConfig = JSON.parse(fs.readFileSync(CONFIG_PATH, "utf8"));
  console.log(`  Loaded config from ${CONFIG_PATH}`);
} catch (err) {
  if (CONFIG_PATH !== path.join(__dirname, "config.json") || err.code !== "ENOENT") {
    // Only warn if user specified a custom path or it's a parse error
    if (err.code === "ENOENT") {
      console.error(`  ERROR: Config file not found: ${CONFIG_PATH}`);
      process.exit(1);
    }
    if (err instanceof SyntaxError) {
      console.error(`  ERROR: Invalid JSON in ${CONFIG_PATH}: ${err.message}`);
      process.exit(1);
    }
  }
  // Default config.json missing is fine, use defaults
}

// --- Configuration (env vars > config.json > defaults) ---
const PROXY_PORT = parseInt(process.env.PROXY_PORT || fileConfig.proxy?.port || "8443", 10);
const TARGET_HOST = process.env.TARGET_HOST || fileConfig.target?.host || "127.0.0.1";
const TARGET_PORT = parseInt(process.env.TARGET_PORT || fileConfig.target?.port || "11434", 10);
const CERT_FILE = process.env.CERT_FILE || fileConfig.tls?.cert || path.join(__dirname, "certs", "cert.pem");
const KEY_FILE = process.env.KEY_FILE || fileConfig.tls?.key || path.join(__dirname, "certs", "key.pem");
const ALLOWED_ORIGINS = process.env.ALLOWED_ORIGINS || fileConfig.proxy?.allowedOrigins || "*";
const TRACKING_ENABLED = fileConfig.tracking?.enabled !== false;
const STATS_FILE = fileConfig.tracking?.statsFile || path.join(__dirname, "stats.json");
const PERSIST_INTERVAL = (fileConfig.tracking?.persistInterval || 30) * 1000;
const HISTORY_SIZE = fileConfig.tracking?.historySize || 100;

// Headers to skip when forwarding (hop-by-hop + browser-specific)
const SKIP_HEADERS = new Set([
  "host", "origin", "referer", "connection", "upgrade", "keep-alive",
  "transfer-encoding", "te", "trailer", "proxy-authorization",
  "proxy-connection", "sec-websocket-key", "sec-websocket-version",
  "sec-websocket-extensions", "sec-fetch-mode", "sec-fetch-site",
  "sec-fetch-dest", "sec-ch-ua", "sec-ch-ua-mobile", "sec-ch-ua-platform",
]);

// --- Token Tracker ---
const stats = {
  startedAt: new Date().toISOString(),
  requests: 0,
  completions: 0,
  errors: 0,
  tokens: {
    prompt: 0,
    completion: 0,
    total: 0,
  },
  byModel: {},     // { "qwen3:8b": { requests, prompt, completion, total } }
  byEndpoint: {},  // { "/v1/chat/completions": { requests, prompt, completion, total } }
  history: [],     // last 100 requests
};

// Load persisted stats on startup
try {
  const saved = JSON.parse(fs.readFileSync(STATS_FILE, "utf8"));
  stats.tokens = saved.tokens || stats.tokens;
  stats.byModel = saved.byModel || {};
  stats.byEndpoint = saved.byEndpoint || {};
  stats.requests = saved.requests || 0;
  stats.completions = saved.completions || 0;
  stats.errors = saved.errors || 0;
  stats.startedAt = saved.startedAt || stats.startedAt;
  console.log(`  Loaded persisted stats: ${stats.tokens.total} total tokens`);
} catch {
  // No saved stats, start fresh
}

function persistStats() {
  try {
    fs.writeFileSync(STATS_FILE, JSON.stringify(stats, null, 2));
  } catch {}
}

// Save stats periodically
if (TRACKING_ENABLED) {
  setInterval(persistStats, PERSIST_INTERVAL);
}

function trackTokens(endpoint, responseBody) {
  try {
    const data = JSON.parse(responseBody);
    const usage = data.usage;
    if (!usage) return;

    // Support both OpenAI format (prompt_tokens) and Anthropic format (input_tokens)
    const prompt = usage.prompt_tokens || usage.input_tokens || 0;
    const completion = usage.completion_tokens || usage.output_tokens || 0;
    const cacheCreation = usage.cache_creation_input_tokens || 0;
    const cacheRead = usage.cache_read_input_tokens || 0;
    const total = usage.total_tokens || (prompt + completion + cacheCreation + cacheRead);
    const model = data.model || "unknown";

    // Global totals
    stats.completions++;
    stats.tokens.prompt += prompt;
    stats.tokens.completion += completion;
    stats.tokens.total += total;

    // Per model
    if (!stats.byModel[model]) {
      stats.byModel[model] = { requests: 0, prompt: 0, completion: 0, total: 0 };
    }
    stats.byModel[model].requests++;
    stats.byModel[model].prompt += prompt;
    stats.byModel[model].completion += completion;
    stats.byModel[model].total += total;

    // Per endpoint
    if (!stats.byEndpoint[endpoint]) {
      stats.byEndpoint[endpoint] = { requests: 0, prompt: 0, completion: 0, total: 0 };
    }
    stats.byEndpoint[endpoint].requests++;
    stats.byEndpoint[endpoint].prompt += prompt;
    stats.byEndpoint[endpoint].completion += completion;
    stats.byEndpoint[endpoint].total += total;

    // History (keep last 100)
    stats.history.push({
      time: new Date().toISOString(),
      model,
      endpoint,
      prompt,
      completion,
      total,
    });
    if (stats.history.length > HISTORY_SIZE) {
      stats.history = stats.history.slice(-HISTORY_SIZE);
    }

    console.log(`  tokens: +${prompt} prompt, +${completion} completion = ${total} (cumulative: ${stats.tokens.total})`);
  } catch {
    // Not a JSON response with usage data, ignore
  }
}

function formatNumber(n) {
  if (n >= 1_000_000) return (n / 1_000_000).toFixed(2) + "M";
  if (n >= 1_000) return (n / 1_000).toFixed(1) + "k";
  return String(n);
}

// --- Load TLS certs ---
let cert, key;
try {
  cert = fs.readFileSync(CERT_FILE);
  key = fs.readFileSync(KEY_FILE);
} catch (err) {
  console.error(`\n  ERROR: Cannot read TLS certificates.`);
  console.error(`  Expected:`);
  console.error(`    cert: ${CERT_FILE}`);
  console.error(`    key:  ${KEY_FILE}`);
  console.error(`\n  Run ./generate-cert.sh first to create self-signed certs.\n`);
  process.exit(1);
}

// --- CORS helper ---
function setCorsHeaders(res, origin) {
  const allowedOrigin =
    ALLOWED_ORIGINS === "*" ? "*" : ALLOWED_ORIGINS.split(",").find((o) => o.trim() === origin) || "";

  if (allowedOrigin) {
    res.setHeader("Access-Control-Allow-Origin", allowedOrigin);
    res.setHeader("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS, PATCH");
    res.setHeader("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Requested-With, x-api-key, anthropic-version, x-stainless-arch, x-stainless-lang, x-stainless-os, x-stainless-package-version, x-stainless-retry-count, x-stainless-runtime, x-stainless-runtime-version");
    res.setHeader("Access-Control-Max-Age", "86400");
    res.setHeader("Access-Control-Allow-Credentials", "true");
  }
}

// --- Stats endpoint ---
function handleStatsRequest(clientRes, origin) {
  setCorsHeaders(clientRes, origin);
  const uptime = Math.floor((Date.now() - new Date(stats.startedAt).getTime()) / 1000);
  const hours = Math.floor(uptime / 3600);
  const mins = Math.floor((uptime % 3600) / 60);

  const response = {
    uptime: `${hours}h ${mins}m`,
    startedAt: stats.startedAt,
    requests: stats.requests,
    completions: stats.completions,
    errors: stats.errors,
    tokens: {
      prompt: stats.tokens.prompt,
      completion: stats.tokens.completion,
      total: stats.tokens.total,
      formatted: formatNumber(stats.tokens.total),
    },
    byModel: stats.byModel,
    byEndpoint: stats.byEndpoint,
    recentRequests: stats.history.slice(-10),
  };

  clientRes.writeHead(200, { "Content-Type": "application/json" });
  clientRes.end(JSON.stringify(response, null, 2));
}

// --- Stats dashboard (HTML) ---
function handleDashboard(clientRes, origin) {
  setCorsHeaders(clientRes, origin);
  clientRes.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
  clientRes.end(`<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>LLM Proxy - Token Tracker</title>
  <style>
    * { margin: 0; padding: 0; box-sizing: border-box; }
    body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, monospace; background: #0f172a; color: #e2e8f0; padding: 24px; }
    h1 { font-size: 18px; color: #94a3b8; margin-bottom: 20px; }
    .grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 16px; margin-bottom: 24px; }
    .card { background: #1e293b; border: 1px solid #334155; border-radius: 8px; padding: 16px; }
    .card .label { font-size: 11px; text-transform: uppercase; letter-spacing: 0.05em; color: #64748b; margin-bottom: 4px; }
    .card .value { font-size: 28px; font-weight: 700; color: #f1f5f9; }
    .card .sub { font-size: 12px; color: #64748b; margin-top: 4px; }
    .section { margin-bottom: 24px; }
    .section h2 { font-size: 14px; color: #94a3b8; margin-bottom: 12px; text-transform: uppercase; letter-spacing: 0.05em; }
    table { width: 100%; border-collapse: collapse; background: #1e293b; border-radius: 8px; overflow: hidden; }
    th { text-align: left; font-size: 10px; text-transform: uppercase; letter-spacing: 0.05em; color: #64748b; padding: 8px 12px; border-bottom: 1px solid #334155; }
    td { padding: 8px 12px; font-size: 13px; border-bottom: 1px solid #1e293b; }
    tr:hover td { background: #1e293b; }
    .mono { font-family: "SF Mono", "Fira Code", monospace; }
    .refresh { font-size: 11px; color: #475569; text-align: center; margin-top: 16px; }
    .tag { display: inline-block; background: #334155; border-radius: 4px; padding: 2px 8px; font-size: 11px; color: #94a3b8; }
  </style>
</head>
<body>
  <h1>LLM Proxy - Token Tracker</h1>
  <div class="grid" id="cards"></div>
  <div class="section" id="models-section"><h2>By Model</h2><table id="models"><thead><tr><th>Model</th><th>Requests</th><th>Prompt</th><th>Completion</th><th>Total</th></tr></thead><tbody></tbody></table></div>
  <div class="section" id="history-section"><h2>Recent Requests</h2><table id="history"><thead><tr><th>Time</th><th>Model</th><th>Endpoint</th><th>Prompt</th><th>Completion</th><th>Total</th></tr></thead><tbody></tbody></table></div>
  <div class="refresh">Auto-refreshes every 5s</div>
  <script>
    function fmt(n) {
      if (n >= 1e6) return (n/1e6).toFixed(2) + 'M';
      if (n >= 1e3) return (n/1e3).toFixed(1) + 'k';
      return n;
    }
    async function refresh() {
      try {
        const res = await fetch('/_proxy/stats');
        const d = await res.json();
        document.getElementById('cards').innerHTML =
          card('Total Tokens', fmt(d.tokens.total), d.tokens.prompt + ' prompt + ' + d.tokens.completion + ' completion') +
          card('Completions', d.completions, d.requests + ' total requests') +
          card('Errors', d.errors, '') +
          card('Uptime', d.uptime, 'since ' + new Date(d.startedAt).toLocaleString());
        const mb = document.querySelector('#models tbody');
        mb.innerHTML = Object.entries(d.byModel).map(([m, s]) =>
          '<tr><td><span class="tag">' + m + '</span></td><td>' + s.requests + '</td><td class="mono">' + fmt(s.prompt) + '</td><td class="mono">' + fmt(s.completion) + '</td><td class="mono"><b>' + fmt(s.total) + '</b></td></tr>'
        ).join('') || '<tr><td colspan="5" style="color:#475569">No completions yet</td></tr>';
        const hb = document.querySelector('#history tbody');
        hb.innerHTML = (d.recentRequests || []).reverse().map(r =>
          '<tr><td>' + new Date(r.time).toLocaleTimeString() + '</td><td><span class="tag">' + r.model + '</span></td><td class="mono" style="font-size:11px">' + r.endpoint + '</td><td class="mono">' + r.prompt + '</td><td class="mono">' + r.completion + '</td><td class="mono"><b>' + r.total + '</b></td></tr>'
        ).join('') || '<tr><td colspan="6" style="color:#475569">No completions yet</td></tr>';
      } catch(e) { console.error(e); }
    }
    refresh();
    setInterval(refresh, 5000);
    function card(label, value, sub) {
      return '<div class="card"><div class="label">' + label + '</div><div class="value">' + value + '</div>' + (sub ? '<div class="sub">' + sub + '</div>' : '') + '</div>';
    }
  </script>
</body>
</html>`);
}

// --- Proxy a single request ---
function proxyRequest(clientReq, clientRes) {
  const origin = clientReq.headers.origin || "";
  const startTime = Date.now();
  const logPrefix = `[${new Date().toISOString()}]`;

  // Handle CORS preflight
  if (clientReq.method === "OPTIONS") {
    setCorsHeaders(clientRes, origin);
    clientRes.writeHead(204);
    clientRes.end();
    return;
  }

  // Serve stats API
  if (clientReq.url === "/_proxy/stats" && clientReq.method === "GET") {
    return handleStatsRequest(clientRes, origin);
  }

  // Serve dashboard
  if (clientReq.url === "/_proxy/dashboard" && clientReq.method === "GET") {
    return handleDashboard(clientRes, origin);
  }

  // Reset stats
  if (clientReq.url === "/_proxy/reset" && clientReq.method === "POST") {
    stats.tokens = { prompt: 0, completion: 0, total: 0 };
    stats.byModel = {};
    stats.byEndpoint = {};
    stats.requests = 0;
    stats.completions = 0;
    stats.errors = 0;
    stats.history = [];
    stats.startedAt = new Date().toISOString();
    persistStats();
    setCorsHeaders(clientRes, origin);
    clientRes.writeHead(200, { "Content-Type": "application/json" });
    clientRes.end(JSON.stringify({ ok: true, message: "Stats reset" }));
    console.log(`${logPrefix} Stats reset`);
    return;
  }

  stats.requests++;

  console.log(`${logPrefix} ${clientReq.method} ${clientReq.url} -> http://${TARGET_HOST}:${TARGET_PORT}${clientReq.url}`);

  // Step 1: Collect the full request body first
  const bodyChunks = [];
  clientReq.on("data", (chunk) => bodyChunks.push(chunk));

  clientReq.on("end", () => {
    const body = Buffer.concat(bodyChunks);

    // Step 2: Build clean headers
    const headers = {
      host: `${TARGET_HOST}:${TARGET_PORT}`,
    };
    for (const [k, v] of Object.entries(clientReq.headers)) {
      if (!SKIP_HEADERS.has(k.toLowerCase()) && v !== undefined) {
        headers[k] = v;
      }
    }
    if (body.length > 0) {
      headers["content-length"] = body.length;
    }

    // Step 3: Send the proxy request
    const proxyReq = http.request(
      {
        hostname: TARGET_HOST,
        port: TARGET_PORT,
        path: clientReq.url,
        method: clientReq.method,
        headers,
      },
      (proxyRes) => {
        // Set CORS headers on response
        setCorsHeaders(clientRes, origin);

        // Copy response headers, skipping CORS ones (we set our own)
        const resHeaders = {};
        for (const [k, v] of Object.entries(proxyRes.headers)) {
          if (!k.startsWith("access-control-")) {
            resHeaders[k] = v;
          }
        }

        // Merge our CORS headers into the response
        for (const name of clientRes.getHeaderNames()) {
          resHeaders[name] = clientRes.getHeader(name);
        }

        clientRes.writeHead(proxyRes.statusCode, resHeaders);

        // Collect response body for token tracking, then pipe through
        const isCompletion =
          clientReq.url.includes("/chat/completions") ||
          clientReq.url.includes("/api/generate") ||
          clientReq.url.includes("/api/chat") ||
          clientReq.url.includes("/v1/messages");

        if (isCompletion && clientReq.method === "POST") {
          // Collect response to extract token usage
          const resChunks = [];
          proxyRes.on("data", (chunk) => {
            resChunks.push(chunk);
            clientRes.write(chunk);
          });
          proxyRes.on("end", () => {
            clientRes.end();
            const elapsed = Date.now() - startTime;
            console.log(`${logPrefix} -> ${proxyRes.statusCode} (${elapsed}ms)`);
            if (TRACKING_ENABLED) {
              trackTokens(clientReq.url, Buffer.concat(resChunks).toString());
            }
          });
        } else {
          // Non-completion: just pipe through
          proxyRes.pipe(clientRes, { end: true });
          proxyRes.on("end", () => {
            const elapsed = Date.now() - startTime;
            console.log(`${logPrefix} -> ${proxyRes.statusCode} (${elapsed}ms)`);
          });
        }
      }
    );

    proxyReq.on("error", (err) => {
      stats.errors++;
      const elapsed = Date.now() - startTime;
      console.error(`${logPrefix} -> ERROR (${elapsed}ms): ${err.message}`);

      if (!clientRes.headersSent) {
        setCorsHeaders(clientRes, origin);
        clientRes.writeHead(502, { "Content-Type": "application/json" });
        clientRes.end(
          JSON.stringify({
            error: "Bad Gateway",
            message: `Cannot connect to http://${TARGET_HOST}:${TARGET_PORT}: ${err.message}`,
            hint: "Make sure your LLM service (e.g. Ollama) is running.",
          })
        );
      }
    });

    // Write body and send
    if (body.length > 0) {
      proxyReq.write(body);
    }
    proxyReq.end();
  });

  // Handle client disconnect
  clientReq.on("error", () => {});
}

// --- Create HTTPS server ---
const server = https.createServer({ cert, key }, proxyRequest);

// --- Start ---
server.listen(PROXY_PORT, "0.0.0.0", () => {
  const trackingLine = TRACKING_ENABLED
    ? `  ║  Tracking:   enabled (${path.basename(STATS_FILE)})${" ".repeat(Math.max(0, 23 - path.basename(STATS_FILE).length))}║`
    : "  ║  Tracking:   disabled                        ║";
  console.log(`
  ╔══════════════════════════════════════════════════════╗
  ║         LLM HTTPS-to-HTTP Proxy                     ║
  ╠══════════════════════════════════════════════════════╣
  ║                                                      ║
  ║  Listening:  https://0.0.0.0:${String(PROXY_PORT).padEnd(25)}║
  ║  Proxying:   http://${TARGET_HOST}:${String(TARGET_PORT).padEnd(21)}║
  ║  CORS:       ${ALLOWED_ORIGINS.substring(0, 38).padEnd(38)} ║
${trackingLine}
  ║                                                      ║
  ║  Dashboard:  https://localhost:${String(PROXY_PORT).padEnd(22)}║
  ║              /_proxy/dashboard                       ║
  ║  Stats API:  /_proxy/stats                           ║
  ║  Reset:      POST /_proxy/reset                      ║
  ║                                                      ║
  ╚══════════════════════════════════════════════════════╝
  `);
});

// Persist stats on shutdown
function shutdown() {
  console.log("\nShutting down... saving stats.");
  persistStats();
  server.close(() => process.exit(0));
}
process.on("SIGINT", shutdown);
process.on("SIGTERM", shutdown);