fix(app): close directory bootstrap races

packages/app/src/runtime/directory-sync/agent-replica.test.ts

@@ -57,6 +57,10 @@ describe("AgentDirectoryReplica", () => {
     store.initializeSession(serverId, null as unknown as DaemonClient);
     const replica = new AgentDirectoryReplica(serverId, () => undefined);
     replica.commitSnapshot([entry(payload("directory"))], []);
+    const directoryPlacement = useSessionStore
+      .getState()
+      .sessions[serverId]?.agents.get("agent")?.projectPlacement;
+    expect(directoryPlacement).toBeDefined();
     const staleToken = replica.captureTimeline("agent");
 
     replica.remove("agent");
@@ -74,6 +78,9 @@ describe("AgentDirectoryReplica", () => {
     expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.title).toBe(
       "current",
     );
+    expect(
+      useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.projectPlacement,
+    ).toEqual(directoryPlacement);
     store.clearSession(serverId);
   });
 });

packages/app/src/runtime/directory-sync/agent-replica.ts

@@ -40,10 +40,15 @@ export class AgentDirectoryReplica {
     ) {
       return false;
     }
-    const normalized = applyLegacyDaemonWorkspaceOwnership({
+    const existing = useSessionStore.getState().sessions[this.serverId]?.agents.get(token.agentId);
+    const timelineAgent = applyLegacyDaemonWorkspaceOwnership({
       serverId: this.serverId,
       agent: normalizeAgentSnapshot(payload, this.serverId),
     });
+    const normalized: Agent = {
+      ...timelineAgent,
+      projectPlacement: timelineAgent.projectPlacement ?? existing?.projectPlacement,
+    };
     const accepted = upsertAgentReplica(this.serverId, normalized);
     replaceAgentPendingPermissions(this.serverId, accepted);
     useSessionStore.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt);

packages/app/src/runtime/directory-sync/index.test.ts

@@ -0,0 +1,113 @@
+import { afterEach, describe, expect, it } from "vitest";
+import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
+import { useSessionStore } from "@/stores/session-store";
+import { DirectoryRefreshSupersededError, DirectorySync } from "./index";
+
+class FakeDirectoryClient {
+  fetchAgentsCalls = 0;
+  fetchWorkspacesCalls = 0;
+
+  on(): () => void {
+    return () => undefined;
+  }
+
+  async fetchAgents(): Promise<Awaited<ReturnType<DaemonClient["fetchAgents"]>>> {
+    this.fetchAgentsCalls += 1;
+    return {
+      requestId: "agents",
+      entries: [],
+      pageInfo: { hasMore: false, nextCursor: null, prevCursor: null },
+    };
+  }
+
+  async fetchWorkspaces(): Promise<Awaited<ReturnType<DaemonClient["fetchWorkspaces"]>>> {
+    this.fetchWorkspacesCalls += 1;
+    return {
+      requestId: "workspaces",
+      entries: [],
+      emptyProjects: [],
+      pageInfo: { hasMore: false, nextCursor: null, prevCursor: null },
+    };
+  }
+}
+
+const serverIds = new Set<string>();
+
+function createDirectory(serverId: string): {
+  client: FakeDirectoryClient;
+  directory: DirectorySync;
+} {
+  serverIds.add(serverId);
+  const client = new FakeDirectoryClient();
+  const directory = new DirectorySync(serverId, {
+    drainQueuedAgentMessage: () => undefined,
+    markAgentLoading: () => undefined,
+    markAgentReady: () => undefined,
+    markAgentError: () => undefined,
+  });
+  directory.connectionChanged({
+    client: client as unknown as DaemonClient,
+    status: "online",
+    source: { clientGeneration: 1, connectionEpoch: 1 },
+  });
+  return { client, directory };
+}
+
+afterEach(() => {
+  for (const serverId of serverIds) useSessionStore.getState().clearSession(serverId);
+  serverIds.clear();
+});
+
+describe("DirectorySync session readiness", () => {
+  it("waits for workspace capability metadata before choosing the workspace protocol", async () => {
+    const serverId = "workspace-metadata";
+    const { client, directory } = createDirectory(serverId);
+
+    const refresh = directory.refreshWorkspaces({ subscribe: true });
+    await Promise.resolve();
+    expect(client.fetchWorkspacesCalls).toBe(0);
+
+    const store = useSessionStore.getState();
+    store.initializeSession(serverId, client as unknown as DaemonClient, 1);
+    await Promise.resolve();
+    expect(client.fetchWorkspacesCalls).toBe(0);
+
+    store.updateSessionServerInfo(serverId, {
+      serverId,
+      hostname: null,
+      version: "test",
+      features: { workspaceMultiplicity: true },
+    });
+    await refresh;
+
+    expect(client.fetchWorkspacesCalls).toBe(1);
+    expect(useSessionStore.getState().sessions[serverId]?.hasHydratedWorkspaces).toBe(true);
+    directory.dispose();
+  });
+
+  it("rejects a session wait on disconnect so the reconnect can refresh", async () => {
+    const serverId = "session-wait-reconnect";
+    const { client, directory } = createDirectory(serverId);
+    const staleRefresh = directory.refreshAgents();
+    await Promise.resolve();
+
+    directory.connectionChanged({
+      client: null,
+      status: "offline",
+      source: { clientGeneration: 1, connectionEpoch: 1 },
+    });
+    await expect(staleRefresh).rejects.toBeInstanceOf(DirectoryRefreshSupersededError);
+
+    directory.connectionChanged({
+      client: client as unknown as DaemonClient,
+      status: "online",
+      source: { clientGeneration: 1, connectionEpoch: 2 },
+    });
+    const currentRefresh = directory.refreshAgents();
+    useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient, 1);
+    await currentRefresh;
+
+    expect(client.fetchAgentsCalls).toBe(1);
+    directory.dispose();
+  });
+});

