From 76830db0cf16bd3d3c54700c69747398be097707 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 17:28:01 +0100 Subject: [PATCH] feat: add warm-up indicator, new thread feature, and runtime improvements in chat UI --- .../src/components/assistant-ui/thread.tsx | 40 +++++++++- .../src/features/chat/api/chat-adapter.ts | 80 +++++++++++-------- .../frontend/src/features/chat/chat-page.tsx | 25 ++++-- .../src/features/chat/runtime-provider.tsx | 30 ++++++- .../chat/stores/chat-runtime-store.ts | 13 +++ studio/frontend/src/features/chat/types.ts | 2 +- 6 files changed, 149 insertions(+), 41 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 3bcaaea105..df77d3ea28 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/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 }> = ({ + !thread.isEmpty}> {!hideComposer && } @@ -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 ( +
+
+ + Warming up model... + +
+
+ ); +}; + const ThreadScrollToBottom: FC = () => { return ( @@ -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 ( @@ -384,7 +420,9 @@ const EditComposer: FC = () => { } if (aui.thread().getState().isRunning) { + resendAfterCancelRef.current = true; aui.thread().cancelRun(); + return; } aui.composer().send(); }} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c7bc827e3..fa670ecda2 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/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); } } }, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c787cf471c..543ae9e865 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/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 ( - +
@@ -197,7 +202,10 @@ function TopBarActions({ } export function ChatPage(): ReactElement { - const [view, setView] = useState({ mode: "single" }); + const [view, setView] = useState({ + 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 { {view.mode === "single" ? ( - + ) : ( )} diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index dade143141..bf1b658c3d 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/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 { 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 ( {initialThreadId && } + {!initialThreadId && newThreadNonce && ( + + )} {children} ); diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 76bf079ce7..74efc57134 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -10,10 +10,12 @@ type ChatRuntimeStore = { params: InferenceParams; models: ChatModelSummary[]; loras: ChatLoraSummary[]; + warmingByThreadId: Record; 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((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) => ({ diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 45d4f9e195..b0dccab307 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/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 {