feat(runtimes/edge): cancelled message state (#495)

Yonom · Jul 14, 2024 · d711b3f · d711b3f
1 parent e220617
commit d711b3f
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 25 deletions.
diff --git a/packages/react/src/runtimes/edge/streams/toolResultStream.ts b/packages/react/src/runtimes/edge/streams/toolResultStream.ts
@@ -10,7 +10,8 @@ export type ToolResultStreamPart =
       toolCallType: "function";
       toolCallId: string;
       toolName: string;
-      result: any;
+      result: unknown;
+      isError?: boolean;
     };
 
 export function toolResultStream(tools: Record<string, Tool> | undefined) {
@@ -34,8 +35,14 @@ export function toolResultStream(tools: Record<string, Tool> | undefined) {
             const result = tool.parameters.safeParse(args);
             if (!result.success) {
               controller.enqueue({
-                type: "error",
-                error: new Error("Invalid tool call arguments"),
+                type: "tool-result",
+                toolCallType,
+                toolCallId,
+                toolName,
+                result:
+                  "Function parameter validation failed. " +
+                  JSON.stringify(result.error.issues),
+                isError: true,
               });
               return;
             } else {
@@ -53,9 +60,14 @@ export function toolResultStream(tools: Record<string, Tool> | undefined) {
                       result,
                     });
                   } catch (error) {
+                    console.error("Error: ", error);
                     controller.enqueue({
-                      type: "error",
-                      error,
+                      type: "tool-result",
+                      toolCallType,
+                      toolCallId,
+                      toolName,
+                      result: "Error: " + error,
+                      isError: true,
                     });
                   } finally {
                     toolCallExecutions.delete(toolCallId);
@@ -81,6 +93,7 @@ export function toolResultStream(tools: Record<string, Tool> | undefined) {
         }
       }
     },
+
     async flush() {
       await Promise.all(toolCallExecutions.values());
     },

diff --git a/packages/react/src/runtimes/local/LocalRuntime.tsx b/packages/react/src/runtimes/local/LocalRuntime.tsx
@@ -118,52 +118,56 @@ class LocalThreadRuntime implements ThreadRuntime {
     const messages = this.repository.getMessages();
 
     // add assistant message
-    const message: ThreadAssistantMessage = {
+    let message: ThreadAssistantMessage = {
       id: generateId(),
       role: "assistant",
       status: { type: "in_progress" },
       content: [{ type: "text", text: "" }],
       createdAt: new Date(),
     };
-    this.repository.addOrUpdateMessage(parentId, { ...message });
 
     // abort existing run
     this.abortController?.abort();
     this.abortController = new AbortController();
 
+    this.repository.addOrUpdateMessage(parentId, { ...message });
     this.notifySubscribers();
 
-    try {
-      const updateHandler = ({ content }: ChatModelRunResult) => {
-        message.content = content;
-        const newMessage = { ...message };
-        this.repository.addOrUpdateMessage(parentId, newMessage);
-        this.notifySubscribers();
-        return newMessage;
+    const updateMessage = (m: Partial<ChatModelRunResult>) => {
+      message = {
+        ...message,
+        ...m,
       };
+      this.repository.addOrUpdateMessage(parentId, message);
+      this.notifySubscribers();
+      return message;
+    };
+
+    try {
       const result = await this.adapter.run({
         messages,
         abortSignal: this.abortController.signal,
         config: this.configProvider.getModelConfig(),
-        onUpdate: updateHandler,
+        onUpdate: updateMessage,
       });
-      if (result !== undefined) {
-        updateHandler(result);
-      }
       if (result.status?.type === "in_progress")
         throw new Error(
           "Unexpected in_progress status returned from ChatModelAdapter",
         );
 
-      message.status = result.status ?? { type: "done" };
+      this.abortController = null;
+      updateMessage({ status: { type: "done" }, ...result });
       this.repository.addOrUpdateMessage(parentId, { ...message });
     } catch (e) {
-      message.status = { type: "error", error: e };
-      this.repository.addOrUpdateMessage(parentId, { ...message });
-      throw e;
-    } finally {
+      const isAbortError = e instanceof Error && e.name === "AbortError";
       this.abortController = null;
-      this.notifySubscribers();
+      updateMessage({
+        status: isAbortError
+          ? { type: "cancelled" }
+          : { type: "error", error: e },
+      });
+
+      if (!isAbortError) throw e;
     }
   }
 

diff --git a/packages/react/src/types/AssistantTypes.ts b/packages/react/src/types/AssistantTypes.ts
@@ -61,7 +61,7 @@ type MessageCommonProps = {
 
 export type MessageStatus =
   | {
-      type: "in_progress";
+      type: "in_progress" | "cancelled";
     }
   | {
       type: "done";