+ Data
Overview
Columns
- Data
Raw
From a29b4e23fd92c804ce2c46ddc4acf73d756ec823 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Fri, 3 Apr 2026 13:48:24 +0100
Subject: [PATCH 09/15] studio: reuse HF cached repo casing to prevent
duplicate downloads (#4822)
* fix(studio): reuse HF cached repo casing to prevent duplicate downloads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Move cache case resolution tests to separate PR
Tests for resolve_cached_repo_id_case and get_model_config case resolution
belong in their own PR to keep this change focused on the runtime fix.
* fix(studio): debug-log HF_HUB_CACHE fallback in path_utils
* Fix stale memoization in resolve_cached_repo_id_case
- Check exact-case path before memo to ensure a newly-appeared exact
match always wins over a previously memoized variant
- Validate memoized entries still exist on disk before returning them
to prevent stale results when cache dirs are deleted/recreated
* Minor cleanups for cache case resolution
- Use .is_dir() instead of .exists() for exact-case cache check
(cache entries are always directories)
- Remove redundant fallback in _detect_audio_from_tokenizer since
get_cache_path already handles case resolution and returns None
when the model is not cached
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/backend/core/inference/worker.py | 3 +
studio/backend/routes/models.py | 15 ++-
studio/backend/utils/models/model_config.py | 31 +++--
studio/backend/utils/paths/__init__.py | 13 ++-
studio/backend/utils/paths/path_utils.py | 118 +++++++++++++++++++-
unsloth/models/loader.py | 23 ++--
unsloth/models/loader_utils.py | 2 +-
7 files changed, 180 insertions(+), 25 deletions(-)
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index e2513f43de..7f7291a56d 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -145,6 +145,8 @@ def _get_hf_download_state(
blobs_dirs: list[Path] = []
if model_names:
+ from utils.paths import resolve_cached_repo_id_case
+
for name in model_names:
if not name:
continue
@@ -154,6 +156,7 @@ def _get_hf_download_state(
# relative paths, and Windows paths.
if name.startswith(("/", ".", "~")) or "\\" in name:
continue
+ name = resolve_cached_repo_id_case(name)
# HF cache dir format: models--org--name (slashes -> --)
cache_dir_name = "models--" + name.replace("/", "--")
blobs_dir = cache / cache_dir_name / "blobs"
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 445cf0e7f4..1e31a91e26 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -49,8 +49,10 @@ try:
)
from core.inference import get_inference_backend
from utils.paths import (
+ is_local_path,
outputs_root,
exports_root,
+ resolve_cached_repo_id_case,
resolve_output_dir,
resolve_export_dir,
)
@@ -77,8 +79,10 @@ except ImportError:
)
from core.inference import get_inference_backend
from utils.paths import (
+ is_local_path,
outputs_root,
exports_root,
+ resolve_cached_repo_id_case,
resolve_output_dir,
resolve_export_dir,
)
@@ -597,10 +601,15 @@ async def get_model_config(
This endpoint wraps the backend load_model_defaults function.
"""
try:
- from utils.models.model_config import is_local_path
-
if not is_local_path(model_name):
- model_name = model_name.lower()
+ resolved = resolve_cached_repo_id_case(model_name)
+ if resolved != model_name:
+ logger.info(
+ "Using cached repo_id casing '%s' for requested '%s'",
+ resolved,
+ model_name,
+ )
+ model_name = resolved
logger.info(f"Getting model config for: {model_name}")
from utils.models.model_config import detect_audio_type
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 6cffc534aa..8c25da4f43 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -11,6 +11,8 @@ from utils.paths import (
normalize_path,
is_local_path,
is_model_cached,
+ get_cache_path,
+ resolve_cached_repo_id_case,
outputs_root,
exports_root,
resolve_output_dir,
@@ -711,12 +713,8 @@ def _detect_audio_from_tokenizer(
# 1) Check local HF cache first (works for gated/offline models)
try:
- from huggingface_hub.constants import HF_HUB_CACHE
-
- cache_dir = Path(HF_HUB_CACHE)
- repo_dir_name = f"models--{model_name.replace('/', '--')}"
- repo_dir = cache_dir / repo_dir_name
- if repo_dir.exists():
+ repo_dir = get_cache_path(model_name)
+ if repo_dir is not None and repo_dir.exists():
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.exists():
for snapshot in snapshots_dir.iterdir():
@@ -1627,11 +1625,18 @@ class ModelConfig:
identifier = f"unsloth/{identifier}"
path = identifier
- # Enforce lowercase for remote Hugging Face identifiers to prevent cache duplication
- # Hugging Face Hub APIs are case-insensitive remotely, but case-sensitive locally (repo_folder_name).
+ # Preserve requested casing, but if a case-variant already exists in local HF cache,
+ # reuse that exact repo_id spelling to avoid one-time re-downloads after #2592.
if not is_local:
- identifier = identifier.lower()
- path = path.lower()
+ resolved_identifier = resolve_cached_repo_id_case(identifier)
+ if resolved_identifier != identifier:
+ logger.info(
+ "Using cached repo_id casing '%s' for requested '%s'",
+ resolved_identifier,
+ identifier,
+ )
+ identifier = resolved_identifier
+ path = resolved_identifier
# Auto-detect GGUF models (check before LoRA/vision detection)
if is_local:
@@ -1852,6 +1857,12 @@ class ModelConfig:
identifier = f"unsloth/{identifier}"
path = identifier
+ if not is_local:
+ resolved_identifier = resolve_cached_repo_id_case(identifier)
+ if resolved_identifier != identifier:
+ identifier = resolved_identifier
+ path = resolved_identifier
+
# --- Logic for Base Model and Vision Detection ---
base_model = None
is_vision = False
diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py
index 44a7c8e287..11709ae56e 100644
--- a/studio/backend/utils/paths/__init__.py
+++ b/studio/backend/utils/paths/__init__.py
@@ -5,7 +5,15 @@
Path utilities for model and dataset handling
"""
-from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path
+from .path_utils import (
+ normalize_path,
+ is_local_path,
+ is_model_cached,
+ get_cache_path,
+ resolve_cached_repo_id_case,
+ get_cache_case_resolution_stats,
+ reset_cache_case_resolution_state,
+)
from .storage_roots import (
studio_root,
assets_root,
@@ -40,6 +48,9 @@ __all__ = [
"is_local_path",
"is_model_cached",
"get_cache_path",
+ "resolve_cached_repo_id_case",
+ "get_cache_case_resolution_stats",
+ "reset_cache_case_resolution_state",
"studio_root",
"assets_root",
"datasets_root",
diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py
index b38db18286..9ef9a2dd92 100644
--- a/studio/backend/utils/paths/path_utils.py
+++ b/studio/backend/utils/paths/path_utils.py
@@ -14,6 +14,20 @@ from loggers import get_logger
logger = get_logger(__name__)
+# Per-process cache to avoid repeated cache-dir scans for the same identifier.
+_CACHE_CASE_RESOLUTION_MEMO: dict[str, str] = {}
+
+# Lightweight instrumentation counters for operational visibility.
+_CACHE_CASE_RESOLUTION_STATS: dict[str, int] = {
+ "calls": 0,
+ "memo_hits": 0,
+ "exact_hits": 0,
+ "variant_hits": 0,
+ "tie_breaks": 0,
+ "fallbacks": 0,
+ "errors": 0,
+}
+
def _is_wsl() -> bool:
"""Detect if we are running inside WSL (Windows Subsystem for Linux)."""
@@ -94,8 +108,9 @@ def is_local_path(path: str) -> bool:
def get_cache_path(model_name: str) -> Optional[Path]:
"""Get HuggingFace cache path for a model if it exists."""
- cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
- model_cache_name = model_name.replace("/", "--")
+ cache_dir = _hf_hub_cache_dir()
+ resolved_name = resolve_cached_repo_id_case(model_name)
+ model_cache_name = resolved_name.replace("/", "--")
model_cache_path = cache_dir / f"models--{model_cache_name}"
return model_cache_path if model_cache_path.exists() else None
@@ -113,3 +128,102 @@ def is_model_cached(model_name: str) -> bool:
return True
return False
+
+
+def _hf_hub_cache_dir() -> Path:
+ """Return HF cache root honoring HF_HUB_CACHE when available."""
+ try:
+ from huggingface_hub.constants import HF_HUB_CACHE
+
+ return Path(HF_HUB_CACHE)
+ except Exception as exc:
+ logger.debug(
+ "Could not read huggingface_hub HF_HUB_CACHE, using default hub path: %s",
+ exc,
+ )
+ return Path.home() / ".cache" / "huggingface" / "hub"
+
+
+def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
+ """Resolve repo_id to the exact casing already present in local HF cache.
+
+ Policy: prefer the requested/canonical repo_id, but if a case-variant already
+ exists in local HF cache, reuse that exact cached spelling. This avoids
+ duplicate downloads while preserving user intent whenever possible.
+ """
+ _CACHE_CASE_RESOLUTION_STATS["calls"] += 1
+
+ if not model_name or "/" not in model_name:
+ _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+ return model_name
+
+ cache_dir = _hf_hub_cache_dir()
+ if not cache_dir.exists():
+ _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+ return model_name
+
+ expected_dir = f"models--{model_name.replace('/', '--')}"
+
+ # Always check the exact-case path first so a newly-appeared exact match
+ # wins over any previously memoized variant.
+ exact_path = cache_dir / expected_dir
+ if exact_path.is_dir():
+ if use_memo:
+ _CACHE_CASE_RESOLUTION_MEMO[model_name] = model_name
+ _CACHE_CASE_RESOLUTION_STATS["exact_hits"] += 1
+ return model_name
+
+ # Validate memoized entries still exist on disk before returning them.
+ # This prevents stale results when cache dirs are deleted/recreated.
+ if use_memo:
+ cached = _CACHE_CASE_RESOLUTION_MEMO.get(model_name)
+ if cached is not None:
+ cached_path = cache_dir / f"models--{cached.replace('/', '--')}"
+ if cached_path.is_dir():
+ _CACHE_CASE_RESOLUTION_STATS["memo_hits"] += 1
+ return cached
+ # Stale entry -- drop it and re-scan below.
+ _CACHE_CASE_RESOLUTION_MEMO.pop(model_name, None)
+
+ expected_lower = expected_dir.lower()
+ try:
+ candidates: list[str] = []
+ for entry in cache_dir.iterdir():
+ if not entry.is_dir():
+ continue
+ if entry.name.lower() != expected_lower:
+ continue
+ if not entry.name.startswith("models--"):
+ continue
+ repo_part = entry.name[len("models--") :]
+ if not repo_part:
+ continue
+ candidates.append(repo_part.replace("--", "/"))
+
+ if candidates:
+ # Deterministic tie-break if multiple case variants coexist.
+ resolved = sorted(candidates)[0]
+ if len(candidates) > 1:
+ _CACHE_CASE_RESOLUTION_STATS["tie_breaks"] += 1
+ _CACHE_CASE_RESOLUTION_STATS["variant_hits"] += 1
+ if use_memo:
+ _CACHE_CASE_RESOLUTION_MEMO[model_name] = resolved
+ return resolved
+ except Exception as exc:
+ _CACHE_CASE_RESOLUTION_STATS["errors"] += 1
+ logger.debug(f"Could not resolve cached repo_id case for '{model_name}': {exc}")
+
+ _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+ return model_name
+
+
+def get_cache_case_resolution_stats() -> dict[str, int]:
+ """Return a copy of case-resolution instrumentation counters."""
+ return dict(_CACHE_CASE_RESOLUTION_STATS)
+
+
+def reset_cache_case_resolution_state() -> None:
+ """Clear resolver memo and counters (primarily for tests)."""
+ _CACHE_CASE_RESOLUTION_MEMO.clear()
+ for key in _CACHE_CASE_RESOLUTION_STATS:
+ _CACHE_CASE_RESOLUTION_STATS[key] = 0
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index cfe3acb656..a6cb3eb529 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -114,6 +114,17 @@ FORCE_FLOAT32 = [
global DISABLE_COMPILE_MODEL_NAMES
# Must be alphabetically sorted for each entry
+
+
+def _strip_unsloth_bnb_4bit_suffix(model_name: str) -> str:
+ """Remove Unsloth 4bit suffixes without lowercasing (HF cache dirs are case-sensitive)."""
+ s = model_name
+ for suffix in ("-unsloth-bnb-4bit", "-bnb-4bit"):
+ if len(s) >= len(suffix) and s.lower().endswith(suffix.lower()):
+ s = s[: -len(suffix)]
+ return s
+
+
DISABLE_COMPILE_MODEL_NAMES = [
"aya_vision",
"modernbert",
@@ -404,8 +415,7 @@ class FastLanguageModel(FastLlamaModel):
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
("-unsloth-bnb-4bit", "-bnb-4bit")
):
- model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
- model_name = model_name.lower().removesuffix("-bnb-4bit")
+ model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
# Change -BF16 to all False for 4bit, 8bit etc
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
@@ -551,8 +561,7 @@ class FastLanguageModel(FastLlamaModel):
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
("-unsloth-bnb-4bit", "-bnb-4bit")
):
- model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
- model_name = model_name.lower().removesuffix("-bnb-4bit")
+ model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
# Change -BF16 to all False for 4bit, 8bit etc
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
@@ -1019,8 +1028,7 @@ class FastModel(FastBaseModel):
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
("-unsloth-bnb-4bit", "-bnb-4bit")
):
- model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
- model_name = model_name.lower().removesuffix("-bnb-4bit")
+ model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
# Change -BF16 to all False for 4bit, 8bit etc
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
@@ -1320,8 +1328,7 @@ class FastModel(FastBaseModel):
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
("-unsloth-bnb-4bit", "-bnb-4bit")
):
- model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
- model_name = model_name.lower().removesuffix("-bnb-4bit")
+ model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
# Change -BF16 to all False for 4bit, 8bit etc
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index cf5af983a6..99da5f799e 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -162,7 +162,7 @@ def __get_model_name(
# Support returning original full -bnb-4bit name if specified specifically
# since we'll map it to the dynamic version instead
if lower_model_name.endswith("-bnb-4bit"):
- return lower_model_name
+ return model_name
new_model_name = FLOAT_TO_INT_MAPPER[lower_model_name]
# logger.warning_once(
From c027ec192ef5e4e69f4c02973081f5b3da67aa20 Mon Sep 17 00:00:00 2001
From: Neodon <82944+neodon@users.noreply.github.com>
Date: Fri, 3 Apr 2026 13:44:22 -0500
Subject: [PATCH 10/15] fix(studio): ensure first chat tool call starts in
session sandbox (#4810)
Fixes #4809
On a new Studio chat, the first tool call could start before the frontend
initializes the thread ID. That meant the first request could go out without
a session_id, so the backend started the tool in the shared sandbox root
instead of the chat's session sandbox.
Frontend:
- Eagerly initialize the thread when switching to a new chat
- Resolve the thread ID once at request time and keep it stable through
async model-load waits
- Disable ActiveThreadSync during new-chat initialization to prevent
stale thread IDs from being written back
- Add error handling for thread initialization failures
- Clear activeThreadId on all compare-mode entry paths to prevent
cross-session leakage
- Fix exitCompare to restore context usage from the saved view
- Coerce falsy thread IDs to undefined for consistent backend/frontend
fallback behavior
- Use _default as the image sessionId fallback to match the backend
Backend:
- Use ~/studio_sandbox/_default when a request arrives without a session_id
---
studio/backend/core/inference/tools.py | 5 ++--
.../src/features/chat/api/chat-adapter.ts | 16 ++++++----
.../frontend/src/features/chat/chat-page.tsx | 15 ++++++++--
.../src/features/chat/runtime-provider.tsx | 30 +++++++++++++++++--
4 files changed, 54 insertions(+), 12 deletions(-)
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index b23372b766..86c22ce25a 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -34,7 +34,8 @@ _MAX_OUTPUT_CHARS = 8000 # truncate long output
_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
# Per-session working directories so each chat thread gets its own sandbox.
-# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
+# Falls back to a shared ~/studio_sandbox/_default for API callers without a
+# session_id.
_workdirs: dict[str, str] = {}
@@ -55,7 +56,7 @@ def _get_workdir(session_id: str | None = None) -> str:
if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
workdir = os.path.join(sandbox_root, "_invalid")
else:
- workdir = sandbox_root
+ workdir = os.path.join(sandbox_root, "_default")
os.makedirs(workdir, exist_ok = True)
_workdirs[key] = workdir
return _workdirs[key]
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index e287daf33a..3d1bce2905 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -421,6 +421,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
return {
async *run({ messages, abortSignal, unstable_threadId }) {
let runtime = useChatRuntimeStore.getState();
+ // Capture the thread ID once at the start so it stays stable even if
+ // the user switches chats while waiting for model load / auto-load.
+ const resolvedThreadId =
+ (unstable_threadId ?? runtime.activeThreadId) || undefined;
// Wait for in-progress model load to finish before inferring
if (runtime.modelLoading) {
@@ -473,14 +477,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
runtime.clearPendingAudio();
}
- const useAdapter = await resolveUseAdapter(unstable_threadId);
+ const useAdapter = await resolveUseAdapter(resolvedThreadId);
// ── Audio model path (non-streaming) ─────────────────────
const activeModel = runtime.models.find(
(m) => m.id === params.checkpoint,
);
if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
- const threadKey = unstable_threadId || "__default";
+ const threadKey = resolvedThreadId || "__default";
runtime.setThreadRunning(threadKey, true);
try {
yield {
@@ -527,7 +531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
return;
}
- const threadKey = unstable_threadId || "__default";
+ const threadKey = resolvedThreadId || "__default";
let waitingFirstChunk = true;
let firstTokenSettled = false;
const streamStartTime = Date.now();
@@ -600,7 +604,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
const mins = useChatRuntimeStore.getState().toolCallTimeout;
return mins >= 9999 ? 9999 : mins * 60;
})(),
- session_id: unstable_threadId || undefined,
+ session_id: resolvedThreadId,
}
: {}),
},
@@ -641,7 +645,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
let parsedResult: string | { text: string; images: string[]; sessionId: string };
if (imgIdx !== -1) {
const text = rawResult.slice(0, imgIdx);
- const sessionId = unstable_threadId || "";
+ // Fall back to "_default" to match the backend sandbox directory
+ // used when no session_id is provided (see tools.py _get_workdir).
+ const sessionId = resolvedThreadId || "_default";
try {
const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[];
parsedResult = { text, images, sessionId };
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 08450c7ec7..1dbff145ee 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -592,6 +592,9 @@ export function ChatPage(): ReactElement {
}, []);
const handleNewCompare = useCallback(() => {
setView({ mode: "compare", pairId: crypto.randomUUID() });
+ // Clear activeThreadId so compare panes do not inherit the single-chat
+ // thread ID as a fallback for session_id routing.
+ useChatRuntimeStore.getState().setActiveThreadId(null);
useChatRuntimeStore.getState().setContextUsage(null);
}, []);
@@ -619,6 +622,9 @@ export function ChatPage(): ReactElement {
const enterCompare = useCallback(() => {
setViewBeforeCompare((prev) => prev ?? view);
setView({ mode: "compare", pairId: crypto.randomUUID() });
+ // Clear activeThreadId so compare panes do not inherit the single-chat
+ // thread ID as a fallback for session_id routing.
+ useChatRuntimeStore.getState().setActiveThreadId(null);
useChatRuntimeStore.getState().setContextUsage(null);
}, [view]);
@@ -626,9 +632,13 @@ export function ChatPage(): ReactElement {
if (!viewBeforeCompare) return;
setView(viewBeforeCompare);
setViewBeforeCompare(null);
- // Restore context usage from the active thread's last assistant message
+ // Restore context usage from the active thread's last assistant message.
+ // Use the thread ID from the saved view rather than the store, because
+ // activeThreadId may have been cleared on compare entry.
const store = useChatRuntimeStore.getState();
- const threadId = store.activeThreadId;
+ const threadId =
+ ("threadId" in viewBeforeCompare ? viewBeforeCompare.threadId : null) ??
+ store.activeThreadId;
if (threadId) {
void db.messages
.where("threadId")
@@ -735,6 +745,7 @@ export function ChatPage(): ReactElement {
await selectModelRef.current({ id: targetLora.id, isLora: true });
if (canceled) return;
setView({ mode: "compare", pairId: crypto.randomUUID() });
+ useChatRuntimeStore.getState().setActiveThreadId(null);
useChatRuntimeStore.getState().setContextUsage(null);
clearHandoff();
console.info("[chat-handoff] loaded lora + opened compare");
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index 02f792b509..7e6cd8e1dd 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -679,9 +679,33 @@ function ThreadNewChatSwitch({
const isLoading = useAuiState(({ threads }) => threads.isLoading);
useEffect(() => {
- if (!isLoading) {
- aui.threads().switchToNewThread();
+ if (isLoading) {
+ return;
}
+
+ let cancelled = false;
+ // Clear immediately so the adapter never picks up a stale thread ID
+ // from a previous chat while we initialize the new one.
+ useChatRuntimeStore.getState().setActiveThreadId(null);
+
+ void (async () => {
+ try {
+ aui.threads().switchToNewThread();
+ const { remoteId } = await aui.threadListItem().initialize();
+ if (!cancelled) {
+ useChatRuntimeStore.getState().setActiveThreadId(remoteId);
+ }
+ } catch (error) {
+ if (!cancelled) {
+ useChatRuntimeStore.getState().setActiveThreadId(null);
+ }
+ console.error("Failed to initialize new chat thread", error);
+ }
+ })();
+
+ return () => {
+ cancelled = true;
+ };
}, [aui, isLoading, nonce]);
return null;
@@ -730,7 +754,7 @@ export function ChatRuntimeProvider({
return (
-
+
{initialThreadId && }
{!initialThreadId && newThreadNonce && (
From 2c73ab7871849d50f85399d0cf8a453ce55c1baa Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Fri, 3 Apr 2026 13:33:42 -0700
Subject: [PATCH 11/15] fix(studio): harden sandbox security for terminal and
python tools (#4827)
* fix(studio): harden sandbox security for terminal and python tools
The existing command blocklist used naive str.split() which is trivially
bypassable via quoting, full paths, nested shells, variable expansion,
and cross-tool pivoting through Python os.system/subprocess. Fixes #4818.
Changes:
- Replace str.split() blocklist with shlex.split() + os.path.basename()
tokenization and regex scanning at shell command boundaries
- Add sanitized subprocess environment (_build_safe_env) that strips
credentials (HF_TOKEN, WANDB_API_KEY, GH_TOKEN, AWS_*, etc.) and
restricts PATH to /usr/local/bin:/usr/bin:/bin
- Add PR_SET_NO_NEW_PRIVS via prctl on Linux so sudo/su/pkexec fail
at the kernel level regardless of how they are invoked
- Add RLIMIT_NPROC (256) and RLIMIT_FSIZE (100MB) to prevent fork
bombs and disk filling attacks
- Extend AST safety checker to detect os.system(), os.popen(),
subprocess.run/Popen/call/check_output, os.exec*, os.spawn* calls
containing blocked commands or dynamic (non-literal) arguments
- Add cross-platform support: cmd.exe on Windows, bash on Unix;
CREATE_NO_WINDOW flag on Windows, preexec_fn on Unix
- Expand blocklist from 7 to 14 commands: add su, chown, passwd,
mount, umount, fdisk, kill, killall, pkill
- Apply all layers to both _bash_exec and _python_exec
Zero measurable performance overhead -- shlex parsing and a single
prctl syscall per subprocess fork.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix review findings: exception_catching dead code, false positives, process substitution
- Include exception_catching reasons in _check_code_safety so bare
except-in-loop timeout evasion is actually blocked (was computed in
_check_signal_escape_patterns but never read by the caller)
- Remove base.split() inner loop that caused false positives on quoted
text arguments containing blocked words (e.g. echo "kill this process")
- Add targeted nested shell detection for bash/sh/zsh -c arguments
instead, which catches bash -c 'sudo whoami' without false positives
- Add <() process substitution to the regex character class so
diff <(rm -rf /path) is also caught
- Fix error message to say "unsafe patterns" instead of specifically
mentioning signal manipulation when other categories trigger
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review feedback: regex paths, keyword args, list element scanning
- Regex now matches blocked commands after optional path prefix at shell
boundaries (catches ls; /usr/bin/sudo and similar)
- Nested shell detection uses os.path.basename so bash -c "/bin/rm" is
caught
- AST checker now inspects keyword arguments (not just positional) so
subprocess.run(args="sudo ...", shell=True) is detected
- List elements in subprocess calls are now checked via
_find_blocked_commands for consistency (catches subprocess.run(["bash",
"-c", "rm -rf /"]))
- Dynamic argument check uses _is_safe_literal that validates list
contents are all string literals
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix nested shell scan to only check the script body, not positional args
bash -c 'script' arg0 arg1 -- only tokens[i+1] is the script body;
subsequent tokens are $0, $1 positional parameters passed to the script
and are not executed as shell commands. Scanning all remaining tokens
caused false positives.
* Add subshell parentheses to regex command boundary detection
(sudo whoami) was not caught because ( was not in the regex character
class for shell command boundaries. Add ( to the set alongside ;, &,
|, backtick, newline.
* Address high-priority review findings from 7 parallel reviewers
- Track from-imports of dangerous functions (from os import system,
from subprocess import run as r, etc.) via shell_exec_aliases dict
so bare-name calls are detected by the AST checker
- Include the active Python interpreter and virtualenv directories
in the sanitized PATH so pip, uv, and Studio packages remain
accessible in the sandbox
- Add Windows-specific blocked commands (rmdir, takeown, icacls,
runas, powershell, pwsh) only on win32 platform
- Add os.posix_spawn and os.posix_spawnp to _SHELL_EXEC_FUNCS
- Handle tuple literals same as list literals in AST argument
inspection (both _extract_strings_from_list and _is_safe_literal)
* Fix false positive on check=True kwargs and recursive nested shell scanning
- Only inspect command-carrying keyword arguments (args, command,
executable, path, file) in the AST checker, not control flags like
check=True, text=True, capture_output=True which are booleans and
were incorrectly flagged as non-literal dynamic arguments
- Replace split() in nested shell detection with recursive call to
_find_blocked_commands so that quoted commands (bash -c '"sudo"
whoami') and semicolons (bash -c "sudo;ls") within nested shells
are properly detected through the full shlex + regex pipeline
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Move preexec_fn imports to module level and use find_library for libc
Addresses two Gemini review findings:
1. preexec_fn thread safety: _sandbox_preexec previously imported ctypes
and resource inside the function body, which runs between fork() and
exec() in the child process. In a multi-threaded server, this could
deadlock if the import machinery locks were held by another thread at
fork time. Now all imports and the libc handle are resolved once at
module load time, so _sandbox_preexec only calls C-level functions
(prctl, setrlimit) with no Python import activity.
2. Hardcoded libc.so.6 path: replaced with ctypes.util.find_library("c")
which works on glibc (libc.so.6), musl (libc.musl-*.so.1), and other
Linux distributions where libc has a different soname.
* Apply Gemini style suggestions: combined regex, dict.fromkeys, constant hoisting
- Combine per-word regex loop into a single re.findall with alternation
pattern, avoiding repeated regex compilation and searching
- Replace manual dedup loop with dict.fromkeys for PATH entries
- Hoist _CMD_KWARGS frozenset out of visit_Call to avoid recreating it
on every AST node visit
* Add cmd /c nested shell detection for Windows parity
The nested shell scan only checked for Unix shells (bash -c, sh -c, etc).
Add cmd /c and cmd.exe /c detection so that Windows nested shell
invocations are also recursively scanned for blocked commands. The token
scan already catches blocked commands at any position, so this is
defense-in-depth for consistency across platforms.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle combined shell flags (-lc, -xc) and interleaved flags (--login -c)
The nested shell scan only matched token == "-c" with the immediately
preceding token being a shell name. This missed:
- Combined flags: bash -lc 'rm ...' (-lc ends with c, is a valid
combined flag meaning -l -c)
- Interleaved flags: bash --login -c 'sudo ...' (--login sits between
bash and -c)
Now matches any short flag ending in 'c' (e.g. -lc, -xc, -ic) and
walks backwards past intermediate flags to find the shell binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix /bin/bash bypass, remove RLIMIT_NPROC, reduce AST false positives
Addresses three high-consensus findings from 20-reviewer pass:
1. /bin/bash -c 'sudo whoami' bypassed nested shell scan because the
backwards flag-skip logic treated paths starting with / as flags.
Now only skips tokens starting with - as Unix flags; on Windows
only skips short /X flags (not /bin/bash style paths). [9/20]
2. RLIMIT_NPROC=256 caused subprocess.run to fail with EAGAIN because
Linux enforces NPROC per real UID, not per process tree. Removed
RLIMIT_NPROC entirely; RLIMIT_FSIZE and PR_SET_NO_NEW_PRIVS remain
as the primary resource and privilege controls. [5/20]
3. AST checker rejected safe dynamic subprocess usage like
cmd=["git","status"]; subprocess.run(cmd) as shell_escape_dynamic.
Now only flags dynamic args for shell-string functions (os.system,
os.popen, subprocess.getoutput, etc.) or when shell=True is
explicitly set. List-based subprocess calls with shell=False (the
default) do not pass through a shell and are not flagged. [12/20]
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle Windows drive letter paths and .exe extensions in command detection
Gemini review found that Windows absolute paths (C:\Windows\System32\
shutdown.exe) and executable extensions (.exe, .com, .bat, .cmd) were
not handled:
- Token scan now strips .exe/.com/.bat/.cmd extensions before checking
the blocklist, so sudo.exe matches sudo, shutdown.bat matches shutdown
- Regex pattern now includes optional Windows drive letter prefix
([a-zA-Z]:[/\\]) and optional executable extension suffix, so commands
after shell metacharacters with full Windows paths are also caught
* Handle **kwargs dict expansion, non-literal shell=, and except Exception false positive
Addresses three findings from second 20-reviewer pass:
1. **kwargs dict expansion (9/20): subprocess.run(**{"args": "rm ...",
"shell": True}) bypassed the AST checker because **kwargs were
treated as opaque. Now expands literal dict **kwargs to inspect
their keys, and flags opaque **kwargs (variable dicts) as unsafe.
2. Non-literal shell= values (7/20): shell=variable was treated as
shell=False (safe). Now any shell= value that is not literally
False is treated as potentially True (conservative default).
3. except Exception false positive (1/20): except Exception in a loop
was flagged as timeout evasion, but Exception does not catch
SystemExit or KeyboardInterrupt which are used for timeout
enforcement. Narrowed to only flag except BaseException and
except TimeoutError in loops.
* [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/tools.py | 481 ++++++++++++++++++++++++-
1 file changed, 466 insertions(+), 15 deletions(-)
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index 86c22ce25a..87cc933d4b 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -14,6 +14,8 @@ import os
os.environ["UNSLOTH_IS_PRESENT"] = "1"
import random
+import re
+import shlex
import ssl
import subprocess
import sys
@@ -27,11 +29,235 @@ logger = get_logger(__name__)
_EXEC_TIMEOUT = 300 # 5 minutes
+# Pre-import modules used in _sandbox_preexec at module level so that
+# the preexec_fn closure does not trigger the import machinery in the
+# forked child (which can deadlock in multi-threaded servers).
+_libc = None
+if sys.platform == "linux":
+ try:
+ import ctypes
+ import ctypes.util
+
+ _libc_name = ctypes.util.find_library("c")
+ if _libc_name:
+ _libc = ctypes.CDLL(_libc_name, use_errno = True)
+ except (OSError, AttributeError):
+ pass
+
+_resource = None
+if sys.platform != "win32":
+ try:
+ import resource as _resource
+ except ImportError:
+ pass
+
# Strict raster-image allowlist for sandbox file serving.
# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
_MAX_OUTPUT_CHARS = 8000 # truncate long output
-_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
+_BLOCKED_COMMANDS_COMMON = frozenset(
+ {
+ "rm",
+ "sudo",
+ "su",
+ "dd",
+ "chmod",
+ "chown",
+ "mkfs",
+ "shutdown",
+ "reboot",
+ "passwd",
+ "mount",
+ "umount",
+ "fdisk",
+ "kill",
+ "killall",
+ "pkill",
+ }
+)
+_BLOCKED_COMMANDS_WIN = frozenset(
+ {
+ "rmdir",
+ "takeown",
+ "icacls",
+ "runas",
+ "powershell",
+ "pwsh",
+ }
+)
+_BLOCKED_COMMANDS = (
+ _BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN
+ if sys.platform == "win32"
+ else _BLOCKED_COMMANDS_COMMON
+)
+
+
+def _find_blocked_commands(command: str) -> set[str]:
+ """Detect blocked commands using shlex tokenization and regex scanning.
+
+ Catches: full paths (/usr/bin/sudo), quoted strings ("sudo"),
+ split-quotes (su""do), backslash escapes (\\rm), and command-position
+ words after ;, |, &&, $().
+ """
+ blocked = set()
+
+ # 1. shlex tokenization (handles quotes, escapes, concatenation)
+ try:
+ tokens = (
+ shlex.split(command)
+ if sys.platform != "win32"
+ else shlex.split(command, posix = False)
+ )
+ except ValueError:
+ tokens = command.split()
+
+ for token in tokens:
+ base = os.path.basename(token).lower()
+ # Strip common Windows executable extensions so that
+ # runas.exe, shutdown.bat, etc. match the blocklist.
+ stem, ext = os.path.splitext(base)
+ if ext in {".exe", ".com", ".bat", ".cmd"}:
+ base = stem
+ if base in _BLOCKED_COMMANDS:
+ blocked.add(base)
+
+ # 2. Regex: catch blocked words at shell command boundaries
+ # (semicolons, pipes, &&, ||, backticks, $(), <(), subshells, newlines)
+ # Uses a single combined pattern for all blocked words.
+ # Handles optional Unix path prefix (/usr/bin/) and Windows drive
+ # letter prefix (C:\Windows\...\).
+ lowered = command.lower()
+ if _BLOCKED_COMMANDS:
+ words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
+ pattern = (
+ rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)"
+ rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?"
+ rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b"
+ )
+ blocked.update(re.findall(pattern, lowered))
+
+ # 3. Check for nested shell invocations (bash -c 'sudo whoami',
+ # bash -lc '...', bash --login -c '...', cmd /c '...').
+ # When a -c or /c flag is found, look backwards for a shell name
+ # (skipping intermediate flags like --login, -l, -x) and recursively
+ # scan the nested command string.
+ _SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"}
+ _SHELLS_WIN = {"cmd", "cmd.exe"}
+ for i, token in enumerate(tokens):
+ tok_lower = token.lower()
+ # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc)
+ is_unix_c = tok_lower == "-c" or (
+ tok_lower.startswith("-")
+ and tok_lower.endswith("c")
+ and not tok_lower.startswith("--")
+ )
+ is_win_c = tok_lower == "/c"
+ if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
+ continue
+ # Look backwards past any flags to find the shell binary.
+ # On Unix, flags start with - (skip those). On Windows, flags
+ # start with / but so do absolute paths, so only skip short
+ # single-char /X flags (not /bin/bash style paths).
+ for j in range(i - 1, -1, -1):
+ prev = tokens[j]
+ if prev.startswith("-"):
+ continue # skip Unix flags like --login, -l
+ if is_win_c and prev.startswith("/") and len(prev) <= 3:
+ continue # skip Windows flags like /s, /q (not /bin/bash)
+ prev_base = os.path.basename(prev).lower()
+ if is_unix_c and prev_base in _SHELLS:
+ blocked |= _find_blocked_commands(tokens[i + 1])
+ elif is_win_c and prev_base in _SHELLS_WIN:
+ blocked |= _find_blocked_commands(tokens[i + 1])
+ break # stop at first non-flag token
+
+ return blocked
+
+
+def _build_safe_env(workdir: str) -> dict[str, str]:
+ """Build a minimal, credential-free environment for sandboxed subprocesses.
+
+ Strips HF_TOKEN, WANDB_API_KEY, AWS_*, GH_TOKEN, LD_PRELOAD, DYLD_*, etc.
+ Preserves the active Python interpreter and virtualenv directories in PATH
+ so that pip, uv, and packages installed in the Studio runtime remain
+ accessible.
+ """
+ # Start with the directory containing the running Python interpreter
+ # so that subprocess calls to 'python', 'pip', etc. resolve to the
+ # same environment the Studio server is running in.
+ exe_dir = os.path.dirname(sys.executable)
+ path_entries = [exe_dir] if exe_dir else []
+
+ # If a virtualenv is active, include its bin/Scripts directory.
+ venv = os.environ.get("VIRTUAL_ENV")
+ if venv:
+ venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin")
+ if venv_bin not in path_entries:
+ path_entries.append(venv_bin)
+
+ if sys.platform == "win32":
+ sysroot = os.environ.get("SystemRoot", r"C:\Windows")
+ path_entries.extend([os.path.join(sysroot, "System32"), sysroot])
+ else:
+ path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"])
+
+ # Deduplicate while preserving order
+ deduped = list(dict.fromkeys(p for p in path_entries if p))
+
+ env = {
+ "PATH": os.pathsep.join(deduped),
+ "HOME": workdir,
+ "TMPDIR": workdir,
+ "LANG": os.environ.get("LANG", "C.UTF-8"),
+ "TERM": "dumb",
+ "PYTHONIOENCODING": "utf-8",
+ }
+ if venv:
+ env["VIRTUAL_ENV"] = venv
+ # Windows needs SystemRoot for Python/subprocess to work
+ if sys.platform == "win32":
+ env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows")
+ return env
+
+
+def _sandbox_preexec():
+ """Pre-exec hook: drop privilege escalation ability and set resource limits.
+
+ On Linux, applies PR_SET_NO_NEW_PRIVS so sudo/su/pkexec fail at the
+ kernel level. On Linux and macOS, sets RLIMIT_FSIZE.
+ No-op on Windows (use creationflags instead).
+
+ Note: RLIMIT_NPROC is intentionally NOT set because Linux enforces it
+ per real UID, not per process tree, so it would starve the Studio
+ server and other sessions sharing the same user account.
+
+ All modules and handles are resolved at import time (module level) so
+ this function does not trigger Python imports in the forked child,
+ avoiding potential deadlocks in multi-threaded servers.
+ """
+ if _libc is not None:
+ try:
+ # PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable)
+ _libc.prctl(38, 1, 0, 0, 0)
+ except (OSError, AttributeError):
+ pass # Not available (container, old kernel, etc.)
+
+ if _resource is not None:
+ try:
+ # Limit file size to 100MB (prevents disk filling)
+ _resource.setrlimit(
+ _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)
+ )
+ except (ValueError, OSError):
+ pass
+
+
+def _get_shell_cmd(command: str) -> list[str]:
+ """Return the platform-appropriate shell invocation for a command string."""
+ if sys.platform == "win32":
+ return ["cmd", "/c", command]
+ return ["bash", "-c", command]
+
# Per-session working directories so each chat thread gets its own sandbox.
# Falls back to a shared ~/studio_sandbox/_default for API callers without a
@@ -429,6 +655,7 @@ def _check_signal_escape_patterns(code: str):
signal_tampering = []
exception_catching = []
+ shell_escapes = []
warnings = []
def _ast_name_matches(node, names):
@@ -446,10 +673,84 @@ def _check_signal_escape_patterns(code: str):
return full_name in names
return False
+ # Dangerous os/subprocess functions that can execute shell commands
+ _SHELL_EXEC_FUNCS = frozenset(
+ {
+ "os.system",
+ "os.popen",
+ "os.popen2",
+ "os.popen3",
+ "os.popen4",
+ "os.execl",
+ "os.execle",
+ "os.execlp",
+ "os.execlpe",
+ "os.execv",
+ "os.execve",
+ "os.execvp",
+ "os.execvpe",
+ "os.spawnl",
+ "os.spawnle",
+ "os.spawnlp",
+ "os.spawnlpe",
+ "os.spawnv",
+ "os.spawnve",
+ "os.spawnvp",
+ "os.spawnvpe",
+ "os.posix_spawn",
+ "os.posix_spawnp",
+ "subprocess.run",
+ "subprocess.call",
+ "subprocess.check_call",
+ "subprocess.check_output",
+ "subprocess.Popen",
+ "subprocess.getoutput",
+ "subprocess.getstatusoutput",
+ }
+ )
+
+ def _extract_string_from_node(node):
+ """Extract a plain string value from an AST node, if it is a constant."""
+ if isinstance(node, ast.Constant) and isinstance(node.value, str):
+ return node.value
+ return None
+
+ def _extract_strings_from_list(node):
+ """Extract string elements from an AST List or Tuple node."""
+ if isinstance(node, (ast.List, ast.Tuple)):
+ parts = []
+ for elt in node.elts:
+ s = _extract_string_from_node(elt)
+ if s is not None:
+ parts.append(s)
+ return parts
+ return []
+
+ # Keyword argument names that carry command content (as opposed to
+ # control flags like check=True, text=True, capture_output=True).
+ _CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"})
+
+ def _check_args_for_blocked(args_nodes):
+ """Check if any call arguments contain blocked commands."""
+ found = set()
+ for arg in args_nodes:
+ s = _extract_string_from_node(arg)
+ if s is not None:
+ found |= _find_blocked_commands(s)
+ strs = _extract_strings_from_list(arg)
+ for s in strs:
+ found |= _find_blocked_commands(s)
+ return found
+
class SignalEscapeVisitor(ast.NodeVisitor):
def __init__(self):
self.imports_signal = False
self.signal_aliases = {"signal"}
+ self.os_aliases = {"os"}
+ self.subprocess_aliases = {"subprocess"}
+ # Maps bare function names to their fully-qualified form
+ # for from-import tracking (e.g. "system" -> "os.system")
+ self.shell_exec_aliases: dict[str, str] = {}
self.loop_depth = 0
def visit_Import(self, node):
@@ -458,6 +759,10 @@ def _check_signal_escape_patterns(code: str):
self.imports_signal = True
if alias.asname:
self.signal_aliases.add(alias.asname)
+ elif alias.name == "os":
+ self.os_aliases.add(alias.asname or "os")
+ elif alias.name == "subprocess":
+ self.subprocess_aliases.add(alias.asname or "subprocess")
self.generic_visit(node)
def visit_ImportFrom(self, node):
@@ -475,6 +780,16 @@ def _check_signal_escape_patterns(code: str):
"alarm",
):
self.signal_aliases.add(alias.asname or alias.name)
+ elif node.module in ("os", "subprocess"):
+ if node.module == "os":
+ self.os_aliases.add("os")
+ else:
+ self.subprocess_aliases.add("subprocess")
+ # Track from-imports of dangerous functions
+ for alias in node.names:
+ fq = f"{node.module}.{alias.name}"
+ if fq in _SHELL_EXEC_FUNCS:
+ self.shell_exec_aliases[alias.asname or alias.name] = fq
self.generic_visit(node)
def visit_While(self, node):
@@ -539,6 +854,111 @@ def _check_signal_escape_patterns(code: str):
"description": "Modifies signal mask (may block SIGALRM)",
}
)
+
+ # --- Shell escape detection ---
+ # Resolve the fully qualified function name for os.*/subprocess.*
+ shell_func = None
+ if isinstance(func, ast.Attribute):
+ if isinstance(func.value, ast.Name):
+ if func.value.id in self.os_aliases:
+ shell_func = f"os.{func.attr}"
+ elif func.value.id in self.subprocess_aliases:
+ shell_func = f"subprocess.{func.attr}"
+ elif isinstance(func, ast.Name):
+ # Check from-import aliases: from os import system; system(...)
+ shell_func = self.shell_exec_aliases.get(func.id)
+
+ if shell_func and shell_func in _SHELL_EXEC_FUNCS:
+ # Expand **kwargs dicts to inspect their keys
+ expanded_kwargs: dict[str, ast.AST] = {}
+ has_opaque_kwargs = False
+ for kw in node.keywords:
+ if kw.arg is not None:
+ expanded_kwargs[kw.arg] = kw.value
+ elif isinstance(kw.value, ast.Dict):
+ for k, v in zip(kw.value.keys, kw.value.values):
+ key = _extract_string_from_node(k) if k else None
+ if key is not None:
+ expanded_kwargs[key] = v
+ else:
+ has_opaque_kwargs = True
+
+ cmd_kw_values = [
+ v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS
+ ]
+ all_call_args = list(node.args) + cmd_kw_values
+ blocked_in_args = _check_args_for_blocked(all_call_args)
+
+ if has_opaque_kwargs:
+ # Can't inspect dynamic **kwargs -- flag as unsafe
+ shell_escapes.append(
+ {
+ "type": "shell_escape_dynamic",
+ "line": node.lineno,
+ "description": (
+ f"{shell_func}() called with dynamic **kwargs"
+ ),
+ }
+ )
+ elif blocked_in_args:
+ shell_escapes.append(
+ {
+ "type": "shell_escape",
+ "line": node.lineno,
+ "description": (
+ f"{shell_func}() invokes blocked command(s): "
+ f"{', '.join(sorted(blocked_in_args))}"
+ ),
+ }
+ )
+ else:
+ # Only flag dynamic args for functions that interpret
+ # strings as shell commands, or when shell= might be
+ # enabled. Treat any non-literal-False shell= value
+ # as potentially True (conservative).
+ _STRING_SHELL_FUNCS = frozenset(
+ {
+ "os.system",
+ "os.popen",
+ "os.popen2",
+ "os.popen3",
+ "os.popen4",
+ "subprocess.getoutput",
+ "subprocess.getstatusoutput",
+ }
+ )
+ shell_node = expanded_kwargs.get("shell")
+ shell_safe = shell_node is None or (
+ isinstance(shell_node, ast.Constant)
+ and shell_node.value is False
+ )
+ if shell_func in _STRING_SHELL_FUNCS or not shell_safe:
+
+ def _is_safe_literal(n):
+ if _extract_string_from_node(n) is not None:
+ return True
+ if isinstance(n, (ast.List, ast.Tuple)):
+ return all(
+ _extract_string_from_node(e) is not None
+ for e in n.elts
+ )
+ return False
+
+ has_non_literal = any(
+ not _is_safe_literal(a) for a in all_call_args
+ )
+ if has_non_literal:
+ shell_escapes.append(
+ {
+ "type": "shell_escape_dynamic",
+ "line": node.lineno,
+ "description": (
+ f"{shell_func}() called with non-literal "
+ f"shell command (potential shell escape)"
+ ),
+ }
+ )
+
self.generic_visit(node)
def visit_ExceptHandler(self, node):
@@ -554,7 +974,12 @@ def _check_signal_escape_patterns(code: str):
}
)
elif isinstance(node.type, ast.Name):
- if node.type.id in ("TimeoutError", "BaseException", "Exception"):
+ # Only flag BaseException and TimeoutError, NOT Exception.
+ # except Exception does not catch SystemExit or
+ # KeyboardInterrupt, so it cannot suppress timeout
+ # enforcement. Flagging Exception causes false positives
+ # on normal error-handling patterns.
+ if node.type.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{node.type.id}_in_loop",
@@ -565,7 +990,7 @@ def _check_signal_escape_patterns(code: str):
elif isinstance(node.type, ast.Tuple):
for elt in node.type.elts:
if isinstance(elt, ast.Name):
- if elt.id in ("TimeoutError", "BaseException", "Exception"):
+ if elt.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{elt.id}_in_loop",
@@ -581,10 +1006,15 @@ def _check_signal_escape_patterns(code: str):
if visitor.imports_signal and not signal_tampering:
warnings.append("Code imports 'signal' module - review manually for safety")
- is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0
+ is_safe = (
+ len(signal_tampering) == 0
+ and len(exception_catching) == 0
+ and len(shell_escapes) == 0
+ )
return is_safe, {
"signal_tampering": signal_tampering,
"exception_catching": exception_catching,
+ "shell_escapes": shell_escapes,
"warnings": warnings,
}
@@ -605,10 +1035,18 @@ def _check_code_safety(code: str) -> str | None:
reasons = [
item.get("description", "") for item in info.get("signal_tampering", [])
]
- return (
- f"Error: unsafe code detected ({'; '.join(reasons)}). "
- f"Please remove signal manipulation from your code."
- )
+ shell_reasons = [
+ item.get("description", "") for item in info.get("shell_escapes", [])
+ ]
+ exception_reasons = [
+ item.get("description", "") for item in info.get("exception_catching", [])
+ ]
+ all_reasons = [r for r in reasons + shell_reasons + exception_reasons if r]
+ if all_reasons:
+ return (
+ f"Error: unsafe code detected ({'; '.join(all_reasons)}). "
+ f"Please remove unsafe patterns from your code."
+ )
return None
@@ -663,13 +1101,20 @@ def _python_exec(
with os.fdopen(fd, "w") as f:
f.write(code)
- proc = subprocess.Popen(
- [sys.executable, tmp_path],
+ safe_env = _build_safe_env(workdir)
+ popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
+ env = safe_env,
)
+ if sys.platform != "win32":
+ popen_kwargs["preexec_fn"] = _sandbox_preexec
+ else:
+ popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
+
+ proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs)
# Spawn cancel watcher if we have a cancel event
if cancel_event is not None:
@@ -735,21 +1180,27 @@ def _bash_exec(
if not command or not command.strip():
return "No command provided."
- # Block dangerous commands
- tokens = set(command.lower().split())
- blocked = tokens & _BASH_BLOCKED_WORDS
+ # Block dangerous commands (shlex + regex based)
+ blocked = _find_blocked_commands(command)
if blocked:
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
try:
workdir = _get_workdir(session_id)
- proc = subprocess.Popen(
- ["bash", "-c", command],
+ safe_env = _build_safe_env(workdir)
+ popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
+ env = safe_env,
)
+ if sys.platform != "win32":
+ popen_kwargs["preexec_fn"] = _sandbox_preexec
+ else:
+ popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
+
+ proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs)
if cancel_event is not None:
watcher = threading.Thread(
From a32b871f0e2e2f94cb88fec81596eff0796f1254 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Fri, 3 Apr 2026 13:56:59 -0700
Subject: [PATCH 12/15] studio: add speculative decoding support (ngram-mod, on
by default) (#4836)
* studio: add speculative decoding support (ngram-mod, on by default)
Enable n-gram speculative decoding for GGUF models in Unsloth Studio.
Uses llama.cpp's ngram-mod mode which gives 10-40% faster generation
with zero VRAM cost via a 4MB fixed hash table that auto-resets on
low acceptance rates.
Backend:
- Add speculative_type field to LoadRequest, LoadResponse, and
InferenceStatusResponse pydantic models
- Add speculative_type parameter to LlamaCppBackend.load_model()
with allowlist validation (ngram-simple, ngram-mod)
- Pass --spec-type, --spec-ngram-size-n 16, --draft-max 24 flags
to llama-server when ngram-mod is active
- Default to ngram-mod for non-vision GGUF models server-side
- Silently skip speculative decoding for vision models (unsupported
in llama.cpp server-context.cpp)
Frontend:
- Add speculative_type to TS API types
- Add speculativeType/loadedSpeculativeType to chat runtime store
with default value of "ngram-mod"
- Add On/Off toggle in Model settings section (GGUF only, hidden
for vision models), included in dirty check for Apply/Reset
- Wire speculative_type through model load request and response
- Restore speculative type state on page refresh/reconnect
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: remove server-side speculative decoding override
The backend was overriding speculative_type=None to "ngram-mod" for
non-vision GGUF models, which prevented users from disabling spec
decoding via the UI toggle. The frontend store already defaults to
"ngram-mod", so the backend fallback was redundant and blocked the
explicit "Off" setting.
* fix: use recommended ngram-mod params from llama.cpp docs
Update speculative decoding params to match the recommended values
from llama.cpp docs (docs/speculative.md):
--spec-ngram-size-n 24 (was 16, docs say small n not recommended)
--draft-min 48 (was 0)
--draft-max 64 (was 24, docs note MoEs need long drafts)
Also fix comment: ngram-mod uses ~16 MB (4M entries * 4 bytes),
not 4 MB.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add benchmark table and references to speculative decoding comment
Include speedup numbers from llama.cpp PRs #18471 and #19164 as an
inline comment so future readers understand the expected gains.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
studio/backend/core/inference/llama_cpp.py | 47 +++++++++++++++++++
studio/backend/models/inference.py | 12 +++++
studio/backend/routes/inference.py | 5 ++
.../src/features/chat/chat-settings-sheet.tsx | 39 ++++++++++++++-
.../chat/hooks/use-chat-model-runtime.ts | 9 +++-
.../chat/stores/chat-runtime-store.ts | 8 ++++
.../frontend/src/features/chat/types/api.ts | 3 ++
7 files changed, 121 insertions(+), 2 deletions(-)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 4da62e23c3..c84ac640df 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -109,6 +109,7 @@ class LlamaCppBackend:
self._supports_tools: bool = False
self._cache_type_kv: Optional[str] = None
self._reasoning_default: bool = True
+ self._speculative_type: Optional[str] = None
# KV-cache estimation fields (populated by _read_gguf_metadata)
self._n_layers: Optional[int] = None
self._n_kv_heads: Optional[int] = None
@@ -198,6 +199,10 @@ class LlamaCppBackend:
def cache_type_kv(self) -> Optional[str]:
return self._cache_type_kv
+ @property
+ def speculative_type(self) -> Optional[str]:
+ return self._speculative_type
+
# ── Binary discovery ──────────────────────────────────────────
@staticmethod
@@ -1055,6 +1060,7 @@ class LlamaCppBackend:
n_ctx: int = 4096,
chat_template_override: Optional[str] = None,
cache_type_kv: Optional[str] = None,
+ speculative_type: Optional[str] = None,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
) -> bool:
@@ -1315,6 +1321,46 @@ class LlamaCppBackend:
else:
self._cache_type_kv = None
+ # Speculative decoding (n-gram self-speculation, zero VRAM cost)
+ # ngram-mod: ~16 MB shared hash pool, constant memory/complexity,
+ # variable draft lengths. Helps most when the model repeats
+ # existing text (code refactoring, summarization, reasoning).
+ # For general chat with low repetition, overhead is ~5 ms.
+ #
+ # Benchmarks from llama.cpp PRs #18471, #19164:
+ # Scenario | Without | With | Speedup
+ # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
+ # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
+ # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
+ #
+ # Params from llama.cpp docs (docs/speculative.md):
+ # --spec-ngram-size-n 24 (small n not recommended)
+ # --draft-min 48 --draft-max 64 (MoEs need long drafts;
+ # dense models can reduce these)
+ # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
+ # ref: https://github.com/ggml-org/llama.cpp/pull/19164
+ # ref: https://github.com/ggml-org/llama.cpp/pull/18471
+ _valid_spec_types = {"ngram-simple", "ngram-mod"}
+ if speculative_type and speculative_type in _valid_spec_types:
+ if not is_vision: # spec decoding disabled for vision models
+ cmd.extend(["--spec-type", speculative_type])
+ if speculative_type == "ngram-mod":
+ cmd.extend(
+ [
+ "--spec-ngram-size-n",
+ "24",
+ "--draft-min",
+ "48",
+ "--draft-max",
+ "64",
+ ]
+ )
+ self._speculative_type = speculative_type
+ else:
+ self._speculative_type = None
+ else:
+ self._speculative_type = None
+
# Apply custom chat template override if provided
if chat_template_override:
import tempfile
@@ -1553,6 +1599,7 @@ class LlamaCppBackend:
self._reasoning_always_on = False
self._supports_tools = False
self._cache_type_kv = None
+ self._speculative_type = None
self._n_layers = None
self._n_kv_heads = None
self._n_heads = None
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 3094df4169..cf08ecbc12 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -48,6 +48,10 @@ class LoadRequest(BaseModel):
None,
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
)
+ speculative_type: Optional[str] = Field(
+ None,
+ description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
+ )
class UnloadRequest(BaseModel):
@@ -163,6 +167,10 @@ class LoadResponse(BaseModel):
None,
description = "Jinja2 chat template string (from GGUF metadata or tokenizer)",
)
+ speculative_type: Optional[str] = Field(
+ None,
+ description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
+ )
class UnloadResponse(BaseModel):
@@ -225,6 +233,10 @@ class InferenceStatusResponse(BaseModel):
None,
description = "Model's native context length from GGUF metadata (not capped by VRAM)",
)
+ speculative_type: Optional[str] = Field(
+ None,
+ description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
+ )
# =====================================================================
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index ced24c1d5f..30ff7da49c 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -179,6 +179,7 @@ async def load_model(
supports_reasoning = llama_backend.supports_reasoning,
reasoning_always_on = llama_backend.reasoning_always_on,
chat_template = llama_backend.chat_template,
+ speculative_type = llama_backend.speculative_type,
)
else:
if (
@@ -263,6 +264,7 @@ async def load_model(
n_ctx = request.max_seq_length,
chat_template_override = request.chat_template_override,
cache_type_kv = request.cache_type_kv,
+ speculative_type = request.speculative_type,
)
else:
# Local mode: llama-server loads via -m
@@ -275,6 +277,7 @@ async def load_model(
n_ctx = request.max_seq_length,
chat_template_override = request.chat_template_override,
cache_type_kv = request.cache_type_kv,
+ speculative_type = request.speculative_type,
)
if not success:
@@ -317,6 +320,7 @@ async def load_model(
supports_tools = llama_backend.supports_tools,
cache_type_kv = llama_backend.cache_type_kv,
chat_template = llama_backend.chat_template,
+ speculative_type = llama_backend.speculative_type,
)
# ── Standard path: load via Unsloth/transformers ──────────
@@ -652,6 +656,7 @@ async def get_status(
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
+ speculative_type = llama_backend.speculative_type,
)
# Otherwise, report Unsloth backend status
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index e5b0814343..550df2bf7c 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -280,6 +280,15 @@ export function ChatSettingsPanel({
}: ChatSettingsPanelProps) {
const isMobile = useIsMobile();
const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+ const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
+ const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType);
+ const loadedSpeculativeType = useChatRuntimeStore(
+ (s) => s.loadedSpeculativeType,
+ );
+ const currentModels = useChatRuntimeStore((s) => s.models);
+ const currentCheckpoint = params.checkpoint;
+ const currentModelIsVision =
+ currentModels.find((m) => m.id === currentCheckpoint)?.isVision ?? false;
const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
const ggufMaxContextLength = useChatRuntimeStore(
(s) => s.ggufMaxContextLength,
@@ -299,7 +308,8 @@ export function ChatSettingsPanel({
const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null;
const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
const ctxDirty = customContextLength !== null;
- const modelSettingsDirty = kvDirty || ctxDirty;
+ const specDirty = speculativeType !== loadedSpeculativeType;
+ const modelSettingsDirty = kvDirty || ctxDirty || specDirty;
const [customPresets, setCustomPresets] = useState(() =>
loadSavedCustomPresets(),
);
@@ -580,6 +590,32 @@ export function ChatSettingsPanel({
+ {!currentModelIsVision && (
+
+
+
+ Speculative Decoding
+
+
+ Speed up generation with no VRAM cost.
+
+
+
+
+ )}
{modelSettingsDirty && (