From 369db3fea552838d87a701fbb052b7436f8dfe0a Mon Sep 17 00:00:00 2001
From: Peng Xiao <pengxiao@outlook.com>
Date: Wed, 27 Sep 2023 09:37:00 +0800
Subject: [PATCH] fix(component): cmdk flaky (#4512)

---
 .../block-suite-page-list/utils.tsx           |  4 ++--
 apps/core/src/components/pure/cmdk/data.tsx   | 20 ++++++++-----------
 .../favorite/add-favourite-button.tsx         |  5 +++--
 .../quick-search-main.stories.tsx             | 12 +++++++++--
 4 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/apps/core/src/components/blocksuite/block-suite-page-list/utils.tsx b/apps/core/src/components/blocksuite/block-suite-page-list/utils.tsx
index 940d57f10a..8180d0add1 100644
--- a/apps/core/src/components/blocksuite/block-suite-page-list/utils.tsx
+++ b/apps/core/src/components/blocksuite/block-suite-page-list/utils.tsx
@@ -19,14 +19,14 @@ export const usePageHelper = (blockSuiteWorkspace: BlockSuiteWorkspace) => {
   );
   const setPageMode = useSetAtom(setPageModeAtom);
   const createPageAndOpen = useCallback(
-    (id?: string, mode?: 'page' | 'edgeless'): string => {
+    (id?: string, mode?: 'page' | 'edgeless') => {
       const page = createPage(id);
       initEmptyPage(page).catch(error => {
         toast(`Failed to initialize Page: ${error.message}`);
       });
       setPageMode(page.id, mode || 'page');
       openPage(blockSuiteWorkspace.id, page.id);
-      return page.id;
+      return page;
     },
     [blockSuiteWorkspace.id, createPage, openPage, setPageMode]
   );
diff --git a/apps/core/src/components/pure/cmdk/data.tsx b/apps/core/src/components/pure/cmdk/data.tsx
index 00dbcc86cf..a03a1d9b6b 100644
--- a/apps/core/src/components/pure/cmdk/data.tsx
+++ b/apps/core/src/components/pure/cmdk/data.tsx
@@ -243,12 +243,10 @@ export const usePageCommands = () => {
           ),
           value: 'affine::create-page' + query, // hack to make the page always showing in the search result
           category: 'affine:creation',
-          run: () => {
-            const pageId = pageHelper.createPage();
-            // need to wait for the page to be created
-            setTimeout(() => {
-              pageMetaHelper.setPageTitle(pageId, query);
-            });
+          run: async () => {
+            const page = pageHelper.createPage();
+            await page.waitForLoaded();
+            pageMetaHelper.setPageTitle(page.id, query);
           },
           icon: <PageIcon />,
         });
@@ -265,12 +263,10 @@ export const usePageCommands = () => {
           ),
           value: 'affine::create-edgeless' + query, // hack to make the page always showing in the search result
           category: 'affine:creation',
-          run: () => {
-            const pageId = pageHelper.createEdgeless();
-            // need to wait for the page to be created
-            setTimeout(() => {
-              pageMetaHelper.setPageTitle(pageId, query);
-            });
+          run: async () => {
+            const page = pageHelper.createEdgeless();
+            await page.waitForLoaded();
+            pageMetaHelper.setPageTitle(page.id, query);
           },
           icon: <EdgelessIcon />,
         });
diff --git a/apps/core/src/components/pure/workspace-slider-bar/favorite/add-favourite-button.tsx b/apps/core/src/components/pure/workspace-slider-bar/favorite/add-favourite-button.tsx
index d585b80aec..919677db5d 100644
--- a/apps/core/src/components/pure/workspace-slider-bar/favorite/add-favourite-button.tsx
+++ b/apps/core/src/components/pure/workspace-slider-bar/favorite/add-favourite-button.tsx
@@ -14,8 +14,9 @@ export const AddFavouriteButton = ({ workspace }: AddFavouriteButtonProps) => {
   const { createPage } = usePageHelper(workspace);
   const { setPageMeta } = usePageMetaHelper(workspace);
   const handleAddFavorite = useCallback(async () => {
-    const id = createPage();
-    setPageMeta(id, { favorite: true });
+    const page = createPage();
+    await page.waitForLoaded();
+    setPageMeta(page.id, { favorite: true });
   }, [createPage, setPageMeta]);
 
   return (
diff --git a/apps/storybook/src/stories/quick-search/quick-search-main.stories.tsx b/apps/storybook/src/stories/quick-search/quick-search-main.stories.tsx
index 45c1384172..06427b35c1 100644
--- a/apps/storybook/src/stories/quick-search/quick-search-main.stories.tsx
+++ b/apps/storybook/src/stories/quick-search/quick-search-main.stories.tsx
@@ -3,6 +3,7 @@ import { WorkspaceFlavour } from '@affine/env/workspace';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { getOrCreateWorkspace } from '@affine/workspace/manager';
+import type { Page } from '@blocksuite/store';
 import type { Meta, StoryFn } from '@storybook/react';
 import { currentWorkspaceIdAtom } from '@toeverything/infra/atom';
 import {
@@ -21,6 +22,13 @@ export default {
   },
 } satisfies Meta;
 
+const createMockedPage = () => {
+  return {
+    id: 'test-page',
+    waitForLoaded: () => Promise.resolve(),
+  } as any as Page;
+};
+
 function useRegisterCommands() {
   const t = useAFFiNEI18N();
   const store = useStore();
@@ -39,8 +47,8 @@ function useRegisterCommands() {
         t,
         store,
         pageHelper: {
-          createEdgeless: () => 'noop',
-          createPage: () => 'noop',
+          createEdgeless: createMockedPage,
+          createPage: createMockedPage,
           importFile: () => Promise.resolve(),
           isPreferredEdgeless: () => false,
         },