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 socket = new WebSocket(`${proto}://${window.location.host}/api/ws`);

  const ping = window.setInterval(() => {
    if (socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify({ type: "ping" }));
    }
  }, 30000);

  socket.onmessage = (event) => {
    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") {
      queryClient.invalidateQueries({ queryKey: ["accounts"] });
    }
    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"] });
    }
  };

  return () => {
    window.clearInterval(ping);
    socket.close();
  };
}