Cyrene_v2/frontend/web/src/lib/ws.ts

import { useAppStore } from "../store/useAppStore";
import { useModeStore } from "../store/useModeStore";
import { getAuthToken, isTokenExpired, onAuthExpired, setAuthToken } from "./api";
import { toRecord } from "./wsParse";
import type { WsUpEvent } from "./types";
import { handleStreamEvent } from "./ws/handlers/stream";
import { handleThinkingEvent } from "./ws/handlers/thinking";
import { handleToolEvent } from "./ws/handlers/tool";
import { handleMessageEvent } from "./ws/handlers/message";
import { handleAgentEvent } from "./ws/handlers/agent";
import { handleSessionEvent } from "./ws/handlers/session";
import { handleDeviceEvent } from "./ws/handlers/device";
import { handleApprovalEvent } from "./ws/handlers/approval";

/**
 * Realtime client for `/ws/chat`. Reconnects with exponential backoff and
 * decodes the downstream event stream into the app store. All `data` payloads
 * are `unknown` and narrowed here; nothing leaks `any` into components.
 */

let ws: WebSocket | null = null;
let shouldRun = false;
let reconnectTimer: number | null = null;
let attempt = 0;

const MAX_BACKOFF = 30_000;
const BASE_BACKOFF = 1_000;

function wsUrl(): string {
  const proto = window.location.protocol === "https:" ? "wss" : "ws";
  const token = getAuthToken() ?? "";
  const qs = token ? `?token=${encodeURIComponent(token)}` : "";
  return `${proto}://${window.location.host}/ws/chat${qs}`;
}

/** Route one decoded downstream frame into the store. Exported (with a `frame`
 *  shaped first argument) so the frame→store mapping can be exercised as a pure
 *  logic test without a browser: this project has no jsdom/test runner, so the
 *  socket lifecycle around it stays untested while the mapping does not have to.
 *  帧按类型分组在 lib/ws/handlers/* 中处理：任一组命中（含组内主动丢弃）即结束，
 *  全部未命中走下方未知帧兜底 —— 与拆分前的单个 switch 等价。 */
export function dispatchFrame(ev: { type: string; sessionId?: string; data: unknown }): void {
  const s = useAppStore.getState();
  const r = toRecord(ev.data);

  if (handleStreamEvent(ev, r, s)) return;
  if (handleThinkingEvent(ev, r, s)) return;
  if (handleToolEvent(ev, r, s)) return;
  if (handleMessageEvent(ev, r, s)) return;
  if (handleAgentEvent(ev, r, s)) return;
  if (handleSessionEvent(ev, r, s)) return;
  if (handleDeviceEvent(ev, r, s)) return;
  if (handleApprovalEvent(ev, r, s)) return;

  // Unknown future event: keep it in the context log so no frame is lost.
  const text = typeof ev.data === "string" ? ev.data : JSON.stringify(ev.data);
  s.addDeviceUpdate(`[${ev.type}] ${text}`);
}

function scheduleReconnect(): void {
  if (!shouldRun) return;
  if (reconnectTimer !== null) return;
  // Only a live, non-expired credential may retry; otherwise this isn't a network
  // blip but a dead session we should surface, not loop on.
  const token = getAuthToken();
  if (!token || isTokenExpired(token)) {
    markAuthExpired();
    return;
  }
  const delay = Math.min(BASE_BACKOFF * 2 ** attempt, MAX_BACKOFF);
  attempt += 1;
  useAppStore.getState().setWsStatus("reconnecting");
  reconnectTimer = window.setTimeout(() => {
    reconnectTimer = null;
    openSocket();
  }, delay);
}

