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

// @vitest-environment jsdom
import { QueryClient } from "@tanstack/react-query";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { CheckoutStatusUpdate } from "@getpaseo/protocol/messages";
import {
  checkoutCommitsQueryKey,
  checkoutPrStatusQueryKey,
  checkoutStatusQueryKey,
} from "@/git/query-keys";
import {
  prPanePipelineQueryKey,
  prPaneTimelineQueryKey,
} from "@/git/pull-request-panel/query-keys";
import {
  resetWorkingDiffComparisons,
  resolveWorkingDiffComparison,
  selectWorkingDiffComparison,
} from "@/git/working-diff-comparison";
import { draftAgentCommandsQueryKey } from "@/hooks/agent-commands-query";
import {
  applyCheckoutStatusUpdateFromEvent,
  ensureCheckoutStatus,
  type CheckoutPrStatusPayload,
  type CheckoutStatusPayload,
  fetchCheckoutStatus,
} from "./checkout-status-cache";

vi.mock("@react-native-async-storage/async-storage", () => ({
  default: {
    getItem: vi.fn(async () => null),
    setItem: vi.fn(async () => undefined),
    removeItem: vi.fn(async () => undefined),
  },
}));

const serverId = "server-1";
const cwd = "/repo";

function checkoutStatus(overrides: Partial<CheckoutStatusPayload> = {}): CheckoutStatusPayload {
  return {
    cwd,
    error: null,
    requestId: "checkout-status-1",
    isGit: true,
    isPaseoOwnedWorktree: false,
    repoRoot: cwd,
    currentBranch: "main",
    isDirty: false,
    baseRef: "origin/main",
    aheadBehind: { ahead: 0, behind: 0 },
    aheadOfOrigin: 0,
    behindOfOrigin: 0,
    hasRemote: true,
    remoteUrl: "git@github.com:getpaseo/paseo.git",
    ...overrides,
  } as CheckoutStatusPayload;
}

function prStatus(overrides: Partial<CheckoutPrStatusPayload> = {}): CheckoutPrStatusPayload {
  return {
    cwd,
    status: {
      forge: "github",
      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,
    authState: "authenticated",
    forge: "github",
    error: null,
    requestId: "pr-status-1",
    ...overrides,
  };
}

function checkoutStatusUpdate(
  payload: CheckoutStatusPayload,
  extraPrStatus?: NonNullable<CheckoutStatusUpdate["payload"]["prStatus"]>,
): CheckoutStatusUpdate {
  return {
    type: "checkout_status_update",
    payload: extraPrStatus ? { ...payload, prStatus: extraPrStatus } : payload,
  };
}

function selectBaseComparison(): void {
  selectWorkingDiffComparison({
    serverId,
    cwd,
    comparison: "base",
  });
}

function createQueryClient(): QueryClient {
  return new QueryClient({ defaultOptions: { queries: { retry: false } } });
}

beforeEach(() => {
  resetWorkingDiffComparisons();
});

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 fetchCheckoutStatus({ client, serverId, cwd });

    expect(result).toEqual(fetched);
    expect(client.getCheckoutStatus).toHaveBeenCalledExactlyOnceWith(cwd);
  });

  it("keeps a manual working-diff comparison when the fetched dirty state flips", async () => {
    selectBaseComparison();
    const client = { getCheckoutStatus: vi.fn(async () => checkoutStatus({ isDirty: false })) };

    await fetchCheckoutStatus({ client, serverId, cwd });

    expect(resolveWorkingDiffComparison({ serverId, cwd })).toBe("base");
  });
});

describe("ensureCheckoutStatus", () => {
  it("awaits the canonical checkout-status query and reuses its cached result", async () => {
    const queryClient = createQueryClient();
    const fetched = checkoutStatus({ currentBranch: "feature/current" });
    const client = { getCheckoutStatus: vi.fn(async () => fetched) };

    const first = await ensureCheckoutStatus({ queryClient, client, serverId, cwd });
    const second = await ensureCheckoutStatus({ queryClient, client, serverId, cwd });

    expect(first).toEqual(fetched);
    expect(second).toEqual(fetched);
    expect(client.getCheckoutStatus).toHaveBeenCalledExactlyOnceWith(cwd);
  });

  it("awaits a refetch when the canonical cached status was invalidated", async () => {
    const queryClient = createQueryClient();
    queryClient.setQueryData(
      checkoutStatusQueryKey(serverId, cwd),
      checkoutStatus({ currentBranch: "feature/stale" }),
    );
    await queryClient.invalidateQueries({
      queryKey: checkoutStatusQueryKey(serverId, cwd),
      refetchType: "none",
    });
    const fetched = checkoutStatus({ currentBranch: "feature/current" });
    const client = { getCheckoutStatus: vi.fn(async () => fetched) };

    const result = await ensureCheckoutStatus({ queryClient, client, serverId, cwd });

    expect(result.currentBranch).toBe("feature/current");
    expect(client.getCheckoutStatus).toHaveBeenCalledExactlyOnceWith(cwd);
  });
});

