State Management (Zustand)

State Management (Zustand)

Client-side stores for the MetaForge Digital Twin Dashboard. Each store is a self-contained Zustand slice with typed state, actions, and optional localStorage persistence.

src/store/ui-store.ts

Controls sidebar visibility, theme preference, and the command palette. Persisted to localStorage so preferences survive page reloads.

import { create } from "zustand";
import { persist, createJSONStorage } from "zustand/middleware";

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export type Theme = "light" | "dark" | "system";

export interface UIState {
  /** Whether the sidebar drawer is open (mobile) or visible (desktop). */
  sidebarOpen: boolean;
  /** When true the sidebar renders in narrow icon-only mode. */
  sidebarCollapsed: boolean;
  /** Active colour-scheme preference. */
  theme: Theme;
  /** Whether the Cmd+K command palette is showing. */
  commandPaletteOpen: boolean;
}

export interface UIActions {
  toggleSidebar: () => void;
  setSidebarCollapsed: (collapsed: boolean) => void;
  setTheme: (theme: Theme) => void;
  toggleCommandPalette: () => void;
}

export type UIStore = UIState & UIActions;

// ---------------------------------------------------------------------------
// Resolved theme helper
// ---------------------------------------------------------------------------

/**
 * Resolve the effective theme, respecting the `system` option by reading
 * the user's OS preference via `prefers-color-scheme`.
 */
export function resolveTheme(theme: Theme): "light" | "dark" {
  if (theme !== "system") return theme;
  if (typeof window === "undefined") return "light";
  return window.matchMedia("(prefers-color-scheme: dark)").matches
    ? "dark"
    : "light";
}

// ---------------------------------------------------------------------------
// Store
// ---------------------------------------------------------------------------

export const useUIStore = create<UIStore>()(
  persist(
    (set) => ({
      // -- state defaults ---------------------------------------------------
      sidebarOpen: true,
      sidebarCollapsed: false,
      theme: "system",
      commandPaletteOpen: false,

      // -- actions ----------------------------------------------------------
      toggleSidebar: () =>
        set((s) => ({ sidebarOpen: !s.sidebarOpen })),

      setSidebarCollapsed: (collapsed) =>
        set({ sidebarCollapsed: collapsed }),

      setTheme: (theme) => {
        // Apply the resolved class to <html> immediately so Tailwind picks it up.
        const resolved = resolveTheme(theme);
        document.documentElement.classList.remove("light", "dark");
        document.documentElement.classList.add(resolved);
        set({ theme });
      },

      toggleCommandPalette: () =>
        set((s) => ({ commandPaletteOpen: !s.commandPaletteOpen })),
    }),
    {
      name: "metaforge-ui",
      storage: createJSONStorage(() => localStorage),
      partialize: (state) => ({
        sidebarCollapsed: state.sidebarCollapsed,
        theme: state.theme,
      }),
    },
  ),
);

src/store/websocket-store.ts

Manages a single WebSocket connection to the Gateway, including automatic reconnection with exponential back-off.

import { create } from "zustand";

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export type WSStatus = "connecting" | "connected" | "disconnected" | "error";

export interface WebSocketMessage<T = unknown> {
  event: string;
  data: T;
  timestamp?: string;
}

export interface WebSocketState {
  status: WSStatus;
  socket: WebSocket | null;
  reconnectAttempts: number;
  lastMessage: WebSocketMessage | null;
}

export interface WebSocketActions {
  connect: (url: string) => void;
  disconnect: () => void;
  send: <T>(data: T) => void;
  setStatus: (status: WSStatus) => void;
}

export type WebSocketStore = WebSocketState & WebSocketActions;

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const INITIAL_BACKOFF_MS = 1_000;
const MAX_BACKOFF_MS = 30_000;
const BACKOFF_MULTIPLIER = 2;

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/** Calculate exponential back-off delay capped at MAX_BACKOFF_MS. */
function backoffDelay(attempt: number): number {
  return Math.min(INITIAL_BACKOFF_MS * BACKOFF_MULTIPLIER ** attempt, MAX_BACKOFF_MS);
}

// ---------------------------------------------------------------------------
// Store
// ---------------------------------------------------------------------------

let reconnectTimer: ReturnType<typeof setTimeout> | null = null;

export const useWebSocketStore = create<WebSocketStore>()((set, get) => ({
  // -- state defaults -------------------------------------------------------
  status: "disconnected",
  socket: null,
  reconnectAttempts: 0,
  lastMessage: null,

  // -- actions --------------------------------------------------------------

  connect: (url: string) => {
    const { socket: existing } = get();
    // Prevent duplicate connections
    if (
      existing &&
      (existing.readyState === WebSocket.OPEN ||
        existing.readyState === WebSocket.CONNECTING)
    ) {
      return;
    }

    set({ status: "connecting" });

    const ws = new WebSocket(url);

    ws.addEventListener("open", () => {
      set({ status: "connected", socket: ws, reconnectAttempts: 0 });
    });

    ws.addEventListener("message", (event: MessageEvent) => {
      try {
        const parsed: WebSocketMessage = JSON.parse(event.data as string);
        set({ lastMessage: parsed });
      } catch {
        // Non-JSON frames are silently ignored.
      }
    });

    ws.addEventListener("close", () => {
      set({ status: "disconnected", socket: null });

      // --- auto-reconnect with exponential back-off -----------------------
      const attempts = get().reconnectAttempts;
      const delay = backoffDelay(attempts);

      if (reconnectTimer) clearTimeout(reconnectTimer);

      reconnectTimer = setTimeout(() => {
        set({ reconnectAttempts: attempts + 1 });
        get().connect(url);
      }, delay);
    });

    ws.addEventListener("error", () => {
      set({ status: "error" });
      // The subsequent "close" event will trigger reconnection.
    });

    set({ socket: ws });
  },

  disconnect: () => {
    if (reconnectTimer) {
      clearTimeout(reconnectTimer);
      reconnectTimer = null;
    }

    const { socket } = get();
    if (socket) {
      // Remove all listeners before closing to prevent auto-reconnect.
      socket.onclose = null;
      socket.onerror = null;
      socket.onmessage = null;
      socket.close();
    }

    set({
      status: "disconnected",
      socket: null,
      reconnectAttempts: 0,
      lastMessage: null,
    });
  },

  send: <T>(data: T) => {
    const { socket, status } = get();
    if (status !== "connected" || !socket) {
      console.warn("[ws] Cannot send — socket is not connected.");
      return;
    }
    socket.send(JSON.stringify(data));
  },

  setStatus: (status: WSStatus) => set({ status }),
}));

chat-store (sidebar visibility, streaming content, typing indicators) is specified in Agent Chat Channel.