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>
This commit is contained in:
Lee Jackson
GitHub
parent
31943d3c21
commit
b53cc957dc
3 changed files with 668 additions and 134 deletions
|
|
@ -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<typeof useAui> = () => {
|
||||
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<string, PromptQueueUIEntry>;
|
||||
}
|
||||
|
||||
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<PromptQueueUIState>(() => ({
|
||||
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<typeof setTimeout> | null = null;
|
||||
|
||||
function compactIds(ids: Array<string | null | undefined>) {
|
||||
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<string>;
|
||||
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<string>(),
|
||||
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<string, PromptQueueUIEntry> = {};
|
||||
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<string, boolean>,
|
||||
) {
|
||||
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<string, boolean>,
|
||||
) {
|
||||
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<PromptQueueCallbacks>({
|
||||
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<string | number | null>(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<NonNullable<ComponentProps<"form">["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<{
|
|||
<PendingAudioChip />
|
||||
<ThreadDocumentsBar
|
||||
threadId={referenceThreadId}
|
||||
onIndexingChange={setIndexingActive}
|
||||
onIndexingChange={handleIndexingChange}
|
||||
/>
|
||||
<ToolStatusDisplay />
|
||||
<div
|
||||
|
|
@ -1269,9 +1708,22 @@ const Composer: FC<{
|
|||
isComposing ||
|
||||
hasPendingAttachments
|
||||
}
|
||||
queueDisabled={!canQueueCurrentPrompt}
|
||||
onQueueClick={() => {
|
||||
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}
|
||||
/>
|
||||
</div>
|
||||
</>
|
||||
|
|
@ -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 (
|
||||
<div className="aui-composer-action-wrapper flex shrink-0 items-center gap-1.5">
|
||||
<ReasoningToggle side={menuSide} />
|
||||
|
|
@ -2550,57 +3027,85 @@ const ComposerRightControls: FC<{
|
|||
</ComposerPrimitive.StopDictation>
|
||||
</ComposerPrimitive.If>
|
||||
{isQueueRunning ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={stopQueue}
|
||||
aria-label="Stop prompt queue"
|
||||
className="ml-1.5 flex items-center gap-1.5 rounded-full border border-border/60 bg-muted/60 px-2.5 py-1 text-xs font-semibold text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
<span
|
||||
className="ml-1 flex h-7 items-center rounded-full bg-primary/10 px-2 text-[11px] font-semibold text-primary"
|
||||
aria-live="polite"
|
||||
>
|
||||
<SquareIcon className="size-2.5 shrink-0 fill-current" />
|
||||
<span className="tabular-nums">
|
||||
Stop queue {queueCurrent}/{queueTotal}
|
||||
</span>
|
||||
</button>
|
||||
) : (
|
||||
<>
|
||||
<AuiIf condition={({ thread }) => !thread.isRunning}>
|
||||
<ComposerPrimitive.Send asChild={true}>
|
||||
<TooltipIconButton
|
||||
tooltip={pendingSend ? "Waiting for documents…" : "Send message"}
|
||||
side="bottom"
|
||||
type="submit"
|
||||
variant="default"
|
||||
size="icon"
|
||||
// Stay clickable while docs index so a click can queue the send;
|
||||
// disabled only once a send is parked.
|
||||
disabled={disabled || pendingSend}
|
||||
onClick={(event) => onSendClick?.(event)}
|
||||
className="aui-composer-send ml-1.5 size-8 rounded-full"
|
||||
aria-label="Send message"
|
||||
>
|
||||
{pendingSend ? (
|
||||
<Spinner className="size-[18px]" />
|
||||
) : (
|
||||
<ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
|
||||
)}
|
||||
</TooltipIconButton>
|
||||
</ComposerPrimitive.Send>
|
||||
</AuiIf>
|
||||
<AuiIf condition={({ thread }) => thread.isRunning}>
|
||||
<span className="tabular-nums">Queue {queueCurrent}/{queueTotal}</span>
|
||||
</span>
|
||||
) : null}
|
||||
<AuiIf condition={({ thread }) => !thread.isRunning && !isQueueRunning}>
|
||||
<ComposerPrimitive.Send asChild={true}>
|
||||
<TooltipIconButton
|
||||
tooltip={pendingSend ? "Waiting for documents…" : "Send message"}
|
||||
side="bottom"
|
||||
type="submit"
|
||||
variant="default"
|
||||
size="icon"
|
||||
// Stay clickable while docs index so a click can queue the send;
|
||||
// disabled only once a send is parked.
|
||||
disabled={disabled || pendingSend}
|
||||
onClick={(event) => onSendClick?.(event)}
|
||||
className="aui-composer-send ml-1.5 size-8 rounded-full"
|
||||
aria-label="Send message"
|
||||
>
|
||||
{pendingSend ? (
|
||||
<Spinner className="size-[18px]" />
|
||||
) : (
|
||||
<ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
|
||||
)}
|
||||
</TooltipIconButton>
|
||||
</ComposerPrimitive.Send>
|
||||
</AuiIf>
|
||||
{isQueueRunning ? (
|
||||
<AuiIf condition={({ thread }) => !thread.isRunning}>
|
||||
<TooltipIconButton
|
||||
tooltip="Queue message"
|
||||
side="bottom"
|
||||
type="button"
|
||||
variant="default"
|
||||
size="icon"
|
||||
disabled={disabled || queueDisabled}
|
||||
onClick={onQueueClick}
|
||||
className="aui-composer-send ml-1.5 size-8 rounded-full"
|
||||
aria-label="Queue message"
|
||||
>
|
||||
<ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
|
||||
</TooltipIconButton>
|
||||
</AuiIf>
|
||||
) : null}
|
||||
<AuiIf condition={({ thread }) => thread.isRunning}>
|
||||
<div className="ml-1.5 flex items-center">
|
||||
{queueDisabled ? (
|
||||
<ComposerPrimitive.Cancel asChild={true}>
|
||||
<Button
|
||||
type="button"
|
||||
variant="default"
|
||||
size="icon"
|
||||
className="aui-composer-cancel ml-1.5 size-8 rounded-full"
|
||||
className="aui-composer-cancel size-8 rounded-full"
|
||||
aria-label="Stop generating"
|
||||
onClick={isQueueRunning ? onStopClick : undefined}
|
||||
>
|
||||
<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
|
||||
</Button>
|
||||
</ComposerPrimitive.Cancel>
|
||||
</AuiIf>
|
||||
</>
|
||||
)}
|
||||
) : (
|
||||
<TooltipIconButton
|
||||
tooltip="Queue message"
|
||||
side="bottom"
|
||||
type="button"
|
||||
variant="default"
|
||||
size="icon"
|
||||
disabled={queueDisabled}
|
||||
onClick={onQueueClick}
|
||||
className="aui-composer-send size-8 rounded-full"
|
||||
aria-label="Queue message"
|
||||
>
|
||||
<ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
|
||||
</TooltipIconButton>
|
||||
)}
|
||||
</div>
|
||||
</AuiIf>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<string, Promise<void>>();
|
||||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue