fix(core): cleanup layout when switch page (#3794)

apps/core/src/bootstrap/plugins/setup.ts

@@ -106,6 +106,11 @@ const deleteLayoutAtom = atom<null, [string], void>(null, (_, set, id) => {
   });
 });
 
+// clean up plugin windows when switching to other pages
+rootStore.sub(currentPageAtom, () => {
+  rootStore.set(contentLayoutAtom, 'editor');
+});
+
 // module -> importName -> updater[]
 export const _rootImportsMap = new Map<string, Map<string, any>>();
 const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {

apps/core/src/components/page-detail-editor.tsx

@@ -17,7 +17,15 @@ import { contentLayoutAtom, rootStore } from '@toeverything/infra/atom';
 import clsx from 'clsx';
 import { useAtomValue, useSetAtom } from 'jotai';
 import type { CSSProperties, ReactElement } from 'react';
-import { memo, startTransition, Suspense, useCallback, useMemo } from 'react';
+import {
+  memo,
+  startTransition,
+  Suspense,
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+} from 'react';
 import { Panel, PanelGroup, PanelResizeHandle } from 'react-resizable-panels';
 
 import { pageSettingFamily } from '../atoms';
@@ -134,27 +142,42 @@ interface PluginContentAdapterProps {
 
 const PluginContentAdapter = memo<PluginContentAdapterProps>(
   function PluginContentAdapter({ windowItem, pluginName }) {
-    return (
-      <div
-        className={pluginContainer}
-        ref={useCallback(
-          (ref: HTMLDivElement | null) => {
-            if (ref) {
-              startTransition(() => {
-                const div = document.createElement('div');
-                const cleanup = windowItem(div);
-                ref.appendChild(div);
-                addCleanup(pluginName, () => {
-                  cleanup();
-                  ref.removeChild(div);
-                });
+    const rootRef = useRef<HTMLDivElement | null>(null);
+    useEffect(() => {
+      const abortController = new AbortController();
+      const root = rootRef.current;
+      if (root) {
+        startTransition(() => {
+          if (abortController.signal.aborted) {
+            return;
+          }
+          const div = document.createElement('div');
+          const cleanup = windowItem(div);
+          root.appendChild(div);
+          if (abortController.signal.aborted) {
+            cleanup();
+            root.removeChild(div);
+          } else {
+            const cl = () => {
+              cleanup();
+              root.removeChild(div);
+            };
+            const dispose = addCleanup(pluginName, cl);
+            abortController.signal.addEventListener('abort', () => {
+              setTimeout(() => {
+                dispose();
+                cl();
               });
-            }
-          },
-          [pluginName, windowItem]
-        )}
-      />
-    );
+            });
+          }
+        });
+        return () => {
+          abortController.abort();
+        };
+      }
+      return;
+    }, [pluginName, windowItem]);
+    return <div className={pluginContainer} ref={rootRef} />;
   }
 );
 

packages/infra/src/__internal__/plugin.ts

@@ -1,4 +1,5 @@
 import type { CallbackMap } from '@affine/sdk/entry';
+import { assertExists } from '@blocksuite/global/utils';
 import { atomWithStorage } from 'jotai/utils';
 import { atom } from 'jotai/vanilla';
 import type { z } from 'zod';
@@ -14,13 +15,21 @@ export const builtinPluginPaths = new Set([
   '/plugins/outline',
 ]);
 
-const pluginCleanupMap = new Map<string, (() => void)[]>();
+const pluginCleanupMap = new Map<string, Set<() => void>>();
 
-export function addCleanup(pluginName: string, cleanup: () => void) {
+export function addCleanup(
+  pluginName: string,
+  cleanup: () => void
+): () => void {
   if (!pluginCleanupMap.has(pluginName)) {
-    pluginCleanupMap.set(pluginName, []);
+    pluginCleanupMap.set(pluginName, new Set());
   }
-  pluginCleanupMap.get(pluginName)?.push(cleanup);
+  const cleanupSet = pluginCleanupMap.get(pluginName);
+  assertExists(cleanupSet);
+  cleanupSet.add(cleanup);
+  return () => {
+    cleanupSet.delete(cleanup);
+  };
 }
 
 export function invokeCleanup(pluginName: string) {

plugins/outline/src/app.tsx

@@ -1,42 +1,30 @@
 import { Tooltip } from '@affine/component';
-import { deleteLayoutAtom, pushLayoutAtom } from '@affine/sdk/entry';
+import {
+  currentPageAtom,
+  deleteLayoutAtom,
+  pushLayoutAtom,
+} from '@affine/sdk/entry';
 import { TOCNotesPanel } from '@blocksuite/blocks';
+import { assertExists } from '@blocksuite/global/utils';
 import { RightSidebarIcon } from '@blocksuite/icons';
-import type { Page } from '@blocksuite/store';
 import { IconButton } from '@toeverything/components/button';
-import { useAtom, useSetAtom } from 'jotai';
+import { useAtomValue, useSetAtom } from 'jotai';
 import type { ComponentType, PropsWithChildren } from 'react';
-import { useCallback, useEffect, useRef, useState } from 'react';
+import { useCallback, useRef, useState } from 'react';
 import { createRoot } from 'react-dom/client';
 
-import { blocksuiteRootAtom } from './atom';
-
 const Outline = () => {
-  const ref = useRef<HTMLDivElement>(null);
   const tocPanelRef = useRef<TOCNotesPanel | null>(null);
-  const [blocksuite] = useAtom(blocksuiteRootAtom);
+  const currentPage = useAtomValue(currentPageAtom);
 
   if (!tocPanelRef.current) {
     tocPanelRef.current = new TOCNotesPanel();
   }
 
-  if (blocksuite?.page !== tocPanelRef.current?.page) {
-    (tocPanelRef.current as TOCNotesPanel).page = blocksuite?.page as Page;
+  if (currentPage !== tocPanelRef.current?.page) {
+    (tocPanelRef.current as TOCNotesPanel).page = currentPage;
   }
 
-  useEffect(() => {
-    if (!ref.current || !tocPanelRef.current) return;
-
-    const container = ref.current;
-    const tocPanel = tocPanelRef.current as TOCNotesPanel;
-
-    container.appendChild(tocPanel);
-
-    return () => {
-      container.removeChild(tocPanel);
-    };
-  }, []);
-
   return (
     <div
       className={`outline-wrapper`}
@@ -44,7 +32,12 @@ const Outline = () => {
         height: '100%',
         borderLeft: `1px solid var(--affine-border-color)`,
       }}
-      ref={ref}
+      ref={useCallback((container: HTMLDivElement | null) => {
+        if (container) {
+          assertExists(tocPanelRef.current);
+          container.appendChild(tocPanelRef.current);
+        }
+      }, [])}
     />
   );
 };

plugins/outline/src/atom.ts

@@ -1,5 +0,0 @@
-import { atom } from 'jotai';
-
-export const blocksuiteRootAtom = atom(() =>
-  document.querySelector('block-suite-root')
-);