Keep git review state fresh after reconnect (#1491)

packages/app/src/git/checkout-status-cache.test.ts

@@ -1,11 +1,17 @@
+// @vitest-environment jsdom
+// The review draft store persists through AsyncStorage's web shim, which needs window.
+import "@/test/window-local-storage";
 import { QueryClient } from "@tanstack/react-query";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { CheckoutStatusUpdate } from "@getpaseo/protocol/messages";
-import { checkoutStatusQueryKey } from "@/git/query-keys";
+import { checkoutPrStatusQueryKey, checkoutStatusQueryKey } from "@/git/query-keys";
+import { prPaneTimelineQueryKey } from "@/git/pull-request-panel/query-keys";
+import { resetReviewDraftStore, useReviewDraftStore } from "@/review/store";
 import {
-  applyCheckoutStatusUpdate,
+  applyCheckoutStatusUpdateFromEvent,
+  type CheckoutPrStatusPayload,
   type CheckoutStatusPayload,
-  peekOrFetchCheckoutStatus,
+  fetchCheckoutStatus,
 } from "./checkout-status-cache";
 
 const serverId = "server-1";
@@ -31,94 +37,179 @@ function checkoutStatus(overrides: Partial<CheckoutStatusPayload> = {}): Checkou
   } as CheckoutStatusPayload;
 }
 
