fix: Phase 2 review issues

packages/app/src/app/settings.tsx

@@ -20,6 +20,7 @@ import { formatConnectionStatus, getConnectionStatusTone } from "@/utils/daemons
 import { theme as defaultTheme } from "@/styles/theme";
 import { BackHeader } from "@/components/headers/back-header";
 import { useSessionForServer } from "@/hooks/use-session-directory";
+import type { SessionContextValue } from "@/contexts/session-context";
 
 const delay = (ms: number) =>
   new Promise<void>((resolve) => {
@@ -797,11 +798,18 @@ function DaemonCard({
           : theme.colors.mutedForeground;
   const badgeText = statusLabel;
   const connectionError = typeof lastError === "string" && lastError.trim().length > 0 ? lastError.trim() : null;
-  const daemonSession = useSessionForServer(daemon.id);
+  const daemonWs = useSessionForServer<SessionContextValue["ws"] | null>(
+    daemon.id,
+    useCallback((session) => session?.ws ?? null, [])
+  );
+  const restartServerFn = useSessionForServer<SessionContextValue["restartServer"] | null>(
+    daemon.id,
+    useCallback((session) => session?.restartServer ?? null, [])
+  );
   const [isRestarting, setIsRestarting] = useState(false);
-  const wsIsConnectedRef = useRef(daemonSession?.ws.isConnected ?? false);
+  const wsIsConnectedRef = useRef(daemonWs?.isConnected ?? false);
   const isTesting = testState?.status === "testing";
-  const sessionIsConnected = daemonSession?.ws.isConnected ?? false;
+  const sessionIsConnected = daemonWs?.isConnected ?? false;
 
   useEffect(() => {
     wsIsConnectedRef.current = sessionIsConnected;
@@ -853,7 +861,7 @@ function DaemonCard({
   }, [daemon.label, daemon.wsUrl, isScreenMountedRef, testServerConnection, waitForCondition]);
 
   const beginServerRestart = useCallback(() => {
-    if (!daemonSession) {
+    if (!restartServerFn) {
       Alert.alert(
         "Host unavailable",
         `${daemon.label} is not connected. Wait for it to come online before restarting.`
@@ -871,7 +879,7 @@ function DaemonCard({
 
     setIsRestarting(true);
     try {
-      daemonSession.restartServer(`settings_daemon_restart_${daemon.id}`);
+      restartServerFn(`settings_daemon_restart_${daemon.id}`);
     } catch (error) {
       console.error(`[Settings] Failed to restart daemon ${daemon.label}`, error);
       setIsRestarting(false);
@@ -883,10 +891,10 @@ function DaemonCard({
     }
 
     void waitForDaemonRestart();
-  }, [daemon.id, daemon.label, daemonSession, waitForDaemonRestart]);
+  }, [daemon.id, daemon.label, restartServerFn, waitForDaemonRestart]);
 
   const handleRestartPress = useCallback(() => {
-    if (!daemonSession) {
+    if (!restartServerFn) {
       Alert.alert(
         "Host unavailable",
         `${daemon.label} is not connected. Wait for it to come online before restarting.`
@@ -917,7 +925,7 @@ function DaemonCard({
         onPress: beginServerRestart,
       },
     ]);
-  }, [beginServerRestart, daemon.label, daemonSession, restartConfirmationMessage]);
+  }, [beginServerRestart, daemon.label, restartConfirmationMessage, restartServerFn]);
 
   return (
     <View style={[styles.settingCard, styles.daemonCard]}>

packages/app/src/components/create-agent-modal.tsx

@@ -68,6 +68,7 @@ import type { SessionContextValue } from "@/contexts/session-context";
 import { useSessionForServer } from "@/hooks/use-session-directory";
 import type { UseWebSocketReturn } from "@/hooks/use-websocket";
 import { useSessionStore } from "@/stores/session-store";
+import { shallow } from "zustand/shallow";
 
 export type CreateAgentInitialValues = {
   workingDir?: string;
@@ -101,6 +102,18 @@ interface ModalWrapperProps {
   serverId?: string | null;
 }
 
+type CreateAgentSessionSlice = Pick<
+  SessionContextValue,
+  | "serverId"
+  | "ws"
+  | "createAgent"
+  | "resumeAgent"
+  | "sendAgentAudio"
+  | "agents"
+  | "providerModels"
+  | "requestProviderModels"
+>;
+
 const providerDefinitions = AGENT_PROVIDER_DEFINITIONS;
 const providerDefinitionMap = new Map<AgentProvider, AgentProviderDefinition>(
   providerDefinitions.map((definition) => [definition.id, definition])
@@ -238,8 +251,25 @@ function AgentFlowModal({
     return exists ? serverId : null;
   }, [serverId, daemonEntries]);
   const [selectedServerId, setSelectedServerId] = useState<string | null>(initialServerId);
-  const selectedSession = useSessionForServer(selectedServerId);
+  const selectSessionSlice = useCallback(
-  const session = selectedSession ?? null;
+    (session: SessionContextValue | null): CreateAgentSessionSlice | null => {
+      if (!session) {
+        return null;
+      }
+      return {
+        serverId: session.serverId,
+        ws: session.ws,
+        createAgent: session.createAgent,
+        resumeAgent: session.resumeAgent,
+        sendAgentAudio: session.sendAgentAudio,
+        agents: session.agents,
+        providerModels: session.providerModels,
+        requestProviderModels: session.requestProviderModels,
+      } satisfies CreateAgentSessionSlice;
+    },
+    []
+  );
+  const session = useSessionForServer<CreateAgentSessionSlice | null>(selectedServerId, selectSessionSlice, shallow);
   const getSession = useSessionStore((state) => state.getSession);
 
   useEffect(() => {
@@ -574,7 +604,7 @@ function AgentFlowModal({
   const queueProviderModelFetch = useCallback(
     (
       serverId: string | null,
-      targetSession: SessionContextValue | null,
+      targetSession: Pick<SessionContextValue, "providerModels" | "requestProviderModels"> | null,
       options?: { cwd?: string; delayMs?: number }
     ) => {
       if (!serverId || !targetSession?.requestProviderModels) {

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

@@ -1,4 +1,5 @@
-import { useMemo } from "react";
+import { useCallback, useMemo } from "react";
+import { useStoreWithEqualityFn } from "zustand/traditional";
 import { useSessionStore, type SessionData } from "@/stores/session-store";
 
 export function useSessionDirectory(): Map<string, SessionData> {
@@ -9,6 +10,33 @@ export function useSessionDirectory(): Map<string, SessionData> {
   }, [sessions]);
 }
 
-export function useSessionForServer(serverId: string | null): SessionData | null {
+type SessionSelector<T> = (session: SessionData | null) => T;
-  return useSessionStore((state) => (serverId ? state.sessions[serverId] ?? null : null));
+type EqualityFn<T> = ((left: T, right: T) => boolean) | undefined;
+
+export function useSessionForServer(serverId: string | null): SessionData | null;
+export function useSessionForServer<T>(
+  serverId: string | null,
+  selector: SessionSelector<T>,
+  equalityFn?: EqualityFn<T>
+): T;
+export function useSessionForServer<T>(
+  serverId: string | null,
+  selector?: SessionSelector<T>,
+  equalityFn?: EqualityFn<T>
+): SessionData | null | T {
+  const baseSelector = useCallback(
+    (state: ReturnType<typeof useSessionStore.getState>) =>
+      (serverId ? state.sessions[serverId] ?? null : null),
+    [serverId]
+  );
+
+  if (selector) {
+    const derivedSelector = useCallback(
+      (state: ReturnType<typeof useSessionStore.getState>) => selector(baseSelector(state)),
+      [baseSelector, selector]
+    );
+    return useStoreWithEqualityFn(useSessionStore, derivedSelector, equalityFn);
+  }
+
+  return useSessionStore(baseSelector);
 }

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

@@ -77,7 +77,10 @@ function areAgentDirectoriesEqual(left: AgentDirectoryEntry[] | undefined, right
       entry.id === other.id &&
       entry.status === other.status &&
       entry.serverId === other.serverId &&
-      entry.lastActivityAt.getTime() === other.lastActivityAt.getTime()
+      entry.lastActivityAt.getTime() === other.lastActivityAt.getTime() &&
+      entry.title === other.title &&
+      entry.cwd === other.cwd &&
+      entry.provider === other.provider
     );
   });
 }