Studio: fix worker ownership on a raced reroute, and the stop-chats prompt

Four review findings on the parallel-chats work, all reproduced first.

- _direct_reader hands a foreign response to its own mailbox, but skipped the
  ownership move the dispatcher makes. A _gen_lock reader already blocked on
  resp_queue can beat the compare dispatcher to that request's first response,
  and the compare consumer opts out of marking, so nothing promoted it: the
  direct request stayed the recorded executor, its late reset cancelled the
  compare generation, and the compare chat's own Stop was ignored.
- A chat stopped while queued on _gen_lock was still claimed and sent once the
  lock freed. Cancellation is only checked on a token, so a long prefill, or a
  generation reaching gen_done without one, occupied the worker after Stop.
  Same hole in the audio-input path, which shares the lock.
- The stop-chats prompt counted generation handles, not conversations. One chat
  holds several while a tool continuation registers its next leg before the
  previous unwinds, so it offered to stop two chats and listed one title.
- Ejecting a model confirms through that dialog, which told the user
  "Unloading the model reloads the model" and offered "Stop and reload".
  Confirming calls /unload and leaves nothing loaded.

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

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)

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

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}?
           </AlertDialogTitle>
           <AlertDialogDescription>
-            {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.
           </AlertDialogDescription>
         </AlertDialogHeader>
@@ -77,7 +83,7 @@ export function StopRunningChatsDialog() {
             Keep generating
           </AlertDialogCancel>
           <AlertDialogAction onClick={() => resolve(true)}>
-            Stop and reload
+            {unloads ? "Stop and unload" : "Stop and reload"}
           </AlertDialogAction>
         </AlertDialogFooter>
       </AlertDialogContent>

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.

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<boolean>;
   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<boolean>((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);
     },
   }));

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<StopRunningChatsDecision> {
   // 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 };

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