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

apps/web/src/components/ComposerPromptEditorTiptap.citationDraft.test.tsx

@@ -0,0 +1,123 @@
+// @vitest-environment jsdom
+
+import { EnvironmentId, MessageId, ThreadId } from "@t3tools/contracts";
+import {
+  serializeAssistantCitation,
+  withAssistantCitationComment,
+} from "@t3tools/shared/assistantCitations";
+import { act, createRef, useState, type ReactNode } from "react";
+import { createRoot, type Root } from "react-dom/client";
+import { afterEach, beforeEach, expect, it, vi } from "vite-plus/test";
+
+import type { AssistantCitationSourceAnchor } from "~/lib/assistantTextSelection";
+
+vi.mock("./chat/AssistantCitationSource", () => ({
+  observeAssistantCitationCommentSource: () => () => {},
+}));
+vi.mock("@tanstack/react-router", () => ({
+  useNavigate: () => () => {},
+  Link: ({ children }: { children: ReactNode }) => <a>{children}</a>,
+}));
+// Keep the real chip and comment editor while replacing floating layers.
+vi.mock("./ui/popover", () => ({
+  Popover: ({ children }: { children: ReactNode }) => <>{children}</>,
+  PopoverTrigger: ({ "aria-label": label }: { "aria-label": string }) => (
+    <button aria-label={label} />
+  ),
+  PopoverPopup: ({ children }: { children: ReactNode }) => <>{children}</>,
+  PopoverClose: ({ children }: { children: ReactNode }) => <>{children}</>,
+}));
+
+import { ComposerPromptEditor, type ComposerPromptEditorHandle } from "./ComposerPromptEditor";
+
+const citation = {
+  version: 1 as const,
+  environmentId: EnvironmentId.make("environment"),
+  threadId: ThreadId.make("thread"),
+  messageId: MessageId.make("source"),
+  text: "hello",
+  start: 0,
+  end: 5,
+  prefix: "",
+  suffix: "",
+};
+const promptBefore = `See ${serializeAssistantCitation(citation)} `;
+
+let root: Root;
+let container: HTMLDivElement;
+
+beforeEach(() => {
+  vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+  container = document.createElement("div");
+  document.body.append(container);
+  root = createRoot(container);
+});
+
+afterEach(async () => {
+  await act(async () => root.unmount());
+  container.remove();
+  vi.unstubAllGlobals();
+});
+
+// Mirrors ChatComposer: a pending question shows its answer in the same editor.
+function Composer({
+  question,
+  editorRef,
+}: {
+  question: boolean;
+  editorRef: React.RefObject<ComposerPromptEditorHandle | null>;
+}) {
+  const [prompt, setPrompt] = useState(promptBefore);
+  const [answer, setAnswer] = useState("");
+  return (
+    <>
+      <output data-prompt={prompt} data-answer={answer} />
+      <ComposerPromptEditor
+        editorRef={editorRef}
+        value={question ? answer : prompt}
+        cursor={0}
+        contextRecords={new Map()}
+        skills={[]}
+        disabled={false}
+        placeholder=""
+        onChange={(next) => (question ? setAnswer(next) : setPrompt(next))}
+        onCitationDraftSave={(replacedValue, savedValue) => {
+          if (replacedValue === prompt) setPrompt(savedValue);
+        }}
+        onPaste={() => {}}
+      />
+    </>
+  );
+}
+
+it("keeps an unsaved citation comment when a question takes the composer", async () => {
+  const editorRef = createRef<ComposerPromptEditorHandle>();
+  const drafts = () => container.querySelector("output")?.dataset;
+
+  await act(async () => root.render(<Composer question={false} editorRef={editorRef} />));
+  await act(async () =>
+    editorRef.current?.requestCitationComment({
+      previousValue: "See ",
+      value: promptBefore,
+      citationStart: "See ".length,
+      sourceAnchor: {} as AssistantCitationSourceAnchor,
+      insertedSpaces: { before: false, after: true },
+    }),
+  );
+  const textarea = container.querySelector("textarea");
+  if (!textarea) throw new Error("Comment editor did not open");
+  await act(async () => {
+    const setValue = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, "value")?.set;
+    setValue?.call(textarea, "check this");
+    textarea.dispatchEvent(new Event("input", { bubbles: true }));
+  });
+
+  await act(async () => root.render(<Composer question editorRef={editorRef} />));
+  expect(drafts()?.answer).toBe("");
+  await act(async () => root.render(<Composer question={false} editorRef={editorRef} />));
+
+  expect(drafts()?.prompt).toBe(
+    `See ${serializeAssistantCitation(withAssistantCitationComment(citation, "check this"))} `,
+  );
+  expect(container.querySelector('[aria-label="Edit citation comment"]')).not.toBeNull();
+});

