diff --git a/blocksuite/affine/block-embed/src/common/render-linked-doc.ts b/blocksuite/affine/block-embed/src/common/render-linked-doc.ts
index 1c312ca664..928f0c46d7 100644
--- a/blocksuite/affine/block-embed/src/common/render-linked-doc.ts
+++ b/blocksuite/affine/block-embed/src/common/render-linked-doc.ts
@@ -1,4 +1,4 @@
-import type { SurfaceBlockModel } from '@blocksuite/affine-block-surface';
+import { getSurfaceBlock } from '@blocksuite/affine-block-surface';
 import {
   type DocMode,
   type ImageBlockModel,
@@ -263,11 +263,6 @@ export function isEmptyNote(note: BlockModel) {
   });
 }
 
-function getSurfaceBlock(doc: Doc) {
-  const blocks = doc.getBlocksByFlavour('affine:surface');
-  return blocks.length !== 0 ? (blocks[0].model as SurfaceBlockModel) : null;
-}
-
 /**
  * Gets the document content with a max length.
  */
diff --git a/blocksuite/affine/block-surface/src/index.ts b/blocksuite/affine/block-surface/src/index.ts
index b15f8521c6..486cc0571b 100644
--- a/blocksuite/affine/block-surface/src/index.ts
+++ b/blocksuite/affine/block-surface/src/index.ts
@@ -104,7 +104,7 @@ import {
   tryMoveNode,
 } from './utils/mindmap/utils';
 export * from './extensions';
-export { getLastPropsKey } from './utils/get-last-props-key';
+export { getLastPropsKey, getSurfaceBlock } from './utils';
 export type { Options } from './utils/rough/core';
 export { sortIndex } from './utils/sort';
 export { updateXYWH } from './utils/update-xywh.js';
diff --git a/blocksuite/affine/block-surface/src/surface-service.ts b/blocksuite/affine/block-surface/src/surface-service.ts
index 030a34c902..eea96f973e 100644
--- a/blocksuite/affine/block-surface/src/surface-service.ts
+++ b/blocksuite/affine/block-surface/src/surface-service.ts
@@ -2,6 +2,7 @@ import { BlockService } from '@blocksuite/block-std';
 import { GfxControllerIdentifier } from '@blocksuite/block-std/gfx';
 
 import { type SurfaceBlockModel, SurfaceBlockSchema } from './surface-model.js';
