Persist in-progress question answers per question

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

@@ -60,7 +60,9 @@ export function PermissionSelector({
     case "switch_mode":
       return <SwitchModePermission {...props} />;
     case "question":
-      return <QuestionPermission {...props} />;
+      // Key per question so answer state resets (and its draft is read) when a
+      // new question replaces the current one without unmounting.
+      return <QuestionPermission key={toolCall.toolCallId} {...props} />;
     default:
       return <DefaultPermission {...props} />;
   }

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

@@ -0,0 +1,108 @@
+import type { PermissionOption } from "@agentclientprotocol/sdk";
+import { Theme } from "@radix-ui/themes";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { QuestionPermission } from "./QuestionPermission";
+import { useQuestionDraftStore } from "./questionDraftStore";
+import type { PermissionToolCall } from "./types";
+
+function questionToolCall(toolCallId: string): PermissionToolCall {
+  return {
+    toolCallId,
+    title: "Question",
+    _meta: {
+      codeToolKind: "question",
+      questions: [
+        {
+          question: "What should the button say?",
+          header: "Button label",
+          options: [{ label: "Submit" }, { label: "Continue" }],
+        },
+      ],
+    },
+  } as unknown as PermissionToolCall;
+}
+
+const options: PermissionOption[] = [];
+
+function renderQuestion(toolCall: PermissionToolCall) {
+  const onSelect = vi.fn();
+  const onCancel = vi.fn();
+  const view = render(
+    <Theme>
+      <QuestionPermission
+        key={toolCall.toolCallId}
+        toolCall={toolCall}
+        options={options}
+        onSelect={onSelect}
+        onCancel={onCancel}
+      />
+    </Theme>,
+  );
+  return { onSelect, onCancel, view };
+}
+
+describe("QuestionPermission draft persistence", () => {
+  beforeEach(() => {
+    useQuestionDraftStore.setState({ drafts: {} });
+  });
+
+  // Options are Submit (1), Continue (2), and the free-text "Other" row (3).
+  const OTHER_KEY = "3";
+
+  it("restores a half-typed answer after unmount and remount", async () => {
+    const user = userEvent.setup();
+    const toolCall = questionToolCall("q-1");
+
+    const { view } = renderQuestion(toolCall);
+
+    // Focus the free-text option and type an in-progress answer.
+    await user.keyboard(OTHER_KEY);
+    const input = await screen.findByPlaceholderText("Type your answer...");
+    await user.type(input, "Buy now");
+
+    // Simulate switching to another session (component unmounts).
+    view.unmount();
+
+    // Return to the session: a fresh mount for the same question id.
+    renderQuestion(toolCall);
+
+    const restored = await screen.findByPlaceholderText<HTMLInputElement>(
+      "Type your answer...",
+    );
+    expect(restored.value).toBe("Buy now");
+  });
+
+  it("does not leak a draft to a different question id", async () => {
+    const user = userEvent.setup();
+
+    const { view } = renderQuestion(questionToolCall("q-1"));
+    await user.keyboard(OTHER_KEY);
+    await user.type(
+      await screen.findByPlaceholderText("Type your answer..."),
+      "for q-1",
+    );
+    view.unmount();
+
+    renderQuestion(questionToolCall("q-2"));
+    const input = await screen.findByPlaceholderText<HTMLInputElement>(
+      "Type your answer...",
+    );
+    expect(input.value).toBe("");
+  });
+
+  it("clears the draft once the question is answered", async () => {
+    const user = userEvent.setup();
+    const toolCall = questionToolCall("q-1");
+
+    const { onSelect } = renderQuestion(toolCall);
+    await user.keyboard(OTHER_KEY);
+    const input = await screen.findByPlaceholderText("Type your answer...");
+    await user.type(input, "Done");
+    await user.keyboard("{Enter}");
+
+    expect(onSelect).toHaveBeenCalled();
+    expect(useQuestionDraftStore.getState().actions.getDraft("q-1")).toBeNull();
+  });
+});

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

@@ -14,7 +14,8 @@ import {
   SUBMIT_OPTION_ID,
 } from "@posthog/ui/primitives/ActionSelector";
 import { Box, Flex, Text } from "@radix-ui/themes";
-import { useCallback, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
+import { useQuestionDraftStore } from "./questionDraftStore";
 import { type BasePermissionProps, toSelectorOptions } from "./types";
 
 function parseQuestionMeta(raw: unknown): QuestionMeta | undefined {
@@ -98,11 +99,42 @@ export function QuestionPermission({
   const allQuestions = meta?.questions ?? [];
   const totalQuestions = allQuestions.length;
 
-  const [activeStep, setActiveStep] = useState(0);
+  // Drafts are keyed per question (the tool-call id), so a half-typed answer
+  // survives switching to another session and back.
+  const questionId = toolCall.toolCallId;
+  const { getDraft, setDraft, clearDraft } = useQuestionDraftStore(
+    (s) => s.actions,
+  );
+
+  const [activeStep, setActiveStep] = useState(
+    () => getDraft(questionId)?.activeStep ?? 0,
+  );
   const [stepAnswers, setStepAnswers] = useState<Map<number, StepAnswer>>(
-    () => new Map(),
+    () => {
+      const saved = getDraft(questionId)?.stepAnswers;
+      return saved
+        ? new Map(
+            Object.entries(saved).map(([step, answer]) => [
+              Number(step),
+              answer,
+            ]),
+          )
+        : new Map();
+    },
   );
 
+  // Persist the draft on every change; cleared when the question is resolved.
+  useEffect(() => {
+    setDraft(questionId, {
+      activeStep,
+      stepAnswers: Object.fromEntries(stepAnswers),
+    });
+  }, [questionId, activeStep, stepAnswers, setDraft]);
+
+  // Snapshot of persisted answers used to seed the selector's per-step state on
+  // mount, so restored values are shown when navigating between steps.
+  const [initialStepAnswers] = useState(() => Object.fromEntries(stepAnswers));
+
   const isOnSubmitStep = activeStep >= totalQuestions;
 
   const activeQuestion = isOnSubmitStep ? undefined : allQuestions[activeStep];
@@ -136,6 +168,23 @@ export function QuestionPermission({
     [activeStep, totalQuestions],
   );
 
+  const resolveSelect = useCallback(
+    (
+      optionId: string,
+      customInput?: string,
+      answers?: Record<string, string>,
+    ) => {
+      clearDraft(questionId);
+      onSelect(optionId, customInput, answers);
+    },
+    [clearDraft, questionId, onSelect],
+  );
+
+  const handleCancel = useCallback(() => {
+    clearDraft(questionId);
+    onCancel();
+  }, [clearDraft, questionId, onCancel]);
+
   const handleMultiSelect = useCallback(
     (optionIds: string[], customInput?: string) => {
       if (totalQuestions === 1) {
@@ -144,23 +193,23 @@ export function QuestionPermission({
           [0, { selectedIds: optionIds, customInput: customInput ?? "" }],
         ]);
         const answers = buildAnswersRecord(singleAnswer, allQuestions);
-        onSelect(filteredIds[0] ?? "other", customInput, answers);
+        resolveSelect(filteredIds[0] ?? "other", customInput, answers);
         return;
       }
       advanceStep(optionIds, customInput);
     },
-    [totalQuestions, onSelect, advanceStep, allQuestions],
+    [totalQuestions, resolveSelect, advanceStep, allQuestions],
   );
 
   const handleSelect = useCallback(
     (optionId: string, customInput?: string) => {
       if (isOnSubmitStep) {
         if (optionId === CANCEL_OPTION_ID) {
-          onCancel();
+          handleCancel();
           return;
         }
         const answers = buildAnswersRecord(stepAnswers, allQuestions);
-        onSelect(SUBMIT_OPTION_ID, undefined, answers);
+        resolveSelect(SUBMIT_OPTION_ID, undefined, answers);
         return;
       }
 
@@ -169,7 +218,7 @@ export function QuestionPermission({
           [0, { selectedIds: [optionId], customInput: customInput ?? "" }],
         ]);
         const answers = buildAnswersRecord(singleAnswer, allQuestions);
-        onSelect(optionId, customInput, answers);
+        resolveSelect(optionId, customInput, answers);
         return;
       }
 
@@ -180,8 +229,8 @@ export function QuestionPermission({
       stepAnswers,
       allQuestions,
       totalQuestions,
-      onSelect,
-      onCancel,
+      resolveSelect,
+      handleCancel,
       advanceStep,
     ],
   );
@@ -200,6 +249,19 @@ export function QuestionPermission({
     [],
   );
 
+  // Persist in-progress edits to the current step before it is committed via a
+  // step change or submit, so half-typed text is not lost on a session switch.
+  const handleDraftChange = useCallback(
+    (stepIndex: number, optionIds: string[], customInput: string) => {
+      setStepAnswers((prev) => {
+        const next = new Map(prev);
+        next.set(stepIndex, { selectedIds: optionIds, customInput });
+        return next;
+      });
+    },
+    [],
+  );
+
   const handleStepChange = useCallback((stepIndex: number) => {
     setActiveStep(stepIndex);
   }, []);
@@ -272,11 +334,14 @@ export function QuestionPermission({
       currentStep={activeStep}
       steps={steps}
       initialSelections={currentStepAnswer?.selectedIds}
+      initialCustomInput={currentStepAnswer?.customInput}
+      initialStepAnswers={initialStepAnswers}
       onSelect={handleSelect}
       onMultiSelect={handleMultiSelect}
-      onCancel={onCancel}
+      onCancel={handleCancel}
       onStepChange={handleStepChange}
       onStepAnswer={handleStepAnswer}
+      onDraftChange={handleDraftChange}
     />
   );
 }

packages/ui/src/features/permissions/questionDraftStore.test.ts

@@ -0,0 +1,85 @@
+import { registerRendererStateStorage } from "@posthog/ui/shell/rendererStorage";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useQuestionDraftStore } from "./questionDraftStore";
+
+const getItem = vi.fn();
+const setItem = vi.fn();
+const removeItem = vi.fn();
+
+registerRendererStateStorage({ getItem, setItem, removeItem });
+
+describe("feature questionDraftStore", () => {
+  beforeEach(() => {
+    getItem.mockReset();
+    setItem.mockReset();
+    removeItem.mockReset();
+    getItem.mockResolvedValue(null);
+    setItem.mockResolvedValue(undefined);
+    removeItem.mockResolvedValue(undefined);
+
+    useQuestionDraftStore.setState({ drafts: {} });
+  });
+
+  it("keeps drafts separate per question id", () => {
+    const { setDraft, getDraft } = useQuestionDraftStore.getState().actions;
+
+    setDraft("tool-call-a", {
+      activeStep: 0,
+      stepAnswers: { 0: { selectedIds: [], customInput: "answer for a" } },
+    });
+    setDraft("tool-call-b", {
+      activeStep: 1,
+      stepAnswers: { 0: { selectedIds: ["option_1"], customInput: "" } },
+    });
+
+    expect(getDraft("tool-call-a")?.stepAnswers[0]?.customInput).toBe(
+      "answer for a",
+    );
+    expect(getDraft("tool-call-b")?.activeStep).toBe(1);
+    expect(getDraft("tool-call-b")?.stepAnswers[0]?.selectedIds).toEqual([
+      "option_1",
+    ]);
+  });
+
+  it("returns null for an unknown question id", () => {
+    expect(
+      useQuestionDraftStore.getState().actions.getDraft("nope"),
+    ).toBeNull();
+  });
+
+  it("clears a resolved question's draft without touching others", () => {
+    const { setDraft, clearDraft, getDraft } =
+      useQuestionDraftStore.getState().actions;
+
+    setDraft("tool-call-a", {
+      activeStep: 0,
+      stepAnswers: { 0: { selectedIds: [], customInput: "keep" } },
+    });
+    setDraft("tool-call-b", {
+      activeStep: 0,
+      stepAnswers: { 0: { selectedIds: [], customInput: "remove" } },
+    });
+
+    clearDraft("tool-call-b");
+
+    expect(getDraft("tool-call-b")).toBeNull();
+    expect(getDraft("tool-call-a")?.stepAnswers[0]?.customInput).toBe("keep");
+  });
+
+  it("persists drafts to the storage backend", async () => {
+    useQuestionDraftStore.getState().actions.setDraft("tool-call-a", {
+      activeStep: 0,
+      stepAnswers: { 0: { selectedIds: [], customInput: "persist me" } },
+    });
+
+    await vi.waitFor(() => {
+      expect(setItem).toHaveBeenCalled();
+    });
+
+    const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
+    const persisted = JSON.parse(lastCall[1]);
+    expect(
+      persisted.state.drafts["tool-call-a"].stepAnswers[0].customInput,
+    ).toBe("persist me");
+  });
+});