amalgm 0.0.0 → 0.0.32

package/lib/tunnel-chat.js

@@ -0,0 +1,328 @@
+'use strict';
+
+const http = require('http');
+const fs = require('fs');
+const os = require('os');
+const { WebSocket } = require('ws');
+const { RUNTIME_STATE_FILE } = require('./paths');
+
+const DEFAULT_TARGET_PORT = 8084;
+const DEFAULT_CHAT_TUNNEL_URL = 'wss://amalgm-chat-gateway.fly.dev/wire';
+const CONNECT_TIMEOUT_MS = 20_000;
+const WATCHDOG_INTERVAL_MS = 15_000;
+const HEARTBEAT_TIMEOUT_MS = 75_000;
+const ALLOWED_TARGET_PORTS = new Set([8083, 8084]);
+const HOP_BY_HOP_HEADERS = new Set([
+  'connection',
+  'keep-alive',
+  'proxy-authenticate',
+  'proxy-authorization',
+  'te',
+  'trailer',
+  'trailers',
+  'transfer-encoding',
+  'upgrade',
+  'host',
+  'content-length',
+]);
+
+function tunnelUrl(rawUrl) {
+  const url = new URL(rawUrl || DEFAULT_CHAT_TUNNEL_URL);
+  url.searchParams.delete('token');
+  return url.toString();
+}
+
+function tunnelHeaders(token) {
+  return {
+    authorization: `Bearer ${token}`,
+    'user-agent': `amalgm-npm/${require('../package.json').version} (${os.platform()}; ${os.arch()})`,
+  };
+}
+
+function normalizeHeaders(headers = {}) {
+  const out = {};
+  for (const [key, value] of Object.entries(headers || {})) {
+    const lower = key.toLowerCase();
+    if (HOP_BY_HOP_HEADERS.has(lower)) continue;
+    if (Array.isArray(value)) out[key] = value.join(', ');
+    else if (typeof value === 'string') out[key] = value;
+    else if (value != null) out[key] = String(value);
+  }
+  return out;
+}
+
+function responseHeaders(headers = {}) {
+  const out = {};
+  for (const [key, value] of Object.entries(headers)) {
+    if (Array.isArray(value)) out[key] = value.join(', ');
+    else if (typeof value === 'string') out[key] = value;
+    else if (value != null) out[key] = String(value);
+  }
+  return out;
+}
+
+function readRuntimePorts() {
+  const envPort = (name) => {
+    const port = Number(process.env[name]);
+    return Number.isInteger(port) && port > 0 && port <= 65535 ? port : undefined;
+  };
+  try {
+    const parsed = JSON.parse(fs.readFileSync(RUNTIME_STATE_FILE, 'utf8'));
+    return {
+      gateway: Number(parsed?.gateway_port || parsed?.ports?.gateway) || envPort('AMALGM_GATEWAY_PORT'),
+      mcp: Number(parsed?.ports?.amalgm_mcp) || envPort('AMALGM_MCP_PORT'),
+      chat: Number(parsed?.ports?.chat_server) || envPort('CHAT_SERVER_PORT'),
+    };
+  } catch {
+    return {
+      gateway: envPort('AMALGM_GATEWAY_PORT'),
+      mcp: envPort('AMALGM_MCP_PORT'),
+      chat: envPort('CHAT_SERVER_PORT'),
+    };
+  }
+}
+
+function allowedTargetPorts() {
+  const ports = new Set(ALLOWED_TARGET_PORTS);
+  const runtimePorts = readRuntimePorts();
+  for (const port of [runtimePorts.gateway, runtimePorts.mcp, runtimePorts.chat]) {
+    if (Number.isInteger(port) && port > 0 && port <= 65535) ports.add(port);
+  }
+  return ports;
+}
+
+function defaultTargetPort() {
+  const port = readRuntimePorts().chat;
+  return Number.isInteger(port) && port > 0 && port <= 65535 ? port : DEFAULT_TARGET_PORT;
+}
+
+function runtimeGatewayPort() {
+  const port = readRuntimePorts().gateway;
+  return Number.isInteger(port) && port > 0 && port <= 65535 ? port : null;
+}
+
+function createChatTunnel({ record, foreground = false }) {
+  let stopped = false;
+  let ws = null;
+  let reconnectTimer = null;
+  let watchdogTimer = null;
+  let backoffMs = 1000;
+  let connectStartedAt = 0;
+  let lastGatewayFrameAt = 0;
+
+  function log(message) {
+    if (foreground) console.log(`[chat-tunnel] ${message}`);
+  }
+
+  function warn(message) {
+    if (foreground) console.warn(`[chat-tunnel] ${message}`);
+  }
+
+  function send(frame) {
+    if (!ws || ws.readyState !== WebSocket.OPEN) return false;
+    ws.send(JSON.stringify(frame));
+    return true;
+  }
+
+  function terminateSocket(reason) {
+    if (!ws) return;
+    warn(`${reason}; reconnecting`);
+    try {
+      ws.terminate();
+    } catch {
+      try {
+        ws.close();
+      } catch {
+        // noop
+      }
+    }
+  }
+
+  function startWatchdog() {
+    if (watchdogTimer) return;
+
+    watchdogTimer = setInterval(() => {
+      if (stopped || !ws) return;
+      const now = Date.now();
+
+      if (ws.readyState === WebSocket.CONNECTING) {
+        if (connectStartedAt && now - connectStartedAt > CONNECT_TIMEOUT_MS) {
+          terminateSocket('connect timeout');
+        }
+        return;
+      }
+
+      if (ws.readyState !== WebSocket.OPEN) return;
+
+      send({
+        type: 'hello',
+        app_version: require('../package.json').version,
+        runtime_gateway_port: runtimeGatewayPort(),
+      });
+
+      try {
+        ws.ping();
+      } catch {
+        terminateSocket('ping failed');
+        return;
+      }
+
+      if (lastGatewayFrameAt && now - lastGatewayFrameAt > HEARTBEAT_TIMEOUT_MS) {
+        terminateSocket('heartbeat timeout');
+      }
+    }, WATCHDOG_INTERVAL_MS);
+
+    if (typeof watchdogTimer.unref === 'function') watchdogTimer.unref();
+  }
+
+  function stopWatchdog() {
+    if (!watchdogTimer) return;
+    clearInterval(watchdogTimer);
+    watchdogTimer = null;
+  }
+
+  function localHeaders(headers) {
+    const out = normalizeHeaders(headers);
+    const runtimeToken = record?.runtime_token || process.env.AMALGM_RUNTIME_TOKEN || '';
+    if (runtimeToken) out['x-amalgm-runtime-token'] = runtimeToken;
+    return out;
+  }
+
+  function resolveTargetPort(frame) {
+    const candidate = frame.target_port == null ? defaultTargetPort() : Number(frame.target_port);
+    if (!Number.isInteger(candidate) || !allowedTargetPorts().has(candidate)) return null;
+    return candidate;
+  }
+
+  function scheduleReconnect() {
+    if (stopped || reconnectTimer) return;
+    const delay = backoffMs;
+    backoffMs = Math.min(backoffMs * 1.5, 15000);
+    reconnectTimer = setTimeout(() => {
+      reconnectTimer = null;
+      connect();
+    }, delay);
+  }
+
+  function forwardStream(frame) {
+    const targetPort = resolveTargetPort(frame);
+    if (!targetPort) {
+      send({
+        type: 'stream_res_error',
+        req_id: frame.req_id,
+        status: 403,
+        message: 'Target port is not allowed',
+      });
+      return;
+    }
+    const body = frame.body_b64 ? Buffer.from(frame.body_b64, 'base64') : Buffer.alloc(0);
+    const req = http.request(
+      {
+        hostname: '127.0.0.1',
+        port: targetPort,
+        method: frame.method || 'POST',
+        path: frame.path || '/api/agent-stream',
+        headers: localHeaders(frame.headers),
+      },
+      (res) => {
+        send({
+          type: 'stream_res_start',
+          req_id: frame.req_id,
+          status: res.statusCode || 200,
+          headers: responseHeaders(res.headers),
+        });
+        res.on('data', (chunk) => {
+          send({
+            type: 'stream_res_data',
+            req_id: frame.req_id,
+            chunk_b64: Buffer.from(chunk).toString('base64'),
+          });
+        });
+        res.on('end', () => {
+          send({ type: 'stream_res_end', req_id: frame.req_id });
+        });
+      },
+    );
+
+    req.on('error', (error) => {
+      send({
+        type: 'stream_res_error',
+        req_id: frame.req_id,
+        status: 502,
+        message: error.message || 'Local chat request failed',
+      });
+    });
+
+    req.end(body);
+  }
+
+  function handleFrame(frame) {
+    if (!frame || typeof frame.type !== 'string') return;
+    if (frame.type === 'ping') {
+      send({ type: 'pong' });
+      return;
+    }
+    if (frame.type === 'stream_req') {
+      forwardStream(frame);
+    }
+  }
+
+  function connect() {
+    if (stopped) return;
+    const url = tunnelUrl(record.chat_tunnel_url);
+    ws = new WebSocket(url, {
+      headers: tunnelHeaders(record.tunnel_token),
+    });
+    connectStartedAt = Date.now();
+    startWatchdog();
+    ws.on('open', () => {
+      backoffMs = 1000;
+      connectStartedAt = 0;
+      lastGatewayFrameAt = Date.now();
+      log(`connected ${record.computer_id}`);
+      send({
+        type: 'hello',
+        app_version: require('../package.json').version,
+        runtime_gateway_port: runtimeGatewayPort(),
+      });
+    });
+    ws.on('pong', () => {
+      lastGatewayFrameAt = Date.now();
+    });
+    ws.on('message', (data) => {
+      lastGatewayFrameAt = Date.now();
+      try {
+        handleFrame(JSON.parse(data.toString()));
+      } catch {
+        // Ignore malformed gateway frames.
+      }
+    });
+    ws.on('close', (code, reason) => {
+      connectStartedAt = 0;
+      lastGatewayFrameAt = 0;
+      warn(`closed code=${code} reason=${reason ? reason.toString() : ''}`);
+      scheduleReconnect();
+    });
+    ws.on('error', (error) => {
+      warn(error.message);
+    });
+  }
+
+  return {
+    start: connect,
+    stop() {
+      stopped = true;
+      if (reconnectTimer) clearTimeout(reconnectTimer);
+      stopWatchdog();
+      if (ws) {
+        try {
+          ws.close();
+        } catch {
+          // noop
+        }
+      }
+    },
+  };
+}
+
+module.exports = { createChatTunnel };