agent-runtime/static/js/ws.js

/** WebSocket connections: /ws (chat), /ws/test, /ws/trace. */

import { authToken, isAuthFailed, setAuthFailed, showLogin } from './auth.js';
import { addTrace } from './trace.js';
import { addMsg, handleDelta, handleDone, setWs as setChatWs } from './chat.js';
import { dockControls, dockArtifacts, setWs as setDashWs } from './dashboard.js';
import { graphAnimate } from './graph.js';
import { updateMeter, updateNodeFromHud, updateAwarenessState, updateAwarenessSensors } from './awareness.js';
import { updateTestStatus } from './tests.js';
import { truncate, esc } from './util.js';

let ws, wsTest, wsTrace;
let _testPollInterval = null;
let _lastTestResultCount = 0;

export function connect() {
  if (isAuthFailed()) return;
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
  let wsUrl = proto + '//' + location.host + '/ws';
  if (authToken) {
    const accessToken = localStorage.getItem('cog_access_token') || '';
    wsUrl += '?token=' + encodeURIComponent(authToken) + '&access_token=' + encodeURIComponent(accessToken);
  }
  ws = new WebSocket(wsUrl);

  ws.onopen = () => {
    document.getElementById('status').textContent = 'connected';
    document.getElementById('status').style.color = '#22c55e';
    addTrace('runtime', 'connected', 'ws open');
    setChatWs(ws);
    setDashWs(ws);
    connectDebugSockets();
    restoreHistory();
  };

  ws.onerror = () => {};

  ws.onclose = (e) => {
    // 4001 = explicit auth rejection from server
    if (e.code === 4001) {
      setAuthFailed(true);
      localStorage.removeItem('cog_token');
      localStorage.removeItem('cog_access_token');
      document.getElementById('status').textContent = 'session expired';
      document.getElementById('status').style.color = '#ef4444';
      showLogin();
      return;
    }
    // 1006 = abnormal close (deploy, network), just reconnect
    document.getElementById('status').textContent = 'reconnecting...';
    document.getElementById('status').style.color = '#f59e0b';
    addTrace('runtime', 'disconnected', `code ${e.code}, reconnecting...`);
    setTimeout(connect, 2000);
  };

  ws.onmessage = (e) => {
    const data = JSON.parse(e.data);
    if (data.type === 'hud') {
      handleHud(data);
    } else if (data.type === 'delta') {
      handleDelta(data.content);
    } else if (data.type === 'done') {
      handleDone();
    } else if (data.type === 'artifacts') {
      dockArtifacts(data.artifacts);
    } else if (data.type === 'controls') {
      dockControls(data.controls);
    } else if (data.type === 'cleared') {
      addTrace('runtime', 'cleared', 'session reset');
    }
  };
}

async function restoreHistory() {
  try {
    const headers = {};
    if (authToken) headers['Authorization'] = 'Bearer ' + authToken;
    const r = await fetch('/api/history?last=20', { headers });
    if (!r.ok) return;
    const data = await r.json();
    const messages = data.messages || [];
    if (!messages.length) return;
    // Only restore if chat is empty (fresh load)
    if (document.getElementById('messages').children.length > 0) return;
    for (const msg of messages) {
      const el = addMsg(msg.role, '');
      if (msg.role === 'assistant') {
        // Render as markdown
        const { renderMarkdown } = await import('./util.js');
        el.innerHTML = renderMarkdown(msg.content || '');
      } else {
        el.textContent = msg.content || '';
      }
    }
    addTrace('runtime', 'restored', `${messages.length} messages`);
  } catch (e) {}
}

