Split transactions when applying suggestions (#465)

components/custom/canvas.tsx

@@ -180,13 +180,18 @@ export function Canvas({
   );
 
   const handleEditorChange = useCallback(
-    (updatedContent: string) => {
+    (updatedContent: string, debounce: boolean) => {
       if (document && updatedContent !== document.content) {
-        debouncedHandleEditorChange(updatedContent);
+        if (debounce) {
+          debouncedHandleEditorChange(updatedContent);
+        } else {
+          handleContentChange(updatedContent);
+        }
+
         setIsContentDirty(true);
       }
     },
-    [document, debouncedHandleEditorChange]
+    [document, debouncedHandleEditorChange, handleContentChange]
   );
 
   function getDocumentContentById(index: number) {

components/custom/editor.tsx

@@ -44,7 +44,7 @@ interface WidgetRoot {
 
 type EditorProps = {
   content: string;
-  onChange: (updatedContent: string) => void;
+  onChange: (updatedContent: string, debounce: boolean) => void;
   status: 'streaming' | 'idle';
   currentVersionIndex: number;
   suggestions: Array<Suggestion>;
@@ -95,7 +95,12 @@ function PureEditor({
 
             if (transaction.docChanged) {
               const content = defaultMarkdownSerializer.serialize(newState.doc);
-              onChange(content);
+
+              if (transaction.getMeta('no-debounce')) {
+                onChange(content, false);
+              } else {
+                onChange(content, true);
+              }
             }
           },
         });
@@ -137,19 +142,34 @@ function PureEditor({
 
     suggestions.forEach((suggestion) => {
       decorations.push(
-        Decoration.inline(suggestion.selectionStart, suggestion.selectionEnd, {
-          class:
-            'suggestion-highlight bg-yellow-100 hover:bg-yellow-200 dark:hover:bg-yellow-400/50 dark:text-yellow-50 dark:bg-yellow-400/40',
-        })
+        Decoration.inline(
+          suggestion.selectionStart,
+          suggestion.selectionEnd,
+          {
+            class:
+              'suggestion-highlight bg-yellow-100 hover:bg-yellow-200 dark:hover:bg-yellow-400/50 dark:text-yellow-50 dark:bg-yellow-400/40',
+          },
+          {
+            suggestionId: suggestion.id,
+            type: 'highlight',
+          }
+        )
       );
 
       decorations.push(
-        Decoration.widget(suggestion.selectionStart, (view) => {
-          const { dom, destroy } = createSuggestionWidget(suggestion, view);
-          const key = `widget-${suggestion.id}`;
-          widgetRootsRef.current.set(key, { destroy });
-          return dom;
-        })
+        Decoration.widget(
+          suggestion.selectionStart,
+          (view) => {
+            const { dom, destroy } = createSuggestionWidget(suggestion, view);
+            const key = `widget-${suggestion.id}`;
+            widgetRootsRef.current.set(key, { destroy });
+            return dom;
+          },
+          {
+            suggestionId: suggestion.id,
+            type: 'widget',
+          }
+        )
       );
     });
 

components/custom/suggestion.tsx

@@ -6,6 +6,7 @@ import { useState } from 'react';
 import { UISuggestion } from '@/lib/editor/suggestions';
 
 import { CrossIcon, MessageIcon } from './icons';
+import useWindowSize from './use-window-size';
 import { Button } from '../ui/button';
 
 export const Suggestion = ({
@@ -16,15 +17,16 @@ export const Suggestion = ({
   onApply: () => void;
 }) => {
   const [isExpanded, setIsExpanded] = useState(false);
+  const { width: windowWidth } = useWindowSize();
 
   return !isExpanded ? (
     <div
-      className="absolute cursor-pointer text-muted-foreground mt-1 -right-8"
+      className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
       onClick={() => {
         setIsExpanded(true);
       }}
     >
-      <MessageIcon size={14} />
+      <MessageIcon size={windowWidth && windowWidth < 768 ? 16 : 14} />
     </div>
   ) : (
     <motion.div

lib/editor/suggestions.tsx

@@ -1,6 +1,7 @@
+import { defaultMarkdownSerializer } from 'prosemirror-markdown';
 import { Node } from 'prosemirror-model';
 import { PluginKey, Plugin } from 'prosemirror-state';
-import { DecorationSet, EditorView } from 'prosemirror-view';
+import { Decoration, DecorationSet, EditorView } from 'prosemirror-view';
 import { createRoot } from 'react-dom/client';
 
 import { Suggestion as PreviewSuggestion } from '@/components/custom/suggestion';
@@ -69,15 +70,42 @@ export function createSuggestionWidget(
   const dom = document.createElement('span');
   const root = createRoot(dom);
 
+  dom.addEventListener('mousedown', (event) => {
+    event.preventDefault();
+    view.dom.blur();
+  });
+
   const onApply = () => {
     const { state, dispatch } = view;
-    const tr = state.tr.replaceWith(
+
+    let decorationTransaction = state.tr;
+    const currentState = suggestionsPluginKey.getState(state);
+    const currentDecorations = currentState?.decorations;
+
+    if (currentDecorations) {
+      const newDecorations = DecorationSet.create(
+        state.doc,
+        currentDecorations.find().filter((decoration: Decoration) => {
+          return decoration.spec.suggestionId !== suggestion.id;
+        })
+      );
+
+      decorationTransaction.setMeta(suggestionsPluginKey, {
+        decorations: newDecorations,
+        selected: null,
+      });
+      dispatch(decorationTransaction);
+    }
+
+    const textTransaction = view.state.tr.replaceWith(
       suggestion.selectionStart,
       suggestion.selectionEnd,
       state.schema.text(suggestion.suggestedText)
     );
 
-    dispatch(tr);
+    textTransaction.setMeta('no-debounce', true);
+
+    dispatch(textTransaction);
   };
 
   root.render(<PreviewSuggestion suggestion={suggestion} onApply={onApply} />);