fix(web): an unsaved citation comment survives a provider question

apps/web/src/components/ChatView.tsx

@@ -347,6 +347,8 @@ import {
   composerDraftHasUserContent,
   type ComposerFileAttachment,
   type ComposerImageAttachment,
+  composerTargetKey,
+  type ComposerThreadTarget,
   type DraftThreadEnvMode,
   useComposerDraftStore,
   DraftId,
@@ -409,6 +411,11 @@ import { isTimelineScrollTarget } from "./chat/timelineScrollTarget";
 import { DraftHeroHeadline } from "./chat/DraftHeroHeadline";
 import { ExpandedImageDialog } from "./chat/ExpandedImageDialog";
 import { PullRequestThreadDialog } from "./PullRequestThreadDialog";
+import {
+  type AssistantCitationCommentDraftEntries,
+  restoreAssistantCitationCommentDrafts,
+  takeAssistantCitationCommentDraftsForComposer,
+} from "./chat/assistantCitationCommentDrafts";
 import type { AssistantCitationRequest } from "./chat/AssistantCitationSource";
 import { MessagesTimeline, type MessagesTimelineHistoryControls } from "./chat/MessagesTimeline";
 import { ProviderSubagentBar } from "./chat/ProviderSubagentBar";
@@ -1739,7 +1746,16 @@ export default function ChatView(props: ChatViewProps) {
   const setComposerDraftInteractionMode = useComposerDraftStore(
     (store) => store.setInteractionMode,
   );
-  const clearComposerDraftContent = useComposerDraftStore((store) => store.clearComposerContent);
+  const clearComposerContent = useComposerDraftStore((store) => store.clearComposerContent);
+  // A sent prompt takes its unsaved citation comments with it; a send that
+  // fails hands them back with the prompt (see assistantCitationCommentDrafts).
+  const clearComposerDraftContent = useCallback(
+    (target: ComposerThreadTarget) => {
+      clearComposerContent(target);
+      return takeAssistantCitationCommentDraftsForComposer(composerTargetKey(target));
+    },
+    [clearComposerContent],
+  );
   const setDraftThreadContext = useComposerDraftStore((store) => store.setDraftThreadContext);
   const getDraftSessionByLogicalProjectKey = useComposerDraftStore(
     (store) => store.getDraftSessionByLogicalProjectKey,
@@ -8442,7 +8458,7 @@ export default function ChatView(props: ChatViewProps) {
       const followUpPreviewAnnotations = [...composerPreviewAnnotations];
       const followUpThreadContexts = [...composerThreadContexts];
       promptRef.current = "";
-      clearComposerDraftContent(composerDraftTarget);
+      const followUpCitationDrafts = clearComposerDraftContent(composerDraftTarget);
       composerRef.current?.resetCursorState();
       const followUpSent = await onSubmitPlanFollowUp({
         text: followUp.text,
@@ -8456,6 +8472,7 @@ export default function ChatView(props: ChatViewProps) {
       });
       if (!followUpSent) {
         promptRef.current = followUpPromptSnapshot;
+        restoreAssistantCitationCommentDrafts(followUpCitationDrafts);
         composerTerminalContextsRef.current = [...followUpTerminalContexts];
         restorePlanFollowUpComposer({
           snapshot: {
@@ -8746,6 +8763,7 @@ export default function ChatView(props: ChatViewProps) {
     if (multipleModelSelections !== null) {
       const failedSelections: ModelSelection[] = [];
       let clearedDraft = false;
+      let clearedCitationDrafts: AssistantCitationCommentDraftEntries = [];
       let releasedComposer = false;
       let canRestoreDraft = () => false;
       let startedCount = 0;
@@ -8766,7 +8784,7 @@ export default function ChatView(props: ChatViewProps) {
             "New thread",
         );
         promptRef.current = "";
-        clearComposerDraftContent(composerDraftTarget);
+        clearedCitationDrafts = clearComposerDraftContent(composerDraftTarget);
         composerRef.current?.resetCursorState();
         clearedDraft = true;
         const clearedDraftSnapshot = useComposerDraftStore
@@ -8935,6 +8953,7 @@ export default function ChatView(props: ChatViewProps) {
           setMultipleModelSelections(failedSelections);
           if (clearedDraft) {
             setComposerDraftPrompt(composerDraftTarget, messageTextForSend);
+            restoreAssistantCitationCommentDrafts(clearedCitationDrafts);
             addComposerDraftImages(
               composerDraftTarget,
               composerImagesSnapshot.map(cloneComposerImageForRetry),
@@ -9076,7 +9095,7 @@ export default function ChatView(props: ChatViewProps) {
       );
     }
     promptRef.current = "";
-    clearComposerDraftContent(composerDraftTarget);
+    const sentCitationDrafts = clearComposerDraftContent(composerDraftTarget);
     composerRef.current?.resetCursorState();
 
     let firstComposerImageName: string | null = null;
@@ -9321,6 +9340,7 @@ export default function ChatView(props: ChatViewProps) {
           return next.length === existing.length ? existing : next;
         });
         promptRef.current = messageTextForSend;
+        restoreAssistantCitationCommentDrafts(sentCitationDrafts);
         const retryComposerImages = composerImagesSnapshot.map(cloneComposerImageForRetry);
         composerImagesRef.current = retryComposerImages;
         composerFilesRef.current = composerFilesSnapshot;

apps/web/src/components/ComposerPromptEditorTiptap.tsx

@@ -68,6 +68,7 @@ import { basenameOfPath } from "~/pierre-icons";
 import { FileTagChipContent } from "./chat/FileTagChip";
 import { SkillChipIcon } from "./chat/SkillInlineText";
 import { AssistantCitationChip } from "./chat/AssistantCitationChip";
+import { assistantCitationDraftKey } from "./chat/assistantCitationCommentDrafts";
 import { getTimelinePageScrollKey } from "./chat/pageScrollController";
 import { ContextChipPopover } from "./contextChipParts";
 import { Button } from "./ui/button";
@@ -148,6 +149,8 @@ export interface ComposerPromptEditorProps {
   onPageScrollKeyUp?: (key: string) => void;
   onPageScrollRelease?: () => void;
   onCitationSubmitAndSend?: () => void;
+  /** Keeps this composer's unsaved citation comments apart from another thread's. */
+  citationDraftScope?: string;
   onPaste: React.ClipboardEventHandler<HTMLElement>;
   editorRef: React.RefObject<ComposerPromptEditorHandle | null>;
 }
@@ -169,7 +172,8 @@ const ComposerCitationCommentContext = createContext<{
   openComment: OpenCitationComment | null;
   onOpenChange: (citeKey: string, open: boolean) => void;
   onSubmitAndSend: () => void;
-}>({ openComment: null, onOpenChange: () => {}, onSubmitAndSend: () => {} });
+  draftScope: string;
+}>({ openComment: null, onOpenChange: () => {}, onSubmitAndSend: () => {}, draftScope: "" });
 
 const RichComposerSkillsContext = createContext<ReadonlyArray<ServerProviderSkill>>([]);
 
@@ -338,6 +342,20 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
   const commentContext = use(ComposerCitationCommentContext);
   const citation = node.attrs.citation as AssistantCitation;
   const citeKey = node.attrs.citeKey as string;
+  // Unsaved comments are keyed by the citation and its place among identical
+  // ones, which survives the rebuild a provider question causes; citeKey does not.
+  const draftKey = useMemo(() => {
+    const pos = typeof getPos === "function" ? getPos() : undefined;
+    const citationsBefore: Array<AssistantCitation> = [];
+    if (typeof pos === "number") {
+      editor.state.doc.nodesBetween(0, pos, (earlier) => {
+        if (earlier.type.name === "composer-citation") {
+          citationsBefore.push(earlier.attrs.citation as AssistantCitation);
+        }
+      });
+    }
+    return assistantCitationDraftKey(citation, citationsBefore, commentContext.draftScope);
+  }, [citation, commentContext.draftScope, editor.state.doc, getPos]);
   const commentTarget =
     commentContext.openComment?.key === citeKey ? commentContext.openComment : null;
 
@@ -425,6 +443,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
             commentContext.onOpenChange(citeKey, open);
           },
           ...(commentTarget?.removeOnCancel ? { onCancel: onRemove } : {}),
+          draftKey,
           onSave: onSaveComment,
           onSaveAndSend: (comment) => {
             if (!onSaveComment(comment)) return false;
@@ -632,6 +651,7 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
     onPageScrollKeyUp,
     onPageScrollRelease,
     onCitationSubmitAndSend,
+    citationDraftScope,
     onPaste,
     editorRef,
   } = props;
@@ -713,8 +733,9 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
         });
       },
       onSubmitAndSend: onCitationSubmitAndSend ?? (() => {}),
+      draftScope: citationDraftScope ?? "",
     }),
-    [onCitationSubmitAndSend, openCitation],
+    [citationDraftScope, onCitationSubmitAndSend, openCitation],
   );
 
   const handleEditorChange = useCallback((updated: TiptapEditor) => {

apps/web/src/components/chat/AssistantCitationChip.test.tsx

@@ -56,13 +56,19 @@ const sourceAnchor = {
 
 let renderer: ReactTestRenderer;
 
-function mount(onSave = vi.fn(() => true)) {
+function mount(onSave = vi.fn(() => true), draftKey?: string) {
   function Composer() {
     const [open, setOpen] = useState(true);
     return (
       <AssistantCitationChip
         citation={citation}
-        commentEditor={{ open, sourceAnchor, onOpenChange: setOpen, onSave }}
+        commentEditor={{
+          open,
+          sourceAnchor,
+          onOpenChange: setOpen,
+          onSave,
+          ...(draftKey === undefined ? {} : { draftKey }),
+        }}
       />
     );
   }
@@ -101,6 +107,65 @@ afterEach(() => {
   vi.unstubAllGlobals();
 });
 
+describe("citation comment draft across a remount", () => {
+  it("resumes an unsaved comment when the chip is mounted again", () => {
+    // The composer editor replaces its document when a provider question
+    // borrows it, which unmounts the node view without any dismissal.
+    mount();
+    typeComment("still typing");
+    act(() => renderer.unmount());
+
+    const onSave = mount();
+    expect(renderer.root.findByType("textarea").props.value).toBe("still typing");
+    clickButton("Save");
+    expect(onSave).toHaveBeenCalledWith("still typing");
+  });
+
+  it("keeps the drafts of two identical citations apart by their node key", () => {
+    mount(
+      vi.fn(() => true),
+      "cite-1",
+    );
+    typeComment("for the first");
+    act(() => renderer.unmount());
+
+    mount(
+      vi.fn(() => true),
+      "cite-2",
+    );
+    expect(renderer.root.findByType("textarea").props.value).toBe("");
+    clickButton("Cancel");
+    act(() => renderer.unmount());
+
+    mount(
+      vi.fn(() => true),
+      "cite-1",
+    );
+    expect(renderer.root.findByType("textarea").props.value).toBe("for the first");
+    clickButton("Cancel");
+  });
+
+  it("forgets the draft once it was saved", () => {
+    mount();
+    typeComment("saved now");
+    clickButton("Save");
+    act(() => renderer.unmount());
+
+    mount();
+    expect(renderer.root.findByType("textarea").props.value).toBe("");
+  });
+
+  it("forgets the draft once it was cancelled", () => {
+    mount();
+    typeComment("never mind");
+    clickButton("Cancel");
+    act(() => renderer.unmount());
+
+    mount();
+    expect(renderer.root.findByType("textarea").props.value).toBe("");
+  });
+});
+
 describe("citation comment source disappearance", () => {
   it("preserves an over-length draft at the composer until it can be shortened and saved", () => {
     const onSave = mount();

apps/web/src/components/chat/AssistantCitationChip.tsx

@@ -21,6 +21,11 @@ import { ContextChip, ContextChipAction, ContextChipLabel } from "../ContextChip
 import { Tooltip, TooltipPopup, TooltipTrigger } from "../ui/tooltip";
 import { Popover, PopoverPopup, PopoverTrigger } from "../ui/popover";
 import { AssistantCitationCommentEditor } from "./AssistantCitationCommentEditor";
+import {
+  clearAssistantCitationCommentDraft,
+  readAssistantCitationCommentDraft,
+  writeAssistantCitationCommentDraft,
+} from "./assistantCitationCommentDrafts";
 import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal";
 import { observeAssistantCitationCommentSource } from "./AssistantCitationSource";
 import { composerFloatingLayerProps } from "./composerEventScope";
@@ -41,6 +46,8 @@ export function AssistantCitationChip({
     onSaveAndSend?: (comment: string) => boolean;
     /** Returns focus to the host editor when the popover closes instead of to the pencil trigger. */
     onRestoreFocus?: () => void;
+    /** Where the unsaved comment is kept across a remount; the serialized citation is the fallback. */
+    draftKey?: string;
   };
 }) {
   const navigate = useNavigate();
@@ -52,17 +59,25 @@ export function AssistantCitationChip({
   const commentOpen = commentEditor?.open ?? false;
   const sourceAnchor = commentEditor?.sourceAnchor;
   const activeSourceAnchor = sourceAnchor === unavailableSourceAnchor ? undefined : sourceAnchor;
+  // The draft outlives this chip (see assistantCitationCommentDrafts), so an
+  // editor that comes back after the composer was borrowed for a question
+  // resumes it, and the dismissal rules see it as typed.
+  const draftKey = commentEditor?.draftKey ?? serializeAssistantCitation(citation);
   useEffect(() => {
-    if (!commentOpen) draftCommentRef.current = null;
-  }, [commentOpen]);
+    draftCommentRef.current = commentOpen ? readAssistantCitationCommentDraft(draftKey) : null;
+  }, [commentOpen, draftKey]);
   const settleDraftOnClose = (reason: string): boolean => {
     const dismissal = resolveAssistantCitationCommentDismissal({
       reason,
       draft: draftCommentRef.current,
       savedComment: citation.comment,
     });
-    if (dismissal.kind === "commit") return commentEditor?.onSave(dismissal.comment) ?? true;
-    return dismissal.kind !== "keep-open";
+    if (dismissal.kind === "keep-open") return false;
+    if (dismissal.kind === "commit" && !(commentEditor?.onSave(dismissal.comment) ?? true)) {
+      return false;
+    }
+    clearAssistantCitationCommentDraft(draftKey);
+    return true;
   };
   const onSourceUnavailable = useEffectEvent(() => {
     if (!sourceAnchor) return;
@@ -195,27 +210,32 @@ export function AssistantCitationChip({
               onPointerDown={(event) => event.stopPropagation()}
             >
               <AssistantCitationCommentEditor
-                key={serializeAssistantCitation(citation)}
+                key={draftKey}
                 citation={citation}
+                draft={readAssistantCitationCommentDraft(draftKey)}
                 inputRef={commentInputRef}
                 onDraftChange={(comment) => {
                   draftCommentRef.current = comment;
+                  writeAssistantCitationCommentDraft(draftKey, comment);
                 }}
                 onSubmit={(comment) => {
                   if (!commentEditor.onSave(comment)) return false;
+                  clearAssistantCitationCommentDraft(draftKey);
                   commentEditor.onOpenChange(false);
                   return true;
                 }}
                 {...(commentEditor.onSaveAndSend
                   ? {
                       onSubmitAndSend: (comment: string) => {
                         if (!commentEditor.onSaveAndSend?.(comment)) return false;
+                        clearAssistantCitationCommentDraft(draftKey);
                         commentEditor.onOpenChange(false);
                         return true;
                       },
                     }
                   : {})}
                 onCancel={() => {
+                  clearAssistantCitationCommentDraft(draftKey);
                   if (commentEditor.onCancel) {
                     commentEditor.onCancel();
                   } else {

apps/web/src/components/chat/AssistantCitationCommentEditor.tsx

@@ -10,15 +10,18 @@ export function AssistantCitationCommentEditor({
   onSubmitAndSend,
   onCancel,
   onDraftChange,
+  draft,
 }: {
   citation: AssistantCitation;
+  /** An unsaved comment from an earlier mount of this editor, shown instead of the saved one. */
+  draft?: string | null;
   inputRef?: Ref<HTMLTextAreaElement>;
   onSubmit: (comment: string) => boolean;
   onSubmitAndSend?: (comment: string) => boolean;
   onCancel: () => void;
   onDraftChange?: (comment: string) => void;
 }) {
-  const [comment, setComment] = useState(citation.comment ?? "");
+  const [comment, setComment] = useState(draft ?? citation.comment ?? "");
   const commentTooLong = comment.length > ASSISTANT_CITATION_MAX_COMMENT_LENGTH;
   const submit = () => {
     if (!commentTooLong) onSubmit(comment);