fix error handling

Author: YousefED

packages/xl-ai/src/api/formats/json/errorHandling.test.ts

@@ -12,8 +12,7 @@ import { testAIModels } from "../../../testUtil/testAIModels.js";
 import { buildAIRequest } from "../../aiRequest/builder.js";
 
 // Separate test suite for error handling with its own server
-// skipping because it throws a (false) unhandled promise rejection in vitest
-describe.skip("Error handling", () => {
+describe("Error handling", () => {
   // Create a separate server for error tests with custom handlers
   const errorServer = setupServer();
 

packages/xl-ai/src/streamTool/ChunkExecutionError.ts

@@ -0,0 +1,10 @@
+export class ChunkExecutionError extends Error {
+  constructor(
+    message: string,
+    public readonly chunk: any,
+    options?: { cause?: unknown },
+  ) {
+    super(message, options);
+    this.name = "ChunkExecutionError";
+  }
+}

packages/xl-ai/src/streamTool/StreamToolExecutor.ts

@@ -1,18 +1,8 @@
 import { getErrorMessage } from "@ai-sdk/provider-utils";
 import { parsePartialJson } from "ai";
+import { ChunkExecutionError } from "./ChunkExecutionError.js";
 import { StreamTool, StreamToolCall } from "./streamTool.js";
 
-export class ChunkExecutionError extends Error {
-  constructor(
-    message: string,
-    public readonly chunk: any,
-    options?: { cause?: unknown },
-  ) {
-    super(message, options);
-    this.name = "ChunkExecutionError";
-  }
-}
-
 /**
  * The Operation types wraps a StreamToolCall with metadata on whether
  * it's an update to an existing and / or or a possibly partial (i.e.: incomplete, streaming in progress) operation

packages/xl-ai/src/streamTool/preprocess.ts

@@ -1,3 +1,4 @@
+import { ChunkExecutionError } from "./ChunkExecutionError.js";
 import { filterValidOperations } from "./filterValidOperations.js";
 import { StreamTool, StreamToolCall } from "./streamTool.js";
 import { toValidatedOperations } from "./toValidatedOperations.js";
@@ -35,8 +36,8 @@ export async function* preprocessOperationsStreaming<
     (chunk) => {
       if (!chunk.isPossiblyPartial) {
         // only throw if the operation is not possibly partial
-        // TODO: I think there's a bug here in unit tests, for example if operations don't include $. validate with main
-        throw new Error("invalid operation: " + chunk.operation.error);
+        
+        throw new ChunkExecutionError("invalid operation: " + chunk.operation.error, chunk);
       }
     },
   );

packages/xl-ai/src/streamTool/vercelAiSdk/util/chatHandlers.test.ts

@@ -0,0 +1,71 @@
+import { Chat } from "@ai-sdk/react";
+import { BlockNoteEditor } from "@blocknote/core";
+import { ChatTransport, DefaultChatTransport, UIMessage, UIMessageChunk } from "ai";
+import { describe, expect, it } from "vitest";
+import { aiDocumentFormats } from "../../../server.js";
+import { setupToolCallStreaming } from "./chatHandlers.js";
+
+class FakeTransport extends DefaultChatTransport<any> {
+  constructor(private chunks: UIMessageChunk[]) {
+    super();
+  }
+
+  override async sendMessages({ }: Parameters<ChatTransport<UIMessage>['sendMessages']>[0]) {
+   const chunks = this.chunks;
+    return new ReadableStream<UIMessageChunk>({
+      start(controller) {
+        for (const chunk of chunks) {
+          controller.enqueue(chunk);
+        }
+        controller.close();
+      },
+    });
+  }
+}
+
+describe("setupToolCallStreaming", () => {
+  it("should handle missing tool error gracefully", async () => {
+    const editor = BlockNoteEditor.create();
+    const chat = new Chat({
+      transport: new FakeTransport([
+        { type: "start" },
+        { type: "start-step" },
+        {
+          type: "tool-input-start",
+          toolCallId: "call_1",
+          toolName: "applyDocumentOperations",
+        },
+        {
+          type: "tool-input-available",
+          toolCallId: "call_1",
+          toolName: "applyDocumentOperations",
+          input: { operations: [{ type: "testTool", value: "hello" }] },
+        },
+        { type: "finish-step" },
+        { type: "finish", finishReason: "stop" },
+      ]),
+    });
+
+    const streamTools = aiDocumentFormats.html
+      .getStreamToolsProvider({ withDelays: false })
+      .getStreamTools(editor, false);
+
+    const streaming = setupToolCallStreaming(streamTools, chat);
+
+    await chat.sendMessage({
+      role: "user",
+      parts: [{ type: "text", text: "ignored" }],
+    });
+
+    const ret  = await streaming;
+
+    expect(chat.status).toBe("ready");
+    expect(ret.ok).toBe(false);
+    if (!ret.ok) {
+      expect(ret.error).toBeDefined();
+      // We can check if the error message contains "No matching function" or similar
+      // The error is likely wrapped or is the ChunkExecutionError
+      // console.log(ret.error);
+    }
+  });
+});

packages/xl-ai/src/streamTool/vercelAiSdk/util/chatHandlers.ts

@@ -1,11 +1,9 @@
 import { getErrorMessage } from "@ai-sdk/provider-utils";
 import type { Chat } from "@ai-sdk/react";
 import { DeepPartial, isToolUIPart, UIMessage } from "ai";
+import { ChunkExecutionError } from "../../ChunkExecutionError.js";
 import { Result, StreamTool, StreamToolCall } from "../../streamTool.js";
-import {
-  ChunkExecutionError,
-  StreamToolExecutor,
-} from "../../StreamToolExecutor.js";
+import { StreamToolExecutor } from "../../StreamToolExecutor.js";
 import { objectStreamToOperationsResult } from "./UIMessageStreamToOperationsResult.js";
 
 /**
@@ -42,7 +40,7 @@ export async function setupToolCallStreaming(
 
   const appendableStream = createAppendableStream<any>();
 
-  appendableStream.output.pipeTo(executor.writable);
+  const pipeToPromise = appendableStream.output.pipeTo(executor.writable);
 
   const toolCallStreams = new Map<string, ToolCallStreamData>();
 
@@ -112,10 +110,17 @@ export async function setupToolCallStreaming(
   await appendableStream.finalize();
   // let all stream executors finish, this can take longer due to artificial delays
   // (e.g. to simulate human typing behaviour)
-  const result = (await Promise.allSettled([executor.finish()]))[0];
+  const results = await Promise.allSettled([executor.finish(), pipeToPromise]); // awaiting pipeToPromise as well to prevent unhandled promises
+  const result = results[0];
+
+  if (results[1].status === "rejected" && (results[0].status !== "rejected" || results[0].reason !== results[1].reason)){
+    throw new Error("unexpected, pipeToPromise rejected but executor.finish() doesn't have same error!?")
+  }
 
   let error: ChunkExecutionError | undefined;
+  
   if (result.status === "rejected") {
+    
     if (result.reason instanceof ChunkExecutionError) {
       error = result.reason;
     } else {