fix(desktop): preserve workstream PR loading state

Signed-off-by: loganj <loganj@squareup.com>
Co-authored-by: Claude Code <noreply@anthropic.com>
Signed-off-by: loganj <loganj@squareup.com>

desktop/src/features/workstream-board/lib/workstreamPullRequestStatus.test.mjs

@@ -1,5 +1,21 @@
 import assert from "node:assert/strict";
-import test from "node:test";
+import { after, before, test } from "node:test";
+import { JSDOM } from "jsdom";
+
+const dom = new JSDOM("<!doctype html><html><body></body></html>", {
+  url: "http://localhost",
+});
+
+before(() => {
+  Object.assign(globalThis, {
+    document: dom.window.document,
+    HTMLElement: dom.window.HTMLElement,
+    IS_REACT_ACT_ENVIRONMENT: true,
+    window: dom.window,
+  });
+});
+
+after(() => dom.window.close());
 
 import {
   createWorkstreamPollingRegistry,
@@ -155,3 +171,101 @@ test("polling deduplicates canonical identities and tears down", async () => {
   assert.equal(aborted, true);
   resolveRequest();
 });
+
+test("workstream status hook exposes loading before a supported PR resolves", async () => {
+  const { cleanup, renderHook } = await import("@testing-library/react");
+  const { useWorkstreamPullRequestStatuses } = await import(
+    "./workstreamPullRequestStatus.ts"
+  );
+  const reference = parseWorkstreamPullRequestReference(
+    "https://github.com/block/buzz/pull/42",
+  );
+  const registry = createWorkstreamPollingRegistry({
+    fetcher: () => new Promise(() => {}),
+  });
+  const references = [reference];
+  const { result, unmount } = renderHook(() =>
+    useWorkstreamPullRequestStatuses(references, registry),
+  );
+  assert.deepEqual(result.current.get(reference.identity), {
+    status: "loading",
+  });
+  unmount();
+  cleanup();
+});
+
+test("workstream status hook preserves a synchronous live callback", async () => {
+  const { act, cleanup, renderHook } = await import("@testing-library/react");
+  const { useWorkstreamPullRequestStatuses } = await import(
+    "./workstreamPullRequestStatus.ts"
+  );
+  const reference = parseWorkstreamPullRequestReference(
+    "https://github.com/block/buzz/pull/42",
+  );
+  const registry = createWorkstreamPollingRegistry({
+    fetcher: async (url) =>
+      String(url).endsWith("/reviews?per_page=100")
+        ? response([])
+        : response({ number: 42, title: "Board", state: "open", draft: false }),
+    maxPolls: 1,
+  });
+  const prime = registry.subscribe(reference, () => {});
+  await act(async () => {
+    await Promise.resolve();
+    await Promise.resolve();
+  });
+  const references = [reference];
+  const { result, unmount } = renderHook(() =>
+    useWorkstreamPullRequestStatuses(references, registry),
+  );
+  await act(async () => {
+    await Promise.resolve();
+  });
+  assert.deepEqual(result.current.get(reference.identity), {
+    status: "live",
+    provider: "github",
+    href: "https://github.com/block/buzz/pull/42",
+    repository: "block/buzz",
+    number: 42,
+    title: "Board",
+    lifecycle: "open",
+    reviewState: "no-reviews",
+  });
+  unmount();
+  prime();
+  cleanup();
+});
+
+test("workstream status hook removes stale references and preserves unsupported state", async () => {
+  const { act, cleanup, renderHook } = await import("@testing-library/react");
+  const { useWorkstreamPullRequestStatuses } = await import(
+    "./workstreamPullRequestStatus.ts"
+  );
+  const supported = parseWorkstreamPullRequestReference(
+    "https://github.com/block/buzz/pull/42",
+  );
+  const unsupported = parseWorkstreamPullRequestReference(
+    "https://example.com/pr/1",
+  );
+  const registry = createWorkstreamPollingRegistry({
+    fetcher: () => new Promise(() => {}),
+  });
+  const { result, rerender, unmount } = renderHook(
+    ({ refs }) => useWorkstreamPullRequestStatuses(refs, registry),
+    { initialProps: { refs: [supported, unsupported] } },
+  );
+  assert.deepEqual(
+    result.current.get(unsupported.identity)?.status,
+    "unsupported",
+  );
+  await act(async () => {
+    rerender({ refs: [unsupported] });
+  });
+  assert.equal(result.current.has(supported.identity), false);
+  assert.deepEqual(
+    result.current.get(unsupported.identity)?.status,
+    "unsupported",
+  );
+  unmount();
+  cleanup();
+});

desktop/src/features/workstream-board/lib/workstreamPullRequestStatus.ts

@@ -507,8 +507,14 @@ export function createWorkstreamPollingRegistry(
 
 const defaultWorkstreamPollingRegistry = createWorkstreamPollingRegistry();
 
+type WorkstreamPollingRegistry = Pick<
+  ReturnType<typeof createWorkstreamPollingRegistry>,
+  "subscribe"
+>;
+
 export function useWorkstreamPullRequestStatuses(
   references: readonly WorkstreamPullRequestReference[],
+  pollingRegistry: WorkstreamPollingRegistry = defaultWorkstreamPollingRegistry,
 ): ReadonlyMap<string, WorkstreamPullRequestDisplayState> {
   const [states, setStates] = React.useState<
     ReadonlyMap<string, WorkstreamPullRequestDisplayState>
@@ -520,34 +526,37 @@ export function useWorkstreamPullRequestStatuses(
   }, [references]);
 
   React.useEffect(() => {
-    const next = new Map<string, WorkstreamPullRequestDisplayState>();
+    const initial = new Map<string, WorkstreamPullRequestDisplayState>();
+    for (const reference of referencesByIdentity.values()) {
+      initial.set(
+        reference.identity,
+        reference.kind === "unsupported"
+          ? {
+              status: "unsupported",
+              href: reference.href,
+              reason: reference.reason,
+            }
+          : { status: "loading" },
+      );
+    }
+    setStates(initial);
+
     const unsubscribers: Array<() => void> = [];
     for (const reference of referencesByIdentity.values()) {
-      if (reference.kind === "unsupported") {
-        next.set(reference.identity, {
-          status: "unsupported",
-          href: reference.href,
-          reason: reference.reason,
+      if (reference.kind === "unsupported") continue;
+      const unsubscribe = pollingRegistry.subscribe(reference, (state) => {
+        setStates((current) => {
+          const updated = new Map(current);
+          updated.set(reference.identity, state);
+          return updated;
         });
-        continue;
-      }
-      const unsubscribe = defaultWorkstreamPollingRegistry.subscribe(
-        reference,
-        (state) => {
-          setStates((current) => {
-            const updated = new Map(current);
-            updated.set(reference.identity, state);
-            return updated;
-          });
-        },
-      );
+      });
       unsubscribers.push(unsubscribe);
     }
-    setStates(next);
     return () => {
       for (const unsubscribe of unsubscribers) unsubscribe();
     };
-  }, [referencesByIdentity]);
+  }, [pollingRegistry, referencesByIdentity]);
 
   return states;
 }

desktop/tests/e2e/workstream-board-b3-pr.capture.spec.ts

@@ -9,13 +9,34 @@ test("continuous relay-backed Workstream Board bullet 3 PR-status demo", async (
   await installRelayBridge(page, "tyler", { seedPreviewFeatures: true });
   await page.routeWebSocket(/localhost:3000/, (socket) => {
     const server = socket.connectToServer();
+    let delayedPrSubscription: string | null = null;
+    let delayedFrames: string[] = [];
+    let releaseScheduled = false;
+    const release = () => {
+      const frames = delayedFrames;
+      delayedFrames = [];
+      delayedPrSubscription = null;
+      releaseScheduled = false;
+      for (const frame of frames) socket.send(frame);
+    };
     server.onMessage((message) => {
       if (typeof message === "string") {
         try {
           const payload = JSON.parse(message) as unknown[];
+          const type = payload[0];
+          const subId = typeof payload[1] === "string" ? payload[1] : null;
           const event = payload[2] as { kind?: unknown } | undefined;
-          if (payload[0] === "EVENT" && event?.kind === 1618) {
-            setTimeout(() => socket.send(message), 2_500);
+          if (type === "EVENT" && event?.kind === 1618 && subId) {
+            delayedPrSubscription = subId;
+            delayedFrames.push(message);
+            if (!releaseScheduled) {
+              releaseScheduled = true;
+              setTimeout(release, 2_500);
+            }
+            return;
+          }
+          if (type === "EOSE" && subId === delayedPrSubscription) {
+            delayedFrames.push(message);
             return;
           }
         } catch {