function openSocket(): void {
  const s = useAppStore.getState();
  // A locally-expired token means the handshake will be rejected with 401 before
  // it ever upgrades; don't even attempt it — flag the expiry instead.
  const token = getAuthToken();
  if (token && isTokenExpired(token)) {
    markAuthExpired();
    return;
  }
  try {
    ws = new WebSocket(wsUrl());
  } catch {
    s.setWsStatus("disconnected");
    scheduleReconnect();
    return;
  }

  // Tracks whether this socket ever completed the handshake (reached OPEN). A
  // browser surfaces a handshake 401 only as onerror+onclose (code 1006, no HTTP
  // status), so the only way to tell "auth rejected" from "network blip" is to
  // compare the staged close against the token's own expiry.
  let opened = false;
  s.setWsStatus(attempt === 0 ? "connecting" : "reconnecting");

  ws.onopen = () => {
    opened = true;
    attempt = 0;
    useAppStore.getState().setWsStatus("connected");
    // Sync the current mode on every (re)connect so the agent's shell matches.
    const mode = useModeStore.getState().mode;
    sendUp({ type: "mode_switch", data: { mode } });
  };

  ws.onmessage = (msg) => {
    try {
      const parsed = JSON.parse(String(msg.data)) as unknown;
      const r = toRecord(parsed);
      if (!r) return;
      dispatchFrame({
        type: typeof r.type === "string" ? r.type : "",
        // The frame envelope's sessionId is what scopes a payload that does not
        // carry its own (see the `message_deleted` case in dispatch).
        sessionId: typeof r.sessionId === "string" ? r.sessionId : undefined,
        data: r.data,
      });
    } catch {
      // ignore malformed frames
    }
  };

  ws.onerror = () => {
    // Handshake/auth rejects land here first. An expired credential is not a
    // transient outage, so stop the loop and prompt re-login rather than looping.
    const cur = getAuthToken();
    if (cur && isTokenExpired(cur)) {
      markAuthExpired();
      return;
    }
    useAppStore.getState().setWsStatus("disconnected");
  };

  ws.onclose = () => {
    ws = null;
    const cur = getAuthToken();
    if (!opened && cur && isTokenExpired(cur)) {
      markAuthExpired();
      return;
    }
    useAppStore.getState().setWsStatus("disconnected");
    scheduleReconnect();
  };
}

/** Open (or reopen) the realtime channel. Safe to call repeatedly. */
export function connectWs(): void {
  if (shouldRun && (ws?.readyState === WebSocket.OPEN || ws?.readyState === WebSocket.CONNECTING)) {
    return;
  }
  shouldRun = true;
  openSocket();
}

/** Close the channel and stop reconnecting. */
export function disconnectWs(): void {
  shouldRun = false;
  if (reconnectTimer !== null) {
    window.clearTimeout(reconnectTimer);
    reconnectTimer = null;
  }
  const cur = ws;
  ws = null;
  if (cur) {
    cur.onclose = null;
    cur.close();
  }
  useAppStore.getState().setWsStatus("idle");
}

/** Terminal "auth expired" path: stop reconnecting, drop the stale token, and
 *  flag the UI so it prompts re-login instead of burning backoff on a dead
 *  credential. Fired from (a) a handshake reject on an expired token and (b) the
 *  REST 401 handler (via onAuthExpired). */
function markAuthExpired(): void {
  setAuthToken(null);
  disconnectWs();
  useAppStore.getState().setWsStatus("expired");
}

/** Send an upstream frame; returns false when the socket is not open. */
export function sendUp(ev: WsUpEvent): boolean {
  const cur = ws;
  if (!cur || cur.readyState !== WebSocket.OPEN) return false;
  cur.send(JSON.stringify(ev));
  return true;
}

/** Is the realtime channel open? */
export function isWsOpen(): boolean {
  return ws?.readyState === WebSocket.OPEN;
}

// When the REST layer hits a 401/unauthorized on a logged-in call it clears the
// token and fires this handler; mirror it into the same terminal expiry path so
// the reconnect loop stops and the UI shows "登录过期" for both WS and REST 401s.
onAuthExpired(() => {
  markAuthExpired();
});