feat: migrate to real Zustand store (Phase 2)

packages/app/package.json

@@ -63,7 +63,7 @@
     "react-native-web": "~0.21.0",
     "react-native-worklets": "0.5.1",
     "zod": "^3.23.8",
-    "zustand": "^5.0.8"
+    "zustand": "^5.0.9"
   },
   "devDependencies": {
     "@types/react": "~19.2.0",

packages/app/src/hooks/use-daemon-session.ts

@@ -1,8 +1,8 @@
-import { useRef } from "react";
+import { useCallback, useRef } from "react";
 import { Alert } from "react-native";
 import type { SessionContextValue } from "@/contexts/session-context";
 import { useDaemonConnections } from "@/contexts/daemon-connections-context";
-import { useSessionDirectory } from "./use-session-directory";
+import { useSessionStore } from "@/stores/session-store";
 
 export class DaemonSessionUnavailableError extends Error {
   serverId: string;
@@ -14,17 +14,6 @@ export class DaemonSessionUnavailableError extends Error {
   }
 }
 
-export function getSessionForServer(
-  serverId: string,
-  directory: Map<string, SessionContextValue>
-): SessionContextValue {
-  const session = directory.get(serverId) ?? null;
-  if (!session) {
-    throw new DaemonSessionUnavailableError(serverId);
-  }
-  return session;
-}
-
 type UseDaemonSessionOptions = {
   suppressUnavailableAlert?: boolean;
   allowUnavailable?: boolean;
@@ -39,7 +28,16 @@ export function useDaemonSession(
   options: UseDaemonSessionOptions & { allowUnavailable: true }
 ): SessionContextValue | null;
 export function useDaemonSession(serverId?: string | null, options?: UseDaemonSessionOptions) {
-  const sessionDirectory = useSessionDirectory();
+  const selectSession = useCallback(
+    (state: ReturnType<typeof useSessionStore.getState>) => {
+      if (!serverId) {
+        return null;
+      }
+      return state.sessions[serverId] ?? null;
+    },
+    [serverId]
+  );
+  const session = useSessionStore(selectSession);
   const { connectionStates } = useDaemonConnections();
   const alertedDaemonsRef = useRef<Set<string>>(new Set());
   const loggedDaemonsRef = useRef<Set<string>>(new Set());
@@ -50,7 +48,10 @@ export function useDaemonSession(serverId?: string | null, options?: UseDaemonSe
   }
 
   try {
-    return getSessionForServer(serverId, sessionDirectory);
+    if (!session) {
+      throw new DaemonSessionUnavailableError(serverId);
+    }
+    return session;
   } catch (error) {
     if (error instanceof DaemonSessionUnavailableError) {
       const connection = connectionStates.get(serverId);

packages/app/src/stores/session-store.ts

@@ -1,9 +1,9 @@
-import { useSyncExternalStore } from "react";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 import type { SessionContextValue } from "@/contexts/session-context";
 import type { AgentDirectoryEntry } from "@/types/agent-directory";
 import { isPerfLoggingEnabled, measurePayload, perfLog } from "@/utils/perf";
 
-// SessionData mirrors SessionContextValue so consumers can subscribe to a single source of truth.
 export type SessionData = SessionContextValue;
 
 interface SessionStoreState {
@@ -21,24 +21,14 @@ interface SessionStore extends SessionStoreState {
   getAgentDirectory: (serverId: string) => AgentDirectoryEntry[] | undefined;
 }
 
-type SessionListener = () => void;
-
-let storeState: SessionStoreState = { sessions: {}, agentDirectory: {} };
-const listeners = new Set<SessionListener>();
 const SESSION_STORE_LOG_TAG = "[SessionStore]";
 let sessionStoreUpdateCount = 0;
 
-const emit = () => {
-  for (const listener of listeners) {
-    listener();
-  }
-};
-
-const logSessionStoreUpdate = (
+function logSessionStoreUpdate(
   type: "setSession" | "updateSession" | "clearSession" | "setAgentDirectory" | "clearAgentDirectory",
   serverId: string,
   payload?: unknown
-) => {
+) {
   if (!isPerfLoggingEnabled()) {
     return;
   }
@@ -52,128 +42,7 @@ const logSessionStoreUpdate = (
     payloadFieldCount: metrics?.fieldCount ?? 0,
     timestamp: Date.now(),
   });
-};
-
-const updateStoreState = (updater: (prev: SessionStoreState) => SessionStoreState) => {
-  const next = updater(storeState);
-  if (next === storeState) {
-    return;
-  }
-  storeState = next;
-  emit();
-};
-
-const setSession: SessionStore["setSession"] = (serverId, data) => {
-  updateStoreState((prev) => {
-    if (prev.sessions[serverId] === data) {
-      return prev;
-    }
-    logSessionStoreUpdate("setSession", serverId, data);
-    return {
-      ...prev,
-      sessions: {
-        ...prev.sessions,
-        [serverId]: data,
-      },
-    };
-  });
-};
-
-const updateSession: SessionStore["updateSession"] = (serverId, partial) => {
-  updateStoreState((prev) => {
-    const existing = prev.sessions[serverId];
-    const next: SessionData | undefined = existing
-      ? { ...existing, ...partial, serverId }
-      : partial.serverId
-        ? ({ ...(partial as SessionData), serverId })
-        : undefined;
-
-    if (!next) {
-      return prev;
-    }
-
-    if (existing && shallowEqual(existing, next)) {
-      return prev;
-    }
-
-    logSessionStoreUpdate("updateSession", serverId, partial);
-    return {
-      ...prev,
-      sessions: {
-        ...prev.sessions,
-        [serverId]: next,
-      },
-    };
-  });
-};
-
-const clearSession: SessionStore["clearSession"] = (serverId) => {
-  updateStoreState((prev) => {
-    if (!(serverId in prev.sessions)) {
-      return prev;
-    }
-    logSessionStoreUpdate("clearSession", serverId);
-    const nextSessions = { ...prev.sessions };
-    delete nextSessions[serverId];
-    return { ...prev, sessions: nextSessions };
-  });
-};
-
-const getSession: SessionStore["getSession"] = (serverId) => {
-  return storeState.sessions[serverId];
-};
-
-const setAgentDirectory: SessionStore["setAgentDirectory"] = (serverId, agents) => {
-  updateStoreState((prev) => {
-    const existing = prev.agentDirectory[serverId];
-    if (existing && areAgentDirectoriesEqual(existing, agents)) {
-      return prev;
-    }
-    logSessionStoreUpdate("setAgentDirectory", serverId, { agentCount: agents.length });
-    return {
-      ...prev,
-      agentDirectory: {
-        ...prev.agentDirectory,
-        [serverId]: agents,
-      },
-    };
-  });
-};
-
-const clearAgentDirectory: SessionStore["clearAgentDirectory"] = (serverId) => {
-  updateStoreState((prev) => {
-    if (!(serverId in prev.agentDirectory)) {
-      return prev;
-    }
-    logSessionStoreUpdate("clearAgentDirectory", serverId);
-    const nextDirectory = { ...prev.agentDirectory };
-    delete nextDirectory[serverId];
-    return { ...prev, agentDirectory: nextDirectory };
-  });
-};
-
-const getAgentDirectory: SessionStore["getAgentDirectory"] = (serverId) => {
-  return storeState.agentDirectory[serverId];
-};
-
-const subscribe = (listener: SessionListener) => {
-  listeners.add(listener);
-  return () => {
-    listeners.delete(listener);
-  };
-};
-
-const buildSnapshot = (): SessionStore => ({
-  sessions: storeState.sessions,
-  agentDirectory: storeState.agentDirectory,
-  setSession,
-  updateSession,
-  clearSession,
-  getSession,
-  setAgentDirectory,
-  clearAgentDirectory,
-  getAgentDirectory,
-});
+}
 
 const shallowEqual = (left: SessionData, right: SessionData): boolean => {
   if (left === right) {
@@ -192,40 +61,116 @@ const shallowEqual = (left: SessionData, right: SessionData): boolean => {
   return true;
 };
 
-export function useSessionStore<T>(selector: (state: SessionStore) => T): T {
-  return useSyncExternalStore(
-    subscribe,
-    () => selector(buildSnapshot()),
-    () => selector(buildSnapshot())
-  );
-}
-
-function areAgentDirectoriesEqual(
-  left: AgentDirectoryEntry[] | undefined,
-  right: AgentDirectoryEntry[]
-): boolean {
+function areAgentDirectoriesEqual(left: AgentDirectoryEntry[] | undefined, right: AgentDirectoryEntry[]): boolean {
   if (!left) {
     return false;
   }
   if (left.length !== right.length) {
     return false;
   }
-  const leftById = new Map(left.map((entry) => [entry.id, entry]));
-  for (const entry of right) {
-    const previous = leftById.get(entry.id);
-    if (!previous) {
+  return left.every((entry, index) => {
+    const other = right[index];
+    if (!other) {
       return false;
     }
-    if (
-      previous.serverId !== entry.serverId ||
-      previous.title !== entry.title ||
-      previous.status !== entry.status ||
-      previous.provider !== entry.provider ||
-      previous.cwd !== entry.cwd ||
-      previous.lastActivityAt.getTime() !== entry.lastActivityAt.getTime()
-    ) {
-      return false;
-    }
-  }
-  return true;
+    return (
+      entry.id === other.id &&
+      entry.status === other.status &&
+      entry.serverId === other.serverId &&
+      entry.lastActivityAt.getTime() === other.lastActivityAt.getTime()
+    );
+  });
 }
+
+export const useSessionStore = create<SessionStore>()(
+  subscribeWithSelector((set, get) => ({
+    sessions: {},
+    agentDirectory: {},
+    setSession: (serverId, data) => {
+      set((prev) => {
+        if (prev.sessions[serverId] === data) {
+          return prev;
+        }
+        logSessionStoreUpdate("setSession", serverId, data);
+        return {
+          ...prev,
+          sessions: {
+            ...prev.sessions,
+            [serverId]: data,
+          },
+        };
+      });
+    },
+    updateSession: (serverId, partial) => {
+      set((prev) => {
+        const existing = prev.sessions[serverId];
+        const next: SessionData | undefined = existing
+          ? { ...existing, ...partial, serverId }
+          : partial.serverId
+            ? ({ ...(partial as SessionData), serverId })
+            : undefined;
+
+        if (!next) {
+          return prev;
+        }
+
+        if (existing && shallowEqual(existing, next)) {
+          return prev;
+        }
+
+        logSessionStoreUpdate("updateSession", serverId, partial);
+        return {
+          ...prev,
+          sessions: {
+            ...prev.sessions,
+            [serverId]: next,
+          },
+        };
+      });
+    },
+    clearSession: (serverId) => {
+      set((prev) => {
+        if (!(serverId in prev.sessions)) {
+          return prev;
+        }
+        logSessionStoreUpdate("clearSession", serverId);
+        const nextSessions = { ...prev.sessions };
+        delete nextSessions[serverId];
+        return { ...prev, sessions: nextSessions };
+      });
+    },
+    getSession: (serverId) => {
+      return get().sessions[serverId];
+    },
+    setAgentDirectory: (serverId, agents) => {
+      set((prev) => {
+        const existing = prev.agentDirectory[serverId];
+        if (existing && areAgentDirectoriesEqual(existing, agents)) {
+          return prev;
+        }
+        logSessionStoreUpdate("setAgentDirectory", serverId, { agentCount: agents.length });
+        return {
+          ...prev,
+          agentDirectory: {
+            ...prev.agentDirectory,
+            [serverId]: agents,
+          },
+        };
+      });
+    },
+    clearAgentDirectory: (serverId) => {
+      set((prev) => {
+        if (!(serverId in prev.agentDirectory)) {
+          return prev;
+        }
+        logSessionStoreUpdate("clearAgentDirectory", serverId);
+        const nextDirectory = { ...prev.agentDirectory };
+        delete nextDirectory[serverId];
+        return { ...prev, agentDirectory: nextDirectory };
+      });
+    },
+    getAgentDirectory: (serverId) => {
+      return get().agentDirectory[serverId];
+    },
+  }))
+);