Address review: fix step-nav race, avoid empty drafts

- Report the current-step draft imperatively from the edit handlers (typing/toggling) instead of a state effect. The effect keyed on internalStep could fire with the previous step's selections before restoreStepAnswer corrected them, briefly persisting wrong step data.
- Only persist a draft while it holds real content, and clear it otherwise, so untouched questions (or ones the agent removes without a submit/cancel) don't leave empty entries accumulating in storage.
- Tighten the test tool-call mock to `satisfies PermissionToolCall` and add a test that an untouched question leaves no draft.

Generated-By: PostHog Code
Task-Id: 624a18e6-9eac-4d18-8bd4-62491e117547

Author: jonathanlab

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

@@ -21,7 +21,7 @@ function questionToolCall(toolCallId: string): PermissionToolCall {
         },
       ],
     },
-  } as unknown as PermissionToolCall;
+  } satisfies PermissionToolCall;
 }
 
 const options: PermissionOption[] = [];
@@ -105,4 +105,11 @@ describe("QuestionPermission draft persistence", () => {
     expect(onSelect).toHaveBeenCalled();
     expect(useQuestionDraftStore.getState().actions.getDraft("q-1")).toBeNull();
   });
+
+  it("does not persist a draft for an untouched question", async () => {
+    const { view } = renderQuestion(questionToolCall("q-1"));
+    // Mounting alone must not leave an empty draft behind on unmount.
+    view.unmount();
+    expect(useQuestionDraftStore.getState().actions.getDraft("q-1")).toBeNull();
+  });
 });

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

@@ -123,13 +123,23 @@ export function QuestionPermission({
     },
   );
 
-  // Persist the draft on every change; cleared when the question is resolved.
+  // Persist the draft on every change, but only while it holds real content, so
+  // questions the user never answers (or that the agent later removes without a
+  // submit/cancel) don't leave empty entries accumulating in storage.
   useEffect(() => {
-    setDraft(questionId, {
-      activeStep,
-      stepAnswers: Object.fromEntries(stepAnswers),
-    });
-  }, [questionId, activeStep, stepAnswers, setDraft]);
+    const hasContent = Array.from(stepAnswers.values()).some(
+      (answer) =>
+        answer.selectedIds.length > 0 || answer.customInput.trim() !== "",
+    );
+    if (hasContent) {
+      setDraft(questionId, {
+        activeStep,
+        stepAnswers: Object.fromEntries(stepAnswers),
+      });
+    } else {
+      clearDraft(questionId);
+    }
+  }, [questionId, activeStep, stepAnswers, setDraft, clearDraft]);
 
   // Snapshot of persisted answers used to seed the selector's per-step state on
   // mount, so restored values are shown when navigating between steps.

packages/ui/src/primitives/action-selector/useActionSelectorState.ts

@@ -97,6 +97,23 @@ export function useActionSelectorState({
   const numSteps = steps?.length ?? 0;
   const showSubmitButton = !hideSubmitButton && (multiSelect || hasSteps);
 
+  // Refs let the edit handlers report the current draft imperatively without
+  // re-subscribing on every keystroke.
+  const activeStepRef = useRef(activeStep);
+  activeStepRef.current = activeStep;
+  const checkedOptionsRef = useRef(checkedOptions);
+  checkedOptionsRef.current = checkedOptions;
+  const onDraftChangeRef = useRef(onDraftChange);
+  onDraftChangeRef.current = onDraftChange;
+
+  const reportDraft = useCallback((checked: Set<string>, input: string) => {
+    onDraftChangeRef.current?.(
+      activeStepRef.current,
+      Array.from(checked),
+      input,
+    );
+  }, []);
+
   const allOptions = useMemo(() => {
     const opts = allowCustomInput
       ? [...options, { id: OTHER_OPTION_ID, label: "Other", description: "" }]
@@ -206,24 +223,6 @@ export function useActionSelectorState({
     }
   }, [initialCustomInput, allOptions]);
 
-  // Report every edit to the current step's draft so callers can persist it.
-  // A ref skips the initial mount, so restoring a draft does not immediately
-  // echo it back.
-  const onDraftChangeRef = useRef(onDraftChange);
-  onDraftChangeRef.current = onDraftChange;
-  const draftReportMountedRef = useRef(false);
-  useEffect(() => {
-    if (!draftReportMountedRef.current) {
-      draftReportMountedRef.current = true;
-      return;
-    }
-    onDraftChangeRef.current?.(
-      internalStep,
-      Array.from(checkedOptions),
-      customInput,
-    );
-  }, [checkedOptions, customInput, internalStep]);
-
   const moveUp = useCallback(() => {
     setHoveredIndex(null);
     setSelectedIndex((prev) => (prev > 0 ? prev - 1 : numOptions - 1));
@@ -250,26 +249,25 @@ export function useActionSelectorState({
 
   const toggleCheck = useCallback(
     (optionId: string) => {
-      setCheckedOptions((prev) => {
-        const next = new Set(prev);
-        if (multiSelect) {
-          if (next.has(optionId)) {
-            next.delete(optionId);
-          } else {
-            next.add(optionId);
-          }
+      const next = new Set(checkedOptionsRef.current);
+      if (multiSelect) {
+        if (next.has(optionId)) {
+          next.delete(optionId);
         } else {
-          if (next.has(optionId)) {
-            next.clear();
-          } else {
-            next.clear();
-            next.add(optionId);
-          }
+          next.add(optionId);
         }
-        return next;
-      });
+      } else {
+        if (next.has(optionId)) {
+          next.clear();
+        } else {
+          next.clear();
+          next.add(optionId);
+        }
+      }
+      setCheckedOptions(next);
+      reportDraft(next, customInputRef.current);
     },
-    [multiSelect],
+    [multiSelect, reportDraft],
   );
 
   const handleSubmitMulti = useCallback(() => {
@@ -475,28 +473,29 @@ export function useActionSelectorState({
   const handleCustomInputChange = useCallback(
     (value: string) => {
       setCustomInput(value);
+      let nextChecked = checkedOptionsRef.current;
       if (
         showSubmitButton &&
         selectedOption &&
         needsCustomInput(selectedOption)
       ) {
-        setCheckedOptions((prev) => {
-          const next = new Set(prev);
-          if (value.trim()) {
-            if (!prev.has(selectedOption.id)) {
-              if (!multiSelect) {
-                next.clear();
-              }
-              next.add(selectedOption.id);
+        const next = new Set(checkedOptionsRef.current);
+        if (value.trim()) {
+          if (!next.has(selectedOption.id)) {
+            if (!multiSelect) {
+              next.clear();
             }
-          } else {
-            next.delete(selectedOption.id);
+            next.add(selectedOption.id);
           }
-          return next;
-        });
+        } else {
+          next.delete(selectedOption.id);
+        }
+        nextChecked = next;
+        setCheckedOptions(next);
       }
+      reportDraft(nextChecked, value);
     },
-    [showSubmitButton, selectedOption, multiSelect],
+    [showSubmitButton, selectedOption, multiSelect, reportDraft],
   );
 
   const ensureChecked = useCallback((optionId: string) => {