function connectDebugSockets() {
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
  const base = proto + '//' + location.host;
  const tokenParam = authToken ? '?token=' + encodeURIComponent(authToken) : '';

  if (!wsTest || wsTest.readyState > 1) {
    wsTest = new WebSocket(base + '/ws/test' + tokenParam);
    wsTest.onopen = () => addTrace('runtime', 'ws/test', 'connected');
    wsTest.onclose = () => setTimeout(connectDebugSockets, 3000);
    wsTest.onerror = () => {};
    wsTest.onmessage = (e) => {
      const data = JSON.parse(e.data);
      if (data.type === 'test_status') updateTestStatus(data);
    };
  }

  if (!wsTrace || wsTrace.readyState > 1) {
    wsTrace = new WebSocket(base + '/ws/trace' + tokenParam);
    wsTrace.onopen = () => addTrace('runtime', 'ws/trace', 'connected');
    wsTrace.onclose = () => {};
    wsTrace.onerror = () => {};
    wsTrace.onmessage = (e) => {
      const data = JSON.parse(e.data);
      if (data.event === 'frame_trace' && data.trace) {
        const t = data.trace;
        const frames = t.frames || [];
        const summary = frames.map(f => `F${f.frame}:${f.node}(${f.duration_ms}ms)`).join(' > ');
        addTrace('frame_engine', 'trace', `${t.path} ${t.total_frames}F ${t.total_ms}ms`, 'instruction',
          summary + '\n' + JSON.stringify(t, null, 2));
      } else if (data.node && data.event) {
        handleHud(data);
      }
    };
  }

  // Polling fallback for test status
  if (!_testPollInterval) {
    _testPollInterval = setInterval(async () => {
      try {
        const headers = authToken ? { 'Authorization': 'Bearer ' + authToken } : {};
        const r = await fetch('/api/test/status', { headers });
        const data = await r.json();
        const count = (data.results || []).length;
        if (count !== _lastTestResultCount || data.running) {
          _lastTestResultCount = count;
          updateTestStatus(data);
        }
      } catch (e) {}
    }, 500);
  }
}

function handleHud(data) {
  const node = data.node || 'unknown';
  const event = data.event || '';

  graphAnimate(event, node);
  updateNodeFromHud(node, event, data);

  if (event === 'context') {
    const count = (data.messages || []).length;
    const tokenInfo = data.tokens ? ` [${data.tokens}/${data.max_tokens}t ${data.fill_pct}%]` : '';
    const summary = count + ' msgs' + tokenInfo;
    const detail = (data.messages || []).map((m, i) =>
      i + ' [' + m.role + '] ' + m.content
    ).join('\n');
    addTrace(node, 'context', summary, 'context', detail);
    if (data.tokens !== undefined) updateMeter(node, data.tokens, data.max_tokens, data.fill_pct);

  } else if (event === 'perceived') {
    const text = data.analysis
      ? Object.entries(data.analysis).map(([k,v]) => k + '=' + v).join(' ')
      : '';
    addTrace(node, 'perceived', text, 'instruction', data.analysis ? JSON.stringify(data.analysis, null, 2) : null);

  } else if (event === 'decided' || event === 'routed') {
    addTrace(node, event, data.instruction || data.goal || data.job || '', 'instruction');

  } else if (event === 'updated' && data.state) {
    const pairs = Object.entries(data.state).map(([k, v]) => {
      const val = Array.isArray(v) ? v.join(', ') : (v || 'null');
      return k + '=' + truncate(val, 25);
    }).join('  ');
    addTrace(node, 'state', pairs, 'state', JSON.stringify(data.state, null, 2));
    updateAwarenessState(data.state);

  } else if (event === 'tool_call') {
    const tool = data.tool || '';
    const argsStr = data.args ? JSON.stringify(data.args) : '';
    addTrace(node, 'tool_call', tool + ' ' + truncate(argsStr, 60), '', argsStr);

  } else if (event === 'tool_result') {
    const output = data.output || '';
    addTrace(node, 'tool_result', truncate(output, 80), '', output);

  } else if (event === 'interpreted') {
    addTrace(node, 'interpreted', data.summary || '', 'instruction');

  } else if (event === 'tick') {
    updateAwarenessSensors(data.tick || 0, data.deltas || {});

  } else if (event === 'started' || event === 'stopped') {
    addTrace(node, event, '');

  } else {
    addTrace(node, event, '', '', JSON.stringify(data, null, 2));
  }
}