diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 400fcb6c81..3bcaaea105 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -19,9 +19,7 @@ import {
   MessagePrimitive,
   SuggestionPrimitive,
   ThreadPrimitive,
-  useComposerRuntime,
-  useMessageRuntime,
-  useThreadRuntime,
+  useAui,
 } from "@assistant-ui/react";
 import { motion } from "framer-motion";
 import {
@@ -359,9 +357,7 @@ const UserActionBar: FC = () => {
 };
 
 const EditComposer: FC = () => {
-  const threadRuntime = useThreadRuntime();
-  const composerRuntime = useComposerRuntime();
-  const messageRuntime = useMessageRuntime();
+  const aui = useAui();
 
   return (
     <MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
@@ -379,18 +375,18 @@ const EditComposer: FC = () => {
           <Button
             size="sm"
             onClick={() => {
-              const newText = composerRuntime.getState().text;
-              const originalText = messageRuntime.unstable_getCopyText();
+              const newText = aui.composer().getState().text;
+              const originalText = aui.message().getCopyText();
 
               if (newText === originalText) {
-                composerRuntime.cancel();
+                aui.composer().cancel();
                 return;
               }
 
-              if (threadRuntime.getState().isRunning) {
-                threadRuntime.cancelRun();
+              if (aui.thread().getState().isRunning) {
+                aui.thread().cancelRun();
               }
-              composerRuntime.send();
+              aui.composer().send();
             }}
           >
             Update
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index 5858de1cac..8e4d242269 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -12,8 +12,8 @@ import {
   type ThreadUserMessagePart,
   WebSpeechDictationAdapter,
   type unstable_RemoteThreadListAdapter,
-  useAssistantRuntime,
   useAui,
+  useAuiState,
   useLocalRuntime,
   unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime,
 } from "@assistant-ui/react";
@@ -222,28 +222,16 @@ const useRuntimeHook = (): ReturnType<typeof useLocalRuntime> =>
 function ThreadAutoSwitch({
   threadId,
 }: { threadId: string }): ReactElement | null {
-  const runtime = useAssistantRuntime();
+  const aui = useAui();
+  const isLoading = useAuiState(({ threads }) => threads.isLoading);
+  const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
+
   useEffect(() => {
-    function trySwitch() {
-      const s = runtime.threadList.getState();
-      if (s.isLoading || s.mainThreadId === threadId) {
-        return false;
-      }
-      runtime.threadList.switchToThread(threadId);
-      return true;
+    if (!isLoading && mainThreadId !== threadId) {
+      aui.threads().switchToThread(threadId);
     }
+  }, [aui, isLoading, mainThreadId, threadId]);
 
-    if (trySwitch()) {
-      return;
-    }
-
-    const unsub = runtime.threadList.subscribe(() => {
-      if (trySwitch()) {
-        unsub();
-      }
-    });
-    return unsub;
-  }, [runtime, threadId]);
   return null;
 }
 
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index f990d42b5a..185a6ac823 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -1,6 +1,6 @@
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { Button } from "@/components/ui/button";
-import { useAssistantRuntime } from "@assistant-ui/react";
+import { useAui } from "@assistant-ui/react";
 import { ArrowUpIcon, SquareIcon } from "lucide-react";
 import {
   type KeyboardEvent,
@@ -44,7 +44,7 @@ export function RegisterCompareHandle({
   name: string;
 }): ReactElement | null {
   const handlesRef = useContext(CompareHandlesContext);
-  const runtime = useAssistantRuntime();
+  const aui = useAui();
 
   useEffect(() => {
     if (!handlesRef) {
@@ -52,14 +52,14 @@ export function RegisterCompareHandle({
     }
     const currentHandles = handlesRef.current;
     currentHandles[name] = {
-      append: (content) => runtime.thread.append({ role: "user", content }),
-      cancel: () => runtime.thread.cancelRun(),
-      isRunning: () => runtime.thread.getState().isRunning,
+      append: (content) => aui.thread().append({ role: "user", content }),
+      cancel: () => aui.thread().cancelRun(),
+      isRunning: () => aui.thread().getState().isRunning,
     };
     return () => {
       delete currentHandles[name];
     };
-  }, [handlesRef, name, runtime]);
+  }, [handlesRef, name, aui]);
 
   return null;
 }