diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
index 383a7c64d6..56c0324146 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
@@ -1,7 +1,6 @@
 import type { BlockSpec } from '@blocksuite/block-std';
 import {
   BookmarkBlockSpec,
-  CodeBlockSpec,
   DatabaseBlockSpec,
   DataViewBlockSpec,
   DividerBlockSpec,
@@ -12,11 +11,14 @@ import {
   EmbedLoomBlockSpec,
   EmbedSyncedDocBlockSpec,
   EmbedYoutubeBlockSpec,
-  ImageBlockSpec,
   ListBlockSpec,
   NoteBlockSpec,
 } from '@blocksuite/blocks';
-import { AIParagraphBlockSpec } from '@blocksuite/presets';
+import {
+  AICodeBlockSpec,
+  AIImageBlockSpec,
+  AIParagraphBlockSpec,
+} from '@blocksuite/presets';
 
 import { CustomAttachmentBlockSpec } from './custom/attachment-block';
 
@@ -26,8 +28,6 @@ export const CommonBlockSpecs: BlockSpec[] = [
   DatabaseBlockSpec,
   DataViewBlockSpec,
   DividerBlockSpec,
-  CodeBlockSpec,
-  ImageBlockSpec,
   BookmarkBlockSpec,
   EmbedFigmaBlockSpec,
   EmbedGithubBlockSpec,
@@ -38,5 +38,7 @@ export const CommonBlockSpecs: BlockSpec[] = [
   EmbedLinkedDocBlockSpec,
   // special
   CustomAttachmentBlockSpec,
+  AICodeBlockSpec,
+  AIImageBlockSpec,
   AIParagraphBlockSpec,
 ];
diff --git a/scripts/setup/vi-mock.ts b/scripts/setup/vi-mock.ts
index bc7d985cfc..ed922b2805 100644
--- a/scripts/setup/vi-mock.ts
+++ b/scripts/setup/vi-mock.ts
@@ -23,6 +23,8 @@ vi.mock('@blocksuite/presets', () => ({
     provide: vi.fn(),
   },
   AIEdgelessRootBlockSpec: {},
+  AICodeBlockSpec: {},
+  AIImageBlockSpec: {},
   AIParagraphBlockSpec: {},
   AIPageRootBlockSpec: {},
 }));