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.
This commit is contained in:
Mohamed Boudra
2026-07-21 13:33:36 +02:00
parent 8aa55db1e8
commit d0456b1943
4 changed files with 39 additions and 84 deletions

View File

@@ -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`:

View File

@@ -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;
}

View File

@@ -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",

View File

@@ -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) {