From d0456b1943358fd73775c7b14bfc705e837e4af3 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 21 Jul 2026 13:33:36 +0200 Subject: [PATCH] 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. --- docs/timeline-sync.md | 10 +++ packages/app/src/contexts/session-context.tsx | 23 ------- packages/app/src/runtime/host-runtime.test.ts | 63 ++----------------- packages/app/src/runtime/host-runtime.ts | 27 ++++++-- 4 files changed, 39 insertions(+), 84 deletions(-) diff --git a/docs/timeline-sync.md b/docs/timeline-sync.md index e96250f0d..79e85d0b8 100644 --- a/docs/timeline-sync.md +++ b/docs/timeline-sync.md @@ -51,6 +51,16 @@ When a client resumes with a known cursor, it catches up after that cursor to co When a client resumes without a cursor, it fetches the latest tail page. +## Client replica lifetime + +The host runtime owns each session replica for as long as the host remains registered. React +providers attach message handlers and UI integrations to that replica, but mounting or unmounting a +provider must not create or clear it. A provider can remount during Fast Refresh or ordinary UI +recomposition while the runtime still owns the same directory snapshot and timeline cursors. + +Removing the host from the registry is the destructive boundary: it stops the runtime and clears the +session and host-scoped setup state together. + ## Selective and legacy delivery The app chooses one delivery policy from `server_info.features.selectiveAgentTimeline`: diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index ed2f62fb0..606ea94c0 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/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; } diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index 156f554dc..0f4946a4b 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/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", diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index af76ba4fa..144632817 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/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) {