fix(web): keep citation comment when popover is dismissed

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

@@ -0,0 +1,156 @@
+import {
+  ASSISTANT_CITATION_MAX_COMMENT_LENGTH,
+  EnvironmentId,
+  MessageId,
+  ThreadId,
+} from "@t3tools/contracts";
+import { act, useState, type ReactNode } from "react";
+import { create, type ReactTestRenderer } from "react-test-renderer";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test";
+
+import type { AssistantCitationSourceAnchor } from "~/lib/assistantTextSelection";
+
+const mocks = vi.hoisted(() => ({ observeSource: vi.fn(), dispose: vi.fn() }));
+vi.mock("./AssistantCitationSource", () => ({
+  observeAssistantCitationCommentSource: mocks.observeSource,
+}));
+vi.mock("@tanstack/react-router", () => ({
+  useNavigate: () => vi.fn(),
+  Link: ({ children }: { children: ReactNode }) => <a>{children}</a>,
+}));
+// Keep the real chip/editor lifecycle while replacing DOM positioning and floating layers.
+vi.mock("../ui/tooltip", () => ({
+  Tooltip: ({ children }: { children: ReactNode }) => <>{children}</>,
+  TooltipTrigger: ({ render }: { render: ReactNode }) => render,
+  TooltipPopup: ({ children }: { children: ReactNode }) => <>{children}</>,
+}));
+vi.mock("../ui/popover", () => ({
+  Popover: ({ children }: { children: ReactNode }) => <>{children}</>,
+  PopoverTrigger: ({ children }: { children: ReactNode }) => <button>{children}</button>,
+  PopoverPopup: ({ children }: { children: ReactNode }) => <>{children}</>,
+}));
+vi.mock("../ui/button", () => ({
+  Button: (props: React.ComponentProps<"button">) => <button {...props} />,
+}));
+
+import { PopoverPopup } from "../ui/popover";
+import { AssistantCitationChip } from "./AssistantCitationChip";
+
+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: "",
+};
+// The observer owns DOM access; these identities let us check which anchor survives.
+const sourceAnchor = {
+  source: {},
+  range: {},
+  viewport: {},
+} as AssistantCitationSourceAnchor;
+
+let renderer: ReactTestRenderer;
+
+function mount(onSave = vi.fn(() => true)) {
+  function Composer() {
+    const [open, setOpen] = useState(true);
+    return (
+      <AssistantCitationChip
+        citation={citation}
+        commentEditor={{ open, sourceAnchor, onOpenChange: setOpen, onSave }}
+      />
+    );
+  }
+  act(() => {
+    renderer = create(<Composer />);
+  });
+  return onSave;
+}
+
+function typeComment(value: string) {
+  act(() => renderer.root.findByType("textarea").props.onChange({ currentTarget: { value } }));
+}
+
+function removeSource() {
+  const { onUnavailable } = mocks.observeSource.mock.lastCall![0];
+  act(() => onUnavailable());
+}
+
+function clickButton(label: string) {
+  act(() =>
+    renderer.root
+      .findAllByType("button")
+      .find((button) => button.props.children === label)!
+      .props.onClick(),
+  );
+}
+
+beforeEach(() => {
+  vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+  mocks.observeSource.mockReset().mockReturnValue(mocks.dispose);
+  mocks.dispose.mockClear();
+});
+
+afterEach(() => {
+  act(() => renderer?.unmount());
+  vi.unstubAllGlobals();
+});
+
+describe("citation comment source disappearance", () => {
+  it("preserves an over-length draft at the composer until it can be shortened and saved", () => {
+    const onSave = mount();
+    const draft = "x".repeat(ASSISTANT_CITATION_MAX_COMMENT_LENGTH + 1);
+    typeComment(draft);
+    removeSource();
+
+    expect(renderer.root.findByType("textarea").props.value).toBe(draft);
+    expect(renderer.root.findByType("textarea").props["aria-invalid"]).toBe(true);
+    expect(renderer.root.findByType(PopoverPopup).props.anchor).toBeUndefined();
+    expect(renderer.root.findByType(PopoverPopup).props.side).toBe("top");
+    expect(onSave).not.toHaveBeenCalled();
+    expect(mocks.dispose).toHaveBeenCalled();
+    expect(mocks.observeSource).toHaveBeenCalledTimes(1);
+
+    typeComment("shortened comment");
+    clickButton("Save");
+    expect(onSave).toHaveBeenCalledWith("shortened comment");
+    expect(renderer.root.findAllByType("textarea")).toHaveLength(0);
+  });
+
+  it("preserves a rejected save and allows a later retry", () => {
+    const onSave = mount(vi.fn(() => false));
+    typeComment("keep this draft");
+    removeSource();
+
+    expect(onSave).toHaveBeenCalledWith("keep this draft");
+    expect(renderer.root.findByType("textarea").props.value).toBe("keep this draft");
+    expect(renderer.root.findByType(PopoverPopup).props.anchor).toBeUndefined();
+    onSave.mockReturnValue(true);
+    clickButton("Save");
+    expect(onSave).toHaveBeenLastCalledWith("keep this draft");
+    expect(renderer.root.findAllByType("textarea")).toHaveLength(0);
+  });
+
+  it("still allows explicit cancellation after source disappearance", () => {
+    const onSave = mount(vi.fn(() => false));
+    typeComment("discard this draft");
+    removeSource();
+    onSave.mockClear();
+    clickButton("Cancel");
+    expect(onSave).not.toHaveBeenCalled();
+    expect(renderer.root.findAllByType("textarea")).toHaveLength(0);
+  });
+
+  it("saves and closes when the source disappears with a valid draft", () => {
+    const onSave = mount();
+    typeComment("saved comment");
+    removeSource();
+    expect(onSave).toHaveBeenCalledWith("saved comment");
+    expect(renderer.root.findAllByType("textarea")).toHaveLength(0);
+  });
+});

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

