From 633d6fefcacabd201e2ec96cde868bebdc8d0f32 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Sun, 28 Jul 2024 11:50:35 -0700
Subject: [PATCH] feat: support new streaming mechanism in the playground
 (#601)

---
 packages/react-playground/package.json        |  2 +-
 .../src/lib/playground-runtime.ts             | 28 +++++++++++++------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/packages/react-playground/package.json b/packages/react-playground/package.json
index fb383edc7..c5bf97de1 100644
--- a/packages/react-playground/package.json
+++ b/packages/react-playground/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react-playground",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "license": "MIT",
   "exports": {
     ".": {
diff --git a/packages/react-playground/src/lib/playground-runtime.ts b/packages/react-playground/src/lib/playground-runtime.ts
index a5d062ac6..ff1841a48 100644
--- a/packages/react-playground/src/lib/playground-runtime.ts
+++ b/packages/react-playground/src/lib/playground-runtime.ts
@@ -193,21 +193,31 @@ export class PlaygroundThreadRuntime implements ReactThreadRuntime {
     };
 
     try {
-      const result = await this.adapter.run({
+      const promiseOrGenerator = this.adapter.run({
         messages,
         abortSignal: this.abortController.signal,
         config: this.configProvider.getModelConfig(),
         onUpdate: updateMessage,
       });
-      if (result.status?.type === "running")
-        throw new Error(
-          "Unexpected running status returned from ChatModelAdapter",
-        );
+
+      if (Symbol.asyncIterator in promiseOrGenerator) {
+        for await (const r of promiseOrGenerator) {
+          updateMessage(r);
+        }
+      } else {
+        updateMessage(await promiseOrGenerator);
+      }
+
       this.abortController = null;
-      updateMessage({
-        status: { type: "complete", reason: "unknown" },
-        ...result,
-      });
+
+      if (message.status.type === "running") {
+        updateMessage({
+          status: { type: "complete", reason: "unknown" },
+        });
+      } else {
+        // notify subscribers that isRunning is now false
+        this.notifySubscribers();
+      }
     } catch (e) {
       this.abortController = null;