add abort system

TypeCellOS · YousefED · Jun 12, 2025 · Jul 1, 2025 · Jul 1, 2025 · Jul 1, 2025
commit ec8a62b7cc146737d6b77a0598cdd27901009cf1
diff --git a/packages/xl-ai/src/AIExtension.ts b/packages/xl-ai/src/AIExtension.ts
@@ -314,13 +314,6 @@ export class AIExtension extends BlockNoteExtension {
 
     // Abort the current request
     this.currentAbortController?.abort();
-
-    // Keep partial changes as requested by user
-    // If in collaboration mode, merge the changes back into the original yDoc
-    this.editor.forkYDocPlugin?.merge({ keepChanges: true });
-
-    // Close the AI menu and clean up
-    this.closeAIMenu();
   }
 
   /**
@@ -375,6 +368,7 @@ export class AIExtension extends BlockNoteExtension {
    * Execute a call to an LLM and apply the result to the editor
    */
   public async callLLM(opts: MakeOptional<LLMRequestOptions, "model">) {
+    const startState = this.store.getState().aiMenuState;
     this.setAIResponseStatus("thinking");
     this.editor.forkYDocPlugin?.fork();
 
@@ -413,18 +407,33 @@ export class AIExtension extends BlockNoteExtension {
         llmResponse: ret,
       });
 
