Commit graph

12 commits

Author SHA1 Message Date
Daniel Han
aab371a068
studio: tighten sandbox blocklist precision (bash, hf upload, NOFILE) (#5487)
* studio: tighten sandbox blocklist precision (bash, hf upload, NOFILE)

Three precision fixes in core/inference/tools.py. Same security
boundary; fewer false positives that broke legitimate sandbox use.

bash blocklist:
The per-token loop introduced in #5375 fired on any blocklist word in
any token position, so the entirely benign `grep -r curl .`,
`echo source the data`, and `ls /usr/bin/curl` were rejected with
"blocked command 'curl'". The position-anchored regex already covers
real command-position invocations, including `;rm`, `&&wget`, `$(rm)`,
`<(rm)`, backticked subshells, and `/usr/bin/sudo`. The token loop is
re-scoped: it only fires when the previous shlex token is a shell
separator (or at start of line), so split-quoting obfuscations like
`r''m -rf /` are still caught (shlex collapses them to a single
command-position token) while argument-position blocklist words pass
through. Trailing meta-chars glued to a shlex token (`rm;`) are
stripped before basename matching.

hf upload AST gate:
`_method_call_is_hf_upload` previously matched any method named
`upload_file` / `upload_folder` / `upload_large_folder` / `create_commit`
on any receiver, so paramiko.SFTPClient.upload_file, boto3.create_commit,
and similar non-HF SDK methods were rejected. The fallback now requires
an `import huggingface_hub` / `import hf_api` / `from huggingface_hub
import ...` somewhere in the same module. Fully-qualified
huggingface_hub.upload_file(...) calls are unchanged.

NOFILE env knob:
`RLIMIT_NOFILE = (1024, 1024)` was the only sandbox rlimit without an
env override. 1024 is below Linux's typical soft default and below
what multi-shard safetensors mmap chains need on Llama-3 70B-class
loads. Default is now 16384 with UNSLOTH_STUDIO_SANDBOX_NOFILE, parity
with the other rlimits.

15 new bash-blocklist-position tests pin both the false-positive
fixes and the still-blocked invariants (semicolon, &&, subshell,
backtick, split-quote, /usr/bin/ prefix, nested bash -c).
4 new hf-upload-import-gate tests pin both the false-positive
allowances and that HF-imported uses are still blocked.
1 new pin asserts the NOFILE env var is wired.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: cover command wrappers, find -exec, dynamic HF imports, NOFILE clamp

Reviewer follow-ups to the sandbox blocklist precision change.

Command-position scanner missed Bash command-prefix wrappers and inline
shell assignments. shlex tokenised `env curl`, `time curl`, `nohup rm`,
`FOO=bar curl`, `sudo rm`, etc. with the prefix at command position and
the real command at argument position, so the position-anchored check
returned set() while pre-PR's per-token scan caught them. Likewise the
position-anchored regex requires `^` or a shell separator before the
command, so `env curl` slipped through.

Reworked the scanner to track an expect_command flag plus a
prefix_pending flag:
  - assignments (FOO=bar) keep expect_command=True for the next token,
  - flags ('-oL', '--') keep it intact while prefix_pending is set,
  - numeric duration args ('timeout 1 cmd') skip without breaking
    expect_command,
  - known wrappers (env, command, builtin, exec, time, nohup, nice,
    setsid, stdbuf, timeout, ionice, chroot, sudo, doas, su, xargs)
    set prefix_pending so the wrapper's command is still checked,
  - shell separators now include `{`, `}`, `)`, `then`, `do`,
    `else`, `elif` so brace groups and if/then/while/do bodies are
    recognised as command positions.

Also lex with `shlex.shlex(punctuation_chars=";&|()`")` so split-quote
forms like `echo done; r''m -rf /tmp/x` and `echo done;r''m` tokenise
as `[..., ';', 'rm', ...]` and the command position check fires.

Added a small `find -exec CMD ... ;` / `-execdir CMD ... ;` pass so
`find . -exec rm -f {} +` and friends are caught even though the
direct token is at argument position to `find`.

Dynamic Hugging Face imports were treated as no-HF-in-scope. The
upload-method gate now also resolves `__import__('huggingface_hub')`,
`importlib.import_module('huggingface_hub')`, and bare
`import_module('huggingface_hub')` (via `from importlib import
import_module`) as HF imports, so HfApi().upload_file via dynamic
import is still blocked.

RLIMIT_NOFILE: setrlimit(NOFILE, (16384, 16384)) silently failed if
the parent's hard cap is below the requested value; the broad
except swallowed the OSError and left the sandbox at the parent's
default. Clamp the requested value to the inherited hard limit
before calling setrlimit.

Test cleanup: the existing test_cat_with_word_source_allowed had
`assert ... or True` so it could not fail; rewrote it to assert the
actual return value plus the two membership checks. Added
parametrised coverage for shell prefix wrappers, find -exec / xargs,
brace groups, if/then, while/do, split-quote command-name forms, and
dynamic HF import upload patterns.

Test:
  - python -m pytest studio/backend/tests/test_sandbox_tools.py -q
    -> 90 passed (was 67 before this commit)
  - full studio/backend/tests/ minus llama_cpp_load_progress_live and
    GPU CUDA_VISIBLE_DEVICES tests (pre-existing isolation flake)
    -> 1063 passed

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: catch bare-name HF upload calls in AST gate

`from huggingface_hub import upload_file; upload_file(...)` is a
canonical HF call shape that the previous Attribute-only check missed:
the bare-name call lands as ast.Name (not ast.Attribute), so the
fuzzy gate skipped it.

Extend _method_call_is_hf_upload to also match ast.Name when HF is in
scope. Same import-gating discipline as the Attribute branch, so
paramiko/boto3 and locally-defined `def upload_file(...)` helpers
without HF imports still pass.

Pins: 4 new TestHfUploadImportGate cases (upload_file/folder/create_commit
bare-name imports blocked; local upload_file without HF import allowed).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: scope HF uploads to sandbox-local literals; block env / token leaks

The previous gate dropped every HF upload call. Two refinements make it
precise enough to allow legitimate sandbox->HF uploads while still
catching credential / file exfil:

- path_or_fileobj / folder_path / create_commit operation paths must be
  sandbox-local relative-path literals (no '/', '~', drive letter, or
  '..' segments). Variable / dynamic paths are rejected.

- Any positional or keyword argument that statically resolves to
  os.environ / os.environ.get / os.getenv / bare getenv / subprocess
  shape readers is rejected (env-var exfil).

- token / hf_token / api_token / api_key / auth_token / access_token /
  password / secret kwargs are always rejected; sandbox env strips all
  parent credentials by construction, so any value here is hard-coded
  or lifted.

Recursive subtree walk in _reads_env_or_secret catches wrapper shapes
(str(os.environ), json.dumps(os.environ.items()), etc.).

Add TestSandboxEnvIsolation: pin that _build_safe_env builds the env
from a whitelist, not by stripping. Cover Linux/macOS/WSL/Windows
secret shapes. The whitelist is PATH / HOME / TMPDIR / LANG / TERM /
PYTHONIOENCODING (+ VIRTUAL_ENV / SystemRoot when applicable); HOME
points at the sandbox workdir, so HF / wandb / aws SDKs cannot reach
the operator's ~/.cache credentials.

Test classes added:
- TestHfUploadSandboxLocalPaths (relative literals allowed; absolute,
  drive-letter, '~', '..', mid-path traversal, dynamic vars, and
  open() of unsafe paths blocked, including create_commit recursion).
- TestHfUploadEnvAndSecretLeakBlock (os.environ subscript/get/getenv,
  bare getenv, subprocess.check_output, str(os.environ), token=,
  hf_token=, api_key=, and create_commit operations referencing env).
- TestSandboxEnvIsolation (no parent secret leaks into sandbox env).

131 tests in test_sandbox_tools.py pass.

* [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>
2026-05-18 00:01:17 -07:00
Daniel Han
0881a7a5d7
studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375)
* studio: contain export and dataset paths under their configured roots

resolve_under_root and resolve_dataset_path previously returned absolute
paths unchanged, so an authenticated client could supply
save_directory="/tmp/escape" (or any other absolute path) and have the
exporter drop adapter files anywhere the server user could write. This
turned up during a recent audit pass where an authenticated POST to
/api/export/export/lora with save_directory="/tmp/lora_escape_test"
returned 200 and wrote adapter_model.safetensors, adapter_config.json,
and tokenizer files under /tmp.

The fix is two-layered:

storage_roots.py adds an _assert_contained(resolved, root) helper that
runs after path resolution and rejects any result whose realpath does
not sit under realpath(root). resolve_under_root now rejects '..'
segments and null bytes outright, and only accepts absolute inputs when
they are already inside the configured root (internal call sites that
re-resolve a stored absolute path stay idempotent;
worker.py:resolve_output_dir(output_dir) etc. continue to work).
resolve_dataset_path picks up the same containment rule, scoped to the
three dataset roots.

models/export.py adds field_validator("save_directory", mode="before")
to ExportCommonOptions and ExportGGUFRequest so bad input fails fast at
422 with a clear message rather than a 500 deep inside the resolver.
The validator rejects empty/whitespace, null bytes, control chars,
strings longer than 255 chars, absolute paths, and '..' segments.

routes/export.py:_export_details now returns os.path.relpath(output_path,
exports_root()) so the Export Complete dialog and /api/models/loras no
longer leak the absolute install prefix to the UI; the basename is
used as a last-resort fallback.

Verified end to end:
- POST /api/export/export/lora {"save_directory":"/tmp/foo"} -> 422
  "save_directory must be a name or relative path under the export
  root; absolute paths are rejected". /tmp/foo is not created.
- "../../etc/escape" -> 422 "may not contain '..' segments".
- save_directory="my_subdir" -> still accepted (400 only because the
  test had no checkpoint loaded yet, not because of validation).
- Internal idempotent re-resolve via resolve_export_dir(absolute path
  that is already under exports_root) returns the same path unchanged.

* studio/sandbox: harden bash + python tool execution

The sandboxed Bash and Python tool channels in Chat ran with a thin
preexec hook (PR_SET_NO_NEW_PRIVS + RLIMIT_FSIZE only). Bash had a
small word blocklist; Python had an AST safety pass aimed at
signal-tampering and shell-escape primitives. An audit pass showed
several gaps that a tool-calling model could trigger inadvertently:

- bash curl/wget/nc reached AWS IMDSv2 and returned live STS
  credentials for the instance role.
- python "import socket; s.connect((169.254.169.254, 80))"
  reached the same endpoint regardless of the bash blocklist.
- "cat /etc/passwd" was blocked at the bash side (because "passwd"
  is in the blocklist), but "open('/etc/passwd').read()" in Python
  happily returned its contents.
- "chr(115)+chr(117)+chr(100)+chr(111)" style dynamic-arg
  construction slipped through the AST shell-escape check.
- The supervisor used proc.kill() on timeout, which only signals
  the immediate pid; bash-backgrounded children survived. A fork
  bomb could spawn for the full 300s timeout window.
- Session work directories under ~/studio_sandbox/<id>/ were
  created with default umask (0o755), so any other UID on the host
  could enumerate them.
- session_id sanitisation used a one-shot str.replace("..",""),
  which is non-iterative and a small footgun.

This commit takes a conservative middle path: the sandbox still
runs as the Studio UID with no namespace tricks where the kernel
disallows them, but every chokepoint is tightened.

_sandbox_preexec now:
- calls os.setsid() so children share a process group; the
  supervisor uses os.killpg(SIGKILL) on timeout/cancel so
  backgrounded children die with the parent (new _kill_process_tree
  helper, wired into _cancel_watcher and both _bash_exec /
  _python_exec timeout branches).
- calls os.umask(0o077) so files the child writes default to 0o600.
- applies PR_SET_PDEATHSIG=SIGKILL so an orphaned child dies if
  Studio exits.
- best-effort unshare(CLONE_NEWNET) for a private network namespace
  (failure is logged and swallowed; defense-in-depth is still in
  place via the bash blocklist and the AST checker below).
- sets RLIMIT_NPROC=10000 (tunable via UNSLOTH_STUDIO_SANDBOX_NPROC),
  RLIMIT_AS=8GB, RLIMIT_CPU=300, RLIMIT_NOFILE=1024. The 10k NPROC
  figure is chosen to sit well above the ~500 LWPs a healthy Studio
  + llama-server combination already uses while still capping a
  runaway fork bomb. NPROC counts LWPs per real UID, so a lower
  figure (e.g. 256) starves legitimate bash forks
  ("bash: fork: retry: Resource temporarily unavailable").

_get_workdir:
- rejects session_id that doesn't match [A-Za-z0-9_-]{1,64};
  non-matching values bucket into a shared "_invalid" dir.
- chmod 0o700 on both the workdir and on ~/studio_sandbox/ so
  other UIDs cannot read another session's contents.

_BLOCKED_COMMANDS_COMMON gains: doas, pkexec, halt, poweroff, curl,
wget, nc, ncat, netcat, socat, ssh, scp, sftp, rsync, eval, source.
The intent is to keep general bash usage working (echo, ls, pipes,
loops, for, head, etc.) while denying the obvious egress and
escalation paths.

The AST checker (_check_signal_escape_patterns) is split into the
existing shell/signal/loop checks plus a new narrow IO denylist:
- Always flag non-literal args to anything in _SHELL_EXEC_FUNCS,
  not just _STRING_SHELL_FUNCS. Closes the dynamic-arg bypass.
- Reject calls to socket.create_connection, socket.socket().connect,
  urllib.request.urlopen, http.client.HTTP*Connection, requests.*,
  httpx.* whose literal host argument is in a cloud-metadata
  denylist (169.254.169.254 + 169.254.* + 100.64.*, plus the
  GCP/Alibaba/ECS metadata hostnames and IPv6 link-local). Public
  hosts (example.com, huggingface.co, ...) still work. Dynamic
  hosts cannot be statically blocked; mitigated by the bash
  blocklist + the netns where the kernel allows it.
- Reject literal open("/etc/passwd"), /etc/shadow, /etc/sudoers,
  /etc/ssh/*, and /proc/<pid>/environ. Other files
  (/etc/os-release, /etc/hostname, /tmp/*, user dirs) still work.

The _check_code_safety summariser is updated to include the new
network_calls and sensitive_file_reads buckets in its error string.

Regression-checked: echo, sleep, ls /tmp, for loops, piped helpers
(echo a | tr a A), urllib.request.urlopen("http://example.com"),
socket.getaddrinfo("example.com",80), open("/etc/os-release"),
open("/tmp/...","w") all still succeed. curl, wget, nc, ssh, rm,
socket.create_connection(("169.254.169.254",80)),
open("/etc/passwd"), open("/proc/self/environ") all correctly
blocked.

* studio: rate-limit login, rotate refresh tokens, add logout, security headers, gate bootstrap injection

A pass over the auth surface found a cluster of related issues that this
commit closes together.

Login (routes/auth.py):
- Add an in-memory per-IP login rate limiter. Five failed POSTs to
  /api/auth/login inside a 60s window produce 429 with Retry-After.
  A successful login clears the bucket. Previously 30 wrong passwords
  in under one second was accepted as 30x 401, which combined with
  the (now fixed) admin-username leak from /api/auth/status made
  brute-force trivial against a small password.

Logout (routes/auth.py):
- New POST /api/auth/logout returns 204 and calls
  storage.revoke_user_refresh_tokens(subject) so the refresh token
  is no longer valid. Previously POST /api/auth/logout returned 405
  and there was no way to invalidate refresh tokens short of
  changing the password. Frontend session.ts already calls
  clearAuthTokens() to drop localStorage; the new endpoint lets the
  client also tell the server to revoke server-side state.

Refresh-token rotation (routes/auth.py + auth/storage.py):
- New storage.consume_refresh_token(token) atomically validates +
  deletes a refresh token, returning (username, is_desktop). The
  /api/auth/refresh handler now mints both a new access AND a new
  refresh token; the supplied token becomes invalid. Replaying a
  consumed refresh returns 401 "Invalid or expired refresh token".
  The previous refresh_access_token helper is left in place for
  callers that intentionally want the non-rotating shape; nothing
  in the route layer uses it now.

/api/auth/status no longer leaks default_username (models/auth.py +
routes/auth.py):
- AuthStatusResponse.default_username becomes Optional[str] with a
  None default; the handler always returns None. The frontend already
  hardcodes HIDDEN_LOGIN_USERNAME = "unsloth" (auth-form.tsx:82), so
  no UI change is required.

window.__UNSLOTH_BOOTSTRAP__ no longer auto-injects (main.py):
- _inject_bootstrap is now opt-in via the
  UNSLOTH_STUDIO_INJECT_BOOTSTRAP env var. The previous default
  (inject whenever requires_password_change is true) embedded the
  plaintext bootstrap password into the first-boot HTML for any
  caller that hit /, /change-password, or any unknown SPA path.
  Browser extensions and any XSS payload on the page could read it
  trivially. With the new gate the bootstrap password lives only in
  the auth/.bootstrap_password file (mode 0o600) where it has always
  been; users typing it into a current-password field is the right
  UX. routes/auth.py:change_password also clears
  app.state.bootstrap_password defensively.

Security headers + server fingerprint (main.py + run.py):
- New SecurityHeadersMiddleware adds Content-Security-Policy,
  X-Frame-Options: DENY, X-Content-Type-Options: nosniff,
  Referrer-Policy: no-referrer,
  Permissions-Policy: camera=(), microphone=(), geolocation=(),
  interest-cohort=(), and stamps server: unsloth-studio so the
  generic uvicorn banner no longer fingerprints the stack. The
  uvicorn.Config gains server_header=False so it stops emitting its
  own Server header.

/api/health minimisation (main.py):
- Unauthenticated GET /api/health returns just
  {"status":"healthy","timestamp":...} so load-balancer liveness
  probes keep working without leaking version, device_type,
  chat_only, desktop_protocol_version, or studio_root_id to
  arbitrary callers. A request that presents a valid Bearer token
  still gets the full diagnostic payload so internal launchers and
  sibling-Studio detection (which compares studio_root_id) keep
  working.

Verification:
- 30 wrong-password POSTs to /api/auth/login -> first 5 = 401, 6th
  through 30th = 429.
- POST /api/auth/logout with a fresh token -> 204. The matching
  refresh token then fails 401.
- Login -> R1; /api/auth/refresh with R1 -> new access + R2 (R2 !=
  R1); /api/auth/refresh with R1 again -> 401; /api/auth/refresh
  with R2 -> still succeeds once and rotates again.
- curl /api/auth/status -> default_username: null.
- curl http://127.0.0.1/ does not contain __UNSLOTH_BOOTSTRAP__.
- curl -I / shows CSP, X-Frame-Options: DENY,
  X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
  Permissions-Policy, and server: unsloth-studio.
- curl /api/health unauthenticated -> {status, timestamp} only.
  curl with Authorization: Bearer <valid> -> full payload.
- Existing /api/system, /api/models/list, /api/train/status,
  /api/inference/status, /api/auth/api-keys, login flow, SPA root
  all still return 200 after the changes (regression smoke).

* studio: add SecurityHeadersMiddleware, MaxBodyMiddleware, /recipes redirect, gate _inject_bootstrap, minimise /api/health

This commit lands the main.py-side changes that share a single
middleware-registration spot. They are kept together because every
change here is either (a) a top-level middleware definition that has
to be added next to LoggingMiddleware, or (b) a route handler at the
same file-level.

SecurityHeadersMiddleware (Content-Security-Policy, X-Frame-Options:
DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
Permissions-Policy, server: unsloth-studio). The previous responses
emitted no CSP, no XFO, no Referrer-Policy and were stamped
server: uvicorn.

MaxBodyMiddleware rejects POST/PUT/PATCH on the inference / dataset /
data-recipe / train / export prefixes when Content-Length exceeds
UNSLOTH_STUDIO_MAX_BODY_MB (default 100). The audit hit this by
attaching a 50 MB plain-text file to a chat message and watching
Studio base64-encode it into the JSON body; uvicorn has no enforced
cap so the only previous guard was the per-file 50 MB ceiling that
data-recipe upload routes already enforce. The new middleware extends
that ceiling to the OpenAI-compat path that the Chat attachments
flow through. Verified: a 200 MB JSON POST to /v1/chat/completions
returns HTTP 413 "Request body too large (209,715,264 bytes; max
104,857,600)". A small valid request continues to reach the handler.

_inject_bootstrap is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP.
The previous default was to inline window.__UNSLOTH_BOOTSTRAP__ =
{username, password} into the first-boot HTML whenever
requires_password_change was true, which exposed the plaintext
bootstrap password to any browser extension, page script, or LAN
caller on -H 0.0.0.0. The bootstrap password remains in the on-disk
.bootstrap_password file (mode 0o600) where it has always lived;
users typing it into a current-password field is the right UX.

/api/health unauthenticated returns {"status":"healthy","timestamp":
...} only; the previous payload (version, device_type, chat_only,
desktop_protocol_version, supports_desktop_auth, studio_root_id,
native_path_leases_supported) is preserved for callers that present
a valid Bearer token, so internal launchers and sibling-Studio
detection (which compares studio_root_id) keep working.

/recipes -> /data-recipes 308 redirect. The Data Recipes page lives
at /data-recipes; users typing /recipes hit the SPA catch-all and
saw "Not Found". The redirect also preserves any tail path, so
/recipes/<rest> -> /data-recipes/<rest>.

Verified end to end with curl: CSP / XFO / X-Content-Type-Options /
Referrer-Policy / Permissions-Policy all present on /, server header
is now unsloth-studio (uvicorn's own banner is suppressed via
server_header=False in run.py from the auth-batch commit). Followed
the /recipes redirect lands on the SPA HTML.

* studio: bound TrainingStartRequest hyperparameters at the schema level

POST /api/train/start accepted any value for learning_rate, batch_size,
max_steps, max_seq_length, warmup_steps, warmup_ratio, num_epochs,
save_steps, weight_decay, gradient_accumulation_steps, lora_r,
lora_alpha and lora_dropout, including -1, 0, 1e9, and non-numeric
strings like 'abc' or 'two' (which silently coerce to 0 in the
trainer). Probing showed the API returning 200 to learning_rate=-1
and batch_size=0; only max_steps had any partial clamping.

This commit adds field_validator on every numeric hyperparameter.
Bounds are chosen wide enough to span realistic single-host
configurations (B200 with 180 GB of memory comfortably fits the
upper end) while rejecting the values that always produce broken
training:

- learning_rate: parses str/float, requires 0 < lr < 1.0. Non-numeric
  input raises with "learning_rate must be parseable as float (got
  'abc')" instead of silently coercing to 0.
- batch_size: [1, 1024].
- gradient_accumulation_steps: [1, 4096].
- num_epochs: [1, 1000].
- max_steps: [1, 1_000_000].
- max_seq_length: [1, 131072].
- warmup_steps: [0, max_steps].
- warmup_ratio: [0.0, 1.0].
- save_steps: [0, 1_000_000].
- weight_decay: [0, 10] (typical 0..0.1).
- lora_r: [1, 512].
- lora_alpha: [1, 1024].
- lora_dropout: [0.0, 1.0).

Each validator names the offending field in its ValueError message
so the 422 response body identifies which input is bad. The
learning_rate validator returns its result as str (the schema field
type is str("2e-4") for backwards compatibility) so existing call
sites that float() the value continue to work.

Verified:
- learning_rate=-1 -> 422 "learning_rate must be > 0 (got -1.0);
  typical range is 1e-6 .. 1e-3".
- learning_rate='abc' -> 422 "must be parseable as float".
- batch_size=-1 / 0 / 999999 -> 422 "batch_size must be in [1, 1024]".
- batch_size='two' -> 422 (pydantic int parser).
- max_steps=0 / -5 -> 422 "must be a positive int".
- max_seq_length=200000 -> 422 "must be in [1, 131072]".
- warmup_ratio=2.5 -> 422 "must be in [0.0, 1.0]".
- lora_dropout=1.5 -> 422 "must be in [0.0, 1.0)".
- Valid request with learning_rate='2e-4', batch_size=1, max_steps=5
  passes validation and the training run starts as normal.

* studio: redact image-decode errors, clean checkpoint dirs on cancel, tolerate Stop-button + tool-result message shapes

Three small fixes that fall under "do not let the audit findings
become user-visible papercuts".

routes/inference.py - image-decode error redaction (the audit hit
this with a 0-byte / malformed / wrong-extension image upload). The
three image-normalise sites previously raised HTTPException(400,
detail=f"Failed to process image: {e}"). When PIL raised
UnidentifiedImageError(io.BytesIO(raw)) the message string included
"<_io.BytesIO object at 0x7e40a5d7bf60>", leaking both the Python
class name (confirming the PIL/io stack) and a heap address (mildly
useful for ASLR-bypass chaining if another memory-corruption bug is
ever found). Each site now catches UnidentifiedImageError and
returns the generic "Unsupported or corrupt image format"; the
fall-through generic except returns "Failed to process image". No
exception-repr is interpolated into a response body anywhere along
these paths.

core/training/training.py - checkpoint cleanup on cancel. When a
user clicks Cancel Training, the trainer flips _cancel_requested=True
and the supervisor force-terminates the subprocess. The trainer
writes checkpoint-<step> directories under output_dir every
save_steps; previously these survived the cancel and accumulated on
disk (the audit recorded ~67 MB stuck after a 200-step cancel with
save_steps=20). New helper _cleanup_cancelled_checkpoints(output_dir)
globs checkpoint-<int> entries and removes them. It is gated by a
realpath containment check against outputs_root() so it cannot
accidentally rmtree anything outside the configured outputs root.
force_terminate() invokes the helper after the subprocess join when
_cancel_requested is true. Stop-and-Save runs are unaffected because
that path keeps _cancel_requested=False.

models/inference.py - chat message shape tolerance. Two related
frontend interactions used to crash the request validator:

- After the Stop button truncates a generation, the frontend
  retained {role:"assistant", content:""} in the conversation
  history and replayed it on the next send. ChatMessage previously
  required role="assistant" to have non-empty content or tool_calls,
  so the next message returned 422 and the thread was permanently
  broken. The validator now normalises empty assistant content to
  None so the request round-trips and the trailing empty turn can
  be ignored downstream.

- The frontend's second-round tool POST drops the streamed
  tool_call_id, hitting the strict-spec check "role=tool requires
  tool_call_id". The validator now synthesises an opaque id
  (call_<8 hex>) when missing, so the request reaches the handler
  and the model's final summarising response gets generated. The
  proper fix lives in the frontend (carry the streamed id through
  the second POST) and will follow.

Verified end to end with curl: HTTP 400 (model not loaded) on both
the empty-assistant history shape and the tool-result-without-id
shape, instead of HTTP 422 from the schema validator.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten code comments from security-hardening pass

Trim verbose docstrings and inline finding references added in the
previous commits in this branch. Functionality unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: await get_current_subject in /api/health and make refresh-token consumption atomic

The /api/health auth probe called get_current_subject(creds) without
awaiting it. The coroutine object is truthy, so any caller presenting a
Bearer header (valid or not) received the full diagnostic payload
including version, device_type, studio_root_id, etc. Await the coroutine
and treat HTTPException as 'fall back to the minimal liveness payload'.

consume_refresh_token did SELECT then DELETE WHERE id under default
autocommit isolation. Two concurrent POST /api/auth/refresh requests
could both win the SELECT before either DELETE ran, defeating
single-use refresh-token rotation. Replace with a single
DELETE ... WHERE token_hash = ? AND expires_at >= ? RETURNING ...
statement so the validate-and-delete lands as one atomic op under
SQLite's write lock (3.45.1 supports RETURNING; min was 3.35).

* studio: enforce body cap on chunked uploads and drop unsafe-inline from script-src

MaxBodyMiddleware previously only inspected the declared Content-Length
header; clients omitting it or sending Transfer-Encoding: chunked
bypassed the cap and could still drive an OOM via the downstream
JSON / file readers on /v1/chat/completions, /api/inference, /api/data-recipe,
/api/datasets, /api/train, /api/export. Rewrite as a raw ASGI middleware
that drains and counts http.request frames, replies 413 once the running
total exceeds UNSLOTH_STUDIO_MAX_BODY_MB before invoking the FastAPI
handler, and replays the buffered body to downstream so route code that
calls request.json() / await request.body() works unchanged.

CSP previously included 'unsafe-inline' on script-src, which defeats the
main XSS protection. The frontend bundle does not need inline scripts;
the only inline <script> the backend ever emits is _inject_bootstrap,
which is opt-in via UNSLOTH_STUDIO_INJECT_BOOTSTRAP. Drop 'unsafe-inline'
from script-src by default; when _inject_bootstrap fires, generate a
per-response nonce, embed it on the inlined <script>, and have
SecurityHeadersMiddleware splice 'nonce-XXX' into the CSP for that one
response (the internal x-internal-script-nonce header is popped before
the response leaves the server). 'unsafe-inline' stays on style-src for
Vite-injected styles.

* studio: drop empty assistant sentinel before passthrough

ChatMessage._validate_role_shape normalises role="assistant", content=""
(the post-Stop sentinel emitted by the frontend) to content=None so the
in-process path can drop it via _extract_content_parts. The passthrough
path then ran m.model_dump(exclude_none=True), which strips the now-None
content key entirely, sending {"role":"assistant"} to llama-server / the
OpenAI-compat backend. That fails upstream and leaves the user without a
recoverable Stop->resume.

Add _drop_empty_assistant_sentinels and call it at both passthrough
message origins: _openai_messages_for_passthrough (covers
/v1/chat/completions and the Responses API which routes through it) and
the anthropic_messages_to_openai output before
_anthropic_passthrough_*. Assistant messages that carry only tool_calls
(no content) are preserved.

* studio/tests: cover audit-fix surfaces and rebase pre-existing tests

Adds and updates pytest coverage for the four bot-flagged audit fixes
landed earlier in this branch and rebases two pre-existing tests that
were broken by the relaxed-validator and /api/health auth-gate changes.

studio/backend/tests/test_middleware.py (new)
  MaxBodyMiddleware: small protected, large declared, unprotected
  passthrough, chunked-upload-over-cap rejection (the regression for
  the original Content-Length-only gap), and chunked-under-cap replay.
  SecurityHeadersMiddleware: script-src no longer carries
  'unsafe-inline', style-src still does, default headers
  (XFO/XCTO/Referrer-Policy/Permissions-Policy/server), and the
  internal x-internal-script-nonce header is consumed by the
  middleware and converted to 'nonce-XXX' in the CSP.
  /api/health: no auth -> minimal, invalid Bearer -> minimal
  (the await regression), valid Bearer -> full diagnostic payload.

studio/backend/tests/test_desktop_auth.py
  consume_refresh_token: second-call returns None, expired returns
  None, and a 64-thread concurrent pile-up against the same hash
  produces exactly one successful consumer (regression for the
  SELECT-then-DELETE race).
  test_health_response_reports_desktop_capability_fields: rebase
  against the new health_check(request) signature by going through
  TestClient with a real bearer instead of asyncio.run-ing the
  handler directly.

studio/backend/tests/test_openai_tool_passthrough.py
  Pin the new ChatMessage tolerance: assistant without content or
  tool_calls is tolerated (normalises content -> None), empty-string
  and empty-list assistant content normalise to None, and a missing
  / empty tool_call_id on role='tool' is synthesised as call_<hex>
  rather than raising. Tests for _drop_empty_assistant_sentinels
  cover the three drop shapes (empty string, empty list, missing
  content key), preservation of assistant text and tool_calls-only
  messages, and end-to-end through
  _openai_messages_for_passthrough.

studio/backend/main.py
  SecurityHeadersMiddleware.dispatch used response.headers.pop(...)
  for the nonce-header handoff; Starlette's MutableHeaders has no
  pop. Read-then-del so the internal handoff header is still
  stripped before the response leaves the server.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio/tests: rebase three more pre-existing CI tests against this branch

CI on PR #5375 was red on three tests that were tuned for behaviour
predating this branch. Updates each so the assertions match what the
audit fixes intentionally changed; no production code touched.

studio/backend/tests/test_trained_model_scan.py
  test_scan_trained_models_includes_lora_and_full_finetune_outputs
  passed an absolute tmp_path through scan_trained_models, which now
  runs resolve_output_dir / _assert_contained against outputs_root().
  Repoint outputs_root() at tmp_path via monkeypatch so the fixture
  dirs land under the configured root and the realpath containment
  check passes.

tests/test_studio_install_workspace_guard.py
  test_health_endpoint_exposes_studio_root_id_not_raw_path read
  the first 1500 bytes after @app.get("/api/health") and asserted on
  the studio_root_id literal. The handler grew (unauth short-circuit
  + await dependency gate) and the literal slid past the byte window.
  Replace the fixed window with a slice up to the next top-level
  @app.* decorator so the test surveys the whole handler regardless
  of size.

tests/studio/studio_api_smoke.py
  The "login burst (5x wrong pw) -> 401 each" assertion was tagged
  "When/if we add one, this assertion updates in the same PR." We
  added the per-IP rate-limit in routes/auth.py
  (_LOGIN_MAX_FAILS=5/60s) but missed the assertion update. Rewrite
  the burst probe to observe the new invariant: at least one 401,
  eventual transition to 429, and Retry-After present on the 429.
  Adds a small _login_with_headers helper since the existing login()
  helper drops response headers.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ci(studio-ui): set UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 for Playwright Studios

The Chat UI Playwright test drives the first-boot change-password
form, which (per playwright_chat_ui.py step "1. Change-password
through the UI") pre-seeds the hidden current_password field from
window.__UNSLOTH_BOOTSTRAP__. That global is only emitted when the
backend's _inject_bootstrap path fires, which since the security
pass on this branch is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP
and defaults to off. Without the global, the React form's
current_password validator never satisfies, the submit button stays
disabled, and the composer.wait_for() probe times out on
/change-password.

Re-enable injection only for the CI Studios that drive the chat UI
across linux/mac/windows. Production deployments are unaffected: the
env var has to be explicitly opted into, and the on-disk
auth/.bootstrap_password remains the source of truth for human users
typing the password in by hand.

Covers all eight Studio launch sites: the primary chat-ui boot and
the "extra UI tests" boot for each of the three OSes, plus the
pipeTransport JSON-crash retry relaunches in the macOS workflow that
re-spawn Studio mid-job.

A follow-up frontend PR will add a visible current_password input so
the form satisfies its own validator without needing the bootstrap
auto-fill at all; once that lands this CI knob can come back out.

* studio/sandbox: drop unshare(CLONE_NEWNET); add trusted-host allowlist; block sandbox file uploads; raise CPU rlimit default to 600 s

CLONE_NEWNET inside _sandbox_preexec silently killed every outbound
HTTP request from sandboxed Python whenever the kernel allowed
unprivileged user namespaces. requests.get('https://huggingface.co'),
urllib.request.urlopen('https://en.wikipedia.org/wiki/...'),
socket.connect(('arxiv.org', 443)) all failed despite the AST visitor
intending to allow them. The bash blocklist (curl / wget / nc / ssh /
scp / sftp / rsync / socat / eval / source) plus the AST-level
metadata-host denylist still carry the network policy after this
change; CLONE_NEWNET was redundant with both.

Add _TRUSTED_PUBLIC_HOST_LITERALS + _TRUSTED_PUBLIC_HOST_SUFFIXES
(~100 informational hosts: Wikipedia language subdomains, Wikimedia,
Wikidata, Google search, Bing, DuckDuckGo, HuggingFace, GitHub,
raw.githubusercontent.com, arXiv, StackOverflow / Stack Exchange,
MDN, docs.python.org, PyTorch / TensorFlow / NumPy / pandas docs,
pypi / files.pythonhosted.org / npmjs / crates.io, ReadTheDocs,
arXiv, Britannica, BBC / Reuters / Nature / Science, NASA / CDC /
NIH / WHO open data, api.weather.gov). The visitor now blocks
literal hosts that are neither metadata nor trusted with a short
LLM-readable string so the model can retry with an allowed source
instead of choking on a multi-line error.

Block upload-shape calls regardless of host: requests.post / put /
patch / delete / request with files= or data=open(...) /
data=bytes_literal; httpx equivalents; urllib.request.urlopen /
Request with data=...; HuggingFace upload_file / upload_folder /
upload_large_folder / create_commit (module-level FQ paths AND
method-name match on any receiver). Message: "Blocked: file upload
disallowed in sandbox".

Bump UNSLOTH_STUDIO_SANDBOX_CPU_S default 300 -> 600 s so long
agentic chains that span multiple tool calls don't get SIGXCPU'd
mid-stride. Env-var override path is unchanged.

Host normalisation now strips trailing dot, userinfo @, and explicit
port before allowlist / denylist comparison so trailing-DNS-dot,
userinfo-smuggling, and explicit-:443 URLs are decided correctly.

* studio: raise default request-body cap from 100 MB to 500 MB

UNSLOTH_STUDIO_MAX_BODY_MB default goes 100 -> 500 to comfortably
cover vision + audio + multi-recipe-batch JSON payloads. The
MaxBodyMiddleware stream-counting logic from this branch's earlier
06ec088 already handles chunked bodies up to the new cap; env-var
override path is unchanged for callers that want a tighter limit.

* studio/auth: restore /api/auth/status.default_username to 'unsloth'

This branch's earlier b39e9a4 changed default_username to None on the
public /api/auth/status endpoint so the username field didn't leak to
unauthenticated callers. In practice this regressed third-party
clients (and the in-tree React login form's pre-fill UX) without
adding meaningful security: the bootstrap password is the actual
secret, and the username 'unsloth' is the documented default.

Pin default_username to storage.DEFAULT_ADMIN_USERNAME ('unsloth')
and tighten the response model so the field is required rather than
Optional. Anyone who needs anonymisation can still reach for an
allow-list deployment with auth disabled.

* studio/training: raise max_seq_length / batch_size / lora_r / lora_alpha caps

This branch's 7102815 introduced field validators with conservative
caps. The follow-up loosens them so long-context experiments and
high-rank LoRA exploration aren't gated at the schema layer:

  _MAX_BATCH_SIZE   1024     -> 4096
  _MAX_SEQ_LENGTH   131_072  -> 2_000_000   (2M tokens)
  lora_r cap        512      -> 16_384      (_MAX_LORA_R)
  lora_alpha cap    1024     -> 32_768      (_MAX_LORA_ALPHA)

_MAX_GRAD_ACCUM / _MAX_STEPS / _MAX_EPOCHS / lora_dropout /
warmup_ratio / weight_decay are unchanged. Hardware (VRAM, host
RAM, kernel launch latency) is now the binding constraint at the
new caps, which is the correct ordering -- the validator stays a
sanity check on -1 / 0 / 'abc' style garbage, not a usability gate.

* studio/tests: cover sandbox allowlist + upload block + raised training caps

studio/backend/tests/test_sandbox_tools.py (new):
  TestMetadataHostDenylist     -- short "Blocked: cloud-metadata host"
                                  message on AWS IMDS, GCP metadata,
                                  Alibaba ECS, AWS IPv6 IMDS, 169.254/16.
  TestTrustedHostAllowlist     -- Wikipedia (any language subdomain),
                                  Google, DuckDuckGo, HF, raw GitHub,
                                  arXiv, StackOverflow / family,
                                  MDN, docs.python.org, pypi, BBC,
                                  api.weather.gov, NumPy / PyTorch docs.
  TestUntrustedHostBlock       -- example.com / random unlisted host
                                  rejected with the short "Blocked: host
                                  not in sandbox allowlist; use an
                                  allowed informational source" message.
                                  Dynamic URLs (computed var) still pass
                                  -- documented limit of static analysis.
  TestHostNormalization        -- trailing dot, explicit :443, uppercase,
                                  userinfo-@-smuggle all decided
                                  correctly without false-block /
                                  false-pass.
  TestUploadDenylist           -- requests / httpx / urllib.urlopen with
                                  files= / data=open / data=bytes,
                                  HfApi().upload_file / upload_folder /
                                  create_commit, module-level
                                  huggingface_hub.upload_folder. POST
                                  json= to trusted host still passes.
  TestSandboxCpuRlimitDefault  -- pin UNSLOTH_STUDIO_SANDBOX_CPU_S=600
                                  default and confirm CLONE_NEWNET
                                  source line is gone.
  TestMaxBodyDefault           -- pin UNSLOTH_STUDIO_MAX_BODY_MB=500
                                  default.

studio/backend/tests/test_studio_train_validation.py (new):
  Pin at-cap-accepts / over-cap-rejects boundaries for
  max_seq_length=2_000_000, batch_size=4_096, lora_r=16_384,
  lora_alpha=32_768 so a future regression that tightens them back
  without explicit user opt-in is caught.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten code comments across the security-hardening pass

* studio: always inject bootstrap credentials on first boot

The UNSLOTH_STUDIO_INJECT_BOOTSTRAP gate added an extra
terminal-to-browser copy-paste on every fresh install. In practice
the LAN credential leak it guarded against is narrow: the password
is one-time, the user rotates it on the very next click, the
default Studio bind is 127.0.0.1, and -H 0.0.0.0 already exposes
the entire API surface. Drop the gate so the inject fires whenever
a bootstrap password is still pending. The CSP nonce wiring stays
in place; the inline script remains the only inline script the
backend ever emits.

The three Playwright UI smoke workflows lose their
UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 lines along with the explanatory
comment blocks since the inject now happens by default.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
2026-05-13 06:12:18 -07:00
Daniel Han
2c73ab7871
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>
2026-04-03 13:33:42 -07:00
Neodon
c027ec192e
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
2026-04-03 11:44:22 -07:00
Daniel Han
c8d311a053
feat(studio): display images from Python tool execution in chat UI (#4778)
* feat(studio): display images from Python tool execution in chat UI

When the model calls the Python tool to create a matplotlib plot or
other image file, the image now displays inline in the chat output
instead of being invisible to the user.

Backend:
- Detect new image files (png/jpg/gif/webp/bmp) after Python subprocess
  completes by diffing os.listdir before/after execution
- Append __IMAGES__ sentinel to tool result for frontend consumption
- Strip sentinel before injecting result into LLM context (role: tool)
  so the model never sees file paths
- Add GET /sandbox/{session_id}/{filename} endpoint with JWT auth
  (header or query param), path traversal protection, extension
  allowlist, realpath containment check, and nosniff header

Frontend:
- Parse __IMAGES__ sentinel in tool_end SSE events, create structured
  result with text/images/sessionId
- Render <img> tags in Python tool UI pointing at the sandbox endpoint

Also fixes a bug where SyntaxError in user code was misreported as
"unsafe code detected" instead of showing the actual Python traceback.
The _check_code_safety function now lets SyntaxError pass through to
the subprocess for a proper error message.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): improve SVG detection and strip XML preamble

Handle <?xml ...?> declarations before <svg> tags in code fences,
strip XML declaration from SVGs before data URI rendering, and
update the sloth suggestion prompt to request showing code.

* fix(studio): persist parentId so retries survive reload

The append() handler was destructuring only { message } from
ExportedMessageRepositoryItem and discarding parentId. When loading
a saved thread, load() used ExportedMessageRepository.fromArray()
which chains all messages sequentially, flattening retry branches
into a linear list.

Now append() writes parentId to the MessageRecord, and load()
reconstructs the tree when parentIds are present. Old threads
without parentId fall back to the existing fromArray() behavior.

* fix(studio): address review findings for image display and retry persistence

Image detection:
- Use mtime comparison instead of filename-only diff so overwritten
  files (e.g. plt.savefig("chart.png") called twice) are detected

Sentinel parsing:
- Use rsplit/lastIndexOf instead of split/indexOf so user code that
  prints __IMAGES__: does not collide with the backend sentinel

Mixed legacy/new threads:
- For old messages without a stored parentId, infer sequential parent
  from the previous message instead of null, preventing multiple roots

Sandbox endpoint:
- Change Cache-Control from "public, max-age=3600" to "private,
  no-store" since these are authenticated responses

---------

Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-02 05:08:16 -07:00
Daniel Han
76cb48be0b
fix: studio web search SSL failures and empty page content (#4754)
- Fix SSL handshake failures (SSLV3_ALERT_HANDSHAKE_FAILURE, CERTIFICATE_VERIFY_FAILED) when fetching HTTPS pages by introducing _PinnedHTTPSConnection that separates TCP connect (to pinned IP) from TLS handshake (with real hostname for SNI/cert verification)
- Fix SSRF DNS-rebinding vulnerability: previous impl swapped conn.host before connect(), causing fresh DNS resolution; new subclass keeps TCP pinned to validated IP
- Fix SPA/JS-rendered doc sites returning empty content by rotating real browser User-Agents (Chrome/Firefox/Safari)
- Strip nav/footer from HTML-to-Markdown output so article content is not buried under navigation chrome
- Increase raw fetch cap from 64KB to 512KB so SSR article content is reached on GitBook/Docusaurus/Next.js pages
- Fix IPv6 address bracketing in URL netloc construction
- Hoist SSL context, handler classes, and stdlib imports to module level (created once, not per-call)
- Use consistent UA across redirect hops to avoid breaking session-aware bot detection
2026-04-01 06:12:02 -07:00
Daniel Han
9a8b622306
Studio: simplify tool-call dedup and replace html2text with builtin converter (#4722)
* Simplify tool-call dedup: drop hashlib, inline helpers

The duplicate tool-call detector only compares calls within a single
request from the same JSON parser, so dict key order is guaranteed
identical for identical calls (Python 3.7+ insertion-ordered dicts).

- Replace hashlib.md5(json.dumps(...)) with name + str(args)
- Inline _tool_call_key, _is_duplicate_call, _record_tool_call
  since each was a one-liner used once
- Remove unused hashlib import

* Remove tool_calling_benchmark_results.md from repo

* Replace html2text with builtin HTML-to-Markdown converter

Drop the external html2text (GPL-3.0) dependency and its regex
fallback. Add _html_to_md.py (~190 lines, stdlib only) using
html.parser.HTMLParser that handles headings, links, bold/italic,
lists, tables, blockquotes, code blocks, and entity decoding.
Strips script/style/head tags entirely.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use json.dumps(sort_keys=True) for tool-call dedup key

str(dict) is sensitive to insertion order, so semantically identical
calls with different key ordering would bypass duplicate detection.
Switch to json.dumps with sort_keys=True for a canonical representation.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Revert dedup key to str(arguments)

json.dumps(sort_keys=True) is unnecessary here -- the arguments dict
always comes from the same JSON parser within a single request, so
key insertion order is deterministic (Python 3.7+).  str() is faster
and sufficient for consecutive-call dedup.

* Address review comments on _html_to_md.py

- Remove "hr" from _BLOCK_TAGS so the dedicated hr handler is reachable
- Prefix all newlines with ">" inside blockquotes (multi-line support)
- Emit full ![alt](url) for images instead of alt text only
- Replace newlines with spaces inside table cells
- Track header cells per-row (_row_has_th) instead of last-cell-only
- Strip trailing tabs in addition to spaces in cleanup regex

* Fix blockquote rendering, truncated-HTML buffer flush, and dedup key canonicalization

_html_to_md.py:
- Rewrite blockquote handling with stack-based buffer approach so nested
  blockquotes, pre blocks inside blockquotes, and multi-paragraph quotes
  all render correctly with proper "> " prefix on every line.
- Add flush_pending() to recover content from truncated HTML where closing
  tags are missing (common when _fetch_page_text caps the download size).
  Flushes open <a>, <td>, <pre>, and blockquote buffers.
- Skip <img> tags to match prior html2text ignore_images=True behavior
  and avoid data-URI amplification consuming the output budget.
- Collapse all whitespace (including newlines) in non-pre content per
  standard HTML whitespace rules: \s+ -> single space.
- Escape pipe characters in table cell content to prevent column breakage.
- Emit separator row after the first row for tables without <th> headers.
- Guard against IndexError on _ol_counter for orphan <li> elements.
- Normalize CRLF line endings before parsing.

llama_cpp.py:
- Restore canonical dedup key with json.dumps(sort_keys=True) so that
  semantically identical tool calls with different JSON key order are
  correctly detected as duplicates.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix table optional end tags, inline code whitespace, and link text normalization

_html_to_md.py:
- Extract _finish_cell() and _finish_row() helpers to handle HTML tables
  that omit optional </td>, </th>, or </tr> end tags. This is valid HTML
  and common on real web pages -- previously the parser would silently
  drop earlier cells and entire rows.
- Call _finish_cell()/_finish_row() from handle_starttag for <tr>/<td>/<th>,
  handle_endtag for </tr>/<td>/<th>/<table>, and flush_pending() so all
  three paths (normal close, implicit close, truncated HTML) use the same
  row-finalization logic including header separator emission.
- Add _in_inline_code flag so handle_data() preserves literal whitespace
  inside <code> spans instead of collapsing it. Source like
  <code>pip  install   unsloth</code> now correctly renders as
  `pip  install   unsloth` rather than `pip install unsloth`.
- Extract _finish_link() helper that normalizes accumulated link text with
  \s+ -> single space before building the Markdown link. Prevents block-
  level content inside <a> tags (e.g. <a><div>one</div><div>two</div></a>)
  from producing multiline [one\n\ntwo](href) link labels.
- Empty blockquotes now produce no output instead of a stray ">".
- Remove unused _bq_depth field (all routing uses _bq_stack).
- Flush open cells and rows in handle_endtag("table") for robustness.

* Support <ol start=N>, <dl>/<dt>/<dd>, and preserve code block whitespace

_html_to_md.py:
- Honor <ol start="N"> attribute so ordered lists preserve their original
  numbering instead of always restarting from 1. Important for docs/tutorials
  that continue numbering across sections.
- Add dl, dt, dd to _BLOCK_TAGS so definition lists (common on MDN, Python
  docs, Django docs) produce separated text instead of concatenated blobs.
- Rewrite _cleanup() to be fence-aware: content inside fenced code blocks
  is now preserved verbatim (intentional blank lines in <pre> content are
  no longer collapsed). Outside code blocks, blank runs are limited to one
  and trailing whitespace is stripped.
- Fix _prefix_blockquote() to strip trailing whitespace before collapsing
  blank lines, preventing the "\n\n \n\n" pattern from sneaking through.

* Suppress whitespace-only text nodes between table structural elements

Indented HTML tables (nearly all real-world pages) produce whitespace
text nodes between <table>, <tr>, </tr> etc. that land in the output
as leading spaces before table rows, breaking Markdown table alignment.

Skip whitespace-only text nodes when inside a table but not inside a
cell, so indentation from source HTML does not leak into the output.

* Revert dedup key to str(arguments) with explanatory comment

json.dumps(sort_keys=True) is unnecessary overhead here: arguments
always comes from json.loads on model output within a single request,
so dict insertion order is deterministic in Python 3.7+. A repeated
call from the model produces the same JSON, which parses to the same
dict repr. str() avoids re-serialization on every tool call.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-31 06:15:18 -07:00
Daniel Han
e159b93b97
studio: improve GGUF tool calling accuracy and reliability (#4700)
* studio: improve GGUF tool calling accuracy and reliability

- Add URL fetching to web_search tool so models can read full page
  content instead of only getting search snippets. Uses html2text for
  clean markdown conversion with regex fallback.
- Inject current date and behavioral guidance (URL fetch workflow,
  no repeated queries, use code for data processing) into the
  tool-use system prompt.
- Append error recovery nudge to tool results that indicate failure,
  helping small models avoid looping on the same broken call.
- Strip leaked <tool_call> XML from assistant messages in conversation
  history and from the outgoing SSE stream.
- Raise default max tool iterations from 10 to 25 across backend,
  model schema, and frontend defaults.
- Increase _MAX_PAGE_CHARS from 4k to 16k so fetched pages contain
  enough content for the model to extract useful information.
- Add "IMPORTANT: These are only short snippets" hint to search
  results so models know to fetch full pages when needed.

Tested with Qwen3.5-4B-GGUF (UD-Q4_K_XL), 10 runs before/after:
- XML leaks in responses: 10/10 -> 0/10
- URL fetch usage: 0 -> 4/10 runs
- Runs producing actual correct answers: 0/10 -> 2/10
- Average tool calls per query: 5.5 -> 3.8 (more efficient)
- Average response time: 12.3s -> 9.8s

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add tool calling benchmark results across model sizes and quants

Tested 16 configurations (4 models x 2 quants x 2 KV cache types)
with 10 runs each on NVIDIA B200.

Best config: 27B UD-Q4_K_XL + bf16 KV -- 6/10 runs found all 4
correct songs, 0 XML leaks, 131s average response time.

* Add duplicate tool-call detection and final-answer synthesis

When the model repeats the exact same tool call (same name + arguments)
twice in a row, skip execution and return a redirect message telling it
to try a different approach. This prevents the 8x-repeated-query loops
observed on 27B and 35B models.

When the tool iteration cap (25) is reached, inject a "provide your
final answer now" message before the final streaming pass. This lets
the model synthesize a useful answer from everything it gathered
instead of being silently cut off.

Tested on Qwen3.5-27B UD-Q4_K_XL (10 runs):
- Repeated query runs: 4/10 -> 2/10
- Cap hits: 1/10 -> 0/10
- All 4/4 accuracy: 5/10 -> 7/10

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix CodeQL alert: handle whitespace in script/style closing tags

The regex fallback for HTML stripping did not match closing tags
with whitespace before the angle bracket (e.g. </script >).
Use \s* before > in both script and style patterns.

* Address reviewer findings: SSRF, timeout crash, XML regex, dedup

- SSRF: resolve hostname via getaddrinfo and reject private, loopback,
  link-local, multicast, and reserved addresses before fetching
- Timeout: handle timeout=None (unlimited mode) in URL fetch path
  by defaulting to 60s instead of crashing on min(None, 60)
- Download cap: read at most max_chars*4+1 bytes instead of the
  full response body before truncating
- XML regex: match both <tool_call> and <function=...> markup in
  the history/stream cleanup (inference.py)
- CodeQL: use [^>]* in closing script/style tags to handle any
  whitespace or attributes before >
- Dedup: track whether each tool call failed so retries after
  transient errors are allowed; only block consecutive identical
  calls that both succeeded
- Final-answer synthesis: guard on max_tool_iterations > 0 so
  callers who disable tools do not get a false "used all calls" turn

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix redirect SSRF, SSE streaming regression, dedup off-by-one

- SSRF redirect bypass: disable auto-redirect in urllib, manually
  follow up to 5 hops with host validation at each step. Prevents
  public URLs from redirecting to loopback/private targets.
- SSE streaming: track prev_text on the raw cumulative and strip
  XML from the delta only, so completed tool_call tags do not cause
  the cumulative to shrink and drop trailing real text.
- Dedup off-by-one: check the immediately previous call (window=1)
  instead of requiring 2 matching history entries, so the second
  identical successful call is blocked rather than the third.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix redirect HTTPError handling and tighten error prefixes

- Redirect fix: urllib raises HTTPError (not a normal response) when
  the redirect handler returns None. Catch HTTPError for 3xx codes
  and extract the Location header from the exception object.
- Error prefixes: remove overly broad "No " prefix that matched
  "No results found." (a valid empty-search outcome, not an error).
  Replace with specific prefixes like "Blocked:", "No query provided",
  "Failed to resolve". This ensures empty search results are correctly
  classified as non-errors for duplicate-call tracking.

* Fix SSE cross-chunk XML leaks, cleanup review findings

- SSE streaming: sanitize the full cumulative text before diffing
  against the previous sanitized snapshot, so XML tags that span
  chunk boundaries are stripped correctly. The previous delta-based
  approach leaked split tags.
- DRAINING fallback: use _strip_tool_markup() helper instead of a
  manual regex that only handled <tool_call> but not <function=...>.
- Move hashlib import, _TOOL_XML_RE compile, and datetime import to
  module level per style guide.
- Remove unused _hit_tool_cap variable.

* Fix DNS rebinding, charset detection, HTTPError handling, dedup double-record

- DNS rebinding: resolve hostname once via getaddrinfo, pin the
  returned IP, rewrite the URL to connect to the pinned IP with
  a Host header. Each redirect hop re-resolves and re-validates.
  Closes the TOCTOU window between validation and connection.
- Charset: use resp.headers.get_content_charset() instead of
  hardcoding utf-8, so pages with other encodings decode correctly.
- HTTPError: return descriptive "HTTP {code} {reason}" instead of
  re-raising into a generic "Search failed" message.
- Dedup: remove redundant _record_tool_call in the duplicate branch;
  the single call at the end of the loop handles all cases.

* [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>
2026-03-31 03:06:44 -07:00
Daniel Han
8582ce3e9c
Fix studio chat crash on Mac: vendor check_signal_escape_patterns (#4431)
* Fix studio crash on Mac: vendor check_signal_escape_patterns from unsloth_zoo

Vendor the `check_signal_escape_patterns` function from
`unsloth_zoo.rl_environments` directly into `tools.py`. The function is
pure Python (only uses stdlib `ast`) and has zero GPU dependencies, but
importing it from unsloth_zoo triggers `unsloth_zoo.__init__` which calls
`get_device_type()` at module scope -- raising NotImplementedError on
Apple Silicon Macs.

By vendoring the code, the safety checks still run on all platforms
(Mac, Linux, Windows) without needing unsloth_zoo at all.

* [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>
2026-03-18 09:10:13 -07:00
Daniel Han
9c95148045
Fix tool call parsing, add tool outputs panel and UI improvements (#4416)
* Add elapsed timer to tool status pill in Studio

Show a count-up seconds timer (0s, 1s, 2s, ...) next to the tool status
text in the composer area. Helps users gauge how long a tool call (web
search, code execution) has been running. Timer resets when a new tool
starts and disappears when all tools finish.

* Fix tool call parsing, add tool outputs panel and reasoning copy button

Backend:
- Rewrite tool call XML parser to use balanced-brace JSON extraction
  instead of greedy regex, fixing truncation on nested braces in
  code/JSON arguments
- Handle optional closing tags (</tool_call>, </function>, </parameter>)
  that models frequently omit
- Support bare <function=...> tags without <tool_call> wrapper
- Strip tool call markup from streamed content so raw XML never leaks
  into the chat UI
- Use a persistent ~/studio_sandbox/ working directory for tool
  execution so files persist across calls within a session
- Emit tool_start/tool_end SSE events so the frontend can display
  tool inputs and outputs

Frontend:
- Add collapsible "Tool Outputs" panel below assistant messages showing
  each tool call's input and output with copy buttons
- Add copy button to reasoning blocks
- Add elapsed timer to tool status pill
- Update project URLs in pyproject.toml (http -> https, add docs link)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add interactive HTML preview with fullscreen toggle for code blocks

HTML code fences now render an interactive sandboxed iframe preview
below the syntax-highlighted code, similar to how SVG fences show
an image preview. The iframe uses sandbox="allow-scripts" to allow
JavaScript execution while blocking access to the parent page.

Includes a fullscreen toggle (enlarge/minimize button) that expands
the preview into a viewport overlay, dismissible via button, Escape
key, or backdrop click. A streaming placeholder prevents partial
HTML from rendering mid-stream.

* Add tool call settings: auto-heal toggle, max iterations, timeout

Add three user-configurable tool call settings to the Studio Settings panel:

- Auto Heal Tool Calls: toggle to control fallback XML parsing of malformed
  tool calls from model output (default: on)
- Max Tool Calls Per Message: slider 0-40 + Max to cap tool call iterations
  per message (default: 10)
- Max Tool Call Duration: slider 1-30 minutes + Max to set per-tool-call
  execution timeout (default: 5 minutes)

All settings persist to localStorage and flow through the full stack:
frontend store -> API request -> Pydantic model -> route -> llama_cpp -> tools.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tool call timeout: respect no-limit and apply to web search

- Use a sentinel to distinguish timeout=None (no limit) from the default
  (300s). Previously None was silently replaced with _EXEC_TIMEOUT.
- Pass the configured timeout to DDGS() for web searches so the setting
  applies uniformly to all tool types.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add input validation bounds and per-thread sandbox isolation

- Add ge=0 constraint to max_tool_calls_per_message (rejects negative values)
- Add ge=1 constraint to tool_call_timeout (minimum 1 second)
- Thread session_id from frontend through backend to tool execution
- Scope sandbox directories per conversation: ~/studio_sandbox/{thread_id}/
- Backwards compatible: API callers without session_id use ~/studio_sandbox/

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix non-monotonic streaming and Python temp script path

- Split tool markup stripping into closed-only (mid-stream) and full
  (final flush) to prevent cumulative text from shrinking mid-stream
- Enforce monotonicity: only emit when cleaned text grows, so the
  proxy's delta logic (cumulative[len(prev_text):]) never breaks
- Place Python temp scripts in the sandbox workdir instead of /tmp so
  sys.path[0] points to the sandbox and cross-call imports work

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Sanitize session_id to prevent path traversal in sandbox

Strip path separators and parent-dir references from session_id before
using it as a directory name. Verify the resolved path stays under
~/studio_sandbox/ as a second guard.

* feat(chat): proper assistant-ui tool call UIs with sources

Replace custom metadata-based ToolOutputsGroup with native assistant-ui
tool-call content parts. Backend SSE tool_start/tool_end events now emit
proper { type: "tool-call" } parts from the adapter, enabling per-tool
UIs registered via tools.by_name in MessagePrimitive.Parts.

- Web search: Globe icon, Source badges with favicons, auto-collapse
  when LLM starts responding
- Python: Code icon, syntax-highlighted code via Streamdown/shiki,
  output block with copy
- Terminal: Terminal icon, command in trigger, output with copy
- ToolGroup wraps consecutive tool calls (skips for single calls)
- Sources component renders URL badges at end of message
- Flattened code block CSS (single border, no nested boxes)

* fix(inference): respect empty enabled_tools allowlist

`if payload.enabled_tools:` is falsy for [], falling through to
ALL_TOOLS. Use `is not None` so an explicit empty list disables
all tools as intended.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Shine1i <wasimysdev@gmail.com>
2026-03-18 08:28:02 -07:00
Daniel Han
0acd1c7eec
studio: improve onboarding UX, tooltips, and training defaults (#4355)
* studio: improve onboarding UX, tooltips, and training defaults

- Change splash text to "Train and run LLMs locally"
- Add "Chat Only" card with BubbleChatIcon to skip directly to chat
- Add Skip/Skip to Chat buttons in sidebar and footer
- Back button on step 1 returns to splash screen instead of being disabled
- Change "Watch video guide" to "Get started with our guide" with new URL
- Update intro text to mention all model types + chat
- Make all tooltips clickable (in addition to hover) via React context
- Strip surrounding quotes from pasted HF tokens
- Rename "Eval Split" to "Evaluation Split"
- Add SparklesIcon to "Auto Detect" format option
- Change step 4 heading to "Choose your training parameters"
- Default max_steps to 60
- Learning rate displayed in scientific notation with +/- stepper
- Context length options capped by model's max_position_embeddings (via AutoConfig)
- Fix "QLORA"/"LORA" to "QLoRA"/"LoRA" in summary step
- Backend: add max_position_embeddings to model config endpoint

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* compare for 2 diff models

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* resolving gemini comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: disable thinking for Qwen3.5 <9B and always for AI Assist

- Change Qwen3.5 thinking threshold from <=2B to <9B (0.8B, 2B, 4B
  all disable thinking by default; 9B+ enables it)
- Always pass enable_thinking=False in AI Assist helper calls
  (_run_with_helper and _generate_with_backend) regardless of chat
  thinking settings

* studio: address PR review comments

- Extract _get_max_position_embeddings helper to DRY config extraction
- Fix "Skip to Chat" to navigate to /chat on step 1 (was /studio)

* fix: comment out debug print statements

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: skip Shiki highlighting for incomplete SVG code fences

While streaming SVG content, the syntax highlighter (Shiki) re-parses
the entire growing SVG on every token, blocking the main thread and
freezing the code area until the fence closes. Show a plain-text
preview for incomplete SVG fences instead, similar to how Mermaid
diagrams show a placeholder while streaming.

* studio: fix default top_k from 50/40 to 20 for chat inference

Per Qwen3.5 docs (unsloth.ai/docs/models/qwen3.5), top_k should be 20
for both thinking and non-thinking modes. The model-specific config in
inference_defaults.json already had top_k=20 for Qwen3.5, but the
generic fallback defaults were wrong:
- Frontend DEFAULT_INFERENCE_PARAMS.topK: 50 -> 20
- Backend generate_chat_completion top_k: 40 -> 20
- Backend generate_chat_completion_with_tools top_k: 40 -> 20
- Frontend title generation top_k: 40 -> 20

* studio: set universal inference defaults for unknown models

Default params for any model without specific config:
  temperature=0.6, top_p=0.95, top_k=20, min_p=0.01,
  presence_penalty=0.0, repetition_penalty=1.0

Models with entries in inference_defaults.json (Qwen3.5, Gemma-3,
Llama, etc.) override these with their recommended values.

Updated in: frontend DEFAULT_INFERENCE_PARAMS, backend Pydantic
request models, and backend generate_chat_completion defaults.

* studio: only trust_remote_code for unsloth/ models in AutoConfig

Only set trust_remote_code=True when the model name starts with
"unsloth/". All other models default to False for safety.

* studio: move Generating spinner above the composer

The "Generating" spinner was below the send message bar, causing
the bar to jump up and down. Move it above the composer in both
the regular thread view and the welcome/empty view.

* studio: adjust toast close button position away from edge

Move the X close button on toasts (like "Starting model...") from
top-1.5 to top-3 and add right-3, giving more breathing room from
the top-right corner.

* studio: make Think button smaller with tighter icon-text gap

Reduce gap from 1.5 to 0.5, padding from px-2.5/py-1 to px-2/py-0.5,
and icon from size-3.5 to size-3.

* studio: multiple onboarding and chat UX improvements

- Move Generating spinner above composer (fixes jumping send bar)
- Make Think button smaller with tighter icon-text gap
- Chat card now inside grid (same size as Audio/Embeddings cards)
- Rename "Chat Only" to "Chat"
- Chat card requires Continue to proceed (no auto-advance)
- Continue on Chat selection skips onboarding and goes to /chat
- Tooltip (i) click on Chat card doesn't trigger navigation
- Step 1 footer Back button goes back to splash (label is "Back")
- Splash "Skip Onboarding" renamed to "Skip to Chat", navigates to /chat
- Toast close button moved away from edge

* studio: align Skip to Chat button, add Skip to footer

- Sidebar "Skip to Chat" now uses primary (green) Button style with
  arrow icon, full width, aligned like step items. Shows on all steps.
- Footer: added "Skip" outline button next to Continue that goes
  directly to /studio with progress saved (markOnboardingDone)

* studio: change default max steps from 30 to 60 in toggle hook

The DEFAULT_MAX_STEPS in use-max-steps-epochs-toggle.ts was still 30,
used as fallback when toggling from epochs back to max steps.

* studio: extend context length options to 262K

CONTEXT_LENGTHS now includes 65536, 131072, 262144 in addition to
the existing 512-32768 range. The onboarding step filters these by
the model's max_position_embeddings (e.g. Nemotron-3-Nano-4B has
262144), showing powers of 2 up to the model's maximum.

* studio: auto-select LoRA vs QLoRA based on model size and GPU memory

After selecting a model in onboarding, detect the total model weight
file size from HF Hub (safetensors/bin files). Then estimate memory
needed: model_size_gb * 1.5 * context_scale, where context_scale is:
  - <=8192 tokens: 1.0x
  - >8192 tokens: 1.7x
  - >=16384 tokens: 2.0x
  - >=32768 tokens: 4.0x

If the estimate fits in free GPU VRAM, default to LoRA (16-bit).
Otherwise default to QLoRA (4-bit).

Backend changes:
- Add model_size_bytes to ModelDetails (models.py)
- Add _get_model_size_bytes() using HfApi.repo_info (routes/models.py)
- Add vram_free_gb to get_gpu_summary (hardware.py)

Frontend changes:
- Add autoSelectTrainingMethod() in training-config-store.ts
- Called after model defaults are loaded
- Add model_size_bytes to ModelConfigResponse type
- Add vramFreeGb to HardwareInfo hook

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: rename "Importing ML libraries..." to "Importing Unsloth..."

* studio: show model/dataset in training status, fix LoRA/QLoRA casing

- Training status now shows 'Training "model_name"' and 'Dataset = ...'
  instead of generic "Starting training..."
- Fix Studio progress section to show QLoRA/LoRA instead of QLORA/LORA

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: rename 'Skip to Chat' to 'Skip Onboarding' on splash screen

* studio: add presence_penalty support for chat inference

Add presence_penalty as a parameter across the full stack:
- Backend: llama_cpp.py generate_chat_completion/with_tools, Pydantic
  models (inference.py), routes/inference.py pass-through
- Frontend: InferenceParams type, DEFAULT_INFERENCE_PARAMS (0.0),
  chat-adapter.ts payload, chat-settings-sheet.tsx slider (0-2),
  model defaults loading from inference_defaults.json
- Set Qwen3.5 default presence_penalty to 1.5 per official docs
- Default for unknown models is 0.0 (off)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix Chat card deselecting Text and aligning with other cards

* studio: fix presence_penalty not loading from inference defaults

The inference_config.py load_inference_config() was not including
presence_penalty in the returned config dict, so the Qwen3.5
default of 1.5 from inference_defaults.json never reached the
frontend. Added it to the config builder.

* studio: add delete button for cached models in model selector

Add trash icon on each downloaded model row (GGUF and safetensors) with
confirmation dialog. Backend DELETE /api/models/delete-cached endpoint
uses huggingface_hub scan_cache_dir + delete_revisions to cleanly remove
cached repos, refusing if the model is currently loaded.

* studio: restore inference defaults, reasoning, and tools on page refresh

On page refresh with a model already loaded, the frontend was not
re-applying model-specific inference defaults (presence_penalty,
temperature, etc.) or restoring reasoning/tools support flags.

Backend: Add inference config, supports_reasoning, supports_tools,
and context_length to InferenceStatusResponse.

Frontend: In the refresh callback, when an active model is detected,
apply mergeRecommendedInference and restore reasoning/tools flags
with proper Qwen3.5 size-based defaults.

* studio: fix delete dialog closing before async completes

Prevent AlertDialogAction's default close behavior with
e.preventDefault() so the dialog stays open during deletion.
Also block onOpenChange dismiss while deleting is in progress.

* fix: add Dict and Any imports to inference models

* studio: fix Qwen3.5 reasoning threshold in frontend load path

The frontend loadModel handler had the old threshold (<=2) for
disabling reasoning on small Qwen3.5 models. Changed to <9 to
match the backend. This was causing 4B to not properly disable
thinking by default when auto-loaded.

* studio: move GGUF delete to per-variant level

For GGUF repos, the trash icon now appears on each downloaded variant
row inside the quantization expander instead of on the repo-level row.
Backend accepts optional variant param to delete specific GGUF files
(blob + symlink) rather than the entire repo cache.

* studio: restore ggufContextLength on page refresh

The Max Tokens slider was capped at 32768 on page refresh because
ggufContextLength was not restored from the status response.
Now set it from statusRes.context_length on reconnect.

* fix: remove <think> from Qwen3.5 response template marker

The train-on-responses-only feature uses template markers to find
where the assistant response starts. The Qwen3.5 response marker
included '<think>\n' which is only present when thinking mode is
enabled. With thinking disabled (default for <9B), the marker
never matched, causing 100% of samples to be dropped.

Changed response marker from '<|im_start|>assistant\n<think>\n'
to '<|im_start|>assistant\n' which works regardless of thinking mode.

* studio: fix sloth ASCII art alignment in training overlay

* fix: correct sloth ASCII art alignment to match Unsloth banner

* studio: add Python and terminal tool calling to chat

Register python and terminal tools alongside web search. Python
executor validates imports (stdlib only) via unsloth_zoo
rl_environments, runs code in a subprocess sandbox with 5-min
timeout and cancel support. Terminal executor blocks dangerous
commands (rm, sudo, etc.) and runs in a temp directory.

Update llama_cpp tool loop to show tool-specific status messages
and pass cancel_event through to executors. Rename composer
toggle from "Search" to "Tools" and show TerminalIcon for
execution status pills.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix Nemotron/transformers 5.x support, onboarding navigation, port binding

Backend:
- Dynamic transformers 5.x detection via tokenizer_config.json fetch
  (checks for TokenizersBackend class, cached per-model)
- Bump transformers 5.x version from 5.2.0 to 5.3.0 across all workers,
  setup scripts (setup.sh, setup.ps1)
- Auto-enable trust_remote_code for unsloth/* models needing transformers 5.x
  (workaround for NemotronH config parsing bug in transformers)
- Auto-install mamba-ssm/causal-conv1d for SSM models (NemotronH, Falcon-H1)
  with --no-build-isolation --no-deps to avoid torch version conflicts
- Add SO_REUSEADDR to port check in run.py (fixes Colab proxy stale connection
  falsely reporting port as in-use)

Frontend:
- Fix "Skip to Chat" navigation: use window.location.href instead of React
  Router navigate() to bypass useEffect redirect race
- Fix "Skip Onboarding" on splash: navigates to /studio (not /chat)
- Fix onboarding guard: only check isOnboardingDone() on initial mount
- Fix Chat card on step 1: add sr-only spacer for consistent alignment
- Fix Chat+Text both selected: clear RadioGroup value when Chat is selected

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: split tools toggle into Search and Code buttons

Replace the single "Tools" toggle with two independent toggles:
- "Search" (globe icon) enables web search only
- "Code" (terminal icon) enables Python and terminal execution

Add enabled_tools list field to the inference payload so the
backend only registers the tools the user has toggled on. Both
toggles appear in the main composer and the compare composer.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix tool calling import validation and error logging

Replace unsloth_zoo-dependent import checker with a standalone
ast-based validator using sys.stdlib_module_names. This properly
blocks non-stdlib imports (numpy, requests, etc.) and returns a
clear error message to the model so it can rewrite using only
stdlib.

Add full traceback to tool streaming error logs for debugging.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: parse gpt-oss harmony channels for clean safetensors chat output

gpt-oss models emit multi-channel output via harmony protocol tokens
(<|channel|>analysis<|message|>... and <|channel|>final<|message|>...).
TextIteratorStreamer with skip_special_tokens=True strips the special
tokens but leaves channel names concatenated with content, producing
garbled output like "analysisWe need to...assistantfinalHello!".

Add HarmonyTextStreamer that decodes with skip_special_tokens=False,
parses harmony markup via regex, and emits <think>analysis</think>
for the analysis channel and plain text for the final channel --
reusing the existing frontend reasoning UI.

Also expose supports_reasoning=True for non-GGUF gpt-oss models in
the /status endpoint so the frontend enables the Think toggle.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: use unsloth_zoo for Python sandbox validation

Set UNSLOTH_IS_PRESENT=1 and import check_python_modules and
check_signal_escape_patterns directly from unsloth_zoo instead
of a standalone fallback. This gives us the full Unsloth
validation including stdlib-only import checks and signal/timeout
escape pattern detection.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: allow all imports in Python tool sandbox

Remove stdlib-only import restriction. Keep signal escape
pattern detection via unsloth_zoo for safety.

* studio: fix ReadTimeout on tool streaming final pass

The 0.5s read timeout used for cancel-checking during streaming
also fires when waiting for the first response from llama-server
(e.g. reasoning model thinking for 15+ seconds). Add
_stream_with_retry() context manager that retries on ReadTimeout
while checking cancel_event, so the model has unlimited time to
think before producing the first token. Applied to both the
regular streaming path and the tool-calling final pass.

* fix: rewrite HarmonyTextStreamer with stateful incremental parsing

The delta-on-transformed approach had two critical bugs:

1. Before the full <|channel|>X<|message|> pattern was complete, the
   strip-tokens fallback emitted "analysis" as plain text. Then when
   the regex matched, _transform returned a completely different format
   (<think>...</think>) and the delta was computed against the wrong
   base string, producing fragments like "think>", "nk>", ">".

2. Even with full matches, the closing </think> tag shifted position
   as content grew, so text[prev_len:] produced garbled deltas.

Replace with stateful incremental parsing that:
- Buffers until a complete channel+message pair is seen
- Emits <think> once when analysis channel first appears
- Streams analysis content deltas (computed on channel content directly)
- Emits </think> once when final channel first appears
- Streams final content deltas
- Closes open think tags in end()

Also skip the generic all_special_tokens stripping in
_clean_generated_text for gpt-oss since HarmonyTextStreamer already
produces clean output and the generic stripping was mangling <think>
tags.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: strip all <|...|> tokens in gpt-oss cleanup, not just harmony subset

The gpt-oss tokenizer has added tokens like <|return|> (id=200002) that
are not part of the harmony channel protocol but can leak into output.
The previous regex only stripped channel|message|start|end tokens.

Broaden the _clean_generated_text regex for gpt-oss to <\|[a-z_]+\|>
which catches all pipe-delimited tokens (return, constrain, reserved,
etc.) without matching <think>/<\/think> tags.

Verified: gpt-oss all_special_tokens are only <|return|>,
<|reserved_200017|>, <|startoftext|> -- none overlap with <think>.
The harmony tokens (channel, message, start, end) are added_tokens
but not in all_special_tokens.

* fix: hide config-only model repos from cached models list

Repos that only have metadata/config files cached (no .safetensors or
.bin weight files) were showing up in the Downloaded list with tiny
sizes like "1.8 KB" or "24 KB". These are just leftover config
snapshots from architecture checks, not usable models.

Filter the cached-models endpoint to only include repos that contain
actual model weight files (.safetensors or .bin).

* studio: fix toast description text contrast in dark mode

Add explicit !text-muted-foreground to toast description classNames
so secondary text (e.g. "Releases VRAM and resets inference state.")
is readable in dark mode.

* studio: fix Chat card icon alignment with size-4 spacer

Replace sr-only span (takes no space) with a size-4 shrink-0 div
matching the RadioGroupItem dimensions in other cards, so the Chat
icon aligns vertically with Text/Audio/Vision/Embeddings icons.

---------

Co-authored-by: workspace <user@workspace.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Manan17 <shahmanan170602@gmail.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-03-17 07:46:07 -07:00
Daniel Han
eeffa4c065
studio: web search, KV cache dtype, training progress, inference fixes
## Summary
- Add web search tool calling for GGUF models (Search toggle, DuckDuckGo via ddgs)
- Add KV cache dtype dropdown (f16/bf16/q8_0/q5_1/q4_1) in Chat Settings
- Fix Qwen3/3.5 inference defaults per official docs (thinking on/off params)
- Enable reasoning by default for Qwen3.5 4B and 9B
- Replace "Generating" toast with inline spinner
- Fix stop button via asyncio.to_thread (event loop no longer blocked)
- Fix CUDA 12 compat lib paths for llama-server on CUDA 13 systems
- Fix auto-load model name not appearing in selector
- Training progress messages + dataset_num_proc fix

Integrated PRs:
- #4327 (imagineer99): BETA badge alignment (already in tree)
- #4340 (Manan Shah): prioritize training models in model selection
- #4344 (Roland Tannous): setup.sh macOS python version compatibility
- #4345 (Manan Shah): revamp model+dataset checking logic
2026-03-17 00:30:01 -07:00