packages/app/src/runtime/directory-sync/index.ts

@@ -74,6 +74,7 @@ export class DirectorySync {
     source: { clientGeneration: 0, connectionEpoch: 0 },
   };
   private unsubscribe: (() => void) | null = null;
+  private readonly abortSessionWaits = new Set<() => void>();
 
   constructor(
     private readonly serverId: string,
@@ -88,7 +89,7 @@ export class DirectorySync {
     this.workspaces = new WorkspaceDirectoryReplica(serverId);
   }
 
-  connectionChanged(connection: DirectoryConnection): void {
+  connectionChanged(connection: DirectoryConnection): boolean {
     const changed =
       this.connection.client !== connection.client ||
       this.connection.source.clientGeneration !== connection.source.clientGeneration ||
@@ -96,13 +97,14 @@ export class DirectorySync {
     const wentOffline = this.connection.status === "online" && connection.status === "offline";
     if (!changed && !wentOffline) {
       this.connection = connection;
-      return;
+      return false;
     }
     this.flushAbortedTransactions();
     this.unsubscribe?.();
     this.unsubscribe = null;
     this.connection = connection;
-    if (!connection.client || connection.status !== "online") return;
+    this.abortPendingSessionWaits();
+    if (!connection.client || connection.status !== "online") return true;
     const client = connection.client;
     const source = connection.source;
     const subscriptions = [
@@ -131,10 +133,12 @@ export class DirectorySync {
     this.unsubscribe = () => {
       for (const unsubscribe of subscriptions) unsubscribe();
     };
+    return true;
   }
 
   dispose(): void {
     this.flushAbortedTransactions();
+    this.abortPendingSessionWaits();
     this.unsubscribe?.();
     this.unsubscribe = null;
   }
@@ -216,14 +220,19 @@ export class DirectorySync {
   }
 
   async refreshWorkspaces(input?: { subscribe?: boolean }): Promise<void> {
-    const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo;
-    if (serverInfo?.features?.workspaceMultiplicity !== true) return;
     const { client, source } = this.requireOnline();
     const transaction = this.workspaceTransactions.begin(source, () => ({
       workspaces: new Map(),
       emptyProjects: new Map(),
     }));
     try {
+      await this.waitForSessionMetadata(client, source);
+      const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo;
+      if (serverInfo?.features?.workspaceMultiplicity !== true) {
+        const deltas = this.workspaceTransactions.fail(transaction);
+        if (deltas) for (const delta of deltas) this.workspaces.applyDelta(delta);
+        return;
+      }
       await this.fetchWorkspaceSnapshot(client, source, transaction, input?.subscribe === true);
       if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) {
         throw new DirectoryRefreshSupersededError("workspace completion no longer current");
@@ -343,18 +352,47 @@ export class DirectorySync {
   }
 
   private async waitForSession(client: DaemonClient, source: DirectorySourceToken): Promise<void> {
-    const matches = () => this.hasMatchingSession(client, source);
+    await this.waitForSessionState(client, source, () => this.hasMatchingSession(client, source));
+  }
+
+  private async waitForSessionMetadata(
+    client: DaemonClient,
+    source: DirectorySourceToken,
+  ): Promise<void> {
+    await this.waitForSessionState(client, source, () => {
+      const session = useSessionStore.getState().sessions[this.serverId];
+      return this.hasMatchingSession(client, source) && session?.serverInfo !== null;
+    });
+  }
+
+  private async waitForSessionState(
+    client: DaemonClient,
+    source: DirectorySourceToken,
+    matches: () => boolean,
+  ): Promise<void> {
     if (matches()) return;
     await new Promise<void>((resolve, reject) => {
-      const unsubscribe = useSessionStore.subscribe(() => {
+      let settled = false;
+      let unsubscribe: () => void = () => undefined;
+      const finish = (result: "ready" | "aborted") => {
+        if (settled) return;
+        settled = true;
+        unsubscribe();
+        this.abortSessionWaits.delete(abort);
+        if (result === "ready") resolve();
+        else reject(new DirectoryRefreshSupersededError("session wait no longer current"));
+      };
+      const abort = () => finish("aborted");
+      const check = () => {
         if (matches()) {
-          unsubscribe();
-          resolve();
+          finish("ready");
         } else if (!this.isCurrent(client, source)) {
-          unsubscribe();
-          reject(new DirectoryRefreshSupersededError("session wait no longer current"));
+          finish("aborted");
         }
-      });
+      };
+      this.abortSessionWaits.add(abort);
+      unsubscribe = useSessionStore.subscribe(check);
+      check();
     });
   }
 
@@ -367,6 +405,10 @@ export class DirectorySync {
     for (const delta of this.agentTransactions.abort()) this.agents.applyDelta(delta);
     for (const delta of this.workspaceTransactions.abort()) this.workspaces.applyDelta(delta);
   }
+
+  private abortPendingSessionWaits(): void {
+    for (const abort of this.abortSessionWaits) abort();
+  }
 }
 
 export class DirectoryRefreshSupersededError extends Error {}

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

@@ -1875,7 +1875,10 @@ describe("HostRuntimeStore", () => {
   });
 
   it("buffers updates until the matching session generation exists", async () => {
-    const host = makeHost({ serverId: "srv_pre_session" });
+    const host = makeHost({
+      serverId: "srv_pre_session",
+      connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }],
+    });
     const fakeClient = new FakeDaemonClient();
     fakeClient.setConnectionState({ status: "connected" });
     const snapshotEntry = makeFetchAgentsEntry({
@@ -1904,9 +1907,15 @@ 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, 1);
+      .initializeSession(host.serverId, fakeClient as unknown as DaemonClient, generation);
+    useSessionStore.getState().updateSessionServerInfo(host.serverId, {
+      serverId: host.serverId,
+      hostname: null,
+      version: "test",
+    });
     await fakeClient.waitForFetches(1);
     await waitForDirectoryReady(store, host.serverId);
 
@@ -1918,6 +1927,50 @@ 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

@@ -1962,14 +1962,15 @@ export class HostRuntimeStore {
       return;
     }
     const snapshot = controller.getSnapshot();
-    this.directorySyncByServer.get(serverId)?.connectionChanged({
-      client: snapshot.client,
-      status: snapshot.connectionStatus === "online" ? "online" : "offline",
-      source: {
-        clientGeneration: snapshot.clientGeneration,
-        connectionEpoch: snapshot.connectionEpoch,
-      },
-    });
+    const directorySourceChanged =
+      this.directorySyncByServer.get(serverId)?.connectionChanged({
+        client: snapshot.client,
+        status: snapshot.connectionStatus === "online" ? "online" : "offline",
+        source: {
+          clientGeneration: snapshot.clientGeneration,
+          connectionEpoch: snapshot.connectionEpoch,
+        },
+      }) ?? false;
     const previousStatus = this.lastConnectionStatusByServer.get(serverId);
     this.lastConnectionStatusByServer.set(serverId, snapshot.connectionStatus);
     const didTransitionOnline =
@@ -1992,7 +1993,7 @@ export class HostRuntimeStore {
     if (!didTransitionOnline && snapshot.hasEverLoadedAgentDirectory) {
       return;
     }
-    if (this.agentDirectoryBootstrapInFlight.has(serverId)) {
+    if (this.agentDirectoryBootstrapInFlight.has(serverId) && !directorySourceChanged) {
       return;
     }