diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
index b852d4b541..a2891cdf75 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
@@ -24,7 +24,6 @@ import {
   type DocCustomPropertyInfo,
   DocService,
   DocsService,
-  FeatureFlagService,
   useFramework,
   useLiveData,
   useService,
@@ -50,6 +49,7 @@ import {
 } from '../../doc-properties';
 import { BiDirectionalLinkPanel } from './bi-directional-link-panel';
 import { BlocksuiteEditorJournalDocTitle } from './journal-doc-title';
+import { extendEdgelessPreviewSpec } from './specs/custom/root-block';
 import {
   patchDocModeService,
   patchEdgelessClipboard,
@@ -96,14 +96,12 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
     docsService,
     editorService,
     workspaceService,
-    featureFlagService,
   } = useServices({
     PeekViewService,
     DocService,
     DocsService,
     WorkspaceService,
     EditorService,
-    FeatureFlagService,
   });
   const framework = useFramework();
   const referenceRenderer: ReferenceReactRenderer = useMemo(() => {
@@ -130,12 +128,15 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
     };
   }, [workspaceService]);
 
+  useMemo(() => {
+    extendEdgelessPreviewSpec(framework);
+  }, [framework]);
+
   const specs = useMemo(() => {
-    const enableAI = featureFlagService.flags.enable_ai.value;
     return mode === 'edgeless'
-      ? createEdgelessModeSpecs(framework, !!enableAI)
-      : createPageModeSpecs(framework, !!enableAI);
-  }, [featureFlagService.flags.enable_ai.value, mode, framework]);
+      ? createEdgelessModeSpecs(framework)
+      : createPageModeSpecs(framework);
+  }, [mode, framework]);
 
   const confirmModal = useConfirmModal();
   const patchedSpecs = useMemo(() => {
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
index 5bc15753c1..d89f519809 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
@@ -30,11 +30,13 @@ import {
 import {
   createSignalFromObservable,
   type Signal,
+  SpecProvider,
 } from '@blocksuite/affine-shared/utils';
 import type { Container } from '@blocksuite/global/di';
 import {
   DocService,
   DocsService,
+  FeatureFlagService,
   type FrameworkProvider,
 } from '@toeverything/infra';
 import type { Observable } from 'rxjs';
@@ -60,7 +62,7 @@ function getTelemetryExtension(): ExtensionType {
   };
 }
 
-function createThemeExtension(framework: FrameworkProvider) {
+function getThemeExtension(framework: FrameworkProvider) {
   class AffineThemeExtension
     extends LifeCycleWatcher
     implements ThemeExtension
@@ -153,14 +155,32 @@ function getEditorConfigExtension(
   ];
 }
 
+export const extendEdgelessPreviewSpec = (function () {
+  let _extension: ExtensionType;
+  let _framework: FrameworkProvider;
+  return function (framework: FrameworkProvider) {
+    if (framework === _framework && _extension) {
+      return _extension;
+    } else {
+      _extension &&
+        SpecProvider.getInstance().omitSpec('edgeless:preview', _extension);
+      _extension = getThemeExtension(framework);
+      _framework = framework;
+      SpecProvider.getInstance().extendSpec('edgeless:preview', [_extension]);
+      return _extension;
+    }
+  };
+})();
+
 export function createPageRootBlockSpec(
-  framework: FrameworkProvider,
-  enableAI: boolean
+  framework: FrameworkProvider
 ): ExtensionType[] {
+  const featureFlagService = framework.get(FeatureFlagService);
+  const enableAI = featureFlagService.flags.enable_ai.value;
   return [
     enableAI ? AIPageRootBlockSpec : PageRootBlockSpec,
     FontLoaderService,
-    createThemeExtension(framework),
+    getThemeExtension(framework),
     getFontConfigExtension(),
     getTelemetryExtension(),
     getEditorConfigExtension(framework),
@@ -168,13 +188,14 @@ export function createPageRootBlockSpec(
 }
 
 export function createEdgelessRootBlockSpec(
-  framework: FrameworkProvider,
-  enableAI: boolean
+  framework: FrameworkProvider
 ): ExtensionType[] {
+  const featureFlagService = framework.get(FeatureFlagService);
+  const enableAI = featureFlagService.flags.enable_ai.value;
   return [
     enableAI ? AIEdgelessRootBlockSpec : EdgelessRootBlockSpec,
     FontLoaderService,
-    createThemeExtension(framework),
+    getThemeExtension(framework),
     EdgelessToolExtension,
     EdgelessBuiltInManager,
     getFontConfigExtension(),
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
index 50a6ffc87a..a1c13dda0a 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
@@ -10,15 +10,19 @@ import {
   EdgelessTextBlockSpec,
   FrameBlockSpec,
 } from '@blocksuite/affine/blocks';
-import type { FrameworkProvider } from '@toeverything/infra';
+import {
+  FeatureFlagService,
+  type FrameworkProvider,
+} from '@toeverything/infra';
 
 import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createEdgelessRootBlockSpec } from './custom/root-block';
 
 export function createEdgelessModeSpecs(
-  framework: FrameworkProvider,
-  enableAI: boolean
+  framework: FrameworkProvider
 ): ExtensionType[] {
+  const featureFlagService = framework.get(FeatureFlagService);
+  const enableAI = featureFlagService.flags.enable_ai.value;
   return [
     ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     EdgelessSurfaceBlockSpec,
@@ -27,7 +31,7 @@ export function createEdgelessModeSpecs(
     EdgelessTextBlockSpec,
     EdgelessNoteBlockSpec,
     // special
-    createEdgelessRootBlockSpec(framework, enableAI),
+    createEdgelessRootBlockSpec(framework),
   ].flat();
 }
 
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
index b7ea0e5f4f..102daeb76b 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
@@ -4,21 +4,25 @@ import {
   PageSurfaceBlockSpec,
   PageSurfaceRefBlockSpec,
 } from '@blocksuite/affine/blocks';
-import { type FrameworkProvider } from '@toeverything/infra';
+import {
+  FeatureFlagService,
+  type FrameworkProvider,
+} from '@toeverything/infra';
 
 import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createPageRootBlockSpec } from './custom/root-block';
 
 export function createPageModeSpecs(
-  framework: FrameworkProvider,
-  enableAI: boolean
+  framework: FrameworkProvider
 ): ExtensionType[] {
+  const featureFlagService = framework.get(FeatureFlagService);
+  const enableAI = featureFlagService.flags.enable_ai.value;
   return [
     ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     PageSurfaceBlockSpec,
     PageSurfaceRefBlockSpec,
     NoteBlockSpec,
     // special
-    createPageRootBlockSpec(framework, enableAI),
+    createPageRootBlockSpec(framework),
   ].flat();
 }
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 aa27d0f468..81e0d94dac 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
@@ -1,16 +1,33 @@
 import { Skeleton } from '@affine/component';
 import type { EditorSettingSchema } from '@affine/core/modules/editor-setting';
 import { EditorSettingService } from '@affine/core/modules/editor-setting';
+import { AppThemeService } from '@affine/core/modules/theme';
 import type { EditorHost } from '@blocksuite/affine/block-std';
-import { BlockStdScope } from '@blocksuite/affine/block-std';
+import {
+  BlockStdScope,
+  LifeCycleWatcher,
+  StdIdentifier,
+} from '@blocksuite/affine/block-std';
 import type { GfxPrimitiveElementModel } from '@blocksuite/affine/block-std/gfx';
-import type { EdgelessRootService } from '@blocksuite/affine/blocks';
-import { SpecProvider } from '@blocksuite/affine/blocks';
+import type {
+  EdgelessRootService,
+  ThemeExtension,
+} from '@blocksuite/affine/blocks';
+import {
+  ColorScheme,
+  SpecProvider,
+  ThemeExtensionIdentifier,
+} from '@blocksuite/affine/blocks';
 import { Bound } from '@blocksuite/affine/global/utils';
 import type { Block, Doc } from '@blocksuite/affine/store';
+import { createSignalFromObservable } from '@blocksuite/affine-shared/utils';
+import type { Container } from '@blocksuite/global/di';
+import type { Signal } from '@preact/signals-core';
+import type { FrameworkProvider } from '@toeverything/infra';
 import { useFramework } from '@toeverything/infra';
 import { isEqual } from 'lodash-es';
 import { useCallback, useEffect, useRef } from 'react';
+import type { Observable } from 'rxjs';
 import { map, pairwise } from 'rxjs';
 
 import {
@@ -74,7 +91,10 @@ export const EdgelessSnapshot = (props: Props) => {
 
     const editorHost = new BlockStdScope({
       doc,
-      extensions: SpecProvider.getInstance().getSpec('edgeless:preview').value,
+      extensions: [
+        ...SpecProvider.getInstance().getSpec('edgeless:preview').value,
+        getThemeExtension(framework),
+      ],
     }).render();
     docRef.current = doc;
     editorHostRef.current?.remove();
@@ -106,7 +126,7 @@ export const EdgelessSnapshot = (props: Props) => {
 
     // append to dom node
     wrapperRef.current.append(editorHost);
-  }, [docName, firstUpdate, updateElements]);
+  }, [docName, firstUpdate, framework, updateElements]);
 
   useEffect(() => {
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
@@ -150,3 +170,58 @@ export const EdgelessSnapshot = (props: Props) => {
     </div>
   );
 };
+
+function getThemeExtension(framework: FrameworkProvider) {
+  class AffineThemeExtension
+    extends LifeCycleWatcher
+    implements ThemeExtension
+  {
+    static override readonly key = 'affine-settings-theme';
+
+    private readonly theme: Signal<ColorScheme>;
+
+    protected readonly disposables: (() => void)[] = [];
+
+    static override setup(di: Container) {
+      super.setup(di);
+      di.override(ThemeExtensionIdentifier, AffineThemeExtension, [
+        StdIdentifier,
+      ]);
+    }
+
+    constructor(std: BlockStdScope) {
+      super(std);
+      const theme$: Observable<ColorScheme> = framework
+        .get(AppThemeService)
+        .appTheme.theme$.map(theme => {
+          return theme === ColorScheme.Dark
+            ? ColorScheme.Dark
+            : ColorScheme.Light;
+        });
+      const { signal, cleanup } = createSignalFromObservable<ColorScheme>(
+        theme$,
+        ColorScheme.Light
+      );
+      this.theme = signal;
+      this.disposables.push(cleanup);
+    }
+
+    getAppTheme() {
+      return this.theme;
+    }
+
+    getEdgelessTheme() {
+      return this.theme;
+    }
+
+    override unmounted() {
+      this.dispose();
+    }
+
+    dispose() {
+      this.disposables.forEach(dispose => dispose());
+    }
+  }
+
+  return AffineThemeExtension;
+}