fix

udecode · Oct 24, 2024 · 6a2b3af · 6a2b3af
1 parent 2e9ced7
commit 6a2b3af
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/.changeset/stale-zebras-invent.md b/.changeset/stale-zebras-invent.md
@@ -0,0 +1,5 @@
+---
+'@udecode/plate-ai': patch
+---
+
+Fix undo after accept
diff --git a/packages/ai/src/lib/transforms/withAIBatch.ts b/packages/ai/src/lib/transforms/withAIBatch.ts
@@ -1,9 +1,12 @@
 import {
+  type History,
   type SlateEditor,
   withMerging,
   withNewBatch,
 } from '@udecode/plate-common';
 
+export type AIBatch = History['undos'][number] & { ai?: boolean };
+
 export const withAIBatch = (
   editor: SlateEditor,
   fn: () => void,
@@ -13,15 +16,15 @@ export const withAIBatch = (
     split?: boolean;
   } = {}
 ) => {
-  if (!split && (editor.history.undos?.at(-1) as any)?.ai) {
+  const lastBatch = editor.history.undos?.at(-1) as AIBatch | undefined;
+
+  if (!split && lastBatch?.ai) {
     withMerging(editor, fn);
   } else {
     withNewBatch(editor, fn);
 
-    const batch = editor.history.undos?.at(-1) as any;
-
-    if (batch) {
-      batch.ai = true;
+    if (lastBatch) {
+      lastBatch.ai = true;
     }
   }
 };
diff --git a/packages/ai/src/react/ai-chat/AIChatPlugin.ts b/packages/ai/src/react/ai-chat/AIChatPlugin.ts
@@ -14,6 +14,8 @@ import {
 } from '@udecode/plate-common/react';
 import { BlockSelectionPlugin } from '@udecode/plate-selection/react';
 
+import type { AIBatch } from '../../lib';
+
 import { AIPlugin } from '../ai/AIPlugin';
 import { acceptAIChat } from './transforms/acceptAIChat';
 import { insertBelowAIChat } from './transforms/insertBelowAIChat';
@@ -139,6 +141,12 @@ export const AIChatPlugin = createTPlatePlugin<AIChatPluginConfig>({
       } else {
         focusEditor(editor);
       }
+
+      const lastBatch = editor.history.undos.at(-1) as AIBatch;
+
+      if (lastBatch?.ai) {
+        delete lastBatch.ai;
+      }
     },
     show: () => {
       api.aiChat.reset();