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"
+ )