@@ -2,7 +2,13 @@ import type { AssistantCitation } from "@t3tools/contracts";
 import { serializeAssistantCitation } from "@t3tools/shared/assistantCitations";
 import { Link, useNavigate } from "@tanstack/react-router";
 import { PencilIcon, QuoteIcon } from "lucide-react";
-import { useEffect, useEffectEvent, useRef, type MouseEvent as ReactMouseEvent } from "react";
+import {
+  useEffect,
+  useEffectEvent,
+  useRef,
+  useState,
+  type MouseEvent as ReactMouseEvent,
+} from "react";
 import {
   findAssistantCitationSourceAnchor,
   type AssistantCitationSourceAnchor,
@@ -23,6 +29,7 @@ import {
 import { Tooltip, TooltipPopup, TooltipTrigger } from "../ui/tooltip";
 import { Popover, PopoverPopup, PopoverTrigger } from "../ui/popover";
 import { AssistantCitationCommentEditor } from "./AssistantCitationCommentEditor";
+import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal";
 import { observeAssistantCitationCommentSource } from "./AssistantCitationSource";
 import { composerFloatingLayerProps } from "./composerEventScope";
 
@@ -49,29 +56,51 @@ export function AssistantCitationChip({
 }) {
   const navigate = useNavigate();
   const commentInputRef = useRef<HTMLTextAreaElement>(null);
+  const draftCommentRef = useRef<string | null>(null);
+  const [unavailableSourceAnchor, setUnavailableSourceAnchor] =
+    useState<AssistantCitationSourceAnchor | null>(null);
   const commentOpen = commentEditor?.open ?? false;
   const sourceAnchor = commentEditor?.sourceAnchor;
+  const activeSourceAnchor = sourceAnchor === unavailableSourceAnchor ? undefined : sourceAnchor;
+  useEffect(() => {
+    if (!commentOpen) draftCommentRef.current = null;
+  }, [commentOpen]);
+  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";
+  };
   const onSourceUnavailable = useEffectEvent(() => {
-    if (sourceAnchor) commentEditor?.onOpenChange(false);
+    if (!sourceAnchor) return;
+    if (settleDraftOnClose("none")) {
+      commentEditor?.onOpenChange(false);
+    } else {
+      // Keep the draft mounted, positioned at the composer trigger instead of a detached range.
+      setUnavailableSourceAnchor(sourceAnchor);
+    }
   });
   useEffect(() => {
-    if (!commentOpen) return;
+    if (!commentOpen || sourceAnchor === unavailableSourceAnchor) return;
     const anchor = sourceAnchor ?? findAssistantCitationSourceAnchor(document, citation);
     if (!anchor) return;
     return observeAssistantCitationCommentSource({
       anchor,
       citation,
       onUnavailable: onSourceUnavailable,
     });
-  }, [citation, commentOpen, sourceAnchor]);
+  }, [citation, commentOpen, sourceAnchor, unavailableSourceAnchor]);
   // A multi-line selection's bounding box spans the full message width; anchor
   // the bubble to the selection's last line, where the pointer released.
-  const popupAnchor = sourceAnchor
+  const popupAnchor = activeSourceAnchor
     ? {
-        contextElement: sourceAnchor.source,
+        contextElement: activeSourceAnchor.source,
         getBoundingClientRect: () => {
-          const rects = sourceAnchor.range.getClientRects();
-          return rects.item(rects.length - 1) ?? sourceAnchor.range.getBoundingClientRect();
+          const rects = activeSourceAnchor.range.getClientRects();
+          return rects.item(rects.length - 1) ?? activeSourceAnchor.range.getBoundingClientRect();
         },
       }
     : undefined;