+import { getSurfaceBlock } from './utils/get-surface-block.js';
 
 export class SurfaceBlockService extends BlockService {
   static override readonly flavour = SurfaceBlockSchema.model.flavour;
@@ -15,9 +16,10 @@ export class SurfaceBlockService extends BlockService {
   override mounted(): void {
     super.mounted();
 
-    this.surface = this.doc.getBlockByFlavour(
-      'affine:surface'
-    )[0] as SurfaceBlockModel;
+    const surface = getSurfaceBlock(this.doc);
+
+    // FIXME: BS-2271
+    this.surface = surface!;
 
     if (!this.surface) {
       const disposable = this.doc.slots.blockUpdated.on(payload => {
diff --git a/blocksuite/affine/block-surface/src/utils/get-surface-block.ts b/blocksuite/affine/block-surface/src/utils/get-surface-block.ts
new file mode 100644
index 0000000000..dc7a0cf6f0
--- /dev/null
+++ b/blocksuite/affine/block-surface/src/utils/get-surface-block.ts
@@ -0,0 +1,8 @@
+import type { Doc } from '@blocksuite/store';
+
+import type { SurfaceBlockModel } from '../surface-model';
+
+export function getSurfaceBlock(doc: Doc) {
+  const blocks = doc.getBlocksByFlavour('affine:surface');
+  return blocks.length !== 0 ? (blocks[0].model as SurfaceBlockModel) : null;
+}
diff --git a/blocksuite/affine/block-surface/src/utils/index.ts b/blocksuite/affine/block-surface/src/utils/index.ts
index d3cb8efd03..e9ae255b37 100644
--- a/blocksuite/affine/block-surface/src/utils/index.ts
+++ b/blocksuite/affine/block-surface/src/utils/index.ts
@@ -34,3 +34,4 @@ export function normalizeWheelDeltaY(delta: number, zoom = 1) {
 }
 
 export { getLastPropsKey } from './get-last-props-key';
+export { getSurfaceBlock } from './get-surface-block';
diff --git a/blocksuite/blocks/src/_common/utils/render-linked-doc.ts b/blocksuite/blocks/src/_common/utils/render-linked-doc.ts
index 7a323840d9..5c9c1be342 100644
--- a/blocksuite/blocks/src/_common/utils/render-linked-doc.ts
+++ b/blocksuite/blocks/src/_common/utils/render-linked-doc.ts
@@ -1,3 +1,4 @@
+import { getSurfaceBlock } from '@blocksuite/affine-block-surface';
 import type { FrameBlockModel, NoteBlockModel } from '@blocksuite/affine-model';
 import { NoteDisplayMode } from '@blocksuite/affine-model';
 import { DocModeProvider } from '@blocksuite/affine-shared/services';
@@ -15,7 +16,6 @@ import {
   isFrameBlock,
   isNoteBlock,
 } from '../../root-block/edgeless/utils/query.js';
-import { getSurfaceBlock } from '../../surface-ref-block/utils.js';
 
 export function addBlocksToDoc(
   targetDoc: Doc,
diff --git a/blocksuite/blocks/src/edgeless-text-block/commands/insert-edgeless-text.ts b/blocksuite/blocks/src/edgeless-text-block/commands/insert-edgeless-text.ts
index b7849be591..1e14219a10 100644
--- a/blocksuite/blocks/src/edgeless-text-block/commands/insert-edgeless-text.ts
+++ b/blocksuite/blocks/src/edgeless-text-block/commands/insert-edgeless-text.ts
@@ -1,10 +1,12 @@
-import { EdgelessCRUDIdentifier } from '@blocksuite/affine-block-surface';
+import {
+  EdgelessCRUDIdentifier,
+  getSurfaceBlock,
+} from '@blocksuite/affine-block-surface';
 import { focusTextModel } from '@blocksuite/affine-components/rich-text';
 import type { Command } from '@blocksuite/block-std';
 import { GfxControllerIdentifier } from '@blocksuite/block-std/gfx';
 import { Bound } from '@blocksuite/global/utils';
 
-import { getSurfaceBlock } from '../../surface-ref-block/utils.js';
 import {
   EDGELESS_TEXT_BLOCK_MIN_HEIGHT,
   EDGELESS_TEXT_BLOCK_MIN_WIDTH,
diff --git a/blocksuite/blocks/src/root-block/edgeless/edgeless-root-service.ts b/blocksuite/blocks/src/root-block/edgeless/edgeless-root-service.ts
index ad01a5b22a..5da8b12fee 100644
--- a/blocksuite/blocks/src/root-block/edgeless/edgeless-root-service.ts
+++ b/blocksuite/blocks/src/root-block/edgeless/edgeless-root-service.ts
@@ -3,6 +3,7 @@ import {
   EdgelessLegacySlotIdentifier,
   type ElementRenderer,
   elementRenderers,
+  getSurfaceBlock,
   type SurfaceBlockModel,
   type SurfaceContext,
 } from '@blocksuite/affine-block-surface';
@@ -30,7 +31,6 @@ import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import { Bound, getCommonBound } from '@blocksuite/global/utils';
 import { effect } from '@preact/signals-core';
 
-import { getSurfaceBlock } from '../../surface-ref-block/utils.js';
 import { RootService } from '../root-service.js';
 import { GfxBlockModel } from './block-model.js';
 import type { EdgelessFrameManager } from './frame-manager.js';
diff --git a/blocksuite/blocks/src/root-block/edgeless/utils/snap-manager.ts b/blocksuite/blocks/src/root-block/edgeless/utils/snap-manager.ts
index 86ccc214a1..8254243531 100644
--- a/blocksuite/blocks/src/root-block/edgeless/utils/snap-manager.ts
+++ b/blocksuite/blocks/src/root-block/edgeless/utils/snap-manager.ts
@@ -2,9 +2,10 @@ import type {
   SurfaceBlockComponent,
   SurfaceBlockModel,
 } from '@blocksuite/affine-block-surface';
-import { Overlay } from '@blocksuite/affine-block-surface';
+import { getSurfaceBlock, Overlay } from '@blocksuite/affine-block-surface';
 import type { ConnectorElementModel } from '@blocksuite/affine-model';
 import type { GfxController } from '@blocksuite/block-std/gfx';
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import { Bound, Point } from '@blocksuite/global/utils';
 
 import { isConnectable } from '../utils/query.js';
@@ -52,9 +53,13 @@ export class EdgelessSnapManager extends Overlay {
   };
 
   private get _surface() {
-    const surfaceModel = this.gfx.doc.getBlockByFlavour(
-      'affine:surface'
-    )[0] as SurfaceBlockModel;
+    const surfaceModel = getSurfaceBlock(this.gfx.doc);
+    if (!surfaceModel) {
+      throw new BlockSuiteError(
+        ErrorCode.ValueNotExists,
+        'Surface block not found in doc when creating snap manager'
+      );
+    }
 
     return this.gfx.std.view.getBlock(surfaceModel.id) as SurfaceBlockComponent;
   }
diff --git a/blocksuite/blocks/src/root-block/widgets/keyboard-toolbar/config.ts b/blocksuite/blocks/src/root-block/widgets/keyboard-toolbar/config.ts
index e79898795d..cb8ec0da55 100644
--- a/blocksuite/blocks/src/root-block/widgets/keyboard-toolbar/config.ts
+++ b/blocksuite/blocks/src/root-block/widgets/keyboard-toolbar/config.ts
@@ -1,4 +1,5 @@
 import { addSiblingAttachmentBlocks } from '@blocksuite/affine-block-attachment';
+import { getSurfaceBlock } from '@blocksuite/affine-block-surface';
 import {
   getInlineEditorByModel,
   insertContent,
@@ -61,7 +62,6 @@ import { cssVarV2 } from '@toeverything/theme/v2';
 import type { TemplateResult } from 'lit';
 
 import { toggleEmbedCardCreateModal } from '../../../_common/components/embed-card/modal/embed-card-create-modal.js';
-import { getSurfaceBlock } from '../../../surface-ref-block/utils.js';
 import type { PageRootBlockComponent } from '../../page/page-root-block.js';
 import { formatDate, formatTime } from '../../utils/misc.js';
 import type { AffineLinkedDocWidget } from '../linked-doc/index.js';
diff --git a/blocksuite/blocks/src/root-block/widgets/slash-menu/config.ts b/blocksuite/blocks/src/root-block/widgets/slash-menu/config.ts
index 2f695bbde4..b014263ba3 100644
--- a/blocksuite/blocks/src/root-block/widgets/slash-menu/config.ts
+++ b/blocksuite/blocks/src/root-block/widgets/slash-menu/config.ts
@@ -5,6 +5,7 @@ import {
   LoomIcon,
   YoutubeIcon,
 } from '@blocksuite/affine-block-embed';
+import { getSurfaceBlock } from '@blocksuite/affine-block-surface';
 import {
   ArrowDownBigIcon,
   ArrowUpBigIcon,
@@ -51,7 +52,6 @@ import type { TemplateResult } from 'lit';
 
 import { toggleEmbedCardCreateModal } from '../../../_common/components/embed-card/modal/embed-card-create-modal.js';
 import type { DataViewBlockComponent } from '../../../data-view-block/index.js';
-import { getSurfaceBlock } from '../../../surface-ref-block/utils.js';
 import type { RootBlockComponent } from '../../types.js';
 import { formatDate, formatTime } from '../../utils/misc.js';
 import type { AffineLinkedDocWidget } from '../linked-doc/index.js';
diff --git a/blocksuite/blocks/src/surface-ref-block/commands.ts b/blocksuite/blocks/src/surface-ref-block/commands.ts
index 202818e90a..203cbe74b7 100644
--- a/blocksuite/blocks/src/surface-ref-block/commands.ts
+++ b/blocksuite/blocks/src/surface-ref-block/commands.ts
@@ -1,9 +1,8 @@
+import { getSurfaceBlock } from '@blocksuite/affine-block-surface';
 import type { SurfaceRefProps } from '@blocksuite/affine-model';
 import { matchFlavours } from '@blocksuite/affine-shared/utils';
 import type { BlockCommands, Command } from '@blocksuite/block-std';
 
-import { getSurfaceBlock } from './utils.js';
-
 export const insertSurfaceRefBlockCommand: Command<
   'selectedModels',
   'insertedSurfaceRefBlockId',
diff --git a/blocksuite/blocks/src/surface-ref-block/surface-ref-block.ts b/blocksuite/blocks/src/surface-ref-block/surface-ref-block.ts
index 9da959dcca..813881ba2a 100644
--- a/blocksuite/blocks/src/surface-ref-block/surface-ref-block.ts
+++ b/blocksuite/blocks/src/surface-ref-block/surface-ref-block.ts
@@ -1,4 +1,5 @@
 import {
+  getSurfaceBlock,
   type SurfaceBlockModel,
   SurfaceElementModel,
 } from '@blocksuite/affine-block-surface';
@@ -306,10 +307,7 @@ export class SurfaceRefBlockComponent extends BlockComponent<SurfaceRefBlockMode
   }
 
   private _initReferencedModel() {
-    const surfaceModel: SurfaceBlockModel | null =
-      (this.doc.getBlocksByFlavour('affine:surface')[0]?.model as
-        | SurfaceBlockModel
-        | undefined) ?? null;
+    const surfaceModel = getSurfaceBlock(this.doc);
     this._surfaceModel = surfaceModel;
 
     const findReferencedModel = (): [
@@ -338,15 +336,11 @@ export class SurfaceRefBlockComponent extends BlockComponent<SurfaceRefBlockMode
         .find(
           doc =>
             doc.getBlock(this.model.reference) ||
-            (
-              doc.getBlocksByFlavour('affine:surface')[0]
-                .model as SurfaceBlockModel
-            ).getElementById(this.model.reference)
+            getSurfaceBlock(doc)!.getElementById(this.model.reference)
         );
 
       if (doc) {
-        this._surfaceModel = doc.getBlocksByFlavour('affine:surface')[0]
-          .model as SurfaceBlockModel;
+        this._surfaceModel = getSurfaceBlock(doc);
       }
 
       if (doc && doc.getBlock(this.model.reference)) {
@@ -356,12 +350,9 @@ export class SurfaceRefBlockComponent extends BlockComponent<SurfaceRefBlockMode
         ];
       }
 
-      if (doc && doc.getBlocksByFlavour('affine:surface')[0]) {
+      if (doc && getSurfaceBlock(doc)) {
         return [
-          (
-            doc.getBlocksByFlavour('affine:surface')[0]
-              .model as SurfaceBlockModel
-          ).getElementById(this.model.reference),
+          getSurfaceBlock(doc)!.getElementById(this.model.reference),
           doc.id,
         ];
       }
diff --git a/blocksuite/blocks/src/surface-ref-block/utils.ts b/blocksuite/blocks/src/surface-ref-block/utils.ts
index 01fa911a52..238c931e93 100644
--- a/blocksuite/blocks/src/surface-ref-block/utils.ts
+++ b/blocksuite/blocks/src/surface-ref-block/utils.ts
@@ -1,12 +1,5 @@
-import type { SurfaceBlockModel } from '@blocksuite/affine-block-surface';
-import type { Doc } from '@blocksuite/store';
 import { html } from 'lit';
 
-export function getSurfaceBlock(doc: Doc) {
-  const blocks = doc.getBlocksByFlavour('affine:surface');
-  return blocks.length !== 0 ? (blocks[0].model as SurfaceBlockModel) : null;
-}
-
 export const noContentPlaceholder = html`
   <svg
     width="182"
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts b/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts
index e50b20b8db..909d9f553e 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts
@@ -7,7 +7,6 @@ import type {
   EdgelessRootService,
   MindmapElementModel,
   ShapeElementModel,
-  SurfaceBlockModel,
 } from '@blocksuite/affine/blocks';
 import {
   addImages,
@@ -17,6 +16,7 @@ import {
   EDGELESS_TEXT_BLOCK_MIN_WIDTH,
   EdgelessTextBlockModel,
   fitContent,
+  getSurfaceBlock,
   ImageBlockModel,
   InsertBelowIcon,
   LightLoadingIcon,
@@ -363,9 +363,8 @@ function responseToCreateImage(host: EditorHost) {
 }
 
 export function responseToExpandMindmap(host: EditorHost, ctx: AIContext) {
-  const [surface] = host.doc.getBlockByFlavour(
-    'affine:surface'
-  ) as SurfaceBlockModel[];
+  const surface = getSurfaceBlock(host.doc);
+  if (!surface) return;
 
   const elements = ctx.get().selectedElements;
   const mindmapNode = ctx.get().node;
@@ -414,9 +413,8 @@ function responseToBrainstormMindmap(host: EditorHost, ctx: AIContext) {
   const edgelessService = getEdgelessService(host);
   const edgelessCopilot = getEdgelessCopilotWidget(host);
   const selectionRect = edgelessCopilot.selectionModelRect;
-  const [surface] = host.doc.getBlockByFlavour(
-    'affine:surface'
-  ) as SurfaceBlockModel[];
+  const surface = getSurfaceBlock(host.doc);
+  if (!surface) return;
 
   const { node, style, selectedElements } = ctx.get();
   if (!node) return;
@@ -478,9 +476,8 @@ function responseToMakeItReal(host: EditorHost, ctx: AIContext) {
   html = preprocessHtml(html);
 
   const edgelessCopilot = getEdgelessCopilotWidget(host);
-  const [surface] = host.doc.getBlockByFlavour(
-    'affine:surface'
-  ) as SurfaceBlockModel[];
+  const surface = getSurfaceBlock(host.doc);
+  if (!surface) return;
 
   const data = ctx.get();
   const bounds = edgelessCopilot.determineInsertionBounds(