-function checkoutStatusUpdate(payload: CheckoutStatusPayload): CheckoutStatusUpdate {
-  return { type: "checkout_status_update", payload };
+function prStatus(overrides: Partial<CheckoutPrStatusPayload> = {}): CheckoutPrStatusPayload {
+  return {
+    cwd,
+    status: {
+      url: "https://github.com/getpaseo/paseo/pull/42",
+      title: "My PR",
+      state: "open",
+      baseRefName: "main",
+      headRefName: "feature",
+      isMerged: false,
+      isDraft: false,
+      mergeable: "MERGEABLE",
+      checks: [],
+      checksStatus: "success",
+      reviewDecision: null,
+    },
+    githubFeaturesEnabled: true,
+    error: null,
+    requestId: "pr-status-1",
+    ...overrides,
+  };
+}
+
+function checkoutStatusUpdate(
+  payload: CheckoutStatusPayload,
+  extraPrStatus?: CheckoutPrStatusPayload,
+): CheckoutStatusUpdate {
+  return {
+    type: "checkout_status_update",
+    payload: extraPrStatus ? { ...payload, prStatus: extraPrStatus } : payload,
+  };
+}
+
+function setDiffModeOverride(isDirtyAtSelection: boolean): void {
+  useReviewDraftStore.getState().setDiffModeOverride({
+    scopeKey: "review:scope",
+    override: { serverId, cwd, mode: "base", isDirtyAtSelection },
+  });
 }
 
 function createQueryClient(): QueryClient {
-  return new QueryClient({
-    defaultOptions: { queries: { retry: false } },
-  });
+  return new QueryClient({ defaultOptions: { queries: { retry: false } } });
 }
 
-describe("peekOrFetchCheckoutStatus", () => {
-  it("fetches from the client and writes the result to the cache on a cold read", async () => {
-    const queryClient = createQueryClient();
-    const fetched = checkoutStatus({ requestId: "cold-read" });
+beforeEach(() => {
+  resetReviewDraftStore();
+});
+
+describe("fetchCheckoutStatus", () => {
+  it("fetches from the client and returns the payload", async () => {
+    const fetched = checkoutStatus({ requestId: "fetch-1" });
     const client = { getCheckoutStatus: vi.fn(async () => fetched) };
 
-    const result = await peekOrFetchCheckoutStatus({ queryClient, client, serverId, cwd });
+    const result = await fetchCheckoutStatus({ client, serverId, cwd });
 
     expect(result).toEqual(fetched);
     expect(client.getCheckoutStatus).toHaveBeenCalledExactlyOnceWith(cwd);
-    expect(queryClient.getQueryData(checkoutStatusQueryKey(serverId, cwd))).toEqual(fetched);
   });
 
-  it("returns the cached snapshot without calling the client when the cache already has data", async () => {
-    const queryClient = createQueryClient();
-    const cached = checkoutStatus({ requestId: "cached" });
-    queryClient.setQueryData(checkoutStatusQueryKey(serverId, cwd), cached);
-    const client = {
-      getCheckoutStatus: vi.fn(async () => checkoutStatus({ requestId: "uncached" })),
-    };
+  it("expires a manual diff-mode override when the fetched dirty state flipped", async () => {
+    setDiffModeOverride(true);
+    const client = { getCheckoutStatus: vi.fn(async () => checkoutStatus({ isDirty: false })) };
 
-    const result = await peekOrFetchCheckoutStatus({ queryClient, client, serverId, cwd });
+    await fetchCheckoutStatus({ client, serverId, cwd });
 
-    expect(result).toEqual(cached);
-    expect(client.getCheckoutStatus).not.toHaveBeenCalled();
+    expect(useReviewDraftStore.getState().diffModeOverrides["review:scope"]).toBeUndefined();
   });
 });
 
-describe("applyCheckoutStatusUpdate", () => {
-  it("writes the pushed payload into the cache key for the matching cwd", () => {
+describe("applyCheckoutStatusUpdateFromEvent", () => {
+  it("writes the checkout status to the cache using the cwd from the payload", () => {
     const queryClient = createQueryClient();
-    const pushed = checkoutStatus({
-      requestId: "server-push",
-      currentBranch: "pushed-branch",
-      isDirty: true,
-      aheadBehind: { ahead: 2, behind: 1 },
-      aheadOfOrigin: 2,
-      behindOfOrigin: 1,
-    });
+    const pushed = checkoutStatus({ requestId: "push-1", isDirty: true });
 
-    applyCheckoutStatusUpdate({
+    applyCheckoutStatusUpdateFromEvent({
       queryClient,
       serverId,
-      cwd,
       message: checkoutStatusUpdate(pushed),
     });
 
     expect(queryClient.getQueryData(checkoutStatusQueryKey(serverId, cwd))).toEqual(pushed);
   });
 
-  it("ignores updates whose payload cwd does not match the subscribed cwd", () => {
+  it("writes the PR status cache when prStatus is present, and skips it otherwise", () => {
     const queryClient = createQueryClient();
-    const otherCwd = "/other-repo";
-    const otherCached = checkoutStatus({
-      cwd: otherCwd,
-      repoRoot: otherCwd,
-      requestId: "other-cached",
-      currentBranch: "other-main",
-    });
-    queryClient.setQueryData(checkoutStatusQueryKey(serverId, otherCwd), otherCached);
+    const pushedPr = prStatus({ requestId: "pr-1" });
 
-    applyCheckoutStatusUpdate({
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus(), pushedPr),
+    });
+    expect(queryClient.getQueryData(checkoutPrStatusQueryKey(serverId, cwd))).toEqual(pushedPr);
+
+    const otherCwd = "/repo2";
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus({ cwd: otherCwd, repoRoot: otherCwd })),
+    });
+    expect(queryClient.getQueryData(checkoutPrStatusQueryKey(serverId, otherCwd))).toBeUndefined();
+  });
+
+  it("expires a manual diff-mode override when the pushed dirty state flipped", () => {
+    const queryClient = createQueryClient();
+    setDiffModeOverride(false);
+
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus({ isDirty: true })),
+    });
+
+    expect(useReviewDraftStore.getState().diffModeOverrides["review:scope"]).toBeUndefined();
+  });
+
+  it("keeps a manual diff-mode override while the pushed dirty state still matches", () => {
+    const queryClient = createQueryClient();
+    setDiffModeOverride(true);
+
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus({ isDirty: true })),
+    });
+
+    expect(useReviewDraftStore.getState().diffModeOverrides["review:scope"]).toBeDefined();
+  });
+
+  it("invalidates the PR timeline when the prStatus changes, ignoring the volatile requestId", () => {
+    const queryClient = createQueryClient();
+    queryClient.setQueryData(
+      checkoutPrStatusQueryKey(serverId, cwd),
+      prStatus({ requestId: "pr-v1" }),
+    );
+    const timelineKey = prPaneTimelineQueryKey({ serverId, cwd, prNumber: 42 });
+    queryClient.setQueryData(timelineKey, { items: [] });
+
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus(), prStatus({ requestId: "pr-v2" })),
+    });
+    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(false);
+
+    applyCheckoutStatusUpdateFromEvent({
       queryClient,
       serverId,
-      cwd,
       message: checkoutStatusUpdate(
-        checkoutStatus({
-          cwd: otherCwd,
-          repoRoot: otherCwd,
-          requestId: "server-push",
-          currentBranch: "pushed-branch",
+        checkoutStatus(),
+        prStatus({
+          requestId: "pr-v3",
+          status: { ...prStatus().status!, state: "closed" },
         }),
       ),
     });
+    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(true);
+  });
 
-    expect(queryClient.getQueryData(checkoutStatusQueryKey(serverId, cwd))).toBeUndefined();
-    expect(queryClient.getQueryData(checkoutStatusQueryKey(serverId, otherCwd))).toEqual(
-      otherCached,
-    );
+  it("invalidates the PR timeline on the first prStatus emission, scoped to its cwd", () => {
+    const queryClient = createQueryClient();
+    const timelineKey = prPaneTimelineQueryKey({ serverId, cwd, prNumber: 42 });
+    const otherTimelineKey = prPaneTimelineQueryKey({ serverId, cwd: "/repo2", prNumber: 42 });
+    queryClient.setQueryData(timelineKey, { items: [] });
+    queryClient.setQueryData(otherTimelineKey, { items: [] });
+
+    applyCheckoutStatusUpdateFromEvent({
+      queryClient,
+      serverId,
+      message: checkoutStatusUpdate(checkoutStatus(), prStatus()),
+    });
+
+    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(true);
+    expect(queryClient.getQueryState(otherTimelineKey)?.isInvalidated).toBe(false);
   });
 });

packages/app/src/git/checkout-status-cache.ts

@@ -1,48 +1,86 @@
 import type { QueryClient } from "@tanstack/react-query";
 import type { CheckoutStatusResponse, CheckoutStatusUpdate } from "@getpaseo/protocol/messages";
-import { checkoutStatusQueryKey } from "@/git/query-keys";
+import equal from "fast-deep-equal/es6";
+import {
+  checkoutPrStatusQueryKey,
+  checkoutStatusQueryKey,
+  invalidatePrPaneTimelineForCheckout,
+} from "@/git/query-keys";
+import { expireStaleDiffModeOverrides } from "@/review/store";
 
 export type CheckoutStatusPayload = CheckoutStatusResponse["payload"];
+export type CheckoutPrStatusPayload = NonNullable<CheckoutStatusUpdate["payload"]["prStatus"]>;
 
 export interface CheckoutStatusClient {
   getCheckoutStatus: (cwd: string) => Promise<CheckoutStatusPayload>;
 }
 
-export async function peekOrFetchCheckoutStatus({
-  queryClient,
+// Checkout status enters the app through exactly two doors: daemon pushes
+// (applyCheckoutStatusUpdateFromEvent) and query fetches (fetchCheckoutStatus). Both run
+// the dirty-state reactions, so they hold regardless of which screens are mounted.
+
+export async function fetchCheckoutStatus({
   client,
   serverId,
   cwd,
 }: {
-  queryClient: QueryClient;
   client: CheckoutStatusClient;
   serverId: string;
   cwd: string;
 }): Promise<CheckoutStatusPayload> {
-  const queryKey = checkoutStatusQueryKey(serverId, cwd);
-  const cached = queryClient.getQueryData<CheckoutStatusPayload>(queryKey);
-  if (cached) {
-    return cached;
-  }
-
-  const snapshot = await client.getCheckoutStatus(cwd);
-  queryClient.setQueryData(queryKey, snapshot);
-  return snapshot;
+  const payload = await client.getCheckoutStatus(cwd);
+  expireStaleDiffModeOverrides({ serverId, cwd, isDirty: payload.isGit && payload.isDirty });
+  return payload;
 }
 
-export function applyCheckoutStatusUpdate({
+export function applyCheckoutStatusUpdateFromEvent({
   queryClient,
   serverId,
-  cwd,
   message,
 }: {
   queryClient: QueryClient;
   serverId: string;
-  cwd: string;
   message: CheckoutStatusUpdate;
 }): void {
-  if (message.payload.cwd !== cwd) {
+  const { payload } = message;
+  queryClient.setQueryData(checkoutStatusQueryKey(serverId, payload.cwd), payload);
+  expireStaleDiffModeOverrides({
+    serverId,
+    cwd: payload.cwd,
+    isDirty: payload.isGit && payload.isDirty,
+  });
+
+  const prStatus = payload.prStatus;
+  if (!prStatus) {
     return;
   }
-  queryClient.setQueryData(checkoutStatusQueryKey(serverId, cwd), message.payload);
+
+  const previous = queryClient.getQueryData<CheckoutPrStatusPayload>(
+    checkoutPrStatusQueryKey(serverId, prStatus.cwd),
+  );
+  queryClient.setQueryData(checkoutPrStatusQueryKey(serverId, prStatus.cwd), prStatus);
+
+  // The PR activity timeline has no push channel; mark it stale when the pushed PR status
+  // meaningfully changed. Active panes refetch immediately, evicted ones on next mount.
+  if (hasPrStatusChanged(previous, prStatus)) {
+    void invalidatePrPaneTimelineForCheckout(queryClient, { serverId, cwd: prStatus.cwd });
+  }
+}
+
+// requestId changes on every emission and carries no PR state.
+function prStatusWithoutVolatileFields(
+  prStatus: CheckoutPrStatusPayload,
+): Omit<CheckoutPrStatusPayload, "requestId"> {
+  const { requestId: _requestId, ...rest } = prStatus;
+  return rest;
+}
+
+function hasPrStatusChanged(
+  previous: CheckoutPrStatusPayload | undefined,
+  next: CheckoutPrStatusPayload,
+): boolean {
+  if (!previous) {
+    return true;
+  }
+  return !equal(prStatusWithoutVolatileFields(previous), prStatusWithoutVolatileFields(next));
 }

packages/app/src/git/diff-pane.tsx

@@ -101,11 +101,10 @@ import {
 import {
   buildReviewDraftScopeKey,
   buildReviewDraftKey,
-  useActiveReviewDraftMode,
   useReviewAttachmentSnapshot,
-  useSetActiveReviewDraftMode,
+  useResolvedDiffMode,
+  useSetDiffModeOverride,
   type ReviewDraftComment,
-  type ReviewDraftMode,
   getInlineReviewThreadState,
   getSplitInlineReviewThreadState,
   InlineReviewGutterCell,
@@ -1593,7 +1592,6 @@ export function GitDiffPane({
   const isMobile = useIsCompactFormFactor();
   const showDesktopWebScrollbar = isWeb && !isMobile;
   const canUseSplitLayout = isWeb && !isMobile;
-  const [diffModeOverride, setDiffModeOverride] = useState<ReviewDraftMode | null>(null);
   const { preferences: changesPreferences, updatePreferences: updateChangesPreferences } =
     useChangesPreferences();
   const wrapLines = changesPreferences.wrapLines;
@@ -1614,9 +1612,6 @@ export function GitDiffPane({
     void updateChangesPreferences({ hideWhitespace: !changesPreferences.hideWhitespace });
   }, [changesPreferences.hideWhitespace, updateChangesPreferences]);
 
-  // handleSelectUncommitted/handleSelectBase are defined later, after reviewDraftScopeKey
-  // and setActiveReviewMode are available, so they can record the active review mode.
-
   const handleLayoutUnified = useCallback(() => {
     handleLayoutChange("unified");
   }, [handleLayoutChange]);
@@ -1691,8 +1686,6 @@ export function GitDiffPane({
   const statusState = deriveStatusState({ status, isStatusLoading, isStatusError, statusError });
   const { isGit, notGit, statusErrorMessage, baseRef, hasUncommittedChanges } = statusState;
 
-  // Auto-select diff mode based on state: uncommitted when dirty, base when clean
-  const autoDiffMode: ReviewDraftMode = hasUncommittedChanges ? "uncommitted" : "base";
   const reviewDraftScopeKey = useMemo(
     () =>
       buildReviewDraftScopeKey({
@@ -1704,8 +1697,11 @@ export function GitDiffPane({
       }),
     [baseRef, changesPreferences.hideWhitespace, cwd, serverId, workspaceId],
   );
-  const activeReviewMode = useActiveReviewDraftMode({ scopeKey: reviewDraftScopeKey });
-  const diffMode = diffModeOverride ?? activeReviewMode ?? autoDiffMode;
+  const diffMode = useResolvedDiffMode({
+    scopeKey: reviewDraftScopeKey,
+    hasUncommittedChanges,
+  });
+  const setDiffModeOverride = useSetDiffModeOverride();
 
   const {
     files,
@@ -1731,17 +1727,20 @@ export function GitDiffPane({
       }),
     [baseRef, changesPreferences.hideWhitespace, cwd, diffMode, serverId, workspaceId],
   );
-  const setActiveReviewMode = useSetActiveReviewDraftMode();
 
   const handleSelectUncommitted = useCallback(() => {
-    setDiffModeOverride("uncommitted");
-    setActiveReviewMode({ scopeKey: reviewDraftScopeKey, mode: "uncommitted" });
-  }, [reviewDraftScopeKey, setActiveReviewMode]);
+    setDiffModeOverride({
+      scopeKey: reviewDraftScopeKey,
+      override: { serverId, cwd, mode: "uncommitted", isDirtyAtSelection: hasUncommittedChanges },
+    });
+  }, [cwd, hasUncommittedChanges, reviewDraftScopeKey, serverId, setDiffModeOverride]);
 
   const handleSelectBase = useCallback(() => {
-    setDiffModeOverride("base");
-    setActiveReviewMode({ scopeKey: reviewDraftScopeKey, mode: "base" });
-  }, [reviewDraftScopeKey, setActiveReviewMode]);
+    setDiffModeOverride({
+      scopeKey: reviewDraftScopeKey,
+      override: { serverId, cwd, mode: "base", isDirtyAtSelection: hasUncommittedChanges },
+    });
+  }, [cwd, hasUncommittedChanges, reviewDraftScopeKey, serverId, setDiffModeOverride]);
 
   const reviewActions = useInlineReviewController({
     reviewDraftKey,
@@ -1994,11 +1993,6 @@ export function GitDiffPane({
     }
   }, [allExpanded, files, setDiffExpandedPathsForWorkspace, workspaceStateKey]);
 
-  // Clear diff mode override when auto mode changes (e.g., after commit)
-  useEffect(() => {
-    setDiffModeOverride(null);
-  }, [autoDiffMode]);
-
   const renderFlatItem = useCallback(
     ({ item }: { item: DiffFlatItem }) => {
       if (item.type === "header") {

packages/app/src/git/pull-request-panel/use-data.ts

@@ -244,7 +244,9 @@ export function usePrPaneData({
     },
     enabled: shouldFetchTimeline,
     staleTime: Infinity,
-    refetchOnMount: false,
+    // Refetch on mount only after explicit invalidation (reconnect, or a pushed PR status
+    // change) — see useCheckoutStatusQuery for the rationale.
+    refetchOnMount: true,
     refetchOnReconnect: false,
     refetchOnWindowFocus: false,
     placeholderData: keepPreviousData,

packages/app/src/git/query-keys.test.ts

@@ -5,6 +5,7 @@ import {
   checkoutPrStatusQueryKey,
   checkoutStatusQueryKey,
   invalidateCheckoutGitQueriesForClient,
+  invalidateCheckoutGitQueriesForServer,
 } from "@/git/query-keys";
 import { prPaneTimelineQueryKey } from "@/git/pull-request-panel/query-keys";
 
@@ -59,4 +60,47 @@ describe("checkout query keys", () => {
 
     queryClient.clear();
   });
+
+  it("invalidates fetch-based checkout queries server-wide without touching other servers", async () => {
+    const queryClient = new QueryClient();
+    const otherServerId = "server-2";
+    const otherCwd = "/tmp/repo-2";
+
+    queryClient.setQueryData(checkoutStatusQueryKey(serverId, cwd), { isGit: true });
+    queryClient.setQueryData(checkoutStatusQueryKey(serverId, otherCwd), { isGit: true });
+    queryClient.setQueryData(checkoutPrStatusQueryKey(serverId, cwd), { status: { number: 12 } });
+    queryClient.setQueryData(prPaneTimelineQueryKey({ serverId, cwd, prNumber: 12 }), {
+      items: [],
+    });
+    // Subscription-fed diff queries are deliberately not part of the server-wide sweep.
+    queryClient.setQueryData(checkoutDiffQueryKey(serverId, cwd, "base", "main", true), {
+      files: [],
+    });
+    queryClient.setQueryData(checkoutStatusQueryKey(otherServerId, cwd), { isGit: true });
+
+    await invalidateCheckoutGitQueriesForServer(queryClient, serverId);
+
+    expect(queryClient.getQueryState(checkoutStatusQueryKey(serverId, cwd))?.isInvalidated).toBe(
+      true,
+    );
+    expect(
+      queryClient.getQueryState(checkoutStatusQueryKey(serverId, otherCwd))?.isInvalidated,
+    ).toBe(true);
+    expect(queryClient.getQueryState(checkoutPrStatusQueryKey(serverId, cwd))?.isInvalidated).toBe(
+      true,
+    );
+    expect(
+      queryClient.getQueryState(prPaneTimelineQueryKey({ serverId, cwd, prNumber: 12 }))
+        ?.isInvalidated,
+    ).toBe(true);
+    expect(
+      queryClient.getQueryState(checkoutDiffQueryKey(serverId, cwd, "base", "main", true))
+        ?.isInvalidated,
+    ).toBe(false);
+    expect(
+      queryClient.getQueryState(checkoutStatusQueryKey(otherServerId, cwd))?.isInvalidated,
+    ).toBe(false);
+
+    queryClient.clear();
+  });
 });

packages/app/src/git/query-keys.ts

@@ -6,6 +6,11 @@ interface CheckoutQueryIdentity {
   cwd: string;
 }
 
+interface CheckoutQueryScope {
+  serverId: string;
+  cwd?: string;
+}
+
 type CheckoutQueryKey = readonly unknown[];
 
 export function checkoutStatusQueryKey(serverId: string, cwd: string) {
@@ -46,17 +51,41 @@ export async function invalidateCheckoutGitQueriesForClient(
   ]);
 }
 
+// checkoutDiff is excluded: diff queries are subscription-fed (queryFn: skipToken) and
+// receive a fresh snapshot on every resubscribe, so invalidation cannot and need not
+// refetch them.
+export async function invalidateCheckoutGitQueriesForServer(
+  queryClient: QueryClient,
+  serverId: string,
+) {
+  const kinds = ["checkoutStatus", "checkoutPrStatus", prPaneTimelineQueryKind];
+  await Promise.all(
+    kinds.map((kind) =>
+      queryClient.invalidateQueries({ predicate: checkoutQueryPredicate(kind, { serverId }) }),
+    ),
+  );
+}
+
+export async function invalidatePrPaneTimelineForCheckout(
+  queryClient: QueryClient,
+  identity: CheckoutQueryIdentity,
+) {
+  await queryClient.invalidateQueries({
+    predicate: checkoutQueryPredicate(prPaneTimelineQueryKind, identity),
+  });
+}
+
 function checkoutQueryPredicate(
   queryKind: CheckoutQueryKey[0],
-  identity: CheckoutQueryIdentity,
+  scope: CheckoutQueryScope,
 ): (query: Query) => boolean {
   return (query) => {
     const key = query.queryKey;
     return (
       isCheckoutQueryKey(key) &&
       key[0] === queryKind &&
-      key[1] === identity.serverId &&
-      key[2] === identity.cwd
+      key[1] === scope.serverId &&
+      (scope.cwd === undefined || key[2] === scope.cwd)
     );
   };
 }

packages/app/src/git/use-pr-status-query.ts

@@ -1,5 +1,4 @@
-import { useQuery, useQueryClient } from "@tanstack/react-query";
-import { useEffect } from "react";
+import { useQuery } from "@tanstack/react-query";
 import { useTranslation } from "react-i18next";
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
 import type { CheckoutPrStatusResponse } from "@getpaseo/protocol/messages";
@@ -81,24 +80,9 @@ export function useCheckoutPrStatusQuery({
   enabled = true,
 }: UseCheckoutPrStatusQueryOptions) {
   const { t } = useTranslation();
-  const queryClient = useQueryClient();
   const client = useHostRuntimeClient(serverId);
   const isConnected = useHostRuntimeIsConnected(serverId);
 
-  useEffect(() => {
-    if (!client || !isConnected || !cwd) {
-      return;
-    }
-
-    return client.on("checkout_status_update", (message) => {
-      const prStatus = message.payload.prStatus;
-      if (!prStatus || prStatus.cwd !== cwd) {
-        return;
-      }
-      queryClient.setQueryData(checkoutPrStatusQueryKey(serverId, cwd), prStatus);
-    });
-  }, [client, isConnected, cwd, queryClient, serverId]);
-
   const query = useQuery({
     queryKey: checkoutPrStatusQueryKey(serverId, cwd),
     queryFn: async () => {
@@ -109,7 +93,9 @@ export function useCheckoutPrStatusQuery({
     },
     enabled: !!client && isConnected && !!cwd && enabled,
     staleTime: Infinity,
-    refetchOnMount: false,
+    // Refetch on mount only after explicit invalidation (e.g. reconnect) — see
+    // useCheckoutStatusQuery for the rationale.
+    refetchOnMount: true,
     refetchOnReconnect: false,
     refetchOnWindowFocus: false,
   });
@@ -131,24 +117,9 @@ export function useWorkspacePrHint({
   enabled = true,
 }: UseCheckoutPrStatusQueryOptions): PrHint | null {
   const { t } = useTranslation();
-  const queryClient = useQueryClient();
   const client = useHostRuntimeClient(serverId);
   const isConnected = useHostRuntimeIsConnected(serverId);
 
-  useEffect(() => {
-    if (!client || !isConnected || !cwd) {
-      return;
-    }
-
-    return client.on("checkout_status_update", (message) => {
-      const prStatus = message.payload.prStatus;
-      if (!prStatus || prStatus.cwd !== cwd) {
-        return;
-      }
-      queryClient.setQueryData(checkoutPrStatusQueryKey(serverId, cwd), prStatus);
-    });
-  }, [client, isConnected, cwd, queryClient, serverId]);
-
   const query = useQuery<CheckoutPrStatusPayload, Error, PrHint | null>({
     queryKey: checkoutPrStatusQueryKey(serverId, cwd),
     queryFn: async () => {
@@ -159,7 +130,9 @@ export function useWorkspacePrHint({
     },
     enabled: !!client && isConnected && !!cwd && enabled,
     staleTime: Infinity,
-    refetchOnMount: false,
+    // Refetch on mount only after explicit invalidation (e.g. reconnect) — see
+    // useCheckoutStatusQuery for the rationale.
+    refetchOnMount: true,
     refetchOnReconnect: false,
     refetchOnWindowFocus: false,
     select: selectWorkspacePrHint,

packages/app/src/git/use-status-query.ts

@@ -1,14 +1,8 @@
-import { useQuery, useQueryClient } from "@tanstack/react-query";
-import { useEffect } from "react";
+import { useQuery } from "@tanstack/react-query";
 import { useTranslation } from "react-i18next";
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
 import { checkoutStatusQueryKey } from "@/git/query-keys";
-import {
-  applyCheckoutStatusUpdate,
-  type CheckoutStatusClient,
-  type CheckoutStatusPayload,
-  peekOrFetchCheckoutStatus,
-} from "./checkout-status-cache";
+import { fetchCheckoutStatus } from "./checkout-status-cache";
 
 export type { CheckoutStatusPayload } from "./checkout-status-cache";
 
@@ -19,40 +13,25 @@ interface UseCheckoutStatusQueryOptions {
   cwd: string;
 }
 
-function fetchCheckoutStatus(
-  client: CheckoutStatusClient,
-  cwd: string,
-): Promise<CheckoutStatusPayload> {
-  return client.getCheckoutStatus(cwd);
-}
-
 export function useCheckoutStatusQuery({ serverId, cwd }: UseCheckoutStatusQueryOptions) {
   const { t } = useTranslation();
-  const queryClient = useQueryClient();
   const client = useHostRuntimeClient(serverId);
   const isConnected = useHostRuntimeIsConnected(serverId);
 
-  useEffect(() => {
-    if (!client || !isConnected || !cwd) {
-      return;
-    }
-
-    return client.on("checkout_status_update", (message) => {
-      applyCheckoutStatusUpdate({ queryClient, serverId, cwd, message });
-    });
-  }, [client, isConnected, cwd, queryClient, serverId]);
-
   const query = useQuery({
     queryKey: checkoutStatusQueryKey(serverId, cwd),
     queryFn: async () => {
       if (!client) {
         throw new Error(t("common.errors.daemonClientUnavailable"));
       }
-      return await peekOrFetchCheckoutStatus({ queryClient, client, serverId, cwd });
+      return await fetchCheckoutStatus({ client, serverId, cwd });
     },
     enabled: !!client && isConnected && !!cwd,
     staleTime: Infinity,
-    refetchOnMount: false,
+    // Freshness is push-driven (checkout_status_update applied globally); with
+    // staleTime: Infinity, refetchOnMount only fires after an explicit invalidation
+    // (e.g. reconnect), which is exactly when the push stream may have been missed.
+    refetchOnMount: true,
     refetchOnReconnect: false,
     refetchOnWindowFocus: false,
   });
@@ -81,7 +60,7 @@ export function useCheckoutStatusCacheOnly({ serverId, cwd }: UseCheckoutStatusQ
       if (!client) {
         throw new Error(t("common.errors.daemonClientUnavailable"));
       }
-      return await fetchCheckoutStatus(client, cwd);
+      return await fetchCheckoutStatus({ client, serverId, cwd });
     },
     enabled: false,
     staleTime: CHECKOUT_STATUS_STALE_TIME,