AetherForge/server/web/src/context/WebSocketProvider.tsx

import React, { useEffect, useRef, useCallback, useState, useMemo } from 'react';
import { WS_LATEST_MESSAGE_TYPES } from '../help/wsStatsCoalesce';
import { applyStatsUpdates } from '../help/applyStatsUpdate';
import type {
  WSDashboardInit,
  WSAgentOffline,
  WSStatsUpdate,
  WSStatsBatch,
  WSCommandResult,
  WSAgentLog,
  WSPolicyAck,
} from '../types/ws';
import type { Agent, Share, FleetAlert, PoolStatus, AIActivityEntry, WSMessage } from '../types';
import { WebSocketContext } from './WebSocketContext';
import type { SeqCommandResult } from './WebSocketContext';
import type { SeqPolicyAck } from './WebSocketContext';
import { authHeaders, getStoredAuth } from '../api/auth';

/**
 * WebSocketProvider mounts a SINGLE WebSocket connection for the whole app.
 * All components call useWebSocket() and receive data from this one connection
 * — fixes M12 (duplicate connections when multiple components called the hook).
 */
export function WebSocketProvider({ children }: { children: React.ReactNode }) {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const unmounted = useRef(false);
  // Guard: true while an openSocket() invocation is in-flight (awaiting ticket fetch).
  // A second connect() call while one is already opening is a no-op — prevents duplicate
  // sockets when 'aetherforge-auth' fires during the async fetch.
  const openingRef = useRef(false);
  // AbortController for the current in-flight ws-ticket fetch; replaced on each open attempt.
  const ticketAbortRef = useRef<AbortController | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [agents, setAgents] = useState<Agent[]>([]);
  const [recentShares, setRecentShares] = useState<Share[]>([]);
  const [fleetAlerts, setFleetAlerts] = useState<FleetAlert[]>([]);
  const [poolStatus, setPoolStatus] = useState<PoolStatus[]>([]);
  const [aiActivity, setAiActivity] = useState<AIActivityEntry[]>([]);
  const [agentLogs, setAgentLogs] = useState<Record<string, string>>({});
  const [commandResults, setCommandResults] = useState<SeqCommandResult[]>([]);
  const [policyAcks, setPolicyAcks] = useState<SeqPolicyAck[]>([]);
  const [latestMessage, setLatestMessage] = useState<WSMessage | null>(null);
  // Monotonic counter so consumers can detect new entries even after the ring buffer trims old ones
  const cmdSeqRef = useRef(0);
  const policyAckSeqRef = useRef(0);

  const sendDashboardMessage = useCallback((type: string, payload: Record<string, unknown>) => {
    const ws = wsRef.current;
    if (!ws || ws.readyState !== WebSocket.OPEN) return;
    ws.send(JSON.stringify({ type, payload }));
  }, []);

  const connect = useCallback(() => {
    if (unmounted.current) return;

    if (reconnectTimer.current) {
      clearTimeout(reconnectTimer.current);
      reconnectTimer.current = null;
    }

    const token = getStoredAuth();
    if (!token) {
      const existing = wsRef.current;
      if (existing) {
        existing.onclose = null;
        existing.close();
        wsRef.current = null;
      }
      setIsConnected(false);
      return;
    }

    const existing = wsRef.current;
    if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) {
      existing.onclose = null;
      existing.close();
    }

    const openSocket = async () => {
      if (openingRef.current) return;
      openingRef.current = true;

      // Cancel any previously in-flight ticket fetch before starting a new one.
      if (ticketAbortRef.current) {
        ticketAbortRef.current.abort();
      }
      const abortCtrl = new AbortController();
      ticketAbortRef.current = abortCtrl;

      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      let wsQuery = `token=${encodeURIComponent(token)}`;
      try {
        const resp = await fetch('/api/v1/auth/ws-ticket', {
          method: 'POST',
          headers: authHeaders(),
          signal: abortCtrl.signal,
        });
        if (resp.ok) {
          const data = (await resp.json()) as { ticket?: string };
          if (data.ticket) {
            wsQuery = `ticket=${encodeURIComponent(data.ticket)}`;
          }
        }
      } catch {
        /* fall back to legacy token query param — also catches AbortError */
      } finally {
        if (ticketAbortRef.current === abortCtrl) {
          ticketAbortRef.current = null;
        }
        openingRef.current = false;
      }
      if (unmounted.current) return;
      // If this invocation was aborted by a newer call, bail out — the newer
      // call will (or already has) opened its own socket.
      if (abortCtrl.signal.aborted) return;

      const wsUrl = `${protocol}//${window.location.host}/ws/dashboard?${wsQuery}`;
      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onopen = () => { if (!unmounted.current) setIsConnected(true); };

      ws.onclose = () => {
        if (unmounted.current) return;
        setIsConnected(false);
        if (reconnectTimer.current) clearTimeout(reconnectTimer.current);
        if (!getStoredAuth()) return;
        reconnectTimer.current = setTimeout(connect, 3000);
      };

      ws.onerror = () => { ws.close(); };

    ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data) as WSMessage;
        if (WS_LATEST_MESSAGE_TYPES.has(msg.type)) {
          setLatestMessage(msg);
        }

        switch (msg.type) {
          case 'init': {
            const data = msg.payload as WSDashboardInit;
            if (data.agents) setAgents(data.agents);
            break;
          }
          case 'agent_online': {
            const agent = msg.payload as Agent;
            setAgents((prev) => {
              const idx = prev.findIndex((a) => a.id === agent.id);
              if (idx >= 0) {
                const updated = [...prev];
                updated[idx] = { ...updated[idx], ...agent };
                return updated;
              }
              return [...prev, agent];
            });
            break;
          }
          case 'agent_offline': {
            const { agent_id } = msg.payload as WSAgentOffline;
            setAgents((prev) =>
              prev.map((a) => a.id === agent_id ? { ...a, status: 'offline' as const } : a)
            );
            break;
          }
          case 'agent_removed':
          case 'agent_deleted': {
            const { agent_id } = msg.payload as { agent_id: string };
            setAgents((prev) => prev.filter((a) => a.id !== agent_id));
            // Also purge any cached alerts for this agent so the dashboard
            // alert banner stops showing errors for machines that no longer exist.
            setFleetAlerts((prev) => prev.filter((a) => a.agent_id !== agent_id));
            break;
          }
          case 'stats_update': {
            const update = msg.payload as WSStatsUpdate;
            setAgents((prev) => applyStatsUpdates(prev, [update]));
            break;
          }
          case 'stats_batch': {
            const batch = msg.payload as WSStatsBatch;
            if (Array.isArray(batch?.updates) && batch.updates.length > 0) {
              setAgents((prev) => applyStatsUpdates(prev, batch.updates));
            }
            break;
          }
          case 'clearance_elevated': {
            const p = msg.payload as {
              agent_id: string;
              from_level: number;
              to_level: number;
              reason?: string;
              source?: string;
            };
            setAgents((prev) =>
              prev.map((a) =>
                a.id === p.agent_id ? { ...a, clearance_level: p.to_level } : a,
              ),
            );
            break;
          }
          case 'new_share': {
            const share = msg.payload as Share;
            setRecentShares((prev) => [share, ...prev].slice(0, 50));
            break;
          }
          case 'fleet_alert': {
            const alert = msg.payload as FleetAlert;
            setFleetAlerts((prev) => [alert, ...prev].slice(0, 20));
            break;
          }
          case 'pool_status': {
            const pools = msg.payload as PoolStatus[];
            if (Array.isArray(pools)) setPoolStatus(pools);
            break;
          }
          case 'ai_activity': {
            const entry = msg.payload as AIActivityEntry;
            setAiActivity((prev) => {
              const idx = prev.findIndex((a) => a.agent_id === entry.agent_id);
              if (idx >= 0) {
                const next = [...prev];
                next[idx] = entry;
                return next;
              }
              return [...prev, entry];
            });
            break;
          }
          case 'command_result': {
            // Payload must be an object — a double-encoded string would spread to
            // char indices and drop agent_id, breaking Crucible terminal routing.
            let p = msg.payload as WSCommandResult | string;
            if (typeof p === 'string') {
              try {
                p = JSON.parse(p) as WSCommandResult;
              } catch {
                break;
              }
            }
            const seq = ++cmdSeqRef.current;
            // Cap at 2000; command results are rare (operator-triggered) so this is plenty.
            // Consumers MUST use _seq for change detection — NOT array index — because the
            // slice trims old entries and makes absolute indices stale.
            setCommandResults((prev) => [...prev, { ...p, _seq: seq }].slice(-2000));
            if (p.agent_id && p.action === 'get_log' && p.success && p.message) {
              setAgentLogs((prev) => ({ ...prev, [p.agent_id!]: p.message! }));
            }
            break;
          }
          case 'policy_ack': {
            const ack = msg.payload as WSPolicyAck;
            policyAckSeqRef.current += 1;
            setPolicyAcks((prev) => [...prev.slice(-49), { ...ack, _seq: policyAckSeqRef.current }]);
            break;
          }
          case 'agent_log': {
            const { agent_id, content } = msg.payload as WSAgentLog;
            if (agent_id) setAgentLogs((prev) => ({ ...prev, [agent_id]: content }));
            break;
          }
          case 'agent_capabilities': {
            const { agent_id, capabilities } = msg.payload as {
              agent_id: string;
              capabilities: Agent['capabilities'];
            };
            if (!agent_id || !capabilities) break;
            setAgents((prev) =>
              prev.map((a) =>
                a.id === agent_id
                  ? {
                      ...a,
                      capabilities: { ...a.capabilities, ...capabilities },
                    }
                  : a
              )
            );
            break;
          }
        }
      } catch (err) {
        console.error('Failed to parse WebSocket message:', err);
      }
    };
    };

    void openSocket();
  }, []);

  useEffect(() => {
    unmounted.current = false;
    connect();
    const onAuthChange = () => connect();
    window.addEventListener('aetherforge-auth', onAuthChange);
    return () => {
      unmounted.current = true;
      window.removeEventListener('aetherforge-auth', onAuthChange);
      if (reconnectTimer.current) clearTimeout(reconnectTimer.current);
      const ws = wsRef.current;
      if (ws) { ws.onclose = null; ws.close(); }
    };
  }, [connect]);

  const value = useMemo(
    () => ({
      isConnected,
      agents,
      recentShares,
      fleetAlerts,
      poolStatus,
      aiActivity,
      agentLogs,
      commandResults,
      policyAcks,
      latestMessage,
      sendDashboardMessage,
    }),
    [
      isConnected,
      agents,
      recentShares,
      fleetAlerts,
      poolStatus,
      aiActivity,
      agentLogs,
      commandResults,
      policyAcks,
      latestMessage,
      sendDashboardMessage,
    ],
  );

  return (
    <WebSocketContext.Provider value={value}>
      {children}
    </WebSocketContext.Provider>
  );
}