feat(ai): enhance loading states for writer and re-writer

packages/super-editor/src/assets/styles/elements/_all.css

@@ -1,3 +1,3 @@
 @import './prosemirror.css';
 @import './toolbar.css';
-@import './toolbar-custom.css';
+@import './ai.css';

packages/super-editor/src/assets/styles/elements/toolbar-custom.css → packages/super-editor/src/assets/styles/elements/ai.css

@@ -35,6 +35,8 @@
   filter: brightness(1.3);
 } 
 
+
+/* AI text appear animation */
 @keyframes aiTextAppear {
   from {
     opacity: 0;
@@ -54,4 +56,30 @@
   will-change: opacity, transform;
   /* Ensure each mark is treated as a separate animation context */
   contain: content;
-}
+}
+
+.sd-ai-loader {
+  display: flex;
+  justify-content: flex-start;
+}
+
+.sd-ai-loader > img {
+  width: fit-content;
+  height: 40px;
+}
+
+@keyframes ai-pulse {
+  0% {
+    background-color: rgba(99, 102, 241, 0.1);
+  }
+  50% {
+    background-color: rgba(99, 102, 241, 0.375);
+  }
+  100% {
+    background-color: rgba(99, 102, 241, 0.1);
+  }
+}
+
+.sd-ai-highlight-pulse {
+  animation: ai-pulse 1.5s ease-in-out infinite;
+}

packages/super-editor/src/components/toolbar/AIWriter.vue

@@ -1,10 +1,9 @@
 <script setup>
 import { computed, ref, onMounted, onUnmounted, nextTick } from 'vue';
 import { writeStreaming, rewriteStreaming } from './ai-helpers';
+import { TextSelection } from 'prosemirror-state';
 import edit from '@harbour-enterprises/common/icons/edit-regular.svg?raw';
 import paperPlane from '@harbour-enterprises/common/icons/paper-plane-regular.svg?raw';
