Fix plan-approval mode not sticking across app restarts (#3737)

Author: haacked

packages/ui/src/features/permissions/PlanApprovalSelector.test.tsx

@@ -1,7 +1,7 @@
 import type { PermissionOption } from "@agentclientprotocol/sdk";
 import { useSettingsStore } from "@posthog/ui/features/settings/settingsStore";
 import { Theme } from "@radix-ui/themes";
-import { render, screen } from "@testing-library/react";
+import { act, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { PlanApprovalSelector } from "./PlanApprovalSelector";
@@ -92,6 +92,91 @@ describe("PlanApprovalSelector", () => {
     expect(useSettingsStore.getState().lastPlanApprovalMode).toBe("auto");
   });
 
+  it("picks up a remembered mode that loads after mount", async () => {
+    // Settings persist asynchronously (an IPC round trip on desktop), so the
+    // selector can mount before `lastPlanApprovalMode` has loaded from disk.
+    const user = userEvent.setup();
+    const { onSelect } = renderSelector([AUTO, ACCEPT_EDITS, DEFAULT_MODE]);
+
+    // The remembered choice loads in after mount.
+    act(() => {
+      useSettingsStore.setState({ lastPlanApprovalMode: "acceptEdits" });
+    });
+
+    await user.click(screen.getByText("Approve and proceed"));
+
+    expect(onSelect).toHaveBeenCalledWith("acceptEdits");
+  });
+
+  it("does not clobber a mode the user already picked", async () => {
+    const user = userEvent.setup();
+    const { onSelect } = renderSelector([AUTO, ACCEPT_EDITS, DEFAULT_MODE]);
+
+    await user.click(screen.getByRole("button", { name: "Mode" }));
+    await user.click(await screen.findByText("Manually approve edits"));
+    // The dropdown applies the pick when its close animation completes, not
+    // synchronously with the click, so wait for it to actually land before
+    // moving on — otherwise this race decides the test's outcome.
+    await waitFor(() =>
+      expect(screen.getByRole("button", { name: "Mode" })).toHaveTextContent(
+        "Manually approve edits",
+      ),
+    );
+
+    // The remembered choice loads in after the user already picked a mode.
+    act(() => {
+      useSettingsStore.setState({ lastPlanApprovalMode: "acceptEdits" });
+    });
+
+    await user.click(screen.getByText("Approve and proceed"));
+
+    expect(onSelect).toHaveBeenCalledWith("default");
+  });
+
+  it("drops a manual pick when a later approval request reuses this instance", async () => {
+    const user = userEvent.setup();
+    const onSelect = vi.fn();
+    const options = [AUTO, ACCEPT_EDITS, DEFAULT_MODE];
+    const { rerender } = render(
+      <Theme>
+        <PlanApprovalSelector
+          toolCall={toolCall}
+          options={options}
+          onSelect={onSelect}
+          onCancel={vi.fn()}
+        />
+      </Theme>,
+    );
+
+    await user.click(screen.getByRole("button", { name: "Mode" }));
+    await user.click(await screen.findByText("Manually approve edits"));
+    // The dropdown applies the pick when its close animation completes, not
+    // synchronously with the click — wait for it to land before rerendering,
+    // or the pending pick can apply after (and survive) the reset below.
+    await waitFor(() =>
+      expect(screen.getByRole("button", { name: "Mode" })).toHaveTextContent(
+        "Manually approve edits",
+      ),
+    );
+
+    // The component isn't guaranteed to unmount between requests; a new
+    // toolCallId means a new request even if this instance is reused.
+    rerender(
+      <Theme>
+        <PlanApprovalSelector
+          toolCall={{ ...toolCall, toolCallId: "plan-2" } as PermissionToolCall}
+          options={options}
+          onSelect={onSelect}
+          onCancel={vi.fn()}
+        />
+      </Theme>,
+    );
+
+    await user.click(screen.getByText("Approve and proceed"));
+
+    expect(onSelect).toHaveBeenCalledWith("auto");
+  });
+
   it("rejects with the typed feedback", async () => {
     const user = userEvent.setup();
     const { onSelect } = renderSelector([DEFAULT_MODE, REJECT]);

packages/ui/src/features/permissions/PlanApprovalSelector.tsx

@@ -61,6 +61,7 @@ function isInteractiveElementInDifferentCell(
  * `onSelect(<rejectOptionId>, feedback)`.
  */
 export function PlanApprovalSelector({
+  toolCall,
   options,
   onSelect,
   onCancel,
@@ -80,6 +81,11 @@ export function PlanApprovalSelector({
 
   // Resolution order: the mode last approved with (remembered preference),
   // then "auto", then manual-approve, then any single-use mode, then the first.
+  // Settings persist asynchronously (an IPC round trip on desktop), so
+  // `lastApprovalMode` can still be its pre-hydration default on mount — e.g.
+  // resuming a task with an already-pending plan approval. Recomputing this
+  // via `useMemo` (rather than seeding a `useState` once) means it stays
+  // correct once the store finishes hydrating.
   const initialMode = useMemo(() => {
     const has = (id: string) => approveOptions.some((o) => o.optionId === id);
     return (
@@ -93,7 +99,23 @@ export function PlanApprovalSelector({
     );
   }, [approveOptions, lastApprovalMode]);
 
-  const [selectedMode, setSelectedMode] = useState(initialMode);
+  // Only the user's own pick lives in state; everything else derives from
+  // `initialMode` so it tracks `lastApprovalMode` live instead of freezing it
+  // at mount — derive it, don't duplicate it.
+  const [explicitMode, setExplicitMode] = useState<string | undefined>(
+    undefined,
+  );
+  // This component can survive to a later approval request without
+  // remounting, so a pick made for the previous request must not leak into
+  // (and potentially not exist in) this one. Reset during render rather than
+  // in an effect: it takes effect before this render paints instead of one
+  // render later, avoiding a flash of the stale mode.
+  const lastToolCallIdRef = useRef(toolCall.toolCallId);
+  if (lastToolCallIdRef.current !== toolCall.toolCallId) {
+    lastToolCallIdRef.current = toolCall.toolCallId;
+    setExplicitMode(undefined);
+  }
+  const selectedMode = explicitMode ?? initialMode;
   const [selectedIndex, setSelectedIndex] = useState(0);
   const [hoveredIndex, setHoveredIndex] = useState<number | null>(null);
   const [feedback, setFeedback] = useState("");
@@ -285,7 +307,7 @@ export function PlanApprovalSelector({
                   <Box onClick={(e) => e.stopPropagation()}>
                     <ModeSelector
                       modeOption={modeConfigOption}
-                      onChange={(value) => setSelectedMode(value)}
+                      onChange={(value) => setExplicitMode(value)}
                       allowBypassPermissions
                     />
                   </Box>