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",
]);

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 || !ALLOWED_EVENTS.has(result.data.event)) return;
    const { event: name, payload } = result.data;
    // 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"] });
      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();
  };
}