Unify app connection state around host runtime

packages/app/src/screens/agent/agent-ready-screen.tsx

@@ -41,8 +41,11 @@ import {
 } from "@/contexts/explorer-sidebar-animation-context";
 import { usePanelStore } from "@/stores/panel-store";
 import { useDaemonConnections } from "@/contexts/daemon-connections-context";
-import type { ConnectionStatus } from "@/contexts/daemon-connections-context";
 import { useSessionStore } from "@/stores/session-store";
+import {
+  useHostRuntimeSession,
+  type HostRuntimeConnectionStatus,
+} from "@/runtime/host-runtime";
 import { useCreateFlowStore } from "@/stores/create-flow-store";
 import type { Agent } from "@/contexts/session-context";
 import type { StreamItem } from "@/types/stream";
@@ -75,7 +78,7 @@ import {
 } from "@/utils/agent-snapshots";
 import { mergePendingCreateImages } from "@/utils/pending-create-images";
 import { shouldClearAgentAttentionOnView } from "@/utils/agent-attention";
-import type { FetchAgentsEntry } from "@server/client/daemon-client";
+import type { DaemonClient, FetchAgentsEntry } from "@server/client/daemon-client";
 import { useExplorerOpenGesture } from "@/hooks/use-explorer-open-gesture";
 import {
   DropdownMenu,
@@ -109,11 +112,12 @@ export function AgentReadyScreen({
   const resolvedAgentId = agentId?.trim() || undefined;
   const resolvedServerId = serverId?.trim() || undefined;
   const { connectionStates } = useDaemonConnections();
-
-  // Check if session exists for this serverId
-  const hasSession = useSessionStore((state) =>
-    resolvedServerId ? !!state.sessions[resolvedServerId] : false
-  );
+  const runtimeServerId = resolvedServerId ?? "";
+  const {
+    snapshot: runtimeSnapshot,
+    client: runtimeClient,
+    isConnected: runtimeIsConnected,
+  } = useHostRuntimeSession(runtimeServerId);
 
   const connectionServerId = resolvedServerId ?? null;
   const connection = connectionServerId
@@ -122,9 +126,11 @@ export function AgentReadyScreen({
   const serverLabel =
     connection?.daemon.label ?? connectionServerId ?? "Selected host";
   const isUnknownDaemon = Boolean(connectionServerId && !connection);
-  const connectionStatus =
-    connection?.status ?? (isUnknownDaemon ? "offline" : "idle");
-  const lastConnectionError = connection?.lastError ?? null;
+  const connectionStatus: HostRuntimeConnectionStatus =
+    runtimeSnapshot?.connectionStatus ??
+    (isUnknownDaemon ? "offline" : "idle");
+  const lastConnectionError = runtimeSnapshot?.lastError ?? null;
+  const isRuntimeSessionAvailable = Boolean(resolvedServerId && runtimeClient);
 
   const handleBackToHome = useCallback(() => {
     const targetServerId = resolvedServerId;
@@ -154,7 +160,9 @@ export function AgentReadyScreen({
   }, [resolvedAgentId, resolvedServerId, router]);
 
   const focusServerId = resolvedServerId;
-  const navigationStatus = hasSession ? "ready" : "session_unavailable";
+  const navigationStatus = isRuntimeSessionAvailable
+    ? "ready"
+    : "session_unavailable";
 
   useFocusEffect(
     useCallback(() => {
@@ -172,7 +180,7 @@ export function AgentReadyScreen({
     }, [focusServerId, navigationStatus, resolvedAgentId])
   );
 
-  if (!hasSession || !resolvedServerId) {
+  if (!resolvedServerId || !runtimeClient) {
     return (
       <AgentSessionUnavailableState
         onBack={handleBackToHome}
@@ -189,6 +197,8 @@ export function AgentReadyScreen({
       <AgentScreenContent
         serverId={resolvedServerId}
         agentId={resolvedAgentId}
+        client={runtimeClient}
+        isConnected={runtimeIsConnected}
         connectionStatus={connectionStatus}
       />
     </ExplorerSidebarAnimationProvider>
@@ -198,7 +208,9 @@ export function AgentReadyScreen({
 type AgentScreenContentProps = {
   serverId: string;
   agentId?: string;
-  connectionStatus: ConnectionStatus;
+  client: DaemonClient;
+  isConnected: boolean;
+  connectionStatus: HostRuntimeConnectionStatus;
 };
 
 function toErrorMessage(error: unknown): string {
@@ -215,6 +227,8 @@ function isNotFoundErrorMessage(message: string): boolean {
 function AgentScreenContent({
   serverId,
   agentId,
+  client,
+  isConnected,
   connectionStatus,
 }: AgentScreenContentProps) {
   const { theme } = useUnistyles();
@@ -465,16 +479,16 @@ function AgentScreenContent({
     return filtered;
   }, [allPendingPermissions, resolvedAgentId]);
 
-  const client = useSessionStore(
-    (state) => state.sessions[serverId]?.client ?? null
-  );
-  const isConnected = useSessionStore(
-    (state) => state.sessions[serverId]?.connection.isConnected ?? false
+  const hasSession = useSessionStore(
+    (state) => Boolean(state.sessions[serverId])
   );
   const focusedAgentId = useSessionStore(
     (state) => state.sessions[serverId]?.focusedAgentId ?? null
   );
-  const { ensureAgentIsInitialized, refreshAgent } = useAgentInitialization(serverId);
+  const { ensureAgentIsInitialized, refreshAgent } = useAgentInitialization({
+    serverId,
+    client: hasSession ? client : null,
+  });
   const [missingAgentState, setMissingAgentState] = useState<AgentScreenMissingState>({
     kind: "idle",
   });
@@ -529,7 +543,7 @@ function AgentScreenContent({
 
   useFocusEffect(
     useCallback(() => {
-      if (!resolvedAgentId || !isConnected) {
+      if (!resolvedAgentId || !isConnected || !hasSession) {
         return;
       }
       ensureAgentIsInitialized(resolvedAgentId).catch((error) => {
@@ -538,7 +552,7 @@ function AgentScreenContent({
           error,
         });
       });
-    }, [ensureAgentIsInitialized, isConnected, resolvedAgentId])
+    }, [ensureAgentIsInitialized, hasSession, isConnected, resolvedAgentId])
   );
 
   const isGitCheckout = activeExplorerCheckout?.isGit ?? false;
@@ -745,7 +759,7 @@ function AgentScreenContent({
       return;
     }
 
-    if (!isConnected) {
+    if (!isConnected || !hasSession) {
       return;
     }
     // On native clients, daemon stream forwarding is focused-agent only, so switching
@@ -764,6 +778,7 @@ function AgentScreenContent({
   }, [
     resolvedAgentId,
     ensureAgentIsInitialized,
+    hasSession,
     isConnected,
     needsAuthoritativeSync,
   ]);
@@ -784,7 +799,7 @@ function AgentScreenContent({
       }
       return;
     }
-    if (!isConnected) {
+    if (!isConnected || !hasSession) {
       return;
     }
     if (missingAgentState.kind === "resolving" || missingAgentState.kind === "not_found") {
@@ -855,6 +870,7 @@ function AgentScreenContent({
     agent,
     client,
     ensureAgentIsInitialized,
+    hasSession,
     isConnected,
     missingAgentState.kind,
     resolvedAgentId,
@@ -898,13 +914,13 @@ function AgentScreenContent({
   ]);
 
   const handleRefreshAgent = useCallback(() => {
-    if (!resolvedAgentId) {
+    if (!resolvedAgentId || !hasSession) {
       return;
     }
     void refreshAgent(resolvedAgentId).catch((error) => {
       console.warn("[AgentScreen] refreshAgent failed", { agentId: resolvedAgentId, error });
     });
-  }, [resolvedAgentId, refreshAgent]);
+  }, [hasSession, refreshAgent, resolvedAgentId]);
 
   const handleCopyMeta = useCallback(
     async (label: string, value: string | null | undefined) => {
@@ -1262,7 +1278,7 @@ function AgentSessionUnavailableState({
 }: {
   onBack: () => void;
   serverLabel: string;
-  connectionStatus: ConnectionStatus;
+  connectionStatus: HostRuntimeConnectionStatus;
   lastError: string | null;
   isUnknownDaemon?: boolean;
 }) {
@@ -1285,19 +1301,24 @@ function AgentSessionUnavailableState({
   }
 
   const isConnecting = connectionStatus === "connecting";
+  const isPreparingSession = connectionStatus === "online";
 
   return (
     <View style={styles.container}>
       <BackHeader title="Agent" onBack={onBack} />
       <View style={styles.centerState}>
-        {isConnecting ? (
+        {isConnecting || isPreparingSession ? (
           <>
             <ActivityIndicator size="large" />
             <Text style={styles.loadingText}>
-              Connecting to {serverLabel}...
+              {isPreparingSession
+                ? `Preparing ${serverLabel} session...`
+                : `Connecting to ${serverLabel}...`}
             </Text>
             <Text style={styles.statusText}>
-              We will show this agent once the host is online.
+              {isPreparingSession
+                ? "We will show this agent in a moment."
+                : "We will show this agent once the host is online."}
             </Text>
           </>
         ) : (

packages/app/src/screens/agent/draft-agent-screen.tsx

@@ -35,6 +35,7 @@ import { buildWorkingDirectorySuggestions } from '@/utils/working-directory-sugg
 import { useExplorerOpenGesture } from '@/hooks/use-explorer-open-gesture'
 import { useSessionStore } from '@/stores/session-store'
 import { useCreateFlowStore } from '@/stores/create-flow-store'
+import { useHostRuntimeSession } from '@/runtime/host-runtime'
 import { ExplorerSidebarAnimationProvider } from '@/contexts/explorer-sidebar-animation-context'
 import { usePanelStore, type ExplorerCheckoutContext } from '@/stores/panel-store'
 import { MAX_CONTENT_WIDTH } from '@/constants/layout'
@@ -368,18 +369,16 @@ function DraftAgentScreenContent({
     return collectAgentWorkingDirectorySuggestions([...liveSources, ...fetchedSources])
   }, [allAgents, sessionAgents])
 
-  const sessionClient = useSessionStore((state) =>
-    selectedServerId ? (state.sessions[selectedServerId]?.client ?? null) : null
-  )
-  const isConnected = useSessionStore((state) =>
-    selectedServerId ? (state.sessions[selectedServerId]?.connection.isConnected ?? false) : false
+  const { client: runtimeClient, isConnected: isHostOnline } = useHostRuntimeSession(
+    selectedServerId ?? ''
   )
+  const sessionClient = runtimeClient
   const trimmedWorkingDir = workingDir.trim()
   const shouldInspectRepo = trimmedWorkingDir.length > 0
   const daemonAvailabilityError =
-    !selectedServerId || hostEntry?.status !== 'online' ? 'Host is offline' : null
+    !selectedServerId || !isHostOnline ? 'Host is offline' : null
   const repoAvailabilityError =
-    shouldInspectRepo && (!hostEntry || hostEntry.status !== 'online' || !isConnected)
+    shouldInspectRepo && !isHostOnline
       ? (daemonAvailabilityError ??
         'Repository details will load automatically once the selected host is back online.')
       : null
@@ -398,7 +397,7 @@ function DraftAgentScreenContent({
       Boolean(trimmedWorkingDir) &&
       !repoAvailabilityError &&
       Boolean(sessionClient) &&
-      isConnected,
+      isHostOnline,
     retry: false,
     staleTime: CHECKOUT_STATUS_STALE_TIME,
     refetchOnMount: 'always',
@@ -461,7 +460,7 @@ function DraftAgentScreenContent({
       Boolean(worktreeListRoot) &&
       !repoAvailabilityError &&
       Boolean(sessionClient) &&
-      isConnected &&
+      isHostOnline &&
       !isNonGitDirectory,
     retry: false,
     staleTime: 0,
@@ -526,7 +525,7 @@ function DraftAgentScreenContent({
       Boolean(trimmedWorkingDir) &&
       !repoAvailabilityError &&
       Boolean(sessionClient) &&
-      isConnected,
+      isHostOnline,
     retry: false,
     staleTime: 15_000,
   })
@@ -559,7 +558,7 @@ function DraftAgentScreenContent({
       Boolean(selectedServerId) &&
       !daemonAvailabilityError &&
       Boolean(sessionClient) &&
-      isConnected,
+      isHostOnline,
     retry: false,
     staleTime: 15_000,
   })
@@ -625,7 +624,7 @@ function DraftAgentScreenContent({
       Boolean(baseBranch) &&
       Boolean(trimmedWorkingDir) &&
       Boolean(sessionClient) &&
-      isConnected,
+      isHostOnline,
     retry: false,
     staleTime: 30_000,
   })
@@ -810,9 +809,7 @@ function DraftAgentScreenContent({
     })
   }, [baseBranch, branchSearchQuery, branchSuggestionsQuery.data, checkout, worktreeOptions])
 
-  const createAgentClient = useSessionStore((state) =>
-    selectedServerId ? (state.sessions[selectedServerId]?.client ?? null) : null
-  )
+  const createAgentClient = sessionClient
   const draftCommandConfig = useMemo<DraftCommandConfig | undefined>(() => {
     const cwd = (
       isAttachWorktree && selectedWorktreePath ? selectedWorktreePath : workingDir

packages/app/src/screens/settings-screen.tsx

@@ -22,6 +22,7 @@ import { measureConnectionLatency } from "@/utils/test-daemon-connection";
 import { confirmDialog } from "@/utils/confirm-dialog";
 import { MenuHeader } from "@/components/headers/menu-header";
 import { useSessionStore } from "@/stores/session-store";
+import { useHostRuntimeSession } from "@/runtime/host-runtime";
 import { AddHostMethodModal } from "@/components/add-host-method-modal";
 import { AddHostModal } from "@/components/add-host-modal";
 import { PairLinkModal } from "@/components/pair-link-modal";
@@ -1035,10 +1036,11 @@ function HostDetailModal({
   );
 
   // Restart logic (moved from DaemonCard)
-  const daemonClient = useSessionStore((state) => host ? (state.sessions[host.serverId]?.client ?? null) : null);
-  const daemonConnection = useSessionStore((state) => host ? (state.sessions[host.serverId]?.connection ?? null) : null);
+  const { client: runtimeClient, isConnected } = useHostRuntimeSession(
+    host?.serverId ?? ""
+  );
+  const daemonClient = runtimeClient;
   const daemonVersion = useSessionStore((state) => host ? (state.sessions[host.serverId]?.serverInfo?.version ?? null) : null);
-  const isConnected = daemonConnection?.isConnected ?? false;
   const isConnectedRef = useRef(isConnected);
   const [isRestarting, setIsRestarting] = useState(false);