the-matrix/js/websocket.js

/**
 * websocket.js — WebSocket client for The Matrix.
 *
 * Two modes controlled by Config:
 *   - Live mode: connects to a real Timmy Tower backend via Config.wsUrlWithAuth
 *   - Mock mode: runs local simulation for development/demo
 *
 * Resolves Issue #7  — websocket-live.js with reconnection + backoff
 * Resolves Issue #11 — WS auth token sent via query param on connect
 */

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

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;
let heartbeatTimer = null;
let heartbeatTimeout = null;

/* ── Public API ── */

export function initWebSocket(_scene) {
  if (Config.isLive) {
    logEvent('Connecting to ' + Config.wsUrl + '…');
    connect();
  } else {
    connectionState = 'mock';
    logEvent('Mock mode — no live backend');
  }
}

export function getConnectionState() {
  return connectionState;
}

export function getJobCount() {
  return jobCount;
}

/**
 * Send a message to the backend. In mock mode this is a no-op.
 * @param {object} msg — message object (will be JSON-stringified)
 */
export function sendMessage(msg) {
  if (!ws || ws.readyState !== WebSocket.OPEN) return;
  try {
    ws.send(JSON.stringify(msg));
  } catch { /* onclose will fire */ }
}

/* ── Live WebSocket Client ── */

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

  connectionState = 'connecting';

  const url = Config.wsUrlWithAuth;
  if (!url) {
    connectionState = 'disconnected';
    logEvent('No WS URL configured');
    return;
  }

  try {
    ws = new WebSocket(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();
    logEvent('Connected to backend');

    // Subscribe to agent world-state channel
    sendMessage({
      type: 'subscribe',
      channel: 'agents',
      clientId: crypto.randomUUID(),
    });
  };

  ws.onmessage = (event) => {
    resetHeartbeatTimeout();
    try {
      handleMessage(JSON.parse(event.data));
    } catch (err) {
      console.warn('[Matrix WS] Parse 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');
      logEvent('Disconnected (clean)');
      return;
    }

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

/* ── Reconnection with exponential backoff ── */

function scheduleReconnect() {
  clearTimeout(reconnectTimer);
  const delay = Math.min(
    Config.reconnectBaseMs * Math.pow(2, reconnectAttempts),
    Config.reconnectMaxMs,
  );
  reconnectAttempts++;
  console.info('[Matrix WS] Reconnecting in', Math.round(delay / 1000), 's (attempt', reconnectAttempts + ')');
  reconnectTimer = setTimeout(connect, delay);
}

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

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

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

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

/* ── Message dispatcher ── */

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;
    }

    /**
     * Dynamic agent hot-add (Issue #12).
     *
     * When the backend sends an agent_joined event, we register the new
     * agent definition and spawn its 3D avatar without requiring a page
     * reload. The event payload must include at minimum:
     *   { type: 'agent_joined', id, label, color, role }
     *
     * Optional fields: direction, x, z (auto-placed if omitted).
     */
    case 'agent_joined': {
      if (!msg.id || !msg.label) {
        console.warn('[Matrix WS] agent_joined missing required fields:', msg);
        break;
      }

      // Build a definition compatible with AGENT_DEFS format
      const newDef = {
        id: msg.id,
        label: msg.label,
        color: typeof msg.color === 'number' ? msg.color : parseInt(msg.color, 16) || 0x00ff88,
        role: msg.role || 'agent',
        direction: msg.direction || 'north',
        x: msg.x ?? null,
        z: msg.z ?? null,
      };

      // addAgent handles placement, scene insertion, and connection lines
      const added = addAgent(newDef);
      if (added) {
        // Update local lookup for future chat messages
        agentById[newDef.id] = newDef;
        logEvent(`Agent ${newDef.label} joined the swarm`);
      }
      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');
}