[ENG-598] Data persistence for tldraw (#303)

trangdoan982 · trangdoan982 · commit be63ce77c1fc · 2025-10-09T14:38:55.000-04:00
* data persistence to the file

* error handling

* address PR comments

* address some PR comments

* address other PR comments

* address PR comments
diff --git a/apps/obsidian/src/components/TldrawView.tsx b/apps/obsidian/src/components/TldrawView.tsx
@@ -2,15 +2,10 @@ import { TextFileView, TFile, WorkspaceLeaf } from "obsidian";
 import { VIEW_TYPE_TLDRAW_DG_PREVIEW } from "~/constants";
 import { Root, createRoot } from "react-dom/client";
 import { TldrawPreviewComponent } from "./TldrawViewComponent";
-import {
-  TLRecord,
-  TLStore,
-  createTLStore,
-  defaultShapeUtils,
-  loadSnapshot,
-} from "tldraw";
+import { TLStore } from "tldraw";
 import React from "react";
 import DiscourseGraphPlugin from "~/index";
+import { processInitialData, TLData } from "~/utils/tldraw";
 
 export class TldrawView extends TextFileView {
   plugin: DiscourseGraphPlugin;
@@ -71,17 +66,17 @@ export class TldrawView extends TextFileView {
 
     const fileData = await this.app.vault.read(file);
 
-    const store = await this.createStore(fileData);
+    const store = this.createStore(fileData);
 
     if (!store) {
       console.warn("No tldraw data found in file");
       return;
     }
 
-    await this.setStore(store);
+    this.setStore(store);
   }
 
-  private async createStore(fileData: string): Promise<TLStore | undefined> {
+  private createStore(fileData: string): TLStore | undefined {
     try {
       const match = fileData.match(
         /```json !!!_START_OF_TLDRAW_DG_DATA__DO_NOT_CHANGE_THIS_PHRASE_!!!([\s\S]*?)!!!_END_OF_TLDRAW_DG_DATA__DO_NOT_CHANGE_THIS_PHRASE_!!!\n```/,
@@ -92,39 +87,13 @@ export class TldrawView extends TextFileView {
         return;
       }
 
-      const data = JSON.parse(match[1]);
+      const data = JSON.parse(match[1]) as TLData;
       if (!data.raw) {
         console.warn("Invalid tldraw data format - missing raw field");
         return;
       }
 
-      const recordsData = Array.isArray(data.raw.records)
-        ? data.raw.records.reduce(
-            (
-              acc: Record<string, TLRecord>,
-              record: { id: string } & TLRecord,
-            ) => {
-              acc[record.id] = {
-                ...record,
-              };
-              return acc;
-            },
-            {},
-          )
-        : data.raw.records;
-
-      let store: TLStore;
-      if (recordsData) {
-        store = createTLStore({
-          shapeUtils: defaultShapeUtils,
-          initialData: recordsData,
-        });
-      } else {
-        store = createTLStore({
-          shapeUtils: defaultShapeUtils,
-        });
-        loadSnapshot(store, data.raw);
-      }
+      const { store } = processInitialData(data);
 
       return store;
     } catch (e) {
@@ -135,15 +104,21 @@ export class TldrawView extends TextFileView {
 
   private createReactRoot(entryPoint: Element, store: TLStore) {
     const root = createRoot(entryPoint);
+    if (!this.file) return;
+
     root.render(
       <React.StrictMode>
-        <TldrawPreviewComponent store={store} isReadonly={false} />
+        <TldrawPreviewComponent
+          store={store}
+          plugin={this.plugin}
+          file={this.file}
+        />
       </React.StrictMode>,
     );
     return root;
   }
 
-  protected async setStore(store: TLStore) {
+  protected setStore(store: TLStore) {
     if (this.store) {
       try {
         this.store.dispose();
@@ -154,11 +129,11 @@ export class TldrawView extends TextFileView {
 
     this.store = store;
     if (this.tldrawContainer) {
-      await this.refreshView();
+      this.refreshView();
     }
   }
 
-  private async refreshView() {
+  private refreshView() {
     if (!this.store) return;
 
     if (this.reactRoot) {
diff --git a/apps/obsidian/src/components/TldrawViewComponent.tsx b/apps/obsidian/src/components/TldrawViewComponent.tsx
@@ -1,18 +1,38 @@
 import { useCallback, useEffect, useRef, useState } from "react";
 import { Editor, ErrorBoundary, Tldraw, TLStore } from "tldraw";
 import "tldraw/tldraw.css";
+import {
+  getTLDataTemplate,
+  createRawTldrawFile,
+  getUpdatedMdContent,
+  TLData,
+  processInitialData,
+} from "~/utils/tldraw";
+import DiscourseGraphPlugin from "~/index";
+import {
+  DEFAULT_SAVE_DELAY,
+  TLDATA_DELIMITER_END,
+  TLDATA_DELIMITER_START,
+} from "~/constants";
+import { TFile } from "obsidian";
 
 interface TldrawPreviewProps {
   store: TLStore;
-  isReadonly?: boolean;
+  plugin: DiscourseGraphPlugin;
+  file: TFile;
 }
 
 export const TldrawPreviewComponent = ({
   store,
-  isReadonly = false,
+  plugin,
+  file,
 }: TldrawPreviewProps) => {
   const containerRef = useRef<HTMLDivElement>(null);
+  const [currentStore, setCurrentStore] = useState<TLStore>(store);
   const [isReady, setIsReady] = useState(false);
+  const editorRef = useRef<Editor | null>(null);
+  const saveTimeoutRef = useRef<NodeJS.Timeout>();
+  const lastSavedDataRef = useRef<string>("");
 
   useEffect(() => {
     const timer = setTimeout(() => {
@@ -21,14 +41,86 @@ export const TldrawPreviewComponent = ({
     return () => clearTimeout(timer);
   }, []);
 
-  const handleMount = useCallback((editor: Editor) => {
-    editor.setCurrentTool("hand");
-    editor.updateInstanceState({});
+  const saveChanges = useCallback(async () => {
+    const newData = getTLDataTemplate({
+      pluginVersion: plugin.manifest.version,
+      tldrawFile: createRawTldrawFile(currentStore),
+      uuid: window.crypto.randomUUID(),
+    });
+    const stringifiedData = JSON.stringify(newData, null, "\t");
+
+    if (stringifiedData === lastSavedDataRef.current) {
+      return;
+    }
+
+    const currentContent = await plugin.app.vault.read(file);
+    if (!currentContent) {
+      console.error("Could not read file content");
+      return;
+    }
 
-    const shapes = editor.getCurrentPageShapes();
-    if (shapes.length > 0) {
-      editor.zoomToFit();
+    const updatedString = getUpdatedMdContent(currentContent, stringifiedData);
+    if (updatedString === currentContent) {
+      return;
     }
+
+    try {
+      await plugin.app.vault.modify(file, updatedString);
+
+      const verifyContent = await plugin.app.vault.read(file);
+      const verifyMatch = verifyContent.match(
+        new RegExp(
+          `${TLDATA_DELIMITER_START}\\s*([\\s\\S]*?)\\s*${TLDATA_DELIMITER_END}`,
+        ),
+      );
+
+      if (!verifyMatch || verifyMatch[1]?.trim() !== stringifiedData.trim()) {
+        throw new Error("Failed to verify saved TLDraw data");
+      }
+
+      lastSavedDataRef.current = stringifiedData;
+    } catch (error) {
+      console.error("Error saving/verifying TLDraw data:", error);
+      // Reload the editor state from file since save failed
+      const fileContent = await plugin.app.vault.read(file);
+      const match = fileContent.match(
+        new RegExp(
+          `${TLDATA_DELIMITER_START}([\\s\\S]*?)${TLDATA_DELIMITER_END}`,
+        ),
+      );
+      if (match?.[1]) {
+        const data = JSON.parse(match[1]) as TLData;
+        const { store: newStore } = processInitialData(data);
+        setCurrentStore(newStore);
+      }
+    }
+  }, [file, plugin, currentStore]);
+
+  useEffect(() => {
+    const unsubscribe = currentStore.listen(
+      () => {
+        if (saveTimeoutRef.current) {
+          clearTimeout(saveTimeoutRef.current);
+        }
+        saveTimeoutRef.current = setTimeout(
+          () => void saveChanges(),
+          DEFAULT_SAVE_DELAY,
+        );
+      },
+      { source: "user", scope: "document" },
+    );
+
+    return () => {
+      unsubscribe();
+      if (saveTimeoutRef.current) {
+        clearTimeout(saveTimeoutRef.current);
+      }
+    };
+  }, [currentStore, saveChanges]);
+
+  const handleMount = useCallback((editor: Editor) => {
+    editorRef.current = editor;
+    editor.setCurrentTool("select");
   }, []);
 
   return (
@@ -43,7 +135,7 @@ export const TldrawPreviewComponent = ({
             <div>Error in Tldraw component: {JSON.stringify(error)}</div>
           )}
         >
-          <Tldraw store={store} onMount={handleMount} autoFocus={false} />
+          <Tldraw store={currentStore} onMount={handleMount} autoFocus={true} />
         </ErrorBoundary>
       ) : (
         <div>Loading Tldraw...</div>
diff --git a/apps/obsidian/src/utils/tldraw.ts b/apps/obsidian/src/utils/tldraw.ts
@@ -1,11 +1,11 @@
-import { createTLStore, TldrawFile, TLStore } from "tldraw";
+import { createTLStore, defaultShapeUtils, TldrawFile, TLRecord, TLStore } from "tldraw";
 import {
   FRONTMATTER_KEY,
   TLDATA_DELIMITER_END,
   TLDATA_DELIMITER_START,
   TLDRAW_VERSION,
 } from "../constants";
-import DiscourseGraphPlugin from "..";
+import DiscourseGraphPlugin from "~/index";
 import { checkAndCreateFolder, getNewUniqueFilepath } from "./file";
 import { Notice } from "obsidian";
 import { format } from "date-fns";
@@ -16,12 +16,47 @@ export type TldrawPluginMetaData = {
   uuid: string;
 };
 
+export type TldrawRawData = {
+  tldrawFileFormatVersion: number;
+  schema: any;
+  records: any;
+};
+
 export type TLData = {
   meta: TldrawPluginMetaData;
-  raw: {
-    tldrawFileFormatVersion: number;
-    schema: any;
-    records: any;
+  raw: TldrawRawData;
+};
+
+export const processInitialData = (
+  data: TLData,
+): { meta: TldrawPluginMetaData; store: TLStore } => {
+  const recordsData = Array.isArray(data.raw.records)
+    ? data.raw.records.reduce(
+        (acc: Record<string, TLRecord>, record: { id: string } & TLRecord) => {
+          acc[record.id] = {
+            ...record,
+          };
+          return acc;
+        },
+        {},
+      )
+    : data.raw.records;
+
+  let store: TLStore;
+  if (recordsData) {
+    store = createTLStore({
+      shapeUtils: defaultShapeUtils,
+      initialData: recordsData,
+    });
+  } else {
+    store = createTLStore({
+      shapeUtils: defaultShapeUtils,
+    });
+  }
+
+  return {
+    meta: data.meta,
+    store,
   };
 };
 
@@ -121,4 +156,23 @@ export const createCanvas = async (plugin: DiscourseGraphPlugin) => {
     new Notice(e instanceof Error ? e.message : "Failed to create canvas file");
     console.error(e);
   }
-};
+};
+
+/**
+ * Get the updated markdown content with the new TLData
+ * @param currentContent - The current markdown content
+ * @param stringifiedData - The new TLData stringified
+ * @returns The updated markdown content
+ */
+export const getUpdatedMdContent = (
+  currentContent: string,
+  stringifiedData: string,
+) => {
+  const regex = new RegExp(
+    `${TLDATA_DELIMITER_START}([\\s\\S]*?)${TLDATA_DELIMITER_END}`,
+  );
+  return currentContent.replace(
+    regex,
+    `${TLDATA_DELIMITER_START}\n${stringifiedData}\n${TLDATA_DELIMITER_END}`,
+  );
+};