describe("applyCheckoutStatusUpdateFromEvent", () => {
  it("writes the checkout status to the cache using the cwd from the payload", () => {
    const queryClient = createQueryClient();
    const pushed = checkoutStatus({ requestId: "push-1", isDirty: true });

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(pushed),
    });

    expect(queryClient.getQueryData(checkoutStatusQueryKey(serverId, cwd))).toEqual(pushed);
  });

  it("invalidates recent commits when checkout status is pushed", () => {
    const queryClient = createQueryClient();
    queryClient.setQueryData(checkoutCommitsQueryKey(serverId, cwd), { commits: [] });
    queryClient.setQueryData(checkoutCommitsQueryKey(serverId, "/repo2"), { commits: [] });

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus()),
    });

    expect(queryClient.getQueryState(checkoutCommitsQueryKey(serverId, cwd))?.isInvalidated).toBe(
      true,
    );
    expect(
      queryClient.getQueryState(checkoutCommitsQueryKey(serverId, "/repo2"))?.isInvalidated,
    ).toBe(false);
  });

  it("drops the checkout's cached draft slash commands when its branch changes", () => {
    const queryClient = createQueryClient();
    queryClient.setQueryData(
      checkoutStatusQueryKey(serverId, cwd),
      checkoutStatus({ currentBranch: "chore/build-paseo" }),
    );
    const thisCheckout = draftAgentCommandsQueryKey({
      serverId,
      draftConfig: { provider: "claude", cwd, model: "haiku" },
    });
    const otherCheckout = draftAgentCommandsQueryKey({
      serverId,
      draftConfig: { provider: "claude", cwd: "/repo2", model: "haiku" },
    });
    queryClient.setQueryData(thisCheckout, [{ name: "build-paseo" }]);
    queryClient.setQueryData(otherCheckout, [{ name: "build-paseo" }]);

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus({ currentBranch: "main" })),
    });

    expect(queryClient.getQueryData(thisCheckout)).toBeUndefined();
    expect(queryClient.getQueryData(otherCheckout)).toEqual([{ name: "build-paseo" }]);
  });

  it("keeps the checkout's draft slash commands when a push leaves its branch unchanged", () => {
    const queryClient = createQueryClient();
    queryClient.setQueryData(checkoutStatusQueryKey(serverId, cwd), checkoutStatus());
    const thisCheckout = draftAgentCommandsQueryKey({
      serverId,
      draftConfig: { provider: "claude", cwd, model: "haiku" },
    });
    queryClient.setQueryData(thisCheckout, [{ name: "build-paseo" }]);

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus({ isDirty: true })),
    });

    expect(queryClient.getQueryData(thisCheckout)).toEqual([{ name: "build-paseo" }]);
  });

  it("writes the PR status cache when prStatus is present, and skips it otherwise", () => {
    const queryClient = createQueryClient();
    const pushedPr = prStatus({ requestId: "pr-1" });

    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("normalizes legacy PR auth state at the pushed-cache boundary", () => {
    const queryClient = createQueryClient();
    const { authState: _authState, ...legacyPrStatus } = prStatus({
      githubFeaturesEnabled: false,
    });

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus(), legacyPrStatus),
    });

    expect(
      queryClient.getQueryData<CheckoutPrStatusPayload>(checkoutPrStatusQueryKey(serverId, cwd))
        ?.authState,
    ).toBe("unauthenticated");
    expect(
      queryClient.getQueryData<CheckoutStatusUpdate["payload"]>(
        checkoutStatusQueryKey(serverId, cwd),
      )?.prStatus?.authState,
    ).toBe("unauthenticated");
  });

  it.each(["base", "uncommitted"] as const)(
    "keeps %s selected across dirty and clean pushes",
    (comparison) => {
      const queryClient = createQueryClient();
      selectWorkingDiffComparison({ serverId, cwd, comparison });
      applyCheckoutStatusUpdateFromEvent({
        queryClient,
        serverId,
        message: checkoutStatusUpdate(checkoutStatus({ isDirty: true })),
      });
      expect(resolveWorkingDiffComparison({ serverId, cwd })).toBe(comparison);
      applyCheckoutStatusUpdateFromEvent({
        queryClient,
        serverId,
        message: checkoutStatusUpdate(checkoutStatus({ isDirty: false })),
      });
      expect(resolveWorkingDiffComparison({ serverId, cwd })).toBe(comparison);
    },
  );

  it("invalidates PR detail queries 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 });
    const pipelineKey = prPanePipelineQueryKey({
      serverId,
      cwd,
      pipelineId: 9001,
      changeRequestNumber: 1,
    });
    queryClient.setQueryData(timelineKey, { items: [] });
    queryClient.setQueryData(pipelineKey, { stages: [] });

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus(), prStatus({ requestId: "pr-v2" })),
    });
    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(false);
    expect(queryClient.getQueryState(pipelineKey)?.isInvalidated).toBe(false);

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(
        checkoutStatus(),
        prStatus({
          requestId: "pr-v3",
          status: { ...prStatus().status!, state: "closed" },
        }),
      ),
    });
    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(true);
    expect(queryClient.getQueryState(pipelineKey)?.isInvalidated).toBe(true);
  });

  it("invalidates PR detail queries 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 });
    const pipelineKey = prPanePipelineQueryKey({
      serverId,
      cwd,
      pipelineId: 9001,
      changeRequestNumber: 1,
    });
    const otherPipelineKey = prPanePipelineQueryKey({
      serverId,
      cwd: "/repo2",
      pipelineId: 9001,
      changeRequestNumber: 1,
    });
    queryClient.setQueryData(timelineKey, { items: [] });
    queryClient.setQueryData(otherTimelineKey, { items: [] });
    queryClient.setQueryData(pipelineKey, { stages: [] });
    queryClient.setQueryData(otherPipelineKey, { stages: [] });

    applyCheckoutStatusUpdateFromEvent({
      queryClient,
      serverId,
      message: checkoutStatusUpdate(checkoutStatus(), prStatus()),
    });

    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(true);
    expect(queryClient.getQueryState(otherTimelineKey)?.isInvalidated).toBe(false);
    expect(queryClient.getQueryState(pipelineKey)?.isInvalidated).toBe(true);
    expect(queryClient.getQueryState(otherPipelineKey)?.isInvalidated).toBe(false);
  });
});