diff --git a/blocksuite/affine/blocks/block-root/src/edgeless/components/frame/frame-preview.ts b/blocksuite/affine/blocks/block-root/src/edgeless/components/frame/frame-preview.ts
index 93360b2147..2c4617d6d0 100644
--- a/blocksuite/affine/blocks/block-root/src/edgeless/components/frame/frame-preview.ts
+++ b/blocksuite/affine/blocks/block-root/src/edgeless/components/frame/frame-preview.ts
@@ -1,4 +1,5 @@
 import type { FrameBlockModel } from '@blocksuite/affine-model';
+import { ViewportElementExtension } from '@blocksuite/affine-shared/services';
 import { SpecProvider } from '@blocksuite/affine-shared/utils';
 import {
   BlockStdScope,
@@ -16,8 +17,6 @@ import { property, query, state } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';
 import debounce from 'lodash-es/debounce';
 
-import type { EdgelessRootPreviewBlockComponent } from '../../edgeless-root-preview-block.js';
-
 const DEFAULT_PREVIEW_CONTAINER_WIDTH = 280;
 const DEFAULT_PREVIEW_CONTAINER_HEIGHT = 166;
 
@@ -118,7 +117,8 @@ export class FramePreview extends WithDisposable(ShadowlessElement) {
       static override key = 'frame-preview-watcher';
 
       override mounted() {
-        const { view } = this.std;
+        const std = this.std;
+        const { view } = std;
         view.viewUpdated.on(payload => {
           if (
             payload.type !== 'block' ||
@@ -127,16 +127,17 @@ export class FramePreview extends WithDisposable(ShadowlessElement) {
           ) {
             return;
           }
-          const edgelessBlock =
-            payload.view as EdgelessRootPreviewBlockComponent;
-          edgelessBlock.editorViewportSelector = 'frame-preview-viewport';
-          edgelessBlock.service.viewport.sizeUpdated.once(() => {
+          const viewport = std.get(GfxControllerIdentifier).viewport;
+          viewport.sizeUpdated.once(() => {
             refreshViewport();
           });
         });
       }
     }
-    this._previewSpec.extend([FramePreviewWatcher]);
+    this._previewSpec.extend([
+      ViewportElementExtension('.frame-preview-viewport'),
+      FramePreviewWatcher,
+    ]);
   }
 
   private _refreshViewport() {
diff --git a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-block.ts b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-block.ts
index ab19f5bdc9..e379019cc3 100644
--- a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-block.ts
+++ b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-block.ts
@@ -20,6 +20,7 @@ import {
   EditPropsStore,
   FontLoaderService,
   ThemeProvider,
+  ViewportElementProvider,
 } from '@blocksuite/affine-shared/services';
 import type { Viewport } from '@blocksuite/affine-shared/types';
 import {
@@ -39,7 +40,6 @@ import {
   type GfxViewportElement,
 } from '@blocksuite/block-std/gfx';
 import { IS_WINDOWS } from '@blocksuite/global/env';
-import { BlockSuiteError } from '@blocksuite/global/exceptions';
 import { Bound, Point, Vec } from '@blocksuite/global/gfx';
 import { effect } from '@preact/signals-core';
 import { css, html } from 'lit';
@@ -125,8 +125,6 @@ export class EdgelessRootBlockComponent extends BlockComponent<
 
   private _resizeObserver: ResizeObserver | null = null;
 
-  private _viewportElement: HTMLElement | null = null;
-
   clipboardController = new EdgelessClipboardController(this);
 
   keyboardManager: EdgelessPageKeyboardManager | null = null;
@@ -165,39 +163,11 @@ export class EdgelessRootBlockComponent extends BlockComponent<
    * This refers to the wrapper element of the EditorHost.
    */
   get viewport(): Viewport {
-    const {
-      scrollLeft,
-      scrollTop,
-      scrollWidth,
-      scrollHeight,
-      clientWidth,
-      clientHeight,
-    } = this.viewportElement;
-    const { top, left } = this.viewportElement.getBoundingClientRect();
-    return {
-      top,
-      left,
-      scrollLeft,
-      scrollTop,
-      scrollWidth,
-      scrollHeight,
-      clientWidth,
-      clientHeight,
-    };
+    return this.std.get(ViewportElementProvider).viewport;
   }
 
   get viewportElement(): HTMLElement {
-    if (this._viewportElement) return this._viewportElement;
-    this._viewportElement = this.host.closest(
-      '.affine-edgeless-viewport'
-    ) as HTMLElement | null;
-    if (!this._viewportElement) {
-      throw new BlockSuiteError(
-        BlockSuiteError.ErrorCode.ValueNotExists,
-        'EdgelessRootBlockComponent.viewportElement: viewport element is not found'
-      );
-    }
-    return this._viewportElement;
+    return this.std.get(ViewportElementProvider).viewportElement;
   }
 
   private _initFontLoader() {
diff --git a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-preview-block.ts b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-preview-block.ts
index 683319f925..38c3a1e3ae 100644
--- a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-preview-block.ts
+++ b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-preview-block.ts
@@ -3,12 +3,12 @@ import {
   type SurfaceBlockComponent,
   type SurfaceBlockModel,
 } from '@blocksuite/affine-block-surface';
-import type { EdgelessPreviewer } from '@blocksuite/affine-block-surface-ref';
 import type { RootBlockModel } from '@blocksuite/affine-model';
 import {
   EditorSettingProvider,
   FontLoaderService,
   ThemeProvider,
+  ViewportElementProvider,
 } from '@blocksuite/affine-shared/services';
 import { requestThrottledConnectedFrame } from '@blocksuite/affine-shared/utils';
 import {
@@ -17,7 +17,6 @@ import {
   SurfaceSelection,
 } from '@blocksuite/block-std';
 import type { GfxViewportElement } from '@blocksuite/block-std/gfx';
-import { BlockSuiteError } from '@blocksuite/global/exceptions';
 import { css, html } from 'lit';
 import { query, state } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';
@@ -26,14 +25,11 @@ import type { EdgelessRootBlockWidgetName } from '../types.js';
 import type { EdgelessRootService } from './edgeless-root-service.js';
 import { isCanvasElement } from './utils/query.js';
 
-export class EdgelessRootPreviewBlockComponent
-  extends BlockComponent<
-    RootBlockModel,
-    EdgelessRootService,
-    EdgelessRootBlockWidgetName
-  >
-  implements EdgelessPreviewer
-{
+export class EdgelessRootPreviewBlockComponent extends BlockComponent<
+  RootBlockModel,
+  EdgelessRootService,
+  EdgelessRootBlockWidgetName
+> {
   static override styles = css`
     affine-edgeless-root-preview {
       pointer-events: none;
@@ -87,8 +83,6 @@ export class EdgelessRootPreviewBlockComponent
 
   private _resizeObserver: ResizeObserver | null = null;
 
-  private _viewportElement: HTMLElement | null = null;
-
   get dispatcher() {
     return this.service?.uiEventDispatcher;
   }
@@ -100,17 +94,7 @@ export class EdgelessRootPreviewBlockComponent
   }
 
   get viewportElement(): HTMLElement {
-    if (this._viewportElement) return this._viewportElement;
-    this._viewportElement = this.host.closest(
-      this.editorViewportSelector
-    ) as HTMLElement | null;
-    if (!this._viewportElement) {
-      throw new BlockSuiteError(
-        BlockSuiteError.ErrorCode.ValueNotExists,
-        'EdgelessRootPreviewBlockComponent.viewportElement: viewport element is not found'
-      );
-    }
-    return this._viewportElement;
+    return this.std.get(ViewportElementProvider).viewportElement;
   }
 
   private _initFontLoader() {
@@ -159,23 +143,24 @@ export class EdgelessRootPreviewBlockComponent
   }
 
   private _initResizeEffect() {
-    if (!this._viewportElement) {
-      return;
-    }
-
     const resizeObserver = new ResizeObserver((_: ResizeObserverEntry[]) => {
       // FIXME: find a better way to get rid of empty check
       if (!this.service || !this.service.selection || !this.service.viewport) {
-        console.error('Service not ready');
+        console.error('Service is not ready');
         return;
       }
       this.service.selection.set(this.service.selection.surfaceSelections);
       this.service.viewport.onResize();
     });
 
-    resizeObserver.observe(this.viewportElement);
-    this._resizeObserver?.disconnect();
-    this._resizeObserver = resizeObserver;
+    try {
+      resizeObserver.observe(this.viewportElement);
+      this._resizeObserver?.disconnect();
+      this._resizeObserver = resizeObserver;
+    } catch {
+      // viewport is not ready
+      console.error('Viewport is not ready');
+    }
   }
 
   private _initSlotEffects() {
@@ -265,18 +250,9 @@ export class EdgelessRootPreviewBlockComponent
     `;
   }
 
-  override willUpdate(_changedProperties: Map<PropertyKey, unknown>): void {
-    if (_changedProperties.has('editorViewportSelector')) {
-      this._initResizeEffect();
-    }
-  }
-
   @state()
   accessor overrideBackground: string | undefined = undefined;
 
-  @state()
-  accessor editorViewportSelector = '.affine-edgeless-viewport';
-
   @query('gfx-viewport')
   accessor gfxViewportElm!: GfxViewportElement;
 
diff --git a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-spec.ts b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-spec.ts
index 47adb4b19a..1ca4dc44e3 100644
--- a/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-spec.ts
+++ b/blocksuite/affine/blocks/block-root/src/edgeless/edgeless-root-spec.ts
@@ -1,3 +1,4 @@
+import { ViewportElementExtension } from '@blocksuite/affine-shared/services';
 import { autoConnectWidget } from '@blocksuite/affine-widget-edgeless-auto-connect';
 import { frameTitleWidget } from '@blocksuite/affine-widget-frame-title';
 import { edgelessRemoteSelectionWidget } from '@blocksuite/affine-widget-remote-selection';
@@ -72,6 +73,7 @@ const EdgelessCommonExtension: ExtensionType[] = [
   CommonSpecs,
   ToolController,
   EdgelessRootService,
+  ViewportElementExtension('.affine-edgeless-viewport'),
 ].flat();
 
 export const EdgelessRootBlockSpec: ExtensionType[] = [
diff --git a/blocksuite/affine/blocks/block-root/src/page/page-root-block.ts b/blocksuite/affine/blocks/block-root/src/page/page-root-block.ts
index 3f9e26682b..7456f378d9 100644
--- a/blocksuite/affine/blocks/block-root/src/page/page-root-block.ts
+++ b/blocksuite/affine/blocks/block-root/src/page/page-root-block.ts
@@ -8,7 +8,10 @@ import {
   type RootBlockModel,
 } from '@blocksuite/affine-model';
 import { focusTextModel } from '@blocksuite/affine-rich-text';
-import { PageViewportService } from '@blocksuite/affine-shared/services';
+import {
+  PageViewportService,
+  ViewportElementProvider,
+} from '@blocksuite/affine-shared/services';
 import type { Viewport } from '@blocksuite/affine-shared/types';
 import {
   focusTitle,
@@ -112,8 +115,6 @@ export class PageRootBlockComponent extends BlockComponent<
     }
   `;
 
-  private _viewportElement: HTMLDivElement | null = null;
-
   clipboardController = new PageClipboard(this);
 
   /**
@@ -182,12 +183,8 @@ export class PageRootBlockComponent extends BlockComponent<
     };
   }
 
-  get viewportElement(): HTMLDivElement | null {
-    if (this._viewportElement) return this._viewportElement;
-    this._viewportElement = this.host.closest<HTMLDivElement>(
-      '.affine-page-viewport'
-    );
-    return this._viewportElement;
+  get viewportElement(): HTMLElement {
+    return this.std.get(ViewportElementProvider).viewportElement;
   }
 
   private _createDefaultNoteBlock() {
diff --git a/blocksuite/affine/blocks/block-root/src/page/page-root-spec.ts b/blocksuite/affine/blocks/block-root/src/page/page-root-spec.ts
index 263d1b6ef7..18076587b6 100644
--- a/blocksuite/affine/blocks/block-root/src/page/page-root-spec.ts
+++ b/blocksuite/affine/blocks/block-root/src/page/page-root-spec.ts
@@ -1,3 +1,4 @@
+import { ViewportElementExtension } from '@blocksuite/affine-shared/services';
 import { BlockViewExtension, WidgetViewExtension } from '@blocksuite/block-std';
 import type { ExtensionType } from '@blocksuite/store';
 import { literal, unsafeStatic } from 'lit/static-html.js';
@@ -23,6 +24,7 @@ const PageCommonExtension: ExtensionType[] = [
   CommonSpecs,
   PageRootService,
   pageDraggingAreaWidget,
+  ViewportElementExtension('.affine-page-viewport'),
 ].flat();
 
 export const PageRootBlockSpec: ExtensionType[] = [
diff --git a/blocksuite/affine/blocks/block-surface-ref/src/index.ts b/blocksuite/affine/blocks/block-surface-ref/src/index.ts
index 6903b745d8..eb4c6cf7b3 100644
--- a/blocksuite/affine/blocks/block-surface-ref/src/index.ts
+++ b/blocksuite/affine/blocks/block-surface-ref/src/index.ts
@@ -2,5 +2,4 @@ export * from './commands.js';
 export * from './surface-ref-block.js';
 export * from './surface-ref-block-edgeless.js';
 export * from './surface-ref-spec.js';
-export * from './types.js';
 export * from './utils.js';
diff --git a/blocksuite/affine/blocks/block-surface-ref/src/surface-ref-block.ts b/blocksuite/affine/blocks/block-surface-ref/src/surface-ref-block.ts
index 91e5922daf..ee43c08e6f 100644
--- a/blocksuite/affine/blocks/block-surface-ref/src/surface-ref-block.ts
+++ b/blocksuite/affine/blocks/block-surface-ref/src/surface-ref-block.ts
@@ -15,6 +15,7 @@ import {
   DocModeProvider,
   EditPropsStore,
   ThemeProvider,
+  ViewportElementExtension,
 } from '@blocksuite/affine-shared/services';
 import {
   matchModels,
@@ -238,7 +239,9 @@ export class SurfaceRefBlockComponent extends BlockComponent<SurfaceRefBlockMode
 
   private _previewDoc: Store | null = null;
 
-  private readonly _previewSpec = SpecProvider._.getSpec('preview:edgeless');
+  private readonly _previewSpec = SpecProvider._.getSpec(
+    'preview:edgeless'
+  ).extend([ViewportElementExtension('.ref-viewport')]);
 
   private _referencedModel: GfxModel | null = null;
 
diff --git a/blocksuite/affine/blocks/block-surface-ref/src/types.ts b/blocksuite/affine/blocks/block-surface-ref/src/types.ts
deleted file mode 100644
index 019076467c..0000000000
--- a/blocksuite/affine/blocks/block-surface-ref/src/types.ts
+++ /dev/null
@@ -1,3 +0,0 @@
-export interface EdgelessPreviewer {
-  editorViewportSelector: string;
-}
diff --git a/blocksuite/affine/shared/src/services/index.ts b/blocksuite/affine/shared/src/services/index.ts
index a26d4cd9ac..c89e3ab47f 100644
--- a/blocksuite/affine/shared/src/services/index.ts
+++ b/blocksuite/affine/shared/src/services/index.ts
@@ -21,4 +21,5 @@ export * from './telemetry-service';
 export * from './theme-service';
 export * from './toolbar-service';
 export * from './user-service';
+export * from './viewport-element-service';
 export * from './virtual-keyboard-service';
diff --git a/blocksuite/affine/shared/src/services/viewport-element-service.ts b/blocksuite/affine/shared/src/services/viewport-element-service.ts
new file mode 100644
index 0000000000..c14c3e318f
--- /dev/null
+++ b/blocksuite/affine/shared/src/services/viewport-element-service.ts
@@ -0,0 +1,60 @@
+import { createIdentifier } from '@blocksuite/global/di';
+import { BlockSuiteError } from '@blocksuite/global/exceptions';
+import type { ExtensionType } from '@blocksuite/store';
+
+import type { Viewport } from '../types';
+
+export interface ViewportElementService {
+  get viewportElement(): HTMLElement;
+  get viewport(): Viewport;
+}
+
+export const ViewportElementProvider = createIdentifier<ViewportElementService>(
+  'ViewportElementProvider'
+);
+
+export const ViewportElementExtension = (selector: string): ExtensionType => {
+  return {
+    setup: di => {
+      di.override(ViewportElementProvider, () => {
+        const getViewportElement = (): HTMLElement => {
+          const viewportElement = document.querySelector<HTMLElement>(selector);
+          if (!viewportElement) {
+            throw new BlockSuiteError(
+              BlockSuiteError.ErrorCode.ValueNotExists,
+              `ViewportElementProvider: viewport element is not found`
+            );
+          }
+          return viewportElement;
+        };
+        return {
+          get viewportElement() {
+            return getViewportElement();
+          },
+          get viewport() {
+            const viewportElement = getViewportElement();
+            const {
+              scrollLeft,
+              scrollTop,
+              scrollWidth,
+              scrollHeight,
+              clientWidth,
+              clientHeight,
+            } = viewportElement;
+            const { top, left } = viewportElement.getBoundingClientRect();
+            return {
+              top,
+              left,
+              scrollLeft,
+              scrollTop,
+              scrollWidth,
+              scrollHeight,
+              clientWidth,
+              clientHeight,
+            };
+          },
+        };
+      });
+    },
+  };
+};
diff --git a/packages/frontend/core/src/desktop/dialogs/setting/general-setting/editor/edgeless/snapshot.tsx b/packages/frontend/core/src/desktop/dialogs/setting/general-setting/editor/edgeless/snapshot.tsx
index da2dcc84ad..b5c7420719 100644
--- a/packages/frontend/core/src/desktop/dialogs/setting/general-setting/editor/edgeless/snapshot.tsx
+++ b/packages/frontend/core/src/desktop/dialogs/setting/general-setting/editor/edgeless/snapshot.tsx
@@ -7,9 +7,9 @@ import {
   GfxControllerIdentifier,
   type GfxPrimitiveElementModel,
 } from '@blocksuite/affine/block-std/gfx';
-import type { EdgelessRootPreviewBlockComponent } from '@blocksuite/affine/blocks/root';
 import { EdgelessCRUDIdentifier } from '@blocksuite/affine/blocks/surface';
 import { Bound } from '@blocksuite/affine/global/gfx';
+import { ViewportElementExtension } from '@blocksuite/affine/shared/services';
 import { SpecProvider } from '@blocksuite/affine/shared/utils';
 import type { Block, Store } from '@blocksuite/affine/store';
 import { useFramework } from '@toeverything/infra';
@@ -76,7 +76,9 @@ export const EdgelessSnapshot = (props: Props) => {
 
     const editorHost = new BlockStdScope({
       store: doc,
-      extensions: SpecProvider._.getSpec('preview:edgeless').value,
+      extensions: SpecProvider._.getSpec('preview:edgeless').extend([
+        ViewportElementExtension('.ref-viewport'),
+      ]).value,
     }).render();
     docRef.current = doc;
     editorHostRef.current?.remove();
@@ -98,9 +100,7 @@ export const EdgelessSnapshot = (props: Props) => {
       ) {
         return;
       }
-      const component = payload.view as EdgelessRootPreviewBlockComponent;
       doc.readonly = false;
-      component.editorViewportSelector = 'ref-viewport';
       const frame = getFrameBlock(doc);
       if (frame && docName !== 'frame') {
         // docName with value 'frame' shouldn't be deleted, it is a part of frame settings