feat: add warm-up indicator, new thread feature, and runtime improvements in chat UI

studio/frontend/src/components/assistant-ui/thread.tsx

@@ -7,7 +7,9 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { AnimatedShinyText } from "@/components/ui/animated-shiny-text";
 import { Button } from "@/components/ui/button";
+import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
 import { cn } from "@/lib/utils";
 import {
   ActionBarMorePrimitive,
@@ -20,6 +22,8 @@ import {
   SuggestionPrimitive,
   ThreadPrimitive,
   useAui,
+  useAuiEvent,
+  useAuiState,
 } from "@assistant-ui/react";
 import { motion } from "framer-motion";
 import {
@@ -36,7 +40,7 @@ import {
   RefreshCwIcon,
   SquareIcon,
 } from "lucide-react";
-import type { FC } from "react";
+import { type FC, useRef } from "react";
 
 export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
   hideComposer,
@@ -69,6 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
 
         <ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4 before:pointer-events-none before:absolute before:inset-x-0 before:bottom-full before:h-20 before:bg-gradient-to-t before:from-background before:to-transparent">
           <ThreadScrollToBottom />
+          <WarmupIndicator />
           <AuiIf condition={({ thread }) => !thread.isEmpty}>
             {!hideComposer && <ComposerAnimated />}
           </AuiIf>
@@ -78,6 +83,28 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
   );
 };
 
+const WarmupIndicator: FC = () => {
+  const threadId = useAuiState(({ threads }) => threads.mainThreadId);
+  const isRunning = useAuiState(({ thread }) => thread.isRunning);
+  const isWarmingUp = useChatRuntimeStore((state) =>
+    Boolean(state.warmingByThreadId[threadId ?? "__default"]),
+  );
+
+  if (!isRunning || !isWarmingUp) {
+    return null;
+  }
+
+  return (
+    <div className="mx-auto -mb-2 w-full max-w-(--thread-max-width) px-2">
+      <div className="inline-flex items-center rounded-full border border-border/60 bg-background/90 px-3 py-1.5 text-xs text-muted-foreground shadow-sm">
+        <AnimatedShinyText className="text-xs">
+          Warming up model...
+        </AnimatedShinyText>
+      </div>
+    </div>
+  );
+};
+
 const ThreadScrollToBottom: FC = () => {
   return (
     <ThreadPrimitive.ScrollToBottom asChild={true}>
@@ -358,6 +385,15 @@ const UserActionBar: FC = () => {
 
 const EditComposer: FC = () => {
   const aui = useAui();
+  const resendAfterCancelRef = useRef(false);
+
+  useAuiEvent("thread.runEnd", () => {
+    if (!resendAfterCancelRef.current) {
+      return;
+    }
+    resendAfterCancelRef.current = false;
+    aui.composer().send();
+  });
 
   return (
     <MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
@@ -384,7 +420,9 @@ const EditComposer: FC = () => {
               }
 
               if (aui.thread().getState().isRunning) {
+                resendAfterCancelRef.current = true;
                 aui.thread().cancelRun();
+                return;
               }
               aui.composer().send();
             }}

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -47,7 +47,7 @@ function toOpenAIMessage(message: RunMessage): {
 
 export function createOpenAIStreamAdapter(): ChatModelAdapter {
   return {
-    async *run({ messages, abortSignal }) {
+    async *run({ messages, abortSignal, unstable_threadId }) {
       const state = useChatRuntimeStore.getState();
       const { params } = state;
 
@@ -68,44 +68,58 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         });
       }
 
-      const stream = streamChatCompletions(
-        {
-          model: params.checkpoint,
-          messages: outboundMessages,
-          stream: true,
-          temperature: params.temperature,
-          top_p: params.topP,
-          max_tokens: params.maxTokens,
-          top_k: params.topK,
-          repetition_penalty: params.repetitionPenalty,
-        },
-        abortSignal,
-      );
-
+      const threadKey = unstable_threadId || "__default";
+      let waitingFirstChunk = true;
+      useChatRuntimeStore.getState().setThreadWarming(threadKey, true);
       let cumulativeText = "";
       let reasoningStartAt: number | null = null;
       let reasoningDuration = 0;
 
-      for await (const chunk of stream) {
-        const delta = chunk.choices?.[0]?.delta?.content;
-        if (!delta) {
-          continue;
-        }
-        cumulativeText += delta;
-        const parts = parseAssistantContent(cumulativeText);
+      try {
+        const stream = streamChatCompletions(
+          {
+            model: params.checkpoint,
+            messages: outboundMessages,
+            stream: true,
+            temperature: params.temperature,
+            top_p: params.topP,
+            max_tokens: params.maxTokens,
+            top_k: params.topK,
+            repetition_penalty: params.repetitionPenalty,
+          },
+          abortSignal,
+        );
 
-        if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) {
-          reasoningStartAt = Date.now();
-        }
-        if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) {
-          reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000);
-        }
+        for await (const chunk of stream) {
+          const delta = chunk.choices?.[0]?.delta?.content;
+          if (!delta) {
+            continue;
+          }
+          if (waitingFirstChunk) {
+            waitingFirstChunk = false;
+            useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
+          }
 
-        if (parts.length > 0) {
-          yield {
-            content: parts,
-            metadata: { custom: { reasoningDuration } },
-          };
+          cumulativeText += delta;
+          const parts = parseAssistantContent(cumulativeText);
+
+          if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) {
+            reasoningStartAt = Date.now();
+          }
+          if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) {
+            reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000);
+          }
+
+          if (parts.length > 0) {
+            yield {
+              content: parts,
+              metadata: { custom: { reasoningDuration } },
+            };
+          }
+        }
+      } finally {
+        if (waitingFirstChunk) {
+          useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
         }
       }
     },

