fix(sync): preserve host replicas across provider remounts

Make the host runtime own session and setup replicas for the registered host lifetime. Provider remounts can then reattach without clearing directory snapshots or timeline cursors.

packages/app/src/contexts/session-context.tsx

@@ -387,8 +387,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider
   const toast = useToast();
 
   // Zustand store actions
-  const initializeSession = useSessionStore((state) => state.initializeSession);
-  const clearSession = useSessionStore((state) => state.clearSession);
   const setIsPlayingAudio = useSessionStore((state) => state.setIsPlayingAudio);
   const setMessages = useSessionStore((state) => state.setMessages);
   const setCurrentAssistantMessage = useSessionStore((state) => state.setCurrentAssistantMessage);
@@ -410,11 +408,9 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider
   const setWorkspaces = useSessionStore((state) => state.setWorkspaces);
   const flushAgentLastActivity = useSessionStore((state) => state.flushAgentLastActivity);
   const setPendingPermissions = useSessionStore((state) => state.setPendingPermissions);
-  const updateSessionClient = useSessionStore((state) => state.updateSessionClient);
   const updateSessionServerInfo = useSessionStore((state) => state.updateSessionServerInfo);
   const setViewedTimelineSync = useSessionStore((state) => state.setViewedTimelineSync);
   const upsertWorkspaceSetupProgress = useWorkspaceSetupStore((state) => state.upsertProgress);
-  const clearWorkspaceSetupServer = useWorkspaceSetupStore((state) => state.clearServer);
 
   // Track focused agent for heartbeat
   const focusedAgentId = useSessionStore(
@@ -524,17 +520,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider
     [serverId],
   );
 
