From 9954781d30fb568520204cd581878281b8398258 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Mon, 20 Apr 2026 21:06:59 +0400 Subject: [PATCH] fix(studio/chat): cancel in-flight run when trashing a thread from sidebar (#5067) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trashing a thread mid-stream used to delete the Dexie rows while the model kept generating, because the sidebar has no access to the @assistant-ui aui context. Expose per-thread cancelRun() through the chat runtime store and call it from deleteChatItem so trash behaves like Stop → Trash. Covers compare pairs by cancelling each paired thread. Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../chat/hooks/use-chat-sidebar-items.ts | 31 +++++++++++++------ .../src/features/chat/runtime-provider.tsx | 29 +++++++++++++++++ .../chat/stores/chat-runtime-store.ts | 17 ++++++++++ 3 files changed, 68 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index fb27f990cf..33cdd88081 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -58,23 +58,36 @@ export function useChatSidebarItems() { return { items, canCompare }; } +function cancelIfRunning(threadId: string): void { + const { runningByThreadId, cancelByThreadId } = + useChatRuntimeStore.getState(); + if (!runningByThreadId[threadId]) return; + cancelByThreadId[threadId]?.(); +} + export async function deleteChatItem( item: SidebarItem, activeId: string | undefined, onSelect: (view: { mode: "single"; newThreadNonce: string }) => void, ) { + const threadIds: string[] = + item.type === "single" + ? [item.id] + : (await db.threads.where("pairId").equals(item.id).toArray()).map( + (t) => t.id, + ); + + // Stop any in-flight streams before deleting, so the model doesn't keep + // generating against a thread that no longer exists. + for (const id of threadIds) cancelIfRunning(id); + await db.transaction("rw", db.threads, db.messages, async () => { - if (item.type === "single") { - await db.messages.where("threadId").equals(item.id).delete(); - await db.threads.delete(item.id); - } else { - const paired = await db.threads.where("pairId").equals(item.id).toArray(); - for (const t of paired) { - await db.messages.where("threadId").equals(t.id).delete(); - await db.threads.delete(t.id); - } + for (const id of threadIds) { + await db.messages.where("threadId").equals(id).delete(); + await db.threads.delete(id); } }); + if (activeId === item.id) { useChatRuntimeStore.getState().setActiveThreadId(null); onSelect({ mode: "single", newThreadNonce: crypto.randomUUID() }); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 348a0a85c4..5747e17970 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -730,6 +730,34 @@ function ActiveThreadSync({ return null; } +// Exposes the current thread's cancelRun() via the shared store so external +// surfaces (e.g. the sidebar trash button) can stop an in-flight stream +// before deleting the thread — mirroring the Stop → Trash sequence. +function CancelRegistrar(): ReactElement | null { + const aui = useAui(); + const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); + const isRunning = useChatRuntimeStore((s) => + mainThreadId ? Boolean(s.runningByThreadId[mainThreadId]) : false, + ); + + useEffect(() => { + if (!mainThreadId || !isRunning) return; + const cancel = () => { + try { + aui.thread().cancelRun(); + } catch { + // Run may have already ended between the caller's read and this call. + } + }; + useChatRuntimeStore.getState().registerThreadCancel(mainThreadId, cancel); + return () => { + useChatRuntimeStore.getState().clearThreadCancel(mainThreadId); + }; + }, [aui, mainThreadId, isRunning]); + + return null; +} + export function ChatRuntimeProvider({ children, modelType = "base", @@ -762,6 +790,7 @@ export function ChatRuntimeProvider({ + {initialThreadId && ( ; + cancelByThreadId: Record void>; autoTitle: boolean; hfToken: string; modelsError: string | null; @@ -189,6 +190,8 @@ type ChatRuntimeStore = { setModels: (models: ChatModelSummary[]) => void; setLoras: (loras: ChatLoraSummary[]) => void; setThreadRunning: (threadId: string, running: boolean) => void; + registerThreadCancel: (threadId: string, cancel: () => void) => void; + clearThreadCancel: (threadId: string) => void; setAutoTitle: (enabled: boolean) => void; setHfToken: (token: string) => void; setModelsError: (error: string | null) => void; @@ -218,6 +221,7 @@ export const useChatRuntimeStore = create((set) => ({ models: [], loras: [], runningByThreadId: {}, + cancelByThreadId: {}, autoTitle: loadBool(AUTO_TITLE_KEY, false), hfToken: loadString(HF_TOKEN_KEY, ""), modelsError: null, @@ -277,6 +281,19 @@ export const useChatRuntimeStore = create((set) => ({ } return { runningByThreadId: next }; }), + registerThreadCancel: (threadId, cancel) => + set((state) => { + const next = { ...state.cancelByThreadId }; + next[threadId] = cancel; + return { cancelByThreadId: next }; + }), + clearThreadCancel: (threadId) => + set((state) => { + if (!(threadId in state.cancelByThreadId)) return state; + const next = { ...state.cancelByThreadId }; + delete next[threadId]; + return { cancelByThreadId: next }; + }), setAutoTitle: (autoTitle) => set(() => { saveBool(AUTO_TITLE_KEY, autoTitle);