@@ -130,7 +159,16 @@ export function AssistantCitationChip({
         </Tooltip>
       )}
       {commentEditor ? (
-        <Popover open={commentEditor.open} onOpenChange={commentEditor.onOpenChange}>
+        <Popover
+          open={commentEditor.open}
+          onOpenChange={(open, eventDetails) => {
+            if (!open && !settleDraftOnClose(eventDetails.reason)) {
+              eventDetails.cancel();
+              return;
+            }
+            commentEditor.onOpenChange(open);
+          }}
+        >
           <PopoverTrigger
             aria-label={citation.comment ? "Edit citation comment" : "Add comment to citation"}
             className={CITATION_ACTION_BUTTON_CLASS_NAME}
@@ -140,7 +178,7 @@ export function AssistantCitationChip({
           {commentEditor.open ? (
             <PopoverPopup
               {...composerFloatingLayerProps}
-              side={sourceAnchor ? "bottom" : "top"}
+              side={activeSourceAnchor ? "bottom" : "top"}
               align="end"
               anchor={popupAnchor}
               initialFocus={() => {
@@ -156,6 +194,9 @@ export function AssistantCitationChip({
                 key={serializeAssistantCitation(citation)}
                 citation={citation}
                 inputRef={commentInputRef}
+                onDraftChange={(comment) => {
+                  draftCommentRef.current = comment;
+                }}
                 onSubmit={(comment) => {
                   if (!commentEditor.onSave(comment)) return false;
                   commentEditor.onOpenChange(false);

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

@@ -9,12 +9,14 @@ export function AssistantCitationCommentEditor({
   onSubmit,
   onSubmitAndSend,
   onCancel,
+  onDraftChange,
 }: {
   citation: AssistantCitation;
   inputRef?: Ref<HTMLTextAreaElement>;
   onSubmit: (comment: string) => boolean;
   onSubmitAndSend?: (comment: string) => boolean;
   onCancel: () => void;
+  onDraftChange?: (comment: string) => void;
 }) {
   const [comment, setComment] = useState(citation.comment ?? "");
   const commentTooLong = comment.length > ASSISTANT_CITATION_MAX_COMMENT_LENGTH;
@@ -51,7 +53,10 @@ export function AssistantCitationCommentEditor({
         rows={2}
         className="field-sizing-content block max-h-40 min-h-16 w-full resize-none bg-transparent px-1 py-1.5 text-base outline-none placeholder:text-muted-foreground sm:text-sm"
         value={comment}
-        onChange={(event) => setComment(event.currentTarget.value)}
+        onChange={(event) => {
+          setComment(event.currentTarget.value);
+          onDraftChange?.(event.currentTarget.value);
+        }}
         onKeyDown={(event) => {
           if (
             event.key === "Enter" &&

apps/web/src/components/chat/assistantCitationCommentDismissal.test.ts

@@ -0,0 +1,80 @@
+import { ASSISTANT_CITATION_MAX_COMMENT_LENGTH } from "@t3tools/contracts";
+import { describe, expect, it } from "vite-plus/test";
+
+import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal";
+
+describe("resolveAssistantCitationCommentDismissal", () => {
+  it("commits typed text when the popover is dismissed by clicking away", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "outside-press",
+        draft: "needs a retry",
+        savedComment: undefined,
+      }),
+    ).toEqual({ kind: "commit", comment: "needs a retry" });
+  });
+
+  it("commits an edited comment when focus leaves the popover", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "focus-out",
+        draft: "second thought",
+        savedComment: "first thought",
+      }),
+    ).toEqual({ kind: "commit", comment: "second thought" });
+  });
+
+  it("closes without saving when nothing changed", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "outside-press",
+        draft: null,
+        savedComment: "kept",
+      }),
+    ).toEqual({ kind: "close" });
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "outside-press",
+        draft: "  kept ",
+        savedComment: "kept",
+      }),
+    ).toEqual({ kind: "close" });
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "outside-press",
+        draft: "kept",
+        savedComment: " kept ",
+      }),
+    ).toEqual({ kind: "close" });
+  });
+
+  it("clears a comment when the draft was emptied", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "trigger-press",
+        draft: "",
+        savedComment: "old",
+      }),
+    ).toEqual({ kind: "commit", comment: "" });
+  });
+
+  it("keeps Escape as an explicit discard", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "escape-key",
+        draft: "unsaved",
+        savedComment: undefined,
+      }),
+    ).toEqual({ kind: "close" });
+  });
+
+  it("keeps the popover open instead of dropping an over-length draft", () => {
+    expect(
+      resolveAssistantCitationCommentDismissal({
+        reason: "outside-press",
+        draft: "x".repeat(ASSISTANT_CITATION_MAX_COMMENT_LENGTH + 1),
+        savedComment: undefined,
+      }),
+    ).toEqual({ kind: "keep-open" });
+  });
+});