fix: sync RuntimeSynchronizer (#342)

Yonom · Jun 27, 2024 · 1131680 · 1131680
1 parent 7be2b6c
commit 1131680
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/packages/react/src/context/providers/ThreadProvider.tsx b/packages/react/src/context/providers/ThreadProvider.tsx
@@ -1,5 +1,12 @@
 import type { FC, PropsWithChildren } from "react";
-import { useEffect, useInsertionEffect, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useRef,
+  useState,
+  useSyncExternalStore,
+} from "react";
 import type { ReactThreadRuntime } from "../../runtime/core/ReactThreadRuntime";
 import type { ThreadRuntime } from "../../runtime/core/ThreadRuntime";
 import type { ThreadContextValue } from "../react/ThreadContext";
@@ -52,8 +59,16 @@ export const ThreadProvider: FC<PropsWithChildren<ThreadProviderProps>> = ({
     return runtime.subscribe(onRuntimeUpdate);
   }, [context, runtime]);
 
-  const RuntimeSynchronizer = (runtime as ReactThreadRuntime)
-    .unstable_synchronizer;
+  const subscribe = useCallback(
+    (c: () => void) => runtime.subscribe(c),
+    [runtime],
+  );
+
+  const RuntimeSynchronizer = useSyncExternalStore(
+    subscribe,
+    () => (runtime as ReactThreadRuntime).unstable_synchronizer,
+    () => undefined,
+  );
 
   return (
     <ThreadContext.Provider value={context}>

diff --git a/packages/react/src/runtime/core/BaseAssistantRuntime.tsx b/packages/react/src/runtime/core/BaseAssistantRuntime.tsx
@@ -23,6 +23,7 @@ export abstract class BaseAssistantRuntime<
     this._unsubscribe();
     this._thread = thread;
     this._unsubscribe = this._thread.subscribe(this.subscriptionHandler);
+    this.subscriptionHandler();
   }
 
   public abstract registerModelConfigProvider(