ai-agent/frontend/src/store.ts

import { create } from "zustand";
import { createJSONStorage, persist } from "zustand/middleware";
import { postNotifSeen, saveFile } from "./api";
import { QK } from "./lib/queries";
import { queryClient } from "./lib/queryClient";
import { retryHydration, serverStorage } from "./lib/serverStorage";

/**
 * Small UI state + notification bookkeeping. Server data (the file bundle, the
 * dashboard, conversations…) lives in React Query (`lib/queries.ts`) — cached,
 * deduped and persisted to IndexedDB there. This store keeps only what is
 * genuinely client state: unsaved drafts, the tree search query, the sidebar,
 * and the notification feed's seen/new lists.
 *
 * Persisted server-side (a file via `/api/ui-state`) instead of localStorage,
 * so it follows the user across browsers/devices. `serverStorage` swallows
 * read/write failures so a network blip can never crash a React state update.
 * Hydration is async — `syncNotifs` is gated on
 * `useStore.persist.hasHydrated()` so it never seeds the notification baseline
 * from defaults before the server value lands.
 */
const safeStorage = createJSONStorage(() => serverStorage);

interface State {
  /** Unsaved edits, keyed by file path. */
  edits: Record<string, string>;
  /** File-tree search query. */
  query: string;
  sidebarOpen: boolean;
  online: boolean;
  saving: boolean;
  error: string | null;

  /** Every notification id we've ever observed (the baseline for "new"). */
  notifKnown: string[];
  /** Ids still flagged "new" — shown on the home feed until swiped away. */
  notifNew: string[];
  /** Whether the "known" baseline has been seeded from a first fetch yet. */
  notifSeeded: boolean;

  setQuery: (q: string) => void;
  setSidebar: (open: boolean) => void;
  setOnline: (online: boolean) => void;
  setDraft: (path: string, content: string) => void;
  revert: (path: string) => void;
  save: (path: string) => Promise<void>;
  /** Reconcile the live set of notification ids into the known/new lists. */
  syncNotifs: (ids: string[]) => void;
  /** Mark a notification seen (swiped) — drop it from the "new" list. */
  dismissNotif: (id: string) => void;
}

export const useStore = create<State>()(
  persist(
    (set, get) => ({
      edits: {},
      query: "",
      sidebarOpen: true,
      online: navigator.onLine,
      saving: false,
      error: null,

      notifKnown: [],
      notifNew: [],
      notifSeeded: false,

      setQuery: (query) => set({ query }),
      setSidebar: (open) => set({ sidebarOpen: open }),
      setOnline: (online) => set({ online }),

      setDraft: (path, content) =>
        set((s) => ({ edits: { ...s.edits, [path]: content } })),

      revert: (path) =>
        set((s) => {
          const edits = { ...s.edits };
          delete edits[path];
          return { edits };
        }),

      save: async (path) => {
        const draft = get().edits[path];
        if (draft === undefined) return;
        set({ saving: true, error: null });
        try {
          const updated = await saveFile(path, draft);
          set((s) => {
            const edits = { ...s.edits };
            delete edits[path];
            return { saving: false, edits };
          });
          // Reflect the save into the React Query caches: the editor's file
          // query gets the fresh entry; the bundle's stats revalidate lazily.
          queryClient.setQueryData(QK.file(path), updated);
          queryClient.invalidateQueries({ queryKey: QK.bundle });
        } catch (e) {
          set({ saving: false, error: String(e) });
          throw e;
        }
      },

      syncNotifs: (ids) =>
        set((s) => {
          const known = new Set(s.notifKnown);
          const fresh = ids.filter((id) => !known.has(id));
          // First reconcile ever: adopt everything as the baseline so the
          // history doesn't flood the home feed as "new". Only notifications
          // that arrive *after* this baseline count as new. Mirror the baseline
          // to the server read-ledger too, so the phone's "unread" agrees.
          if (!s.notifSeeded) {
            void postNotifSeen(ids, { seed: true });
            return { notifKnown: ids, notifNew: [], notifSeeded: true };
          }
          if (fresh.length === 0) return {};
          return {
            notifKnown: [...s.notifKnown, ...fresh],
            notifNew: [...s.notifNew, ...fresh],
          };
        }),

      dismissNotif: (id) => {
        // Opening/swiping a card marks it read here AND on the server, so the
        // phone stops reading it out ("listening" on the phone never does this).
        void postNotifSeen([id]);
        set((s) => ({ notifNew: s.notifNew.filter((n) => n !== id) }));
      },
    }),
    {
      name: "ai-agent-store",
      storage: safeStorage,
      partialize: (s) => ({
        edits: s.edits,
        sidebarOpen: s.sidebarOpen,
        notifKnown: s.notifKnown,
        notifNew: s.notifNew,
        notifSeeded: s.notifSeeded,
      }),
    },
  ),
);
retryHydration(useStore);