diff --git a/apps/web/src/components/ComposerCitationNode.tsx b/apps/web/src/components/ComposerCitationNode.tsx index ffaae589343f..532334a481d6 100644 --- a/apps/web/src/components/ComposerCitationNode.tsx +++ b/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; diff --git a/apps/web/src/components/ComposerPromptEditor.serialization.test.tsx b/apps/web/src/components/ComposerPromptEditor.serialization.test.tsx index 980521084f8c..a64c708b53e2 100644 --- a/apps/web/src/components/ComposerPromptEditor.serialization.test.tsx +++ b/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; // 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(); + }); +}); diff --git a/apps/web/src/components/ComposerPromptEditor.tsx b/apps/web/src/components/ComposerPromptEditor.tsx index 6b49ef884734..23ec143c24a4 100644 --- a/apps/web/src/components/ComposerPromptEditor.tsx +++ b/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 }), diff --git a/apps/web/src/components/chat/AssistantCitationChip.tsx b/apps/web/src/components/chat/AssistantCitationChip.tsx index 7582bb19c5c6..fab71843f43e 100644 --- a/apps/web/src/components/chat/AssistantCitationChip.tsx +++ b/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({ )} {commentEditor ? ( - + { + if (open) commentEditor.onOpenChange(true); + else commentEditor.onCancel(); + }} + > commentEditor.onOpenChange(false)} + onCancel={commentEditor.onCancel} /> ) : null}