webmail-client/frontend/src/lib/ws.ts

// WebSocket layer (SPEC §9.2). Validates frames with zod and invalidates
// TanStack Query caches keyed on ids in the payload.

import type { QueryClient } from "@tanstack/react-query";
import { z } from "zod";

const frameSchema = z.object({
  event: z.string(),
  payload: z.record(z.unknown()),
});

const ALLOWED_EVENTS = new Set([
  "message.new",
  "message.changed",
  "message.deleted",
  "message.sent",
  "flags.changed",
  "thread.updated",
  "sync.progress",
  "account.state_changed",
  "draft.saved",
  "outbox.state_changed",
]);

// Lightweight in-process fan-out of WS frames to plugin subscribers (SPEC §11.3).
// Plugins receive plugin.{name}.* events plus any core event they subscribe to.
export type WsListener = (event: string, payload: Record<string, unknown>) => void;

const wsListeners = new Set<WsListener>();

export function subscribeWs(listener: WsListener): () => void {
  wsListeners.add(listener);
  return () => {
    wsListeners.delete(listener);
  };
}

export function connectWs(queryClient: QueryClient): () => void {
  const proto = window.location.protocol === "https:" ? "wss" : "ws";
  const url = `${proto}://${window.location.host}/api/ws`;

  let socket: WebSocket | null = null;
  let ping: number | undefined;
  let reconnectTimer: number | undefined;
  let attempts = 0;
  let stopped = false;

  // Refetch the live-updated lists — used on every (re)connect so a dropped socket
  // (server restart, sleep, network blip) self-heals and catches up on missed events.
  const resync = () => {
    for (const key of [
      ["threads"],
      ["accounts"],
      ["folder-counts"],
      ["quota"],
      ["drafts-remote"],
    ]) {
      queryClient.invalidateQueries({ queryKey: key });
    }
  };

  const handleMessage = (event: MessageEvent) => {
    let parsed: unknown;
    try {
      parsed = JSON.parse(event.data);
    } catch {
      return;
    }
    const result = frameSchema.safeParse(parsed);
    if (!result.success) return;
    const { event: name, payload } = result.data;
    const isCore = ALLOWED_EVENTS.has(name);
    // Fan out core + plugin.* events to plugin subscribers (a failing listener
    // must not break the others or the cache invalidation below).
    if (isCore || name.startsWith("plugin.")) {
      for (const listener of wsListeners) {
        try {
          listener(name, payload);
        } catch (err) {
          console.error("ws listener failed", err);
        }
      }
    }
    if (!isCore) return;
    // Coarse invalidation: refetch affected lists (SPEC §9.2 contract).
    if (name === "sync.progress" || name.startsWith("message.") || name === "thread.updated") {
      queryClient.invalidateQueries({ queryKey: ["threads"] });
      queryClient.invalidateQueries({ queryKey: ["quota"] });
      queryClient.invalidateQueries({ queryKey: ["drafts-remote"] });
      queryClient.invalidateQueries({ queryKey: ["folder-counts"] });
      // Backfill progress counters live on the account row (drives the banner).
      if (name === "sync.progress") {
        queryClient.invalidateQueries({ queryKey: ["accounts"] });
      }
      if (typeof payload.thread_id === "string") {
        queryClient.invalidateQueries({ queryKey: ["thread", payload.thread_id] });
      }
    }
    if (name === "flags.changed") {
      // Read/unread changes affect the Inbox unread badge.
      queryClient.invalidateQueries({ queryKey: ["folder-counts"] });
    }
    if (name === "account.state_changed") {
      // Adding/deleting an account changes threads and (Gmail) labels wholesale —
      // refetch the list, chips, and sidebar so nothing stale lingers.
      queryClient.invalidateQueries({ queryKey: ["accounts"] });
      queryClient.invalidateQueries({ queryKey: ["threads"] });
      queryClient.invalidateQueries({ queryKey: ["thread-labels"] });
      queryClient.invalidateQueries({ queryKey: ["message-labels"] });
      queryClient.invalidateQueries({ queryKey: ["labels-all"] });
      queryClient.invalidateQueries({ queryKey: ["folder-counts"] });
      window.dispatchEvent(new Event("webmail:labels-changed"));
    }
    if (name === "message.sent") {
      // The sent message's draft is deleted server-side; refresh the Drafts list.
      queryClient.invalidateQueries({ queryKey: ["drafts"] });
      queryClient.invalidateQueries({ queryKey: ["folder-counts"] });
    }
  };

  const connect = () => {
    socket = new WebSocket(url);
    socket.onopen = () => {
      attempts = 0;
      ping = window.setInterval(() => {
        if (socket?.readyState === WebSocket.OPEN) {
          socket.send(JSON.stringify({ type: "ping" }));
        }
      }, 30000);
      resync(); // catch up on anything published while we were disconnected
    };
    socket.onmessage = handleMessage;
    socket.onclose = () => {
      if (ping !== undefined) window.clearInterval(ping);
      if (stopped) return;
      // Exponential backoff, capped at 30s, so a downed server is retried politely.
      const delay = Math.min(30000, 1000 * 2 ** attempts);
      attempts += 1;
      reconnectTimer = window.setTimeout(connect, delay);
    };
    socket.onerror = () => socket?.close(); // route errors through onclose -> reconnect
  };

  connect();

  return () => {
    stopped = true;
    if (ping !== undefined) window.clearInterval(ping);
    if (reconnectTimer !== undefined) window.clearTimeout(reconnectTimer);
    socket?.close();
  };
}