From aec27263e701ada32b2a9881c547f35fa413e24c Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Fri, 19 Jun 2026 01:54:33 -0700 Subject: [PATCH 01/16] Polish the Hub card carousel: edge shadows, fades, and drag-to-scroll (#6461) * Fix clipped card shadows in Hub trending carousel The carousel scroller only had vertical padding, so with overflow-x set the first and last cards had their drop shadow clipped on the horizontal edges. Add px-2 with a matching -mx-2 so the shadow has room while the cards stay aligned with the section heading, and scroll-px-2 so snap-start does not scroll the padding away on load. * Align carousel edge fades with the scroll clip edge The shadow fix gave the scroller an -mx-2 bleed, but the left/right fade overlays stayed pinned to the wrapper edges, 8px inside the clip edge. That left a thin strip where a card showed beside the fade, so the fade read as a separate block instead of blending into the background. Offset both fades by the same 8px so their opaque edge sits on the clip edge. * Add click-and-drag panning to the Hub card carousel The rows only scrolled by wheel or trackpad, and grabbing a card started a native drag of its avatar image, so the cards could not be dragged to move the row. Add mouse drag-to-scroll (touch and pen keep native scrolling), swallow the click a drag would otherwise fire on a card, keep plain clicks working, and block the avatar's native drag. * Trim carousel edge fade width from 56px to 44px * Smooth out carousel drag panning Scroll snap was correcting the position on every drag frame, which made the pan feel sticky. Disable snap while a drag is active and restore it on release so the row follows the pointer and then settles on a card. * Drop stale carousel drag when the button is released off-element If a press ended outside the scroller before the drag threshold was crossed, no pointerup reached us and the drag stayed armed, so a later buttonless mousemove would scroll the row. Bail out and clear the drag whenever the primary button is no longer held. --- .../features/hub/catalog/card-carousel.tsx | 70 ++++++++++++++++++- studio/frontend/src/features/hub/hub.css | 9 ++- 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/hub/catalog/card-carousel.tsx b/studio/frontend/src/features/hub/catalog/card-carousel.tsx index 932df88ff6..0816188143 100644 --- a/studio/frontend/src/features/hub/catalog/card-carousel.tsx +++ b/studio/frontend/src/features/hub/catalog/card-carousel.tsx @@ -5,6 +5,8 @@ import { cn } from "@/lib/utils"; import { ArrowLeft01Icon, ArrowRight01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { + type MouseEvent as ReactMouseEvent, + type PointerEvent as ReactPointerEvent, type ReactNode, useCallback, useEffect, @@ -102,13 +104,79 @@ export function CardCarousel({ [stepPx], ); + // Click-and-drag panning (mouse only; touch/pen keep native scrolling). + const drag = useRef<{ id: number; x: number; left: number; moved: boolean } | null>( + null, + ); + const suppressClick = useRef(false); + + const onPointerDown = useCallback((e: ReactPointerEvent) => { + suppressClick.current = false; + const el = scrollerRef.current; + if (!el || e.pointerType !== "mouse" || e.button !== 0) return; + drag.current = { id: e.pointerId, x: e.clientX, left: el.scrollLeft, moved: false }; + }, []); + + const onPointerMove = useCallback((e: ReactPointerEvent) => { + const d = drag.current; + const el = scrollerRef.current; + if (!d || !el || e.pointerId !== d.id) return; + // Primary button no longer held: the press ended off the scroller, so no + // pointerup reached us. Drop the stale drag instead of scrolling on hover. + if ((e.buttons & 1) === 0) { + if (d.moved) el.style.scrollSnapType = ""; + drag.current = null; + return; + } + const dx = e.clientX - d.x; + // Ignore tiny moves so plain clicks still register. + if (!d.moved && Math.abs(dx) < 5) return; + if (!d.moved) { + d.moved = true; + // Snap fights the per-frame scrollLeft writes; disable it while dragging. + el.style.scrollSnapType = "none"; + el.setPointerCapture(d.id); + } + el.scrollLeft = d.left - dx; + }, []); + + const endDrag = useCallback((e: ReactPointerEvent) => { + const d = drag.current; + if (!d || e.pointerId !== d.id) return; + if (d.moved) { + // A drag just happened: swallow the click it would fire on a card. + suppressClick.current = true; + const el = scrollerRef.current; + // Restore snap so the row settles on a card after the drag. + if (el) el.style.scrollSnapType = ""; + el?.releasePointerCapture?.(d.id); + } + drag.current = null; + }, []); + + const onClickCapture = useCallback((e: ReactMouseEvent) => { + if (!suppressClick.current) return; + suppressClick.current = false; + e.preventDefault(); + e.stopPropagation(); + }, []); + return (
e.preventDefault()} aria-label={ariaLabel} - className="hub-carousel flex snap-x gap-4 overflow-x-auto pb-4 pt-2" + // px-2 + -mx-2 give card shadows room so the edge cards aren't clipped; + // scroll-px-2 keeps snap-start aligned with the heading. + className="hub-carousel -mx-2 flex cursor-grab snap-x scroll-px-2 gap-4 overflow-x-auto px-2 pb-4 pt-2 select-none active:cursor-grabbing" > {items.map((item) => (
Date: Fri, 19 Jun 2026 11:07:10 +0100 Subject: [PATCH 02/16] Unify retrieval slider styling (#6436) --- .../components/retrieval-settings-section.tsx | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx index d8f69d93a1..103ed5c8e8 100644 --- a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx +++ b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx @@ -1,7 +1,6 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import type { ReactNode } from "react"; import { Select, SelectContent, @@ -10,22 +9,20 @@ import { SelectValue, } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; -import { - ToggleGroup, - ToggleGroupItem, -} from "@/components/ui/toggle-group"; +import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { Tooltip, TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { InfoIcon } from "lucide-react"; -import { cn } from "@/lib/utils"; import { type RagAutoInject, type RagMode, useChatRuntimeStore, } from "@/features/chat/stores/chat-runtime-store"; +import { cn } from "@/lib/utils"; +import { InfoIcon } from "lucide-react"; +import type { ReactNode } from "react"; const MODE_LABEL: Record = { hybrid: "Hybrid", @@ -92,6 +89,7 @@ function SliderRow({ disabled={disabled} onValueChange={([v]) => onChange(v)} aria-label={label} + className="panel-slider" />
); @@ -152,6 +150,7 @@ export function RetrievalSettingsSection() { step={1} onValueChange={([value]) => setRagTopK(value)} aria-label="Number of passages to retrieve" + className="panel-slider" />
@@ -175,7 +174,9 @@ export function RetrievalSettingsSection() { value={ragAutoInject} onValueChange={(value) => { // Radix clears on re-click; ignore empty so one stays selected. - if (value) setRagAutoInject(value as RagAutoInject); + if (value) { + setRagAutoInject(value as RagAutoInject); + } }} className="w-full" aria-label="Auto-retrieve documents" From b552f2fbc801630ea67350f18e39df25cc4e76a4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 19 Jun 2026 04:33:52 -0700 Subject: [PATCH 03/16] studio: fix two backend CI test failures (capability dict + MTP recovery race) (#6464) test_safetensors_capability_advertise: detect_reasoning_flags now returns a reasoning_effort_levels key, so the none-template expectation must include it. test_tensor_parallel::test_runtime_recovery_reloads_without_mtp: the assertion raced the recovery thread, which sets _spec_fallback_reason just before its finally clears _mtp_runtime_fallback_in_progress. Wait for the flag to clear before asserting. Co-authored-by: danielhanchen --- .../backend/tests/test_safetensors_capability_advertise.py | 1 + studio/backend/tests/test_tensor_parallel.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 1e8fb9e2b2..13cb6bbd46 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -107,6 +107,7 @@ def test_detect_safetensors_features_none_template_returns_all_false(): "supports_reasoning": False, "reasoning_style": "enable_thinking", "reasoning_always_on": False, + "reasoning_effort_levels": [], "supports_preserve_thinking": False, "supports_tools": False, } diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 5b1138ef09..86ac79eda9 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -365,6 +365,11 @@ def test_runtime_recovery_reloads_without_mtp(monkeypatch): while b._spec_fallback_reason != "runtime_error" and time.monotonic() < deadline: time.sleep(0.02) assert b._spec_fallback_reason == "runtime_error" + # The reload thread clears the single-flight flag in its finally, a beat after + # it sets the fallback reason -- wait for that instead of racing the thread. + deadline = time.monotonic() + 2 + while b._mtp_runtime_fallback_in_progress and time.monotonic() < deadline: + time.sleep(0.02) assert b._mtp_runtime_fallback_in_progress is False From 1eb15162d943d8767c857ebe13d0a370bdd2e50e Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 19 Jun 2026 13:06:03 +0100 Subject: [PATCH 04/16] fix: clean up Studio warning log formatting (#6265) * feat: queue chat prompts during generation * fix: address prompt queue review edge cases * fix: harden queued prompt dispatch * fix: track queued prompt run state by thread * fix: preserve prompt queue ordering * fix: isolate prompt queue on new chat * fix: clean up Studio warning log formatting * Fix export log markup --------- Co-authored-by: wasimysaid --- .../export/components/export-run-panel.tsx | 13 ++---- .../src/features/export/lib/log-style.ts | 40 +++++++++++++++++++ studio/install_python_stack.py | 31 +++++++++++++- 3 files changed, 73 insertions(+), 11 deletions(-) create mode 100644 studio/frontend/src/features/export/lib/log-style.ts diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9a4f0ee137..ebd35c2158 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -30,6 +30,7 @@ import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { EXPORT_METHODS, type ExportMethod } from "../constants"; import type { ExportLogEntry } from "../api/export-api"; +import { getExportLogLineClass } from "../lib/log-style"; import { selectExportProgressPercent, useExportRuntimeStore, @@ -520,22 +521,16 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
) : ( - 
+                  

                     {run.logLines.map((entry, idx) => (
                       

                         {formatLogLine(entry)}
                       

                     ))}
-
+ )} diff --git a/studio/frontend/src/features/export/lib/log-style.ts b/studio/frontend/src/features/export/lib/log-style.ts new file mode 100644 index 0000000000..f5c057702b --- /dev/null +++ b/studio/frontend/src/features/export/lib/log-style.ts @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import type { ExportLogEntry } from "../api/export-api"; + +type ExportLogTone = "stdout" | "stderr" | "status" | "warning"; + +const WARNING_LINE_PATTERNS = [ + /Skipping import of cpp extensions due to incompatible torch version/i, + /Please see GitHub issue #2919 for more info/i, + /torch_dtype is deprecated!\s*Use dtype instead!/i, +] as const; + +function isWarningLine(line: string): boolean { + return WARNING_LINE_PATTERNS.some((pattern) => pattern.test(line)); +} + +export function getExportLogTone(entry: ExportLogEntry): ExportLogTone { + if (entry.stream === "status") { + return "status"; + } + if (isWarningLine(entry.line)) { + return "warning"; + } + return entry.stream === "stderr" ? "stderr" : "stdout"; +} + +export function getExportLogLineClass(entry: ExportLogEntry): string { + const tone = getExportLogTone(entry); + if (tone === "stderr") { + return "text-rose-300/90"; + } + if (tone === "status") { + return "text-sky-300/90"; + } + if (tone === "warning") { + return "text-status-warning"; + } + return ""; +} diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 36c2bc05b5..b2de6592c8 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -20,6 +20,7 @@ import subprocess import sys import sysconfig import tempfile +import textwrap import urllib.request from pathlib import Path @@ -1401,6 +1402,7 @@ VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1" # Update _TOTAL if you add/remove steps in install_python_stack(). _STEP: int = 0 _TOTAL: int = 0 # set at runtime in install_python_stack() based on platform +_PROGRESS_LINE_ACTIVE: bool = False # -- Paths -------------------------------------------------------------- SCRIPT_DIR = Path(__file__).resolve().parent @@ -1486,6 +1488,7 @@ _HAS_COLOR = _stdout_supports_color() # 2-space indent, 15-char label (dim), then value. _LABEL = "deps" _COL = 15 +_INDENT = 2 def _green(msg: str) -> str: @@ -1517,15 +1520,38 @@ def _step( color_fn = None, ) -> None: """Print a single step line in the column format.""" + global _PROGRESS_LINE_ACTIVE if color_fn is None: color_fn = _green padded = label[:_COL] - _safe_print(f" {_dim(padded)}{' ' * (_COL - len(padded))}{color_fn(value)}") + plain_prefix_width = _INDENT + _COL + prefix = f"{' ' * _INDENT}{_dim(padded)}{' ' * (_COL - len(padded))}" + wrap_width = max( + 24, + shutil.get_terminal_size((100, 20)).columns - plain_prefix_width, + ) + lines = textwrap.wrap( + value, + width = wrap_width, + break_long_words = False, + break_on_hyphens = False, + ) or [""] + if _PROGRESS_LINE_ACTIVE and not VERBOSE: + try: + sys.stdout.write("\n") + sys.stdout.flush() + except OSError: + pass + _PROGRESS_LINE_ACTIVE = False + _safe_print(f"{prefix}{color_fn(lines[0])}") + continuation_prefix = " " * plain_prefix_width + for line in lines[1:]: + _safe_print(f"{continuation_prefix}{color_fn(line)}") def _progress(label: str) -> None: """Print an in-place progress bar aligned to the step column layout.""" - global _STEP + global _STEP, _PROGRESS_LINE_ACTIVE _STEP += 1 if VERBOSE: return @@ -1537,6 +1563,7 @@ def _progress(label: str) -> None: try: sys.stdout.write(f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}") sys.stdout.flush() + _PROGRESS_LINE_ACTIVE = end == "" except OSError: pass From 420799b61ef35d6cfd87c4f4b02c98152fdf6599 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 19 Jun 2026 09:14:58 -0300 Subject: [PATCH 05/16] Studio: add an Open button to reveal the models folder in the file manager (#6452) * Studio: add an Open button to reveal the models folder in the file manager * Studio: report models folder creation failures --------- Co-authored-by: Wasim Yousef Said --- studio/backend/hub/routes/inventory.py | 6 ++ studio/backend/hub/schemas/inventory.py | 10 +++ .../hub/services/models/local_inventory.py | 25 +++++++ .../backend/hub/tests/test_model_services.py | 49 +++++++++++++ .../hub/catalog/on-device-folders-dialog.tsx | 34 +++++++++ .../src/features/native-intents/api.ts | 6 ++ .../features/settings/api/models-folder.ts | 40 +++++++++++ .../features/settings/tabs/general-tab.tsx | 70 +++++++++++++++++++ studio/frontend/src/i18n/locales/en.ts | 11 +++ studio/src-tauri/src/commands.rs | 28 +++++--- studio/src-tauri/src/main.rs | 1 + studio/src-tauri/src/native_intents.rs | 2 +- 12 files changed, 273 insertions(+), 9 deletions(-) create mode 100644 studio/frontend/src/features/settings/api/models-folder.ts diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py index fcfdd0ad14..4b6c179a2b 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -29,6 +29,7 @@ from hub.schemas.inventory import ( DeleteCachedModelResponse, GgufVariantsResponse, LocalModelListResponse, + ModelsFolderResponse, RecommendedFoldersResponse, RemoveScanFolderResponse, ScanFolderInfo, @@ -91,6 +92,11 @@ def browse_folders( return folder_browser.browse_folders_response(path, show_hidden) +@router.get("/models-folder", response_model = ModelsFolderResponse) +def get_models_folder(current_subject: str = Depends(get_current_subject)): + return local_inventory.get_models_folder_response() + + @router.get("/gguf-variants", response_model = GgufVariantsResponse) async def get_gguf_variants( repo_id: str = Query( diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index c333c7ca89..44ff545e76 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -284,3 +284,13 @@ class BrowseFoldersResponse(BaseModel): "they contain only files, no subdirectories)." ), ) + + +class ModelsFolderResponse(BaseModel): + """The directory where downloaded models are stored (the active HF hub + cache, honoring ``HF_HOME`` / ``HF_HUB_CACHE``).""" + + path: str = Field( + ..., + description = "Path to the model download directory.", + ) diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py index 94e0913ac9..a3782efead 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -670,6 +670,31 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel ) +def get_models_folder_response() -> dict: + """Return the directory where downloaded models are stored. + + This is the active HF hub cache (honors ``HF_HOME`` / ``HF_HUB_CACHE``); + the desktop app reveals it in the OS file manager. + """ + path = _resolve_hf_cache_dir() + # Create it if missing so "Open folder" works before the first download: + # HF builds the cache lazily, and studio only pre-creates the *default* + # dir, not a user's explicit HF_HOME / HF_HUB_CACHE. + try: + path.mkdir(parents = True, exist_ok = True) + except OSError as e: + raise HTTPException( + status_code = 500, + detail = f"Failed to create models folder: {path}: {e}", + ) from e + if not path.is_dir(): + raise HTTPException( + status_code = 500, + detail = f"Models folder path is not a directory: {path}", + ) + return {"path": str(path)} + + def get_scan_folders_response() -> dict: return {"folders": list_scan_folders()} diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index df97b9cf97..1eb7042e4e 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -181,6 +181,55 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): assert ".ssh" not in names +def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): + # The endpoint creates the cache dir on demand so the desktop "Open folder" + # action works even before the first download. + target = tmp_path / "hub" + monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target) + + response = local_inventory.get_models_folder_response() + + assert response == {"path": str(target)} + assert target.is_dir() + + +def test_get_models_folder_response_reports_create_failure(monkeypatch, tmp_path): + target = tmp_path / "hub" + target.write_text("not a directory") + monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target) + + with pytest.raises(HTTPException) as exc_info: + local_inventory.get_models_folder_response() + + assert exc_info.value.status_code == 500 + assert "Failed to create models folder" in exc_info.value.detail + + +def test_get_models_folder_response_requires_directory(monkeypatch, tmp_path): + class MissingPath: + def __init__(self, value: Path): + self.value = value + + def mkdir(self, *, parents: bool, exist_ok: bool): + assert parents is True + assert exist_ok is True + + def is_dir(self): + return False + + def __str__(self): + return str(self.value) + + target = MissingPath(tmp_path / "hub") + monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target) + + with pytest.raises(HTTPException) as exc_info: + local_inventory.get_models_folder_response() + + assert exc_info.value.status_code == 500 + assert "not a directory" in exc_info.value.detail + + def test_contained_link_path_confines_to_link_dir(tmp_path): link_dir = tmp_path / "ollama" / ".studio_links" / "abc123" diff --git a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx index 1769f8861a..20a4eca4dd 100644 --- a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx +++ b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx @@ -23,11 +23,14 @@ import { listScanFolders, removeScanFolder, } from "@/features/hub/inventory"; +import { openModelsDir } from "@/features/native-intents/api"; +import { isTauri } from "@/lib/api-base"; import { cn } from "@/lib/utils"; import { Delete02Icon, FileSearchIcon, FolderAddIcon, + FolderExportIcon, FolderOpenIcon, FolderSearchIcon, PlusSignIcon, @@ -138,6 +141,16 @@ export function OnDeviceFoldersDialog({ [handleInventoryChanged, pending], ); + // Scan folders are arbitrary paths that may be moved or deleted after they + // were registered, so surface the command's failure as a toast. + const handleOpen = useCallback(async (folder: ScanFolderInfo) => { + try { + await openModelsDir(folder.path); + } catch (err) { + toast.error("Couldn't open location", { description: formatError(err) }); + } + }, []); + const handleRemove = useCallback( async (folder: ScanFolderInfo) => { const key = `remove:${folder.id}` as const; @@ -333,6 +346,27 @@ export function OnDeviceFoldersDialog({ + {isTauri ? ( + + + + + + Open in file manager + + + ) : null} + + + + ) : null} + ) -> Vec { } } -/// Open the Unsloth Studio directory in the system file manager. -#[tauri::command] -pub fn open_logs_dir() -> Result<(), String> { - let home = dirs::home_dir().ok_or("Could not determine home directory")?; - let dir = home.join(".unsloth").join("studio"); - - if !dir.exists() { +/// Open an existing directory in the system file manager. Validates the path +/// up front so callers get a clean error instead of a raw OS failure. +fn open_existing_dir(dir: &std::path::Path) -> Result<(), String> { + if !dir.is_dir() { return Err(format!("Directory does not exist: {}", dir.display())); } + open::that(dir).map_err(|e| format!("Failed to open directory: {}", e)) +} - open::that(&dir).map_err(|e| format!("Failed to open directory: {}", e)) +/// Open the Unsloth Studio directory in the system file manager. +#[tauri::command] +pub fn open_logs_dir(window: tauri::WebviewWindow) -> Result<(), String> { + crate::native_intents::ensure_main_window(&window)?; + let home = dirs::home_dir().ok_or("Could not determine home directory")?; + open_existing_dir(&home.join(".unsloth").join("studio")) +} + +/// Open a models directory (resolved by the backend, e.g. the HF cache) in the +/// system file manager. +#[tauri::command] +pub fn open_models_dir(window: tauri::WebviewWindow, path: String) -> Result<(), String> { + crate::native_intents::ensure_main_window(&window)?; + open_existing_dir(std::path::Path::new(&path)) } /// Start the first-launch installation process. diff --git a/studio/src-tauri/src/main.rs b/studio/src-tauri/src/main.rs index 498cd81579..4ed12051ed 100644 --- a/studio/src-tauri/src/main.rs +++ b/studio/src-tauri/src/main.rs @@ -204,6 +204,7 @@ fn main() { commands::check_health, commands::get_server_logs, commands::open_logs_dir, + commands::open_models_dir, commands::start_backend_update, commands::start_managed_repair, commands::cancel_pending_elevation, diff --git a/studio/src-tauri/src/native_intents.rs b/studio/src-tauri/src/native_intents.rs index 30e11aa310..dccbba7083 100644 --- a/studio/src-tauri/src/native_intents.rs +++ b/studio/src-tauri/src/native_intents.rs @@ -258,7 +258,7 @@ fn prune_expired(inner: &mut NativeIntakeInner) { .retain(|intent| intent.path.expires_at_ms > now); } -fn ensure_main_window(window: &WebviewWindow) -> Result<(), String> { +pub(crate) fn ensure_main_window(window: &WebviewWindow) -> Result<(), String> { if window.label() == "main" { Ok(()) } else { From 76a2b9edf160d68208dc30c02c6523bc6551f950 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 19 Jun 2026 05:40:16 -0700 Subject: [PATCH 06/16] Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable (#6468) * Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable Studio's Auto speculative mode promotes any embedded-MTP model >=3B to --spec-type draft-mtp. For MLA models (GLM-5.2/DeepSeek/Kimi) that is a regression: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV context and recomputes the sparse-attention indexer every draft step, so it runs ~2x slower than no speculation (GLM-5.2 UD-IQ1_S bench: 27 vs 45 tok/s, flat across draft depth 1..6 and 96-100% acceptance, on both prose and code). vLLM/SGLang get a speedup from the same model, so this is a llama.cpp implementation gap, not a model property. Auto now drops embedded MTP for MLA models and falls back to ngram-mod (or spec-off when the binary lacks ngram-mod), mirroring the existing sub-3B fallback. The metadata separator is kv_lora_rank: it is present on MLA models and absent on non-MLA embedded-MTP models (Qwen3.x-MTP), whose MTP module is structurally identical but fast, so a "full layer" heuristic cannot tell them apart. Qwen MTP, separate drafters (Gemma, --model-draft), and non-MTP models are unchanged. Explicit overrides still engage the slower MTP route: choosing MTP / MTP+Ngram in Settings, or passing --spec-type in extra args. UNSLOTH_MLA_MTP_ENABLED=1 re-enables Auto promotion for MLA once the upstream path is optimized. A new spec_fallback_reason value "mla_mtp_disabled" surfaces this as an Auto-mode policy downgrade (not a binary/update problem), with a settings banner that points users at the MTP override. It is deliberately kept out of the "Update llama.cpp" affordance since updating does not help. Tests: resolver-matrix rows for MLA->ngram-mod / MLA-no-ngram->off / non-MLA-Qwen->draft-mtp / MLA-separate-drafter->draft-mtp / non-MTP-MLA->default / forced mtp|mtp+ngram on MLA->draft-mtp / env flag; kv_lora_rank metadata fixtures; and reload-skip coverage (Auto ngram-mod is idempotent, forced mtp bounces a reload). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 53 +++- studio/backend/models/inference.py | 8 +- .../tests/test_llama_cpp_mtp_detection.py | 277 ++++++++++++++++++ .../src/features/chat/chat-settings-sheet.tsx | 4 +- .../frontend/src/features/chat/types/api.ts | 5 +- 5 files changed, 342 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 4826403bbc..4b4fe1d7cc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -851,6 +851,23 @@ def _auto_mode_drops_mtp( return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B +def _mla_mtp_auto_enabled() -> bool: + """Whether Auto may pick embedded MTP for an MLA model (GLM-5.2/DeepSeek/Kimi). + + Off by default: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV + context and recomputes the sparse-attention indexer every draft step, so it runs + ~2x slower than no speculation (GLM-5.2 bench: 27 vs 45 tok/s, flat across draft + depth and 96-100% acceptance) -- the opposite of the vLLM/SGLang speedup on the + same model. Set UNSLOTH_MLA_MTP_ENABLED=1 to let Auto promote MLA MTP again once + that path is optimized upstream. Forced mtp / mtp+ngram ignore this gate.""" + return os.environ.get("UNSLOTH_MLA_MTP_ENABLED", "0").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: """User passed --spec-type / --spec-default? llama-server takes one --spec-type (comma-separated to chain), so suppress auto-emit.""" @@ -6169,6 +6186,17 @@ class LlamaCppBackend: _mtp_too_small = ( _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # Embedded MTP head on an MLA model (GLM-5.2/DeepSeek/Kimi, detected by + # kv_lora_rank): llama.cpp's MLA/DSA MTP path is ~2x slower than no spec, + # so Auto drops it (override via the Settings dropdown / forced mtp, or + # UNSLOTH_MLA_MTP_ENABLED=1). Separate drafters (Gemma, mtp_draft_path) and + # non-MLA embedded heads (Qwen, no kv_lora_rank) are unaffected. + _auto_mla_embedded_mtp = ( + bool(self._nextn_predict_layers) + and self._kv_lora_rank is not None + and not bool(mtp_draft_path) + and not _mla_mtp_auto_enabled() + ) if user_owns_spec_type: # User --spec-type wins outright; suppress auto-emit to avoid a @@ -6312,7 +6340,30 @@ class LlamaCppBackend: # effective_mode == "auto": the promotion path. llama.cpp #22673: # MTP is compatible with mmproj, so there's no vision gate. - if is_mtp_model and not _mtp_too_small: + if _auto_mla_embedded_mtp: + # MLA embedded-MTP (GLM-5.2 et al.): the MTP path regresses vs spec-off + # on llama.cpp today, so Auto drops it and falls back to ngram-mod (or + # spec-off if unsupported), mirroring the sub-3B branch. Forced mtp / + # mtp+ngram (handled above) still engage; UNSLOTH_MLA_MTP_ENABLED=1 + # re-enables this promotion once upstream optimizes the path. + self._spec_fallback_reason = "mla_mtp_disabled" + _mla_caps = self.probe_server_capabilities(binary) + if _mla_caps.get("supports_ngram_mod"): + logger.info( + "Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA " + "MTP path is slower than no speculation, so using ngram-mod " + "instead. Override via the Studio Speculative Decoding " + "dropdown or UNSLOTH_MLA_MTP_ENABLED=1." + ) + _emit_ngram_mod() + else: + logger.info( + "Auto: MLA embedded-MTP model detected; disabling speculative " + "decoding (this llama-server does not advertise ngram-mod). " + "Override via the dropdown or UNSLOTH_MLA_MTP_ENABLED=1." + ) + # spec-off: emit nothing, mirroring the sub-3B no-ngram path. + elif is_mtp_model and not _mtp_too_small: # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. _emit_mtp(chain_ngram = not gpus) elif is_mtp_model and _mtp_too_small: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index c9b10fcfc2..0d33cfa976 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -412,8 +412,12 @@ class InferenceStatusResponse(BaseModel): "(auto on an MTP model, or forced mtp / mtp+ngram). " "'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would " "re-enable it (show the update affordance); 'runtime_error' -> the " - "current build could not run it. None when MTP engaged or was not " - "requested." + "current build could not run it. 'mla_mtp_disabled' -> an Auto-mode " + "policy downgrade: the model is MLA (GLM-5.2 et al.) whose llama.cpp " + "MTP path runs slower than no speculation, so Auto used ngram-mod or " + "spec-off instead -- updating won't help; choose MTP in Settings (or " + "set UNSLOTH_MLA_MTP_ENABLED=1) to force it. None when MTP engaged or " + "was not requested." ), ) llama_cpp_prebuilt_stale: bool = Field( diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 063a9ce2cd..a81a49acd5 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -62,6 +62,7 @@ from core.inference.llama_cpp import ( _extra_args_set_any_flag, _extra_args_set_spec_type, _is_mtp_model_name, + _mla_mtp_auto_enabled, ) @@ -1329,6 +1330,282 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch): assert backend.requested_spec_mode == "mtp+ngram" +# ── Auto drops embedded MTP for MLA models (GLM-5.2 et al.) ─────────── +# +# llama.cpp's MLA/DSA MTP path runs ~2x slower than no speculation (GLM-5.2 +# bench), so Auto downgrades it to ngram-mod (or spec-off). The clean +# metadata separator from non-MLA MTP (Qwen, kept on draft-mtp) is +# self._kv_lora_rank. Forced mtp / mtp+ngram and separate drafters (Gemma) +# stay on draft-mtp; UNSLOTH_MLA_MTP_ENABLED=1 re-enables Auto promotion. + +# GLM-5.2's repo name has no "MTP" marker, so its MTP signal is metadata-only +# (nextn_predict_layers) -- exactly the embedded-MLA case we gate. +_GLM_MLA_MODEL = "unsloth/GLM-5.2-GGUF" + + +def _mla_resolver_backend( + monkeypatch, + *, + ngram_supported = True, + kv_lora_rank = 512, + nextn = 1, +): + """Resolver backend posing as an embedded-MTP MLA model (kv_lora_rank set).""" + backend = _resolver_backend(monkeypatch, ngram_supported = ngram_supported) + backend._nextn_predict_layers = nextn + backend._kv_lora_rank = kv_lora_rank + return backend + + +@pytest.mark.parametrize("gpus", [True, False]) +def test_auto_mla_embedded_mtp_falls_back_to_ngram(monkeypatch, gpus): + # Auto + MLA embedded MTP + ngram supported -> ngram-mod on BOTH platforms + # (the CPU chain ngram-mod,draft-mtp is dropped: no draft-mtp for MLA). + backend = _mla_resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _GLM_MLA_MODEL, + model_path = None, + gpus = gpus, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--spec-draft-n-max" not in parsed + assert "--spec-ngram-mod-n-match" in parsed + assert backend.speculative_type == "ngram-mod" + assert backend.requested_spec_mode == "auto" + assert backend.spec_fallback_reason == "mla_mtp_disabled" + assert backend.spec_draft_n_max is None + + +def test_auto_mla_embedded_mtp_no_ngram_disables_spec(monkeypatch): + # Auto + MLA embedded MTP + no ngram-mod support -> emit nothing (spec-off), + # mirroring the sub-3B no-ngram path. Still flagged as a policy downgrade. + backend = _mla_resolver_backend(monkeypatch, ngram_supported = False) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _GLM_MLA_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert "--spec-type" not in flags + assert backend.speculative_type is None + assert backend.requested_spec_mode == "auto" + assert backend.spec_fallback_reason == "mla_mtp_disabled" + + +def test_auto_non_mla_embedded_mtp_keeps_draft_mtp(monkeypatch): + # Auto + embedded MTP + NON-MLA (kv_lora_rank None, e.g. Qwen) -> unchanged: + # still draft-mtp at the platform default. No policy downgrade. + backend = _mla_resolver_backend(monkeypatch, kv_lora_rank = None) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--spec-draft-n-max") == "2" + assert backend.speculative_type == "draft-mtp" + assert backend.spec_fallback_reason is None + + +def test_auto_mla_separate_drafter_keeps_mtp(monkeypatch): + # Auto + MLA + a separate drafter (mtp_draft_path) -> the drafter exemption + # wins over the MLA gate: still draft-mtp (Gemma-style external drafter is + # not the slow embedded MLA/DSA path). + backend = _mla_resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _GLM_MLA_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = "/fake/mtp-draft.gguf", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert backend.speculative_type == "draft-mtp" + assert backend.spec_fallback_reason is None + + +def test_auto_non_mtp_mla_model_unaffected(monkeypatch): + # Auto + MLA but NO embedded MTP head (kv_lora_rank set, nextn None, e.g. + # GLM-4.7-Flash) -> non-MTP default; no accidental ngram drop. + backend = _mla_resolver_backend(monkeypatch, nextn = None) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/GLM-4.7-Flash-GGUF", + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert "--spec-default" in flags + assert "ngram-mod" not in flags + assert backend.speculative_type == "default" + assert backend.spec_fallback_reason is None + + +@pytest.mark.parametrize( + "mode, expect_spec_type, expect_n_max", + [ + ("mtp", "draft-mtp", "2"), + ("mtp+ngram", "ngram-mod,draft-mtp", "2"), + ], +) +def test_forced_mtp_on_mla_still_engages(monkeypatch, mode, expect_spec_type, expect_n_max): + # Explicit override engages the deliberately-slower MTP route on MLA models, + # regardless of the Auto gate. No policy downgrade reason. + backend = _mla_resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = _GLM_MLA_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == expect_spec_type + assert parsed.get("--spec-draft-n-max") == expect_n_max + assert backend.speculative_type == "draft-mtp" + assert backend.requested_spec_mode == mode + assert backend.spec_fallback_reason is None + + +def test_env_flag_reenables_auto_mla_mtp(monkeypatch): + # UNSLOTH_MLA_MTP_ENABLED=1 -> Auto promotes MLA embedded MTP to draft-mtp + # again (the forward hook for when llama.cpp optimizes the path). + monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", "1") + backend = _mla_resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _GLM_MLA_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert backend.speculative_type == "draft-mtp" + assert backend.spec_fallback_reason is None + + +@pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", "On"]) +def test_mla_mtp_auto_enabled_truthy_values(monkeypatch, value): + monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", value) + assert _mla_mtp_auto_enabled() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "off", "", " ", "bogus"]) +def test_mla_mtp_auto_disabled_default_and_falsy(monkeypatch, value): + monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", value) + assert _mla_mtp_auto_enabled() is False + + +def test_mla_mtp_auto_disabled_when_unset(monkeypatch): + monkeypatch.delenv("UNSLOTH_MLA_MTP_ENABLED", raising = False) + assert _mla_mtp_auto_enabled() is False + + +def test_read_gguf_metadata_captures_kv_lora_rank(tmp_path): + # GLM-5.2-style header: MLA (kv_lora_rank) + embedded MTP (nextn) populate + # both fields, so the Auto gate sees an MLA embedded-MTP model. + gguf = _write_minimal_gguf( + tmp_path / "model.gguf", + arch = "glm-dsa", + nextn = 1, + extra_uint32 = { + "glm-dsa.block_count": 4, + "glm-dsa.attention.kv_lora_rank": 512, + }, + ) + backend = LlamaCppBackend() + backend._read_gguf_metadata(str(gguf)) + assert backend._nextn_predict_layers == 1 + assert backend._kv_lora_rank == 512 + + +def test_read_gguf_metadata_qwen_mtp_has_no_kv_lora_rank(tmp_path): + # Qwen MTP header: embedded MTP but non-MLA, so kv_lora_rank stays None and + # Auto keeps it on draft-mtp. + gguf = _write_minimal_gguf( + tmp_path / "model.gguf", + arch = "qwen35moe", + nextn = 1, + extra_uint32 = {"qwen35moe.block_count": 4}, + ) + backend = LlamaCppBackend() + backend._read_gguf_metadata(str(gguf)) + assert backend._nextn_predict_layers == 1 + assert backend._kv_lora_rank is None + + +def test_reload_skip_auto_mla_ngram_is_idempotent(): + # A GLM model resolved to ngram-mod under Auto must not churn: a duplicate + # Auto /load at the same settings is already-satisfied. + backend = _mtp_backend( + _model_identifier = _GLM_MLA_MODEL, + _speculative_type = "ngram-mod", + _requested_spec_mode = "auto", + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = _GLM_MLA_MODEL, + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_reload_forced_mtp_bounces_auto_mla(): + # Overriding Auto (ngram-mod) with a forced mtp request must reload (to the + # slower draft-mtp route), not dedup against the running ngram-mod server. + backend = _mtp_backend( + _model_identifier = _GLM_MLA_MODEL, + _speculative_type = "ngram-mod", + _requested_spec_mode = "auto", + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = _GLM_MLA_MODEL, + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "mtp", + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + # ── Full named-repo resolver matrix (the shipping Studio families) ───── # # Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index a4115d58bb..935f7e1bc6 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1005,7 +1005,9 @@ export function ChatSettingsPanel({ speculativeType === "mtp+ngram") && (

- {specFallbackReason === "runtime_error" + {specFallbackReason === "mla_mtp_disabled" + ? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Select MTP above to force it." + : specFallbackReason === "runtime_error" ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding." : "MTP is not available in the installed llama.cpp build, so this model is running without it." + (llamaUpdateStatus?.update_available diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 4ac70486e5..d9b69a1ef6 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -196,7 +196,10 @@ export interface InferenceStatusResponse { /** * Why MTP was disabled on the loaded model despite being requested. * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable - * it; "runtime_error" -> the current build could not run it. Null otherwise. + * it; "runtime_error" -> the current build could not run it; + * "mla_mtp_disabled" -> an Auto-mode policy downgrade for MLA models + * (GLM-5.2 et al.) whose llama.cpp MTP path is slower than no speculation + * (updating won't help; choose MTP in Settings to force it). Null otherwise. */ spec_fallback_reason?: string | null; } From 52877bba05c48d82582e1874b8da242a22135596 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 19 Jun 2026 05:51:35 -0700 Subject: [PATCH 07/16] Studio: gate the MTP target-KV reserve to MTP spec mode, not just MLA (#6449) _estimate_mtp_overhead_bytes is also reached for the separate-drafter spec modes (draft-simple / draft-eagle3) through _user_draft_via_extras. Those modes load a small distinct drafter with its own KV -- already counted in the draft KV + weights -- and keep no duplicated full target context; only MTP runs a second context over the target model's own KV geometry (llama.cpp ctx_tgt). Charging the ~main-KV-sized f16 copy there over-reserved by tens of GiB on an MLA model and needlessly shrank the advertised context, the same under-advertising #6312 set out to fix. Thread mtp_keeps_target_ctx through _estimate_mtp_overhead_bytes (True for MTP, False for separate-drafter modes) and derive _engaged_is_mtp at the fit call site so the target copy is added only when the engaged mode is actually MTP. MLA + MTP (GLM-5.2 / DeepSeek / Kimi) is unchanged, so the GLM-5.2 OOM fix is preserved; non-MLA and the draft-simple / draft-eagle3 paths no longer pay the copy. test_mtp_mla_target_ctx.py adds a case asserting the separate-drafter reserve collapses to the draft KV (no target copy) while the default MTP path keeps it. --- studio/backend/core/inference/llama_cpp.py | 68 ++++++++++++------- .../backend/tests/test_mtp_mla_target_ctx.py | 18 ++++- 2 files changed, 59 insertions(+), 27 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 4b4fe1d7cc..16b312d8f6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2840,12 +2840,16 @@ class LlamaCppBackend: drafter_path: Optional[str] = None, draft_weights_bytes: int = 0, n_parallel: int = 1, + mtp_keeps_target_ctx: bool = True, ) -> Optional[int]: """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- - drafter weights + (MLA only) a duplicated target KV context. The verify - buffer rides in the ctx-fit headroom (no tuned constant). None when the - draft KV can't be sized (caller keeps the flat fallback). - ``draft_weights_bytes`` is the drafter file size (0 for embedded).""" + drafter weights + (MTP + MLA only) a duplicated target KV context. The + verify buffer rides in the ctx-fit headroom (no tuned constant). None when + the draft KV can't be sized (caller keeps the flat fallback). + ``draft_weights_bytes`` is the drafter file size (0 for embedded). + ``mtp_keeps_target_ctx`` is True for MTP draft modes (which keep the + duplicated target context) and False for separate-drafter spec modes + (draft-simple/draft-eagle3), which do not.""" draft_kv = self._mtp_draft_kv_bytes( n_ctx, drafter_path = drafter_path, @@ -2854,16 +2858,19 @@ class LlamaCppBackend: n_parallel = n_parallel, ) weights = max(0, draft_weights_bytes) - # MLA models (GLM-5.x, DeepSeek, Kimi-K2) keep a *second* full copy of the - # target model's KV context for MTP draft verification -- llama.cpp's + # MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy + # of the target model's KV context for draft verification -- llama.cpp's # `ctx_tgt=yes` -- allocated at f16 regardless of the main cache type. It is # ~the main KV again and dwarfs the embedded draft head (GLM-5.2 @ 1M ctx: # a ~2 GiB head next to a ~89 GiB target copy), so omitting it lets auto-fit # pick a context that fits on paper but OOMs cublasCreate at the first - # decode. Non-MLA MTP (Qwen/Gemma) keeps no such copy, so this is gated - # strictly on MLA (kv_lora_rank present) and leaves those models unchanged. + # decode. Gated on both MLA (kv_lora_rank present) and the engaged mode + # actually being MTP: non-MLA MTP (Qwen/Gemma) keeps no such copy, and the + # separate-drafter spec modes (draft-simple/draft-eagle3) load a small + # distinct drafter with its own KV -- already counted in draft_kv/weights -- + # rather than duplicating the target, so they must not be charged for it. target_ctx_copy = 0 - if self._kv_lora_rank is not None: + if mtp_keeps_target_ctx and self._kv_lora_rank is not None: target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) if draft_kv is None: # KV unsized (exotic/remote drafter): still reserve known weights + any @@ -4835,25 +4842,31 @@ class LlamaCppBackend: except Exception: _mtp_binary_ok = False _mtp_probe_raised = True - _mtp_will_engage = bool( - _user_mtp_via_extras - or _user_draft_via_extras - or ( - not _extra_args_set_spec_type(extra_args) - and _mtp_model_for_fit - and ( - _mtp_effective in ("mtp", "mtp+ngram") - or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) - ) - and ( - _mtp_binary_ok - # Reserve on a raised (uncached) probe too: it re-probes in - # _build_speculative_flags and may still engage MTP (embedded - # head or separate drafter -- _mtp_model_for_fit covers both). - or _mtp_probe_raised - ) + _auto_studio_mtp = ( + not _extra_args_set_spec_type(extra_args) + and _mtp_model_for_fit + and ( + _mtp_effective in ("mtp", "mtp+ngram") + or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) + ) + and ( + _mtp_binary_ok + # Reserve on a raised (uncached) probe too: it re-probes in + # _build_speculative_flags and may still engage MTP (embedded + # head or separate drafter -- _mtp_model_for_fit covers both). + or _mtp_probe_raised ) ) + _mtp_will_engage = bool( + _user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp + ) + # The duplicated full target-KV copy (ctx_tgt) is an MTP-only + # cost: the MTP head runs a second context over the target + # model's own KV geometry. The separate-drafter spec modes + # (draft-simple/draft-eagle3, reached via _user_draft_via_extras) + # load a small distinct drafter with its own KV and keep no such + # copy, so only charge it when the engaged mode is truly MTP. + _engaged_is_mtp = bool(_user_mtp_via_extras or _auto_studio_mtp) # Effective draft depth: extras win (last-wins at launch), else # the field, else the platform default (2 GPU / 3 CPU). @@ -4922,6 +4935,7 @@ class LlamaCppBackend: drafter_path = _mtp_draft_for_budget, draft_weights_bytes = _mtp_draft_weights, n_parallel = n_parallel, + mtp_keeps_target_ctx = _engaged_is_mtp, ) is not None ): @@ -4937,6 +4951,7 @@ class LlamaCppBackend: _dp: Optional[str] = _mtp_draft_for_budget, _w: int = _mtp_draft_weights, _np: int = n_parallel, + _mtp: bool = _engaged_is_mtp, ) -> int: v = self._estimate_mtp_overhead_bytes( ctx, @@ -4946,6 +4961,7 @@ class LlamaCppBackend: drafter_path = _dp, draft_weights_bytes = _w, n_parallel = _np, + mtp_keeps_target_ctx = _mtp, ) return v if v is not None else 0 diff --git a/studio/backend/tests/test_mtp_mla_target_ctx.py b/studio/backend/tests/test_mtp_mla_target_ctx.py index 50866942bf..c38e4c7b62 100644 --- a/studio/backend/tests/test_mtp_mla_target_ctx.py +++ b/studio/backend/tests/test_mtp_mla_target_ctx.py @@ -172,6 +172,22 @@ class TestMlaTargetCtxReserve: ctx = 131072 assert mla._estimate_mtp_overhead_bytes(ctx) > non._estimate_mtp_overhead_bytes(ctx) + def test_separate_drafter_mode_drops_target_copy(self): + # The duplicated target context is MTP-only. draft-simple / draft-eagle3 + # load a small separate drafter with its own KV (counted in the draft KV) + # and keep no target copy, so even on an MLA model the reserve must drop + # the f16 copy when mtp_keeps_target_ctx=False -- which is what the loader + # threads for those modes. The default (True) keeps the MTP copy. + b = _make_mla_backend() + ctx = 262144 + mtp = b._estimate_mtp_overhead_bytes(ctx) # default True == MTP draft + separate = b._estimate_mtp_overhead_bytes(ctx, mtp_keeps_target_ctx = False) + # Separate-drafter overhead is exactly the draft KV (no target copy)... + assert separate == b._mtp_draft_kv_bytes(ctx) + # ...and the MTP reserve is that plus the full f16 target copy. + assert mtp == separate + b._estimate_kv_cache_bytes(ctx, "f16") + assert mtp > separate + class TestMlaFitPreventsOom: """The corrected reserve must actually lower the auto-fit context so the @@ -200,7 +216,7 @@ class TestMlaFitPreventsOom: self.MODEL_BYTES, mtp_engaged = True, total_mib = self.TOTAL_MIB, - mtp_overhead_fn = lambda c: (b._mtp_draft_kv_bytes(c) or 0), + mtp_overhead_fn = lambda c: b._mtp_draft_kv_bytes(c) or 0, ) assert draft_only == self.REQ_CTX # reproduces the over-advertised context assert with_copy < self.REQ_CTX # corrected reserve backs the context off From 17e9714a98a361385511ae2a821899e934638876 Mon Sep 17 00:00:00 2001 From: Parvesh Saini <97528080+parveshsaini@users.noreply.github.com> Date: Fri, 19 Jun 2026 18:49:04 +0530 Subject: [PATCH 08/16] studio: run /generate/stream's sync generator off the event loop to avoid blocking it (#6466) * studio: run /generate/stream's sync generator off the event loop to avoid blocking it * fix: close generator in finally on client disconnect in generate_stream * Fix/adjust generate stream test for PR #6466 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix/adjust generate stream cancellation for PR #6466 * Fix/adjust generate stream cleanup for PR #6466 * fix: cancel incomplete generate stream cleanup --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/routes/inference.py | 29 ++- .../test_stream_cancel_registration_timing.py | 242 ++++++++++++++++++ 2 files changed, 269 insertions(+), 2 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b3f112a944..4b068941b8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3127,9 +3127,13 @@ async def generate_stream( log = logger, ) + cancel_event = threading.Event() + async def stream(): + gen = None + completed = False try: - for chunk in backend.generate_chat_response( + gen = backend.generate_chat_response( messages = request.messages, system_prompt = request.system_prompt, image = image, @@ -3138,14 +3142,35 @@ async def generate_stream( top_k = request.top_k, max_new_tokens = request.max_new_tokens, repetition_penalty = request.repetition_penalty, - ): + cancel_event = cancel_event, + ) + _DONE = object() + while True: + chunk = await asyncio.to_thread(next, gen, _DONE) + if chunk is _DONE: + break yield f"data: {json.dumps({'content': chunk})}\n\n" + completed = True yield "data: [DONE]\n\n" + except asyncio.CancelledError: + cancel_event.set() + backend.reset_generation_state() + raise except Exception as e: + cancel_event.set() backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info = True) yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" + finally: + if not completed and not cancel_event.is_set(): + cancel_event.set() + backend.reset_generation_state() + if gen is not None: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass return StreamingResponse( stream(), diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 8d12fd528e..2e6b5f14da 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -10,6 +10,7 @@ from __future__ import annotations import ast import asyncio +import json import threading import time from pathlib import Path @@ -453,6 +454,147 @@ def test_audio_input_stream_offloads_blocking_next_to_thread(): ) +def test_generate_stream_offloads_blocking_next_to_thread(): + outer = None + for fn in ast.walk(_TREE): + if isinstance(fn, ast.AsyncFunctionDef) and fn.name == "generate_stream": + outer = fn + break + assert outer is not None, "generate_stream handler missing" + + inner = None + for sub in ast.walk(outer): + if isinstance(sub, ast.AsyncFunctionDef) and sub.name == "stream": + inner = sub + break + assert inner is not None, "generate_stream inner stream() generator missing" + + for sub in ast.walk(inner): + if isinstance(sub, (ast.For, ast.AsyncFor)): + it_src = ast.unparse(sub.iter) + assert "generate_chat_response" not in it_src, ( + "generate_stream's inner stream() must not iterate " + "backend.generate_chat_response() directly -- that blocks the event " + "loop on every blocking subprocess read between tokens. Use " + "`await asyncio.to_thread(next, gen, _DONE)` inside a `while True` " + "loop instead" + ) + + found_to_thread_next = False + for sub in ast.walk(inner): + if not isinstance(sub, ast.Call): + continue + fn_expr = sub.func + if not ( + isinstance(fn_expr, ast.Attribute) + and fn_expr.attr == "to_thread" + and isinstance(fn_expr.value, ast.Name) + and fn_expr.value.id == "asyncio" + ): + continue + if sub.args and isinstance(sub.args[0], ast.Name) and sub.args[0].id == "next": + found_to_thread_next = True + break + assert found_to_thread_next, ( + "generate_stream's inner stream() must call " + "`asyncio.to_thread(next, gen, _DONE)` to keep the event loop free while the " + "worker subprocess produces the next token" + ) + + +def test_generate_stream_cancels_backend_on_stream_cancelled_error(): + outer = None + for fn in ast.walk(_TREE): + if isinstance(fn, ast.AsyncFunctionDef) and fn.name == "generate_stream": + outer = fn + break + assert outer is not None, "generate_stream handler missing" + + outer_src = ast.unparse(outer) + assert "cancel_event = threading.Event()" in outer_src + + inner = None + for sub in ast.walk(outer): + if isinstance(sub, ast.AsyncFunctionDef) and sub.name == "stream": + inner = sub + break + assert inner is not None, "generate_stream inner stream() generator missing" + + def _awaits_to_thread_gen_close(node: ast.AST) -> bool: + for sub in ast.walk(node): + if not isinstance(sub, ast.Await): + continue + call = sub.value + if not isinstance(call, ast.Call): + continue + fn_expr = call.func + if not ( + isinstance(fn_expr, ast.Attribute) + and fn_expr.attr == "to_thread" + and isinstance(fn_expr.value, ast.Name) + and fn_expr.value.id == "asyncio" + ): + continue + if not call.args: + continue + close_expr = call.args[0] + if ( + isinstance(close_expr, ast.Attribute) + and close_expr.attr == "close" + and isinstance(close_expr.value, ast.Name) + and close_expr.value.id == "gen" + ): + return True + return False + + found_cancel_kwarg = False + found_cancel_handler = False + found_finally_cleanup = False + for sub in ast.walk(inner): + if isinstance(sub, ast.Call): + call_src = ast.unparse(sub.func) + if call_src.endswith("generate_chat_response"): + found_cancel_kwarg = any( + kw.arg == "cancel_event" + and isinstance(kw.value, ast.Name) + and kw.value.id == "cancel_event" + for kw in sub.keywords + ) + if isinstance(sub, ast.ExceptHandler): + exc_src = ast.unparse(sub.type) if sub.type is not None else "" + if exc_src != "asyncio.CancelledError": + continue + body_src = "\n".join(ast.unparse(stmt) for stmt in sub.body) + found_cancel_handler = ( + "cancel_event.set()" in body_src + and "backend.reset_generation_state()" in body_src + and any(isinstance(stmt, ast.Raise) and stmt.exc is None for stmt in sub.body) + ) + if isinstance(sub, ast.Try) and sub.finalbody: + final_src = "\n".join(ast.unparse(stmt) for stmt in sub.finalbody) + found_finally_cleanup = ( + "not completed" in final_src + and "not cancel_event.is_set()" in final_src + and "cancel_event.set()" in final_src + and "backend.reset_generation_state()" in final_src + and _awaits_to_thread_gen_close(sub) + ) + + assert found_cancel_kwarg, ( + "generate_stream must pass cancel_event into backend.generate_chat_response " + "so cancelled streams can stop backend generation" + ) + assert found_cancel_handler, ( + "generate_stream must catch asyncio.CancelledError, set cancel_event, " + "reset backend state, and re-raise" + ) + assert found_finally_cleanup, ( + "generate_stream cleanup must cancel/reset incomplete streams and " + "offload gen.close() with asyncio.to_thread so backend joins cannot " + "block the event loop" + ) + + def test_stream_chunks_cancel_branch_resets_backend_state(): # The cancel branch must call backend.reset_generation_state() to flush # GPU/KV-cache state, else cancel-via-POST leaves the subprocess dirty. @@ -542,6 +684,106 @@ def test_unsloth_stream_loop_breaks_on_external_cancel_event(): ) +def test_generate_stream_stays_responsive_under_blocking_next(): + # Same sync-generator shape as generate_stream, with resp_queue.get modeled + # by sleep. The output must stay unchanged while next() moves off-loop. + chunks = ["alpha", "beta", "gamma", "delta"] + + def _generate_chat_response(): + for chunk in chunks: + time.sleep(0.08) + yield chunk + + def _sse(chunk): + return f"data: {json.dumps({'content': chunk})}\n\n" + + async def _direct_loop(): + out = [] + for chunk in _generate_chat_response(): + out.append(_sse(chunk)) + out.append("data: [DONE]\n\n") + return out + + async def _to_thread_loop(): + _DONE = object() + gen = _generate_chat_response() + out = [] + try: + while True: + chunk = await asyncio.to_thread(next, gen, _DONE) + if chunk is _DONE: + break + out.append(_sse(chunk)) + out.append("data: [DONE]\n\n") + return out + finally: + try: + gen.close() + except (RuntimeError, ValueError): + pass + + async def _run_with_heartbeat(loop_coro): + ticks = 0 + max_gap = 0.0 + + async def _heartbeat(): + nonlocal ticks, max_gap + last = time.monotonic() + while True: + await asyncio.sleep(0.01) + now = time.monotonic() + max_gap = max(max_gap, now - last) + last = now + ticks += 1 + + heartbeat = asyncio.create_task(_heartbeat()) + await asyncio.sleep(0) + try: + out = await loop_coro() + finally: + heartbeat.cancel() + try: + await heartbeat + except asyncio.CancelledError: + pass + return out, ticks, max_gap + + async def _main(): + direct_out, direct_ticks, direct_max_gap = await _run_with_heartbeat(_direct_loop) + threaded_out, threaded_ticks, threaded_max_gap = await _run_with_heartbeat(_to_thread_loop) + return ( + direct_out, + direct_ticks, + direct_max_gap, + threaded_out, + threaded_ticks, + threaded_max_gap, + ) + + ( + direct_out, + direct_ticks, + direct_max_gap, + threaded_out, + threaded_ticks, + threaded_max_gap, + ) = asyncio.run(_main()) + + assert threaded_out == direct_out == [_sse(chunk) for chunk in chunks] + ["data: [DONE]\n\n"] + assert direct_ticks == 0, ( + f"direct generate_stream loop should block the event loop; " + f"got {direct_ticks} heartbeat ticks and max gap {direct_max_gap:.3f}s" + ) + assert threaded_ticks >= 8, ( + f"to_thread generate_stream loop should let the event loop run; " + f"got {threaded_ticks} heartbeat ticks" + ) + assert threaded_max_gap < 0.06, ( + f"to_thread generate_stream loop should avoid long heartbeat gaps; " + f"got {threaded_max_gap:.3f}s" + ) + + def test_audio_stream_stays_responsive_under_blocking_next(): # Assert the pre-fix `for chunk in audio_input_generate()` pattern blocks the # event loop, then confirm the post-fix pattern exits promptly. From 7ce8dc73ac52cf7a3e87501cce7546c30b66e010 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 19 Jun 2026 11:16:16 -0300 Subject: [PATCH 09/16] Studio: label Apple Silicon as Metal/unified memory instead of CPU-only in installers (#6470) * Studio: label Apple Silicon as Metal/unified memory instead of CPU-only in installers * Studio: drop redundant aarch64 check from the macOS GPU label detection --- install.sh | 3 +++ studio/setup.sh | 3 +++ 2 files changed, 6 insertions(+) diff --git a/install.sh b/install.sh index 7e8df1bc64..18aa922f05 100755 --- a/install.sh +++ b/install.sh @@ -2460,6 +2460,9 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then substep "ROCm: $_rocm_root" [ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver" [ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt" +elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + # Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only. + step "gpu" "Apple Silicon (Metal, unified memory)" else step "gpu" "none (CPU-only)" "$C_WARN" fi diff --git a/studio/setup.sh b/studio/setup.sh index b78029d6e2..75a7fce31e 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -968,6 +968,9 @@ elif [ "$_setup_amd_detected" = true ]; then substep "ROCm: $_setup_rocm_root" [ -n "$_setup_rocm_ver" ] && substep "hipconfig: $_setup_rocm_ver" [ -n "$_setup_mkt" ] && [ -n "$_setup_gfx" ] && substep "GPU: $_setup_mkt" +elif [ "$(uname -s 2>/dev/null)" = "Darwin" ] && [ "$(uname -m 2>/dev/null)" = "arm64" ]; then + # Apple Silicon: llama.cpp builds with Metal over unified memory, so not a CPU-only host. + step "gpu" "Apple Silicon (Metal, unified memory)" else step "gpu" "none (chat-only / GGUF)" "$C_WARN" substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." From e1ae4756d9afe7fc5ba34e939fe5ba75dfb6e94e Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 19 Jun 2026 22:52:27 +0100 Subject: [PATCH 10/16] Studio: rework prompt queue management UI/UX (#6467) * feat: add prompt queue management UI * fix: clean up prompt queue controls * fix: address prompt queue review feedback * fix: scope prompt queue controls to active thread * fix: tighten prompt queue row behavior --- .../src/components/assistant-ui/thread.tsx | 409 ++++++++++++++++-- 1 file changed, 371 insertions(+), 38 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 8789a1bb71..57b9f6c965 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -135,6 +135,7 @@ import { ChevronLeftIcon, ChevronRightIcon, Columns2Icon, + CornerDownRightIcon, GitBranchIcon, GlobeIcon, HeadphonesIcon, @@ -178,12 +179,33 @@ type PromptQueueUIEntry = { total: number; }; +type PromptQueueUIItemStatus = "queued" | "next" | "waiting" | "running"; + +type PromptQueueUIItem = { + id: string; + prompt: string; + position: number; + total: number; + status: PromptQueueUIItemStatus; + threadIds: string[]; + canEdit: boolean; + canRemove: boolean; +}; + interface PromptQueueUIState { byThreadId: Record; + current: number; + total: number; + items: PromptQueueUIItem[]; + isRunning: boolean; } const usePromptQueueUI = create(() => ({ byThreadId: {}, + current: 0, + total: 0, + items: [], + isRunning: false, })); type PromptQueueTarget = { @@ -195,8 +217,10 @@ type PromptQueueTarget = { }; type PromptQueueItem = { + id: string; prompt: string; target: PromptQueueTarget; + dispatched: boolean; }; const PROMPT_QUEUE_INDEXING_RETRY_MS = 500; @@ -214,6 +238,10 @@ function compactIds(ids: Array) { return Array.from(new Set(ids.filter((id): id is string => Boolean(id)))); } +function createPromptQueueItemId() { + return `prompt-queue-${crypto.randomUUID()}`; +} + function stopPromptQueueSubscription({ resetRunningState = true, }: { @@ -228,7 +256,7 @@ function stopPromptQueueSubscription({ } } -function resetPromptQueue(showToast = false) { +function resetPromptQueue() { promptQueueGeneration += 1; promptQueueIsRunning = false; promptQueueItems = []; @@ -240,16 +268,10 @@ function resetPromptQueue(showToast = false) { } stopPromptQueueSubscription(); syncPromptQueueUI(); - if (showToast) { - toast.success("Prompt queue complete"); - } -} - -function queueToastDescription(prompt: string) { - return prompt.length > 80 ? `${prompt.slice(0, 80)}...` : prompt; } function appendQueuedPrompt(item: PromptQueueItem) { + item.dispatched = true; syncPromptQueueUI(); item.target.append(item.prompt); } @@ -335,8 +357,10 @@ async function dispatchQueuedPrompt( function createQueuedPrompt(prompt: string, target: PromptQueueTarget) { return { + id: createPromptQueueItemId(), prompt, target, + dispatched: false, }; } @@ -368,13 +392,62 @@ function findPromptQueueEntry( return null; } +function canEditPromptQueueItem(item: PromptQueueItem) { + return !item.dispatched; +} + +function canRemovePromptQueueItem(item: PromptQueueItem) { + return !item.dispatched; +} + +function promptQueueItemMatchesThreadIds( + item: PromptQueueUIItem, + threadIds: string[], +) { + return item.threadIds.some((threadId) => threadIds.includes(threadId)); +} + function syncPromptQueueUI() { if (!promptQueueIsRunning || promptQueueItems.length === 0) { - usePromptQueueUI.setState({ byThreadId: {} }); + usePromptQueueUI.setState({ + byThreadId: {}, + current: 0, + total: 0, + items: [], + isRunning: false, + }); return; } const activeItemIndex = Math.max(promptQueueIndex, 0); + const total = promptQueueItems.length; + const current = promptQueueIndex >= 0 ? Math.min(activeItemIndex + 1, total) : 0; + const items = promptQueueItems + .map((item, index): PromptQueueUIItem | null => { + if (index < activeItemIndex || item.dispatched) { + return null; + } + const threadIds = getPromptQueueTargetIds(item.target); + const isActive = promptQueueIndex >= 0 && index === activeItemIndex; + const status: PromptQueueUIItemStatus = item.dispatched + ? "running" + : isActive + ? promptQueueWaitingForTargetIdle + ? "waiting" + : "next" + : "queued"; + return { + id: item.id, + prompt: item.prompt, + position: index + 1, + total, + status, + threadIds, + canEdit: canEditPromptQueueItem(item), + canRemove: canRemovePromptQueueItem(item), + }; + }) + .filter((item): item is PromptQueueUIItem => Boolean(item)); const groups: Array<{ ids: Set; current: number; @@ -423,7 +496,80 @@ function syncPromptQueueUI() { }); } - usePromptQueueUI.setState({ byThreadId }); + usePromptQueueUI.setState({ + byThreadId, + current, + total, + items, + isRunning: true, + }); +} + +function editPromptQueueItem(itemId: string, prompt: string) { + const nextPrompt = prompt.trim(); + if (!nextPrompt) { + return false; + } + const itemIndex = promptQueueItems.findIndex( + (candidate) => candidate.id === itemId, + ); + if (itemIndex < 0) { + return false; + } + const item = promptQueueItems[itemIndex]; + if (!canEditPromptQueueItem(item)) { + return false; + } + item.prompt = nextPrompt; + syncPromptQueueUI(); + return true; +} + +function clearPromptQueueRetryTimer() { + if (!promptQueueRetryTimer) { + return; + } + clearTimeout(promptQueueRetryTimer); + promptQueueRetryTimer = null; +} + +function removePromptQueueItem(itemId: string) { + const itemIndex = promptQueueItems.findIndex((item) => item.id === itemId); + if (itemIndex < 0) { + return false; + } + const item = promptQueueItems[itemIndex]; + if (!canRemovePromptQueueItem(item)) { + return false; + } + + const wasActive = + promptQueueIndex >= 0 && itemIndex === Math.max(promptQueueIndex, 0); + promptQueueItems.splice(itemIndex, 1); + if (promptQueueItems.length === 0) { + resetPromptQueue(); + return true; + } + + if (itemIndex < promptQueueIndex) { + promptQueueIndex -= 1; + } + if (wasActive && promptQueueIndex >= promptQueueItems.length) { + resetPromptQueue(); + return true; + } + + syncPromptQueueUI(); + if (wasActive) { + clearPromptQueueRetryTimer(); + promptQueueWaitingForTargetIdle = false; + promptQueuePrevStoreRunning = false; + const next = promptQueueItems[promptQueueIndex]; + if (next) { + scheduleQueuedPromptDispatch(next, 50); + } + } + return true; } function isPromptQueueTargetRunning( @@ -456,15 +602,12 @@ function isActivePromptQueueTargetRunning( function advancePromptQueue() { const nextIndex = promptQueueIndex + 1; if (nextIndex >= promptQueueItems.length) { - resetPromptQueue(true); + resetPromptQueue(); return; } promptQueueIndex = nextIndex; syncPromptQueueUI(); const next = promptQueueItems[nextIndex]; - toast(`Prompt ${nextIndex + 1} / ${promptQueueItems.length}`, { - description: queueToastDescription(next.prompt), - }); promptQueueWaitingForTargetIdle = false; promptQueuePrevStoreRunning = false; scheduleQueuedPromptDispatch(next, 100); @@ -529,9 +672,6 @@ function startPromptQueue( ...filtered.map((prompt) => createQueuedPrompt(prompt, target)), ); syncPromptQueueUI(); - toast.success("Added to prompt queue", { - description: `${filtered.length} prompt${filtered.length === 1 ? "" : "s"} queued.`, - }); return; } @@ -547,12 +687,6 @@ function startPromptQueue( promptQueueIsRunning = true; promptQueuePrevStoreRunning = shouldWaitForCurrentRun; syncPromptQueueUI(); - toast( - shouldWaitForCurrentRun ? "Prompt queued" : `Prompt 1 / ${filtered.length}`, - { - description: queueToastDescription(filtered[0]), - }, - ); startPromptQueueSubscription(); if (!shouldWaitForCurrentRun) { const first = promptQueueItems[0]; @@ -563,10 +697,15 @@ function startPromptQueue( } function stopPromptQueueRun() { - const activeTarget = promptQueueItems[Math.max(promptQueueIndex, 0)]?.target; + const activeItem = promptQueueItems[Math.max(promptQueueIndex, 0)]; + const activeTarget = activeItem?.target; + const shouldCancelActiveRun = Boolean(activeItem?.dispatched); resetPromptQueue(); + if (!shouldCancelActiveRun) { + return; + } try { - activeTarget?.cancel(); + activeTarget.cancel(); } catch { // The active run may have already ended. } @@ -1018,6 +1157,26 @@ const ThreadComposerDock: FC<{ onHeightChange?: (height: number | null) => void; }> = ({ disabled, threadId, onHeightChange }) => { const { overlay } = useGeneratedImageOverlay(); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const threadListItemId = useAuiState( + ({ threadListItem }) => threadListItem.id, + ); + const threadListItemRemoteId = useAuiState( + ({ threadListItem }) => threadListItem.remoteId, + ); + const promptQueueThreadIds = compactIds([ + threadListItemId, + threadListItemRemoteId, + threadId, + activeThreadId, + ]); + const queueVisible = usePromptQueueUI( + (s) => + Boolean(findPromptQueueEntry(s, promptQueueThreadIds)) && + s.items.some((item) => + promptQueueItemMatchesThreadIds(item, promptQueueThreadIds), + ), + ); // Report dock height so the viewport reserves matching scroll space when // attachments or multiline input grow the composer. @@ -1046,7 +1205,12 @@ const ThreadComposerDock: FC<{ {/* Fade the top edge so scrolling text is not cut off by a hard line. */}

@@ -1743,14 +1907,15 @@ const Composer: FC<{ aria-disabled={disabled} onSubmit={handleSubmit} > + {isTauri ? ( // Phase 1 native model owns Tauri local-path drops. Restore browser // attachment drops in Tauri once Phase 1d adds token bridging. -
+
{composerContent}
) : ( - + {composerContent} {/* Gemini-style drop affordance, shown while a file is dragged over the composer. Absolute + pointer-events-none so the outline adds @@ -2991,6 +3156,184 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ ); }; +function promptQueueStatusLabel(status: PromptQueueUIItemStatus) { + switch (status) { + case "running": + return "Running now"; + case "waiting": + return "Waiting"; + case "next": + return "Next"; + case "queued": + return "Queued"; + default: { + const exhaustiveStatus: never = status; + throw new Error(`Unhandled prompt queue status: ${exhaustiveStatus}`); + } + } +} + +const PromptQueueStack: FC<{ queueThreadIds: string[] }> = ({ + queueThreadIds, +}) => { + const queueEntry = usePromptQueueUI((s) => + findPromptQueueEntry(s, queueThreadIds), + ); + const items = usePromptQueueUI((s) => s.items); + const [editingItemId, setEditingItemId] = useState(null); + const [draftPrompt, setDraftPrompt] = useState(""); + const editInputRef = useRef(null); + const visibleItems = items.filter((item) => + promptQueueItemMatchesThreadIds(item, queueThreadIds), + ); + const editingItem = visibleItems.find((item) => item.id === editingItemId); + const editingItemCanEdit = editingItem?.canEdit ?? false; + const activeEditingItemId = editingItem ? editingItemId : null; + + useEffect(() => { + if (!activeEditingItemId) { + return; + } + editInputRef.current?.focus(); + editInputRef.current?.select(); + }, [activeEditingItemId]); + + useEffect(() => { + if (!editingItemId || editingItemCanEdit) { + return; + } + setEditingItemId(null); + setDraftPrompt(""); + }, [editingItemCanEdit, editingItemId]); + + if (!queueEntry || visibleItems.length === 0) { + return null; + } + + const { current, total } = queueEntry; + + const startEditing = (item: PromptQueueUIItem) => { + if (!item.canEdit) { + return; + } + setEditingItemId(item.id); + setDraftPrompt(item.prompt); + }; + const saveEditing = () => { + if (!activeEditingItemId) { + return; + } + if (editPromptQueueItem(activeEditingItemId, draftPrompt)) { + setEditingItemId(null); + setDraftPrompt(""); + } + }; + const cancelEditing = () => { + setEditingItemId(null); + setDraftPrompt(""); + }; + + return ( +
+
+ {visibleItems.map((item, visibleIndex) => { + const isEditing = item.id === activeEditingItemId; + const visiblePosition = visibleIndex + 1; + return ( +
+ {isEditing ? ( +
+