From f372c86039bd38d1e2f561cf688f8143dd68d535 Mon Sep 17 00:00:00 2001
From: Trang Doan <doantranghp2000@gmail.com>
Date: Thu, 2 Apr 2026 16:21:03 -0400
Subject: [PATCH] changes

---
 apps/roam/src/components/canvas/Tldraw.tsx    | 15 ++--
 ...tlesOnLoad.ts => syncCanvasNodesOnLoad.ts} | 79 ++++++++++++++++++-
 2 files changed, 85 insertions(+), 9 deletions(-)
 rename apps/roam/src/utils/{syncCanvasNodeTitlesOnLoad.ts => syncCanvasNodesOnLoad.ts} (64%)

diff --git a/apps/roam/src/components/canvas/Tldraw.tsx b/apps/roam/src/components/canvas/Tldraw.tsx
index d19bfea5e..fce4abd2b 100644
--- a/apps/roam/src/components/canvas/Tldraw.tsx
+++ b/apps/roam/src/components/canvas/Tldraw.tsx
@@ -93,7 +93,6 @@ import ToastListener, { dispatchToastEvent } from "./ToastListener";
 import { CanvasDrawerPanel } from "./CanvasDrawer";
 import { ClipboardPanel, ClipboardProvider } from "./Clipboard";
 import internalError from "~/utils/internalError";
-import { syncCanvasNodeTitlesOnLoad } from "~/utils/syncCanvasNodeTitlesOnLoad";
 import { AUTO_CANVAS_RELATIONS_KEY } from "~/data/userSettings";
 import { getSetting } from "~/utils/extensionSettings";
 import { isPluginTimerReady, waitForPluginTimer } from "~/utils/pluginTimer";
@@ -117,6 +116,7 @@ import {
 } from "./useCanvasStoreAdapterArgs";
 import posthog from "posthog-js";
 import { json, normalizeProps } from "~/utils/getBlockProps";
+import { syncCanvasNodesOnLoad } from "~/utils/syncCanvasNodesOnLoad";
 
 declare global {
   // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
@@ -1003,14 +1003,15 @@ const TldrawCanvasShared = ({
 
               appRef.current = app;
 
-              void syncCanvasNodeTitlesOnLoad(
-                app,
-                allNodes.map((n) => n.type),
-                allRelationIds,
-              ).catch((error) => {
+              void syncCanvasNodesOnLoad({
+                editor: app,
+                nodeTypeIds: allNodes.map((n) => n.type),
+                relationShapeTypeIds: allRelationIds,
+                extensionAPI,
+              }).catch((error) => {
                 internalError({
                   error,
-                  type: "Canvas: Sync node titles on load",
+                  type: "Canvas: Sync nodes on load",
                 });
               });
 
diff --git a/apps/roam/src/utils/syncCanvasNodeTitlesOnLoad.ts b/apps/roam/src/utils/syncCanvasNodesOnLoad.ts
similarity index 64%
rename from apps/roam/src/utils/syncCanvasNodeTitlesOnLoad.ts
rename to apps/roam/src/utils/syncCanvasNodesOnLoad.ts
index acd74df5d..3ae395e91 100644
--- a/apps/roam/src/utils/syncCanvasNodeTitlesOnLoad.ts
+++ b/apps/roam/src/utils/syncCanvasNodesOnLoad.ts
@@ -1,5 +1,7 @@
 import type { Editor } from "tldraw";
+import type { OnloadArgs } from "roamjs-components/types";
 import type { DiscourseNodeShape } from "~/components/canvas/DiscourseNodeUtil";
+import calcCanvasNodeSizeAndImg from "./calcCanvasNodeSizeAndImg";
 
 /**
  * Query Roam for current :node/title or :block/string for each uid.
@@ -56,11 +58,37 @@ const deleteNodeShapeAndRelations = (
  * - Updates shapes whose title changed
  * - Removes shapes whose uid no longer exists in the graph
  */
+export const syncCanvasNodesOnLoad = async ({
+  editor,
+  nodeTypeIds,
+  relationShapeTypeIds,
+  extensionAPI,
+}: {
+  editor: Editor;
+  nodeTypeIds: string[];
+  relationShapeTypeIds: string[];
+  extensionAPI: OnloadArgs["extensionAPI"];
+}): Promise<void> => {
+  const { discourseNodeShapes, uidToTitle } = await syncCanvasNodeTitlesOnLoad(
+    editor,
+    nodeTypeIds,
+    relationShapeTypeIds,
+  );
+  await syncCanvasKeyImagesOnLoad({
+    editor,
+    discourseNodeShapes,
+    uidToTitle,
+    extensionAPI,
+  });
+};
 export const syncCanvasNodeTitlesOnLoad = async (
   editor: Editor,
   nodeTypeIds: string[],
   relationShapeTypeIds: string[],
-): Promise<void> => {
+): Promise<{
+  discourseNodeShapes: DiscourseNodeShape[];
+  uidToTitle: Map<string, string>;
+}> => {
   const nodeTypeSet = new Set(nodeTypeIds);
   const relationIds = new Set(relationShapeTypeIds);
   const allRecords = editor.store.allRecords();
@@ -72,7 +100,8 @@ export const syncCanvasNodeTitlesOnLoad = async (
   ) as DiscourseNodeShape[];
 
   const uids = [...new Set(discourseNodeShapes.map((s) => s.props.uid))];
-  if (uids.length === 0) return;
+  if (uids.length === 0)
+    return { discourseNodeShapes: [], uidToTitle: new Map() };
 
   const uidToTitle = await queryTitlesByUids(uids);
 
@@ -104,4 +133,50 @@ export const syncCanvasNodeTitlesOnLoad = async (
       })),
     );
   }
+
+  return { discourseNodeShapes, uidToTitle };
+};
+
+const syncCanvasKeyImagesOnLoad = async ({
+  editor,
+  discourseNodeShapes,
+  uidToTitle,
+  extensionAPI,
+}: {
+  editor: Editor;
+  discourseNodeShapes: DiscourseNodeShape[];
+  uidToTitle: Map<string, string>;
+  extensionAPI: OnloadArgs["extensionAPI"];
+}): Promise<void> => {
+  const survivingShapes = discourseNodeShapes.filter((s) =>
+    uidToTitle.has(s.props.uid),
+  );
+  const imageUpdates: {
+    id: DiscourseNodeShape["id"];
+    type: string;
+    props: { imageUrl: string; w: number; h: number };
+  }[] = [];
+
+  await Promise.all(
+    survivingShapes.map(async (shape) => {
+      const title = uidToTitle.get(shape.props.uid) ?? shape.props.title ?? "";
+      const { w, h, imageUrl } = await calcCanvasNodeSizeAndImg({
+        nodeText: title,
+        uid: shape.props.uid,
+        nodeType: shape.type,
+        extensionAPI,
+      });
+      if ((shape.props.imageUrl ?? "") !== imageUrl) {
+        imageUpdates.push({
+          id: shape.id,
+          type: shape.type,
+          props: { imageUrl, w, h },
+        });
+      }
+    }),
+  );
+
+  if (imageUpdates.length > 0) {
+    editor.updateShapes(imageUpdates);
+  }
 };