feat(editor): type safe draft model and transformer (#10486)

blocksuite/affine/block-embed/src/common/render-linked-doc.ts

@@ -304,11 +304,12 @@ export function getDocContentWithMaxLength(doc: Store, maxlength = 500) {
 
 export function getTitleFromSelectedModels(selectedModels: DraftModel[]) {
   const firstBlock = selectedModels[0];
-  if (
+  const isParagraph = (
-    matchModels(firstBlock, [ParagraphBlockModel]) &&
+    model: DraftModel
-    firstBlock.type.startsWith('h')
+  ): model is DraftModel<ParagraphBlockModel> =>
-  ) {
+    model.flavour === 'affine:paragraph';
-    return firstBlock.text.toString();
+  if (isParagraph(firstBlock) && firstBlock.type.startsWith('h')) {
+    return firstBlock.text?.toString();
   }
   return undefined;
 }
@@ -394,7 +395,7 @@ export async function convertSelectedBlocksToLinkedDoc(
     'before'
   );
   // delete selected elements
-  models.forEach(model => doc.deleteBlock(model));
+  models.forEach(model => doc.deleteBlock(model.id));
   return linkedDoc;
 }
 

blocksuite/affine/block-note/src/quick-action.ts

@@ -9,6 +9,7 @@ import {
   getSelectedModelsCommand,
 } from '@blocksuite/affine-shared/commands';
 import type { BlockStdScope } from '@blocksuite/block-std';
+import { toDraftModel } from '@blocksuite/store';
 
 export interface QuickActionConfig {
   id: string;
@@ -45,7 +46,9 @@ export const quickActionConfig: QuickActionConfig[] = [
       std.selection.clear();
 
       const doc = std.store;
-      const autofill = getTitleFromSelectedModels(selectedModels);
+      const autofill = getTitleFromSelectedModels(
+        selectedModels.map(toDraftModel)
+      );
       promptDocTitle(std, autofill)
         .then(title => {
           if (title === null) return;

blocksuite/affine/block-root/src/keyboard/keyboard-manager.ts

@@ -17,6 +17,7 @@ import {
   type UIEventHandler,
 } from '@blocksuite/block-std';
 import { IS_MAC, IS_WINDOWS } from '@blocksuite/global/env';
+import { toDraftModel } from '@blocksuite/store';
 
 export class PageKeyboardManager {
   private readonly _handleDelete: UIEventHandler = ctx => {
@@ -143,7 +144,9 @@ export class PageKeyboardManager {
     }
 
     const doc = rootComponent.host.doc;
-    const autofill = getTitleFromSelectedModels(selectedModels);
+    const autofill = getTitleFromSelectedModels(
+      selectedModels.map(toDraftModel)
+    );
     promptDocTitle(rootComponent.std, autofill)
       .then(title => {
         if (title === null) return;

blocksuite/affine/block-root/src/widgets/format-bar/config.ts

@@ -64,7 +64,7 @@ import type {
 import { tableViewMeta } from '@blocksuite/data-view/view-presets';
 import { assertExists } from '@blocksuite/global/utils';
 import { MoreVerticalIcon } from '@blocksuite/icons/lit';
-import { Slice } from '@blocksuite/store';
+import { Slice, toDraftModel } from '@blocksuite/store';
 import { html, type TemplateResult } from 'lit';
 
 import { FormatBarContext } from './context.js';
@@ -230,7 +230,9 @@ export function toolbarDefaultConfig(toolbar: AffineFormatBarWidget) {
         host.selection.clear();
 
         const doc = host.doc;
-        const autofill = getTitleFromSelectedModels(selectedModels);
+        const autofill = getTitleFromSelectedModels(
+          selectedModels.map(toDraftModel)
+        );
         promptDocTitle(std, autofill)
           .then(async title => {
             if (title === null) return;

blocksuite/affine/shared/src/adapters/middlewares/copy.ts

@@ -1,4 +1,4 @@
-import { RootBlockModel } from '@blocksuite/affine-model';
+import type { RootBlockModel } from '@blocksuite/affine-model';
 import {
   type BlockStdScope,
   type EditorHost,
@@ -12,7 +12,9 @@ import type {
   TransformerSlots,
 } from '@blocksuite/store';
 
-import { matchModels } from '../../utils';
+const isRootDraftModel = (
+  model: DraftModel
+): model is DraftModel<RootBlockModel> => model.flavour === 'affine:root';
 
 const handlePoint = (
   point: TextRangePoint,
@@ -20,7 +22,7 @@ const handlePoint = (
   model: DraftModel
 ) => {
   const { index, length } = point;
-  if (matchModels(model, [RootBlockModel])) {
+  if (isRootDraftModel(model)) {
     if (length === 0) return;
     (snapshot.props.title as Record<string, unknown>).delta =
       model.title.sliceToDelta(index, length + index);

blocksuite/affine/shared/src/utils/model/checker.ts

@@ -11,7 +11,7 @@ type ModelList<T> =
 export function matchModels<
   const Model extends ConstructorType<BlockModel>[],
   U extends ModelList<Model>[number] = ModelList<Model>[number],
->(model: unknown, expected: Model): model is U {
+>(model: BlockModel | null, expected: Model): model is U {
   return (
     !!model && expected.some(expectedModel => model instanceof expectedModel)
   );

blocksuite/affine/widget-drag-handle/src/watchers/drag-event-watcher.ts

@@ -1139,7 +1139,7 @@ export class DragEventWatcher {
                   block.flavour === 'affine:bookmark' ||
                   block.flavour.startsWith('affine:embed-'))
               ) {
-                store.updateBlock(block as BlockModel, {
+                store.updateBlock(block.id, {
                   xywh: content[idx].props.xywh,
                   style: content[idx].props.style,
                 });
@@ -1164,7 +1164,7 @@ export class DragEventWatcher {
                 block.flavour === 'affine:attachment' ||
                 block.flavour.startsWith('affine:embed-')
               ) {
-                store.updateBlock(block as BlockModel, {
+                store.updateBlock(block.id, {
                   xywh: content[idx].props.xywh,
                   style: content[idx].props.style,
                 });

blocksuite/framework/store/src/adapter/base.ts

@@ -1,6 +1,11 @@
 import { BlockSuiteError } from '@blocksuite/global/exceptions';
 
-import type { DraftModel, Store } from '../model/index.js';
+import {
+  BlockModel,
+  type DraftModel,
+  type Store,
+  toDraftModel,
+} from '../model/index.js';
 import type { AssetsManager } from '../transformer/assets.js';
 import type { Slice, Transformer } from '../transformer/index.js';
 import type {
@@ -72,9 +77,11 @@ export abstract class BaseAdapter<AdapterTarget = unknown> {
     this.job = job;
   }
 
-  async fromBlock(model: DraftModel) {
+  async fromBlock(model: BlockModel | DraftModel) {
     try {
-      const blockSnapshot = this.job.blockToSnapshot(model);
+      const draftModel =
+        model instanceof BlockModel ? toDraftModel(model) : model;
+      const blockSnapshot = this.job.blockToSnapshot(draftModel);
       if (!blockSnapshot) return;
       return await this.fromBlockSnapshot({
         snapshot: blockSnapshot,

blocksuite/framework/store/src/model/block/draft.ts

@@ -4,12 +4,16 @@ type PropsInDraft = 'version' | 'flavour' | 'role' | 'id' | 'keys' | 'text';
 
 type ModelProps<Model> = Model extends BlockModel<infer U> ? U : never;
 
+const draftModelSymbol = Symbol('draftModel');
+
 export type DraftModel<Model extends BlockModel = BlockModel> = Pick<
   Model,
   PropsInDraft
 > & {
   children: DraftModel[];
-} & ModelProps<Model>;
+} & ModelProps<Model> & {
+    [draftModelSymbol]: true;
+  };
 
 export function toDraftModel<Model extends BlockModel = BlockModel>(
   origin: Model

blocksuite/framework/store/src/model/store/store.ts

@@ -16,7 +16,6 @@ import {
   type BlockModel,
   type BlockOptions,
   type BlockProps,
-  type DraftModel,
 } from '../block/index.js';
 import type { Doc } from '../doc.js';
 import { DocCRUD } from './crud.js';
@@ -100,10 +99,10 @@ export class Store {
   };
 
   updateBlock: {
-    <T extends Partial<BlockProps>>(model: BlockModel, props: T): void;
+    <T extends Partial<BlockProps>>(model: BlockModel | string, props: T): void;
-    (model: BlockModel, callback: () => void): void;
+    (model: BlockModel | string, callback: () => void): void;
   } = (
-    model: BlockModel,
+    modelOrId: BlockModel | string,
     callBackOrProps: (() => void) | Partial<BlockProps>
   ) => {
     if (this.readonly) {
@@ -113,6 +112,17 @@ export class Store {
 
     const isCallback = typeof callBackOrProps === 'function';
 
+    const model =
+      typeof modelOrId === 'string'
+        ? this.getBlock(modelOrId)?.model
+        : modelOrId;
+    if (!model) {
+      throw new BlockSuiteError(
+        ErrorCode.ModelCRUDError,
+        `updating block: ${modelOrId} not found`
+      );
+    }
+
     if (!isCallback) {
       const parent = this.getParent(model);
       this.schema.validate(
@@ -549,7 +559,7 @@ export class Store {
   }
 
   deleteBlock(
-    model: DraftModel,
+    model: BlockModel | string,
     options: {
       bringChildrenTo?: BlockModel;
       deleteChildren?: boolean;
@@ -575,7 +585,10 @@ export class Store {
     };
 
     this.transact(() => {
-      this._crud.deleteBlock(model.id, opts);
+      this._crud.deleteBlock(
+        typeof model === 'string' ? model : model.id,
+        opts
+      );
     });
   }
 

blocksuite/framework/store/src/transformer/middleware.ts

@@ -1,6 +1,6 @@
 import type { Slot } from '@blocksuite/global/utils';
 
-import type { DraftModel, Store } from '../model/index.js';
+import type { BlockModel, DraftModel, Store } from '../model/index.js';
 import type { AssetsManager } from './assets.js';
 import type { Slice } from './slice.js';
 import type {
@@ -48,7 +48,7 @@ export type BeforeExportPayload =
       type: 'info';
     };
 
-export type FinalPayload =
+export type AfterExportPayload =
   | {
       snapshot: BlockSnapshot;
       type: 'block';
@@ -71,11 +71,34 @@ export type FinalPayload =
       type: 'info';
     };
 
+export type AfterImportPayload =
+  | {
+      snapshot: BlockSnapshot;
+      type: 'block';
+      model: BlockModel;
+      parent?: string;
+      index?: number;
+    }
+  | {
+      snapshot: DocSnapshot;
+      type: 'page';
+      page: Store;
+    }
+  | {
+      snapshot: SliceSnapshot;
+      type: 'slice';
+      slice: Slice;
+    }
+  | {
+      snapshot: CollectionInfoSnapshot;
+      type: 'info';
+    };
+
 export type TransformerSlots = {
   beforeImport: Slot<BeforeImportPayload>;
-  afterImport: Slot<FinalPayload>;
+  afterImport: Slot<AfterImportPayload>;
   beforeExport: Slot<BeforeExportPayload>;
-  afterExport: Slot<FinalPayload>;
+  afterExport: Slot<AfterExportPayload>;
 };
 
 type TransformerMiddlewareOptions = {

blocksuite/framework/store/src/transformer/slice.ts

@@ -1,4 +1,9 @@
-import type { DraftModel, Store } from '../model/index';
+import {
+  BlockModel,
+  type DraftModel,
+  type Store,
+  toDraftModel,
+} from '../model/index';
 
 type SliceData = {
   content: DraftModel[];
@@ -21,9 +26,15 @@ export class Slice {
 
   constructor(readonly data: SliceData) {}
 
-  static fromModels(doc: Store, models: DraftModel[]) {
+  static fromModels(doc: Store, models: DraftModel[] | BlockModel[]) {
+    const draftModels = models.map(model => {
+      if (model instanceof BlockModel) {
+        return toDraftModel(model);
+      }
+      return model;
+    });
     return new Slice({
-      content: models,
+      content: draftModels,
       workspaceId: doc.workspace.id,
       pageId: doc.id,
     });

blocksuite/framework/store/src/transformer/transformer.ts

@@ -1,19 +1,21 @@
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import { nextTick, Slot } from '@blocksuite/global/utils';
 
-import type {
+import {
   BlockModel,
-  BlockSchemaType,
+  type BlockSchemaType,
-  DraftModel,
+  type DraftModel,
-  Store,
+  type Store,
+  toDraftModel,
 } from '../model/index.js';
 import type { Schema } from '../schema/index.js';
 import { AssetsManager } from './assets.js';
 import { BaseBlockTransformer } from './base.js';
 import type {
+  AfterExportPayload,
+  AfterImportPayload,
   BeforeExportPayload,
   BeforeImportPayload,
-  FinalPayload,
   TransformerMiddleware,
   TransformerSlots,
 } from './middleware.js';
@@ -66,14 +68,19 @@ export class Transformer {
 
   private readonly _slots: TransformerSlots = {
     beforeImport: new Slot<BeforeImportPayload>(),
-    afterImport: new Slot<FinalPayload>(),
+    afterImport: new Slot<AfterImportPayload>(),
     beforeExport: new Slot<BeforeExportPayload>(),
-    afterExport: new Slot<FinalPayload>(),
+    afterExport: new Slot<AfterExportPayload>(),
   };
 
-  blockToSnapshot = (model: DraftModel): BlockSnapshot | undefined => {
+  blockToSnapshot = (
+    model: DraftModel | BlockModel
+  ): BlockSnapshot | undefined => {
     try {
-      const snapshot = this._blockToSnapshot(model);
+      const draftModel =
+        model instanceof BlockModel ? toDraftModel(model) : model;
+
+      const snapshot = this._blockToSnapshot(draftModel);
 
       if (!snapshot) {
         return;
@@ -103,7 +110,7 @@ export class Transformer {
           'Root block not found in doc'
         );
       }
-      const blocks = this.blockToSnapshot(rootModel);
+      const blocks = this.blockToSnapshot(toDraftModel(rootModel));
       if (!blocks) {
         return;
       }
@@ -286,7 +293,8 @@ export class Transformer {
 
       const contentBlocks = blockTree.children
         .map(tree => doc.getBlockById(tree.draft.id))
-        .filter(Boolean) as DraftModel[];
+        .filter((x): x is BlockModel => x !== null)
+        .map(model => toDraftModel(model));
 
       const slice = new Slice({
         content: contentBlocks,

packages/frontend/core/src/blocksuite/utils/markdown-utils.ts

@@ -24,7 +24,7 @@ import type {
   Store,
   TransformerMiddleware,
 } from '@blocksuite/affine/store';
-import { Transformer } from '@blocksuite/affine/store';
+import { toDraftModel, Transformer } from '@blocksuite/affine/store';
 
 const updateSnapshotText = (
   point: TextRangePoint,
@@ -51,10 +51,10 @@ function processSnapshot(
 
   const modelId = model.id;
   if (text.from.blockId === modelId) {
-    updateSnapshotText(text.from, snapshot, model);
+    updateSnapshotText(text.from, snapshot, toDraftModel(model));
   }
   if (text.to && text.to.blockId === modelId) {
-    updateSnapshotText(text.to, snapshot, model);
+    updateSnapshotText(text.to, snapshot, toDraftModel(model));
   }
 
   // If the snapshot has children, handle them recursively

packages/frontend/core/src/modules/docs-search/worker/in-worker.ts

@@ -147,7 +147,7 @@ function generateMarkdownPreviewBuilder(
       keys: Array.from(yblock.keys())
         .filter(key => key.startsWith('prop:'))
         .map(key => key.substring(5)),
-    };
+    } as DraftModel;
   }
 
   const titleMiddleware: TransformerMiddleware = ({ adapterConfigs }) => {