-  // Initialize session in store
-  useEffect(() => {
-    const generation = getHostRuntimeStore().getSnapshot(serverId)?.clientGeneration ?? 0;
-    initializeSession(serverId, client, generation);
-  }, [serverId, client, initializeSession]);
-
-  useEffect(() => {
-    const generation = getHostRuntimeStore().getSnapshot(serverId)?.clientGeneration ?? 0;
-    updateSessionClient(serverId, client, generation);
-  }, [serverId, client, updateSessionClient]);
-
   useEffect(() => {
     const serverInfo = client.getLastServerInfoMessage();
     if (!serverInfo) {
@@ -1327,13 +1312,5 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider
     [client],
   );
 
-  // Cleanup on unmount
-  useEffect(() => {
-    return () => {
-      clearWorkspaceSetupServer(serverId);
-      clearSession(serverId);
-    };
-  }, [clearSession, clearWorkspaceSetupServer, serverId]);
-
   return children;
 }

packages/app/src/runtime/host-runtime.test.ts

@@ -1530,7 +1530,7 @@ describe("HostRuntimeStore", () => {
     useSessionStore.getState().clearSession(host.serverId);
   });
 
-  it("waits for the matching session replica before committing the connected client bootstrap", async () => {
+  it("owns the session replica for the registered host lifecycle", async () => {
     const host = makeHost({
       serverId: "srv_no_session",
       connections: [
@@ -1557,15 +1557,12 @@ describe("HostRuntimeStore", () => {
 
     store.syncHosts([host]);
 
-    await Promise.resolve();
-    expect(fakeClient.fetchAgentsCalls).toEqual([]);
-    useSessionStore
-      .getState()
-      .initializeSession(host.serverId, fakeClient as unknown as DaemonClient, 1);
-
     await fakeClient.waitForFetches(1);
     await waitForDirectoryReady(store, host.serverId);
 
+    const session = useSessionStore.getState().sessions[host.serverId];
+    expect(session?.client).toBe(fakeClient);
+    expect(session?.clientGeneration).toBe(1);
     expect(fakeClient.fetchAgentsCalls).toHaveLength(1);
     expect(fakeClient.fetchAgentsCalls[0]).toEqual({
       scope: "active",
@@ -1575,7 +1572,7 @@ describe("HostRuntimeStore", () => {
     });
 
     store.syncHosts([]);
-    useSessionStore.getState().clearSession(host.serverId);
+    expect(useSessionStore.getState().sessions[host.serverId]).toBeUndefined();
   });
 
   it("bootstraps legacy daemons from unscoped agents and creates path-backed workspaces", async () => {
@@ -1943,7 +1940,7 @@ describe("HostRuntimeStore", () => {
     useSessionStore.getState().clearSession(host.serverId);
   });
 
-  it("buffers updates until the matching session generation exists", async () => {
+  it("applies agent updates received during initial directory bootstrap", async () => {
     const host = makeHost({
       serverId: "srv_pre_session",
       connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }],
@@ -1976,10 +1973,6 @@ describe("HostRuntimeStore", () => {
       agent: { ...snapshotEntry.agent, title: "before-session" },
       project: snapshotEntry.project,
     });
-    const generation = store.getSnapshot(host.serverId)?.clientGeneration ?? 0;
-    useSessionStore
-      .getState()
-      .initializeSession(host.serverId, fakeClient as unknown as DaemonClient, generation);
     useSessionStore.getState().updateSessionServerInfo(host.serverId, {
       serverId: host.serverId,
       hostname: null,
@@ -1996,50 +1989,6 @@ describe("HostRuntimeStore", () => {
     useSessionStore.getState().clearSession(host.serverId);
   });
 
-  it("restarts directory bootstrap when reconnect supersedes a pending session wait", async () => {
-    const host = makeHost({
-      serverId: "srv_session_wait_reconnect",
-      connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }],
-    });
-    const fakeClient = new FakeDaemonClient();
-    fakeClient.setConnectionState({ status: "connected" });
-    const store = new HostRuntimeStore({
-      deps: {
-        createClient: () => fakeClient as unknown as DaemonClient,
-        connectToDaemon: async () => ({
-          client: fakeClient as unknown as DaemonClient,
-          serverId: host.serverId,
-          hostname: null,
-        }),
-        getClientId: async () => "cid_session_wait_reconnect",
-      },
-    });
-
-    store.syncHosts([host]);
-    await fakeClient.waitForAgentUpdates();
-    fakeClient.setConnectionState({ status: "disconnected", reason: "network" });
-    fakeClient.setConnectionState({ status: "connected" });
-
-    const generation = store.getSnapshot(host.serverId)?.clientGeneration ?? 0;
-    const sessionStore = useSessionStore.getState();
-    sessionStore.initializeSession(
-      host.serverId,
-      fakeClient as unknown as DaemonClient,
-      generation,
-    );
-    sessionStore.updateSessionServerInfo(host.serverId, {
-      serverId: host.serverId,
-      hostname: null,
-      version: "test",
-    });
-    await fakeClient.waitForFetches(1);
-    await waitForDirectoryReady(store, host.serverId);
-
-    expect(fakeClient.fetchAgentsCalls).toHaveLength(1);
-    store.syncHosts([]);
-    sessionStore.clearSession(host.serverId);
-  });
-
   it("rejects a superseded refresh without overwriting the newer replica", async () => {
     const host = makeHost({
       serverId: "srv_overlap",

packages/app/src/runtime/host-runtime.ts

@@ -40,6 +40,7 @@ import { getDesktopHost } from "@/desktop/host";
 import { CLIENT_CAPS } from "@getpaseo/protocol/client-capabilities";
 import { BROWSER_AUTOMATION_COMMAND_NAMES } from "@getpaseo/protocol/browser-automation/rpc-schemas";
 import { useSessionStore } from "@/stores/session-store";
+import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store";
 import { invalidateCheckoutGitQueriesForServer } from "@/git/query-keys";
 import { queryClient } from "@/data/query-client";
 import {
@@ -1583,8 +1584,6 @@ export class HostRuntimeStore {
     }
 
     rekeyMap(this.controllers, oldServerId, newServerId);
-    controller.adoptReconciledServerId(newServerId);
-
     rekeyMap(this.lastConnectionStatusByServer, oldServerId, newServerId);
     rekeyMap(this.directoryBootstrapInFlight, oldServerId, newServerId);
     this.replicaCache.reconcileServerId(oldServerId, newServerId);
@@ -1597,7 +1596,10 @@ export class HostRuntimeStore {
       markAgentError: (error) => controller.markAgentDirectorySyncError(error),
     });
     this.directorySyncByServer.set(newServerId, directory);
+    controller.adoptReconciledServerId(newServerId);
     const snapshot = controller.getSnapshot();
+    this.clearHostReplica(oldServerId);
+    this.syncSessionReplica(newServerId, snapshot);
     directory.connectionChanged({
       client: snapshot.client,
       status: snapshot.connectionStatus === "online" ? "online" : "offline",
@@ -1899,6 +1901,7 @@ export class HostRuntimeStore {
       this.directoryBootstrapInFlight.delete(serverId);
       this.directorySyncByServer.get(serverId)?.dispose();
       this.directorySyncByServer.delete(serverId);
+      this.clearHostReplica(serverId);
       void controller.stop();
       this.emit(serverId);
     }
@@ -1936,8 +1939,10 @@ export class HostRuntimeStore {
         controller.getSnapshot().connectionStatus,
       );
       controller.subscribe(() => {
-        this.maybeAutoBootstrapDirectories(host.serverId);
-        this.emit(host.serverId);
+        const snapshot = controller.getSnapshot();
+        this.syncSessionReplica(snapshot.serverId, snapshot);
+        this.maybeAutoBootstrapDirectories(snapshot.serverId);
+        this.emit(snapshot.serverId);
       });
       void controller
         .start(
@@ -1955,6 +1960,20 @@ export class HostRuntimeStore {
     }
   }
 
+  private syncSessionReplica(serverId: string, snapshot: HostRuntimeSnapshot): void {
+    if (!snapshot.client) {
+      return;
+    }
+    const sessionStore = useSessionStore.getState();
+    sessionStore.initializeSession(serverId, snapshot.client, snapshot.clientGeneration);
+    sessionStore.updateSessionClient(serverId, snapshot.client, snapshot.clientGeneration);
+  }
+
+  private clearHostReplica(serverId: string): void {
+    useSessionStore.getState().clearSession(serverId);
+    useWorkspaceSetupStore.getState().clearServer(serverId);
+  }
+
   private maybeAutoBootstrapDirectories(serverId: string): void {
     const controller = this.controllers.get(serverId);
     if (!controller) {