feat(FN-4370): complete Step 3 — retain badge snapshots across unsubscribes

Fusion-Task-Id: FN-4370
Fusion-Task-Lineage: c18188d5-c478-4233-943f-943b0125b271

packages/dashboard/app/hooks/__tests__/useBadgeWebSocket.test.ts

@@ -207,7 +207,7 @@ describe("useBadgeWebSocket", () => {
     expect(subscribeMsg).toBeDefined();
   });
 
-  it("sends unsubscribe, retains cached state, and closes the socket when the final subscription is removed", () => {
+  it("sends unsubscribe, moves cached state to retention, and closes the socket when the final subscription is removed", () => {
     const { result } = renderHook(() => useBadgeWebSocket());
 
     act(() => {
@@ -240,7 +240,76 @@ describe("useBadgeWebSocket", () => {
     });
     expect(unsubscribeMsg).toBeDefined();
     expect(MockWebSocket.instances[0].close).toHaveBeenCalled();
-    expect(result.current.badgeUpdates.has("default:FN-063")).toBe(true);
+    expect(result.current.badgeUpdates.has("default:FN-063")).toBe(false);
+  });
+
+  it("restores retained snapshot immediately after re-subscribing", () => {
+    const { result } = renderHook(() => useBadgeWebSocket());
+
+    act(() => {
+      result.current.subscribeToBadge("FN-063");
+      MockWebSocket.instances[0].emitOpen();
+      MockWebSocket.instances[0].emitMessage({
+        type: "badge:updated",
+        taskId: "FN-063",
+        prInfo: {
+          url: "https://github.com/owner/repo/pull/1",
+          number: 1,
+          status: "open",
+          title: "Tracked PR",
+          headBranch: "feature/test",
+          baseBranch: "main",
+          commentCount: 0,
+        },
+        timestamp: "2026-03-30T12:00:00.000Z",
+      });
+      result.current.unsubscribeFromBadge("FN-063");
+    });
+
+    expect(result.current.badgeUpdates.has("default:FN-063")).toBe(false);
+
+    act(() => {
+      result.current.subscribeToBadge("FN-063");
+    });
+
+    expect(result.current.badgeUpdates.get("default:FN-063")?.prInfo?.number).toBe(1);
+  });
+
+  it("evicts oldest retained snapshot when retention cap is exceeded", () => {
+    const { result } = renderHook(() => useBadgeWebSocket());
+
+    act(() => {
+      MockWebSocket.instances[0]?.emitOpen();
+      for (let i = 1; i <= 201; i++) {
+        const taskId = `FN-${String(i).padStart(3, "0")}`;
+        result.current.subscribeToBadge(taskId);
+        MockWebSocket.instances[0].emitMessage({
+          type: "badge:updated",
+          taskId,
+          prInfo: {
+            url: `https://github.com/owner/repo/pull/${i}`,
+            number: i,
+            status: "open",
+            title: `Tracked PR ${i}`,
+            headBranch: "feature/test",
+            baseBranch: "main",
+            commentCount: 0,
+          },
+          timestamp: "2026-03-30T12:00:00.000Z",
+        });
+        result.current.unsubscribeFromBadge(taskId);
+      }
+    });
+
+    act(() => {
+      result.current.subscribeToBadge("FN-001");
+    });
+    expect(result.current.badgeUpdates.get("default:FN-001")).toBeUndefined();
+
+    act(() => {
+      result.current.subscribeToBadge("FN-201");
+    });
+    expect(result.current.badgeUpdates.get("default:FN-201")?.prInfo?.number).toBe(201);
   });
 
   it("shares a single websocket and ref-counted subscription across hook instances", () => {
@@ -302,6 +371,41 @@ describe("useBadgeWebSocket", () => {
     expect(unsubscribeMsg).toBeDefined();
   });
 
+  it("reset clears retained snapshots", () => {
+    const { result, unmount } = renderHook(() => useBadgeWebSocket());
+
+    act(() => {
+      result.current.subscribeToBadge("FN-063");
+      MockWebSocket.instances[0].emitOpen();
+      MockWebSocket.instances[0].emitMessage({
+        type: "badge:updated",
+        taskId: "FN-063",
+        prInfo: {
+          url: "https://github.com/owner/repo/pull/1",
+          number: 1,
+          status: "open",
+          title: "Tracked PR",
+          headBranch: "feature/test",
+          baseBranch: "main",
+          commentCount: 0,
+        },
+        timestamp: "2026-03-30T12:00:00.000Z",
+      });
+      result.current.unsubscribeFromBadge("FN-063");
+    });
+
+    unmount();
+    __resetBadgeWebSocketStoreForTests();
+
+    const next = renderHook(() => useBadgeWebSocket());
+    act(() => {
+      next.result.current.subscribeToBadge("FN-063");
+    });
+
+    expect(next.result.current.badgeUpdates.get("default:FN-063")).toBeUndefined();
+    next.unmount();
+  });
+
   describe("projectId support", () => {
     it("includes projectId in WebSocket URL when provided", () => {
       const { result } = renderHook(() => useBadgeWebSocket("proj-123"));
@@ -395,6 +499,32 @@ describe("useBadgeWebSocket", () => {
       expect(newSubscribe).toBeGreaterThanOrEqual(1);
     });
 
+    it("invalidates retained snapshots on project change", () => {
+      const { result, rerender } = renderHook(
+        ({ projectId }: { projectId?: string }) => useBadgeWebSocket(projectId),
+        { initialProps: { projectId: "proj-A" } },
+      );
+
+      act(() => {
+        result.current.subscribeToBadge("FN-063");
+        MockWebSocket.instances[0].emitOpen();
+        MockWebSocket.instances[0].emitMessage({
+          type: "badge:updated",
+          taskId: "FN-063",
+          prInfo: { url: "https://github.com/owner/repo/pull/1", number: 1, status: "open", title: "Test PR", headBranch: "feat", baseBranch: "main", commentCount: 0 },
+          timestamp: "2026-03-30T12:00:00.000Z",
+        });
+        result.current.unsubscribeFromBadge("FN-063");
+      });
+
+      rerender({ projectId: "proj-B" });
+      act(() => {
+        result.current.subscribeToBadge("FN-063");
+      });
+
+      expect(result.current.badgeUpdates.get("proj-B:FN-063")).toBeUndefined();
+    });
+
     it("clears badge updates on project change", async () => {
       const { result, rerender } = renderHook(
         ({ projectId }: { projectId?: string }) => useBadgeWebSocket(projectId),

packages/dashboard/app/hooks/useBadgeWebSocket.ts

@@ -60,6 +60,7 @@ class BadgeWebSocketStore {
   private ws: WebSocket | null = null;
   private listeners = new Set<() => void>();
   private badgeUpdates = new Map<string, BadgeSnapshot>();
+  private retainedBadgeUpdates = new Map<string, BadgeSnapshot>();
   private subscriptionsByTask = new Map<string, Set<string>>();
   private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
   private reconnectDelayMs = 1_000;
@@ -76,6 +77,7 @@ class BadgeWebSocketStore {
   private previousProjectIdRef: string | null = null;
   private projectContextVersionRef = 0;
   private contextVersionAtStart = 0;
+  private static readonly RETAINED_BADGE_CAP = 200;
 
   subscribe(listener: () => void): () => void {
     this.listeners.add(listener);
@@ -127,6 +129,13 @@ class BadgeWebSocketStore {
 
   subscribeTask(hookId: string, taskId: string): void {
     const scopedKey = toScopedKey(this.projectId, taskId);
+    const retainedSnapshot = this.retainedBadgeUpdates.get(scopedKey);
+    if (retainedSnapshot && !this.badgeUpdates.has(scopedKey)) {
+      this.badgeUpdates.set(scopedKey, retainedSnapshot);
+      this.retainedBadgeUpdates.delete(scopedKey);
+      this.emit();
+    }
+
     const subscribers = this.subscriptionsByTask.get(scopedKey) ?? new Set<string>();
     const isNewSubscription = !subscribers.has(hookId);
     subscribers.add(hookId);
@@ -150,6 +159,18 @@ class BadgeWebSocketStore {
     subscribers.delete(hookId);
     if (subscribers.size === 0) {
       this.subscriptionsByTask.delete(scopedKey);
+      const snapshot = this.badgeUpdates.get(scopedKey);
+      if (snapshot) {
+        // Keep a bounded cache of last-known snapshots so cards can restore badge state
+        // immediately after transient viewport unsubscriptions. Cap prevents unbounded growth.
+        this.retainedBadgeUpdates.delete(scopedKey);
+        this.retainedBadgeUpdates.set(scopedKey, snapshot);
+        while (this.retainedBadgeUpdates.size > BadgeWebSocketStore.RETAINED_BADGE_CAP) {
+          const oldestKey = this.retainedBadgeUpdates.keys().next().value;
+          if (!oldestKey) break;
+          this.retainedBadgeUpdates.delete(oldestKey);
+        }
+      }
       this.badgeUpdates.delete(scopedKey);
       this.send({ type: "unsubscribe", taskId, projectId: this.projectId });
       this.emit();
@@ -172,6 +193,7 @@ class BadgeWebSocketStore {
   reset(): void {
     this.disconnect();
     this.badgeUpdates.clear();
+    this.retainedBadgeUpdates.clear();
     this.subscriptionsByTask.clear();
     this.shouldReconnect = false;
     this.emit();