Clean up editor code and reorganize contents (#478)

app/globals.css

@@ -140,3 +140,7 @@
 .ProseMirror {
   outline: none;
 }
+
+.suggestion-highlight {
+  @apply bg-yellow-100 hover:bg-yellow-200 dark:hover:bg-yellow-400/50 dark:text-yellow-50 dark:bg-yellow-400/40;
+}

components/custom/canvas-stream-handler.tsx

@@ -5,7 +5,7 @@ import { UICanvas } from './canvas';
 import { useCanvasStream } from './use-canvas-stream';
 
 interface CanvasStreamHandlerProps {
-  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  setCanvas: Dispatch<SetStateAction<UICanvas>>;
   streamingData: JSONValue[] | undefined;
 }
 

components/custom/canvas.tsx

@@ -1,6 +1,6 @@
 import { Attachment, ChatRequestOptions, CreateMessage, Message } from 'ai';
 import cx from 'classnames';
-import { formatDistance, isAfter } from 'date-fns';
+import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
   Dispatch,
@@ -10,6 +10,7 @@ import {
   useState,
 } from 'react';
 import useSWR, { useSWRConfig } from 'swr';
+import { useDebounceCallback, useWindowSize } from 'usehooks-ts';
 
 import { Document, Suggestion } from '@/db/schema';
 import { fetcher } from '@/lib/utils';
@@ -18,14 +19,12 @@ import { DiffView } from './diffview';
 import { DocumentSkeleton } from './document-skeleton';
 import { Editor } from './editor';
 import { CrossIcon, DeltaIcon, RedoIcon, UndoIcon } from './icons';
+import { Markdown } from './markdown';
 import { Message as PreviewMessage } from './message';
 import { MultimodalInput } from './multimodal-input';
 import { Toolbar } from './toolbar';
-import { useDebounce } from './use-debounce';
 import { useScrollToBottom } from './use-scroll-to-bottom';
-import useWindowSize from './use-window-size';
+import { VersionFooter } from './version-footer';
-import { Button } from '../ui/button';
-
 export interface UICanvas {
   title: string;
   documentId: string;
@@ -61,7 +60,7 @@ export function Canvas({
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
   canvas: UICanvas;
-  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  setCanvas: Dispatch<SetStateAction<UICanvas>>;
   messages: Array<Message>;
   setMessages: Dispatch<SetStateAction<Array<Message>>>;
   append: (
@@ -81,7 +80,6 @@ export function Canvas({
   const {
     data: documents,
     isLoading: isDocumentsFetching,
-    isValidating: isDocumentsValidating,
     mutate: mutateDocuments,
   } = useSWR<Array<Document>>(
     canvas && canvas.status !== 'streaming'
@@ -111,14 +109,10 @@ export function Canvas({
       if (mostRecentDocument) {
         setDocument(mostRecentDocument);
         setCurrentVersionIndex(documents.length - 1);
-        setCanvas((currentCanvas) =>
+        setCanvas((currentCanvas) => ({
-          currentCanvas
+          ...currentCanvas,
-            ? {
+          content: mostRecentDocument.content ?? '',
-                ...currentCanvas,
+        }));
-                content: mostRecentDocument.content ?? '',
-              }
-            : null
-        );
       }
     }
   }, [documents, setCanvas]);
@@ -174,24 +168,24 @@ export function Canvas({
     [canvas, mutate]
   );
 
-  const debouncedHandleEditorChange = useCallback(
+  const debouncedHandleContentChange = useDebounceCallback(
-    useDebounce(handleContentChange, 4000),
+    handleContentChange,
-    [handleContentChange]
+    2000
   );
 
-  const handleEditorChange = useCallback(
+  const saveContent = useCallback(
     (updatedContent: string, debounce: boolean) => {
       if (document && updatedContent !== document.content) {
+        setIsContentDirty(true);
+
         if (debounce) {
-          debouncedHandleEditorChange(updatedContent);
+          debouncedHandleContentChange(updatedContent);
         } else {
           handleContentChange(updatedContent);
         }
-
-        setIsContentDirty(true);
       }
     },
-    [document, debouncedHandleEditorChange, handleContentChange]
+    [document, debouncedHandleContentChange, handleContentChange]
   );
 
   function getDocumentContentById(index: number) {
@@ -200,11 +194,6 @@ export function Canvas({
     return documents[index].content ?? '';
   }
 
-  function getDocumentTimestampById(index: number) {
-    if (!documents) return '';
-    return documents[index]?.createdAt ?? '';
-  }
-
   const handleVersionChange = (type: 'next' | 'prev' | 'toggle' | 'latest') => {
     if (!documents) return;
 
@@ -237,11 +226,9 @@ export function Canvas({
    */
 
   const isCurrentVersion =
-    isDocumentsFetching || isDocumentsValidating
+    documents && documents.length > 0
-      ? true
+      ? currentVersionIndex === documents.length - 1
-      : documents && documents.length > 0
+      : true;
-        ? currentVersionIndex === documents.length - 1
-        : true;
 
   const { width: windowWidth, height: windowHeight } = useWindowSize();
   const isMobile = windowWidth ? windowWidth < 768 : false;
@@ -391,7 +378,10 @@ export function Canvas({
             <div
               className="cursor-pointer hover:bg-muted dark:hover:bg-zinc-700 p-2 rounded-lg text-muted-foreground"
               onClick={() => {
-                setCanvas(null);
+                setCanvas((currentCanvas) => ({
+                  ...currentCanvas,
+                  isVisible: false,
+                }));
               }}
             >
               <CrossIcon size={18} />
@@ -462,9 +452,10 @@ export function Canvas({
                     ? canvas.content
                     : getDocumentContentById(currentVersionIndex)
                 }
+                isCurrentVersion={isCurrentVersion}
                 currentVersionIndex={currentVersionIndex}
-                status={canvas.status ?? 'idle'}
+                status={canvas.status}
-                onChange={handleEditorChange}
+                saveContent={saveContent}
                 suggestions={isCurrentVersion ? (suggestions ?? []) : []}
               />
             ) : (
@@ -482,63 +473,12 @@ export function Canvas({
 
         <AnimatePresence>
           {!isCurrentVersion && (
-            <motion.div
+            <VersionFooter
-              className="absolute flex flex-col gap-4 lg:flex-row bottom-0 bg-background p-4 w-full border-t z-50 justify-between"
+              canvas={canvas}
-              initial={{ y: isMobile ? 200 : 77 }}
+              currentVersionIndex={currentVersionIndex}
-              animate={{ y: 0 }}
+              documents={documents}
-              exit={{ y: isMobile ? 200 : 77 }}
+              handleVersionChange={handleVersionChange}
-              transition={{ type: 'spring', stiffness: 140, damping: 20 }}
+            />
-            >
-              <div>
-                <div>You are viewing a previous version</div>
-                <div className="text-muted-foreground text-sm">
-                  Restore this version to make edits
-                </div>
-              </div>
-
-              <div className="flex flex-row gap-4">
-                <Button
-                  onClick={async () => {
-                    mutate(
-                      `/api/document?id=${canvas.documentId}`,
-                      await fetch(`/api/document?id=${canvas.documentId}`, {
-                        method: 'PATCH',
-                        body: JSON.stringify({
-                          timestamp:
-                            getDocumentTimestampById(currentVersionIndex),
-                        }),
-                      }),
-                      {
-                        optimisticData: documents
-                          ? [
-                              ...documents.filter((d) =>
-                                isAfter(
-                                  new Date(d.createdAt),
-                                  new Date(
-                                    getDocumentTimestampById(
-                                      currentVersionIndex
-                                    )
-                                  )
-                                )
-                              ),
-                            ]
-                          : [],
-                      }
-                    );
-                  }}
-                >
-                  Restore this version
-                </Button>
-                <Button
-                  variant="outline"
-                  onClick={() => {
-                    handleVersionChange('latest');
-                  }}
-                >
-                  Back to latest version
-                </Button>
-              </div>
-            </motion.div>
           )}
         </AnimatePresence>
       </motion.div>

components/custom/chat.tsx

@@ -4,8 +4,8 @@ import { Attachment, Message } from 'ai';
 import { useChat } from 'ai/react';
 import { AnimatePresence } from 'framer-motion';
 import { useState } from 'react';
+import { useWindowSize } from 'usehooks-ts';
 
-import { Model } from '@/ai/models';
 import { ChatHeader } from '@/components/custom/chat-header';
 import { Message as PreviewMessage } from '@/components/custom/message';
 import { useScrollToBottom } from '@/components/custom/use-scroll-to-bottom';
@@ -42,7 +42,22 @@ export function Chat({
     },
   });
 
-  const [canvas, setCanvas] = useState<UICanvas | null>(null);
+  const { width: windowWidth = 1920, height: windowHeight = 1080 } =
+    useWindowSize();
+
+  const [canvas, setCanvas] = useState<UICanvas>({
+    documentId: 'init',
+    content: '',
+    title: '',
+    status: 'idle',
+    isVisible: false,
+    boundingBox: {
+      top: windowHeight / 4,
+      left: windowWidth / 4,
+      width: 250,
+      height: 50,
+    },
+  });
 
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>();

components/custom/document.tsx

@@ -19,8 +19,8 @@ const getActionText = (type: 'create' | 'update' | 'request-suggestions') => {
 interface DocumentToolResultProps {
   type: 'create' | 'update' | 'request-suggestions';
   result: any;
-  canvas: UICanvas | null;
+  canvas: UICanvas;
-  setCanvas: (value: SetStateAction<UICanvas | null>) => void;
+  setCanvas: (value: SetStateAction<UICanvas>) => void;
 }
 
 export function DocumentToolResult({
@@ -42,27 +42,14 @@ export function DocumentToolResult({
           height: rect.height,
         };
 
-        if (!canvas) {
+        setCanvas({
-          setCanvas({
+          documentId: result.id,
-            documentId: result.id,
+          content: '',
-            content: '',
+          title: result.title,
-            title: result.title,
+          isVisible: true,
-            isVisible: true,
+          status: 'idle',
-            status: 'idle',
+          boundingBox,
-            boundingBox,
+        });
-          });
-        } else {
-          if (canvas.documentId !== result.id) {
-            setCanvas({
-              documentId: result.id,
-              content: '',
-              title: result.title,
-              isVisible: true,
-              status: 'idle',
-              boundingBox,
-            });
-          }
-        }
       }}
     >
       <div className="text-muted-foreground mt-1">

components/custom/editor.tsx

@@ -1,213 +1,170 @@
 'use client';
 
 import { exampleSetup } from 'prosemirror-example-setup';
-import { inputRules, textblockTypeInputRule } from 'prosemirror-inputrules';
+import { inputRules } from 'prosemirror-inputrules';
-import { defaultMarkdownSerializer } from 'prosemirror-markdown';
-import { Schema, DOMParser } from 'prosemirror-model';
-import { schema } from 'prosemirror-schema-basic';
-import { addListNodes } from 'prosemirror-schema-list';
 import { EditorState } from 'prosemirror-state';
-import { Decoration, DecorationSet, EditorView } from 'prosemirror-view';
+import { EditorView } from 'prosemirror-view';
 import React, { memo, useEffect, useRef } from 'react';
-import { renderToString } from 'react-dom/server';
 
 import { Suggestion } from '@/db/schema';
 import {
-  createSuggestionWidget,
+  documentSchema,
-  projectWithHighlights,
+  handleTransaction,
+  headingRule,
+} from '@/lib/editor/config';
+import {
+  buildContentFromDocument,
+  buildDocumentFromContent,
+  createDecorations,
+} from '@/lib/editor/functions';
+import {
+  projectWithPositions,
   suggestionsPlugin,
   suggestionsPluginKey,
-  UISuggestion,
 } from '@/lib/editor/suggestions';
 
-import { Markdown } from './markdown';
-
-const mySchema = new Schema({
-  nodes: addListNodes(schema.spec.nodes, 'paragraph block*', 'block'),
-  // @ts-expect-error: TODO need to fix type mismatch
-  marks: {
-    ...schema.spec.marks,
-  },
-});
-
-function headingRule(level: number) {
-  return textblockTypeInputRule(
-    new RegExp(`^(#{1,${level}})\\s$`),
-    mySchema.nodes.heading,
-    () => ({ level })
-  );
-}
-
-interface WidgetRoot {
-  destroy: () => void;
-}
-
 type EditorProps = {
   content: string;
-  onChange: (updatedContent: string, debounce: boolean) => void;
+  saveContent: (updatedContent: string, debounce: boolean) => void;
   status: 'streaming' | 'idle';
+  isCurrentVersion: boolean;
   currentVersionIndex: number;
   suggestions: Array<Suggestion>;
 };
 
 function PureEditor({
   content,
-  onChange,
+  saveContent,
-  suggestions: suggestionsWithoutHighlights,
+  suggestions,
+  status,
 }: EditorProps) {
-  const editorRef = useRef<HTMLDivElement>(null);
+  const containerRef = useRef<HTMLDivElement>(null);
-  const viewRef = useRef<EditorView | null>(null);
+  const editorRef = useRef<EditorView | null>(null);
-  const widgetRootsRef = useRef<Map<string, WidgetRoot>>(new Map());
 
   useEffect(() => {
-    if (editorRef.current && !viewRef.current) {
+    if (containerRef.current && !editorRef.current) {
-      if (mySchema) {
+      const state = EditorState.create({
-        const parser = DOMParser.fromSchema(mySchema);
+        doc: buildDocumentFromContent(content),
+        plugins: [
+          ...exampleSetup({ schema: documentSchema, menuBar: false }),
+          inputRules({
+            rules: [
+              headingRule(1),
+              headingRule(2),
+              headingRule(3),
+              headingRule(4),
+              headingRule(5),
+              headingRule(6),
+            ],
+          }),
+          suggestionsPlugin,
+        ],
+      });
 
-        const htmlContent = renderToString(<Markdown>{content}</Markdown>);
+      editorRef.current = new EditorView(containerRef.current, {
-
+        state,
-        const container = document.createElement('div');
+      });
-        container.innerHTML = htmlContent;
-
-        const state = EditorState.create({
-          doc: parser.parse(container),
-          plugins: [
-            ...exampleSetup({ schema: mySchema, menuBar: false }),
-            inputRules({
-              rules: [
-                headingRule(1),
-                headingRule(2),
-                headingRule(3),
-                headingRule(4),
-                headingRule(5),
-                headingRule(6),
-              ],
-            }),
-            suggestionsPlugin,
-          ],
-        });
-
-        viewRef.current = new EditorView(editorRef.current, {
-          state,
-          dispatchTransaction: (transaction) => {
-            const newState = viewRef.current!.state.apply(transaction);
-            viewRef.current!.updateState(newState);
-
-            if (transaction.docChanged) {
-              const content = defaultMarkdownSerializer.serialize(newState.doc);
-
-              if (transaction.getMeta('no-debounce')) {
-                onChange(content, false);
-              } else {
-                onChange(content, true);
-              }
-            }
-          },
-        });
-      } else {
-        console.error('Schema is not properly initialized');
-      }
     }
 
     return () => {
-      if (viewRef.current) {
+      if (editorRef.current) {
-        viewRef.current.destroy();
+        editorRef.current.destroy();
-        viewRef.current = null;
+        editorRef.current = null;
       }
     };
-  }, [content, onChange]);
+    // NOTE: we only want to run this effect once
+    // eslint-disable-next-line
+  }, []);
 
   useEffect(() => {
-    if (viewRef.current && viewRef.current.state.doc && content) {
+    if (editorRef.current) {
-      const computedSuggestions = projectWithHighlights(
+      editorRef.current.setProps({
-        viewRef.current.state.doc,
+        dispatchTransaction: (transaction) => {
-        suggestionsWithoutHighlights
+          handleTransaction({ transaction, editorRef, saveContent });
+        },
+      });
+    }
+  }, [saveContent]);
+
+  useEffect(() => {
+    if (editorRef.current && content) {
+      const currentContent = buildContentFromDocument(
+        editorRef.current.state.doc
+      );
+
+      if (status === 'streaming') {
+        const newDocument = buildDocumentFromContent(content);
+
+        const transaction = editorRef.current.state.tr.replaceWith(
+          0,
+          editorRef.current.state.doc.content.size,
+          newDocument.content
+        );
+
+        transaction.setMeta('no-save', true);
+        editorRef.current.dispatch(transaction);
+        return;
+      }
+
+      if (currentContent !== content) {
+        const newDocument = buildDocumentFromContent(content);
+
+        const transaction = editorRef.current.state.tr.replaceWith(
+          0,
+          editorRef.current.state.doc.content.size,
+          newDocument.content
+        );
+
+        transaction.setMeta('no-save', true);
+        editorRef.current.dispatch(transaction);
+      }
+    }
+  }, [content, status]);
+
+  useEffect(() => {
+    if (editorRef.current && editorRef.current.state.doc && content) {
+      const projectedSuggestions = projectWithPositions(
+        editorRef.current.state.doc,
+        suggestions
       ).filter(
-        (suggestion) =>
+        (suggestion) => suggestion.selectionStart && suggestion.selectionEnd
-          suggestion.selectionStart !== 0 && suggestion.selectionEnd !== 0
       );
 
       const decorations = createDecorations(
-        computedSuggestions,
+        projectedSuggestions,
-        viewRef.current
+        editorRef.current
       );
-      const transaction = viewRef.current.state.tr;
+
+      const transaction = editorRef.current.state.tr;
       transaction.setMeta(suggestionsPluginKey, { decorations });
-      viewRef.current.dispatch(transaction);
+      editorRef.current.dispatch(transaction);
     }
-  }, [suggestionsWithoutHighlights, content]);
+  }, [suggestions, content]);
 
-  const createDecorations = (suggestions: UISuggestion[], view: EditorView) => {
+  return (
-    const decorations: Decoration[] = [];
+    <div className="relative prose dark:prose-invert" ref={containerRef} />
-
+  );
-    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',
-          },
-          {
-            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;
-          },
-          {
-            suggestionId: suggestion.id,
-            type: 'widget',
-          }
-        )
-      );
-    });
-
-    return DecorationSet.create(view.state.doc, decorations);
-  };
-
-  useEffect(() => {
-    return () => {
-      widgetRootsRef.current.forEach((root) => {
-        root.destroy();
-      });
-
-      widgetRootsRef.current.clear();
-
-      if (viewRef.current) {
-        viewRef.current.destroy();
-        viewRef.current = null;
-      }
-    };
-  }, []);
-
-  return <div className="relative prose dark:prose-invert" ref={editorRef} />;
 }
 
 function areEqual(prevProps: EditorProps, nextProps: EditorProps) {
   if (prevProps.suggestions !== nextProps.suggestions) {
     return false;
-  } else if (prevProps.content === '' && nextProps.content !== '') {
-    return false;
   } else if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex) {
     return false;
-  } else if (prevProps.onChange !== nextProps.onChange) {
+  } else if (prevProps.isCurrentVersion !== nextProps.isCurrentVersion) {
+    return false;
+  } else if (
+    prevProps.status === 'streaming' &&
+    nextProps.status === 'streaming'
+  ) {
+    return false;
+  } else if (prevProps.content !== nextProps.content) {
+    return false;
+  } else if (prevProps.saveContent !== nextProps.saveContent) {
     return false;
-  } else if (prevProps.status === 'idle') {
-    return true;
-  } else {
-    return prevProps.content === nextProps.content;
   }
+
+  return true;
 }
 
 export const Editor = memo(PureEditor, areEqual);

components/custom/message.tsx

@@ -24,8 +24,8 @@ export const Message = ({
   content: string | ReactNode;
   toolInvocations: Array<ToolInvocation> | undefined;
   attachments?: Array<Attachment>;
-  canvas: UICanvas | null;
+  canvas: UICanvas;
-  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  setCanvas: Dispatch<SetStateAction<UICanvas>>;
 }) => {
   return (
     <motion.div

components/custom/multimodal-input.tsx

@@ -13,13 +13,12 @@ import React, {
   ChangeEvent,
 } from 'react';
 import { toast } from 'sonner';
-import { useLocalStorage } from 'usehooks-ts';
+import { useLocalStorage, useWindowSize } from 'usehooks-ts';
 
 import { sanitizeUIMessages } from '@/lib/utils';
 
 import { ArrowUpIcon, PaperclipIcon, StopIcon } from './icons';
 import { PreviewAttachment } from './preview-attachment';
-import useWindowSize from './use-window-size';
 import { Button } from '../ui/button';
 import { Textarea } from '../ui/textarea';
 

components/custom/suggestion.tsx

@@ -1,12 +1,12 @@
 'use client';
 
-import { motion } from 'framer-motion';
+import { AnimatePresence, motion } from 'framer-motion';
 import { useState } from 'react';
+import { useWindowSize } from 'usehooks-ts';
 
 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 = ({
@@ -19,43 +19,52 @@ export const Suggestion = ({
   const [isExpanded, setIsExpanded] = useState(false);
   const { width: windowWidth } = useWindowSize();
 
-  return !isExpanded ? (
+  return (
-    <div
+    <AnimatePresence>
-      className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
+      {!isExpanded ? (
-      onClick={() => {
+        <motion.div
-        setIsExpanded(true);
+          className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
-      }}
-    >
-      <MessageIcon size={windowWidth && windowWidth < 768 ? 16 : 14} />
-    </div>
-  ) : (
-    <motion.div
-      className="absolute bg-background p-3 flex flex-col gap-3 rounded-2xl border text-sm w-56 shadow-xl z-50 -right-12 md:-right-20"
-      transition={{ type: 'spring', stiffness: 500, damping: 30 }}
-      whileHover={{ scale: 1.05 }}
-    >
-      <div className="flex flex-row items-center justify-between">
-        <div className="flex flex-row items-center gap-2">
-          <div className="size-4 bg-muted-foreground/25 rounded-full" />
-          <div className="font-medium">Assistant</div>
-        </div>
-        <div
-          className="text-xs text-gray-500 cursor-pointer"
           onClick={() => {
-            setIsExpanded(false);
+            setIsExpanded(true);
           }}
+          whileHover={{ scale: 1.1 }}
         >
-          <CrossIcon size={12} />
+          <MessageIcon size={windowWidth && windowWidth < 768 ? 16 : 14} />
-        </div>
+        </motion.div>
-      </div>
+      ) : (
-      <div>{suggestion.description}</div>
+        <motion.div
-      <Button
+          key={suggestion.id}
-        variant="outline"
+          className="absolute bg-background p-3 flex flex-col gap-3 rounded-2xl border text-sm w-56 shadow-xl z-50 -right-12 md:-right-16"
-        className="w-fit py-1.5 px-3 rounded-full"
+          transition={{ type: 'spring', stiffness: 500, damping: 30 }}
-        onClick={onApply}
+          initial={{ opacity: 0, y: -10 }}
-      >
+          animate={{ opacity: 1, y: -20 }}
-        Apply
+          exit={{ opacity: 0, y: -10 }}
-      </Button>
+          whileHover={{ scale: 1.05 }}
-    </motion.div>
+        >
+          <div className="flex flex-row items-center justify-between">
+            <div className="flex flex-row items-center gap-2">
+              <div className="size-4 bg-muted-foreground/25 rounded-full" />
+              <div className="font-medium">Assistant</div>
+            </div>
+            <div
+              className="text-xs text-gray-500 cursor-pointer"
+              onClick={() => {
+                setIsExpanded(false);
+              }}
+            >
+              <CrossIcon size={12} />
+            </div>
+          </div>
+          <div>{suggestion.description}</div>
+          <Button
+            variant="outline"
+            className="w-fit py-1.5 px-3 rounded-full"
+            onClick={onApply}
+          >
+            Apply
+          </Button>
+        </motion.div>
+      )}
+    </AnimatePresence>
   );
 };

components/custom/toolbar.tsx

@@ -9,6 +9,7 @@ import {
   useTransform,
 } from 'framer-motion';
 import { Dispatch, SetStateAction, useEffect, useRef, useState } from 'react';
+import { useOnClickOutside } from 'usehooks-ts';
 
 import {
   Tooltip,
@@ -25,7 +26,6 @@ import {
   StopIcon,
   SummarizeIcon,
 } from './icons';
-import useClickOutside from './use-click-outside';
 
 type ToolProps = {
   type: 'final-polish' | 'request-suggestions' | 'adjust-reading-level';
@@ -265,7 +265,7 @@ export const Toolbar = ({
   const [selectedTool, setSelectedTool] = useState<string | null>(null);
   const [isAnimating, setIsAnimating] = useState(false);
 
-  useClickOutside(toolbarRef, () => {
+  useOnClickOutside(toolbarRef, () => {
     setIsToolbarVisible(false);
     setSelectedTool(null);
   });

components/custom/use-canvas-stream.tsx

@@ -5,7 +5,6 @@ import { useSWRConfig } from 'swr';
 import { Suggestion } from '@/db/schema';
 
 import { UICanvas } from './canvas';
-import useWindowSize from './use-window-size';
 
 type StreamingDelta = {
   type: 'text-delta' | 'title' | 'id' | 'suggestion' | 'clear' | 'finish';
@@ -17,7 +16,7 @@ export function useCanvasStream({
   setCanvas,
 }: {
   streamingData: JSONValue[] | undefined;
-  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  setCanvas: Dispatch<SetStateAction<UICanvas>>;
 }) {
   const { mutate } = useSWRConfig();
   const [optimisticSuggestions, setOptimisticSuggestions] = useState<
@@ -32,9 +31,6 @@ export function useCanvasStream({
     }
   }, [optimisticSuggestions, mutate]);
 
-  const { width: windowWidth = 1920, height: windowHeight = 1080 } =
-    useWindowSize();
-
   useEffect(() => {
     const mostRecentDelta = streamingData?.at(-1);
     if (!mostRecentDelta) return;
@@ -42,23 +38,19 @@ export function useCanvasStream({
     const delta = mostRecentDelta as StreamingDelta;
 
     setCanvas((draftCanvas) => {
-      if (!draftCanvas) {
-        return {
-          content: '',
-          title: '',
-          isVisible: false,
-          documentId: delta.type === 'id' ? (delta.content as string) : '',
-          status: 'idle',
-          boundingBox: {
-            top: windowHeight / 4,
-            left: windowWidth / 4,
-            width: 250,
-            height: 50,
-          },
-        };
-      }
-
       switch (delta.type) {
+        case 'id':
+          return {
+            ...draftCanvas,
+            documentId: delta.content as string,
+          };
+
+        case 'title':
+          return {
+            ...draftCanvas,
+            title: delta.content as string,
+          };
+
         case 'text-delta':
           return {
             ...draftCanvas,
@@ -70,12 +62,6 @@ export function useCanvasStream({
             status: 'streaming',
           };
 
-        case 'title':
-          return {
-            ...draftCanvas,
-            title: delta.content as string,
-          };
-
         case 'suggestion':
           setTimeout(() => {
             setOptimisticSuggestions((currentSuggestions) => [

components/custom/use-click-outside.tsx

@@ -1,31 +0,0 @@
-import { useEffect, RefObject } from 'react';
-
-type Handler = (event: MouseEvent | TouchEvent) => void;
-
-function useClickOutside<T extends HTMLElement = HTMLElement>(
-  ref: RefObject<T>,
-  handler: Handler,
-  mouseEvent: 'mousedown' | 'mouseup' = 'mousedown'
-): void {
-  useEffect(() => {
-    const listener = (event: MouseEvent | TouchEvent) => {
-      const el = ref?.current;
-
-      if (!el || el.contains((event?.target as Node) || null)) {
-        return;
-      }
-
-      handler(event);
-    };
-
-    document.addEventListener(mouseEvent, listener);
-    document.addEventListener('touchstart', listener);
-
-    return () => {
-      document.removeEventListener(mouseEvent, listener);
-      document.removeEventListener('touchstart', listener);
-    };
-  }, [ref, handler, mouseEvent]);
-}
-
-export default useClickOutside;

components/custom/use-debounce.tsx

@@ -1,21 +0,0 @@
-import { useCallback, useRef } from "react";
-
-export function useDebounce<T extends (...args: any[]) => any>(
-  callback: T,
-  delay: number,
-): (...args: Parameters<T>) => void {
-  const timeoutRef = useRef<NodeJS.Timeout | null>(null);
-
-  return useCallback(
-    (...args: Parameters<T>) => {
-      if (timeoutRef.current) {
-        clearTimeout(timeoutRef.current);
-      }
-
-      timeoutRef.current = setTimeout(() => {
-        callback(...args);
-      }, delay);
-    },
-    [callback, delay],
-  );
-}

components/custom/use-window-size.tsx

@@ -1,39 +0,0 @@
-"use client";
-
-import { useState, useEffect } from "react";
-
-interface WindowSize {
-  width: number | undefined;
-  height: number | undefined;
-}
-
-function useWindowSize(): WindowSize {
-  const [windowSize, setWindowSize] = useState<WindowSize>({
-    width: undefined,
-    height: undefined,
-  });
-
-  useEffect(() => {
-    // Handler to call on window resize
-    function handleResize() {
-      // Set window width/height to state
-      setWindowSize({
-        width: window.innerWidth,
-        height: window.innerHeight,
-      });
-    }
-
-    // Add event listener
-    window.addEventListener("resize", handleResize);
-
-    // Call handler right away so state gets updated with initial window size
-    handleResize();
-
-    // Remove event listener on cleanup
-    return () => window.removeEventListener("resize", handleResize);
-  }, []); // Empty array ensures that effect is only run on mount and unmount
-
-  return windowSize;
-}
-
-export default useWindowSize;

components/custom/version-footer.tsx

@@ -0,0 +1,107 @@
+'use client';
+
+import { isAfter } from 'date-fns';
+import { motion } from 'framer-motion';
+import { useState } from 'react';
+import { useSWRConfig } from 'swr';
+import { useWindowSize } from 'usehooks-ts';
+
+import { Document } from '@/db/schema';
+import { getDocumentTimestampByIndex } from '@/lib/utils';
+
+import { UICanvas } from './canvas';
+import { LoaderIcon } from './icons';
+import { Button } from '../ui/button';
+
+interface VersionFooterProps {
+  canvas: UICanvas;
+  handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
+  documents: Array<Document> | undefined;
+  currentVersionIndex: number;
+}
+
+export const VersionFooter = ({
+  canvas,
+  handleVersionChange,
+  documents,
+  currentVersionIndex,
+}: VersionFooterProps) => {
+  const { width } = useWindowSize();
+  const isMobile = width < 768;
+
+  const { mutate } = useSWRConfig();
+  const [isMutating, setIsMutating] = useState(false);
+
+  if (!documents) return;
+
+  return (
+    <motion.div
+      className="absolute flex flex-col gap-4 lg:flex-row bottom-0 bg-background p-4 w-full border-t z-50 justify-between"
+      initial={{ y: isMobile ? 200 : 77 }}
+      animate={{ y: 0 }}
+      exit={{ y: isMobile ? 200 : 77 }}
+      transition={{ type: 'spring', stiffness: 140, damping: 20 }}
+    >
+      <div>
+        <div>You are viewing a previous version</div>
+        <div className="text-muted-foreground text-sm">
+          Restore this version to make edits
+        </div>
+      </div>
+
+      <div className="flex flex-row gap-4">
+        <Button
+          disabled={isMutating}
+          onClick={async () => {
+            setIsMutating(true);
+
+            mutate(
+              `/api/document?id=${canvas.documentId}`,
+              await fetch(`/api/document?id=${canvas.documentId}`, {
+                method: 'PATCH',
+                body: JSON.stringify({
+                  timestamp: getDocumentTimestampByIndex(
+                    documents,
+                    currentVersionIndex
+                  ),
+                }),
+              }),
+              {
+                optimisticData: documents
+                  ? [
+                      ...documents.filter((document) =>
+                        isAfter(
+                          new Date(document.createdAt),
+                          new Date(
+                            getDocumentTimestampByIndex(
+                              documents,
+                              currentVersionIndex
+                            )
+                          )
+                        )
+                      ),
+                    ]
+                  : [],
+              }
+            );
+          }}
+        >
+          <div>Restore this version</div>
+          {isMutating && (
+            <div className="animate-spin">
+              <LoaderIcon />
+            </div>
+          )}
+        </Button>
+        <Button
+          variant="outline"
+          onClick={() => {
+            handleVersionChange('latest');
+          }}
+        >
+          Back to latest version
+        </Button>
+      </div>
+    </motion.div>
+  );
+};

lib/editor/config.ts

@@ -0,0 +1,47 @@
+import { textblockTypeInputRule } from 'prosemirror-inputrules';
+import { Schema } from 'prosemirror-model';
+import { schema } from 'prosemirror-schema-basic';
+import { addListNodes } from 'prosemirror-schema-list';
+import { Transaction } from 'prosemirror-state';
+import { EditorView } from 'prosemirror-view';
+import { MutableRefObject } from 'react';
+
+import { buildContentFromDocument } from './functions';
+
+export const documentSchema = new Schema({
+  nodes: addListNodes(schema.spec.nodes, 'paragraph block*', 'block'),
+  marks: schema.spec.marks,
+});
+
+export function headingRule(level: number) {
+  return textblockTypeInputRule(
+    new RegExp(`^(#{1,${level}})\\s$`),
+    documentSchema.nodes.heading,
+    () => ({ level })
+  );
+}
+
+export const handleTransaction = ({
+  transaction,
+  editorRef,
+  saveContent,
+}: {
+  transaction: Transaction;
+  editorRef: MutableRefObject<EditorView | null>;
+  saveContent: (updatedContent: string, debounce: boolean) => void;
+}) => {
+  if (!editorRef || !editorRef.current) return;
+
+  const newState = editorRef.current.state.apply(transaction);
+  editorRef.current.updateState(newState);
+
+  if (transaction.docChanged && !transaction.getMeta('no-save')) {
+    const updatedContent = buildContentFromDocument(newState.doc);
+
+    if (transaction.getMeta('no-debounce')) {
+      saveContent(updatedContent, false);
+    } else {
+      saveContent(updatedContent, true);
+    }
+  }
+};

lib/editor/functions.tsx

@@ -0,0 +1,62 @@
+'use client';
+
+import { defaultMarkdownSerializer } from 'prosemirror-markdown';
+import { DOMParser, Node } from 'prosemirror-model';
+import { Decoration, DecorationSet, EditorView } from 'prosemirror-view';
+import { renderToString } from 'react-dom/server';
+
+import { Markdown } from '@/components/custom/markdown';
+
+import { documentSchema } from './config';
+import { createSuggestionWidget, UISuggestion } from './suggestions';
+
+export const buildDocumentFromContent = (content: string) => {
+  const parser = DOMParser.fromSchema(documentSchema);
+  const stringFromMarkdown = renderToString(<Markdown>{content}</Markdown>);
+  const tempContainer = document.createElement('div');
+  tempContainer.innerHTML = stringFromMarkdown;
+  return parser.parse(tempContainer);
+};
+
+export const buildContentFromDocument = (document: Node) => {
+  return defaultMarkdownSerializer.serialize(document);
+};
+
+export const createDecorations = (
+  suggestions: Array<UISuggestion>,
+  view: EditorView
+) => {
+  const decorations: Array<Decoration> = [];
+
+  suggestions.forEach((suggestion) => {
+    decorations.push(
+      Decoration.inline(
+        suggestion.selectionStart,
+        suggestion.selectionEnd,
+        {
+          class: 'suggestion-highlight',
+        },
+        {
+          suggestionId: suggestion.id,
+          type: 'highlight',
+        }
+      )
+    );
+
+    decorations.push(
+      Decoration.widget(
+        suggestion.selectionStart,
+        (view) => {
+          const { dom } = createSuggestionWidget(suggestion, view);
+          return dom;
+        },
+        {
+          suggestionId: suggestion.id,
+          type: 'widget',
+        }
+      )
+    );
+  });
+
+  return DecorationSet.create(view.state.doc, decorations);
+};

lib/editor/suggestions.tsx

@@ -1,4 +1,3 @@
-import { defaultMarkdownSerializer } from 'prosemirror-markdown';
 import { Node } from 'prosemirror-model';
 import { PluginKey, Plugin } from 'prosemirror-state';
 import { Decoration, DecorationSet, EditorView } from 'prosemirror-view';
@@ -40,7 +39,7 @@ function findPositionsInDoc(doc: Node, searchText: string): Position | null {
   return positions;
 }
 
-export function projectWithHighlights(
+export function projectWithPositions(
   doc: Node,
   suggestions: Array<Suggestion>
 ): Array<UISuggestion> {

lib/utils.ts

@@ -9,7 +9,7 @@ import {
 import { clsx, type ClassValue } from 'clsx';
 import { twMerge } from 'tailwind-merge';
 
-import { Chat } from '@/db/schema';
+import { Chat, Document } from '@/db/schema';
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
@@ -221,3 +221,13 @@ export function sanitizeUIMessages(messages: Array<Message>): Array<Message> {
       (message.toolInvocations && message.toolInvocations.length > 0)
   );
 }
+
+export function getDocumentTimestampByIndex(
+  documents: Array<Document>,
+  index: number
+) {
+  if (!documents) return new Date();
+  if (index > documents.length) return new Date();
+
+  return documents[index].createdAt;
+}