Keep git review state fresh after reconnect (#1491)

packages/app/src/review/index.ts

@@ -2,15 +2,17 @@ export {
   buildReviewAttachmentSnapshot,
   buildReviewDraftKey,
   buildReviewDraftScopeKey,
+  expireStaleDiffModeOverrides,
   getReviewDraftComments,
   resetReviewDraftStore,
-  useActiveReviewDraftMode,
   useClearReviewDraft,
   useReviewAttachmentSnapshot,
-  useSetActiveReviewDraftMode,
+  useResolvedDiffMode,
+  useSetDiffModeOverride,
   addReviewDraftComment,
   type BuildReviewDraftKeyInput,
   type BuildReviewDraftScopeKeyInput,
+  type DiffModeOverride,
   type ReviewDraftCommentInput,
   type ReviewDraftComment,
   type ReviewDraftMode,

packages/app/src/review/state.ts

@@ -11,32 +11,80 @@ export interface ReviewDraftComment {
   updatedAt: string;
 }
 
+// A manual mode selection is valid only while the checkout's dirty state matches the
+// value at the time of selection. serverId/cwd identify the checkout so the override can
+// be expired when its dirty state changes (see expireStaleDiffModeOverridesInState).
+export interface DiffModeOverride {
+  serverId: string;
+  cwd: string;
+  mode: ReviewDraftMode;
+  isDirtyAtSelection: boolean;
+}
+
 export interface ReviewDraftStoreState {
   drafts: Record<string, ReviewDraftComment[]>;
-  activeModesByScope: Record<string, ReviewDraftMode>;
+  // In-memory only — not persisted. Keyed by scope key.
+  diffModeOverrides: Record<string, DiffModeOverride>;
 }
 
+// Only drafts are persisted; diffModeOverrides is intentionally excluded.
 export interface SerializedReviewDraftState {
   drafts: Record<string, ReviewDraftComment[]>;
-  activeModesByScope: Record<string, ReviewDraftMode>;
 }
 
-export function setActiveModeInState(
+export function setDiffModeOverrideInState(
   state: ReviewDraftStoreState,
-  input: { scopeKey: string; mode: ReviewDraftMode },
+  input: { scopeKey: string; override: DiffModeOverride },
 ): ReviewDraftStoreState {
-  if (state.activeModesByScope[input.scopeKey] === input.mode) {
-    return state;
-  }
   return {
     ...state,
-    activeModesByScope: {
-      ...state.activeModesByScope,
-      [input.scopeKey]: input.mode,
+    diffModeOverrides: {
+      ...state.diffModeOverrides,
+      [input.scopeKey]: input.override,
     },
   };
 }
 
+// Drops every override for the checkout whose dirty state no longer matches the value it
+// was selected under. Called whenever a checkout status enters the app (push or fetch),
+// so expiry does not depend on any screen being mounted.
+export function expireStaleDiffModeOverridesInState(
+  state: ReviewDraftStoreState,
+  input: { serverId: string; cwd: string; isDirty: boolean },
+): ReviewDraftStoreState {
+  const staleScopeKeys = Object.entries(state.diffModeOverrides)
+    .filter(
+      ([, override]) =>
+        override.serverId === input.serverId &&
+        override.cwd === input.cwd &&
+        override.isDirtyAtSelection !== input.isDirty,
+    )
+    .map(([scopeKey]) => scopeKey);
+  if (staleScopeKeys.length === 0) {
+    return state;
+  }
+  const next = { ...state.diffModeOverrides };
+  for (const scopeKey of staleScopeKeys) {
+    delete next[scopeKey];
+  }
+  return { ...state, diffModeOverrides: next };
+}
+
+// Pure read — returns the effective mode without mutating state. The staleness check is
+// kept even though stale overrides are expired at the data boundary: a render can observe
+// a fresh dirty state before the expiry lands, and resolution must be correct under any
+// interleaving.
+export function resolveDiffMode(input: {
+  override: DiffModeOverride | undefined;
+  hasUncommittedChanges: boolean;
+}): ReviewDraftMode {
+  const { override, hasUncommittedChanges } = input;
+  if (override && override.isDirtyAtSelection === hasUncommittedChanges) {
+    return override.mode;
+  }
+  return hasUncommittedChanges ? "uncommitted" : "base";
+}
+
 export function addCommentToState(
   state: ReviewDraftStoreState,
   input: { key: string; comment: ReviewDraftComment },
@@ -108,18 +156,18 @@ export function serializeReviewDraftState(
 ): SerializedReviewDraftState {
   return {
     drafts: state.drafts,
-    activeModesByScope: state.activeModesByScope,
   };
 }
 
 export function normalizePersistedState(state: unknown): ReviewDraftStoreState {
   if (!state || typeof state !== "object") {
-    return { drafts: {}, activeModesByScope: {} };
+    return { drafts: {}, diffModeOverrides: {} };
   }
-  const persisted = state as { drafts?: unknown; activeModesByScope?: unknown };
+  // activeModesByScope may be present in old persisted JSON — tolerate and ignore it.
+  const persisted = state as { drafts?: unknown };
   const drafts = persisted.drafts;
   if (!drafts || typeof drafts !== "object" || Array.isArray(drafts)) {
-    return { drafts: {}, activeModesByScope: {} };
+    return { drafts: {}, diffModeOverrides: {} };
   }
 
   const normalized: Record<string, ReviewDraftComment[]> = {};
@@ -132,20 +180,7 @@ export function normalizePersistedState(state: unknown): ReviewDraftStoreState {
     );
   }
 
-  const activeModesByScope: Record<string, ReviewDraftMode> = {};
-  const persistedActiveModes = persisted.activeModesByScope;
-  if (
-    persistedActiveModes &&
-    typeof persistedActiveModes === "object" &&
-    !Array.isArray(persistedActiveModes)
-  ) {
-    for (const [key, mode] of Object.entries(persistedActiveModes)) {
-      if (mode === "base" || mode === "uncommitted") {
-        activeModesByScope[key] = mode;
-      }
-    }
-  }
-  return { drafts: normalized, activeModesByScope };
+  return { drafts: normalized, diffModeOverrides: {} };
 }
 
 export function isReviewDraftComment(value: unknown): value is ReviewDraftComment {

packages/app/src/review/store.test.ts

@@ -9,15 +9,25 @@ import {
   addCommentToState,
   clearReviewInState,
   deleteCommentFromState,
+  type DiffModeOverride,
+  expireStaleDiffModeOverridesInState,
   normalizePersistedState,
+  resolveDiffMode,
   type ReviewDraftComment,
   type ReviewDraftStoreState,
-  setActiveModeInState,
+  serializeReviewDraftState,
+  setDiffModeOverrideInState,
   updateCommentInState,
 } from "./state";
 
 function emptyState(): ReviewDraftStoreState {
-  return { drafts: {}, activeModesByScope: {} };
+  return { drafts: {}, diffModeOverrides: {} };
+}
+
+function makeOverride(
+  input: Pick<DiffModeOverride, "mode" | "isDirtyAtSelection">,
+): DiffModeOverride {
+  return { serverId: "server-1", cwd: "/repo", ...input };
 }
 
 function makeComment(overrides: Partial<ReviewDraftComment> = {}): ReviewDraftComment {
@@ -95,7 +105,7 @@ describe("buildReviewDraftKey", () => {
     ).toBe("review:server=local:cwd=%2Frepo:mode=base:base=main:ignoreWhitespace=false");
   });
 
-  it("builds a mode-free scope key for active review mode sharing", () => {
+  it("builds a mode-free scope key for diff mode override sharing", () => {
     const scope = buildReviewDraftScopeKey({
       serverId: "local",
       workspaceId: "workspace-1",
@@ -112,7 +122,7 @@ describe("buildReviewDraftKey", () => {
 });
 
 describe("normalizePersistedState", () => {
-  it("filters invalid draft comments and drops unknown active modes", () => {
+  it("filters invalid draft comments and ignores activeModesByScope from old persisted JSON", () => {
     const normalized = normalizePersistedState({
       drafts: {
         "review:key": [
@@ -128,17 +138,14 @@ describe("normalizePersistedState", () => {
           { id: "bad", filePath: "src/example.ts" },
         ],
       },
+      // Old persisted field — must be tolerated and ignored, not migrated.
       activeModesByScope: {
         "review:scope:base": "base",
         "review:scope:dirty": "uncommitted",
-        "review:scope:bad": "other",
       },
     });
 
-    expect(normalized.activeModesByScope).toEqual({
-      "review:scope:base": "base",
-      "review:scope:dirty": "uncommitted",
-    });
+    expect(normalized.diffModeOverrides).toEqual({});
     expect(normalized.drafts["review:key"]).toEqual([
       {
         id: "comment-1",
@@ -153,15 +160,144 @@ describe("normalizePersistedState", () => {
   });
 
   it("returns empty state for null, non-object, or malformed inputs", () => {
-    expect(normalizePersistedState(null)).toEqual({ drafts: {}, activeModesByScope: {} });
-    expect(normalizePersistedState("nope")).toEqual({ drafts: {}, activeModesByScope: {} });
+    expect(normalizePersistedState(null)).toEqual({ drafts: {}, diffModeOverrides: {} });
+    expect(normalizePersistedState("nope")).toEqual({ drafts: {}, diffModeOverrides: {} });
     expect(normalizePersistedState({ drafts: [] })).toEqual({
       drafts: {},
-      activeModesByScope: {},
+      diffModeOverrides: {},
     });
   });
 });
 
+describe("serializeReviewDraftState", () => {
+  it("serialized output does not contain activeModesByScope or diffModeOverrides", () => {
+    const state = setDiffModeOverrideInState(
+      addCommentToState(emptyState(), { key: "review:key", comment: makeComment() }),
+      {
+        scopeKey: "review:scope",
+        override: makeOverride({ mode: "uncommitted", isDirtyAtSelection: true }),
+      },
+    );
+
+    const serialized = serializeReviewDraftState(state);
+
+    expect(Object.keys(serialized)).toEqual(["drafts"]);
+    expect("activeModesByScope" in serialized).toBe(false);
+    expect("diffModeOverrides" in serialized).toBe(false);
+    expect(serialized.drafts["review:key"]).toHaveLength(1);
+  });
+});
+
+describe("diff mode override", () => {
+  it("resolves to auto mode from the dirty state when no override exists", () => {
+    expect(resolveDiffMode({ override: undefined, hasUncommittedChanges: true })).toBe(
+      "uncommitted",
+    );
+    expect(resolveDiffMode({ override: undefined, hasUncommittedChanges: false })).toBe("base");
+  });
+
+  it("honors the override while isDirty matches the value at selection, across remounts", () => {
+    const state = setDiffModeOverrideInState(emptyState(), {
+      scopeKey: "review:scope",
+      override: makeOverride({ mode: "base", isDirtyAtSelection: true }),
+    });
+
+    // Resolution is a plain store read, so it gives the same answer on every (re)mount.
+    const override = state.diffModeOverrides["review:scope"];
+    expect(resolveDiffMode({ override, hasUncommittedChanges: true })).toBe("base");
+  });
+
+  it("masks a stale override before its expiry lands", () => {
+    const state = setDiffModeOverrideInState(emptyState(), {
+      scopeKey: "review:scope",
+      override: makeOverride({ mode: "uncommitted", isDirtyAtSelection: true }),
+    });
+
+    const override = state.diffModeOverrides["review:scope"];
+    expect(resolveDiffMode({ override, hasUncommittedChanges: false })).toBe("base");
+  });
+
+  it("expires overrides for the checkout whose dirty state flipped, keeping the rest", () => {
+    let state = setDiffModeOverrideInState(emptyState(), {
+      scopeKey: "review:scope-a",
+      override: makeOverride({ mode: "base", isDirtyAtSelection: true }),
+    });
+    state = setDiffModeOverrideInState(state, {
+      scopeKey: "review:scope-b",
+      override: makeOverride({ mode: "uncommitted", isDirtyAtSelection: false }),
+    });
+    state = setDiffModeOverrideInState(state, {
+      scopeKey: "review:scope-other",
+      override: {
+        serverId: "server-2",
+        cwd: "/other",
+        mode: "base",
+        isDirtyAtSelection: true,
+      },
+    });
+
+    const next = expireStaleDiffModeOverridesInState(state, {
+      serverId: "server-1",
+      cwd: "/repo",
+      isDirty: false,
+    });
+
+    expect(next.diffModeOverrides["review:scope-a"]).toBeUndefined();
+    expect(next.diffModeOverrides["review:scope-b"]).toBeDefined();
+    expect(next.diffModeOverrides["review:scope-other"]).toBeDefined();
+  });
+
+  it("does not resurrect an expired override when the dirty state flips back", () => {
+    let state = setDiffModeOverrideInState(emptyState(), {
+      scopeKey: "review:scope",
+      override: makeOverride({ mode: "base", isDirtyAtSelection: true }),
+    });
+
+    // The checkout goes clean (e.g. agent commits): the override expires at the boundary.
+    state = expireStaleDiffModeOverridesInState(state, {
+      serverId: "server-1",
+      cwd: "/repo",
+      isDirty: false,
+    });
+    // The checkout goes dirty again: nothing to expire, and nothing comes back.
+    state = expireStaleDiffModeOverridesInState(state, {
+      serverId: "server-1",
+      cwd: "/repo",
+      isDirty: true,
+    });
+
+    expect(state.diffModeOverrides["review:scope"]).toBeUndefined();
+    expect(
+      resolveDiffMode({
+        override: state.diffModeOverrides["review:scope"],
+        hasUncommittedChanges: true,
+      }),
+    ).toBe("uncommitted");
+  });
+
+  it("keeps state identity when no override is stale", () => {
+    const state = setDiffModeOverrideInState(emptyState(), {
+      scopeKey: "review:scope",
+      override: makeOverride({ mode: "base", isDirtyAtSelection: true }),
+    });
+
+    expect(
+      expireStaleDiffModeOverridesInState(state, {
+        serverId: "server-1",
+        cwd: "/repo",
+        isDirty: true,
+      }),
+    ).toBe(state);
+    expect(
+      expireStaleDiffModeOverridesInState(emptyState(), {
+        serverId: "server-1",
+        cwd: "/repo",
+        isDirty: false,
+      }),
+    ).toEqual(emptyState());
+  });
+});
+
 describe("review draft reducers", () => {
   it("adds, updates, and deletes draft comments by key", () => {
     let state = emptyState();
@@ -203,18 +339,6 @@ describe("review draft reducers", () => {
     expect(deleteCommentFromState(state, { key: "review:key", id: "missing" })).toBe(state);
     expect(clearReviewInState(state, { key: "other-key" })).toBe(state);
   });
-
-  it("keeps state identity when setActiveMode does not change the mode", () => {
-    const state = setActiveModeInState(emptyState(), {
-      scopeKey: "review:scope",
-      mode: "base",
-    });
-
-    expect(setActiveModeInState(state, { scopeKey: "review:scope", mode: "base" })).toBe(state);
-    expect(setActiveModeInState(state, { scopeKey: "review:scope", mode: "uncommitted" })).not.toBe(
-      state,
-    );
-  });
 });
 
 describe("buildReviewAttachmentSnapshot", () => {

packages/app/src/review/store.ts

@@ -8,22 +8,31 @@ import {
   addCommentToState,
   clearReviewInState,
   deleteCommentFromState,
+  type DiffModeOverride,
+  expireStaleDiffModeOverridesInState,
   normalizePersistedState,
+  resolveDiffMode,
   type ReviewDraftComment,
   type ReviewDraftMode,
   type ReviewDraftSide,
   type ReviewDraftStoreState,
   serializeReviewDraftState,
-  setActiveModeInState,
+  setDiffModeOverrideInState,
   updateCommentInState,
 } from "@/review/state";
 import { generateMessageId } from "@/types/stream";
 import { buildNumberedDiffHunks, type NumberedDiffLine } from "@/utils/diff-layout";
 import type { AgentAttachment } from "@getpaseo/protocol/messages";
 
-export type { ReviewDraftComment, ReviewDraftMode, ReviewDraftSide } from "@/review/state";
+export type {
+  DiffModeOverride,
+  ReviewDraftComment,
+  ReviewDraftMode,
+  ReviewDraftSide,
+} from "@/review/state";
 
-const STORE_VERSION = 1;
+// v2 dropped persisted activeModesByScope (diff mode overrides are in-memory only).
+const STORE_VERSION = 2;
 const CONTEXT_RADIUS = 3;
 const EMPTY_REVIEW_DRAFT_COMMENTS: ReviewDraftComment[] = [];
 
@@ -55,7 +64,7 @@ export type ReviewDraftCommentInput = Omit<ReviewDraftComment, "id" | "createdAt
   Partial<Pick<ReviewDraftComment, "id" | "createdAt" | "updatedAt">>;
 
 interface ReviewDraftStoreActions {
-  setActiveMode: (input: { scopeKey: string; mode: ReviewDraftMode }) => void;
+  setDiffModeOverride: (input: { scopeKey: string; override: DiffModeOverride }) => void;
   addComment: (input: { key: string; comment: ReviewDraftCommentInput }) => ReviewDraftComment;
   updateComment: (input: {
     key: string;
@@ -129,9 +138,9 @@ export const useReviewDraftStore = create<ReviewDraftStore>()(
   persist(
     (set) => ({
       drafts: {},
-      activeModesByScope: {},
-      setActiveMode: (input) => {
-        set((state) => setActiveModeInState(state, input));
+      diffModeOverrides: {},
+      setDiffModeOverride: (input) => {
+        set((state) => setDiffModeOverrideInState(state, input));
       },
       addComment: ({ key, comment }) => {
         const nextComment = createDraftComment(comment);
@@ -271,8 +280,18 @@ export function useReviewDraftComments(key: string): ReviewDraftComment[] {
   return useReviewDraftStore((state) => state.drafts[key] ?? EMPTY_REVIEW_DRAFT_COMMENTS);
 }
 
-export function useSetActiveReviewDraftMode(): ReviewDraftStoreActions["setActiveMode"] {
-  return useReviewDraftStore((state) => state.setActiveMode);
+export function useSetDiffModeOverride(): ReviewDraftStoreActions["setDiffModeOverride"] {
+  return useReviewDraftStore((state) => state.setDiffModeOverride);
+}
+
+// Non-React entry point: called from the checkout-status data boundary, where dirty-state
+// changes enter the app regardless of which screens are mounted.
+export function expireStaleDiffModeOverrides(input: {
+  serverId: string;
+  cwd: string;
+  isDirty: boolean;
+}): void {
+  useReviewDraftStore.setState((state) => expireStaleDiffModeOverridesInState(state, input));
 }
 
 export function useClearReviewDraft(): ReviewDraftStoreActions["clearReview"] {
@@ -291,7 +310,7 @@ export function getReviewDraftComments(key: string): ReviewDraftComment[] | unde
 }
 
 export function resetReviewDraftStore(): void {
-  useReviewDraftStore.setState({ drafts: {}, activeModesByScope: {} });
+  useReviewDraftStore.setState({ drafts: {}, diffModeOverrides: {} });
 }
 
 export function useReviewDraftCommentsForAttachment(input: {
@@ -309,8 +328,16 @@ export function useReviewCommentCount(key: string): number {
   return useReviewDraftStore((state) => state.drafts[key]?.length ?? 0);
 }
 
-export function useActiveReviewDraftMode(input: { scopeKey: string }): ReviewDraftMode | null {
-  return useReviewDraftStore((state) => state.activeModesByScope[input.scopeKey] ?? null);
+export function useResolvedDiffMode(input: {
+  scopeKey: string;
+  hasUncommittedChanges: boolean;
+}): ReviewDraftMode {
+  return useReviewDraftStore((state) =>
+    resolveDiffMode({
+      override: state.diffModeOverrides[input.scopeKey],
+      hasUncommittedChanges: input.hasUncommittedChanges,
+    }),
+  );
 }
 
 export function useReviewAttachmentSnapshot(input: {

packages/app/src/review/surface.test.tsx

@@ -168,7 +168,7 @@ function comment(overrides: Partial<ReviewDraftComment> = {}): ReviewDraftCommen
 
 describe("useInlineReviewController", () => {
   beforeEach(() => {
-    useReviewDraftStore.setState({ drafts: {}, activeModesByScope: {} });
+    useReviewDraftStore.setState({ drafts: {}, diffModeOverrides: {} });
   });
 
   afterEach(() => {