-import timesCircle from '@harbour-enterprises/common/icons/times-circle-regular.svg?raw';
-import sun from '@harbour-enterprises/common/icons/sun-regular.svg?raw';
 
 const props = defineProps({
   selectedText: {
@@ -37,6 +36,8 @@ const props = defineProps({
   },
 });
 
+const emits = defineEmits(['ai-highlight']);
+
 // Store the selection state
 const selectionState = ref(null);
 
@@ -45,6 +46,10 @@ const aiWriterRef = ref(null);
 
 const handleClickOutside = (event) => {
   if (aiWriterRef.value && !aiWriterRef.value.contains(event.target)) {
+    // Only emit 'remove' if we're not in a loading state
+    if (!isLoading.value) {
+      emitAiHighlight('remove');
+    }
     props.handleClose();
   }
 };
@@ -53,16 +58,25 @@ const handleClickOutside = (event) => {
 const editableRef = ref(null);
 
 // Helper function to emit AI highlight events
+// We need to emit through the superToolbar if it exists, otherwise emit through the emits
+// Hoping we can simplify this logic in the future if we combine SE and SD
 const emitAiHighlight = (type, data = null) => {
   if (props.superToolbar) {
     props.superToolbar.emit('ai-highlight', { type, data });
+  } else {
+    emits('ai-highlight', { type, data });
   }
 };
 
 // Helper functions
 const saveSelection = () => {
   if (props.selectedText) {
-    selectionState.value = props.editor.state.selection;
+    // Store the complete selection state
+    selectionState.value = {
+      ...props.editor.state.selection,
+      from: props.editor.state.selection.from,
+      to: props.editor.state.selection.to
+    };
     // Store the selection in the editor's state
     props.editor.commands.setMeta('storedSelection', selectionState.value);
 
@@ -97,9 +111,10 @@ onMounted(() => {
 });
 
 onUnmounted(() => {
-  // emit the ai highlight remove event through the toolbar
-  emitAiHighlight('remove');
-
+  // Only emit 'remove' if we're not in a loading state
+  if (!isLoading.value) {
+    emitAiHighlight('remove');
+  }
   // Remove all event listeners
   removeEventListeners();
 });
@@ -123,6 +138,8 @@ const placeholderText = computed(() =>
 const isLoading = ref(false);
 const isError = ref('');
 const promptText = ref('');
+const textProcessingStarted = ref(false);
+const previousText = ref('');
 
 // Computed property to check if editor is in suggesting mode
 const isInSuggestingMode = computed(() => {
@@ -144,10 +161,39 @@ const getDocumentXml = () => {
 // Handler for processing text chunks from the stream
 const handleTextChunk = (text) => {
   try {
-    // If this is the first chunk and we're rewriting, remove the selected text
+    // Remove the loader node when we start receiving text
+    props.editor.commands.removeAiNode('aiLoaderNode');
+
+    // If this is the first chunk and we're rewriting, handle the selection
     if (props.selectedText && !textProcessingStarted.value) {
+      emitAiHighlight('remove');
+
+      // Clear the pulsing animation when we start inserting text
+      props.editor.commands.clearAiHighlightStyle();
+
+      // Check if we have a valid stored selection
+      if (selectionState.value) {
+        // Apply the stored selection using the TextSelection API
+        const { state } = props.editor;
+        const { from, to } = selectionState.value;
+
+        // Create a transaction to set the selection
+        const tr = state.tr.setSelection(
+          TextSelection.create(state.doc, from, to)
+        );
+
+        // Dispatch the transaction to update the editor state
+        props.editor.view.dispatch(tr);
+      } else {
+        console.warn('[AIWriter] No stored selection to restore');
+      }
+
+      // Now delete the selection
       props.editor.commands.deleteSelection();
+
+      // Mark as processed
       textProcessingStarted.value = true;
+
     }
 
     // If the text is null, undefined or empty, don't process it
@@ -188,22 +234,33 @@ const handleDone = () => {
   // We need to wait for the animation to finish before removing the mark
   setTimeout(() => {
     props.editor.commands.removeAiMark('aiAnimationMark');
+    // Remove the highlight when we're done
+    emitAiHighlight('remove');
   }, 1000);
 };
 
-// Track text processing state
-const textProcessingStarted = ref(false);
-const previousText = ref('');
-
 // Refactored handleSubmit function
 const handleSubmit = async () => {
   // Reset state
-  isLoading.value = true;
   isError.value = '';
   textProcessingStarted.value = false;
   previousText.value = '';
+  isLoading.value = true;  // Set loading to true before starting the operation
 
   try {
+    // Close the AI Writer immediately and transition to loading states
+    props.handleClose();
+
+    // If there is selected text, emit the update event to start pulsing animation
+    if (props.selectedText) {
+      emitAiHighlight('update');
+    } else {
+      // Insert the loader node at the current cursor position
+      props.editor.commands.insertContent({
+        type: 'aiLoaderNode',
+      });
+    }
+
     // Enable track changes if in suggesting mode
     if (isInSuggestingMode.value) {
       props.editor.commands.enableTrackChanges();
@@ -233,18 +290,19 @@ const handleSubmit = async () => {
       await writeStreaming(promptText.value, options, handleTextChunk, handleDone);
     }
 
-    // If all is good, close the AI Writer
-    props.handleClose();
   } catch (error) {
     console.error('AI generation error:', error);
     isError.value = error.message || 'An error occurred';
   } finally {
-    promptText.value = ''; // Clear the input after submission
+    // Clear the input after submission
+    promptText.value = '';
     // Only disable track changes if we enabled it (in suggesting mode)
     if (isInSuggestingMode.value) {
       props.editor.commands.disableTrackChanges();
     }
     isLoading.value = false;
+    // Ensure textProcessingStarted is reset for next operation
+    textProcessingStarted.value = false;
   }
 };
 
@@ -283,11 +341,7 @@ const handleInput = (event) => {
       ></textarea>
     </div>
     <div class="ai-loader">
-      <span v-if="isLoading" class="ai-textarea-icon loading spinner-wrapper">
-        <span v-html="sun"></span>
-      </span>
-      <span v-else-if="isError" class="ai-textarea-icon error" v-html="timesCircle" />
-      <span v-else-if="promptText" class="ai-textarea-icon ai-submit-button" @click.stop="handleSubmit" v-html="paperPlane" />
+      <span v-if="promptText" class="ai-textarea-icon ai-submit-button" @click.stop="handleSubmit" v-html="paperPlane" />
     </div>
   </div>
 </template>

packages/super-editor/src/extensions/ai/ai-constants.js

@@ -1,2 +1,3 @@
 export const AiMarkName = 'aiMark'; 
 export const AiAnimationMarkName = 'aiAnimationMark';
+export const AiLoaderNodeName = 'aiLoaderNode';

packages/super-editor/src/extensions/ai/ai-marks.js

@@ -1,4 +1,4 @@
-import { Mark, Attribute } from '@core/index.js';
+import { Mark, Attribute, Node } from '@core/index.js';
 import { AiMarkName, AiAnimationMarkName } from './ai-constants.js';
 
 export const AiMark = Mark.create({
@@ -10,7 +10,7 @@ export const AiMark = Mark.create({
 
   addOptions() {
     return {
-      htmlAttributes: { class: 'sd-editor-ai' },
+      htmlAttributes: { class: 'sd-ai-highlight' },
     };
   },
 

packages/super-editor/src/extensions/ai/ai-nodes.js

@@ -0,0 +1,44 @@
+import { Attribute, Node } from '@core/index.js';
+import dotsLoader from '@harbour-enterprises/common/icons/dots-loader.svg';
+import { AiLoaderNodeName } from './ai-constants.js';
+
+export const AiLoaderNode = Node.create({
+  name: AiLoaderNodeName,
+
+  group: 'inline',
+
+  inline: true,
+
+  atom: true,
+
+  selectable: false,
+
+  draggable: false,
+
+  addOptions() {
+    return {
+      htmlAttributes: {
+        class: 'sd-ai-loader',
+        contentEditable: 'false',
+      }
+    };
+  },
+
+  parseDOM() {
+    return [{ tag: 'span.sd-ai-loader' }];
+  },
+
+  renderDOM({ htmlAttributes }) {
+    const span = document.createElement('span');
+    Object.entries(Attribute.mergeAttributes(this.options.htmlAttributes, htmlAttributes))
+      .forEach(([k, v]) => span.setAttribute(k, v));
+
+    const img = document.createElement('img');
+    img.src = dotsLoader;
+    img.alt = 'loading...';
+    img.width = 100;
+    img.height = 50;
+    span.appendChild(img);
+    return span;
+  }
+});