diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index 69a204ccd8a3..30887274ac0f 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -374,6 +374,8 @@ import { composerDraftHasUserContent, type ComposerFileAttachment, type ComposerImageAttachment, + composerTargetKey, + type ComposerThreadTarget, type DraftThreadEnvMode, useComposerDraftStore, DraftId, @@ -439,6 +441,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"; @@ -1826,7 +1833,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, @@ -9127,7 +9143,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, @@ -9141,6 +9157,7 @@ export default function ChatView(props: ChatViewProps) { }); if (!followUpSent) { promptRef.current = followUpPromptSnapshot; + restoreAssistantCitationCommentDrafts(followUpCitationDrafts); composerTerminalContextsRef.current = [...followUpTerminalContexts]; restorePlanFollowUpComposer({ snapshot: { @@ -9464,6 +9481,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; @@ -9484,7 +9502,7 @@ export default function ChatView(props: ChatViewProps) { "New thread", ); promptRef.current = ""; - clearComposerDraftContent(composerDraftTarget); + clearedCitationDrafts = clearComposerDraftContent(composerDraftTarget); composerRef.current?.resetCursorState(); clearedDraft = true; const clearedDraftSnapshot = useComposerDraftStore @@ -9653,6 +9671,7 @@ export default function ChatView(props: ChatViewProps) { setMultipleModelSelections(failedSelections); if (clearedDraft) { setComposerDraftPrompt(composerDraftTarget, messageTextForSend); + restoreAssistantCitationCommentDrafts(clearedCitationDrafts); addComposerDraftImages( composerDraftTarget, composerImagesSnapshot.map(cloneComposerImageForRetry), @@ -9794,7 +9813,7 @@ export default function ChatView(props: ChatViewProps) { ); } promptRef.current = ""; - clearComposerDraftContent(composerDraftTarget); + const sentCitationDrafts = clearComposerDraftContent(composerDraftTarget); composerRef.current?.resetCursorState(); let firstComposerImageName: string | null = null; @@ -10062,6 +10081,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; diff --git a/apps/web/src/components/ComposerPromptEditorTiptap.tsx b/apps/web/src/components/ComposerPromptEditorTiptap.tsx index 22002431097f..807485a1184e 100644 --- a/apps/web/src/components/ComposerPromptEditorTiptap.tsx +++ b/apps/web/src/components/ComposerPromptEditorTiptap.tsx @@ -93,6 +93,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"; @@ -176,6 +177,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; editorRef: React.RefObject; } @@ -201,7 +204,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>([]); @@ -370,6 +374,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 = []; + 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; @@ -459,6 +477,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) { commentContext.onOpenChange(citeKey, open); }, ...(removeOnCancel ? { onCancel: onRemove } : {}), + draftKey, onSave: onSaveComment, onSaveAndSend: (comment) => { if (!onSaveComment(comment)) return false; @@ -830,6 +849,7 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) { onPageScrollKeyUp, onPageScrollRelease, onCitationSubmitAndSend, + citationDraftScope, onPaste, editorRef, } = props; @@ -926,8 +946,9 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) { }); }, onSubmitAndSend: onCitationSubmitAndSend ?? (() => {}), + draftScope: citationDraftScope ?? "", }), - [onCitationSubmitAndSend, openCitation], + [citationDraftScope, onCitationSubmitAndSend, openCitation], ); const handleEditorChange = useCallback( diff --git a/apps/web/src/components/chat/AssistantCitationChip.test.tsx b/apps/web/src/components/chat/AssistantCitationChip.test.tsx index 91948cb587b9..acbd121067fc 100644 --- a/apps/web/src/components/chat/AssistantCitationChip.test.tsx +++ b/apps/web/src/components/chat/AssistantCitationChip.test.tsx @@ -51,14 +51,20 @@ 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 ( ); } @@ -97,6 +103,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(); diff --git a/apps/web/src/components/chat/AssistantCitationChip.tsx b/apps/web/src/components/chat/AssistantCitationChip.tsx index 0b4d6a962e61..e926925f0903 100644 --- a/apps/web/src/components/chat/AssistantCitationChip.tsx +++ b/apps/web/src/components/chat/AssistantCitationChip.tsx @@ -27,6 +27,11 @@ import { Button } from "../ui/button"; import { Popover, PopoverClose, PopoverPopup, PopoverTrigger } from "../ui/popover"; import { getVirtualizedScrollFadeClassName } from "../ui/scroll-area"; import { AssistantCitationCommentEditor } from "./AssistantCitationCommentEditor"; +import { + clearAssistantCitationCommentDraft, + readAssistantCitationCommentDraft, + writeAssistantCitationCommentDraft, +} from "./assistantCitationCommentDrafts"; import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal"; import { observeAssistantCitationCommentSource } from "./AssistantCitationSource"; import { composerFloatingLayerProps } from "./composerEventScope"; @@ -47,6 +52,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(); @@ -58,17 +65,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; @@ -204,14 +219,17 @@ export function AssistantCitationChip({ onPointerDown={(event) => event.stopPropagation()} > { draftCommentRef.current = comment; + writeAssistantCitationCommentDraft(draftKey, comment); }} onSubmit={(comment) => { if (!commentEditor.onSave(comment)) return false; + clearAssistantCitationCommentDraft(draftKey); commentEditor.onOpenChange(false); return true; }} @@ -219,12 +237,14 @@ export function AssistantCitationChip({ ? { 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 { diff --git a/apps/web/src/components/chat/AssistantCitationCommentEditor.tsx b/apps/web/src/components/chat/AssistantCitationCommentEditor.tsx index d9fd60645618..ec435dd02aa7 100644 --- a/apps/web/src/components/chat/AssistantCitationCommentEditor.tsx +++ b/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; 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); diff --git a/apps/web/src/components/chat/ChatComposer.tsx b/apps/web/src/components/chat/ChatComposer.tsx index 82a34a840c2f..c68fe18a8561 100644 --- a/apps/web/src/components/chat/ChatComposer.tsx +++ b/apps/web/src/components/chat/ChatComposer.tsx @@ -327,6 +327,7 @@ import { getComposerSubmissionValidationMessage, submitComposerDraft, } from "./composerSubmission"; +import { commitAssistantCitationCommentDrafts } from "./assistantCitationCommentDrafts"; import { ComposerPromptLengthValidation } from "./ComposerPromptLengthValidation"; import { PierreEntryIcon } from "./PierreEntryIcon"; import { pendingDraftWork } from "./pendingDraftWork"; @@ -3041,6 +3042,22 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps) [composerDraftTarget, setComposerDraftPrompt], ); + // A question or an approval borrows the prompt editor, which unmounts the + // citation chips without dismissing an open comment. The unsaved comments are + // written onto their citations first, as a dismissal would have, so the chips + // come back with them and Send includes them. + const editorShowsPrompt = !isComposerApprovalState && !activePendingProgress; + useEffect(() => { + if (editorShowsPrompt) return; + const committed = commitAssistantCitationCommentDrafts( + promptRef.current, + composerDraftTargetKey, + ); + if (committed === promptRef.current) return; + promptRef.current = committed; + setPrompt(committed); + }, [composerDraftTargetKey, editorShowsPrompt, promptRef, setPrompt]); + const addComposerImage = useCallback( (image: ComposerImageAttachment) => addComposerDraftImages(attachmentDraftTarget, [image]), [attachmentDraftTarget, addComposerDraftImages], @@ -7478,6 +7495,7 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps) onPageScrollKeyUp={onPageScrollKeyUp} onPageScrollRelease={onPageScrollRelease} onCitationSubmitAndSend={submitCitationAndSend} + citationDraftScope={composerTargetKey(composerDraftTarget)} onPaste={onComposerPaste} placeholder={ isComposerApprovalState diff --git a/apps/web/src/components/chat/assistantCitationCommentDrafts.test.ts b/apps/web/src/components/chat/assistantCitationCommentDrafts.test.ts new file mode 100644 index 000000000000..1a9f7bd3bc55 --- /dev/null +++ b/apps/web/src/components/chat/assistantCitationCommentDrafts.test.ts @@ -0,0 +1,228 @@ +import { + ASSISTANT_CITATION_MAX_COMMENT_LENGTH, + EnvironmentId, + MessageId, + ThreadId, + type AssistantCitation, +} from "@t3tools/contracts"; +import { + collectAssistantCitations, + serializeAssistantCitation, +} from "@t3tools/shared/assistantCitations"; +import { describe, expect, it } from "vite-plus/test"; + +import { + assistantCitationDraftKey, + commitAssistantCitationCommentDrafts, + readAssistantCitationCommentDraft, + restoreAssistantCitationCommentDrafts, + takeAssistantCitationCommentDraftsForComposer, + writeAssistantCitationCommentDraft, +} from "./assistantCitationCommentDrafts"; + +const citation: AssistantCitation = { + version: 1, + environmentId: EnvironmentId.make("environment"), + threadId: ThreadId.make("thread"), + messageId: MessageId.make("source"), + text: "hello", + start: 0, + end: 5, + prefix: "", + suffix: "", +}; +const other: AssistantCitation = { ...citation, text: "world", start: 6, end: 11 }; + +describe("assistantCitationDraftKey", () => { + it("is the same for the same citation whatever nodes come before it", () => { + expect(assistantCitationDraftKey(citation, [])).toBe( + assistantCitationDraftKey(citation, [other]), + ); + }); + + it("tells identical citations apart by their order in the document", () => { + const first = assistantCitationDraftKey(citation, []); + const second = assistantCitationDraftKey(citation, [citation]); + const third = assistantCitationDraftKey(citation, [other, citation, citation]); + expect(new Set([first, second, third]).size).toBe(3); + // The same prompt rebuilt in the same order yields the same keys again. + expect(assistantCitationDraftKey(citation, [citation])).toBe(second); + }); + + it("keeps the same citation apart across composers", () => { + expect(assistantCitationDraftKey(citation, [], "thread-a")).not.toBe( + assistantCitationDraftKey(citation, [], "thread-b"), + ); + }); +}); + +describe("commitAssistantCitationCommentDrafts", () => { + it("writes a trimmed draft onto its citation, preserves surrounding text, and removes the draft", () => { + const scope = "commit-new-comment"; + const key = assistantCitationDraftKey(citation, [], scope); + const before = "Before the quote.\n\n"; + const after = "\nAfter the quote. "; + const prompt = `${before}${serializeAssistantCitation(citation)}${after}`; + writeAssistantCitationCommentDraft(key, " \n Please explain this. \t "); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + const expected = { ...citation, comment: "Please explain this." }; + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([expected]); + expect(committed).toBe(`${before}${serializeAssistantCitation(expected)}${after}`); + expect(readAssistantCitationCommentDraft(key)).toBeNull(); + }); + + it("replaces a saved comment with a different draft and removes the draft", () => { + const scope = "commit-replacement-comment"; + const saved = { ...citation, comment: "Old comment" }; + const key = assistantCitationDraftKey(saved, [], scope); + const prompt = serializeAssistantCitation(saved); + writeAssistantCitationCommentDraft(key, "New comment"); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([ + { ...citation, comment: "New comment" }, + ]); + expect(readAssistantCitationCommentDraft(key)).toBeNull(); + }); + + it("leaves the prompt unchanged and removes a draft equal to the saved comment after trimming", () => { + const scope = "commit-unchanged-comment"; + const saved = { ...citation, comment: "Saved comment" }; + const key = assistantCitationDraftKey(saved, [], scope); + const prompt = serializeAssistantCitation(saved); + writeAssistantCitationCommentDraft(key, " \n Saved comment \t "); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(committed).toBe(prompt); + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([saved]); + expect(readAssistantCitationCommentDraft(key)).toBeNull(); + }); + + it("leaves the prompt unchanged and retains a draft over the comment length limit", () => { + const scope = "commit-over-limit-comment"; + const key = assistantCitationDraftKey(citation, [], scope); + const prompt = serializeAssistantCitation(citation); + const draft = "x".repeat(ASSISTANT_CITATION_MAX_COMMENT_LENGTH + 1); + writeAssistantCitationCommentDraft(key, draft); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(committed).toBe(prompt); + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([citation]); + expect(readAssistantCitationCommentDraft(key)).toBe(draft); + }); + + it("commits different drafts for identical citations using their original ordinals", () => { + const scope = "commit-duplicate-citations"; + const first = assistantCitationDraftKey(citation, [], scope); + const second = assistantCitationDraftKey(citation, [citation], scope); + const link = serializeAssistantCitation(citation); + const prompt = `Before ${link}\nBetween ${link}\nAfter`; + writeAssistantCitationCommentDraft(first, "First comment"); + writeAssistantCitationCommentDraft(second, "Second comment"); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + const expectedFirst = { ...citation, comment: "First comment" }; + const expectedSecond = { ...citation, comment: "Second comment" }; + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([ + expectedFirst, + expectedSecond, + ]); + expect(committed).toBe( + `Before ${serializeAssistantCitation(expectedFirst)}\nBetween ${serializeAssistantCitation(expectedSecond)}\nAfter`, + ); + expect(readAssistantCitationCommentDraft(first)).toBeNull(); + expect(readAssistantCitationCommentDraft(second)).toBeNull(); + }); + + it("changes only the second of two different citations when only it has a draft", () => { + const scope = "commit-second-citation"; + const first = assistantCitationDraftKey(citation, [], scope); + const second = assistantCitationDraftKey(other, [citation], scope); + const firstLink = serializeAssistantCitation(citation); + const prompt = `${firstLink}\nBetween\n${serializeAssistantCitation(other)}`; + writeAssistantCitationCommentDraft(second, "Comment on world"); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + const expectedSecond = { ...other, comment: "Comment on world" }; + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([ + citation, + expectedSecond, + ]); + expect(committed).toBe(`${firstLink}\nBetween\n${serializeAssistantCitation(expectedSecond)}`); + expect(readAssistantCitationCommentDraft(first)).toBeNull(); + expect(readAssistantCitationCommentDraft(second)).toBeNull(); + }); + + it("neither applies nor removes drafts belonging to another scope", () => { + const scope = "commit-scope-isolation"; + const elsewhere = assistantCitationDraftKey(citation, [], "commit-scope-isolation-other"); + const prompt = serializeAssistantCitation(citation); + writeAssistantCitationCommentDraft(elsewhere, "Another composer's comment"); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(committed).toBe(prompt); + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([citation]); + expect(readAssistantCitationCommentDraft(elsewhere)).toBe("Another composer's comment"); + }); + + it("returns the same prompt when its citations have no drafts", () => { + const scope = "commit-no-drafts"; + const first = assistantCitationDraftKey(citation, [], scope); + const second = assistantCitationDraftKey(other, [citation], scope); + const prompt = `Before ${serializeAssistantCitation(citation)}\n${serializeAssistantCitation(other)} After`; + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(committed).toBe(prompt); + expect(collectAssistantCitations(committed).map((entry) => entry.citation)).toEqual([ + citation, + other, + ]); + expect(readAssistantCitationCommentDraft(first)).toBeNull(); + expect(readAssistantCitationCommentDraft(second)).toBeNull(); + }); + + it("returns the same prompt without citations and retains its scope's unrelated draft", () => { + const scope = "commit-no-citations"; + const key = assistantCitationDraftKey(citation, [], scope); + const prompt = "Plain prompt\n[ordinary link](https://example.com) "; + writeAssistantCitationCommentDraft(key, "Keep this draft"); + + const committed = commitAssistantCitationCommentDrafts(prompt, scope); + + expect(committed).toBe(prompt); + expect(collectAssistantCitations(committed)).toEqual([]); + expect(readAssistantCitationCommentDraft(key)).toBe("Keep this draft"); + }); +}); + +describe("takeAssistantCitationCommentDraftsForComposer", () => { + it("takes the sent composer's drafts, keeps every other composer's, and restores what it took", () => { + const sent = assistantCitationDraftKey(citation, [], "thread-a"); + const sentDuplicate = assistantCitationDraftKey(citation, [citation], "thread-a"); + const elsewhere = assistantCitationDraftKey(citation, [], "thread-b"); + writeAssistantCitationCommentDraft(sent, "first"); + writeAssistantCitationCommentDraft(sentDuplicate, "second"); + writeAssistantCitationCommentDraft(elsewhere, "other"); + + const taken = takeAssistantCitationCommentDraftsForComposer("thread-a"); + + expect(readAssistantCitationCommentDraft(sent)).toBeNull(); + expect(readAssistantCitationCommentDraft(sentDuplicate)).toBeNull(); + expect(readAssistantCitationCommentDraft(elsewhere)).toBe("other"); + + // A failed send gives the prompt back, and its drafts with it. + restoreAssistantCitationCommentDrafts(taken); + + expect(readAssistantCitationCommentDraft(sent)).toBe("first"); + expect(readAssistantCitationCommentDraft(sentDuplicate)).toBe("second"); + }); +}); diff --git a/apps/web/src/components/chat/assistantCitationCommentDrafts.ts b/apps/web/src/components/chat/assistantCitationCommentDrafts.ts new file mode 100644 index 000000000000..c4077abbd981 --- /dev/null +++ b/apps/web/src/components/chat/assistantCitationCommentDrafts.ts @@ -0,0 +1,122 @@ +import type { AssistantCitation } from "@t3tools/contracts"; +import { + collectAssistantCitations, + serializeAssistantCitation, + withAssistantCitationComment, +} from "@t3tools/shared/assistantCitations"; + +import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal"; + +/** + * Unsaved citation comments, keyed by the serialized citation. The comment + * editor lives inside the composer's citation node view, and the editor + * replaces its whole document when a provider question borrows it for the + * answer, which unmounts the node view without any dismissal. The draft is + * kept here across that, and dropped once it is saved or discarded. + */ +const drafts = new Map(); + +/** + * The key a citation's draft lives under. The composer rebuilds its document + * from the prompt text whenever its value changes, and every node gets a fresh + * random key then, so the key cannot be the node's. Two identical citations + * still need separate drafts, so the serialized citation is suffixed with how + * many identical ones come before it; a rebuild from the same prompt keeps + * that order, and so keeps the key. + */ +export function assistantCitationDraftKey( + citation: AssistantCitation, + citationsBefore: ReadonlyArray, + /** The composer the draft belongs to, so another thread's composer never resumes it. */ + scope = "", +): string { + const serialized = serializeAssistantCitation(citation); + let ordinal = 0; + for (const earlier of citationsBefore) { + if (serializeAssistantCitation(earlier) === serialized) ordinal += 1; + } + return `${scope}\n${serialized}#${ordinal}`; +} + +export function readAssistantCitationCommentDraft(key: string): string | null { + return drafts.get(key) ?? null; +} + +export function writeAssistantCitationCommentDraft(key: string, draft: string): void { + drafts.set(key, draft); +} + +export function clearAssistantCitationCommentDraft(key: string): void { + drafts.delete(key); +} + +/** + * Writes a composer's unsaved comments onto their citations in the prompt, the + * way dismissing the popover would have. Called when a question or an approval + * borrows the prompt editor: its chips unmount without a dismissal, and + * without this they would come back uncommented and Send would leave the + * comment out. A draft the dismissal rules would not commit (over the length + * limit) stays a draft and is resumed when its popover is opened again. + */ +export function commitAssistantCitationCommentDrafts(prompt: string, scope: string): string { + const citations = collectAssistantCitations(prompt); + // Keys are read off the prompt as it is: committing one comment changes that + // citation's serialized form, which the later duplicates' ordinals count. + const keys = citations.map((entry, index) => + assistantCitationDraftKey( + entry.citation, + citations.slice(0, index).map((earlier) => earlier.citation), + scope, + ), + ); + let committed = prompt; + for (let index = citations.length - 1; index >= 0; index -= 1) { + const { citation, start, end } = citations[index]!; + const key = keys[index]!; + const draft = drafts.get(key); + if (draft === undefined) continue; + const dismissal = resolveAssistantCitationCommentDismissal({ + reason: "none", + draft, + savedComment: citation.comment, + }); + if (dismissal.kind === "keep-open") continue; + drafts.delete(key); + if (dismissal.kind !== "commit") continue; + committed = + committed.slice(0, start) + + serializeAssistantCitation(withAssistantCitationComment(citation, dismissal.comment)) + + committed.slice(end); + } + return committed; +} + +export type AssistantCitationCommentDraftEntries = ReadonlyArray< + readonly [key: string, draft: string] +>; + +/** + * Takes every draft a composer still holds once its prompt has been sent. A + * popover that stayed open through the send (a comment over the length limit + * keeps it open) would otherwise hand its text to the next prompt that cites + * the same text in the same composer. A send that fails gives the prompt back + * and restores these with it. + */ +export function takeAssistantCitationCommentDraftsForComposer( + scope: string, +): AssistantCitationCommentDraftEntries { + const prefix = `${scope}\n`; + const taken: Array = []; + for (const [key, draft] of drafts) { + if (!key.startsWith(prefix)) continue; + taken.push([key, draft]); + drafts.delete(key); + } + return taken; +} + +export function restoreAssistantCitationCommentDrafts( + entries: AssistantCitationCommentDraftEntries, +): void { + for (const [key, draft] of entries) drafts.set(key, draft); +}