fix(app): restore directory subscriptions after demand returns (#8403)

Invalidate refresh satisfaction when subscriptions are released, and retry
if demand returns before the overtaken refresh completes. Otherwise the
same connection source can stay satisfied with no live list listeners.

Cover both demand sources and completion orderings with live delta checks.

Co-authored-by: Okan YILDIRIM <63511519+ragokan@users.noreply.github.com>

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

@@ -16,7 +16,8 @@ import {
   normalizeWorkspaceDescriptor,
   useSessionStore,
 } from "@/stores/session-store";
-import { normalizeAgentSnapshot } from "@/utils/agent-snapshots";
+import { normalizeAgentSnapshot, projectAgentSnapshot } from "@/utils/agent-snapshots";
+import { deriveProjectPlacementFromCwd } from "@/utils/project-placement";
 import { selectWorkspaceDirectoryServerIds } from "@/stores/session-store-hooks/selectors";
 import { useWorkspaceLabels } from "@/workspace-labels";
 import type { DirectoryReplicaMutation } from "@/runtime/replica-cache";
@@ -41,12 +42,15 @@ class FakeDirectoryClient {
   supportsWorkspaceLabels = false;
   listWorkspaceLabelsCalls = 0;
   fetchAgentsCalls = 0;
+  observeAgentsCalls = 0;
   lastAgentOptions: unknown;
   fetchWorkspacesCalls = 0;
+  observeWorkspacesCalls = 0;
   lastWorkspaceOptions: unknown;
   listProjectsCalls = 0;
   lastProjectOptions: unknown;
   projectResult: ProjectListResult | null = null;
+  readonly listSnapshots: Promise<unknown>[] = [];
   private pendingAgentFetch: Promise<AgentFetchResult> | null = null;
   private pendingWorkspaceFetch: Promise<WorkspaceFetchResult> | null = null;
   private pendingProjectFetch: Promise<ProjectListResult> | null = null;
@@ -151,15 +155,22 @@ class FakeDirectoryClient {
   }
 
   observeAgents(options: Parameters<DaemonClient["observeAgents"]>[0]) {
-    return subscriptionFixture(this.fetchAgents({ ...options, subscribe: {} }), (receive) =>
-      this.on("agent_update", receive),
+    this.observeAgentsCalls += 1;
+    const subscription = subscriptionFixture(
+      this.fetchAgents({ ...options, subscribe: {} }),
+      (receive) => this.on("agent_update", receive),
     );
+    this.listSnapshots.push(subscription.ready);
+    return subscription;
   }
 
   observeWorkspaces(options: Parameters<DaemonClient["observeWorkspaces"]>[0]) {
-    return subscriptionFixture(this.fetchWorkspaces(options), (receive) =>
+    this.observeWorkspacesCalls += 1;
+    const subscription = subscriptionFixture(this.fetchWorkspaces(options), (receive) =>
       this.on("workspace_update", receive),
     );
+    this.listSnapshots.push(subscription.ready);
+    return subscription;
   }
 
   observeEvents(events: readonly SessionOutboundMessage["type"][]) {
@@ -592,6 +603,83 @@ describe("DirectorySync session readiness", () => {
     directory.dispose();
   });
 
+  it.each([
+    ["full", "before"],
+    ["full", "after"],
+    ["route", "before"],
+    ["route", "after"],
+  ] as const)(
+    "restores live directory updates when %s demand returns %s refresh completion",
+    async (demand, timing) => {
+      const serverId = `released-directory-${demand}-${timing}`;
+      const { client, directory } = createDirectory(serverId);
+      const store = useSessionStore.getState();
+      store.initializeSession(serverId, client as unknown as DaemonClient, 1);
+      store.updateSessionServerInfo(serverId, {
+        serverId,
+        hostname: null,
+        version: "test",
+        features: { workspaceMultiplicity: true, projectList: true },
+      });
+      const source = {};
+      function setDemand(demanded: boolean): void {
+        if (demand === "full") directory.setDemand(source, demanded);
+        else directory.setAgentRouteDemand(demanded ? ["agent-1"] : []);
+      }
+      const releaseProjects = client.holdProjectFetch();
+      const projects: ProjectListResult = { requestId: "projects", projects: [] };
+
+      try {
+        setDemand(true);
+        const refresh = directory.refreshDemand();
+        await expect
+          .poll(() => [client.observeAgentsCalls, client.observeWorkspacesCalls])
+          .toEqual([1, 1]);
+
+        await Promise.all(client.listSnapshots);
+        setDemand(false);
+        if (timing === "before") setDemand(true);
+        releaseProjects(projects);
+        await refresh;
+        if (timing === "after") setDemand(true);
+
+        await expect
+          .poll(() => [client.observeAgentsCalls, client.observeWorkspacesCalls])
+          .toEqual([2, 2]);
+        client.emit({
+          type: "agent_update",
+          payload: {
+            kind: "upsert",
+            agent: {
+              ...projectAgentSnapshot(createAgent(serverId, "agent-1")),
+              title: "Live title",
+            },
+            project: deriveProjectPlacementFromCwd("/repo"),
+          },
+        });
+        client.emit({
+          type: "workspace_update",
+          payload: {
+            kind: "upsert",
+            workspace: { ...createWorkspaceEntry("workspace-1"), name: "Live workspace" },
+          },
+        });
+
+        await expect
+          .poll(() => {
+            const session = useSessionStore.getState().sessions[serverId];
+            return [
+              session.agents.get("agent-1")?.title,
+              session.workspaces.get("workspace-1")?.name,
+            ];
+          })
+          .toEqual(["Live title", "Live workspace"]);
+      } finally {
+        directory.dispose();
+      }
+    },
+  );
+
   it("publishes cached directory data before the authoritative request completes", async () => {
     const serverId = "cached-directory";
     serverIds.add(serverId);

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

@@ -160,6 +160,7 @@ export class DirectorySync {
   private readonly fullDemandSources = new Set<object>();
   private readonly scriptStatusDemand = new Map<object, string>();
   private demandRefresh: Promise<void> | null = null;
+  private subscriptionGeneration = 0;
   private satisfiedDemandSource: DirectorySourceToken | null = null;
   private cursors: DirectoryCheckpoint = {};
 
@@ -262,6 +263,7 @@ export class DirectorySync {
   }
 
   private releaseSubscriptions(): void {
+    this.subscriptionGeneration += 1;
     for (const subscription of [
       this.agentSubscription,
       this.workspaceSubscription,
@@ -362,6 +364,7 @@ export class DirectorySync {
     this.observeDirectoryEvents();
     this.observeScriptStatusEvents();
     const source = this.connection.source;
+    const subscriptionGeneration = this.subscriptionGeneration;
     if (
       !force &&
       this.satisfiedDemandSource?.clientGeneration === source.clientGeneration &&
@@ -378,17 +381,19 @@ export class DirectorySync {
           ]).then(() => undefined);
     this.demandRefresh = refresh
       .then(() => {
-        this.satisfiedDemandSource = source;
+        if (this.subscriptionGeneration === subscriptionGeneration) {
+          this.satisfiedDemandSource = source;
+        }
         return undefined;
       })
       .finally(() => {
         this.demandRefresh = null;
         const current = this.connection.source;
-        if (
-          this.hasDemand() &&
-          (current.clientGeneration !== source.clientGeneration ||
-            current.connectionEpoch !== source.connectionEpoch)
-        ) {
+        const sourceChanged =
+          current.clientGeneration !== source.clientGeneration ||
+          current.connectionEpoch !== source.connectionEpoch;
+        const subscriptionsReleased = this.subscriptionGeneration !== subscriptionGeneration;
+        if (this.hasDemand() && (sourceChanged || subscriptionsReleased)) {
           void this.requestDemandRefresh().catch(() => undefined);
         }
       });