From 87f95ca40354ff31d85d9696d2183df2a548566e Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Fri, 16 Aug 2024 23:23:16 -0700
Subject: [PATCH] feat: pass abortController to tool call execute function
 (#669)

---
 .../src/components/ui/assistant-ui/thread.tsx   |  6 +++---
 .../src/context/providers/ThreadProvider.tsx    |  3 +--
 packages/react/src/context/stores/Thread.ts     |  7 -------
 .../composer/useComposerSend.tsx                |  2 +-
 .../src/primitives/composer/ComposerInput.tsx   |  4 ++--
 .../react/src/runtimes/edge/EdgeChatAdapter.ts  |  2 +-
 .../runtimes/edge/converters/toCoreMessages.ts  |  2 +-
 .../edge/converters/toLanguageModelMessages.ts  | 17 ++++++++---------
 .../src/runtimes/edge/createEdgeRuntimeAPI.ts   |  4 +++-
 .../runtimes/edge/streams/toolResultStream.ts   | 10 +++++++---
 packages/react/src/types/ModelConfigTypes.ts    |  3 +++
 11 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/packages/react-playground/src/components/ui/assistant-ui/thread.tsx b/packages/react-playground/src/components/ui/assistant-ui/thread.tsx
index 3aebff8a0..7ff31ad51 100644
--- a/packages/react-playground/src/components/ui/assistant-ui/thread.tsx
+++ b/packages/react-playground/src/components/ui/assistant-ui/thread.tsx
@@ -104,7 +104,7 @@ const Composer: FC = () => {
   const { useThread, useThreadActions, useComposer, useThreadMessages } =
     useThreadContext();
 
-  const unstable_canAppendNew = useThread((t) => t.unstable_canAppendNew);
+  const isRunning = useThread((t) => t.isRunning);
   const hasValue = useComposer((c) => c.value.length > 0);
 
   const performAdd = () => {
@@ -167,7 +167,7 @@ const Composer: FC = () => {
         <div className="flex-grow" />
         <Button
           onClick={handleAdd}
-          disabled={!unstable_canAppendNew || !hasValue}
+          disabled={isRunning || !hasValue}
           variant="default"
           size="sm"
           className="transition-opacity"
@@ -176,7 +176,7 @@ const Composer: FC = () => {
         </Button>
         <Button
           onClick={handleSubmit}
-          disabled={!unstable_canAppendNew}
+          disabled={isRunning}
           variant="default"
           size="sm"
           className="transition-opacity"
diff --git a/packages/react/src/context/providers/ThreadProvider.tsx b/packages/react/src/context/providers/ThreadProvider.tsx
index 7cc56db58..35345f3b4 100644
--- a/packages/react/src/context/providers/ThreadProvider.tsx
+++ b/packages/react/src/context/providers/ThreadProvider.tsx
@@ -57,8 +57,7 @@ export const ThreadProvider: FC<PropsWithChildren<ThreadProviderProps>> = ({
           const state = getThreadStateFromRuntime(thread);
           if (
             oldState.isDisabled !== state.isDisabled ||
-            oldState.isRunning !== state.isRunning ||
-            oldState.unstable_canAppendNew !== state.unstable_canAppendNew
+            oldState.isRunning !== state.isRunning
           ) {
             (context.useThread as unknown as StoreApi<ThreadState>).setState(
               getThreadStateFromRuntime(thread),
diff --git a/packages/react/src/context/stores/Thread.ts b/packages/react/src/context/stores/Thread.ts
index 5db870e4b..c11d8f860 100644
--- a/packages/react/src/context/stores/Thread.ts
+++ b/packages/react/src/context/stores/Thread.ts
@@ -5,8 +5,6 @@ import { ThreadRuntimeStore } from "./ThreadRuntime";
 export type ThreadState = Readonly<{
   isRunning: boolean;
   isDisabled: boolean;
-  // TODO remove this once we have cancel tool call support
-  unstable_canAppendNew: boolean;
 }>;
 
 export const getThreadStateFromRuntime = (
@@ -17,15 +15,10 @@ export const getThreadStateFromRuntime = (
     return Object.freeze({
       isDisabled: runtime.isDisabled,
       isRunning: false,
-      unstable_canAppendNew: !runtime.isDisabled,
     });
   return Object.freeze({
     isDisabled: runtime.isDisabled,
     isRunning: lastMessage.status.type === "running",
-    unstable_canAppendNew:
-      !runtime.isDisabled &&
-      lastMessage.status.type !== "running" &&
-      lastMessage.status.type !== "requires-action",
   });
 };
 
diff --git a/packages/react/src/primitive-hooks/composer/useComposerSend.tsx b/packages/react/src/primitive-hooks/composer/useComposerSend.tsx
index 0c4e82d6d..c1f3c2188 100644
--- a/packages/react/src/primitive-hooks/composer/useComposerSend.tsx
+++ b/packages/react/src/primitive-hooks/composer/useComposerSend.tsx
@@ -12,7 +12,7 @@ export const useComposerSend = () => {
 
   const disabled = useCombinedStore(
     [useThread, useComposer],
-    (t, c) => !t.unstable_canAppendNew || !c.isEditing || c.value.length === 0,
+    (t, c) => t.isRunning || !c.isEditing || c.value.length === 0,
   );
 
   const callback = useCallback(() => {
diff --git a/packages/react/src/primitives/composer/ComposerInput.tsx b/packages/react/src/primitives/composer/ComposerInput.tsx
index 16f5908b7..bf6e1234b 100644
--- a/packages/react/src/primitives/composer/ComposerInput.tsx
+++ b/packages/react/src/primitives/composer/ComposerInput.tsx
@@ -63,9 +63,9 @@ export const ComposerPrimitiveInput = forwardRef<
       if (isDisabled) return;
 
       if (e.key === "Enter" && e.shiftKey === false) {
-        const { unstable_canAppendNew } = useThread.getState();
+        const { isRunning } = useThread.getState();
 
-        if (unstable_canAppendNew) {
+        if (!isRunning) {
           e.preventDefault();
 
           textareaRef.current?.closest("form")?.requestSubmit();
diff --git a/packages/react/src/runtimes/edge/EdgeChatAdapter.ts b/packages/react/src/runtimes/edge/EdgeChatAdapter.ts
index 8774cd0d2..95aecf8db 100644
--- a/packages/react/src/runtimes/edge/EdgeChatAdapter.ts
+++ b/packages/react/src/runtimes/edge/EdgeChatAdapter.ts
@@ -55,7 +55,7 @@ export class EdgeChatAdapter implements ChatModelAdapter {
     const stream = result
       .body!.pipeThrough(streamPartDecoderStream())
       .pipeThrough(assistantDecoderStream())
-      .pipeThrough(toolResultStream(config.tools))
+      .pipeThrough(toolResultStream(config.tools, abortSignal))
       .pipeThrough(runResultStream());
 
     let update: ChatModelRunResult | undefined;
diff --git a/packages/react/src/runtimes/edge/converters/toCoreMessages.ts b/packages/react/src/runtimes/edge/converters/toCoreMessages.ts
index 107fcc228..7a2576f92 100644
--- a/packages/react/src/runtimes/edge/converters/toCoreMessages.ts
+++ b/packages/react/src/runtimes/edge/converters/toCoreMessages.ts
@@ -1,4 +1,4 @@
-import { ThreadMessage, CoreMessage } from "../../../types";
+  import { ThreadMessage, CoreMessage } from "../../../types";
 
 export const toCoreMessages = (message: ThreadMessage[]): CoreMessage[] => {
   return message.map(toCoreMessage);
diff --git a/packages/react/src/runtimes/edge/converters/toLanguageModelMessages.ts b/packages/react/src/runtimes/edge/converters/toLanguageModelMessages.ts
index 447665f0e..0f6ff8124 100644
--- a/packages/react/src/runtimes/edge/converters/toLanguageModelMessages.ts
+++ b/packages/react/src/runtimes/edge/converters/toLanguageModelMessages.ts
@@ -52,15 +52,14 @@ const assistantMessageSplitter = () => {
         toolName: part.toolName,
         args: part.args,
       });
-      if (part.result) {
-        toolMessage.content.push({
-          type: "tool-result",
-          toolCallId: part.toolCallId,
-          toolName: part.toolName,
-          result: part.result,
-          // isError
-        });
-      }
+
+      toolMessage.content.push({
+        type: "tool-result",
+        toolCallId: part.toolCallId,
+        toolName: part.toolName,
+        result: part.result ?? "<no result>",
+        isError: part.isError ?? false,
+      });
     },
     getMessages: () => {
       if (toolMessage.content.length > 0) {
diff --git a/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts b/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
index 06bac509f..23467b2f6 100644
--- a/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
+++ b/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
@@ -130,7 +130,9 @@ export const createEdgeRuntimeAPI = ({
     // add tool results if we have server tools
     const canExecuteTools = hasServerTools && toolChoice?.type !== "none";
     if (canExecuteTools) {
-      stream = stream.pipeThrough(toolResultStream(serverTools));
+      stream = stream.pipeThrough(
+        toolResultStream(serverTools, request.signal),
+      );
     }
 
     if (canExecuteTools || onFinish) {
diff --git a/packages/react/src/runtimes/edge/streams/toolResultStream.ts b/packages/react/src/runtimes/edge/streams/toolResultStream.ts
index 5649e9856..8a291d4c3 100644
--- a/packages/react/src/runtimes/edge/streams/toolResultStream.ts
+++ b/packages/react/src/runtimes/edge/streams/toolResultStream.ts
@@ -14,7 +14,10 @@ export type ToolResultStreamPart =
       isError?: boolean;
     };
 
-export function toolResultStream(tools: Record<string, Tool> | undefined) {
+export function toolResultStream(
+  tools: Record<string, Tool> | undefined,
+  abortSignal: AbortSignal,
+) {
   const toolCallExecutions = new Map<string, Promise<any>>();
 
   return new TransformStream<ToolResultStreamPart, ToolResultStreamPart>({
@@ -49,8 +52,10 @@ export function toolResultStream(tools: Record<string, Tool> | undefined) {
               toolCallExecutions.set(
                 toolCallId,
                 (async () => {
+                  if (!tool.execute) return;
+
                   try {
-                    const result = await tool.execute!(args);
+                    const result = await tool.execute(args, { abortSignal });
 
                     controller.enqueue({
                       type: "tool-result",
@@ -60,7 +65,6 @@ export function toolResultStream(tools: Record<string, Tool> | undefined) {
                       result,
                     });
                   } catch (error) {
-                    console.error("Error: ", error);
                     controller.enqueue({
                       type: "tool-result",
                       toolCallType,
diff --git a/packages/react/src/types/ModelConfigTypes.ts b/packages/react/src/types/ModelConfigTypes.ts
index de9390552..33ea6b40c 100644
--- a/packages/react/src/types/ModelConfigTypes.ts
+++ b/packages/react/src/types/ModelConfigTypes.ts
@@ -25,6 +25,9 @@ export type LanguageModelConfig = z.infer<typeof LanguageModelConfigSchema>;
 
 type ToolExecuteFunction<TArgs, TResult> = (
   args: TArgs,
+  context: {
+    abortSignal: AbortSignal;
+  },
 ) => TResult | Promise<TResult>;
 
 export type Tool<