diff --git a/apps/web/src/components/ComposerCitationNode.tsx b/apps/web/src/components/ComposerCitationNode.tsx index ffaae589343f..2c8f02586688 100644 --- a/apps/web/src/components/ComposerCitationNode.tsx +++ b/apps/web/src/components/ComposerCitationNode.tsx @@ -42,6 +42,7 @@ export type ComposerCitationCommentRequest = { export type ComposerCitationCommentTarget = { nodeKey: NodeKey; sourceAnchor?: AssistantCitationSourceAnchor; + removeOnCancel?: boolean; }; export const ComposerCitationCommentContext = createContext<{ @@ -68,7 +69,11 @@ 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, + removeOnCancel: true, + }; } offset += node.getTextContentSize(); } @@ -78,6 +83,8 @@ export function $consumeComposerCitationCommentRequest(requestRef: { function ComposerCitationDecorator(props: { citation: AssistantCitation; nodeKey: NodeKey }) { const [editor] = useLexicalComposerContext(); const commentContext = use(ComposerCitationCommentContext); + const commentTarget = + commentContext.openComment?.nodeKey === props.nodeKey ? commentContext.openComment : null; const onSaveComment = (comment: string): boolean => { if (!editor.isEditable()) return false; let accepted = false; @@ -118,15 +125,13 @@ function ComposerCitationDecorator(props: { citation: AssistantCitation; nodeKey { if (open && !editor.isEditable()) return; commentContext.onOpenChange(props.nodeKey, open); }, + ...(commentTarget?.removeOnCancel ? { onCancel: onRemove } : {}), onSave: onSaveComment, onSaveAndSend: (comment) => { if (!onSaveComment(comment)) return false; diff --git a/apps/web/src/components/chat/AssistantCitationChip.tsx b/apps/web/src/components/chat/AssistantCitationChip.tsx index 7582bb19c5c6..ccfd746666a4 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; }; @@ -168,7 +169,13 @@ export function AssistantCitationChip({ }, } : {})} - onCancel={() => commentEditor.onOpenChange(false)} + onCancel={() => { + if (commentEditor.onCancel) { + commentEditor.onCancel(); + } else { + commentEditor.onOpenChange(false); + } + }} /> ) : null}