refactor(app): centralize checkout query keys (#844)

packages/app/src/git/actions-store.test.ts

@@ -1,12 +1,10 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { QueryClient } from "@tanstack/react-query";
 import type { DaemonClient } from "@server/client/daemon-client";
 import { queryClient as appQueryClient } from "@/query/query-client";
 import { useSessionStore } from "@/stores/session-store";
 import type { WorkspaceDescriptor } from "@/stores/session-store";
 import {
   __resetCheckoutGitActionsStoreForTests,
-  invalidateCheckoutGitQueriesForClient,
   isLocalWorktreeArchivePending,
   useCheckoutGitActionsStore,
 } from "@/git/actions-store";
@@ -172,32 +170,6 @@ describe("checkout-git-actions-store", () => {
     ).toBe("idle");
   });
 
-  it("invalidates checkout PR status and every PR pane timeline for a checkout", async () => {
-    const queryClient = new QueryClient();
-
-    queryClient.setQueryData(["checkoutPrStatus", serverId, cwd], { status: { number: 12 } });
-    queryClient.setQueryData(["prPaneTimeline", serverId, cwd, 12], { items: [] });
-    queryClient.setQueryData(["prPaneTimeline", serverId, cwd, 13], { items: [] });
-    queryClient.setQueryData(["prPaneTimeline", serverId, "/tmp/other", 12], { items: [] });
-
-    await invalidateCheckoutGitQueriesForClient(queryClient, { serverId, cwd });
-
-    expect(queryClient.getQueryState(["checkoutPrStatus", serverId, cwd])?.isInvalidated).toBe(
-      true,
-    );
-    expect(queryClient.getQueryState(["prPaneTimeline", serverId, cwd, 12])?.isInvalidated).toBe(
-      true,
-    );
-    expect(queryClient.getQueryState(["prPaneTimeline", serverId, cwd, 13])?.isInvalidated).toBe(
-      true,
-    );
-    expect(
-      queryClient.getQueryState(["prPaneTimeline", serverId, "/tmp/other", 12])?.isInvalidated,
-    ).toBe(false);
-
-    queryClient.clear();
-  });
-
   it("hides an archived worktree optimistically while the archive RPC is in flight", async () => {
     const deferred = createDeferred<Record<string, never>>();
     const client = {

packages/app/src/git/actions-store.ts

@@ -1,4 +1,4 @@
-import type { QueryClient, QueryKey } from "@tanstack/react-query";
+import type { QueryKey } from "@tanstack/react-query";
 import type { CheckoutPrMergeMethod } from "@server/shared/messages";
 import { create } from "zustand";
 import { queryClient as appQueryClient } from "@/query/query-client";
@@ -17,6 +17,7 @@ import {
   resolveWorkspaceIdByExecutionDirectory,
   resolveWorkspaceMapKeyByIdentity,
 } from "@/utils/workspace-execution";
+import { invalidateCheckoutGitQueriesForClient } from "@/git/query-keys";
 
 const SUCCESS_DISPLAY_MS = 1000;
 
@@ -74,44 +75,6 @@ function setStatus(
   });
 }
 
-export async function invalidateCheckoutGitQueriesForClient(
-  queryClient: QueryClient,
-  { serverId, cwd }: { serverId: string; cwd: string },
-) {
-  await Promise.all([
-    queryClient.invalidateQueries({
-      queryKey: ["checkoutStatus", serverId, cwd],
-    }),
-    queryClient.invalidateQueries({
-      predicate: (query) => {
-        const key = query.queryKey;
-        return (
-          Array.isArray(key) && key[0] === "checkoutDiff" && key[1] === serverId && key[2] === cwd
-        );
-      },
-    }),
-    queryClient.invalidateQueries({
-      predicate: (query) => {
-        const key = query.queryKey;
-        return (
-          Array.isArray(key) &&
-          key[0] === "checkoutPrStatus" &&
-          key[1] === serverId &&
-          key[2] === cwd
-        );
-      },
-    }),
-    queryClient.invalidateQueries({
-      predicate: (query) => {
-        const key = query.queryKey;
-        return (
-          Array.isArray(key) && key[0] === "prPaneTimeline" && key[1] === serverId && key[2] === cwd
-        );
-      },
-    }),
-  ]);
-}
-
 function invalidateCheckoutGitQueries(serverId: string, cwd: string) {
   return invalidateCheckoutGitQueriesForClient(appQueryClient, { serverId, cwd });
 }

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

@@ -0,0 +1,62 @@
+import { QueryClient } from "@tanstack/react-query";
+import { describe, expect, it } from "vitest";
+import {
+  checkoutDiffQueryKey,
+  checkoutPrStatusQueryKey,
+  checkoutStatusQueryKey,
+  invalidateCheckoutGitQueriesForClient,
+  prPaneTimelineQueryKey,
+} from "@/git/query-keys";
+
+describe("checkout query keys", () => {
+  const serverId = "server-1";
+  const cwd = "/tmp/repo";
+
+  it("invalidates every query for a checkout without touching other checkouts", async () => {
+    const queryClient = new QueryClient();
+
+    queryClient.setQueryData(checkoutStatusQueryKey(serverId, cwd), { isGit: true });
+    queryClient.setQueryData(checkoutDiffQueryKey(serverId, cwd, "base", "main", true), {
+      files: [],
+    });
+    queryClient.setQueryData(checkoutPrStatusQueryKey(serverId, cwd), { status: { number: 12 } });
+    queryClient.setQueryData(prPaneTimelineQueryKey({ serverId, cwd, prNumber: 12 }), {
+      items: [],
+    });
+    queryClient.setQueryData(prPaneTimelineQueryKey({ serverId, cwd, prNumber: 13 }), {
+      items: [],
+    });
+    queryClient.setQueryData(
+      prPaneTimelineQueryKey({ serverId, cwd: "/tmp/other", prNumber: 12 }),
+      { items: [] },
+    );
+
+    await invalidateCheckoutGitQueriesForClient(queryClient, { serverId, cwd });
+
+    expect(queryClient.getQueryState(checkoutStatusQueryKey(serverId, cwd))?.isInvalidated).toBe(
+      true,
+    );
+    expect(
+      queryClient.getQueryState(checkoutDiffQueryKey(serverId, cwd, "base", "main", true))
+        ?.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(prPaneTimelineQueryKey({ serverId, cwd, prNumber: 13 }))
+        ?.isInvalidated,
+    ).toBe(true);
+    expect(
+      queryClient.getQueryState(
+        prPaneTimelineQueryKey({ serverId, cwd: "/tmp/other", prNumber: 12 }),
+      )?.isInvalidated,
+    ).toBe(false);
+
+    queryClient.clear();
+  });
+});

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

@@ -0,0 +1,82 @@
+import type { Query, QueryClient } from "@tanstack/react-query";
+
+interface CheckoutQueryIdentity {
+  serverId: string;
+  cwd: string;
+}
+
+type CheckoutQueryKey = readonly unknown[];
+
+export function checkoutStatusQueryKey(serverId: string, cwd: string) {
+  return ["checkoutStatus", serverId, cwd] as const;
+}
+
+export function checkoutDiffQueryKey(
+  serverId: string,
+  cwd: string,
+  mode: "uncommitted" | "base",
+  baseRef?: string,
+  ignoreWhitespace?: boolean,
+) {
+  return ["checkoutDiff", serverId, cwd, mode, baseRef ?? "", ignoreWhitespace === true] as const;
+}
+
+export function checkoutPrStatusQueryKey(serverId: string, cwd: string) {
+  return ["checkoutPrStatus", serverId, cwd] as const;
+}
+
+export function prPaneTimelineQueryKey({
+  serverId,
+  cwd,
+  prNumber,
+}: {
+  serverId: string;
+  cwd: string;
+  prNumber: number | null;
+}) {
+  return ["prPaneTimeline", serverId, cwd, prNumber] as const;
+}
+
+export async function invalidateCheckoutGitQueriesForClient(
+  queryClient: QueryClient,
+  identity: CheckoutQueryIdentity,
+) {
+  await Promise.all([
+    queryClient.invalidateQueries({
+      queryKey: checkoutStatusQueryKey(identity.serverId, identity.cwd),
+    }),
+    queryClient.invalidateQueries({
+      predicate: checkoutQueryPredicate("checkoutDiff", identity),
+    }),
+    queryClient.invalidateQueries({
+      predicate: checkoutQueryPredicate("checkoutPrStatus", identity),
+    }),
+    queryClient.invalidateQueries({
+      predicate: checkoutQueryPredicate("prPaneTimeline", identity),
+    }),
+  ]);
+}
+
+function checkoutQueryPredicate(
+  queryKind: CheckoutQueryKey[0],
+  identity: CheckoutQueryIdentity,
+): (query: Query) => boolean {
+  return (query) => {
+    const key = query.queryKey;
+    return (
+      isCheckoutQueryKey(key) &&
+      key[0] === queryKind &&
+      key[1] === identity.serverId &&
+      key[2] === identity.cwd
+    );
+  };
+}
+
+function isCheckoutQueryKey(key: readonly unknown[]): key is CheckoutQueryKey {
+  return (
+    key.length >= 3 &&
+    typeof key[0] === "string" &&
+    typeof key[1] === "string" &&
+    typeof key[2] === "string"
+  );
+}

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

@@ -3,16 +3,7 @@ import { useEffect, useId, useMemo } from "react";
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
 import type { SubscribeCheckoutDiffResponse } from "@server/shared/messages";
 import { orderCheckoutDiffFiles } from "@/git/diff-order";
-
+import { checkoutDiffQueryKey } from "@/git/query-keys";
-function checkoutDiffQueryKey(
-  serverId: string,
-  cwd: string,
-  mode: "uncommitted" | "base",
-  baseRef?: string,
-  ignoreWhitespace?: boolean,
-) {
-  return ["checkoutDiff", serverId, cwd, mode, baseRef ?? "", ignoreWhitespace === true] as const;
-}
 
 interface UseCheckoutDiffQueryOptions {
   serverId: string;

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

@@ -2,10 +2,7 @@ import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { useEffect } from "react";
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
 import type { CheckoutPrStatusResponse } from "@server/shared/messages";
-
+import { checkoutPrStatusQueryKey } from "@/git/query-keys";
-export function checkoutPrStatusQueryKey(serverId: string, cwd: string) {
-  return ["checkoutPrStatus", serverId, cwd] as const;
-}
 
 interface UseCheckoutPrStatusQueryOptions {
   serverId: string;

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

@@ -9,7 +9,8 @@ import {
 } from "@tanstack/react-query";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { CheckoutStatusResponse } from "@server/shared/messages";
-import { checkoutStatusQueryKey, useCheckoutStatusQuery } from "./use-status-query";
+import { checkoutStatusQueryKey } from "@/git/query-keys";
+import { useCheckoutStatusQuery } from "./use-status-query";
 
 type CheckoutStatusPayload = CheckoutStatusResponse["payload"];
 

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

@@ -2,13 +2,10 @@ import { type QueryClient, useQuery, useQueryClient } from "@tanstack/react-quer
 import { useEffect } from "react";
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
 import type { CheckoutStatusResponse } from "@server/shared/messages";
+import { checkoutStatusQueryKey } from "@/git/query-keys";
 
 export const CHECKOUT_STATUS_STALE_TIME = 15_000;
 
-export function checkoutStatusQueryKey(serverId: string, cwd: string) {
-  return ["checkoutStatus", serverId, cwd] as const;
-}
-
 interface UseCheckoutStatusQueryOptions {
   serverId: string;
   cwd: string;

packages/app/src/hooks/use-branch-switcher.ts

@@ -3,7 +3,7 @@ import { useQuery, type QueryClient } from "@tanstack/react-query";
 import type { DaemonClient } from "@server/client/daemon-client";
 import type { ComboboxOption } from "@/components/ui/combobox";
 import type { ToastApi } from "@/components/toast-host";
-import { invalidateCheckoutGitQueriesForClient } from "@/git/actions-store";
+import { invalidateCheckoutGitQueriesForClient } from "@/git/query-keys";
 import { confirmDialog } from "@/utils/confirm-dialog";
 
 interface UseBranchSwitcherInput {

packages/app/src/hooks/use-pr-pane-data.test.ts

@@ -12,11 +12,8 @@ import type {
   CheckoutPrStatusResponse,
   PullRequestTimelineResponse,
 } from "@server/shared/messages";
-import {
+import { checkoutPrStatusQueryKey, prPaneTimelineQueryKey } from "@/git/query-keys";
-  prPaneTimelineQueryKey,
+import { usePrPaneData, type UsePrPaneDataResult } from "./use-pr-pane-data";
-  usePrPaneData,
-  type UsePrPaneDataResult,
-} from "./use-pr-pane-data";
 import { useWorkspacePrHint } from "@/git/use-pr-status-query";
 
 type CheckoutPrStatus = NonNullable<CheckoutPrStatusResponse["payload"]["status"]>;
@@ -335,7 +332,7 @@ describe("usePrPaneData", () => {
     ).toBeUndefined();
 
     mockClient.checkoutPrStatus.mockResolvedValue(statusPayload());
-    hook.queryClient.invalidateQueries({ queryKey: ["checkoutPrStatus", serverId, cwd] });
+    hook.queryClient.invalidateQueries({ queryKey: checkoutPrStatusQueryKey(serverId, cwd) });
 
     await waitForExpectation(() => {
       expect(mockClient.pullRequestTimeline).toHaveBeenCalledWith({
@@ -434,7 +431,7 @@ describe("usePrPaneData", () => {
     await waitForExpectation(() => {
       expect(hook.latest.data?.checks[0]?.status).toBe("success");
     });
-    expect(queryClient.getQueryData(["checkoutPrStatus", serverId, cwd])).toEqual(
+    expect(queryClient.getQueryData(checkoutPrStatusQueryKey(serverId, cwd))).toEqual(
       statusPayload({
         requestId: "server-push",
         status: status({
@@ -474,8 +471,10 @@ describe("usePrPaneData", () => {
       }),
     );
 
-    expect(queryClient.getQueryData(["checkoutPrStatus", serverId, cwd])).toEqual(initial);
+    expect(queryClient.getQueryData(checkoutPrStatusQueryKey(serverId, cwd))).toEqual(initial);
-    expect(queryClient.getQueryData(["checkoutPrStatus", serverId, "/other-repo"])).toBeUndefined();
+    expect(
+      queryClient.getQueryData(checkoutPrStatusQueryKey(serverId, "/other-repo")),
+    ).toBeUndefined();
   });
 
   it("passes repoOwner and repoName to the timeline request when present", async () => {
@@ -719,7 +718,7 @@ describe("usePrPaneData", () => {
 
     const refreshDeferred = createDeferred<CheckoutPrStatusPayload>();
     mockClient.checkoutPrStatus.mockReturnValue(refreshDeferred.promise);
-    hook.queryClient.invalidateQueries({ queryKey: ["checkoutPrStatus", serverId, cwd] });
+    hook.queryClient.invalidateQueries({ queryKey: checkoutPrStatusQueryKey(serverId, cwd) });
 
     await waitForExpectation(() => {
       expect(hook.latest.isLoading).toBe(false);

packages/app/src/hooks/use-pr-pane-data.ts

@@ -6,6 +6,7 @@ import type {
 } from "@server/shared/messages";
 import { mapPrPaneData, type PrPaneData } from "@/git/pr-pane-data";
 import { useCheckoutPrStatusQuery } from "@/git/use-pr-status-query";
+import { prPaneTimelineQueryKey } from "@/git/query-keys";
 
 type CheckoutPrStatusPayloadError = CheckoutPrStatusResponse["payload"]["error"];
 type PullRequestTimeline = PullRequestTimelineResponse["payload"];
@@ -152,18 +153,6 @@ export function usePrPaneData({
   };
 }
 
-export function prPaneTimelineQueryKey({
-  serverId,
-  cwd,
-  prNumber,
-}: {
-  serverId: string;
-  cwd: string;
-  prNumber: number | null;
-}) {
-  return ["prPaneTimeline", serverId, cwd, prNumber] as const;
-}
-
 function firstNonSuppressedError({
   statusPayloadError,
   statusError,

packages/app/src/screens/workspace/workspace-screen.tsx

@@ -92,7 +92,8 @@ import { useProvidersSnapshot } from "@/hooks/use-providers-snapshot";
 import { shouldShowWorkspaceSetup, useWorkspaceSetupStore } from "@/stores/workspace-setup-store";
 import { useWorkspace } from "@/stores/session-store-hooks";
 import { useWorkspaceTerminalSessionRetention } from "@/terminal/hooks/use-workspace-terminal-session-retention";
-import { checkoutStatusQueryKey, type CheckoutStatusPayload } from "@/git/use-status-query";
+import type { CheckoutStatusPayload } from "@/git/use-status-query";
+import { checkoutStatusQueryKey } from "@/git/query-keys";
 import type { ListTerminalsResponse } from "@server/shared/messages";
 import { upsertTerminalListEntry } from "@/utils/terminal-list";
 import { confirmDialog } from "@/utils/confirm-dialog";