From b0a22e34f20e7d3600153fe8fafde3e029597f23 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Mon, 16 Sep 2024 13:38:08 -0700
Subject: [PATCH] feat: runtime.switchToNewThread() (#849)

---
 .changeset/early-dragons-float.md             |  5 +++++
 .../src/lib/playground-runtime.ts             | 12 ++++++----
 .../src/runtimes/core/AssistantRuntime.tsx    |  4 ++++
 .../runtimes/core/BaseAssistantRuntime.tsx    |  2 ++
 .../external-store/ExternalStoreRuntime.tsx   | 22 +++++++++++--------
 .../react/src/runtimes/local/LocalRuntime.tsx | 14 +++++++-----
 6 files changed, 41 insertions(+), 18 deletions(-)
 create mode 100644 .changeset/early-dragons-float.md

diff --git a/.changeset/early-dragons-float.md b/.changeset/early-dragons-float.md
new file mode 100644
index 000000000..077436d8e
--- /dev/null
+++ b/.changeset/early-dragons-float.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: runtime.switchToNewThread()
diff --git a/packages/react-playground/src/lib/playground-runtime.ts b/packages/react-playground/src/lib/playground-runtime.ts
index 41d796dcf..67cf0ea88 100644
--- a/packages/react-playground/src/lib/playground-runtime.ts
+++ b/packages/react-playground/src/lib/playground-runtime.ts
@@ -66,10 +66,7 @@ class PlaygroundRuntime
     this._proxyConfigProvider = cp;
   }
 
-  public switchToThread(threadId: string | null) {
-    if (threadId)
-      throw new Error("PlaygroundRuntime does not support switching threads");
-
+  public switchToNewThread() {
     this.thread = new PlaygroundThreadRuntime(
       this._proxyConfigProvider,
       [],
@@ -77,6 +74,13 @@ class PlaygroundRuntime
     );
   }
 
+  public switchToThread(threadId: string | null) {
+    if (threadId !== null)
+      throw new Error("PlaygroundRuntime does not support switching threads");
+
+    this.switchToNewThread();
+  }
+
   public override registerModelConfigProvider(
     provider: ModelConfigProvider,
   ): Unsubscribe {
diff --git a/packages/react/src/runtimes/core/AssistantRuntime.tsx b/packages/react/src/runtimes/core/AssistantRuntime.tsx
index 38cae5b80..ace2e7012 100644
--- a/packages/react/src/runtimes/core/AssistantRuntime.tsx
+++ b/packages/react/src/runtimes/core/AssistantRuntime.tsx
@@ -8,6 +8,10 @@ export type ThreadRuntimeWithSubscribe = {
 };
 
 export type AssistantRuntime = ThreadRuntimeWithSubscribe & {
+  switchToNewThread: () => void;
+  /**
+   * @deprecated Use `switchToNewThread` instead. This will be removed in 0.6.0.
+   */
   switchToThread: (threadId: string | null) => void;
   registerModelConfigProvider: (provider: ModelConfigProvider) => Unsubscribe;
 };
diff --git a/packages/react/src/runtimes/core/BaseAssistantRuntime.tsx b/packages/react/src/runtimes/core/BaseAssistantRuntime.tsx
index 81fb0fed2..f2124108d 100644
--- a/packages/react/src/runtimes/core/BaseAssistantRuntime.tsx
+++ b/packages/react/src/runtimes/core/BaseAssistantRuntime.tsx
@@ -20,6 +20,8 @@ export abstract class BaseAssistantRuntime<
     this.subscriptionHandler();
   }
 
+  public abstract switchToNewThread(): void;
+
   public abstract registerModelConfigProvider(
     provider: ModelConfigProvider,
   ): Unsubscribe;
diff --git a/packages/react/src/runtimes/external-store/ExternalStoreRuntime.tsx b/packages/react/src/runtimes/external-store/ExternalStoreRuntime.tsx
index ebcaf1c39..cc05fb5ef 100644
--- a/packages/react/src/runtimes/external-store/ExternalStoreRuntime.tsx
+++ b/packages/react/src/runtimes/external-store/ExternalStoreRuntime.tsx
@@ -26,8 +26,19 @@ export class ExternalStoreRuntime extends BaseAssistantRuntime<ExternalStoreThre
     return this._proxyConfigProvider.registerModelConfigProvider(provider);
   }
 
+  public async switchToNewThread() {
+    if (!this.store.onNewThread)
+      throw new Error("Runtime does not support switching to new threads.");
+
+    this.thread = new ExternalStoreThreadRuntime({
+      messages: [],
+      onNew: this.store.onNew,
+    });
+    await this.store.onNewThread();
+  }
+
   public async switchToThread(threadId: string | null) {
-    if (threadId) {
+    if (threadId !== null) {
       if (!this.store.onSwitchThread)
         throw new Error("Runtime does not support switching threads.");
 
@@ -37,14 +48,7 @@ export class ExternalStoreRuntime extends BaseAssistantRuntime<ExternalStoreThre
       });
       this.store.onSwitchThread(threadId);
     } else {
-      if (!this.store.onNewThread)
-        throw new Error("Runtime does not support switching to new threads.");
-
-      this.thread = new ExternalStoreThreadRuntime({
-        messages: [],
-        onNew: this.store.onNew,
-      });
-      await this.store.onNewThread();
+      this.switchToNewThread();
     }
   }
 }
diff --git a/packages/react/src/runtimes/local/LocalRuntime.tsx b/packages/react/src/runtimes/local/LocalRuntime.tsx
index f3cc3d736..87d37396d 100644
--- a/packages/react/src/runtimes/local/LocalRuntime.tsx
+++ b/packages/react/src/runtimes/local/LocalRuntime.tsx
@@ -28,11 +28,7 @@ export class LocalRuntime extends BaseAssistantRuntime<LocalThreadRuntime> {
     return this._proxyConfigProvider.registerModelConfigProvider(provider);
   }
 
-  public switchToThread(threadId: string | null) {
-    if (threadId) {
-      throw new Error("LocalRuntime does not yet support switching threads");
-    }
-
+  public switchToNewThread() {
     const { initialMessages, ...options } = this.thread.options;
 
     return (this.thread = new LocalThreadRuntime(
@@ -42,6 +38,14 @@ export class LocalRuntime extends BaseAssistantRuntime<LocalThreadRuntime> {
     ));
   }
 
+  public switchToThread(threadId: string | null) {
+    if (threadId !== null) {
+      throw new Error("LocalRuntime does not yet support switching threads");
+    }
+
+    this.switchToNewThread();
+  }
+
   public reset({
     initialMessages,
   }: {