feat: ThreadRuntime Events API (experimental) (#988)

replace ComposerRuntime.focus with the new Events API

apps/docs/components/shadcn/Shadcn.tsx

@@ -1,6 +1,6 @@
 import { ArchiveIcon, EditIcon, MenuIcon, ShareIcon } from "lucide-react";
 import Link from "next/link";
-import { Thread, useSwitchToNewThread } from "@assistant-ui/react";
+import { Thread, useAssistantRuntime } from "@assistant-ui/react";
 import { makeMarkdownText } from "@assistant-ui/react-markdown";
 import remarkGfm from "remark-gfm";
 import { makePrismAsyncSyntaxHighlighter } from "@assistant-ui/react-syntax-highlighter";
@@ -61,11 +61,11 @@ const ButtonWithTooltip: FC<ButtonWithTooltipProps> = ({
 };
 
 const TopLeft: FC = () => {
-  const switchToNewThread = useSwitchToNewThread();
+  const runtime = useAssistantRuntime();
 
   return (
     <ButtonWithTooltip
-      onClick={switchToNewThread}
+      onClick={() => runtime.switchToNewThread()}
       variant="ghost"
       className="flex w-full justify-between px-3"
       tooltip="New Chat"

packages/react-playground/src/lib/playground-runtime.ts

@@ -292,6 +292,11 @@ export class PlaygroundThreadRuntimeCore implements INTERNAL.ThreadRuntimeCore {
     throw new Error("PlaygroundRuntime does not support feedback.");
   }
 
+  public unstable_on() {
+    // events not supported in playground
+    return () => {};
+  }
+
   public deleteMessage(messageId: string) {
     this.setMessages(this.messages.filter((m) => m.id !== messageId));
   }

packages/react/src/api/ComposerRuntime.ts

@@ -87,12 +87,6 @@ type LegacyThreadComposerState = Readonly<{
   send: () => void;
   /** @deprecated Use `useComposerRuntime().cancel` instead. This will be removed in 0.6.0. */
   cancel: () => void;
-
-  // TODO replace with events
-  /** @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it. */
-  focus: () => void;
-  /** @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it. */
-  onFocus: (listener: () => void) => Unsubscribe;
 }>;
 
 type BaseComposerState = {
@@ -124,8 +118,6 @@ const METHOD_NOT_SUPPORTED = () => {
 const EMPTY_ARRAY = Object.freeze([]);
 const getThreadComposerState = (
   runtime: ThreadComposerRuntimeCore | undefined,
-  focus: () => void,
-  onFocus: (listener: () => void) => Unsubscribe,
 ): ThreadComposerState => {
   return Object.freeze({
     type: "thread",
@@ -142,8 +134,6 @@ const getThreadComposerState = (
     // edit: beginEdit,
     send: runtime?.send.bind(runtime) ?? METHOD_NOT_SUPPORTED,
     cancel: runtime?.cancel.bind(runtime) ?? METHOD_NOT_SUPPORTED,
-    focus: focus,
-    onFocus: onFocus,
     reset: runtime?.reset.bind(runtime) ?? METHOD_NOT_SUPPORTED,
 
     addAttachment: runtime?.addAttachment.bind(runtime) ?? METHOD_NOT_SUPPORTED,
@@ -336,12 +326,6 @@ export type ThreadComposerRuntime = Omit<
    */
   attachments: readonly PendingAttachment[];
 
-  /** @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it. */
-  focus(): void;
-
-  /** @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it. */
-  onFocus(callback: () => void): Unsubscribe;
-
   getAttachmentByIndex(
     idx: number,
   ): AttachmentRuntime & { source: "thread-composer" };
@@ -359,12 +343,7 @@ export class ThreadComposerRuntimeImpl
 
   constructor(core: ThreadComposerRuntimeCoreBinding) {
     const stateBinding = new LazyMemoizeSubject({
-      getState: () =>
-        getThreadComposerState(
-          core.getState(),
-          this.focus.bind(this),
-          this.onFocus.bind(this),
-        ),
+      getState: () => getThreadComposerState(core.getState()),
       subscribe: (callback) => core.subscribe(callback),
     });
     super({
@@ -382,24 +361,6 @@ export class ThreadComposerRuntimeImpl
     return this._getState();
   }
 
-  // TODO replace with events
-  private _focusListeners = new Set<() => void>();
-
-  /**
-   * @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it.
-   */
-  public focus() {
-    this._focusListeners.forEach((callback) => callback());
-  }
-
-  /**
-   * @deprecated This feature is being removed in 0.6.0. Submit feedback if you need it.
-   */
-  public onFocus(callback: () => void) {
-    this._focusListeners.add(callback);
-    return () => this._focusListeners.delete(callback);
-  }
-
   public getAttachmentByIndex(idx: number) {
     return new ThreadComposerAttachmentRuntimeImpl(
       new ShallowMemoizeSubject({

packages/react/src/api/ThreadRuntime.ts

@@ -21,7 +21,10 @@ import {
 } from "./MessageRuntime";
 import { NestedSubscriptionSubject } from "./subscribable/NestedSubscriptionSubject";
 import { ShallowMemoizeSubject } from "./subscribable/ShallowMemoizeSubject";
-import { SubscribableWithState } from "./subscribable/Subscribable";
+import {
+  Subscribable,
+  SubscribableWithState,
+} from "./subscribable/Subscribable";
 import {
   ThreadComposerRuntime,
   ThreadComposerRuntimeImpl,
@@ -64,7 +67,10 @@ const toAppendMessage = (
   } as AppendMessage;
 };
 
-export type ThreadRuntimeCoreBinding = SubscribableWithState<ThreadRuntimeCore>;
+export type ThreadRuntimeCoreBinding =
+  SubscribableWithState<ThreadRuntimeCore> & {
+    outerSubscribe(callback: () => void): Unsubscribe;
+  };
 
 export type ThreadState = Readonly<{
   threadId: string;
@@ -112,6 +118,10 @@ export type ThreadRuntime = {
   import(repository: ExportedMessageRepository): void;
   getMesssageByIndex(idx: number): MessageRuntime;
   stopSpeaking: () => void;
+  unstable_on(
+    event: "switched-to" | "run-start",
+    callback: () => void,
+  ): Unsubscribe;
 
   // Legacy methods with deprecations
 
@@ -271,6 +281,7 @@ export class ThreadRuntimeImpl implements ThreadRuntimeCore, ThreadRuntime {
     this._threadBinding = {
       getState: () => threadBinding.getState(),
       getStateState: () => stateBinding.getState(),
+      outerSubscribe: (callback) => threadBinding.outerSubscribe(callback),
       subscribe: (callback) => threadBinding.subscribe(callback),
     };
   }
@@ -413,4 +424,27 @@ export class ThreadRuntimeImpl implements ThreadRuntimeCore, ThreadRuntime {
       this._threadBinding,
     );
   }
+
+  private _eventListenerNestedSubscriptions = new Map<
+    string,
+    NestedSubscriptionSubject<Subscribable>
+  >();
+
+  public unstable_on(
+    event: "switched-to" | "run-start",
+    callback: () => void,
+  ): Unsubscribe {
+    let subject = this._eventListenerNestedSubscriptions.get(event);
+    if (!subject) {
+      subject = new NestedSubscriptionSubject({
+        getState: () => ({
+          subscribe: (callback) =>
+            this._threadBinding.getState().unstable_on(event, callback),
+        }),
+        subscribe: (callback) => this._threadBinding.outerSubscribe(callback),
+      });
+      this._eventListenerNestedSubscriptions.set(event, subject);
+    }
+    return subject.subscribe(callback);
+  }
 }

packages/react/src/api/subscribable/NestedSubscriptionSubject.ts

@@ -18,6 +18,10 @@ export class NestedSubscriptionSubject<TState extends Subscribable | undefined>
     return this.binding.getState();
   }
 
+  public outerSubscribe(callback: () => void) {
+    return this.binding.subscribe(callback);
+  }
+
   protected _connect(): Unsubscribe {
     const callback = () => {
       this.notifySubscribers();
@@ -36,7 +40,7 @@ export class NestedSubscriptionSubject<TState extends Subscribable | undefined>
       callback();
     };
 
-    const outerUnsubscribe = this.binding.subscribe(onRuntimeUpdate);
+    const outerUnsubscribe = this.outerSubscribe(onRuntimeUpdate);
     return () => {
       outerUnsubscribe?.();
       innerUnsubscribe?.();

packages/react/src/hooks/useAppendMessage.tsx

@@ -1,27 +1,18 @@
 import { useCallback } from "react";
-import { useThreadViewportStore } from "../context";
-import {
-  useThreadComposerStore,
-  useThreadRuntime,
-} from "../context/react/ThreadContext";
+import { useThreadRuntime } from "../context/react/ThreadContext";
 import { CreateAppendMessage } from "../api/ThreadRuntime";
 
 /**
  * @deprecated Use `useThreadRuntime().append()` instead. This will be removed in 0.6.
  */
 export const useAppendMessage = () => {
   const threadRuntime = useThreadRuntime();
-  const threadViewportStore = useThreadViewportStore();
-  const threadComposerStore = useThreadComposerStore();
 
   const append = useCallback(
     (message: CreateAppendMessage) => {
       threadRuntime.append(message);
-
-      threadViewportStore.getState().scrollToBottom();
-      threadComposerStore.getState().focus();
     },
-    [threadRuntime, threadViewportStore, threadComposerStore],
+    [threadRuntime],
   );
 
   return append;

packages/react/src/hooks/useSwitchToNewThread.tsx

@@ -1,17 +1,14 @@
 import { useCallback } from "react";
-import { useThreadComposerStore } from "../context/react/ThreadContext";
 import { useAssistantRuntime } from "../context";
 
 /**
  * @deprecated Use `useRuntimeActions().switchToNewThread()` instead. This will be removed in 0.6.0.
  */
 export const useSwitchToNewThread = () => {
   const assistantRuntime = useAssistantRuntime();
-  const threadComposerStore = useThreadComposerStore();
   const switchToNewThread = useCallback(() => {
     assistantRuntime.switchToNewThread();
-    threadComposerStore.getState().focus();
-  }, [assistantRuntime, threadComposerStore]);
+  }, [assistantRuntime]);
 
   return switchToNewThread;
 };

packages/react/src/primitive-hooks/actionBar/useActionBarReload.tsx

@@ -1,32 +1,20 @@
 import { useCallback } from "react";
-import {
-  useMessageRuntime,
-  useMessageStore,
-} from "../../context/react/MessageContext";
-import {
-  useThreadComposerStore,
-  useThreadStore,
-  useThreadViewportStore,
-} from "../../context/react/ThreadContext";
+import { useMessageRuntime } from "../../context/react/MessageContext";
+import { useThreadRuntime } from "../../context/react/ThreadContext";
 import { useCombinedStore } from "../../utils/combined/useCombinedStore";
 
 export const useActionBarReload = () => {
-  const messageStore = useMessageStore();
-  const threadStore = useThreadStore();
   const messageRuntime = useMessageRuntime();
-  const threadComposerStore = useThreadComposerStore();
-  const threadViewportStore = useThreadViewportStore();
+  const threadRuntime = useThreadRuntime();
 
   const disabled = useCombinedStore(
-    [threadStore, messageStore],
+    [threadRuntime, messageRuntime],
     (t, m) => t.isRunning || t.isDisabled || m.role !== "assistant",
   );
 
   const callback = useCallback(() => {
     messageRuntime.reload();
-    threadViewportStore.getState().scrollToBottom();
-    threadComposerStore.getState().focus();
-  }, [messageRuntime, threadComposerStore, threadViewportStore]);
+  }, [messageRuntime]);
 
   if (disabled) return null;
   return callback;

packages/react/src/primitive-hooks/composer/useComposerSend.tsx

@@ -1,32 +1,22 @@
 import { useCallback } from "react";
-import { useComposerStore } from "../../context";
 import { useCombinedStore } from "../../utils/combined/useCombinedStore";
-import {
-  useThreadComposerStore,
-  useThreadStore,
-  useThreadViewportStore,
-} from "../../context/react/ThreadContext";
+import { useThreadRuntime } from "../../context/react/ThreadContext";
+import { useComposerRuntime } from "../../context";
 
 export const useComposerSend = () => {
-  const threadStore = useThreadStore();
-  const threadViewportStore = useThreadViewportStore();
-  const composerStore = useComposerStore();
-  const threadComposerStore = useThreadComposerStore();
+  const composerRuntime = useComposerRuntime();
+  const threadRuntime = useThreadRuntime();
 
   const disabled = useCombinedStore(
-    [threadStore, composerStore],
+    [threadRuntime, composerRuntime],
     (t, c) => t.isRunning || !c.isEditing || c.isEmpty,
   );
 
   const callback = useCallback(() => {
-    const composerState = composerStore.getState();
-    if (!composerState.isEditing) return;
+    if (!composerRuntime.getState().isEditing) return;
 
-    composerState.send();
-
-    threadViewportStore.getState().scrollToBottom();
-    threadComposerStore.getState().focus();
-  }, [threadComposerStore, composerStore, threadViewportStore]);
+    composerRuntime.send();
+  }, [threadRuntime]);
 
   if (disabled) return null;
   return callback;

packages/react/src/primitive-hooks/thread/useThreadScrollToBottom.tsx

@@ -1,20 +1,15 @@
 import { useCallback } from "react";
 import { useThreadViewport } from "../../context";
-import {
-  useThreadComposerStore,
-  useThreadViewportStore,
-} from "../../context/react/ThreadContext";
+import { useThreadViewportStore } from "../../context/react/ThreadContext";
 
 export const useThreadScrollToBottom = () => {
   const isAtBottom = useThreadViewport((s) => s.isAtBottom);
 
   const threadViewportStore = useThreadViewportStore();
-  const threadComposerStore = useThreadComposerStore();
 
   const handleScrollToBottom = useCallback(() => {
     threadViewportStore.getState().scrollToBottom();
-    threadComposerStore.getState().focus();
-  }, [threadViewportStore, threadComposerStore]);
+  }, [threadViewportStore]);
 
   if (isAtBottom) return null;
   return handleScrollToBottom;

packages/react/src/primitive-hooks/thread/useThreadViewportAutoScroll.tsx

@@ -1,18 +1,23 @@
 "use client";
 import { useComposedRefs } from "@radix-ui/react-compose-refs";
-import { useRef } from "react";
-import { useThreadViewportStore } from "../../context/react/ThreadContext";
+import { useEffect, useRef } from "react";
+import {
+  useThreadRuntime,
+  useThreadViewportStore,
+} from "../../context/react/ThreadContext";
 import { useOnResizeContent } from "../../utils/hooks/useOnResizeContent";
 import { useOnScrollToBottom } from "../../utils/hooks/useOnScrollToBottom";
 import { useManagedRef } from "../../utils/hooks/useManagedRef";
 import { writableStore } from "../../context/ReadonlyStore";
 
 export type UseThreadViewportAutoScrollProps = {
   autoScroll?: boolean | undefined;
+  unstable_scrollToBottomOnRunStart?: boolean | undefined;
 };
 
 export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
   autoScroll = true,
+  unstable_scrollToBottomOnRunStart = true,
 }: UseThreadViewportAutoScrollProps) => {
   const divRef = useRef<TElement>(null);
 
@@ -81,5 +86,13 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
     scrollToBottom("auto");
   });
 
+  // autoscroll on run start
+  const threadRuntime = useThreadRuntime();
+  useEffect(() => {
+    if (!unstable_scrollToBottomOnRunStart) return undefined;
+
+    return threadRuntime.unstable_on("run-start", focus);
+  }, [unstable_scrollToBottomOnRunStart]);
+
   return autoScrollRef;
 };