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);