Clean up editor code and reorganize contents (#478)

2024-11-01 15:31:54 +05:30 · 2024-11-01 15:31:54 +05:30 · 5190b109c9
commit 5190b109c9
parent b3aa3cb18a
19 changed files with 475 additions and 444 deletions
--- a/app/globals.css
+++ b/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;
+}
--- a/components/custom/canvas-stream-handler.tsx
+++ b/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;
 }

--- a/components/custom/canvas.tsx
+++ b/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 { Button } from '../ui/button';
-
+import { VersionFooter } from './version-footer';
 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) =>
-          currentCanvas
-            ? {
-                ...currentCanvas,
-                content: mostRecentDocument.content ?? '',
-              }
-            : null
-        );
+        setCanvas((currentCanvas) => ({
+          ...currentCanvas,
+          content: mostRecentDocument.content ?? '',
+        }));
      }
    }
  }, [documents, setCanvas]);
@ -174,24 +168,24 @@ export function Canvas({
    [canvas, mutate]
  );

-  const debouncedHandleEditorChange = useCallback(
-    useDebounce(handleContentChange, 4000),
-    [handleContentChange]
+  const debouncedHandleContentChange = useDebounceCallback(
+    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
-      ? true
-      : documents && documents.length > 0
-        ? currentVersionIndex === documents.length - 1
-        : true;
+    documents && documents.length > 0
+      ? 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'}
-                onChange={handleEditorChange}
+                status={canvas.status}
+                saveContent={saveContent}
                suggestions={isCurrentVersion ? (suggestions ?? []) : []}
              />
            ) : (
@ -482,63 +473,12 @@ export function Canvas({

        <AnimatePresence>
          {!isCurrentVersion && (
-            <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
-                  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>
+            <VersionFooter
+              canvas={canvas}
+              currentVersionIndex={currentVersionIndex}
+              documents={documents}
+              handleVersionChange={handleVersionChange}
+            />
          )}
        </AnimatePresence>
      </motion.div>
--- a/components/custom/chat.tsx
+++ b/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>();
--- a/components/custom/document.tsx
+++ b/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;
-  setCanvas: (value: SetStateAction<UICanvas | null>) => void;
+  canvas: UICanvas;
+  setCanvas: (value: SetStateAction<UICanvas>) => void;
 }

 export function DocumentToolResult({
@ -42,27 +42,14 @@ export function DocumentToolResult({
          height: rect.height,
        };

-        if (!canvas) {
-          setCanvas({
-            documentId: result.id,
-            content: '',
-            title: result.title,
-            isVisible: true,
-            status: 'idle',
-            boundingBox,
-          });
-        } else {
-          if (canvas.documentId !== result.id) {
-            setCanvas({
-              documentId: result.id,
-              content: '',
-              title: result.title,
-              isVisible: true,
-              status: 'idle',
-              boundingBox,
-            });
-          }
-        }
+        setCanvas({
+          documentId: result.id,
+          content: '',
+          title: result.title,
+          isVisible: true,
+          status: 'idle',
+          boundingBox,
+        });
      }}
    >
      <div className="text-muted-foreground mt-1">
--- a/components/custom/editor.tsx
+++ b/components/custom/editor.tsx
@ -1,213 +1,170 @@
 'use client';

 import { exampleSetup } from 'prosemirror-example-setup';
-import { inputRules, textblockTypeInputRule } 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 { inputRules } from 'prosemirror-inputrules';
 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,
-  projectWithHighlights,
+  documentSchema,
+  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,
-  suggestions: suggestionsWithoutHighlights,
+  saveContent,
+  suggestions,
+  status,
 }: EditorProps) {
-  const editorRef = useRef<HTMLDivElement>(null);
-  const viewRef = useRef<EditorView | null>(null);
-  const widgetRootsRef = useRef<Map<string, WidgetRoot>>(new Map());
+  const containerRef = useRef<HTMLDivElement>(null);
+  const editorRef = useRef<EditorView | null>(null);

  useEffect(() => {
-    if (editorRef.current && !viewRef.current) {
-      if (mySchema) {
-        const parser = DOMParser.fromSchema(mySchema);
+    if (containerRef.current && !editorRef.current) {
+      const state = EditorState.create({
+        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>);
-
-        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');
-      }
+      editorRef.current = new EditorView(containerRef.current, {
+        state,
+      });
    }

    return () => {
-      if (viewRef.current) {
-        viewRef.current.destroy();
-        viewRef.current = null;
+      if (editorRef.current) {
+        editorRef.current.destroy();
+        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) {
-      const computedSuggestions = projectWithHighlights(
-        viewRef.current.state.doc,
-        suggestionsWithoutHighlights
+    if (editorRef.current) {
+      editorRef.current.setProps({
+        dispatchTransaction: (transaction) => {
+          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.selectionStart !== 0 && suggestion.selectionEnd !== 0
+        (suggestion) => suggestion.selectionStart && suggestion.selectionEnd
      );

      const decorations = createDecorations(
-        computedSuggestions,
-        viewRef.current
+        projectedSuggestions,
+        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) => {
-    const decorations: Decoration[] = [];
-
-    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} />;
+  return (
+    <div className="relative prose dark:prose-invert" ref={containerRef} />
+  );
 }

 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);
--- a/components/custom/message.tsx
+++ b/components/custom/message.tsx
@ -24,8 +24,8 @@ export const Message = ({
  content: string | ReactNode;
  toolInvocations: Array<ToolInvocation> | undefined;
  attachments?: Array<Attachment>;
-  canvas: UICanvas | null;
-  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  canvas: UICanvas;
+  setCanvas: Dispatch<SetStateAction<UICanvas>>;
 }) => {
  return (
    <motion.div
--- a/components/custom/multimodal-input.tsx
+++ b/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';

--- a/components/custom/suggestion.tsx
+++ b/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 ? (
-    <div
-      className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
-      onClick={() => {
-        setIsExpanded(true);
-      }}
-    >
-      <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"
+  return (
+    <AnimatePresence>
+      {!isExpanded ? (
+        <motion.div
+          className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
          onClick={() => {
-            setIsExpanded(false);
+            setIsExpanded(true);
          }}
+          whileHover={{ scale: 1.1 }}
        >
-          <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>
+          <MessageIcon size={windowWidth && windowWidth < 768 ? 16 : 14} />
+        </motion.div>
+      ) : (
+        <motion.div
+          key={suggestion.id}
+          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"
+          transition={{ type: 'spring', stiffness: 500, damping: 30 }}
+          initial={{ opacity: 0, y: -10 }}
+          animate={{ opacity: 1, y: -20 }}
+          exit={{ opacity: 0, y: -10 }}
+          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);
+              }}
+            >
+              <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>
  );
 };
--- a/components/custom/toolbar.tsx
+++ b/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);
  });
--- a/components/custom/use-canvas-stream.tsx
+++ b/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) => [
--- a/components/custom/use-click-outside.tsx
+++ b/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;
--- a/components/custom/use-debounce.tsx
+++ b/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],
-  );
-}
--- a/components/custom/use-window-size.tsx
+++ b/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;
--- a/components/custom/version-footer.tsx
+++ b/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>
+  );
+};
--- a/lib/editor/config.ts
+++ b/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);
+    }
+  }
+};
--- a/lib/editor/functions.tsx
+++ b/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);
+};
--- a/lib/editor/suggestions.tsx
+++ b/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> {
--- a/lib/utils.ts
+++ b/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;
+}