From 0ff22a729df146a77f3832b198a5d915b678bb58 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Tue, 8 Oct 2024 16:02:05 -0700
Subject: [PATCH] feat: switch to DataStream transfer protocol for edge runtime
 (#953)

---
 .changeset/rich-trees-care.md                 |  5 +++++
 .../edge/streams/AssistantStreamChunkType.ts  | 21 ++++++++++--------
 .../edge/streams/assistantDecoderStream.ts    | 22 ++++++++++++-------
 .../edge/streams/assistantEncoderStream.ts    | 13 ++++++-----
 4 files changed, 39 insertions(+), 22 deletions(-)
 create mode 100644 .changeset/rich-trees-care.md

diff --git a/.changeset/rich-trees-care.md b/.changeset/rich-trees-care.md
new file mode 100644
index 000000000..e7b6e94a6
--- /dev/null
+++ b/.changeset/rich-trees-care.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: switch to DataStream transfer protocol for edge runtime
diff --git a/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts b/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts
index c62b95836..94915cbca 100644
--- a/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts
+++ b/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts
@@ -2,22 +2,25 @@ import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
 
 export enum AssistantStreamChunkType {
   TextDelta = "0",
-  ToolCallBegin = "1",
-  ToolCallArgsTextDelta = "2",
-  ToolCallResult = "3",
-  Error = "E",
-  Finish = "F",
+  ToolCallBegin = "b",
+  ToolCallDelta = "c",
+  ToolCallResult = "a",
+  Error = "3",
+  Finish = "d",
 }
 
 export type AssistantStreamChunk = {
   [AssistantStreamChunkType.TextDelta]: string;
   [AssistantStreamChunkType.ToolCallBegin]: {
-    id: string;
-    name: string;
+    toolCallId: string;
+    toolName: string;
+  };
+  [AssistantStreamChunkType.ToolCallDelta]: {
+    toolCallId: string;
+    argsTextDelta: string;
   };
-  [AssistantStreamChunkType.ToolCallArgsTextDelta]: string;
   [AssistantStreamChunkType.ToolCallResult]: {
-    id: string;
+    toolCallId: string;
     result: any;
   };
   [AssistantStreamChunkType.Error]: unknown;
diff --git a/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts b/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts
index a96f6afe0..a67c596b9 100644
--- a/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts
+++ b/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts
@@ -18,7 +18,7 @@ export function assistantDecoderStream() {
     transform({ type, value }, controller) {
       if (
         currentToolCall &&
-        type !== AssistantStreamChunkType.ToolCallArgsTextDelta &&
+        type !== AssistantStreamChunkType.ToolCallDelta &&
         type !== AssistantStreamChunkType.Error
       ) {
         controller.enqueue({
@@ -40,20 +40,26 @@ export function assistantDecoderStream() {
           break;
         }
         case AssistantStreamChunkType.ToolCallBegin: {
-          const { id, name } = value;
+          const { toolCallId: id, toolName: name } = value;
           toolCallNames.set(id, name);
           currentToolCall = { id, name, argsText: "" };
           break;
         }
-        case AssistantStreamChunkType.ToolCallArgsTextDelta: {
-          const delta = value;
-          currentToolCall!.argsText += delta;
+        case AssistantStreamChunkType.ToolCallDelta: {
+          const { toolCallId, argsTextDelta } = value;
+          if (currentToolCall?.id !== toolCallId) {
+            throw new Error(
+              `Received tool call delta for unknown tool call "${toolCallId}".`,
+            );
+          }
+
+          currentToolCall!.argsText += argsTextDelta;
           controller.enqueue({
             type: "tool-call-delta",
             toolCallType: "function",
             toolCallId: currentToolCall!.id,
             toolName: currentToolCall!.name,
-            argsTextDelta: delta,
+            argsTextDelta: argsTextDelta,
           });
           break;
         }
@@ -61,8 +67,8 @@ export function assistantDecoderStream() {
           controller.enqueue({
             type: "tool-result",
             toolCallType: "function",
-            toolCallId: value.id,
-            toolName: toolCallNames.get(value.id)!,
+            toolCallId: value.toolCallId,
+            toolName: toolCallNames.get(value.toolCallId)!,
             result: value.result,
           });
           break;
diff --git a/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts b/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts
index 37fe47ef5..38e3a68cf 100644
--- a/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts
+++ b/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts
@@ -27,15 +27,18 @@ export function assistantEncoderStream() {
             controller.enqueue({
               type: AssistantStreamChunkType.ToolCallBegin,
               value: {
-                id: chunk.toolCallId,
-                name: chunk.toolName,
+                toolCallId: chunk.toolCallId,
+                toolName: chunk.toolName,
               },
             });
           }
 
           controller.enqueue({
-            type: AssistantStreamChunkType.ToolCallArgsTextDelta,
-            value: chunk.argsTextDelta,
+            type: AssistantStreamChunkType.ToolCallDelta,
+            value: {
+              toolCallId: chunk.toolCallId,
+              argsTextDelta: chunk.argsTextDelta,
+            },
           });
           break;
         }
@@ -49,7 +52,7 @@ export function assistantEncoderStream() {
           controller.enqueue({
             type: AssistantStreamChunkType.ToolCallResult,
             value: {
-              id: chunk.toolCallId,
+              toolCallId: chunk.toolCallId,
               result: chunk.result,
             },
           });