apps/web/src/components/ComposerPromptEditorTiptap.tsx

@@ -1,5 +1,12 @@
 import { formatProviderSkillDisplayName } from "@t3tools/shared/inlineSkills";
-import { Extension, InputRule, Node, wrappingInputRule, type JSONContent } from "@tiptap/core";
+import {
+  Extension,
+  findChildren,
+  InputRule,
+  Node,
+  wrappingInputRule,
+  type JSONContent,
+} from "@tiptap/core";
 import { ReactNodeViewRenderer, NodeViewWrapper, type NodeViewProps } from "@tiptap/react";
 import StarterKit from "@tiptap/starter-kit";
 import { type Node as ProseMirrorNode } from "@tiptap/pm/model";
@@ -31,6 +38,7 @@ import {
   use,
   useCallback,
   useEffect,
+  useEffectEvent,
   useImperativeHandle,
   useLayoutEffect,
   useMemo,
@@ -107,6 +115,7 @@ import {
 import type { AssistantCitationSourceAnchor } from "~/lib/assistantTextSelection";
 import { Tooltip, TooltipPopup, TooltipTrigger } from "./ui/tooltip";
 import { ComposerCodeBlockNodeView } from "./chat/ComposerCodeBlockNodeView";
+import { resolveAssistantCitationCommentDismissal } from "./chat/assistantCitationCommentDismissal";
 import { composerCodeBlockHighlight } from "./composerCodeBlockHighlight";
 import { importPastedComposerText } from "./composerInlineTokenPaste";
 import { didComposerSelectionChangeVisibly } from "./composerSelection";
@@ -180,6 +189,12 @@ export interface ComposerPromptEditorProps {
   onPageScrollKeyUp?: (key: string) => void;
   onPageScrollRelease?: () => void;
   onCitationSubmitAndSend?: () => void;
+  /**
+   * Called when a new controlled value replaces the document while a citation
+   * comment is unsaved. The comment is saved into `replacedValue`, giving
+   * `savedValue`, so the owner of `replacedValue` can keep it.
+   */
+  onCitationDraftSave?: (replacedValue: string, savedValue: string) => void;
   onPaste: React.ClipboardEventHandler<HTMLElement>;
   editorRef: React.RefObject<ComposerPromptEditorHandle | null>;
 }
@@ -204,8 +219,14 @@ type OpenCitationComment = {
 const ComposerCitationCommentContext = createContext<{
   openComment: OpenCitationComment | null;
   onOpenChange: (citeKey: string, open: boolean) => void;
+  onDraftChange: (citeKey: string, comment: string) => void;
   onSubmitAndSend: () => void;
-}>({ openComment: null, onOpenChange: () => {}, onSubmitAndSend: () => {} });
+}>({
+  openComment: null,
+  onOpenChange: () => {},
+  onDraftChange: () => {},
+  onSubmitAndSend: () => {},
+});
 
 const RichComposerSkillsContext = createContext<ReadonlyArray<ServerProviderSkill>>([]);
 
@@ -472,6 +493,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
           },
           ...(removeOnCancel ? { onCancel: onRemove } : {}),
           onSave: onSaveComment,
+          onDraftChange: (comment) => commentContext.onDraftChange(citeKey, comment),
           onSaveAndSend: (comment) => {
             if (!onSaveComment(comment)) return false;
             commentContext.onSubmitAndSend();
@@ -843,6 +865,7 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
     onPageScrollKeyUp,
     onPageScrollRelease,
     onCitationSubmitAndSend,
+    onCitationDraftSave,
     onPaste,
     editorRef,
   } = props;
@@ -928,12 +951,15 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
   const hasAppliedControlledSelectionRef = useRef(false);
   const citationRequestRef = useRef<ComposerCitationCommentRequest | null>(null);
   const [openCitation, setOpenCitation] = useState<OpenCitationComment | null>(null);
+  // The open citation's unsaved comment, saved before a new value replaces the document.
+  const citationDraftRef = useRef<{ key: string; comment: string } | null>(null);
   const [isEmpty, setIsEmpty] = useState(value.length === 0);
 
   const citationCommentActions = useMemo(
     () => ({
       openComment: openCitation,
       onOpenChange: (nodeKey: string, open: boolean) => {
+        if (!open && citationDraftRef.current?.key === nodeKey) citationDraftRef.current = null;
         setOpenCitation((current) => {
           if (open) {
             if (current?.key === nodeKey) return current;
@@ -942,6 +968,9 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
           return current?.key === nodeKey ? null : current;
         });
       },
+      onDraftChange: (nodeKey: string, comment: string) => {
+        citationDraftRef.current = { key: nodeKey, comment };
+      },
       onSubmitAndSend: onCitationSubmitAndSend ?? (() => {}),
     }),
     [onCitationSubmitAndSend, openCitation],
@@ -1529,6 +1558,39 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
     return next;
   }, [clampEditorCursor, editor, literalText]);
 
+  // A question or approval swaps the controlled value while a citation comment
+  // is open. Its node view unmounts with the old document, so save the comment
+  // the way dismissing the popover would, and hand the saved value back.
+  const saveCitationDraft = useEffectEvent((current: TiptapEditor, replacedValue: string) => {
+    const draft = citationDraftRef.current;
+    citationDraftRef.current = null;
+    if (!draft || !onCitationDraftSave) return;
+    const [target] = findChildren(
+      current.state.doc,
+      (node) => node.type.name === "composer-citation" && node.attrs.citeKey === draft.key,
+    );
+    if (!target) return;
+    const { node, pos } = target;
+    const citation = node.attrs.citation as AssistantCitation;
+    const dismissal = resolveAssistantCitationCommentDismissal({
+      reason: "none",
+      draft: draft.comment,
+      savedComment: citation.comment,
+    });
+    if (dismissal.kind !== "commit") return;
+    const next = withAssistantCitationComment(citation, dismissal.comment);
+    current.view.dispatch(
+      current.state.tr
+        .setNodeMarkup(pos, undefined, {
+          ...node.attrs,
+          citation: next,
+          source: serializeAssistantCitation(next),
+        })
+        .setMeta("addToHistory", false),
+    );
+    onCitationDraftSave(replacedValue, serializeEditorDoc(current.state.doc).value);
+  });
+
   // Controlled value/cursor from the store (history recall, chip insertion…).
   useLayoutEffect(() => {
     if (!editor) return;
@@ -1563,6 +1625,7 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
     const pendingCitation =
       citationRequestRef.current?.value === value ? citationRequestRef.current : null;
     if (previousSnapshot.value !== value) {
+      saveCitationDraft(editor, previousSnapshot.value);
       editor.commands.setContent(
         buildDocJson(value, skillLabelFor, { styling: richText, literalText }),
         {

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

@@ -47,6 +47,8 @@ export function AssistantCitationChip({
     onCancel?: () => void;
     onSave: (comment: string) => boolean;
     onSaveAndSend?: (comment: string) => boolean;
+    /** Reports each keystroke of the unsaved comment. */
+    onDraftChange?: (comment: string) => void;
     /** Returns focus to the host editor when the popover closes instead of to the pencil trigger. */
     onRestoreFocus?: () => void;
   };
@@ -211,6 +213,7 @@ export function AssistantCitationChip({
                 inputRef={commentInputRef}
                 onDraftChange={(comment) => {
                   draftCommentRef.current = comment;
+                  commentEditor.onDraftChange?.(comment);
                 }}
                 onSubmit={(comment) => {
                   if (!commentEditor.onSave(comment)) return false;

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

@@ -4358,6 +4358,16 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps)
       }),
     );
   }, [phase, settings.followUpBehavior, submitComposer]);
+  // A question or approval took the editor while a citation comment was
+  // unsaved. Keep the comment in the prompt that comes back afterward.
+  const saveCitationDraftToPrompt = useCallback(
+    (replacedValue: string, savedValue: string) => {
+      if (replacedValue !== prompt) return;
+      if (promptRef.current === replacedValue) promptRef.current = savedValue;
+      setPrompt(savedValue);
+    },
+    [prompt, promptRef, setPrompt],
+  );
   const compactThreadContext = useCallback(() => {
     if (
       !readEnvironmentScope(environmentId, AuthOrchestrationOperateScope) ||
@@ -7591,6 +7601,7 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps)
                     onPageScrollKeyUp={onPageScrollKeyUp}
                     onPageScrollRelease={onPageScrollRelease}
                     onCitationSubmitAndSend={submitCitationAndSend}
+                    onCitationDraftSave={saveCitationDraftToPrompt}
                     onPaste={onComposerPaste}
                     placeholder={
                       isComposerApprovalState