studio/frontend/src/features/chat/chat-page.tsx

@@ -49,9 +49,14 @@ import type { ChatView } from "./types";
 
 const SingleContent = memo(function SingleContent({
   threadId,
-}: { threadId?: string }): ReactElement {
+  newThreadNonce,
+}: { threadId?: string; newThreadNonce?: string }): ReactElement {
   return (
-    <ChatRuntimeProvider modelType="base" initialThreadId={threadId}>
+    <ChatRuntimeProvider
+      modelType="base"
+      initialThreadId={threadId}
+      newThreadNonce={newThreadNonce}
+    >
       <div className="min-h-0 flex-1">
         <Thread />
       </div>
@@ -197,7 +202,10 @@ function TopBarActions({
 }
 
 export function ChatPage(): ReactElement {
-  const [view, setView] = useState<ChatView>({ mode: "single" });
+  const [view, setView] = useState<ChatView>({
+    mode: "single",
+    newThreadNonce: crypto.randomUUID(),
+  });
   const [settingsOpen, setSettingsOpen] = useState(false);
   const inferenceParams = useChatRuntimeStore((state) => state.params);
   const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
@@ -215,7 +223,10 @@ export function ChatPage(): ReactElement {
   const handleEject = useCallback(() => {
     void ejectModel();
   }, [ejectModel]);
-  const handleNewThread = useCallback(() => setView({ mode: "single" }), []);
+  const handleNewThread = useCallback(
+    () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }),
+    [],
+  );
   const handleNewCompare = useCallback(
     () => setView({ mode: "compare", pairId: crypto.randomUUID() }),
     [],
@@ -300,7 +311,11 @@ export function ChatPage(): ReactElement {
         </div>
 
         {view.mode === "single" ? (
-          <SingleContent key={view.threadId ?? "new"} threadId={view.threadId} />
+          <SingleContent
+            key={view.threadId ?? view.newThreadNonce ?? "new"}
+            threadId={view.threadId}
+            newThreadNonce={view.newThreadNonce}
+          />
         ) : (
           <CompareContent key={view.pairId} pairId={view.pairId} />
         )}

studio/frontend/src/features/chat/runtime-provider.tsx

@@ -7,6 +7,7 @@ import {
   type ExportedMessageRepositoryItem,
   type PendingAttachment,
   RuntimeAdapterProvider,
+  Suggestions,
   SimpleImageAttachmentAdapter,
   SimpleTextAttachmentAdapter,
   type ThreadHistoryAdapter,
@@ -296,7 +297,14 @@ function useRuntimeHook(): ReturnType<typeof useLocalRuntime> {
 function ThreadAutoSwitch({
   threadId,
 }: { threadId: string }): ReactElement | null {
-  const aui = useAui();
+  const aui = useAui({
+    suggestions: Suggestions([
+      "Draw a simple flowchart of a login system using Mermaid",
+      "Solve the integral of x²·sin(x) step by step",
+      "Write a Python function that finds the longest palindrome in a string",
+      "Format a comparison of 3 databases as a markdown table with pros and cons",
+    ]),
+  });
   const isLoading = useAuiState(({ threads }) => threads.isLoading);
   const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
 
@@ -309,16 +317,33 @@ function ThreadAutoSwitch({
   return null;
 }
 
+function ThreadNewChatSwitch({
+  nonce,
+}: { nonce: string }): ReactElement | null {
+  const aui = useAui();
+  const isLoading = useAuiState(({ threads }) => threads.isLoading);
+
+  useEffect(() => {
+    if (!isLoading) {
+      aui.threads().switchToNewThread();
+    }
+  }, [aui, isLoading, nonce]);
+
+  return null;
+}
+
 export function ChatRuntimeProvider({
   children,
   modelType = "base",
   pairId,
   initialThreadId,
+  newThreadNonce,
 }: {
   children: ReactNode;
   modelType?: ModelType;
   pairId?: string;
   initialThreadId?: string;
+  newThreadNonce?: string;
 }): ReactElement {
   const runtime = useRemoteThreadListRuntime({
     runtimeHook: useRuntimeHook,
@@ -333,6 +358,9 @@ export function ChatRuntimeProvider({
   return (
     <AssistantRuntimeProvider runtime={runtime} aui={aui}>
       {initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
+      {!initialThreadId && newThreadNonce && (
+        <ThreadNewChatSwitch nonce={newThreadNonce} />
+      )}
       {children}
     </AssistantRuntimeProvider>
   );

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -10,10 +10,12 @@ type ChatRuntimeStore = {
   params: InferenceParams;
   models: ChatModelSummary[];
   loras: ChatLoraSummary[];
+  warmingByThreadId: Record<string, boolean>;
   modelsError: string | null;
   setParams: (params: InferenceParams) => void;
   setModels: (models: ChatModelSummary[]) => void;
   setLoras: (loras: ChatLoraSummary[]) => void;
+  setThreadWarming: (threadId: string, warming: boolean) => void;
   setModelsError: (error: string | null) => void;
   setCheckpoint: (modelId: string) => void;
   clearCheckpoint: () => void;
@@ -23,10 +25,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   params: DEFAULT_INFERENCE_PARAMS,
   models: [],
   loras: [],
+  warmingByThreadId: {},
   modelsError: null,
   setParams: (params) => set({ params }),
   setModels: (models) => set({ models }),
   setLoras: (loras) => set({ loras }),
+  setThreadWarming: (threadId, warming) =>
+    set((state) => {
+      const next = { ...state.warmingByThreadId };
+      if (warming) {
+        next[threadId] = true;
+      } else {
+        delete next[threadId];
+      }
+      return { warmingByThreadId: next };
+    }),
   setModelsError: (modelsError) => set({ modelsError }),
   setCheckpoint: (modelId) =>
     set((state) => ({

studio/frontend/src/features/chat/types.ts

@@ -1,7 +1,7 @@
 export type ModelType = "base" | "lora";
 
 export type ChatView =
-  | { mode: "single"; threadId?: string }
+  | { mode: "single"; threadId?: string; newThreadNonce?: string }
   | { mode: "compare"; pairId: string };
 
 export interface ThreadRecord {