fix(web): discard new citations when cancelled

apps/web/src/components/ComposerCitationNode.tsx

@@ -42,13 +42,15 @@ export type ComposerCitationCommentRequest = {
 export type ComposerCitationCommentTarget = {
   nodeKey: NodeKey;
   sourceAnchor?: AssistantCitationSourceAnchor;
+  insertion?: ComposerCitationCommentRequest;
 };
 
 export const ComposerCitationCommentContext = createContext<{
   openComment: ComposerCitationCommentTarget | null;
   onOpenChange: (nodeKey: NodeKey, open: boolean) => void;
+  onCancel: (nodeKey: NodeKey) => void;
   onSubmitAndSend: () => void;
-}>({ openComment: null, onOpenChange: () => {}, onSubmitAndSend: () => {} });
+}>({ openComment: null, onOpenChange: () => {}, onCancel: () => {}, onSubmitAndSend: () => {} });
 
 /** Consume a cite action once its controlled prompt has been committed to the editor. */
 export function $consumeComposerCitationCommentRequest(requestRef: {
@@ -68,7 +70,7 @@ export function $consumeComposerCitationCommentRequest(requestRef: {
   let offset = 0;
   for (const node of paragraph.getChildren()) {
     if (offset === request.citationStart && node instanceof ComposerCitationNode) {
-      return { nodeKey: node.getKey(), sourceAnchor: request.sourceAnchor };
+      return { nodeKey: node.getKey(), sourceAnchor: request.sourceAnchor, insertion: request };
     }
     offset += node.getTextContentSize();
   }
@@ -127,6 +129,7 @@ function ComposerCitationDecorator(props: { citation: AssistantCitation; nodeKey
             if (open && !editor.isEditable()) return;
             commentContext.onOpenChange(props.nodeKey, open);
           },
+          onCancel: () => commentContext.onCancel(props.nodeKey),
           onSave: onSaveComment,
           onSaveAndSend: (comment) => {
             if (!onSaveComment(comment)) return false;

apps/web/src/components/ComposerPromptEditor.serialization.test.tsx

@@ -1,11 +1,22 @@
 import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
-import { $copyNode, $getRoot, $isElementNode, PASTE_COMMAND, type LexicalEditor } from "lexical";
-import { act, createRef } from "react";
+import { EnvironmentId, MessageId, ThreadId, type AssistantCitation } from "@t3tools/contracts";
+import { serializeAssistantCitation } from "@t3tools/shared/assistantCitations";
+import {
+  $copyNode,
+  $createTextNode,
+  $getRoot,
+  $isElementNode,
+  PASTE_COMMAND,
+  type LexicalEditor,
+} from "lexical";
+import { act, createRef, use, type ContextType } from "react";
 import { create, type ReactTestRenderer } from "react-test-renderer";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test";
 
 import { collapseExpandedComposerCursor } from "../composer-logic";
 import { ComposerPromptEditor, type ComposerPromptEditorHandle } from "./ComposerPromptEditor";
+import { ComposerCitationCommentContext, ComposerCitationNode } from "./ComposerCitationNode";
+import type { AssistantCitationSourceAnchor } from "~/lib/assistantTextSelection";
 
 vi.mock("./chat/FileTagChip", () => ({
   FILE_TAG_CHIP_CLASS_NAME: "",
@@ -17,11 +28,13 @@ vi.mock("./chat/ComposerPendingTerminalContexts", () => ({
 vi.mock("./chat/AssistantCitationChip", () => ({ AssistantCitationChip: () => null }));
 
 let lexicalEditor: LexicalEditor;
+let citationComments: ContextType<typeof ComposerCitationCommentContext>;
 // Keep the real composer, registered nodes, updates, and snapshot API. Only the
 // DOM view is omitted so Lexical runs headlessly in this component test.
 vi.mock("@lexical/react/LexicalPlainTextPlugin", () => ({
   PlainTextPlugin: function HeadlessEditor() {
     [lexicalEditor] = useLexicalComposerContext();
+    citationComments = use(ComposerCitationCommentContext);
     return null;
   },
 }));
@@ -182,3 +195,86 @@ describe("composer mention serialization", () => {
     expect(lexicalEditor.getEditorState().read(() => $firstMention().isInline())).toBe(true);
   });
 });
+
+describe("composer citation cancellation", () => {
+  const citation: AssistantCitation = {
+    version: 1,
+    environmentId: EnvironmentId.make("local"),
+    threadId: ThreadId.make("thread-1"),
+    messageId: MessageId.make("message-1"),
+    text: "Selected assistant text",
+    start: 0,
+    end: 23,
+    prefix: "",
+    suffix: "",
+  };
+  const source = serializeAssistantCitation(citation);
+  const savedSource = serializeAssistantCitation({ ...citation, comment: "Saved comment" });
+  const sourceAnchor: AssistantCitationSourceAnchor = {
+    source: { nodeType: 1 } as HTMLElement,
+    range: { collapsed: false } as Range,
+    viewport: { nodeType: 1 } as HTMLElement,
+  };
+
+  async function cite(previousValue: string) {
+    await renderPrompt(previousValue);
+    const prefix = previousValue ? `${previousValue} ` : "";
+    const value = `${prefix}${source} `;
+    await act(() => {
+      editorRef.current?.requestCitationComment({
+        previousValue,
+        value,
+        citationStart: prefix.length,
+        sourceAnchor,
+      });
+    });
+    await renderPrompt(value);
+    const target = citationComments.openComment;
+    expect(target).not.toBeNull();
+    return target!.nodeKey;
+  }
+
+  it.each(["", "Keep my draft.", `@README.md\n  雪 👋 ${savedSource}`])(
+    "restores the original draft when a new citation is cancelled: %s",
+    async (previousValue) => {
+      const nodeKey = await cite(previousValue);
+      await act(() => citationComments.onCancel(nodeKey));
+      expect(editorRef.current?.readSnapshot().value).toBe(previousValue);
+      expect(citationComments.openComment).toBeNull();
+    },
+  );
+
+  it("preserves edits made to the draft after inserting the citation", async () => {
+    const previousValue = `Keep ${savedSource}`;
+    const nodeKey = await cite(previousValue);
+    await act(() => {
+      lexicalEditor.update(
+        () => {
+          const paragraph = $getRoot().getFirstChildOrThrow();
+          if (!$isElementNode(paragraph)) throw new Error("Expected composer paragraph");
+          paragraph.append($createTextNode("Typed meanwhile"));
+        },
+        { discrete: true },
+      );
+    });
+    await act(() => citationComments.onCancel(nodeKey));
+    expect(editorRef.current?.readSnapshot().value).toBe(`${previousValue}  Typed meanwhile`);
+    expect(citationComments.openComment).toBeNull();
+  });
+
+  it("keeps a saved citation and its comment when cancelling a comment edit", async () => {
+    const value = `Keep ${savedSource} and this draft.`;
+    await renderPrompt(value);
+    const nodeKey = lexicalEditor.getEditorState().read(() => {
+      const paragraph = $getRoot().getFirstChildOrThrow();
+      if (!$isElementNode(paragraph)) throw new Error("Expected composer paragraph");
+      const node = paragraph.getChildren().find((child) => child instanceof ComposerCitationNode);
+      if (!node) throw new Error("Expected saved citation");
+      return node.getKey();
+    });
+    await act(() => citationComments.onOpenChange(nodeKey, true));
+    await act(() => citationComments.onCancel(nodeKey));
+    expect(editorRef.current?.readSnapshot().value).toBe(value);
+    expect(citationComments.openComment).toBeNull();
+  });
+});

apps/web/src/components/ComposerPromptEditor.tsx

@@ -1678,9 +1678,39 @@ function ComposerPromptEditorInner({
           open ? { nodeKey } : current?.nodeKey === nodeKey ? null : current,
         );
       },
+      onCancel: (nodeKey: NodeKey) => {
+        if (openCitationComment?.nodeKey !== nodeKey) return;
+        const insertion = openCitationComment.insertion;
+        if (insertion && editor.isEditable()) {
+          editor.update(
+            () => {
+              const node = $getNodeByKey(nodeKey);
+              if (!(node instanceof ComposerCitationNode) || !node.isAttached()) return;
+              // Restore the original draft, including spacing added by Cite. If
+              // the draft changed meanwhile, remove only the pending citation.
+              if ($getRoot().getTextContent() === insertion.value) {
+                $setComposerEditorPrompt(
+                  insertion.previousValue,
+                  terminalContexts,
+                  skillMetadataRef.current,
+                );
+                $setSelectionAtComposerOffset(
+                  collapseExpandedComposerCursor(insertion.previousValue, insertion.citationStart),
+                );
+              } else {
+                node.selectPrevious();
+                node.remove();
+              }
+            },
+            { discrete: true, tag: HISTORY_PUSH_TAG },
+          );
+          editor.getRootElement()?.focus({ preventScroll: true });
+        }
+        setOpenCitationComment(null);
+      },
       onSubmitAndSend: onCitationSubmitAndSend ?? (() => {}),
     }),
-    [onCitationSubmitAndSend, openCitationComment],
+    [editor, onCitationSubmitAndSend, openCitationComment, terminalContexts],
   );
   const terminalContextActions = useMemo(
     () => ({ onRemoveTerminalContext }),

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

@@ -41,6 +41,7 @@ export function AssistantCitationChip({
     open: boolean;
     sourceAnchor?: AssistantCitationSourceAnchor | undefined;
     onOpenChange: (open: boolean) => void;
+    onCancel: () => void;
     onSave: (comment: string) => boolean;
     onSaveAndSend?: (comment: string) => boolean;
   };
@@ -50,7 +51,7 @@ export function AssistantCitationChip({
   const commentOpen = commentEditor?.open ?? false;
   const sourceAnchor = commentEditor?.sourceAnchor;
   const onSourceUnavailable = useEffectEvent(() => {
-    if (sourceAnchor) commentEditor?.onOpenChange(false);
+    if (sourceAnchor) commentEditor?.onCancel();
   });
   useEffect(() => {
     if (!commentOpen) return;
@@ -128,7 +129,13 @@ export function AssistantCitationChip({
         </Tooltip>
       )}
       {commentEditor ? (
-        <Popover open={commentEditor.open} onOpenChange={commentEditor.onOpenChange}>
+        <Popover
+          open={commentEditor.open}
+          onOpenChange={(open) => {
+            if (open) commentEditor.onOpenChange(true);
+            else commentEditor.onCancel();
+          }}
+        >
           <PopoverTrigger
             aria-label={citation.comment ? "Edit citation comment" : "Add comment to citation"}
             className={CITATION_ACTION_BUTTON_CLASS_NAME}
@@ -168,7 +175,7 @@ export function AssistantCitationChip({
                       },
                     }
                   : {})}
-                onCancel={() => commentEditor.onOpenChange(false)}
+                onCancel={commentEditor.onCancel}
               />
             </PopoverPopup>
           ) : null}