From a73f0d47d727baef409caa12afb80ebb434aaf77 Mon Sep 17 00:00:00 2001
From: Felix Feng <felixfeng33@gmail.com>
Date: Sat, 19 Oct 2024 22:38:54 +0800
Subject: [PATCH] withNewBatch

---
 .../ai/src/lib/transforms/insertAINodes.ts    | 27 +++++--------------
 .../react/ai-chat/transforms/acceptAIChat.ts  |  4 +--
 .../transforms/replaceSelectionAIChat.ts      |  4 +--
 .../ai/src/react/ai-chat/useAIChatHook.ts     | 12 ++++-----
 .../src/interfaces/history-editor/index.ts    |  1 +
 .../interfaces/history-editor/withNewBatch.ts |  4 +--
 6 files changed, 20 insertions(+), 32 deletions(-)

diff --git a/packages/ai/src/lib/transforms/insertAINodes.ts b/packages/ai/src/lib/transforms/insertAINodes.ts
index 59889471dc..3763c53441 100644
--- a/packages/ai/src/lib/transforms/insertAINodes.ts
+++ b/packages/ai/src/lib/transforms/insertAINodes.ts
@@ -6,8 +6,7 @@ import {
   collapseSelection,
   getEndPoint,
   insertNodes,
-  withMerging,
-  withoutMergingHistory,
+  withNewBatch,
 } from '@udecode/plate-common';
 
 import { AIPlugin } from '../../react/ai/AIPlugin';
@@ -16,10 +15,10 @@ export const insertAINodes = (
   editor: SlateEditor,
   nodes: TDescendant[],
   {
-    history = 'default',
+    splitHistory = false,
     target,
   }: {
-    history?: 'default' | 'merge' | 'withoutMerge';
+    splitHistory?: boolean;
     target?: Path;
   } = {}
 ) => {
@@ -38,21 +37,9 @@ export const insertAINodes = (
     collapseSelection(editor, { edge: 'end' });
   };
 
-  switch (history) {
-    case 'default': {
-      insert();
-
-      break;
-    }
-    case 'merge': {
-      withMerging(editor, insert);
-
-      break;
-    }
-    case 'withoutMerge': {
-      withoutMergingHistory(editor, insert);
-
-      break;
-    }
+  if (splitHistory) {
+    withNewBatch(editor, insert);
+  } else {
+    insert();
   }
 };
diff --git a/packages/ai/src/react/ai-chat/transforms/acceptAIChat.ts b/packages/ai/src/react/ai-chat/transforms/acceptAIChat.ts
index 0c138a900e..b560aa0c9e 100644
--- a/packages/ai/src/react/ai-chat/transforms/acceptAIChat.ts
+++ b/packages/ai/src/react/ai-chat/transforms/acceptAIChat.ts
@@ -1,4 +1,4 @@
-import { withMerging } from '@udecode/plate-common';
+import { withNewBatch } from '@udecode/plate-common';
 import {
   type PlateEditor,
   focusEditor,
@@ -12,7 +12,7 @@ import { AIPlugin } from '../../ai/AIPlugin';
 export const acceptAIChat = (editor: PlateEditor) => {
   const { tf } = getEditorPlugin(editor, AIPlugin);
 
-  withMerging(editor, () => {
+  withNewBatch(editor, () => {
     tf.ai.removeMarks();
   });
 
diff --git a/packages/ai/src/react/ai-chat/transforms/replaceSelectionAIChat.ts b/packages/ai/src/react/ai-chat/transforms/replaceSelectionAIChat.ts
index a929db50ca..454a03eb63 100644
--- a/packages/ai/src/react/ai-chat/transforms/replaceSelectionAIChat.ts
+++ b/packages/ai/src/react/ai-chat/transforms/replaceSelectionAIChat.ts
@@ -1,6 +1,6 @@
 import type { PlateEditor } from '@udecode/plate-common/react';
 
-import { isEditorEmpty, withMerging } from '@udecode/plate-common';
+import { isEditorEmpty, withNewBatch } from '@udecode/plate-common';
 import { focusEditor } from '@udecode/plate-common/react';
 import {
   BlockSelectionPlugin,
@@ -31,7 +31,7 @@ export const replaceSelectionAIChat = (
     editor.withoutNormalizing(() => {
       removeBlockSelectionNodes(editor);
 
-      withMerging(editor, () => {
+      withNewBatch(editor, () => {
         editor
           .getTransforms(BlockSelectionPlugin)
           .blockSelection.insertBlocksAndSelect(
diff --git a/packages/ai/src/react/ai-chat/useAIChatHook.ts b/packages/ai/src/react/ai-chat/useAIChatHook.ts
index cc47037907..622cee23d1 100644
--- a/packages/ai/src/react/ai-chat/useAIChatHook.ts
+++ b/packages/ai/src/react/ai-chat/useAIChatHook.ts
@@ -13,9 +13,9 @@ export const useAIChatHooks = () => {
   const mode = useOption('mode');
 
   useChatChunk({
-    onChunk: ({ isFirst, nodes }) => {
+    onChunk: ({ nodes }) => {
       if (mode === 'insert' && nodes.length > 0) {
-        tf.ai.insertNodes(nodes, { history: isFirst ? 'default' : 'merge' });
+        tf.ai.insertNodes(nodes, { splitHistory: true });
       }
     },
     onFinish: ({ content }) => {
@@ -28,11 +28,11 @@ export const useAIChatHooks = () => {
       editor.undo();
       editor.history.redos.pop();
 
-      setTimeout(() => {
-        const nodes = deserializeInlineMd(editor, content);
+      // setTimeout(() => {
+      const nodes = deserializeInlineMd(editor, content);
 
-        tf.ai.insertNodes(nodes);
-      }, 0);
+      tf.ai.insertNodes(nodes, { splitHistory: true });
+      // }, 0);
     },
   });
 };
diff --git a/packages/slate/src/interfaces/history-editor/index.ts b/packages/slate/src/interfaces/history-editor/index.ts
index b83ae2ed23..4c2aeafd9e 100644
--- a/packages/slate/src/interfaces/history-editor/index.ts
+++ b/packages/slate/src/interfaces/history-editor/index.ts
@@ -8,3 +8,4 @@ export * from './isHistorySaving';
 export * from './withMerging';
 export * from './withoutMergingHistory';
 export * from './withoutSavingHistory';
+export * from './withNewBatch';
diff --git a/packages/slate/src/interfaces/history-editor/withNewBatch.ts b/packages/slate/src/interfaces/history-editor/withNewBatch.ts
index a001407046..39b13c0acb 100644
--- a/packages/slate/src/interfaces/history-editor/withNewBatch.ts
+++ b/packages/slate/src/interfaces/history-editor/withNewBatch.ts
@@ -3,5 +3,5 @@ import type { TEditor } from '../editor';
 import { HistoryEditor } from '../../slate-history';
 
 /** {@link HistoryEditor.withMerging} */
-export const withMerging = (editor: TEditor, fn: () => void) =>
-  HistoryEditor.withMerging(editor as any, fn);
+export const withNewBatch = (editor: TEditor, fn: () => void) =>
+  HistoryEditor.withNewBatch(editor as any, fn);