the-matrix/js/websocket.js

import { AGENT_DEFS, colorToCss } from './agent-defs.js';
import { setAgentState } from './agents.js';
import { appendChatMessage } from './ui.js';

const WS_URL = import.meta.env.VITE_WS_URL || '';

const agentById = Object.fromEntries(AGENT_DEFS.map(d => [d.id, d]));

let ws = null;
let connectionState = 'disconnected';
let jobCount = 0;
let reconnectTimer = null;
let reconnectAttempts = 0;

const RECONNECT_BASE_MS = 2000;
const RECONNECT_MAX_MS = 30000;
const HEARTBEAT_INTERVAL_MS = 30000;
const HEARTBEAT_TIMEOUT_MS = 5000;

let heartbeatTimer = null;
let heartbeatTimeout = null;

export function initWebSocket(_scene) {
  if (!WS_URL) {
    connectionState = 'disconnected';
    return;
  }
  connect();
}

function connect() {
  if (ws) {
    ws.onclose = null;
    ws.close();
  }

  connectionState = 'connecting';

  try {
    ws = new WebSocket(WS_URL);
  } catch (err) {
    console.warn('[Matrix WS] Connection failed:', err.message || err);
    logEvent('WebSocket connection failed');
    connectionState = 'disconnected';
    scheduleReconnect();
    return;
  }

  ws.onopen = () => {
    connectionState = 'connected';
    reconnectAttempts = 0;
    clearTimeout(reconnectTimer);
    startHeartbeat();
    ws.send(JSON.stringify({
      type: 'subscribe',
      channel: 'agents',
      clientId: crypto.randomUUID(),
    }));
  };

  ws.onmessage = (event) => {
    // Any message counts as a heartbeat response
    resetHeartbeatTimeout();
    try {
      handleMessage(JSON.parse(event.data));
    } catch (err) {
      console.warn('[Matrix WS] Message parse/handle error:', err.message, '| raw:', event.data?.slice?.(0, 200));
    }
  };

  ws.onerror = (event) => {
    console.warn('[Matrix WS] Error event:', event);
    connectionState = 'disconnected';
  };

  ws.onclose = (event) => {
    connectionState = 'disconnected';
    stopHeartbeat();

    // Don't reconnect on clean close (1000) or going away (1001)
    if (event.code === 1000 || event.code === 1001) {
      console.info('[Matrix WS] Clean close (code ' + event.code + '), not reconnecting');
      return;
    }

    console.warn('[Matrix WS] Unexpected close — code:', event.code, 'reason:', event.reason || '(none)');
    scheduleReconnect();
  };
}

function scheduleReconnect() {
  clearTimeout(reconnectTimer);
  // Exponential backoff: 2s, 4s, 8s, 16s, capped at 30s
  const delay = Math.min(RECONNECT_BASE_MS * Math.pow(2, reconnectAttempts), RECONNECT_MAX_MS);
  reconnectAttempts++;
  console.info('[Matrix WS] Reconnecting in', Math.round(delay / 1000), 's (attempt', reconnectAttempts + ')');
  reconnectTimer = setTimeout(connect, delay);
}

/* ── Heartbeat / zombie connection detection ── */

function startHeartbeat() {
  stopHeartbeat();
  heartbeatTimer = setInterval(() => {
    if (ws && ws.readyState === WebSocket.OPEN) {
      try {
        ws.send(JSON.stringify({ type: 'ping' }));
      } catch { /* ignore send failures, onclose will fire */ }
      heartbeatTimeout = setTimeout(() => {
        console.warn('[Matrix WS] Heartbeat timeout — closing zombie connection');
        if (ws) ws.close(4000, 'heartbeat timeout');
      }, HEARTBEAT_TIMEOUT_MS);
    }
  }, HEARTBEAT_INTERVAL_MS);
}

function stopHeartbeat() {
  clearInterval(heartbeatTimer);
  clearTimeout(heartbeatTimeout);
  heartbeatTimer = null;
  heartbeatTimeout = null;
}

function resetHeartbeatTimeout() {
  clearTimeout(heartbeatTimeout);
  heartbeatTimeout = null;
}

/* ── Message handler ── */

function handleMessage(msg) {
  switch (msg.type) {
    case 'agent_state': {
      if (msg.agentId && msg.state) {
        setAgentState(msg.agentId, msg.state);
      }
      break;
    }
    case 'job_started': {
      jobCount++;
      if (msg.agentId) setAgentState(msg.agentId, 'active');
      logEvent(`JOB ${(msg.jobId || '').slice(0, 8)} started`);
      break;
    }
    case 'job_completed': {
      if (jobCount > 0) jobCount--;
      if (msg.agentId) setAgentState(msg.agentId, 'idle');
      logEvent(`JOB ${(msg.jobId || '').slice(0, 8)} completed`);
      break;
    }
    case 'chat': {
      const def = agentById[msg.agentId];
      if (def && msg.text) {
        appendChatMessage(def.label, msg.text, colorToCss(def.color));
      }
      break;
    }
    case 'pong':
    case 'agent_count':
      break;
    default:
      console.debug('[Matrix WS] Unhandled message type:', msg.type);
      break;
  }
}

function logEvent(text) {
  appendChatMessage('SYS', text, '#005500');
}

export function getConnectionState() {
  return connectionState;
}

export function getJobCount() {
  return jobCount;
}