From b53cc957dcc43bb105b0fde3ac42c92458eeaba5 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Mon, 15 Jun 2026 16:03:57 +0100 Subject: [PATCH] Studio: queue chat prompts while generation is running (#6244) * feat: queue chat prompts during generation * fix: address prompt queue review edge cases * fix: harden queued prompt dispatch * fix: track queued prompt run state by thread * fix: preserve prompt queue ordering * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: isolate prompt queue on chat switches * fix: harden prompt queue cancellation * fix: address prompt queue review followups --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../src/components/assistant-ui/thread.tsx | 785 ++++++++++++++---- .../src/features/chat/runtime-provider.tsx | 21 + .../chat/utils/prompt-queue-boundary.ts | 8 + 3 files changed, 674 insertions(+), 140 deletions(-) create mode 100644 studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index ce1beb24d7..ee3754bdf6 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -75,6 +75,7 @@ import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled" import { BypassPermissionsMenuItem } from "@/features/chat/bypass-permissions-menu-item"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; +import { PROMPT_QUEUE_STOP_EVENT } from "@/features/chat/utils/prompt-queue-boundary"; import { PLUS_MENU_ORDER, composerDraftKey, @@ -84,6 +85,7 @@ import { writeComposerDraft, } from "@/features/chat"; import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; +import { listThreadDocuments } from "@/features/rag/api/rag-api"; import { ThreadDocumentsBar } from "@/features/rag/components/thread-documents-bar"; import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button"; import { DocumentPreviewMount } from "@/features/rag/components/document-preview-mount"; @@ -158,6 +160,7 @@ import { useRef, useState, } from "react"; +import { create } from "zustand"; // True while a file is dragged anywhere over the chat page, so the composer // can show its "Drop files here" affordance. @@ -165,79 +168,421 @@ const PageDragContext = createContext(false); // Single-chat prompt queue. State lives at module level so it survives the // Composer remount when the first queued message creates a new thread, and -// detection subscribes to the store's runningByThreadId rather than -// aui.thread() (unbound on the welcome screen). - -import { create as _createZustand } from "zustand"; - -// Module-level Zustand so ComposerRightControls re-renders across Composer mounts. -interface _QueueUIState { isRunning: boolean; current: number; total: number; } -const _useQueueUI = _createZustand<_QueueUIState>(() => ({ - isRunning: false, current: 0, total: 0, -})); - -let _qItems: string[] = []; -let _qIndex = 0; -let _qIsRunning = false; -let _qPrevStoreRunning = false; -let _qStoreUnsub: (() => void) | null = null; - -// Points to the current Composer's aui (updated every render), so it stays valid -// after a remount. -let _qGetAui: () => ReturnType = () => { - throw new Error("aui not initialised"); +// detection subscribes to runningByThreadId instead of aui.thread() so the +// welcome-screen composer can queue safely before a thread is bound. +type PromptQueueUIEntry = { + current: number; + total: number; }; -function _qStopSubscription() { - if (_qStoreUnsub) { _qStoreUnsub(); _qStoreUnsub = null; } - _qPrevStoreRunning = false; +interface PromptQueueUIState { + byThreadId: Record; } -function _qAdvance() { - const nextIndex = _qIndex + 1; - if (nextIndex >= _qItems.length) { - _qIsRunning = false; - _qItems = []; - _qIndex = 0; - _qStopSubscription(); - _useQueueUI.setState({ isRunning: false, current: 0, total: 0 }); +const usePromptQueueUI = create(() => ({ + byThreadId: {}, +})); + +type PromptQueueTarget = { + getDocumentThreadId: () => string | null; + getRunningThreadIds: () => string[]; + append: (prompt: string) => void; + cancel: () => void; + isIndexing: () => boolean; +}; + +type PromptQueueItem = { + prompt: string; + target: PromptQueueTarget; +}; + +const PROMPT_QUEUE_INDEXING_RETRY_MS = 500; + +let promptQueueItems: PromptQueueItem[] = []; +let promptQueueIndex = 0; +let promptQueueIsRunning = false; +let promptQueueGeneration = 0; +let promptQueuePrevStoreRunning = false; +let promptQueueWaitingForTargetIdle = false; +let promptQueueStoreUnsub: (() => void) | null = null; +let promptQueueRetryTimer: ReturnType | null = null; + +function compactIds(ids: Array) { + return Array.from(new Set(ids.filter((id): id is string => Boolean(id)))); +} + +function stopPromptQueueSubscription({ + resetRunningState = true, +}: { + resetRunningState?: boolean; +} = {}) { + if (promptQueueStoreUnsub) { + promptQueueStoreUnsub(); + promptQueueStoreUnsub = null; + } + if (resetRunningState) { + promptQueuePrevStoreRunning = false; + } +} + +function resetPromptQueue(showToast = false) { + promptQueueGeneration += 1; + promptQueueIsRunning = false; + promptQueueItems = []; + promptQueueIndex = 0; + promptQueueWaitingForTargetIdle = false; + if (promptQueueRetryTimer) { + clearTimeout(promptQueueRetryTimer); + promptQueueRetryTimer = null; + } + stopPromptQueueSubscription(); + syncPromptQueueUI(); + if (showToast) { toast.success("Prompt queue complete"); + } +} + +function queueToastDescription(prompt: string) { + return prompt.length > 80 ? `${prompt.slice(0, 80)}...` : prompt; +} + +function appendQueuedPrompt(item: PromptQueueItem) { + syncPromptQueueUI(); + item.target.append(item.prompt); +} + +async function targetHasIndexingDocuments(item: PromptQueueItem) { + if (item.target.isIndexing()) { + return true; + } + const state = useChatRuntimeStore.getState(); + if ( + !state.ragEnabled || + state.ragSource.type !== "thread" + ) { + return false; + } + const threadId = item.target.getDocumentThreadId(); + if (!threadId) { + return false; + } + try { + const documents = await listThreadDocuments(threadId); + return documents.some( + (doc) => doc.status === "pending" || doc.status === "running", + ); + } catch { + return item.target.isIndexing(); + } +} + +function isActivePromptQueueItem(item: PromptQueueItem, generation: number) { + if (!promptQueueIsRunning || generation !== promptQueueGeneration) { + return false; + } + return promptQueueItems[Math.max(promptQueueIndex, 0)] === item; +} + +function scheduleQueuedPromptDispatch( + item: PromptQueueItem, + delay: number, + generation = promptQueueGeneration, +) { + if (promptQueueRetryTimer) { + clearTimeout(promptQueueRetryTimer); + } + promptQueueRetryTimer = setTimeout(() => { + promptQueueRetryTimer = null; + void dispatchQueuedPrompt(item, generation); + }, delay); +} + +async function dispatchQueuedPrompt( + item: PromptQueueItem, + generation = promptQueueGeneration, +) { + if (!isActivePromptQueueItem(item, generation)) { return; } - _qIndex = nextIndex; - _useQueueUI.setState({ current: nextIndex + 1, total: _qItems.length }); - const next = _qItems[nextIndex]; - toast(`Prompt ${nextIndex + 1} / ${_qItems.length}`, { - description: next.length > 80 ? next.slice(0, 80) + "…" : next, - }); - _qPrevStoreRunning = false; // catch the next run - setTimeout(() => { - _qGetAui().thread().append({ - role: "user", - content: [{ type: "text", text: next }], - createdAt: new Date(), - } as never); - }, 100); + if ( + isPromptQueueTargetRunning( + item.target, + useChatRuntimeStore.getState().runningByThreadId, + ) + ) { + promptQueueWaitingForTargetIdle = true; + promptQueuePrevStoreRunning = true; + syncPromptQueueUI(); + startPromptQueueSubscription(); + return; + } + const hasIndexingDocuments = await targetHasIndexingDocuments(item); + if (!isActivePromptQueueItem(item, generation)) { + return; + } + if (hasIndexingDocuments) { + scheduleQueuedPromptDispatch(item, PROMPT_QUEUE_INDEXING_RETRY_MS); + return; + } + if (!isActivePromptQueueItem(item, generation)) { + return; + } + appendQueuedPrompt(item); } -function _qStartSubscription() { - _qStopSubscription(); +function createQueuedPrompt(prompt: string, target: PromptQueueTarget) { + return { + prompt, + target, + }; +} + +function appendTextToThread(prompt: string) { + return { + role: "user", + content: [{ type: "text", text: prompt }], + createdAt: new Date(), + } as never; +} + +function getPromptQueueTargetIds(target: PromptQueueTarget) { + return compactIds([ + ...target.getRunningThreadIds(), + target.getDocumentThreadId(), + ]); +} + +function findPromptQueueEntry( + state: PromptQueueUIState, + threadIds: string[], +) { + for (const threadId of threadIds) { + const entry = state.byThreadId[threadId]; + if (entry) { + return entry; + } + } + return null; +} + +function syncPromptQueueUI() { + if (!promptQueueIsRunning || promptQueueItems.length === 0) { + usePromptQueueUI.setState({ byThreadId: {} }); + return; + } + + const activeItemIndex = Math.max(promptQueueIndex, 0); + const groups: Array<{ + ids: Set; + current: number; + total: number; + active: boolean; + }> = []; + + for (const [index, item] of promptQueueItems.entries()) { + const ids = getPromptQueueTargetIds(item.target); + if (ids.length === 0) { + continue; + } + let group = groups.find((candidate) => + ids.some((id) => candidate.ids.has(id)), + ); + if (!group) { + group = { + ids: new Set(), + current: 0, + total: 0, + active: false, + }; + groups.push(group); + } + ids.forEach((id) => group.ids.add(id)); + group.total += 1; + if (promptQueueIndex >= 0 && index <= promptQueueIndex) { + group.current += 1; + } + if (index === activeItemIndex) { + group.active = true; + } + } + + const byThreadId: Record = {}; + for (const group of groups) { + if (!group.active && group.current >= group.total) { + continue; + } + const entry = { + current: Math.min(group.current, group.total), + total: group.total, + }; + group.ids.forEach((id) => { + byThreadId[id] = entry; + }); + } + + usePromptQueueUI.setState({ byThreadId }); +} + +function isPromptQueueTargetRunning( + target: PromptQueueTarget, + runningByThreadId: Record, +) { + const runningIds = Object.keys(runningByThreadId); + if (runningIds.length === 0) { + return false; + } + + const targetIds = target.getRunningThreadIds(); + if (targetIds.length === 0) { + return runningIds.length > 0; + } + + return runningIds.some((threadId) => targetIds.includes(threadId)); +} + +function isActivePromptQueueTargetRunning( + runningByThreadId: Record, +) { + const activeItem = promptQueueItems[Math.max(promptQueueIndex, 0)]; + if (!activeItem) { + return false; + } + return isPromptQueueTargetRunning(activeItem.target, runningByThreadId); +} + +function advancePromptQueue() { + const nextIndex = promptQueueIndex + 1; + if (nextIndex >= promptQueueItems.length) { + resetPromptQueue(true); + return; + } + promptQueueIndex = nextIndex; + syncPromptQueueUI(); + const next = promptQueueItems[nextIndex]; + toast(`Prompt ${nextIndex + 1} / ${promptQueueItems.length}`, { + description: queueToastDescription(next.prompt), + }); + promptQueueWaitingForTargetIdle = false; + promptQueuePrevStoreRunning = false; + scheduleQueuedPromptDispatch(next, 100); +} + +function startPromptQueueSubscription() { + const wasWaitingForRun = promptQueuePrevStoreRunning; + stopPromptQueueSubscription({ resetRunningState: false }); + promptQueuePrevStoreRunning = wasWaitingForRun; // runningByThreadId tracks the actual thread (not aui.thread()), so detection // survives navigation. - _qStoreUnsub = useChatRuntimeStore.subscribe((state) => { - if (!_qIsRunning) { _qStopSubscription(); return; } - const isRunning = Object.keys(state.runningByThreadId).length > 0; - const wasRunning = _qPrevStoreRunning; - _qPrevStoreRunning = isRunning; + promptQueueStoreUnsub = useChatRuntimeStore.subscribe((state) => { + if (!promptQueueIsRunning) { + stopPromptQueueSubscription(); + return; + } + const isRunning = isActivePromptQueueTargetRunning(state.runningByThreadId); + const wasRunning = promptQueuePrevStoreRunning; + promptQueuePrevStoreRunning = isRunning; if (wasRunning && !isRunning) { - _qAdvance(); + if (promptQueueWaitingForTargetIdle) { + promptQueueWaitingForTargetIdle = false; + const activeItem = promptQueueItems[promptQueueIndex]; + if (activeItem) { + scheduleQueuedPromptDispatch(activeItem, 50); + } + return; + } + advancePromptQueue(); } }); + + const isRunningNow = isActivePromptQueueTargetRunning( + useChatRuntimeStore.getState().runningByThreadId, + ); + if (promptQueuePrevStoreRunning && !isRunningNow) { + promptQueuePrevStoreRunning = false; + if (promptQueueWaitingForTargetIdle) { + promptQueueWaitingForTargetIdle = false; + const activeItem = promptQueueItems[promptQueueIndex]; + if (activeItem) { + scheduleQueuedPromptDispatch(activeItem, 50); + } + return; + } + advancePromptQueue(); + } } -interface _QueueCallbacks { startQueue: (items: string[]) => void; stopQueue: () => void; } -const PromptQueueContext = createContext<_QueueCallbacks>({ - startQueue: () => {}, stopQueue: () => {}, +function startPromptQueue( + items: string[], + target: PromptQueueTarget, + waitForCurrentRun = false, +) { + const filtered = items.map((item) => item.trim()).filter(Boolean); + if (filtered.length === 0) { + return; + } + + if (promptQueueIsRunning) { + promptQueueItems.push( + ...filtered.map((prompt) => createQueuedPrompt(prompt, target)), + ); + syncPromptQueueUI(); + toast.success("Added to prompt queue", { + description: `${filtered.length} prompt${filtered.length === 1 ? "" : "s"} queued.`, + }); + return; + } + + const runningByThreadId = useChatRuntimeStore.getState().runningByThreadId; + const shouldWaitForCurrentRun = + waitForCurrentRun && + isPromptQueueTargetRunning(target, runningByThreadId); + promptQueueGeneration += 1; + promptQueueItems = filtered.map((prompt) => + createQueuedPrompt(prompt, target), + ); + promptQueueIndex = shouldWaitForCurrentRun ? -1 : 0; + promptQueueIsRunning = true; + promptQueuePrevStoreRunning = shouldWaitForCurrentRun; + syncPromptQueueUI(); + toast( + shouldWaitForCurrentRun ? "Prompt queued" : `Prompt 1 / ${filtered.length}`, + { + description: queueToastDescription(filtered[0]), + }, + ); + startPromptQueueSubscription(); + if (!shouldWaitForCurrentRun) { + const first = promptQueueItems[0]; + if (first) { + scheduleQueuedPromptDispatch(first, 50); + } + } +} + +function stopPromptQueueRun() { + const activeTarget = promptQueueItems[Math.max(promptQueueIndex, 0)]?.target; + resetPromptQueue(); + try { + activeTarget?.cancel(); + } catch { + // The active run may have already ended. + } +} + +if (typeof window !== "undefined") { + window.addEventListener(PROMPT_QUEUE_STOP_EVENT, () => stopPromptQueueRun()); +} + +interface PromptQueueCallbacks { + startQueue: (items: string[], waitForCurrentRun?: boolean) => void; + stopQueue: () => void; +} +const noopStartPromptQueue: PromptQueueCallbacks["startQueue"] = () => + undefined; +const noopStopPromptQueue: PromptQueueCallbacks["stopQueue"] = () => undefined; +const PromptQueueContext = createContext({ + startQueue: noopStartPromptQueue, + stopQueue: noopStopPromptQueue, }); // Gap (px) between last message and floating composer; bottom spacer tracks @@ -914,7 +1259,7 @@ const Composer: FC<{ (s) => s.setPendingImageEditReference, ); const { inputProps, isComposing, isComposingRef } = - useImeComposerInputHandlers(); + useImeComposerInputHandlers({ submitOnEnter: true }); const composerText = useAuiState(({ composer }) => composer.text); // Expand only once the input wraps to a second line, not on first keystroke. // Latch until cleared so it can't flip-flop at the wrap boundary. @@ -958,19 +1303,49 @@ const Composer: FC<{ const hasPendingAudio = useChatRuntimeStore((s) => Boolean(s.pendingAudioName), ); + const threadIsRunning = useAuiState(({ thread }) => thread.isRunning); + const threadListItemId = useAuiState( + ({ threadListItem }) => threadListItem.id, + ); + const threadListItemRemoteId = useAuiState( + ({ threadListItem }) => threadListItem.remoteId, + ); const referenceThreadId = threadId ?? activeThreadId ?? null; + const promptQueueThreadIds = compactIds([ + threadListItemId, + threadListItemRemoteId, + threadId, + ]); + const promptQueueActive = usePromptQueueUI((s) => + Boolean(findPromptQueueEntry(s, promptQueueThreadIds)), + ); + useEffect(() => { + if (threadId != null || activeThreadId != null) { + return; + } + stopPromptQueueRun(); + }, [activeThreadId, threadId]); const hasSendableContent = composerText.trim().length > 0 || hasAttachments || hasPendingAudio; + const canQueueCurrentPrompt = + composerText.trim().length > 0 && + !hasAttachments && + !hasPendingAudio && + !isComposing && + !hasPendingAttachments && + !disabled && + !overlay; // Per-thread draft autosave: restore on mount, then mirror composer text // into localStorage (debounced) so a half-typed message survives a // navigation or reload. Cleared once empty (i.e. after a send). Setting the // text even when no draft exists keeps a thread from inheriting the // previous thread's composer contents. - const draftKey = composerDraftKey(activeThreadId); + const draftThreadId = referenceThreadId; + const draftKey = draftThreadId ? composerDraftKey(draftThreadId) : null; const lastDraftKeyRef = useRef(draftKey); useEffect(() => { - const draft = readComposerDraft(draftKey) ?? ""; + const draft = draftKey ? (readComposerDraft(draftKey) ?? "") : ""; const composer = aui.composer(); if (composer.getState().isEditing) { composer.setText(draft); @@ -983,6 +1358,9 @@ const Composer: FC<{ lastDraftKeyRef.current = draftKey; return; } + if (!draftKey) { + return; + } const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300); return () => clearTimeout(t); }, [composerText, draftKey]); @@ -1043,10 +1421,53 @@ const Composer: FC<{ // While this thread's docs index, hold the send and fire it once they finish so // retrieval covers all of them. const [indexingActive, setIndexingActive] = useState(false); + const indexingActiveRef = useRef(false); const [pendingSend, setPendingSend] = useState(false); const pendingSendRef = useRef(false); const waitToastRef = useRef(null); + const handleIndexingChange = useCallback((active: boolean) => { + indexingActiveRef.current = active; + setIndexingActive(active); + }, []); + + const createPromptQueueTarget = useCallback((): PromptQueueTarget => { + const thread = aui.thread(); + const threadListItem = aui.threadListItem(); + const initialState = threadListItem.getState(); + const initialRunningThreadIds = [ + initialState.id, + initialState.remoteId, + referenceThreadId, + ].filter((id): id is string => Boolean(id)); + const initialDocumentThreadId = + initialState.remoteId ?? referenceThreadId ?? null; + return { + getDocumentThreadId: () => { + const state = threadListItem.getState(); + return state.remoteId ?? referenceThreadId ?? initialDocumentThreadId; + }, + getRunningThreadIds: () => { + const state = threadListItem.getState(); + return Array.from( + new Set( + [ + ...initialRunningThreadIds, + state.id, + state.remoteId, + referenceThreadId, + ].filter((id): id is string => Boolean(id)), + ), + ); + }, + append: (prompt) => { + thread.append(appendTextToThread(prompt)); + }, + cancel: () => thread.cancelRun(), + isIndexing: () => indexingActiveRef.current, + }; + }, [aui, referenceThreadId]); + const dismissWaitToast = useCallback(() => { if (waitToastRef.current !== null) { toast.dismiss(waitToastRef.current); @@ -1120,6 +1541,39 @@ const Composer: FC<{ const handleSubmit = useCallback( (event: Parameters["onSubmit"]>>[0]) => { + if (disabled || shouldBlockSend()) { + event.preventDefault(); + return; + } + + if (threadIsRunning || promptQueueActive) { + event.preventDefault(); + if (!canQueueCurrentPrompt) { + if (overlay || hasAttachments || hasPendingAudio) { + toast.error( + threadIsRunning + ? "Wait for the current response to finish" + : "Wait for the prompt queue to finish", + { + description: + "Only text prompts can be queued while a response is running or the prompt queue is active.", + }, + ); + } + return; + } + const queuedPrompt = composerText.trim(); + flushResourcesSync(() => { + aui.composer().setText(""); + }); + startPromptQueue( + [queuedPrompt], + createPromptQueueTarget(), + threadIsRunning, + ); + return; + } + if (interceptSend(event)) return; if (overlay) { @@ -1166,51 +1620,36 @@ const Composer: FC<{ }, [ aui, + canQueueCurrentPrompt, closeOverlay, composerText, + createPromptQueueTarget, + disabled, + hasAttachments, + hasPendingAudio, interceptSend, overlay, + promptQueueActive, referenceThreadId, setImageToolsEnabled, setPendingImageEditReference, + shouldBlockSend, + threadIsRunning, ], ); - // Update the getter every render so the queue always calls the current - // Composer's aui (post-remount). - _qGetAui = () => aui; - const stopQueue = useCallback(() => { - _qIsRunning = false; - _qStopSubscription(); - _useQueueUI.setState({ isRunning: false, current: 0, total: 0 }); - _qItems = []; - _qIndex = 0; - try { _qGetAui().thread().cancelRun(); } catch {} + stopPromptQueueRun(); }, []); - const startQueue = useCallback((items: string[]) => { - const filtered = items.filter((p) => p.trim()); - if (!filtered.length) return; - _qItems = filtered; - _qIndex = 0; - _qIsRunning = true; - _useQueueUI.setState({ isRunning: true, current: 1, total: filtered.length }); - toast(`Prompt 1 / ${filtered.length}`, { - description: filtered[0].length > 80 ? filtered[0].slice(0, 80) + "…" : filtered[0], - }); - // Subscribe BEFORE appending so we don't miss a very fast completion. - _qStartSubscription(); - setTimeout(() => { - _qGetAui().thread().append({ - role: "user", - content: [{ type: "text", text: filtered[0] }], - createdAt: new Date(), - } as never); - }, 50); - }, []); + const startQueue = useCallback( + (items: string[], waitForCurrentRun = threadIsRunning) => { + startPromptQueue(items, createPromptQueueTarget(), waitForCurrentRun); + }, + [createPromptQueueTarget, threadIsRunning], + ); - const queueContextValue: _QueueCallbacks = { startQueue, stopQueue }; + const queueContextValue: PromptQueueCallbacks = { startQueue, stopQueue }; const composerContent = ( <> @@ -1218,7 +1657,7 @@ const Composer: FC<{
{ + const queuedPrompt = composerText.trim(); + if (queuedPrompt.length === 0) { + return; + } + flushResourcesSync(() => { + aui.composer().setText(""); + }); + startPromptQueue([queuedPrompt], createPromptQueueTarget(), true); + }} onSendClick={interceptSend} + onStopClick={stopQueue} pendingSend={pendingSend} menuSide={effectiveMenuSide} + queueThreadIds={promptQueueThreadIds} />
@@ -1330,7 +1782,11 @@ function isNativeComposing(event: Event) { // pause but short enough to recover before the user notices Send is stuck. const IME_STUCK_TIMEOUT_MS = 2500; -function useImeComposerInputHandlers() { +function useImeComposerInputHandlers({ + submitOnEnter = false, +}: { + submitOnEnter?: boolean; +} = {}) { const aui = useAui(); const composingRef = useRef(false); const [isComposing, setIsComposing] = useState(false); @@ -1422,7 +1878,9 @@ function useImeComposerInputHandlers() { if (e.nativeEvent.isComposing || e.keyCode === 229) { composingRef.current = true; refreshStuckTimer(); - } else if (composingRef.current) { + return; + } + if (composingRef.current) { // Candidate-confirming Enter can arrive as non-composing; keep it gated. if (e.key === "Enter") { if (!e.shiftKey) { @@ -1438,8 +1896,12 @@ function useImeComposerInputHandlers() { // rather than waiting for the 2500ms watchdog. setCompositionState(false); } + if (submitOnEnter && e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + e.currentTarget.form?.requestSubmit(); + } }, - [refreshStuckTimer, setCompositionState], + [refreshStuckTimer, setCompositionState, submitOnEnter], ); // On macOS, switching input methods (e.g. ABC → Pinyin) while the textarea @@ -2514,14 +2976,29 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ const ComposerRightControls: FC<{ disabled?: boolean; + queueDisabled?: boolean; + onQueueClick?: () => void; onSendClick?: (event: { preventDefault: () => void }) => void; + onStopClick?: () => void; pendingSend?: boolean; menuSide?: "top" | "bottom"; -}> = ({ disabled, onSendClick, pendingSend, menuSide }) => { - const isQueueRunning = _useQueueUI((s) => s.isRunning); - const queueCurrent = _useQueueUI((s) => s.current); - const queueTotal = _useQueueUI((s) => s.total); - const { stopQueue } = useContext(PromptQueueContext); + queueThreadIds: string[]; +}> = ({ + disabled, + queueDisabled, + onQueueClick, + onSendClick, + onStopClick, + pendingSend, + menuSide, + queueThreadIds, +}) => { + const queueEntry = usePromptQueueUI((s) => + findPromptQueueEntry(s, queueThreadIds), + ); + const isQueueRunning = Boolean(queueEntry); + const queueCurrent = queueEntry?.current ?? 0; + const queueTotal = queueEntry?.total ?? 0; return (
@@ -2550,57 +3027,85 @@ const ComposerRightControls: FC<{ {isQueueRunning ? ( - - ) : ( - <> - !thread.isRunning}> - - onSendClick?.(event)} - className="aui-composer-send ml-1.5 size-8 rounded-full" - aria-label="Send message" - > - {pendingSend ? ( - - ) : ( - - )} - - - - thread.isRunning}> + Queue {queueCurrent}/{queueTotal} + + ) : null} + !thread.isRunning && !isQueueRunning}> + + onSendClick?.(event)} + className="aui-composer-send ml-1.5 size-8 rounded-full" + aria-label="Send message" + > + {pendingSend ? ( + + ) : ( + + )} + + + + {isQueueRunning ? ( + !thread.isRunning}> + + + + + ) : null} + thread.isRunning}> +
+ {queueDisabled ? ( - - - )} + ) : ( + + + + )} +
+
); }; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index c15950c655..3de3ef6dab 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -69,6 +69,7 @@ import { import { isChatThreadDeleted } from "./utils/chat-thread-tombstones"; import { syncExportedRepositoryToBackend } from "./utils/delete-thread-message"; import { getImageInputUnavailableReason } from "./utils/image-input-support"; +import { requestPromptQueueStop } from "./utils/prompt-queue-boundary"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; const pendingHistoryAppendByMessageId = new Map>(); @@ -1047,6 +1048,17 @@ function createRuntimeHook(modelType: ModelType, pairId?: string) { }; } +function stopChatRun(threadId: string | null | undefined) { + if (!threadId) { + return; + } + try { + useChatRuntimeStore.getState().cancelByThreadId[threadId]?.(); + } catch { + // The run may have ended while navigation was mounting. + } +} + function ThreadAutoSwitch({ threadId, syncActiveThreadId = true, @@ -1060,6 +1072,10 @@ function ThreadAutoSwitch({ useEffect(() => { if (!isLoading && mainThreadId !== threadId) { + if (syncActiveThreadId) { + requestPromptQueueStop(); + stopChatRun(mainThreadId); + } const switchResult = aui.threads().switchToThread(threadId) as unknown; if ( switchResult && @@ -1089,11 +1105,16 @@ function ThreadNewChatSwitch({ }: { nonce: string }): ReactElement | null { const aui = useAui(); const isLoading = useAuiState(({ threads }) => threads.isLoading); + const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); + const mainThreadIdRef = useRef(mainThreadId); + mainThreadIdRef.current = mainThreadId; useEffect(() => { if (isLoading) { return; } + requestPromptQueueStop(); + stopChatRun(mainThreadIdRef.current); // Switch to a fresh local thread without persisting it yet; persistence // still happens on first message append. void aui.threads().switchToNewThread(); diff --git a/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts b/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts new file mode 100644 index 0000000000..ac1d973bfe --- /dev/null +++ b/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts @@ -0,0 +1,8 @@ +export const PROMPT_QUEUE_STOP_EVENT = "unsloth:prompt-queue-stop"; + +export function requestPromptQueueStop() { + if (typeof window === "undefined") { + return; + } + window.dispatchEvent(new Event(PROMPT_QUEUE_STOP_EVENT)); +}