-      await ret.execute();
+      await ret.execute(this.currentAbortController?.signal);
 
       this.setAIResponseStatus("user-reviewing");
     } catch (e) {
       // Handle abort errors gracefully
       if (e instanceof Error && e.name === "AbortError") {
         // Request was aborted, don't set error status as abort() handles cleanup
+        const state = this.store.getState().aiMenuState;
+        if (state === "closed" || startState === "closed") {
+          throw new Error(
+            "Unexpected: AbortError occurred while the AI menu was closed",
+          );
+        }
+        if (state.status === "ai-writing") {
+          // we were already writing. Set to reviewing to show the user the partial result
+          this.setAIResponseStatus("user-reviewing");
+        } else {
+          // we were not writing yet. Set to the previous state
+          if (startState.status === "error") {
+            this.setAIResponseStatus({ status: startState.status, error: e });
+          } else {
+            this.setAIResponseStatus(startState.status);
+          }
+        }
         return ret;
       }
 
-      // TODO in error state, should we discard the forked document?
-
       this.setAIResponseStatus({
         status: "error",
         error: e,

diff --git a/packages/xl-ai/src/api/LLMResponse.ts b/packages/xl-ai/src/api/LLMResponse.ts
@@ -28,14 +28,14 @@ export class LLMResponse {
    *
    * (this method consumes underlying streams in `llmResult`)
    */
-  async *applyToolCalls() {
+  async *applyToolCalls(abortSignal?: AbortSignal) {
     let currentStream: AsyncIterable<{
       operation: StreamToolCall<StreamTool<any>[]>;
       isUpdateToPreviousOperation: boolean;
       isPossiblyPartial: boolean;
     }> = this.llmResult.operationsSource;
     for (const tool of this.streamTools) {
-      currentStream = tool.execute(currentStream);
+      currentStream = tool.execute(currentStream, abortSignal);
     }
     yield* currentStream;
   }
@@ -45,9 +45,9 @@ export class LLMResponse {
    *
    * (this method consumes underlying streams in `llmResult`)
    */
-  public async execute() {
+  public async execute(abortSignal?: AbortSignal) {
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    for await (const _result of this.applyToolCalls()) {
+    for await (const _result of this.applyToolCalls(abortSignal)) {
       // no op
     }
   }

diff --git a/packages/xl-ai/src/api/formats/base-tools/createAddBlocksTool.ts b/packages/xl-ai/src/api/formats/base-tools/createAddBlocksTool.ts
@@ -169,7 +169,7 @@ export function createAddBlocksTool<T>(config: {
       },
       // Note: functionality mostly tested in jsontools.test.ts
       // would be nicer to add a direct unit test
-      execute: async function* (operationsStream) {
+      execute: async function* (operationsStream, abortSignal?: AbortSignal) {
         // An add operation has some complexity:
         // - it can add multiple blocks in 1 operation
         //   (this is needed because you need an id as reference block - and if you want to insert multiple blocks you can only use an existing block as reference id)
@@ -266,6 +266,11 @@ export function createAddBlocksTool<T>(config: {
             }
 
             for (const step of agentSteps) {
+              if (abortSignal?.aborted) {
+                const error = new Error("Operation was aborted");
+                error.name = "AbortError";
+                throw error;
+              }
               if (options.withDelays) {
                 await delayAgentStep(step);
               }

diff --git a/packages/xl-ai/src/api/formats/base-tools/createUpdateBlockTool.ts b/packages/xl-ai/src/api/formats/base-tools/createUpdateBlockTool.ts
@@ -178,6 +178,7 @@ export function createUpdateBlockTool<T>(config: {
           isUpdateToPreviousOperation: boolean;
           isPossiblyPartial: boolean;
         }>,
+        abortSignal?: AbortSignal,
       ) {
         const STEP_SIZE = 50;
         let minSize = STEP_SIZE;
@@ -254,6 +255,11 @@ export function createUpdateBlockTool<T>(config: {
           const agentSteps = getStepsAsAgent(tr);
 
           for (const step of agentSteps) {
+            if (abortSignal?.aborted) {
+              const error = new Error("Operation was aborted");
+              error.name = "AbortError";
+              throw error;
+            }
             if (options.withDelays) {
               await delayAgentStep(step);
             }

diff --git a/packages/xl-ai/src/api/formats/base-tools/delete.ts b/packages/xl-ai/src/api/formats/base-tools/delete.ts
@@ -76,7 +76,7 @@ export const deleteBlockTool = (
     },
     // Note: functionality mostly tested in jsontools.test.ts
     // would be nicer to add a direct unit test
-    execute: async function* (operationsStream) {
+    execute: async function* (operationsStream, abortSignal?: AbortSignal) {
       for await (const chunk of operationsStream) {
         if (chunk.operation.type !== "delete") {
           // pass through non-delete operations
@@ -93,6 +93,11 @@ export const deleteBlockTool = (
         const agentSteps = getStepsAsAgent(tr);
 
         for (const step of agentSteps) {
+          if (abortSignal?.aborted) {
+            const error = new Error("Operation was aborted");
+            error.name = "AbortError";
+            throw error;
+          }
           if (options.withDelays) {
             await delayAgentStep(step);
           }

diff --git a/packages/xl-ai/src/blocknoteAIClient/client.ts b/packages/xl-ai/src/blocknoteAIClient/client.ts
@@ -18,13 +18,20 @@ const fetchViaBlockNoteAIServer =
         body: init?.body || request.body,
         method: request.method,
         duplex: "half",
+        signal: request.signal,
       } as any,
     );
     try {
       const resp = await fetch(newRequest);
       return resp;
     } catch (e) {
       // Temp fix for https://github.com/vercel/ai/issues/6370
+      if (
+        e instanceof Error &&
+        (e.name === "AbortError" || e.name === "TimeoutError")
+      ) {
+        throw e;
+      }
       throw new TypeError("fetch failed", {
         cause: e,
       });

diff --git a/packages/xl-ai/src/components/AIMenu/AIMenu.tsx b/packages/xl-ai/src/components/AIMenu/AIMenu.tsx
@@ -94,9 +94,12 @@ export const AIMenu = (props: AIMenuProps) => {
   const rightSection = useMemo(() => {
     if (aiResponseStatus === "thinking" || aiResponseStatus === "ai-writing") {
       return (
-        <Components.SuggestionMenu.Loader
-          className={"bn-suggestion-menu-loader bn-combobox-right-section"}
-        />
+        // TODO
+        <div onClick={() => ai.abort()}>
+          <Components.SuggestionMenu.Loader
+            className={"bn-suggestion-menu-loader bn-combobox-right-section"}
+          />
+        </div>
       );
     } else if (aiResponseStatus === "error") {
       return (

diff --git a/packages/xl-ai/src/streamTool/callLLMWithStreamTools.ts b/packages/xl-ai/src/streamTool/callLLMWithStreamTools.ts
@@ -15,6 +15,7 @@ import {
   AsyncIterableStream,
   createAsyncIterableStream,
   createAsyncIterableStreamFromAsyncIterable,
+  withAbort,
 } from "../util/stream.js";
 import { filterNewOrUpdatedOperations } from "./filterNewOrUpdatedOperations.js";
 import {
@@ -118,6 +119,13 @@ export async function generateOperations<T extends StreamTool<any>[]>(
 
   const ret = await generateObject<{ operations: any }>(options);
 
+  if (opts.abortSignal?.aborted) {
+    // throw abort error before stream processing starts and `onStart` is called
+    const error = new Error("Operation was aborted");
+    error.name = "AbortError";
+    throw error;
+  }
+
   // because the rest of the codebase always expects a stream, we convert the object to a stream here
   const stream = operationsToStream(ret.object);
 
@@ -133,7 +141,10 @@ export async function generateOperations<T extends StreamTool<any>[]>(
     get operationsSource() {
       if (!_operationsSource) {
         _operationsSource = createAsyncIterableStreamFromAsyncIterable(
-          preprocessOperationsNonStreaming(stream.value, streamTools),
+          withAbort(
+            preprocessOperationsNonStreaming(stream.value, streamTools),
+            opts.abortSignal,
+          ),
         );
       }
       return _operationsSource;
@@ -259,8 +270,11 @@ export async function streamOperations<T extends StreamTool<any>[]>(
           preprocessOperationsStreaming(
             filterNewOrUpdatedOperations(
               streamOnStartCallback(
-                partialObjectStreamThrowError(
-                  createAsyncIterableStream(fullStream1),
+                withAbort(
+                  partialObjectStreamThrowError(
+                    createAsyncIterableStream(fullStream1),
+                  ),
+                  opts.abortSignal,
                 ),
                 onStart,
               ),

diff --git a/packages/xl-ai/src/streamTool/streamTool.ts b/packages/xl-ai/src/streamTool/streamTool.ts
@@ -51,6 +51,7 @@ export type StreamTool<T extends { type: string }> = {
       isUpdateToPreviousOperation: boolean;
       isPossiblyPartial: boolean;
     }>,
+    abortSignal?: AbortSignal,
   ) => AsyncIterable<{
     operation: StreamToolCall<StreamTool<{ type: string }>[]>;
     isUpdateToPreviousOperation: boolean;

diff --git a/packages/xl-ai/src/util/stream.ts b/packages/xl-ai/src/util/stream.ts
@@ -2,7 +2,7 @@
  * Converts an AsyncIterable to a ReadableStream
  */
 export function asyncIterableToStream<T>(
-  iterable: AsyncIterable<T>
+  iterable: AsyncIterable<T>,
 ): ReadableStream<T> {
   return new ReadableStream({
     async start(controller) {
@@ -29,11 +29,11 @@ export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
  * Creates an AsyncIterableStream from a ReadableStream
  */
 export function createAsyncIterableStream<T>(
-  source: ReadableStream<T>
+  source: ReadableStream<T>,
 ): AsyncIterableStream<T> {
   if (source.locked) {
     throw new Error(
-      "Stream (source) is already locked and cannot be iterated."
+      "Stream (source) is already locked and cannot be iterated.",
     );
   }
 
@@ -60,7 +60,24 @@ export function createAsyncIterableStream<T>(
  * Creates an AsyncIterableStream from an AsyncGenerator
  */
 export function createAsyncIterableStreamFromAsyncIterable<T>(
-  source: AsyncIterable<T>
+  source: AsyncIterable<T>,
 ): AsyncIterableStream<T> {
   return createAsyncIterableStream(asyncIterableToStream(source));
 }
+
+/**
+ * Helper to wrap an async iterable and throw if the abort signal is triggered.
+ */
+export async function* withAbort<T>(
+  iterable: AsyncIterable<T>,
+  signal?: AbortSignal,
+) {
+  for await (const item of iterable) {
+    if (signal?.aborted) {
+      const error = new Error("Operation was aborted");
+      error.name = "AbortError";
+      throw error;
+    }
+    yield item;
+  }
+}