diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 03d63ed37d..4699148a08 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -531,7 +531,17 @@ class InferenceOrchestrator: if rid and rid != request_id: with self._mailbox_lock: other = self._mailboxes.get(rid) or self._direct_mailboxes.get(rid) + owner = self._request_cancel_events.get(rid) if other is not None: + # We beat the dispatcher to this response, so make its ownership move here + # too. The compare consumer opts out of marking, so nothing else promotes + # or retires that request: skipping it left this one recorded as the + # executor, ignoring its Stop and letting a late reset cancel it. + if owner is not None: + if resp.get("type", "") in ("gen_done", "gen_error"): + self._release_worker(owner) + else: + self._mark_worker_started(owner) other.put(resp) return None return resp @@ -1705,6 +1715,11 @@ class InferenceOrchestrator: # Won the lock handoff during a switch; don't start on the outgoing model. yield GenStreamError("Error: model is being unloaded", public = True) return + if cancel_event is not None and cancel_event.is_set(): + # Stopped while queued on the lock. Sending anyway occupied the worker with a + # run the user ended: the cancel is only seen on a token, so a long prefill + # (or a generation that reaches gen_done without one) held up its siblings. + return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None cmd = self._build_generate_cmd( @@ -1980,6 +1995,9 @@ class InferenceOrchestrator: # Won the lock handoff during a switch; don't start on the outgoing model. yield GenStreamError("Error: model is being unloaded", public = True) return + if cancel_event is not None and cancel_event.is_set(): + # Stopped while queued on the lock, same as _generate_inner. + return request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py index 9c5a492647..ea903a6ce0 100644 --- a/studio/backend/tests/test_inference_dispatcher_resilience.py +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -119,3 +119,68 @@ def test_route_llama_streaming_async_clients_disable_proxy_env(): kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False for kw in call.keywords ), f"httpx.AsyncClient at line {call.lineno} must set trust_env=False" + + +def _direct_reader_host(): + """Orchestrator with only what _direct_reader and the ownership helpers touch.""" + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._direct_mailboxes = {} + o._request_cancel_events = {} + o._active_cancel_lock = threading.Lock() + o._active_cancel_events = [] + o._executing_cancel_events = [] + o._dispatcher_thread = None + return o + + +def test_rerouting_a_foreign_response_moves_worker_ownership(): + # A _gen_lock reader already blocked on resp_queue can beat the compare dispatcher to + # that request's first response. The compare consumer passes mark_started=False, so if + # this path does not promote it nothing does: the direct request stays recorded as the + # executor, so the compare chat's Stop is ignored and a late reset from the direct one + # cancels the compare generation instead. + o = _direct_reader_host() + mine, theirs = threading.Event(), threading.Event() + o._request_cancel_events = {"mine": mine, "theirs": theirs} + o._claim_worker(mine) + o._mark_worker_started(mine) + o._claim_worker(theirs) + compare_mailbox = queue.Queue() + o._mailboxes["theirs"] = compare_mailbox + + read_one, _drain, release = _direct_reader_calls(o, "mine") + o._scripted = [{"request_id": "theirs", "type": "token", "text": "hi"}] + + assert read_one(timeout = 0.1) is None, "a foreign response is routed, not returned" + assert compare_mailbox.get_nowait()["text"] == "hi" + assert o._owns_worker(theirs), "the compare request is the one the worker answered" + assert not o._owns_worker(mine), "so a late reset from the direct request must not fire" + release() + + +def test_rerouting_a_foreign_gen_done_retires_that_request(): + # The other half of the dispatcher's move: once its last response is routed, the + # request no longer owns the worker, or a Stop for it would end whatever starts next. + o = _direct_reader_host() + mine, theirs = threading.Event(), threading.Event() + o._request_cancel_events = {"mine": mine, "theirs": theirs} + o._claim_worker(theirs) + o._mark_worker_started(theirs) + o._claim_worker(mine) + o._mailboxes["theirs"] = queue.Queue() + + read_one, _drain, release = _direct_reader_calls(o, "mine") + o._scripted = [{"request_id": "theirs", "type": "gen_done"}] + + assert read_one(timeout = 0.1) is None + assert not o._owns_worker(theirs), "retired once its last response was routed" + assert o._owns_worker(mine), "the next claim takes over" + release() + + +def _direct_reader_calls(o, request_id): + """_direct_reader wired to a scripted _read_resp (o._scripted, popped in order).""" + o._read_resp = lambda timeout = 1.0: o._scripted.pop(0) if o._scripted else None + return o._direct_reader(request_id) diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index d3e9b1e06b..7963b71e8e 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -1991,3 +1991,45 @@ def test_audio_input_claims_the_worker_before_sending(): assert claim < send, "the claim has to happen before the command is enqueued" assert "with self._send_order_lock:" in body, "claim and send must be one critical section" assert "self._release_worker(cancel_event)" in body + + +def test_generation_stopped_while_queued_is_never_sent(monkeypatch): + # Two chats on the serialized backend: the second blocks on _gen_lock, and Stop sets its + # event while it waits. Sending anyway occupied the worker with a run the user ended -- + # the cancel is only checked on a token, so a long prefill (or a generation that reaches + # gen_done without one) still held up its siblings. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda *a, **k: None) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped") + ) + stopped = threading.Event() + stopped.set() + + out = list( + o._generate_inner(messages = [{"role": "user", "content": "hi"}], cancel_event = stopped) + ) + + assert out == [], "a stopped request yields nothing rather than an error banner" + assert o._active_cancel_events == [], "it must not claim the worker either" + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_audio_input_stopped_while_queued_is_never_sent(monkeypatch): + # Same lock, same hole. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped") + ) + stopped = threading.Event() + stopped.set() + + out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1], cancel_event = stopped)) + + assert out == [] + assert o._active_cancel_events == [] + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() diff --git a/studio/frontend/src/features/chat/components/stop-running-chats-dialog.tsx b/studio/frontend/src/features/chat/components/stop-running-chats-dialog.tsx index 3cf9c38cd3..dd9d6c13a1 100644 --- a/studio/frontend/src/features/chat/components/stop-running-chats-dialog.tsx +++ b/studio/frontend/src/features/chat/components/stop-running-chats-dialog.tsx @@ -24,6 +24,7 @@ export function StopRunningChatsDialog() { const titles = useStopRunningChatsDialogStore((s) => s.titles); const action = useStopRunningChatsDialogStore((s) => s.action); const hasNonChat = useStopRunningChatsDialogStore((s) => s.hasNonChat); + const effect = useStopRunningChatsDialogStore((s) => s.effect); const resolve = useStopRunningChatsDialogStore((s) => s.resolve); // Embeddings, raw completions and audio share the model but are not conversations, @@ -35,6 +36,13 @@ export function StopRunningChatsDialog() { : count === 1 ? "chat" : "chats"; + const sharer = hasNonChat ? "request" : "conversation"; + // Ejecting leaves no model loaded. Saying it "reloads the model" and offering "Stop and + // reload" promised the opposite of what confirming does, for the destructive one. + const unloads = effect === "unload"; + const lead = unloads + ? `${action || "Unloading the model"} leaves no model loaded, and every open ${sharer} shares it, ` + : `${action ? `${action} reloads the model, ` : "Reloading the model "}which every open ${sharer} shares, `; const shown = titles.slice(0, 5); const remaining = Math.max(0, titles.length - shown.length); @@ -52,9 +60,7 @@ export function StopRunningChatsDialog() { Stop {count} running {noun}? - {action ? `${action} reloads the model, ` : "Reloading the model "} - which every open {hasNonChat ? "request" : "conversation"} shares, - so {count === 1 ? "this" : "these"} {noun} will stop + {lead}so {count === 1 ? "this" : "these"} {noun} will stop {hasNonChat ? "" : " generating"}. Work produced so far is kept. @@ -77,7 +83,7 @@ export function StopRunningChatsDialog() { Keep generating resolve(true)}> - Stop and reload + {unloads ? "Stop and unload" : "Stop and reload"} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index cf6bedb473..d4057591b0 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1632,9 +1632,11 @@ export function useChatModelRuntime() { return true; } try { - // Ejecting tears down llama-server, so every chat stops. Same prompt. + // Ejecting tears down llama-server, so every chat stops. Same prompt, but it + // leaves no model loaded, so it must not be worded as a reload. const stopDecision = await confirmStopRunningChatsIfNeeded( "Unloading the model", + "unload", ); if (!stopDecision.proceed) return false; // Same window as selectModel: a load may have started during the confirm. diff --git a/studio/frontend/src/features/chat/stores/stop-running-chats-dialog-store.ts b/studio/frontend/src/features/chat/stores/stop-running-chats-dialog-store.ts index 5c4a12d8f3..01ccc76f43 100644 --- a/studio/frontend/src/features/chat/stores/stop-running-chats-dialog-store.ts +++ b/studio/frontend/src/features/chat/stores/stop-running-chats-dialog-store.ts @@ -5,6 +5,9 @@ import { create } from "zustand"; type Resolver = (confirmed: boolean) => void; +/** What confirming does to the model: reload it, or leave none loaded. */ +export type StopRunningChatsEffect = "reload" | "unload"; + // One at a time: a new request declines any pending one so no promise leaks. let pendingResolver: Resolver | null = null; @@ -18,11 +21,14 @@ interface StopRunningChatsDialogStore { action: string; /** The set includes an embeddings/completions/audio request, which is not a chat. */ hasNonChat: boolean; + /** Ejecting leaves no model loaded, so it must not be described as a reload. */ + effect: StopRunningChatsEffect; requestConfirm: (args: { count: number; titles?: string[]; action?: string; hasNonChat?: boolean; + effect?: StopRunningChatsEffect; }) => Promise; resolve: (confirmed: boolean) => void; } @@ -34,16 +40,30 @@ export const useStopRunningChatsDialogStore = titles: [], action: "", hasNonChat: false, - requestConfirm: ({ count, titles = [], action = "", hasNonChat = false }) => + effect: "reload", + requestConfirm: ({ + count, + titles = [], + action = "", + hasNonChat = false, + effect = "reload", + }) => new Promise((resolve) => { pendingResolver?.(false); pendingResolver = resolve; - set({ open: true, count, titles, action, hasNonChat }); + set({ open: true, count, titles, action, hasNonChat, effect }); }), resolve: (confirmed) => { const resolver = pendingResolver; pendingResolver = null; - set({ open: false, count: 0, titles: [], action: "", hasNonChat: false }); + set({ + open: false, + count: 0, + titles: [], + action: "", + hasNonChat: false, + effect: "reload", + }); resolver?.(confirmed); }, })); diff --git a/studio/frontend/src/features/chat/utils/confirm-stop-running-chats.ts b/studio/frontend/src/features/chat/utils/confirm-stop-running-chats.ts index c99ccbcd54..7b20ceb362 100644 --- a/studio/frontend/src/features/chat/utils/confirm-stop-running-chats.ts +++ b/studio/frontend/src/features/chat/utils/confirm-stop-running-chats.ts @@ -3,7 +3,10 @@ import { getActiveGenerations } from "../api/chat-api"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; -import { useStopRunningChatsDialogStore } from "../stores/stop-running-chats-dialog-store"; +import { + type StopRunningChatsEffect, + useStopRunningChatsDialogStore, +} from "../stores/stop-running-chats-dialog-store"; import { listStoredChatThreads } from "./chat-history-storage"; export interface StopRunningChatsDecision { @@ -20,6 +23,7 @@ export interface StopRunningChatsDecision { */ export async function confirmStopRunningChatsIfNeeded( action = "Loading a different model", + effect: StopRunningChatsEffect = "reload", ): Promise { // Local runs only: an external-provider chat is not stopped by the swap, so counting it // would block a safe load behind a dialog. The backend excludes them for the same reason. @@ -36,19 +40,23 @@ export async function confirmStopRunningChatsIfNeeded( // The union stays local-only, since external-provider runs are never in it. try { const active = await getActiveGenerations(); + const entries = active.active ?? []; const merged = new Set(running); for (const threadId of active.thread_ids ?? []) { merged.add(threadId); } running = [...merged]; - // A first turn started before its id was persisted is counted but not named, so never - // claim fewer chats than the backend reports. - count = Math.max(active.count ?? 0, running.length); + // Count conversations, not handles: one chat holds several at once while a tool + // continuation registers its next leg before the previous unwinds, and active.count + // counts those separately. A first turn started before its id was persisted has no + // id to merge, so add those back or the prompt names fewer chats than will stop. + const unnamed = entries.filter((entry) => !entry.thread_id).length; + count = entries.length + ? running.length + unnamed + : Math.max(active.count ?? 0, running.length); // Embeddings / completions / audio share the model but are not conversations, so the // prompt must not offer to stop chats that do not exist. - hasNonChat = (active.active ?? []).some( - (entry) => (entry.kind ?? "chat") !== "chat", - ); + hasNonChat = entries.some((entry) => (entry.kind ?? "chat") !== "chat"); } catch { // Backend unreachable / older build: fall back to the local map only. } @@ -80,7 +88,7 @@ export async function confirmStopRunningChatsIfNeeded( const confirmed = await useStopRunningChatsDialogStore .getState() - .requestConfirm({ count, titles, action, hasNonChat }); + .requestConfirm({ count, titles, action, hasNonChat, effect }); if (!confirmed) { return { proceed: false, forceCancelActive: false }; diff --git a/tests/studio/test_stop_running_chats_prompt_contract.py b/tests/studio/test_stop_running_chats_prompt_contract.py new file mode 100644 index 0000000000..2a7114723a --- /dev/null +++ b/tests/studio/test_stop_running_chats_prompt_contract.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Source contracts for the "stop running chats" confirmation. + +The dialog is what a user reads before losing in-flight work, so two things have +to hold: it counts conversations rather than generation handles, and it describes +what confirming actually does. There is no frontend test runner in this repo, so +these read the source the way the other frontend contracts here do. +""" + +from __future__ import annotations + +from pathlib import Path + +WORKDIR = Path(__file__).resolve().parents[2] +FRONTEND = WORKDIR / "studio" / "frontend" / "src" + + +def _read(rel: str) -> str: + path = FRONTEND / rel + assert path.exists(), f"missing source file: {path}" + return path.read_text(encoding = "utf-8") + + +def test_the_prompt_counts_conversations_not_generation_handles(): + # One chat holds several handles while a tool continuation registers its next leg + # before the previous unwinds (active_generations.ActiveGeneration mints one per + # __enter__), so active.count exceeds the deduplicated thread_ids and the dialog + # offered to stop two chats while listing one title. + src = _read("features/chat/utils/confirm-stop-running-chats.ts") + assert "entry.thread_id" in src, "the unnamed entries have to be counted separately" + # The raw handle count survives only for a backend too old to send the entries. + primary = src.index("running.length + unnamed") + fallback = src.index("Math.max(active.count") + assert primary < fallback, "the handle count must be the fallback, not the primary" + + +def test_an_unload_is_not_described_as_a_reload(): + # ejectModel confirms through the same dialog, but confirming calls /unload and + # leaves no model loaded: "Unloading the model reloads the model" and "Stop and + # reload" promised the opposite for the destructive one. + dialog = _read("features/chat/components/stop-running-chats-dialog.tsx") + assert "Stop and unload" in dialog and "Stop and reload" in dialog + assert "leaves no model loaded" in dialog + + runtime = _read("features/chat/hooks/use-chat-model-runtime.ts") + eject = runtime.index('"Unloading the model"') + assert '"unload"' in runtime[eject : eject + 120], ( + "the eject path must ask for the unload wording" + )