Commit graph

28 commits

Author SHA1 Message Date
Daniel Han
5211b506e1
Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392)
* Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm

The OpenAI-compatible endpoints serve whichever GGUF is loaded and ignore the
request model field, so an OpenAI client that changes model never reloads. Add
an opt-in setting that, when a /v1 request names a downloaded local GGUF
different from the loaded one, loads it before serving by reusing the existing
/load path (its dedup, tensor fallback, and threading apply). Unknown names
still serve the loaded model, so drop-in compatibility is preserved and no
remote download is triggered.

Also add an optional idle auto-unload (TTL keep-warm): a pure-ASGI middleware
tracks in-flight inference requests so a stream is never unloaded mid-response,
and a lifespan loop unloads the model after the configured idle seconds. Both
settings default off and live in the app_settings store, exposed via
GET/PUT /api/settings/openai-auto-switch.

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

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

* Studio: variant-aware auto-switch, /v1/responses coverage, keep-warm load stamp

Follow-ups from review of the opt-in OpenAI auto-switch path:

1. Variant-aware dedup. _maybe_auto_switch_model compared only the repo id, so
   requesting another quant of the loaded repo (e.g. Q4_K_M loaded, Q8_0 asked)
   was served by the old quant. Compare hf_variant too, matching /load dedup.

2. Streaming /v1/responses now calls the auto-switch hook. It went straight into
   _responses_stream and only checked is_loaded, so stream=True could serve the
   old model or 400. Non-streaming already routed through chat completions; the
   hook is idempotent once loaded.

3. resolve_local_gguf tries an exact id match before splitting a trailing
   :VARIANT, so local ids that contain a colon (e.g. a Windows path) resolve
   instead of being cut at the drive letter.

4. Idle keep-warm stamps activity on a load/swap transition. _last_active was
   only refreshed by inference requests, so a model loaded after the server sat
   idle past the TTL could be unloaded before its first request.

Tests cover each case.

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

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

* Studio: make the /v1/responses auto-switch test order-independent

The new streaming-responses test passed in isolation but failed under the CI's
randomized collection order with "object has no attribute 'state'": it passed a
bare object() as the request and stubbed only one dispatcher, so an ordering
where the real dispatcher ran hit request.state. Give the request a state and
stub both dispatchers; the test still asserts the hook fires before dispatch.

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

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

* Studio: assert /v1/responses auto-switch wiring on source, not at runtime

The behavioral version executed openai_responses and relied on stubbing its
callees, which a randomized collection order in CI could defeat (the real
dispatcher ran and hit request attributes). Assert on the function source that
the hook precedes both dispatchers instead; the hook's runtime behavior is
already covered by the direct _maybe_auto_switch_model tests.

* Studio: auto-switch on /v1/embeddings, GGUF-only targets, idle-unload race gate

Second-pass review follow-ups on the opt-in auto-switch path:

1. /v1/embeddings now calls the auto-switch hook before the loaded-state check,
   matching the other model-bearing OpenAI endpoints (the keep-warm middleware
   already treats embeddings as inference).

2. The resolver index is now GGUF-only. The local-model scanners also surface
   Transformers/safetensors repos; without a filter, auto-switch could unload
   the GGUF and route a request into the non-GGUF loader. _has_local_gguf checks
   a direct .gguf, a models-dir folder, and the HF-cache snapshots layout.

3. Idle keep-warm now holds an asyncio gate across the idle check and the
   unload, and a request bumps inflight under the same gate, so the loop can no
   longer unload in the window between "looks idle" and the kill.

Tests cover each. Broader local-model source parity (LM Studio, Ollama, legacy
caches, custom scan folders) is a follow-up; missing one of those today just
falls through to the loaded model.

* Studio: variant-aware local resolver, count_tokens + audio auto-switch coverage

Third-pass review follow-ups on the opt-in auto-switch path:

1. The resolver is now variant-aware via list_local_gguf_variants. It indexes
   only the quants actually on disk, recursing snapshots and quant subdirs such
   as the nested per-quant folders, so a requested repo:VARIANT resolves only
   when that quant is local and a bare repo resolves to a concrete local quant.
   This fixes two gaps: the previous shallow glob rejected nested-variant GGUF
   repos, and a request for an uncached quant could send /load down the remote
   download path, breaking the local-only contract.

2. /v1/messages/count_tokens now auto-switches like its sibling /v1/messages, so
   a count uses the requested model's tokenizer.

3. /api/inference/audio/generate (direct GGUF TTS) is now tracked as in-flight
   inference, so the idle loop cannot unload the model mid-generation.

Tests cover each. Two reviewer items are left as follow-ups: indexing the
remaining local sources (LM Studio, Ollama, legacy/default caches, custom scan
folders), which fails safe today by falling through to the loaded model; and
fully serializing concurrent different-model requests, an inherent limit of the
single-slot llama backend that the opt-in feature is not designed around.

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

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

* Studio: make local GGUF resolver fail-safe so a bad model name cannot 500

The auto-switch hook calls resolve_local_gguf without its own guard, and
/v1/completions and /v1/embeddings pass body.get("model") through unchanged.
A non-string model (e.g. {"model": 123}) or any internal scan failure would
then raise out of the resolver and turn a request that would otherwise be
served by the loaded model into a 500, breaking the drop-in compatibility the
feature is built on.

Guard the resolver at its boundary: reject non-string input up front and wrap
the lookup so any failure returns None (fall through to the loaded model).
Add regression tests for both paths.

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

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

* Studio: per-model launch flags for auto-switched GGUF models

* Studio: list switch-eligible GGUFs in /v1/models when auto-switch is on

* Studio: settings UI for OpenAI model auto-switch and idle auto-unload

* Studio: show save error over the disabled-idle hint in auto-switch settings

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

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

* Studio: address gemini review (case-insensitive /v1/models retrieve, idle-input empty guard)

* Studio: address codex review (deterministic override args, exclude probe/embedding models from discovery)

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

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

* Studio: keep-warm count_tokens, gate idle on auto-switch, drop hidden models

Three hardening fixes to the opt-in auto-switch path surfaced while reviewing
the work that builds on it:

1. count_tokens keep-warm. /v1/messages/count_tokens counts via the loaded
   tokenizer and already auto-switches, but the keep-warm middleware did not
   track it, so idle auto-unload could free the model mid-count. It is now a
   tracked in-flight path.

2. "Off means unchanged" for idle unload. get_auto_unload_idle_seconds now
   reports 0 while auto-switch is disabled. Idle unload only makes sense with
   auto-switch on (an unloaded model returns only via the next request's swap),
   so a stray TTL can no longer trigger a destructive unload while the feature
   is off, keeping the disabled state identical to pre-feature behavior.

3. Hidden models are not switch targets. The resolver index now skips what
   Studio hides from its own pickers (the llama.cpp validation probe, RAG
   embedding weights) via _is_hidden_model, so they can never be auto-switched
   to by name.

Tests added for each.

* Studio: bare-id reuse, responses validation order, in-flight tracking

Review follow-ups after folding in the per-model overrides and discovery work:

1. A bare model id (no :VARIANT) is now satisfied by any loaded quant of that
   repo. Previously a bare name resolved to the largest local quant, so it could
   force a slow reload when a different quant of the same repo was already
   serving. An explicit repo:VARIANT request still honors the quant.

2. /v1/responses now runs the auto-switch hook after the empty-input validation
   so a request that 400s can no longer trigger a multi-minute model load before
   being rejected. The hook still precedes both dispatchers, so streaming
   requests switch.

3. The keep-warm middleware now tracks in-flight requests whenever auto-switch
   is enabled rather than only when the idle TTL is already positive, so a stream
   that starts with the TTL at 0 is still protected if idle-unload is enabled
   mid-stream. Off still passes straight through.

Tests added for each.

* Studio: tighten auto-switch code comments

Comment/docstring-only pass over the OpenAI auto-switch feature: collapse
multi-line blocks, drop a comment that restated the gate it sits next to, and
trim verbose docstrings on internal helpers while keeping the load-bearing
rationale (concurrency, API behavior, drop-in compat, gotchas). No logic
change: verified comment-only with the AST/printer signature check.

* Studio: bind auto-switch locks per running loop

Review follow-up. The auto-switch swap lock and the keep-warm unload gate were
module-level asyncio.Lock objects. That is safe under the single uvicorn loop
and on Python 3.10+ (the Lock resolves the running loop lazily on acquire), but
a module-level Lock binds to one loop on pre-3.10, which can raise a loop
mismatch in multi-loop runners. Resolve each lock through a per-loop accessor
backed by a WeakKeyDictionary so every running loop gets its own Lock and stale
loops are collected. No behavior change under the server's single loop.

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

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

* Studio: auto-switch re-review fixes (body codes, coverage, swap, alias, tracking)

Follow-ups from a second review pass over the opt-in OpenAI auto-switch feature:

1. OFF-state status codes: /v1/completions and /v1/embeddings moved the body
   read ahead of the loaded-state check, so a malformed/empty body with no model
   loaded returned 500 instead of the prior 503. A shared helper reads the body
   defensively (an unparseable/non-dict body yields no model), and the handler
   re-reads after the 503 gate to surface the original parse error exactly as
   before. OFF behavior is unchanged.

2. Local-model coverage: the resolver index only scanned ./models and the active
   HF cache, while the model picker also lists the legacy/default HF caches, LM
   Studio dirs, and user scan folders. A request for one of those named models
   silently served the loaded model instead. _build_index now scans the same
   roots (Ollama's symlink-creating scanner is skipped on the request path), and
   resolution is offloaded with asyncio.to_thread so the wider scan never blocks
   the event loop.

3. Swap vs in-flight stream: a cross-model swap killed the llama-server while
   another client was still streaming from it. The hook now tracks how many
   requests are streaming on the loaded model (in-flight minus those still inside
   the hook) and returns 409 instead of swapping while one is active. Concurrent
   same-model requests never reach this path, so they are unaffected.

4. Idle-unload + alias: after idle-unload freed the model, an unknown/alias name
   resolved to nothing and 503'd, though it served the active model before the
   TTL. Idle-unload now remembers the freed id and an alias request reloads it
   (only an already-local model, so no remote download), cleared once a model is
   loaded again.

5. In-flight tracking: the keep-warm middleware tracked in-flight only while the
   feature was on, so a stream started while off could be unloaded if idle-unload
   was enabled mid-stream. It now tracks on every inference path; counting is
   cheap and invisible to clients.

Tests added for each.

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

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

* Remove stray async_task_outputs files committed by mistake

* Studio: auto-switch review round 3 (revert swap guard, hardening)

Addressing a third review pass:

- Revert the cross-model swap guard. It counted keep-warm in-flight (which
  includes external-provider calls that never touch the local model) and so
  could 409 a local swap spuriously, and it still left a same-model request able
  to start streaming on the model a concurrent swap was unloading. A correct fix
  needs a request-lifetime reader/writer barrier; a partial guard was worse than
  the honest single-slot behavior, so concurrent different-model use is back to
  being serialized (documented), like llama-swap's single slot.
- Non-string request model (e.g. {"model": 123} on a raw-body endpoint) is now
  treated as absent, so it falls through instead of raising in the membership
  checks once an idle-unload stash exists.
- Idle-unload now stashes and replays the freed quant: an alias reload restores
  the exact (id, variant) that was freed rather than the largest local quant.
- Anthropic /v1/messages validates max_tokens before the auto-switch hook, so a
  request that 400s never triggers a model load.
- Keep-warm tracks a pending count for requests waiting on the unload gate, so
  the idle loop cannot unload the model out from under a request that is blocked
  on the gate but not yet counted as in-flight.
- The idle-unload task is awaited after cancel on shutdown to avoid pending-task
  warnings.
- The resolver's HF cache scan is None-safe and logs at debug instead of letting
  a bad root abort the whole index build.
- upsert_app_setting_map_entry rolls back explicitly on error.

Tests updated/added for each.

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

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

* Studio: keep saved idle-unload seconds when auto-switch is toggled off

* Studio: auto-switch hardening (thread-safe lock maps, body validation)

Defensive fixes from review:

- Guard the per-loop WeakKeyDictionary get-or-create for both the unload gate and
  the auto-switch lock with a threading lock, since WeakKeyDictionary mutation is
  not thread-safe when two event loops run on different threads.
- Build the resolver index under the cache lock so concurrent callers with an
  expired cache don't all run the multi-dir scan at once.
- /v1/completions and /v1/embeddings return a clean 400 for a valid JSON body
  that is not an object (e.g. a list), instead of a 500 from body.get(...).
- The keep-warm middleware only tracks POST requests (inference is always POST),
  so CORS preflight (OPTIONS) is not counted, and tolerates a None path.

Tests added for the list-body 400 and the non-POST skip.

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

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

* Studio: auto-switch review round 4 (local-path load, swap guard, idle fixes)

From a 10-reviewer pass:

- HF-cache entries now load by a concrete local path, not the bare repo id. The
  resolver records a load_path (the snapshot dir for a models--* cache repo, the
  file/dir otherwise) so /load takes the local branch and can never trigger a
  download to satisfy a partial cache. The advertised loader_id (repo id) is kept
  as the launch-override key. resolve_local_gguf now returns
  (load_path, variant, loader_id).
- Re-add a single-slot swap guard: a cross-model swap returns 409 model_switch_busy
  while another inference request is active rather than killing its stream (the
  caller is excluded from the count), and holds the keep-warm gate across the load
  so no new inference starts mid-swap. Concurrent same-model requests never reach
  this path. A residual spurious 409 is possible while a concurrent or external-
  provider request is active; that is the documented single-slot tradeoff.
- Idle keep-warm tracks (model_identifier, hf_variant): reloading the same repo at
  a different quant counts as a fresh model, so it is not unloaded before one TTL.
- Track Studio's own /api/inference/generate/stream so the idle loop can't unload
  the model mid-stream on that route.
- A successful manual /load clears the idle-unload reload stash synchronously, not
  only on the next idle poll.

Also merged origin/main (the branch had fallen behind, which would have reverted
unrelated files on merge). Tests added/updated for each.

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

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

* Studio: auto-switch review round 5 (concurrency, identity, load gate)

From a 10-reviewer pass (9 request-changes, 1 approve):

- Concurrent same-target requests load once instead of each returning 409. The
  count-based busy guard could not tell "another request wants the same model"
  (safe, load once) from "another request is using the loaded model" (refuse).
  Track in-flight auto-switch requests per (target, variant) and subtract
  same-target waiters from the busy count; a cross-model swap still 409s while a
  genuinely different request is active.
- Fix the identity confusion introduced when round 4 began loading by concrete
  local path: the backend identifier became a filesystem path. Record the
  advertised repo id on the backend after an auto-switch load and use it so
  (a) a model loaded manually by repo id is recognized as already serving
  (no spurious reswap/409), (b) /v1/models reports the repo id, never a host
  path or a duplicate, and (c) the idle-unload stash keeps the override keyed by
  the repo id, so an alias reload after TTL keeps the user's saved launch flags.
- Gate the manual /load route with the keep-warm lifecycle gate so idle
  auto-unload can't unload a model mid-load. load_model now wraps _load_model_impl
  in the gate; auto-switch calls _load_model_impl directly since it already holds
  the gate.
- Restore default-off parity on Anthropic /v1/messages: an unloaded backend with
  auto-switch disabled 503s before the max_tokens 400 check, as it did pre-feature.
  When the feature is on, request-shape validation still runs before any load.

Tests added for each; full backend suite diff vs baseline is unchanged.

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

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

* Studio: auto-switch review round 6 (concurrency ordering, leaks, unload gate)

From a second 10-reviewer pass (8 request-changes, 2 approve):

- Same-target concurrency: register a waiter by the raw requested model before
  the (slow) resolve, and exclude pending requests from the swap busy count. The
  middleware counts a concurrent same-model request as in-flight before it
  resolves and joins the resolved-target waiter map, so the prior fix could still
  409 it. The guard now subtracts max(same resolved-target, same raw-request)
  waiters and ignores pending (a pending request is blocked in the middleware,
  not generating, so a swap can't interrupt it).
- External-provider requests no longer block a local swap. The keep-warm
  middleware counts every inference-path POST, but external-provider chat returns
  before the auto-switch hook and never touches the local GGUF. The chat handler
  now untracks itself before proxying, so its in-flight stream can't trip
  model_switch_busy on a concurrent local auto-switch. The middleware skips its
  own end-decrement for an untracked request.
- Manual /unload is gated like load and idle-unload: it holds the lifecycle gate
  and returns 409 rather than tearing down llama-server while an inference request
  is in flight.
- Response model id no longer leaks the load path. /v1/models already advertised
  the repo id; chat, completions, embeddings, Anthropic messages, and audio
  response bodies now use the same _llama_public_model_id helper instead of the
  concrete on-disk model_identifier.
- Chat completions validates the non-system-message requirement before the
  auto-switch hook (as /responses and /messages already do), so an invalid
  request can't swap the resident model before returning 400.

Tests added for each; full backend suite diff vs baseline is unchanged.

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

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

* Studio: auto-switch review round 7 (teardown policy, Unsloth-active swap, training)

From a third 10-reviewer pass (9 request-changes, 1 approve), all on the same
asymmetric-teardown theme. Resolved per the intended policy that only automatic
paths defer to an active stream; deliberate user actions stay interrupting:

- Revert the manual /unload in-flight guard added last round. A manual /load or
  /unload is a deliberate action and tears down immediately, as before; only the
  automatic idle-unload loop and auto-switch defer to an active request. This
  removes the asymmetry the reviewers flagged (manual /load, the /unload Unsloth
  branch, and the opposite-backend swaps inside _load_model_impl) by not
  extending the guard to deliberate paths, rather than spreading it.
- Auto-switch now refuses a swap whenever another inference request is in flight,
  not only when a GGUF is already loaded. _load_model_impl also unloads an active
  Unsloth/transformers backend before loading a GGUF, so the busy guard must cover
  that case too; otherwise an Unsloth stream could be killed by an auto-switch.
- Refuse API-initiated training while inference is active. When Studio is driven
  as an inference API (sk-unsloth key auth), POST /api/training/start returns 409
  if a request is in flight, since training frees VRAM by unloading the chat
  model and would kill the stream. The Studio UI (session auth) still starts
  training and coexists/frees VRAM as before. A mixed UI+API session is not yet
  special-cased. Adds auth.authentication.authenticated_via_api_key.

Tests added/updated for each; full backend suite diff vs baseline is unchanged.

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

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

* Studio: add UNSLOTH_MODEL_IDLE_TTL env override for idle-unload

Borrowed from PR 6517: a startup env var that sets the idle-unload TTL without
the settings UI. Unlike the stored setting (gated on auto-switch), the env value
is a standalone default that enables idle-unload even with auto-switch off, for
headless/container deploys. An explicit UI/API value still overrides it and stays
gated. The settings GET reflects the env default when nothing is stored.

* Studio: auto-switch fixes from review (paths, embeddings input, env idle reload)

- /v1/models advertises a client-facing alias instead of a filesystem path:
  the ./models and LM Studio scanners report the on-disk path as the model id,
  so the index now prefers model_id/display_name as the advertised/override id
  and keeps the concrete path internal as load_path, still resolvable by path.
- /v1/embeddings validates input before auto-switch: a request with a model but
  no input now 400s before the hook (like chat/responses/messages), so an
  invalid embeddings request cannot unload or swap the resident model.
- Standalone UNSLOTH_MODEL_IDLE_TTL reloads the freed model: the hook now runs
  when auto-switch or idle-unload is active, and with auto-switch off it skips
  the resolver and only restores the idle-unloaded model, so the first idle
  timeout no longer leaves later /v1 requests with nothing loaded.
- Do not resurrect a stale GGUF over an active Unsloth model: the reload-stash
  path bails when a non-GGUF backend is loaded, so an unknown /v1 name cannot
  tear down a live Transformers/Unsloth model.
- Defensive HF cache scan: each cache root's resolve/dedup is wrapped so a
  missing or malformed root skips that root rather than aborting the index.
- Single-model retrieve checks the id is a string before lowercasing.

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

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

* Studio: fix automatic-load asymmetry, audio reload, preview, idle timer

The standalone UNSLOTH_MODEL_IDLE_TTL reload is a second automatic-load
trigger, but several validate-before-switch guards and reload hooks only
checked the auto-switch toggle. Add a shared _automatic_model_load_may_run()
(auto-switch on, or idle TTL > 0) and route every guard through it.

- /v1/completions validates prompt before any automatic load (it was the one
  model-bearing route with no pre-check).
- /v1/chat/completions and /v1/embeddings pre-checks gate on the shared
  predicate so a standalone idle TTL cannot reload then reject.
- /v1/messages no longer 503s before the reload hook can restore an idle-freed
  model when auto-switch is off.
- Raw completions/embeddings with no model field pass a non-empty sentinel so
  the idle-stash reload runs, restoring the legacy "omit model, use loaded" path.
- /api/inference/audio/generate gains the reload hook (after message validation)
  so an idle-freed audio GGUF is restored.
- Public preview opts out of auto-switch via a request-scope flag, so a caller's
  model field cannot swap away from the pinned checkpoint; preview chat streams
  are now matched by _is_inference_path so idle-unload cannot kill them.
- Keep-warm no longer stamps activity on request start, and external-provider
  untracking decrements without restamping, so periodic external traffic can no
  longer keep the local GGUF warm forever.

Merges origin/main (the branch had fallen behind, which also brought in the
preview route the review flagged).

* Studio: surface model auto-switch in the API tab and demo it in examples

The OpenAI auto-switch toggle previously lived only in Settings -> General.
Add the same toggle to the API tab's usage-examples panel (it shares the
settings cache), and make the examples reflect it: when on, the Python
examples append a second call naming a different downloaded GGUF (so the
model field visibly selects which model serves), and the curl examples gain
a one-line note. Reuses the existing settings API client and i18n keys.

* Studio: harden OpenAI auto-switch reload-only path and Anthropic tool validation

- Omitted-model raw-body requests pass a reload-only sentinel so the idle-stash
  reload still restores an idle-freed model, but the resolver never matches a
  downloaded GGUF literally named "default".
- Reject malformed Anthropic client tools before _maybe_auto_switch_model so an
  invalid request can no longer evict the loaded model.

* Studio: extend auto-switch reload-only and tool validation to schema endpoints

- Schema-backed endpoints (chat completions, responses, count_tokens, messages,
  audio) defaulted an omitted model to "default" and passed it to the switch
  hook, so a downloaded GGUF named "default" could be swapped to. Route the hook
  through a helper that switches only on an explicitly set model, else reload-only.
- Propagate the explicit-set status when building the chat request from a
  Responses request, so the non-streaming chat re-check stays reload-only too.
- Validate Responses function tools before the switch hook so a malformed tool
  returns 400 without evicting the loaded model.

* Studio: serialize auto-switch swaps across event loops with a process-wide gate

The auto-switch lock is a per-event-loop asyncio.Lock, so two /v1 swaps on
different loops in one process could both pass it and race the single model slot
(the backend and _load_model_impl are process-wide). Add a process-wide
threading gate around the swap, acquired off the loop so a cross-loop wait never
blocks it, layered with the existing per-loop lock. Add a cross-loop test that
fails without the gate (two slow loads overlap) and passes with it.

* Studio: make the auto-switch swap gate wait cancellation-safe

_acquire_swap_gate awaited asyncio.to_thread(lock.acquire) when another loop held
the process-wide gate. to_thread cancellation doesn't stop the worker thread, so a
/v1 request cancelled mid-wait (client disconnect during a cross-loop swap) would
have its thread acquire the gate after the fact, while the finally that releases it
never runs -- permanently deadlocking later auto-switch swaps.

Poll a non-blocking acquire off a short asyncio.sleep instead: it still keeps the
wait off the loop and serializes across loops, but a cancel now lands during the
sleep, when the gate is not held, so nothing leaks. Add a test that deadlocks the
to_thread variant (it times out) and passes with the poll.

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

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

* Studio: validate modality and tool-confirmation before auto-switch

Two more request shapes could load a named GGUF and only then 400, evicting the
resident model:

- An image request naming a different text-only GGUF. The switch hook now takes
  require_vision and rejects a swap to a non-vision target before loading it; a
  GGUF's vision capability is its companion mmproj, knowable without a load, and
  matches the post-load guard. Only the resolver branch is checked, never the
  reload-stash restore.
- confirm_tool_calls=true with stream=false and local tools. /v1/chat/completions
  now rejects that shape before the hook, mirroring the local tool path's
  bypass_permissions exemption and intent signal.

The vision probe threads the ambient HF token to keep the capability-probe
invariant. Reload-only and idle-reload paths are unaffected.

* Studio: extend validate-before-switch and make the lifecycle gate process-wide

- /v1/messages/count_tokens now rejects malformed client tools before the switch
  hook, like /messages (shared _validate_anthropic_client_tools helper), so a
  count request can't evict the loaded model.
- /v1/chat/completions rejects a malformed tool_choice forcing object (a
  {"type":"function","function":{}} with no name) before the switch hook.
- The inference lifecycle gate that blocks new inference during a swap is now
  process-wide (a poll-acquired threading lock, cancellation-safe), not a
  per-loop asyncio lock, so a request on another event loop can't start inference
  while a swap tears the single backend down.
- Usage examples no longer hard-code a switch-demo repo most users lack; the
  model is an explicit placeholder the user replaces.

* Studio: extend the auto-switch modality guard to /v1/responses and /v1/messages

The pre-load vision check that guards /v1/chat/completions now also runs on
/v1/responses and /v1/messages, so an image request naming a text-only GGUF is
rejected before the swap and never evicts the resident vision model. Run the
vision capability probe off the event loop. Make the /v1/models retrieve
loaded fast-path case-insensitive, and never advertise a host path from the
resolver. Remove the dead list_switch_eligible_ids helper, superseded by the
/v1/models catalog.

* Studio: filter /v1/models to GGUF, per-loop catalog lock, reject system-only Responses

Address review findings on the auto-switch path:
- /v1/models advertises only GGUF models the API can actually switch to; a
  safetensors/LoRA entry would be selectable but never loadable via llama.cpp.
- The /v1/models catalog cache uses a per-loop lock (like the auto-switch path)
  so a second event loop awaiting it can't hang in a multi-loop process.
- /v1/responses rejects system/developer-only input before the switch, mirroring
  chat, so an invalid request can't evict the resident model.
- _build_index guards each scan source on its own so one bad root drops only
  that source; the vision probe logs a real detection failure instead of
  swallowing it.

* Studio: list cached GGUFs in /v1/models by inspecting files, not model_format

The HF-cache scanner leaves model_format unset for GGUF snapshots, so the
previous model_format == "gguf" filter dropped every downloaded HF-cache GGUF
from /v1/models and the retrieve fallback. Decide GGUF-ness from the on-disk
files via the resolver (info_has_local_gguf) instead, run off the event loop, so
the catalog advertises exactly what /v1 can serve.

* Studio: fix /v1/messages/count_tokens route binding plus auto-switch review fixes

The @router.post decorator for /messages/count_tokens had been separated from
anthropic_count_tokens by the _validate_anthropic_client_tools helper, so the
route bound to the validator and dropped its auth dependency. Move the decorator
back onto the handler. Add route-binding tests asserting each /v1 endpoint maps
to its handler with the auth dependency, so a decorator/handler split is caught
at the route level (the direct-call tests missed it).

Also from review:
- update_openai_auto_switch writes both settings keys in one transaction so a PUT
  can't leave one updated and the other stale (drop the now-unused single setters).
- max_seq_length override rejects 0 at the boundary (ge=1) instead of accepting
  then silently dropping it.
- Document that embeddings auto-switch is best-effort: GGUF pooling has no cheap
  pre-load probe like vision's mmproj, so a guard would false-reject GGUF embedders.
- Add a positive idle-unload test (loop frees the model and stashes it for reload).

* Studio: validate Responses tool_choice + Anthropic mixed tools before switch, filter Ollama from catalog

More auto-switch review findings:
- /v1/responses rejects a forcing-function tool_choice with no name before the
  switch, mirroring chat, so a malformed request can't evict the resident model.
- /v1/messages rejects mixing Anthropic server tools with custom client tools
  before the switch (the check depends only on the payload, so it moves up cleanly).
- /v1/models no longer advertises Ollama-link models: info_has_local_gguf excludes
  .studio_links / ollama_links entries, which the resolver skips and can't switch
  to, so an advertised id never silently falls through.

* Studio: guard chat audio input before switch; surface env-backed idle unload in settings UI

A chat request carrying audio_base64 rides the same companion mmproj
projector as a vision request, so a text-only target cannot serve it
either. Flag require_vision for audio input as well so the multimodal
probe runs before the switch and a rejected request never evicts the
working model. Generalize the reject message to cover image and audio.

The settings response now reports idle_unload_active (effective TTL > 0)
so the UI can distinguish idle-unload that is active via the
UNSLOTH_MODEL_IDLE_TTL env var from the case where it needs the toggle
enabled.

* Studio: harden auto-switch eviction guards (count_tokens vision, TTS reload-only, mmproj/stash)

Four eviction/correctness fixes on the opt-in /v1 auto-switch path:

- /v1/messages/count_tokens now carries the same require_vision guard as
  /messages, so an image count naming a text-only GGUF can't evict a loaded
  vision model for a swap that can't serve the request.

- /audio/generate is now reload-only. A local GGUF's audio-input capability
  is not a cheap pre-load probe (the companion mmproj signal can't tell an
  audio projector from a vision one, and codec TTS ships no projector), so
  resolving the client model could load a text/vision-only target and evict
  the working audio model before the audio check fails. Only the idle-stash
  restore runs here; switching TTS models is an explicit /load.

- The resolver no longer treats a standalone mmproj .gguf as a servable
  model. _scan_models_dir's standalone-file pass does not filter mmproj the
  way its directory scan does, so /v1/models could advertise a projector and
  a switch could load it over the real weights.

- A non-GGUF (Transformers/Unsloth) load and a deliberate /unload now clear
  the idle reload stash, so a manual load/unload is never superseded by a
  stale idle-freed GGUF that the next /v1 request resurrects.

* Studio: report advertised repo id consistently after an auto-switch

Two model-id reporting fixes so an auto-switched cached HF GGUF is named by
its repo id everywhere, not its snapshot path:

- Streamed /v1/responses envelopes now derive the model id from
  _llama_public_model_id (which prefers _openai_advertised_id) instead of the
  raw model_identifier. After an auto-switch the identifier is the snapshot
  path while the repo id lives in _openai_advertised_id, so the stream used to
  report a snapshot basename while /v1/models, chat completions, and
  non-streaming Responses all reported the repo id.

- When an advertised alias already resolves to the loaded model (a model
  loaded by local path, requested by its repo or LM Studio id), the
  already-serving early return now records the alias as the advertised id, so
  /v1/models and responses report the alias and mark it loaded instead of the
  path-derived basename. Resolver branch only; safe lock-free because an
  in-flight request blocks any concurrent swap via the single-slot busy guard.

* Studio: validate request shapes before auto-switch (prompt/input/audio/mcp confirm)

Four more validate-before-switch guards so a deterministic client error never
evicts the resident model on the opt-in /v1 auto-switch path:

- /v1/completions rejects an object/number prompt (only a string or array is
  valid) before the switch, instead of loading the named GGUF and letting
  llama-server reject the shape afterward.

- /v1/embeddings rejects an object/number input the same way.

- Chat rejects an oversized audio_base64 upload (413) before the switch. The
  size cap is a cheap, target-independent length check; the decode itself
  stays post-switch to avoid decoding a valid upload twice.

- The chat confirm-without-stream pre-switch guard now mirrors the tool loop's
  actual enablement: _effective_enable_tools (honoring a CLI --enable-tools
  policy) and mcp_enabled (which opens the tool loop on its own but defers to a
  CLI --disable-tools policy). Previously a confirm+no-stream request with only
  mcp_enabled slipped past and 400'd after the swap.

* Studio: fix model-id retrieval, streaming n>1, resolver cache TTL, keep-warm auth

Four fixes from review:

- GET /v1/models/{id} legacy raw-path fallback now maps the raw identifier to
  the same public id its /v1/models entry uses. After an auto-switch load the
  identifier is the snapshot path while the entry is keyed by the advertised
  repo id, so a client that cached the old absolute path no longer 404s on a
  model that is in fact loaded.

- stream=true with n>1 is now rejected before the switch. Only the
  non-streaming GGUF path returns multiple choices, so streaming n>1 is invalid
  on every local serving path; both fields are known pre-switch, so it must not
  load model B only to 400 and evict model A. Non-streaming n>1 stays
  post-switch where the serving path decides.

- The resolver index cache is stamped after _build_index, not with the pre-scan
  timestamp. On installs with enough local models for the multi-root scan to
  exceed the 5s TTL, the cache was stored already expired and every request
  rebuilt it.

- The keep-warm middleware no longer stamps model activity for 401/403
  responses. It runs before FastAPI auth, so unauthenticated probes used to
  refresh the idle timer without touching llama.cpp; they now decrement the
  in-flight count without keeping the model warm.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-01 06:42:23 -07:00
Daniel Han
2ef394137a
Studio: harden background consumer loops and streaming paths against silent UI freezes (#6653)
* Studio: harden the data-recipe and inference consumer loops against pump death

Follow-up to #6643. The same single-unsupervised-consumer pattern the training
pump had lives in two sibling loops, with the same failure mode: one bad event
kills the only thread that updates the in-memory state every UI surface reads,
while the worker subprocess keeps running.

- data_recipe JobManager._pump_loop: a malformed worker log line that makes
  parse_log_message raise no longer kills the pump. Guard _handle_event, the
  queue read, and the worker-exit finalize, and broaden _drain_queue so a drain
  error still finalizes the job instead of leaving it wedged "active" (which also
  leaked the workflow-scoped API key until its 24h expiry).
- inference InferenceOrchestrator._dispatcher_loop: guard the routing body so a
  malformed response or a mailbox put error can't kill the dispatcher and hang
  every in-flight generation (callers key liveness on the subprocess, not on
  this thread).

Adds regression tests for both.

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

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

* Studio: extend consumer-loop hardening to RAG, hub, auth, and stream-reader paths

Continuation of the data-recipe and inference pump hardening: the same
"background producer updates in-memory state that a single unsupervised
consumer surfaces to the UI" pattern shows up in several more Studio paths,
each able to silently freeze a UI surface while the worker keeps running.

RAG ingestion SSE (core/rag/ingestion.py):
- job_events polled the queue with a blocking get and never noticed client
  disconnect or a dead worker, so a closed tab or a producer that died
  without emitting a terminal event left the stream hanging. It now polls
  with a timeout, emits heartbeats, ends on terminal job status, caps idle
  time, and always pops the job registry in finally.
- Added _reap_finished_jobs() and call it from start_ingestion so finished
  job state does not accumulate.

Startup reconcile (storage/rag_db.py, main.py):
- reconcile_orphaned_ingestion_jobs() marks ingestion jobs (and their
  documents) that were left non-terminal by a previous crash as failed, so
  the UI does not show jobs stuck "running" forever after a restart. Wired
  in at startup next to cleanup_orphaned_runs().

Hub download watcher (hub/services/download_lifecycle.py):
- _watch() could leave a job pinned "running" if finalize raised. Body is
  now guarded: on failure it logs and sets the job to error, and always
  invalidates the hf cache scan in finally.

External provider stream (core/inference/external_provider.py):
- read timeout was None (no stall ceiling); set to 300s so a wedged
  upstream surfaces as an error instead of an indefinitely hung stream.

Auth store (auth/storage.py):
- Enable WAL + busy_timeout on the auth DB so token validation (read on
  every request) and login writes stop serialising on the rollback journal.
  Matches studio_db / rag_db / providers_db.

Login rate limiter (routes/auth.py):
- _LOGIN_IP_BUCKETS could grow unbounded under spoofed-IP traffic; cap it
  and prune stale buckets, mirroring the per-account bucket handling.

Training progress SSE (routes/training.py):
- Break promptly on client disconnect instead of waiting for the next
  yield to fail on a closed socket, matching the export / data-recipe SSE
  routes.

llama-server stdout drain (core/inference/llama_cpp.py):
- Broaden the drain guard so an unexpected decode/read error logs at debug
  and stops the drainer cleanly instead of escaping the thread.

Frontend stream readers (chat-api.ts, rag-api.ts):
- Wrap the SSE read loops in try/finally + reader.cancel() so early return
  ([DONE]), thrown errors, and consumer aborts release the reader lock
  instead of holding it until GC.

Tests:
- test_training_progress_stream_nan: fake request now implements the async
  is_disconnected() the route polls, matching the other SSE route fakes.

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

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

* Studio: address Codex review feedback on the consumer-loop hardening

Four follow-ups from the automated review, all on code this PR introduced:

- Data-recipe pump (manager.py): a queue read that keeps raising an error
  outside the read's narrow catch set (e.g. a broken queue pipe after the
  child died) hit the `continue` guard and skipped the dead-worker finalize
  below, spinning forever and leaving the job wedged "active" with its
  workflow key unretired. On a read failure, fall through to finalize when
  the worker is no longer alive. Added a regression test.

- RAG ingestion SSE (ingestion.py): the 5-minute idle cap could end the
  stream while the job was still pending/running (a large document spends
  minutes in embedding/storing with no per-batch progress event). The route
  then sends [DONE], and the client treats a no-terminal-frame end as
  completion, marking the document indexed mid-ingestion. Drop the idle cap:
  while the worker is alive and non-terminal we keep heartbeating; the stream
  ends only on terminal DB status, the None sentinel, or client disconnect.

- Login rate limiter (auth.py): the per-IP path pruned but then added the
  new IP unconditionally, so a spoofed-source-IP spray kept _LOGIN_IP_BUCKETS
  unbounded and made every new IP pay a full-dict prune scan. Gate the add on
  the cap, mirroring the account path.

- Hub download watcher (download_lifecycle.py): if finalize raised before it
  reaped (proc.wait) and dropped the worker (e.g. an I/O error draining
  stderr), the crash path published a terminal state while the live Popen
  stayed registered and kept writing the cache, and the terminal set_job let
  claim() admit a retry on the same repo. Terminate + drop the worker before
  setting the terminal state.

* Studio: keep login throttling working when the per-IP bucket dict saturates

Review follow-up. The previous cap fix skipped creating a bucket for a new IP
once _LOGIN_IP_BUCKETS was full, returning ip_fails=0. Under a sustained spray
that also fills the account dict, every failure from such an IP then looked
first-seen and _login_blocked had no bucket to enforce, so the cap effectively
disabled throttling once saturated.

Bound the dict with a FIFO eviction instead: if the IP is new and the dict is
full, reclaim expired buckets (rate-limited so a burst of distinct IPs can't
make each failure an O(n) sweep) and, if still full, evict the oldest-inserted
IP. The new IP always gets a real bucket, so a saturating (e.g. spoofed
X-Forwarded-For) spray stays throttled while memory stays bounded. Added a
regression test that saturates the dict and asserts a later IP is still blocked.

* Studio: address Codex review (RAG queue lifecycle, stream error, orphan chunks)

Three follow-ups on the Phase 6 changes:

- RAG ingestion SSE (ingestion.py): job_events removed the per-job queue in its
  finally on ANY exit, including an early client disconnect while the worker is
  still running. That dropped the worker's later events (the queue is the only
  one _emit writes to) and made a reconnect find no queue and receive only
  [DONE], which the client treats as completion. Only drop the queue on a
  terminal exit (None sentinel / terminal DB status); leftover terminal queues
  are still swept by _reap_finished_jobs. Added queue-lifecycle tests.

- External provider stream (routes/inference.py): once the 300s read timeout can
  fire, the stream's except path failed the monitor but ended without an error
  frame or [DONE], so the chat client saw a bare EOF and saved the timed-out
  answer as a successful partial with no error. Emit an SSE error frame (and
  [DONE]) on stream failure so the client surfaces it.

- RAG startup reconcile (storage/rag_db.py): marking a half-ingested document
  failed left its chunks/fts/vec rows intact, and retrieval filters by scope not
  status, so a failed document could still be retrieved and cited. Purge the
  document's chunks when reconciling it to failed (the doc row stays for
  re-ingest).

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

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

* Studio: release the remaining SSE stream readers (training, data-recipe, export)

reviewer.py follow-up. The chat and RAG SSE readers were wrapped in
try/finally + reader.cancel(), but the other three readers built on the same
response.body.getReader() pattern were left without it: streamTrainingProgress,
streamRecipeJobEvents, and streamExportLogs leak the ReadableStreamDefaultReader
lock (held until GC) when the consumer aborts, returns early, or a parse/callback
throws. Wrap each in try/finally + reader.cancel() (export already had a
try/catch, so it only needed the finally). All five frontend SSE readers now
release the reader symmetrically.

* Tighten resilience comments and docstrings

Condense the verbose explanatory comments and internal-helper docstrings added
in this branch to shorter, clearer forms. Comment/whitespace only; verified no
code changed via AST diff. No behaviour change.

* Studio: keep chunks for completed docs during ingestion reconcile

Startup reconciliation flips orphaned (non-terminal) ingestion jobs to failed and
purges the document's chunks so a failed source can't be retrieved. But it dropped
the chunks unconditionally, so a document the worker had already committed as
'completed' before the crash (only its job row left non-terminal) lost every chunk
while still reporting 'completed'. That leaves an empty source that retrieval can't
return and dedup (status != 'failed') blocks from re-ingest.

Only purge chunks when the document UPDATE actually transitions it to failed; an
already-completed document keeps its chunks. Adds reconcile regression tests for
both the completed-doc and genuine in-flight-orphan cases.

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

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

* Studio: drop a finished RAG job's queue when the client disconnects

job_events kept the per-job queue until it consumed the None sentinel, so a UI
that stops on the terminal event (its reader.cancel aborts the stream before
[DONE]) left the queue registered until the next _reap_finished_jobs sweep; a
batch of uploads followed by idling retained them all.

_run writes the terminal DB status before emitting the terminal event, so on
generator exit, drop the queue when the job's DB row is already terminal (worker
done, nothing to resume) and keep it only while the worker is still running. Adds
a disconnect-after-terminal-event regression test.

* Remove stray async task output files committed by mistake

* Studio: harden login IP throttle and end progress stream on disconnect

Two Codex review items:

Login per-IP throttle: when the per-IP bucket dict saturated, FIFO eviction could
drop a still-hot (blocked) bucket, so an IP could flood the dict with distinct
(or spoofed) source IPs to push out its own bucket and retry as first-seen. Stop
evicting hot buckets; a new IP that can't fit now shares a bounded overflow
counter that still trips the per-IP threshold, so a saturating spray stays
throttled and no live counter is reset.

Progress SSE: on client disconnect the polling loop only broke and fell through
to the unconditional final 'complete' frame, so a buffered or proxying consumer
could read a still-active run as completed. Return from the generator instead.

Adds regression tests for both (spray cannot reset a hot bucket; disconnect while
active emits no complete frame).

* Studio: shard the login overflow counter and stop cancelling chat stream after [DONE]

Two Codex review items:

Login throttle overflow: the single shared overflow counter meant that once a
saturating spray pushed it past the per-IP threshold, _login_blocked returned 429
for every new unbucketed source IP, before credentials were checked -- a global
login denial. Shard the overflow into a fixed array of counters keyed by hash(ip),
so a hot shard only throttles the IPs that map to it while a single source's
repeated failures still concentrate in one shard and stay throttled. Memory stays
bounded and no live bucket is evicted. Adds a regression test that a hot overflow
shard does not block an unrelated IP.

Chat stream: the reader.cancel() in the SSE finally fired even after a natural
[DONE]/EOF. The backend finalizes its api-monitor entry right after yielding the
sentinel (the local pass-through finishes after the last yield), so a client
cancel there can be observed as a disconnect and mark a completed request as
cancelled. Track natural completion and only cancel on an early/abnormal exit.
(No frontend unit test: the Studio frontend has no test harness.)

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

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

* Studio: give prep-timeout test fakes an is_disconnected method

The progress stream now ends on client disconnect (await request.is_disconnected()
before falling through to the terminal frame). After merging that into the
prep-timeout tests added later on main, their _FakeRequest/_ReconnectRequest must
provide is_disconnected or the generator raises AttributeError under CI.

* Studio: keep the login overflow throttle when bucket capacity frees up

_login_blocked only consulted the per-IP overflow shard while the bucket dict was
still at capacity. If a slot freed before the 60s window expired (e.g. another
IP's successful login calls _clear_login_bucket), a source counted in a hot shard
stopped being blocked and its next failure got a fresh per-IP bucket, resetting
the throttle the overflow path exists to preserve. Always max in the IP's shard
(shards are empty outside saturation, so it is a no-op in the common case). Adds a
regression test that a hot source stays throttled after a bucket frees.

* Studio: clear a login IP's overflow throttle on successful login

_clear_login_bucket reset the per-IP and per-account buckets on a successful
login but not the overflow shard, so after the dict saturated and an IP was
counted in overflow, a later successful login left those entries behind and the
next failed attempt could immediately return 429.

Store overflow entries as (timestamp, ip) so a source is throttled by its own
count within the shard (also removing cross-IP collateral within a shard), and
drop just that IP's entries in _clear_login_bucket. Adds a regression test that a
successful login clears the overflow throttle.

* Studio: bound the login overflow shard memory under high-cardinality spray

The per-IP overflow tracked failures in a time-pruned deque of (timestamp, ip)
tuples, so a spoofed-X-Forwarded-For spray of distinct one-off IPs grew memory and
the per-check scan with request cardinality for the whole window -- undermining
the bucket cap that exists to bound memory. Replace each shard with a fixed-
capacity dict (ip -> [count, window_start]): O(1) lookups, and when a shard is
full a one-off IP evicts the lowest-count entry (Space-Saving) so memory is hard-
bounded while a persistent attacker keeps a high count and is never evicted. Adds
a regression test that shards stay within the per-shard cap under a 5000-IP spray.

* Studio: purge chunks for already-failed docs during ingestion reconcile

The reconcile chunk-purge was gated on the documents UPDATE actually flipping a
non-terminal doc to failed. A doc the worker had already marked 'failed' before
the crash (job row left non-terminal) was not re-flipped, so its committed chunks
were kept and stayed retrievable/citable, since retrieval filters by scope not
status. Purge chunks whenever the document is not 'completed' (failed, in-flight,
or gone), preserving the completed-doc carve-out. Adds a regression test.

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

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

* Studio: don't inherit an evicted IP's count onto a new overflow source

When a full overflow shard evicted the lowest-count entry, the new source
inherited that count (Space-Saving base + 1). If a shard was saturated with hot
entries, an unrelated new IP could land at/over the threshold and be 429'd after a
single attempt -- cross-IP collateral despite the per-source-isolation intent.
New entries now start clean at count 1; the only cost is that a heavy hitter that
is the lowest-count entry in a fully saturated shard can briefly reset, which is
preferable to blocking a bystander. Adds a regression test.

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

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

* Studio: carry overflow failures into a new IP bucket on transition

_login_blocked took max(per-IP bucket, overflow shard) rather than combining them,
so a source could log (threshold-1) failures in overflow during saturation and,
once a bucket slot freed, another (threshold-1) in a fresh bucket within the same
window -- roughly doubling the per-IP limit. When a saturated-era IP first gets a
real bucket, migrate its windowed overflow count into that bucket (and drop the
overflow entry) so the combined failures throttle at the intended limit. Adds a
regression test.

* Studio: reconcile a completed doc's orphaned job to completed, not failed

When a crash left an ingestion job non-terminal after its document was already
committed as completed, reconcile marked the job failed. After restart the upload
UI has no in-memory SSE queue and falls back to getJob(), which treats a failed
job as an indexing failure and removes/toasts a document that is actually
searchable. Mark the job completed (keeping its chunks) when its document is
completed. Extends the completed-doc reconcile test to assert the job status.

* Studio: clamp the overflow failure count migrated into a login bucket

A saturated source could accrue an unbounded overflow count, then materialize
one deque entry per recorded failure when a bucket slot freed, allocating an
arbitrarily large deque under the login lock. Only at-or-above the per-IP
threshold matters for blocking, so cap the count there at the record and take
sites; the migration is now bounded without weakening the limit.

* Studio: keep the RAG job stream alive on a transient status read

The heartbeat poll read the job row unguarded; a momentarily-locked DB would
raise out of job_events, which the SSE route turns into a terminal error frame,
and the UI drops a document whose worker is still running. Treat a failed status
read as non-terminal: heartbeat and retry, and keep the queue so a reconnect can
resume.

* Studio: set busy_timeout before journal_mode on the auth DB

Switching journal_mode needs a lock, so if a refresh-token write already holds
one, journal_mode=WAL raises SQLITE_BUSY and the shared try leaves the
connection on SQLite's default zero lock wait. Set busy_timeout first so the
switch waits instead of failing.

* [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-06-26 03:31:33 -07:00
Daniel Han
7f45635280
Studio: auto-shut-down an exposed first-run instance if the admin password is never changed (#6651)
* Studio: set the admin password before exposing it on the network

On first run Studio seeds the default `unsloth` admin with a random
bootstrap password and embeds it into index.html (window.__UNSLOTH_BOOTSTRAP__)
so the local user can change it without typing it. A request with no Origin
header counts as same-origin, which is what a normal top-level GET sends, so
the page hands out the password to whoever loads it. That is harmless on the
default 127.0.0.1 bind, but `--secure` (public Cloudflare tunnel) and
`--host 0.0.0.0` (raw port reachable on the network) would serve the plaintext
admin password to remote visitors during the bootstrap window.

Fix this at the source: when launching a network-exposed web UI, prompt the
operator in the terminal for a real admin password (with confirmation) before
the socket binds or the tunnel opens, and persist it via update_password (which
clears must_change_password and deletes the .bootstrap_password file). After
that there is no bootstrap secret to leak. Non-interactive launches can supply
it via UNSLOTH_STUDIO_ADMIN_PASSWORD. The masked reader echoes '*' per
character and works on Linux, macOS, and Windows (PowerShell/cmd). Loopback
binds, --api-only (no web UI), and Colab are unaffected.

As defense in depth, the index handler now embeds the bootstrap object only for
a direct local navigation: same-origin AND a loopback TCP peer with no
proxy/tunnel forwarding headers (cf-ray, cf-connecting-ip, x-forwarded-for,
x-forwarded-host, x-real-ip, forwarded). Colab stays exempt. This keeps the
password off the wire even when the prompt is skipped (no TTY and no env var).

Adds unit coverage for the prompt/confirm/decision logic, an integration test
that provisioning clears the bootstrap state, and regression tests for the
local-direct gate (loopback/IPv6/mapped/localhost peers, LAN/public peers,
missing client, each forwarding header, spoofed XFF, and the Colab exemption).

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

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

* Studio: fail fast on an explicitly empty admin-password env var

resolve_admin_password_source treated UNSLOTH_STUDIO_ADMIN_PASSWORD="" like
the var was unset and fell back to the bootstrap backstop. Treat any set value
(including empty) as the env source so it reaches the minimum-length guard and
refuses to expose the server instead of silently keeping the seeded password.

* Studio: apply repo kwarg-spacing format to the secure-admin-password files

* Studio: drop the pre-exposure password prompt; keep the local-direct gate

Per review, the blocking prompt added friction for --secure / 0.0.0.0 first-run
launches without extra security: the local-direct injection gate in main.py
already keeps the bootstrap password off the network for any remote request.
Remove the prompt module and its tests; the gate plus the existing
must_change_password first-login flow are the fix.

* Studio: shut down an exposed first-run instance if the admin password is never changed

The local-direct gate keeps the seeded bootstrap password off the network, but
it stays a valid credential until first login changes it. For an exposed web UI
(--secure / 0.0.0.0, not --api-only, not Colab), arm a daemon timer: if the
password is still the seeded one after the deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT,
default 3600s, 0 disables), print a message and shut Studio down via the existing
graceful-shutdown path; if it was changed, leave Studio running.

* Studio: revert the local-direct injection gate; keep the 1-hour auto-shutdown

Per maintainer decision, keep the first-run auto-fill behavior unchanged (the
bootstrap password still seeds the login form for convenience) and rely on the
exposed-instance auto-shutdown to bound the window: an exposed web UI that never
changes the seeded admin password is torn down after UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT
(default 1h). Restores studio/backend/main.py and its origin test to upstream.

* Studio: render the bootstrap-timeout shutdown message with a human duration

The message hardcoded 'minute(s)' via timeout//60, so a sub-minute timeout
(e.g. a 30s test value) printed 'within 1 minute(s)'. Add _format_duration so
it reads '30 seconds' / '1 minute 30 seconds' / '60 minutes' as appropriate.
The default 3600s still renders '60 minutes'.

* Studio: drop stale local-direct gate reference from bootstrap_timeout docstring

The gate was reverted (timer-only), so the module docstring should not describe
a main.py gate that no longer exists.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-26 01:27:27 -07:00
Daniel Han
80d3434d61
Studio: require signed capability tokens for /p preview links (#6666)
* Studio: require signed capability tokens for /p preview links

The public /p preview routes added in #6486 run model load and chat
generation as the admin user with no authentication. The only gate is the
preview ref, a deterministic outputs-root path (run or run/checkpoint) that
is guessable rather than secret. On a network-reachable Studio (--secure
tunnel or -H 0.0.0.0), an unauthenticated caller who guesses a ref can
consume GPU and probe a private fine-tuned checkpoint.

Make the share link an unguessable, revocable capability:

- Sign the canonical ref with a dedicated server-side secret (HMAC-SHA256,
  stored in app_secrets, independent of the JWT/login secret).
- Require a valid token on every /p chat, models, and page request before
  resolving a checkpoint or loading a model; missing or invalid tokens get a
  generic 404 so the surface never confirms a ref exists.
- Accept the token via ?k= (browser link and preview page) or
  Authorization: Bearer (OpenAI-compatible clients).
- Rotate the secret to revoke every outstanding link
  (POST /api/settings/preview-links/rotate).
- Clamp preview generation (max_tokens/max_completion_tokens <= 1024, n = 1)
  and set Referrer-Policy: no-referrer on the page so the token is not
  leaked via Referer.

Training history hands the authenticated owner the signed token, and the
copy-link button builds /p/{ref}?k={sig}.

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

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

* Studio: honor a lower caller token limit in the preview clamp

Codex review: when only the legacy max_tokens was sent, the clamp left
max_completion_tokens at the 1024 default, and _effective_max_tokens prefers
max_completion_tokens, so a request like max_tokens=16 could still generate up
to 1024 tokens. Derive one effective limit (max_completion_tokens wins, else the
legacy max_tokens) and pin both fields to it so a caller's lower limit is kept.

* Studio: add preview kill switch, rate limit, and revoke-links UI

Follow-ups to the /p preview capability work:

- Public-sharing kill switch: a persisted setting (default on) gates the public
  /p surface. When off, every preview request 404s even with a valid token, and
  the owner UI stops offering share links. GET/PUT /api/settings/preview-sharing;
  enforced in _verify_or_404.
- Per-IP rate limit on the preview chat route: a coarse in-process sliding-window
  limiter (20 req/min/IP) returns 429 + Retry-After before the GPU lock is taken.
  Client IP honors X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is
  set, matching the login limiter's trust model.
- Settings UI: a "Preview sharing" section with the public-sharing toggle and a
  "Revoke all preview links" button (confirm dialog) that rotates the secret.

Tests cover the kill switch (404 when off), the 429 path, the sliding window,
client-IP trust behavior, and the setting default.

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

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

* Studio: fix preview-fields sharing arg and refresh sigs after revoke

Codex review:
- P1: get_training_run_detail and update_training_run called _preview_fields
  with only output_dir after it gained a required sharing_on parameter, raising
  a 500 TypeError once get_run succeeded. Pass get_preview_sharing_enabled() at
  both sites; add a detail-endpoint regression test.
- P2: after rotating the preview secret from settings, the history grid still
  held stale preview_sig values, so a freshly copied link would 404. Emit
  emitTrainingRunsChanged() after a successful revoke so the grid refetches
  freshly signed refs.

* Studio: harden preview sharing controls (Codex review)

- Fail closed: a read failure on the preview-sharing kill switch now returns
  False instead of defaulting to enabled, so an unavailable settings DB can't
  reopen the public surface. A missing key still defaults to enabled.
- Per-IP rate limit behind the managed Cloudflare tunnel: client_ip now honors
  CF-Connecting-IP when the socket peer is loopback, so tunneled visitors are
  keyed by their real IP instead of collapsing onto the local cloudflared peer.
- GET /p no longer mints key/share_url when sharing is disabled; it returns
  sharing_enabled=false so clients don't distribute links that 404.
- Settings UI: toggling public sharing emits the training-runs-changed event so
  the history grid shows/hides Copy preview link without a manual refresh.

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

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

* Studio: harden preview rate limiter and IP keying (Opus review)

From a two-agent review of the PR:

- Rate limiter no longer evicts an active bucket when the table is full: a flood
  of distinct keys could otherwise cycle out a throttled bucket and reset its
  counter. Evict only aged-out buckets; if the table is full of live clients,
  fail closed (deny the new key) instead.
- client_ip keys on the rightmost (proxy-appended) X-Forwarded-For hop when the
  trust env is set; the leftmost is client-spoofable. Documented the
  append/overwrite-proxy assumption.
- _verify_or_404 checks the capability token before the kill-switch DB read, so
  unauthenticated /p spam can't be used as an unbounded settings-DB sink and the
  response is identical regardless of the sharing on/off state.

Tests: nested run/checkpoint happy path + wrong-ref rejection, the eviction
fail-closed behavior, and route-level coverage for the rotate / preview-sharing
settings endpoints.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-25 21:40:48 -07:00
Daniel Han
9b5c94df32
CLI: stop unsloth connect from leaking Studio credentials to unverified servers (#6479)
* CLI: stop `unsloth connect` from leaking Studio credentials to unverified servers

`unsloth connect` (and `unsloth chat`) discovered a Studio base URL from
UNSLOTH_STUDIO_URL or the default localhost port after only an unauthenticated
/api/health probe, then sent credentials to it:

- keyless connect iterated every cached API key and sent each as a bearer token
  to {base}/v1/models, so a malicious or port-preempting endpoint could harvest
  all of them;
- with no cached key it self-issued a Studio JWT and POSTed it to
  {base}/api/auth/api-keys;
- unsloth chat sent the same self-issued JWT to the discovered base.

The key cache was a flat, global list with no binding to a server identity, so a
key minted for one Studio could be replayed to any other.

Changes:

- Scope the agent key cache per base URL so a key is only ever replayed to the
  exact server it was minted for. Pre-scoping flat caches are ignored rather
  than replayed (at most one extra local mint on the next launch).
- Gate every automatic credential flow to loopback bases. A non-loopback
  UNSLOTH_STUDIO_URL now requires an explicit --api-key and nothing is sent
  automatically. SSH-tunnelled Studios that land on 127.0.0.1 keep working.
- Mint the API key locally against the Studio auth DB instead of POSTing a
  self-issued JWT over the network, so no bearer token leaves the process on the
  local path.
- Apply the same loopback gate to connect_studio_server (used by unsloth chat).

Fully closing same-host loopback-port preemption needs a signed /api/health
handshake so the client can verify the server identity before sending anything;
that is tracked as a server-side follow-up.

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

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

* CLI: verify Studio server identity before auto-sending credentials

Adds a challenge-response so `unsloth connect` and `unsloth chat` can confirm a
discovered loopback endpoint is really this install's Studio (not a process
that preempted the port) before sending it a cached or freshly minted
credential. This closes the same-host loopback-preemption gap left open by the
previous commit, which could only limit the blast radius.

Server:

- storage.get_or_create_identity_secret(): a dedicated server-wide secret in
  app_secrets (kept separate from the per-user JWT secret), readable only by
  the same OS user.
- storage.compute_identity_proof(nonce) = HMAC-SHA256(identity secret, nonce).
- GET /api/auth/identity?nonce=<base64url>: unauthenticated, returns the proof.
  The nonce is opaque to the server and the proof reveals nothing about the
  secret, so answering is safe.

Client:

- verify_studio_identity(base): sends a fresh 32-byte nonce, recomputes the
  expected HMAC from the local same-user secret, and constant-time compares.
  Fails closed on any error.
- connect._agent_api_key gates the loopback cached-key replay and the local
  mint on it; connect_studio_server (used by unsloth chat) gates the
  self-issued JWT on it.

A server that cannot read this install's secret (a different OS user, or a
remote/fake endpoint) cannot produce a matching proof, so the client refuses
and falls back to an explicit --api-key.

Tests: studio/backend/tests/test_identity.py (proof determinism, secret
persistence and caching, route response and nonce validation) and additions to
test_connect.py (the verify gate refuses when unverified, an explicit key skips
the check, and an end-to-end client plus server handshake against a stub HTTP
server).

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

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

* CLI: mint through the verified server instead of the local auth DB

CodeQL (py/clear-text-storage-sensitive-data) flagged the API key written to
the per-server cache and the agent config files once it was sourced from
storage.create_api_key(); the original HTTP-minted key did not trip the query.

Now that the identity handshake cryptographically confirms the loopback
responder really is this Studio before anything is sent, minting through the
server's /api/auth/api-keys endpoint with a self-issued JWT is safe again and
restores the original, CodeQL-clean data flow. The local-DB mint path is
removed.

The security properties are unchanged: discovery is still loopback-gated and
identity-verified, the key cache is still scoped per server, and a credential
reaches the server only after the handshake has proven its identity. The only
difference from the previous commit is that the self-issued JWT is sent to the
already-verified loopback server rather than the key being minted in-process.

Tests updated to mint through the fake server again.

* CLI: address review feedback on connect credential handling

- Reuse a saved per-server key before the loopback/identity gate. Keys are
  scoped per base URL, so a key the user saved with --api-key for a remote or
  SSH-tunnelled Studio (whose identity secret the local handshake can't match)
  is replayed only to that exact server. The loopback + identity-handshake gate
  now guards just auto-minting (self-issuing a JWT and creating a new key),
  which is the path that needs a cryptographically verified local Studio. Fixes
  keyless reuse being impossible for remote/tunnelled Studios the user had
  saved a key for.

- connect_studio_server (unsloth chat / inference): when the user explicitly set
  UNSLOTH_STUDIO_URL but the server can't be safely attached (non-loopback, or
  identity unverifiable), fail with a clear message instead of silently loading
  the model locally. Opportunistic discovery of the local default still falls
  back to a local load.

- Harden cache parsing: tolerate a corrupt or hand-edited cache where a base
  maps to a non-list (which would otherwise iterate a string into
  single-character "keys"), and read the cache as UTF-8.

Tests updated and added: saved-key replay without the handshake for both local
and remote bases, keyless mint still refused when the loopback server is
unverified, and connect_studio_server erroring on an explicit remote while
falling back locally on default discovery.

* CLI: harden connect handshake against relay and gate cached minted keys

Addresses review feedback on the credential handshake:

- Refuse HTTP redirects on credential-bearing requests (the identity handshake,
  /v1/models, key minting, and the chat HTTP backend). A process squatting the
  discovered port could 302 /api/auth/identity to the real Studio and relay its
  valid proof, or bounce a bearer-token request to another base, and urllib
  follows redirects by default. A shared no-redirect opener now treats any 3xx
  as an error.

- Give cached keys provenance. Keys the user supplied with --api-key are "saved"
  and replay without the handshake (needed for remote or SSH-tunnelled Studios
  whose secret the local handshake can't match). Keys we auto-mint are "minted"
  and replay only after the identity handshake, so a port squatter can't collect
  a previously minted localhost key just by answering the health check. New cache
  shape: servers[base] = {"saved": [...], "minted": [...]}.

Known residual: a different-OS-user process that squats the port and can also
reach a genuine same-secret Studio elsewhere on loopback can still manually relay
the identity challenge. Fully closing that needs the proof bound to the server's
real listening port, or OS-level peer-credential checks; tracked as follow-up. A
same-user attacker is out of scope, since it can already read the 0600 key cache.

Tests: redirect rejection in the handshake, minted-cache requiring the handshake
while saved-cache bypasses it, and the existing suites updated for the new cache
shape. unsloth_cli (206) and test_identity.py (5) pass.

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

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

* CLI: keep urllib imports function-local in the no-redirect opener

The repo's import-hoist safety linter (scripts/verify_import_hoist.py) flags
hoisting urllib to module level because it re-points the 'urllib' name in the
pre-existing HttpChatBackend._request scope. Build the no-redirect opener lazily
with function-local urllib imports instead, matching this module's convention,
and restore the local 'import urllib.request' in _request and
verify_studio_identity.

* test(identity): skip route tests when routes.auth import chain is unavailable

The identity route tests build a TestClient from routes.auth, which pulls the
whole routes package (routes/__init__ -> inference -> llama_cpp, ...). In a
minimal test matrix without the heavy backend deps, or when another test in the
same process has already broken that import chain, importing it raised and the
two route tests hard-failed. Skip in that case instead: the proof crypto is
covered by the storage-level tests, and the full backend CI still exercises the
route. No behaviour change where the deps are present (5 passed in isolation).

* test(connect): make connect tests pass on native Windows

unsloth connect supports Windows: --no-launch prints PowerShell ($env:X =
"v" / Remove-Item Env:X) instead of POSIX (export/unset), and the launch
path bridges env into a Windows agent .exe over WSLENV. The tests hardcoded
the POSIX shell forms, so on a real windows-latest runner 12 of them failed on
the assertion string even though every command exited 0.

Add OS-aware assertion helpers (_assert_env_set / _assert_env_unset) that check
the right shell syntax for the host OS, and skip the two WSL-from-Linux shim
tests on native Windows (os.name is 'posix' inside WSL, so that path can't run
there). No change on Linux/macOS (57 passed); the connect command's behaviour
is untouched. Validated on a windows-latest staging runner.

* style(connect): tighten comments in the credential-leak fix

Condense the verbose explanatory comments and multi-line docstrings added by
this PR to one or two lines each, drop a few that just restated the code, and
keep the security rationale where it is load-bearing. Comment/whitespace only;
verified with unslothai/scripts comment_tools.py (check --strip-docstrings:
6/6 'code unchanged'). Tests unchanged: connect 57 passed, identity 5 passed.

* CLI/Studio: harden the identity handshake (review round)

Addresses the latest Codex/Gemini review of the handshake:

- Store the identity secret privately. sqlite3.connect created the auth DB
  world-readable under a 022 umask, so another OS user could read app_secrets
  and forge proofs, defeating the same-user assumption the handshake rests on.
  The auth dir and DB are now restricted to owner-only (0700/0600); the JWT
  secret and password hashes there get the same protection.

- Bind the proof to the server's listening port. The stateless HMAC(secret,
  nonce) was relayable: a process squatting the discovered port could proxy the
  challenge to the real Studio on another port and pass it back. The proof now
  covers the port the server actually listens on (from the socket, never the
  Host header) and the client checks it against the port it connected to, so a
  relayed proof from a different port no longer matches. Closes the manual-relay
  residual left after the redirect fix.

- Cap the identity response read (the server is still unverified at that point)
  and serve the identity route from a sync def so its first-call SQLite read
  runs in the threadpool instead of the event loop.

Tests: port-bound proof + relayed-proof rejection added; identity (5) and
unsloth_cli (58) suites pass. Verified end to end against a real backend
(auth DB owner-only, handshake + mint still succeed).

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

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

* CLI/Studio: bind the identity proof to the connection address, not just port

Follow-up to the port binding from the last review. Binding only to the port
left a cross-address relay: a squatter on a different loopback address but the
same port (for example localhost resolving to a squatter on ::1 while the real
Studio is on 127.0.0.1) could proxy the nonce to the real Studio and pass back
a proof that still matched, since both share the port.

The proof now covers the address and the port the connection landed on:

- Server: takes the address+port from request.scope, which uvicorn populates
  from getsockname, so it is the real local address the client reached even
  when Studio is bound to 0.0.0.0 (verified empirically), never the
  client-controlled Host header.

- Client: resolves the base host to one concrete IP, talks to exactly that IP,
  and binds the proof to (IP, port). A proof relayed from a Studio on a
  different address or port was computed for that other endpoint and no longer
  matches the one the client dialed.

Both sides normalise the address through ipaddress so equivalent forms compare
equal. Together with the private-secret and redirect fixes, this closes the
cross-user loopback relay an attacker can mount without reading the secret.

Tests: proof now bound to host+port; relayed-proof rejection retained; identity
(5) and unsloth_cli (58) suites pass. Verified end to end against a real backend
(handshake + mint still succeed).

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

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

* CLI: pick the loopback address at discovery so localhost does not regress

find_studio_server now resolves a bare localhost base to its concrete loopback
addresses and returns the first that answers /api/health, IPv4 127.0.0.1 first
(where unsloth studio binds by default). The whole flow (health probe, identity
check, credential send) then targets that one address instead of racing
IPv4/IPv6 resolution, where localhost could resolve ::1-first and hide a Studio
bound to 127.0.0.1. A literal IP or remote name is unchanged.

* [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-06-21 21:28:38 -07:00
Daniel Han
40c8ad78b9
Studio: add --secure Cloudflare-only mode and revamp API usage examples (#6300)
* Studio: add --secure Cloudflare-only mode and revamp API usage examples

--secure / --not-secure on `unsloth studio` and `unsloth studio run`:
- --secure binds 127.0.0.1, requires the Cloudflare tunnel, and advertises only
  the Cloudflare link. cloudflared reaches the server over localhost, so the raw
  port is never exposed on a public interface.
- If the tunnel cannot start, fail closed with a clear message instead of
  silently leaving a raw 0.0.0.0 link.
- Default stays not-secure (no behavior change); coexists with the existing
  --cloudflare/--no-cloudflare flag. Host defaults are unchanged.
- /api/health (authed) now reports the live tunnel URL.

API usage examples (Profile > API):
- Example tabs for curl, Python, curl + tools, Python + tools, plus an OS row
  (Linux/macOS/WSL vs Windows) auto-detected from the platform.
- Windows curl passes the JSON body via a file so PowerShell does not strip the
  quotes when calling curl.exe.
- Python + tools forwards enable_tools/enabled_tools through extra_body and
  guards chunk.choices, since tool-lifecycle events carry no choices.
- Shows the loaded model name and the real API key while it is still revealed.
- A Cloudflare Tunnel toggle (default on) shows the public tunnel URL and uses
  it as the base_url in the examples when a tunnel is running.

Tests cover the tunnel start gate and the --secure flag on both commands.

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

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

* Studio: gate --secure tools on public exposure and harden API examples

In secure mode the server binds loopback but is reachable via the public
Cloudflare tunnel, so resolve the tool policy against the public exposure
(0.0.0.0) rather than the loopback bind. This keeps server-side tools off by
default and prompts before enabling them, instead of inheriting the loopback
default of on. The startup tool notice now names the public surface.

Also reject --secure with --no-cloudflare directly in run_server and the
run.py argparse (not only the CLI), JSON-encode interpolated model names so
Windows paths and quotes cannot produce invalid JSON or broken snippets, and
force-refresh /api/health on the API panel so a tunnel that starts after the
first health read still surfaces its URL.

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

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

* Studio: API examples show direct host when tunnel toggle is off; move Copy onto code

The Cloudflare Tunnel toggle had no visible effect when Studio was opened
through the tunnel: the off state fell back to window.location.origin, which
equals the tunnel URL in that case. /api/health now reports the direct
host:port (server_url), and the API panel uses it for the off state so it shows
the real non-tunnel base. Also move the Copy button out of the tab row and onto
the code block.

* Studio: highlight API examples, add advanced tabs, fix tunnel toggle row

Syntax-highlight the curl/PowerShell/Python snippets with the app's shared
shiki plugin (bash/powershell/python). Add 'curl + advanced' and
'Python + advanced' tabs that set temperature/top_p/top_k/min_p/
repetition_penalty/max_tokens, enable thinking, and turn on all tools.

The Cloudflare Tunnel row no longer shifts the code block: the tunnel URL is
always rendered (dimmed when off) so toggling keeps the row height constant.
Key the highlighted block on its content so it remounts when only the base URL
changes (the renderer's block memo otherwise kept a stale URL).

* Studio: rename API tunnel toggle to Secure HTTPS, hint --secure when exposed

Rename the API examples toggle from Cloudflare Tunnel to Secure HTTPS. When the
server was not launched with --secure, show an info tooltip noting the raw
0.0.0.0 port is still globally reachable and pointing at --secure. /api/health
now reports whether --secure was used so the hint is hidden in secure mode.

* Studio: force tools off for plain network/secure launches

The plain 'unsloth studio --secure' (and '-H 0.0.0.0') launcher re-execs run.py
and never installed a tool policy, so the process default (honor per-request
enable_tools) let any API-key holder run Python/terminal tools over the public
endpoint. Force the policy off at the run.py entrypoint when network-reachable
(0.0.0.0 or --secure); 'unsloth studio run' still installs its own resolved
policy and does not go through this path.

* Studio: apply default tool policy in run_server, not the run.py entrypoint

The plain launcher runs from the studio venv and calls run_server directly, so
it never hit the run.py __main__ guard. Move the network/secure default-off tool
policy into run_server so every launch path (plain, --secure, direct run.py)
gets it; the run subcommand still overrides it with its resolved policy.

* Studio: clarify --secure help text on the network exposure tradeoff

Spell out in --help (both unsloth studio and unsloth studio run, plus the
run.py argparse) that --not-secure also serves the raw 0.0.0.0 port reachable
from anywhere on the network, matching the API panel's Secure HTTPS hint.

* Studio: cache API-key PBKDF2 derivation to cut per-request /v1 auth overhead

validate_api_key re-ran the 100k-round PBKDF2 on every authenticated
request, adding ~15ms to each /v1 call made with an sk-unsloth- key.
Benchmarked against the bare llama-server it proxies to, API-key requests
carried ~22ms of fixed overhead vs ~7ms for the JWT path; the gap was
entirely this redundant key derivation (Pydantic validation measured
0.005ms, so it is not a factor).

The raw-key to hash mapping is a pure deterministic function of the fixed
server salt, so memoize it per process, keyed by a salted HMAC of the key
(never the key or a recoverable digest). The cached value equals what is
already stored at rest. Revocation and expiry remain enforced by the
SQLite read on every call, so a cache hit only skips the KDF, never the
active or expiry checks. Only keys that exist in the DB are cached, so
unknown-key spam cannot grow it.

After the change the API-key /v1 overhead drops to ~8ms, at parity with
JWT, while the at-rest PBKDF2 hashing is unchanged.

Adds test_api_key_expiry.py covering API-key and JWT expiry enforcement
and the new cache: it skips the KDF on repeat and still rejects revoked
or expired keys.

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

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

* Studio: tighten comments across the secure-tunnel and API-key changes

Condense multi-line comments and docstrings to one or two lines, drop the
ones that restate obvious code, and remove an orphaned test section header.
Comment-only: verified with comment_tools.py check (9/9 code unchanged), the
auth/secure-tunnel/CLI test suites, and a clean frontend typecheck and build.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-15 04:18:15 -07:00
Daniel Han
187144d4e7
Reduce and tighten code comments and docstrings repo-wide (#6095)
Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:09:51 -07:00
Daniel Han
8292e699e4
Studio: make code comments and docstrings more succinct (#6029)
Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:07:28 -07:00
Daniel Han
3ce187da02
Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
2026-06-08 04:24:13 -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
b09aa82a3a
Studio: add github_repo seed reader and GitHub Support Bot recipe (#5169)
* Studio: add github_repo seed reader and GitHub Support Bot recipe

Adds a first-party Data Designer seed reader that scrapes GitHub issues,
pull requests, and commits from one or more repositories via the GraphQL
API, and a learning recipe (GitHub Support Bot) that turns those rows into
synthetic support Q&A pairs for fine-tuning.

Backend (new plugin studio/backend/plugins/data-designer-github-repo-seed):
* GitHubRepoSeedSource config: repos, token (falls back to GH_TOKEN /
  GITHUB_TOKEN env var), item_types (issues / pulls / commits),
  per-resource limit (0 means all), max_comments_per_item.
* Rate-limit-aware GraphQL client (GitHubClient + RepoScraper) shared
  across repos; flattens each item into a uniform row with columns
  item_type, repo, number, title, body, state, author, created_at,
  closed_at, url, labels, comments.
* Registered via the data_designer.plugins entry point.

Frontend:
* New seed_github block variant so the seed node card shows
  "GitHub repositories" instead of the generic "Document file"
  placeholder, with its own icon and inline summary (repo count +
  item-type list).
* Rewritten seed dialog github_repo form: repos textarea pre-filled with
  unslothai/unsloth + unslothai/unsloth-zoo, password input for the GH
  token, items-per-repo number with an "All" toggle, and the noisier
  options (item types, max comments, include comments) tucked under an
  Advanced collapsible.
* Local model auto-load on Run: if a recipe uses an is_local provider
  and the inference server is not already serving that model, the
  executions hook calls /api/inference/load first. Removes the "open
  /chat to load a model" prerequisite that users kept tripping on.
* Honor the recipe's run.rows value in the Run dialog (previously the
  store reset to 5 regardless of what the template shipped).

Recipe (studio/frontend/src/features/data-recipes/learning-recipes/
github-support-bot.json):
* Defaults to the Local Model provider + unsloth/gemma-4-E2B-it-GGUF.
* Scrapes unslothai/unsloth and unslothai/unsloth-zoo, issues and pulls,
  up to 100 items per resource.
* Two LLM blocks: normalized_question (llm-text) rewrites each thread
  into a clean support question, support_answer (llm-structured)
  produces JSON with answer / diagnosis_questions / cites / confidence.
* Run defaults to 10 rows for a quick smoke test.

Verified end-to-end on a running Studio: card renders, source-data
dialog is pre-populated, All toggle disables the limit input, the
recipe executes and produces rows against a loaded local GGUF.

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

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

* fix: improve GitHub recipe support

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

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

* Studio: speed up GitHub scraper and harden the support-bot recipe

Addresses a perf issue found while demoing the github_repo seed reader:

Scraper is too slow at scale. The PRs GraphQL query pulls deeply nested
fields (reviewThreads, reviews, commits, timelineItems, etc.) so the
page size was pinned at 3 to stay under GitHub's node-count ceiling. 100
PRs meant 34 serial round trips. Added lighter query variants
(PRS_PAGE_QUERY_LIGHT, ISSUES_PAGE_QUERY_LIGHT) that drop the fields the
Studio flatten layer does not use (it only reads title, body, state,
author, labels, comments). With the light query PR pages can safely go
to 25 per page and issues to 50. The plugin scraper now passes
light=True to RepoScraper so Studio always uses the fast path; the heavy
query remains available for other callers.

Recipe defaults are now demo-ready with production knobs called out:
- max_parallel_requests: 1 and max_tokens: 800 so small local models
  stay stable when running the support_answer structured column.
- support_answer prompt trimmed to 80-200 words so gemma-4-E2B GGUF can
  actually comply with the schema. The canonical 150-300 word codex
  prompt is still documented in the node3 markdown note for
  production upgrades.

* Studio: rename GitHub recipe to 'GitHub Scraper' and add Easy mode

Changes the recipe framing from a single-purpose 'Support Bot' pipeline
to a general-purpose scraper that produces {user_request,
grounded_response} training pairs. Aligns with the canonical
github_data_gatherer dataset (11 enrichment tasks mirrored in pr_requests_20
/ issue_requests_20 on the input side and explain_pr / issue_fix_plan /
issue_solution on the output side).

Recipe JSON changes:
- columns[0] renamed normalized_question -> user_request, prompt now
  inverts a GitHub thread into a realistic user ask instead of
  normalising it.
- columns[1] renamed support_answer -> coauthor_response, emits
  {response, followups, cites, task, confidence} and branches on
  issue vs PR thread type.
- Notes rewritten to document the 11-task catalog and the canonical
  production prompt to paste in for a full dataset backfill.

Frontend: Easy mode for github_repo recipes. The drag-and-drop canvas is
hidden behind an 'Advanced' tab; Easy mode is the default for any recipe
whose seed_source_type is github_repo. The Easy form reuses the existing
GithubRepoSeedForm (promoted to exported), adds a rows input bound to
previewRows, a model field bound to the model_config, and a single Run
button that calls runPreview() directly (no modal). Non-github recipes
see the same Editor / Runs tabs as before.

View mode persists per-recipe-id in localStorage under
recipe-studio:view-mode:<recipeId>.

* Studio: auto-detect server GH_TOKEN and widen Easy-mode detection

The GitHub seed form now fetches /api/data-recipe/seed/github/env-token
on mount and, when the server exposes a GH_TOKEN / GITHUB_TOKEN env var
and the token field is blank, shows a small 'Using server env var' badge
and swaps the placeholder text. The token value itself is never returned
to the UI.

Widens Easy-mode detection in recipe-studio-page.tsx so that recipes
saved before ui.seed_source_type was persisted also get the Easy tab:
falls back to recipe.seed_config.source.seed_type, which is always
present for github_repo seeds.

* fix: polish GitHub recipe UI

* Studio: default llama-server --threads to -1 (auto)

Previously we passed --threads only when the caller set an explicit
value, which meant llama-server fell back to its internal default.
That default has varied across llama.cpp builds (some versions use
hardware concurrency including hyperthreads, which hurts throughput on
CPU-heavy inference). Always passing --threads -1 pins the behaviour
to llama.cpp's auto-detect (physical cores).

Caller-supplied n_threads still wins when non-None.

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

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

* Studio: auto-switch Easy mode to Runs pane on run start

Easy mode had no progress island or canvas overlay, so after clicking Run
the only visible state was the button label flipping to "Running..." while
the screen otherwise stayed identical. This reads as stuck even though the
job is progressing.

Wire an onExecutionStart callback from recipe-studio-page.tsx through to
useRecipeExecutions so that when a run is kicked off from easy mode, the
page flips to the executions view where the Runs sidebar, progress bar,
rate/ETA panel, and live log are rendered. Advanced/editor mode keeps its
existing behavior and stays on the canvas (it already has the floating
ExecutionProgressIsland).

* fix: clean up GitHub scraper layout

* Studio: forward llm-structured output_format as llama-server response_format

Local GGUF runs of llm-structured columns used to generate the full
max_tokens budget before the prompt-level "return JSON in a ```json
fence" instruction got parsed. Small models (e.g. gemma-4-E2B-it)
routinely broke format, so each row took ~65s and frequently failed
with "No parsable JSON structure within ```json markdown fence".

For any local-provider model_config referenced by an llm-structured
column, clone the model_config and inject response_format into the
clone's inference_parameters. Uses llama.cpp server's flat shape
(tools/server/README.md):

    {"type": "json_schema", "schema": <output_format>}

Not the OpenAI-nested form; data_designer's OpenAI adapter forwards
response_format verbatim via facade._COMPLETION_REQUEST_FIELDS, and
llama-server's documented schema path expects the flat variant.

The clone is per (model_alias, column) so:
- llm-text / llm-judge columns that share the same alias keep
  free-form sampling.
- Each structured column gets its own schema, so columns with
  different output_formats don't collide.

Effect on gemma-4-E2B-it demos: every row parses cleanly, and the
model terminates immediately after the closing brace instead of
running to max_tokens. Net wall-clock is usually faster even though
grammar-constrained sampling is slightly slower per token.

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

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

* Studio: flip Easy to Runs pane before validation scrape, not after

Previously onExecutionStart fired inside runExecution, which runs AFTER
validateRecipe() -- and validation re-invokes the seed reader. For the
github_repo reader that is a full GraphQL scrape, so the user sat on a
"Running..." button with an otherwise unchanged Easy form for 10-15s
before anything moved.

Call onExecutionStart at the top of runWithValidation, right after we
have a payload to send. The view flips immediately; ensureLocalModelLoaded
+ validateRecipe now run against the Runs pane instead of a frozen Easy
form. runExecution still calls onExecutionStart downstream, but the
callback is idempotent (the page's easy -> executions guard skips the
second call), so no behaviour change for runs that pass validation.

If validation fails the toast + runErrors path still fires; the Easy
form's error banner still reads runErrors when the user switches back.

* Studio: unify data-recipe workflow auth on sk-unsloth-* keys

The previous commit (a61b4cc9) assumed storage.create_api_key(..., internal=True)
and storage.revoke_internal_api_key(key_id) existed, but those helpers were
only in the working tree, never committed. Recipe runs in local-model mode
were therefore crashing with 500 when _inject_local_providers tried to mint
a workflow key. This commit ships the missing pieces.

auth/storage.py:
- api_keys schema gains is_internal INTEGER DEFAULT 0 (with a guarded
  ALTER TABLE migration so existing auth.db files upgrade in place).
- create_api_key takes an internal=False kwarg; internal keys are flagged
  so they can be hidden from user-facing listings.
- list_api_keys takes include_internal=False so UIs never see workflow keys.
- New revoke_internal_api_key(key_id): id-only revoke for keys minted by
  non-user subjects (the JobManager does not know a username).

core/data_recipe/jobs/manager.py:
- JobManager.start accepts internal_api_key_id and stores it on Job so
  lifecycle handlers can revoke eagerly.
- _handle_event revokes on EVENT_JOB_COMPLETED / _ERROR / _CANCELLED.
- _pump_loop subprocess-died fallback also retires the key so a crashed
  worker cannot leak a live sk-unsloth-* beyond its TTL.
- Revocation is best-effort (swallow exceptions) -- the 24h TTL is the
  safety net if storage hiccups.

core/data_recipe/jobs/types.py:
- Job dataclass gains internal_api_key_id: int | None = None.

Replaces the bespoke 24h JWT path that jobs.py used to mint for local
providers. One mint/revoke/verify surface for every API key the server
issues, and revocation is now eager (seconds, not 24h) instead of TTL-only.

* Studio: plug workflow-key leak on unexpected create_job errors

Review follow-up on the sk-unsloth-* workflow-key lifecycle in
create_job. Previously the revoke handlers wrapped mgr.start(...) but
only caught RuntimeError and ValueError, and get_job_manager() sat
outside the try block entirely. Any other exception type (TypeError
from a mismatched kwarg, OSError from the queue write, etc.) would
bubble up to FastAPI and leave the minted key live until its 24h TTL.

Fix: one try block covers both get_job_manager() and mgr.start(), with
a trailing except Exception that revokes and re-raises. The
RuntimeError -> 409 and ValueError -> 400 paths are unchanged so
specific client-facing status codes still surface. Revocation is still
best-effort (_revoke_internal_api_key_safe swallows errors) because we
never want revoke failures to mask the original crash.

Severity is low -- the key can't bootstrap longer access and the 24h
TTL bounds the window -- but the reviewer's point stands: eager revoke
on every failure path is the right invariant.

* Studio: nest response_format under extra_body so pydantic accepts it

The previous commit dropped response_format at the top level of a cloned
model_config's inference_parameters, which BuilderConfig rejected with:

  ValidationError: Extra inputs are not permitted [type=extra_forbidden]
  data_designer.model_configs.1.inference_parameters.response_format

data_designer's BaseInferenceParams is a pydantic model with extra=forbid
and only a fixed set of fields (temperature, top_p, max_tokens,
max_parallel_requests, timeout, extra_body). The pass-through path for
anything the schema doesn't know about is `extra_body`, which the
OpenAI SDK spreads into the chat-completions request body at the top
level -- which is exactly where llama-server reads response_format from.

Inject under extra_body (merging with any existing extra_body contents)
so the clone validates. llama-server still receives
{"type": "json_schema", "schema": <output_format>} at the top level of
the request body, which is the flat shape llama.cpp's server expects.

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

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

* Studio: forward response_format to llama-server and fence-wrap the reply

Two-part fix for the llm-structured data-recipe path:

(1) The /v1/chat/completions proxy was dropping response_format. The
route's passthrough branch only triggered on tools / tool messages, so
requests carrying a JSON schema fell into the non-passthrough GGUF path
which calls generate_chat_completion (no response_format kwarg). The
schema never reached llama-server, so guided decoding was a no-op and
the model emitted free-form text that happened to parse a fraction of
the time. Widen the passthrough trigger and teach _build_passthrough_payload
to forward response_format so llama-server's GBNF grammar actually runs.

Guided decoding does not require supports_tools, so split the condition:
a request is now passthrough-routed if it carries tools/tool messages
(existing behavior) OR carries response_format (new). The vision guard,
streaming fork, and tools-choice defaulting are unchanged.

(2) data_designer's llm-structured parser looks for a ```json ... ```
markdown fence and discards anything else. Guided decoding emits only
the JSON object (the GBNF grammar has no fence tokens), so a
100%-valid schema-constrained run still ended up 0 ok / N failed with
"No parsable JSON structure within ```json markdown fence". In
_openai_passthrough_non_streaming, wrap each choice's content in the
expected fence when the caller asked for guided decoding. Already-fenced
content is left alone so other clients that prefer raw JSON are not
affected; the wrap is scoped to requests that carried response_format.

Net effect on the GitHub Support Bot recipe on a local GGUF: schema
actually binds during sampling, content arrives wrapped in the fence
data_designer expects, and generation terminates immediately after the
closing brace instead of running out to max_tokens.

* Studio: Easy mode runs a full run, capped at the user's row count

Easy mode used to call runPreview, which produces a test run: no
artifact persisted, reduced progress tracking, and framed in the Runs
pane as "Test run". The whole point of the form is to let a user kick
off a real dataset build with one click, so wire it to runFull instead
and bind the Rows input to fullRows (not previewRows).

runFull requires a non-empty fullRunName. The Easy form has no run-name
input, so seed a default on mount whenever Easy is active and
fullRunName is still empty. Uses `<recipe name> <iso-timestamp>` so
each Easy run gets a stable-ish default that still sorts chronologically
in the Runs pane. User can override it from the Advanced run dialog
before clicking Run.

Rename GithubScraperEasyView's rows props from previewRows/setPreviewRows
to rows/setRows so the view stays agnostic to which hook state the page
chooses to bind. Loading indicator now follows fullLoading.

* Studio: clamp GitHub scrape page size and memoize the materialization

Two wins for the "before Generating fires" gap on small previews:

(1) scrape_{issues,prs,commits} hardcoded per_page (50 / 25 / 100) and
only checked the trial limit AFTER the page was written, so a 1-row
Easy run still asked GitHub for a full 50-issue + 25-PR page, wrote
them all to JSONL, and then stopped because total_new already exceeded
the trial cap. Cap per_page at min(page_cap, trial_limit) so
github_limit=1 actually asks for first:1.

(2) GitHubRepoSeedReader.get_dataset_uri used to scrape fresh on every
invocation. data_designer calls the seed reader multiple times per
recipe job (validation, preview, per-column sampling), so a 2-repo
Easy preview ran the full GraphQL scrape three times back-to-back,
burning ~15s of dead air before any LLM generation began.

Added a module-level in-process cache keyed on
(repos, item_types, limit, include_comments, max_comments_per_item,
sha256(token)[:16]) that stores the JSONL path of the first
materialization. Subsequent calls with the same signature return the
cached path, guarded by a staleness check that drops the entry if the
file was tmp-cleaned. Raw token values never land in the key.

Net effect on a 1-row Easy run, 2 repos, limit=1: 2 GraphQL round
trips instead of ~12, and the first-to-Generating gap collapses from
~15s to roughly 2-3s.

* Studio: make Easy mode Rows input editable instead of snapping to 1

The Rows to generate input used type="number" with value bound directly
to the rows state and an onChange that coerced any non-positive parse
result back to 1. The moment the user pressed backspace to clear the
field, the parent re-rendered with value=1 and the caret jumped, making
it impossible to change the value without arrowing the browser's +/-
spinner.

Switch to a text input with inputMode="numeric" and pattern="[0-9]*"
(so mobile still shows a numeric keyboard, and the browser drops the
spinner buttons the user did not want). Add a local rowsText buffer so
the field can hold transient empty / partial digit strings while
editing without fighting the parent state; the canonical rows value
only advances when the buffer parses to a valid integer in [1, 10000],
and onBlur clamps back to 1 or 10000 if the user left it out of range.

No behavior change for valid numeric edits - the downstream runFull()
still sees a clean positive integer.

* Studio: expand dataset cells horizontally by column on click

Click a long cell to expand that whole column. Click again to collapse.
Replaces the prior row-level vertical expansion which made it hard to
compare cells across columns. State is scoped per execution and per
column; the row itself is no longer a click target.

* Studio: force expanded dataset column to grow wide enough to read

* Studio: disable thinking for local recipe inference and plumb the kwarg

Reasoning-capable models (gemma-3n, qwen3.5, etc.) emit a
<think>...</think> preamble ahead of the answer by default, which
roughly doubles the generated token count per row on a local GGUF
and pushes the actual answer past data_designer's json-fence regex
on llm-structured columns. Recipes want the terse answer, not the
scratchpad.

Two halves of the fix:

(1) routes/data_recipe/jobs.py: when _inject_local_providers walks
the recipe's model_configs to point them at the local endpoint, also
stash chat_template_kwargs={"enable_thinking": false} under each
config's inference_parameters.extra_body. OpenAI SDK spreads
extra_body into the top-level request body, so llama-server and the
Studio /v1/chat/completions route both see it.

(2) routes/inference.py: the chat-completions route previously
dropped chat_template_kwargs on the floor because the whitelist
body builder only forwarded known fields.

    - At the top of openai_chat_completions, lift
      chat_template_kwargs.enable_thinking from payload.model_extra
      onto the typed payload.enable_thinking field when the caller
      did not set the latter, so the non-passthrough GGUF path's
      generate_chat_completion(...) call honors the override.
    - Teach _build_passthrough_payload to forward a
      chat_template_kwargs dict, and have _build_openai_passthrough_body
      derive that dict from payload.enable_thinking so
      response_format requests (structured columns) also land at
      llama-server with the reasoning preamble suppressed.

Net effect on a 10-row support-bot run with gemma-4-E2B-it-GGUF:
responses arrive without <think> tags, wall-clock per call drops
roughly in half, and structured columns stop leaking reasoning
tokens through the GBNF-constrained output.

* Studio: update GitHub Support Bot learning recipe with maintainer layout

Replace the template with the hand-laid-out export from the maintainer
so note nodes ship with real x/y positions (scattered around the
graph instead of all stacked at x=480) and the edges / canvas pan look
correct on first load. Also picks up the maintainer's prompt tweaks and
output schema names (coauthor_response / user_request / followups / task /
cites / confidence).

Diff is mostly ui.nodes positions and prompt bodies; runtime shape is
unchanged (seed_config / columns still target model_1 against the Local
Model provider).

* Studio: auto-size dataset sample columns; wide text gets a wide column

Drop the per-column click-to-expand toggle and the 180-char truncation.
Every column now renders its full value. Columns with long text get a
min-w of 48rem so the text is readable without wrapping into a tall
block; narrow-content columns get a 12rem min-w. The table wrapper
already has overflow-x-auto, so wide-column totals cause a horizontal
scrollbar instead of cramming everything into the viewport.

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

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

* fix GitHub scrape progress

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

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

* add resetApiBase export for test setup

* Studio: rename github-support-bot output columns to User / Assistant

Previously emitted user_request and coauthor_response, which did not
match the canonical User / Assistant chat-pair shape that downstream
SFT consumers expect. Renamed the columns in the recipe JSON (columns,
UI node ids, edges, notes, prompt Jinja refs) and the matching copy in
the learning-recipes index, data-recipes-page, and easy view.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-04-24 12:02:03 -07:00
Wasim Yousef Said
a5eb2e3d50
Add tauri (#5144)
* add unsloth studio desktop app

* Fix review findings

- studio/src-tauri/tauri.conf.json: retarget updater to staging repo
  (danielhanchen/unsloth-staging-2); switch to unslothai/unsloth on upstream merge.
- studio/src-tauri/linux/postremove.sh: drop the interactive read loop and the
  /home/* iteration. Package maintainer scripts must stay non-interactive and
  must not touch other users' data.
- studio/frontend/src/app/auth-guards.ts: honor tauriAutoAuth() boolean. Failed
  auto-auth now redirects to /login; requireGuest/requirePasswordChangeFlow
  only redirect to /chat when auth succeeds. The new early-return on failed
  auth is intentional so the login / change-password flows remain reachable
  when desktop auth is not yet established.
- studio/frontend/src/config/env.ts: keep fetched=false on health failure so
  later calls retry instead of caching the client-side platform guess.
- studio/src-tauri/src/install.rs: pick the available system package manager
  (apt-get, dnf, zypper, pacman); AppImage bundles run on non-Debian distros.
- studio/frontend/src/lib/open-link.ts + markdown-text/sources callers: return
  boolean from openLink so callers only preventDefault on handled URLs; relative
  hrefs now navigate natively.
- studio/frontend/src/features/settings/tabs/about-tab.tsx: fetch(apiUrl(...))
  so the version request targets the backend port in desktop mode. The bare
  /api/health predates the Tauri webview (blame: the earlier onboarding commit,
  which ran with same-origin frontend/backend); in desktop mode the webview
  origin is tauri://localhost so the bare path fails.
- install.ps1: gate the install_python_stack.py hotfix on a sentinel comment
  instead of a content regex; append the sentinel after applying so reruns
  are unambiguous.
- unsloth_cli/commands/studio.py _write_auth_secret: use the atomic mkstemp +
  os.replace path on Windows too; chmod calls are wrapped in try/except OSError.
- studio/src-tauri/src/preflight.rs probe_existing_backends: fan out the health
  probes concurrently; desktop-auth status still runs sequentially per candidate.
  reqwest::Client is internally Arc-wrapped so the in-loop .clone() is a
  refcount bump, not a deep clone; annotated inline.
- studio/src-tauri/src/preflight.rs run_cli_probe: wait() after kill() to reap
  the child, matching probe_cli_capability.
- studio/src-tauri/src/process.rs + main.rs: add stop_backend_detached and use
  it from the tray quit handler so the 5s graceful-wait does not block the
  Tauri main loop. RunEvent::Exit keeps the synchronous safety-net call.
- studio/backend/main.py: drop the permissive localhost CORS regex in
  api-only mode; the explicit allow_origins list is sufficient.
- .github/workflows/release-desktop.yml: drop max-parallel: 1 so platform
  builds run in parallel, and lift releaseBody to an env var so the three
  tauri-action invocations share one source of truth.

* Fix review findings (loop 2)

- studio/backend/auth/storage.py update_password: clear_desktop_secret()
  alongside clear_bootstrap_password() so rotating the admin password
  also revokes any previously provisioned .desktop_secret. Without this,
  an old local desktop credential keeps minting fresh admin tokens via
  /api/auth/desktop-login after a password rotation.
- studio/src-tauri/src/desktop_auth.rs provision_desktop_auth: wrap
  cmd.output().await in tokio::time::timeout(30s). DESKTOP_AUTH_LOCK is
  held across the whole desktop_auth flow, and previously a hanging
  `unsloth studio provision-desktop-auth` subprocess would pin the lock
  indefinitely and freeze every subsequent desktop_auth call.

* Add review tests

* Consolidate review tests

Merge review-added tests into the existing studio/backend/tests/test_desktop_auth.py
(the PR's authoritative desktop-auth test file). Drops three scaffolding files under
tests/python/ in favor of five focused tests next to the tests they extend:
- test_update_password_clears_desktop_secret (runtime)
- test_update_password_on_unknown_user_leaves_desktop_secret_intact (runtime)
- test_cli_provisioning_delegates_to_storage_create_desktop_secret (source-level)
- test_cli_connect_auth_db_reads_storage_db_path (source-level)
- test_desktop_auth_provision_has_bounded_timeout (Rust source-level)

* Revert auth-guards.ts Tauri branches to unconditional form

The review loop on PR 5144 introduced a regression: the isTauri branch of
requireAuth redirected to /login when tauriAutoAuth() returned false, and
requireGuest / requirePasswordChangeFlow silently fell through on the same
condition. The Tauri desktop app authenticates via a local auto-generated
secret; it must never surface /login or /change-password to the user. A
failed auto-auth should let the startup layer retry, not expose a password
form.

Restore the three Tauri branches to the author's original unconditional
form (requireAuth: return; requireGuest / requirePasswordChangeFlow: throw
redirect({to: '/chat'})). Keep the rest of the review fixes -- the
apiUrl() fetch wrapping, authRedirect helper, and fetchAuthStatus refactor
are all legitimate improvements and are preserved.

* Revert release-desktop.yml to author's version

The review loop's workflow-file tweaks (drop max-parallel: 1, lift releaseBody
to an env var) are cosmetic. OAuth tokens cannot push workflow-file changes,
and fine-grained PATs cannot honor maintainerCanModify on a third-party fork.
Reverting the workflow file to wasimysaid's version lets the push go through
without needing a classic PAT with both repo and workflow scopes.

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

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

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-23 04:50:10 -07:00
Daniel Han
9a261aec5f
Studio: Expose openai and anthropic compatible external API end points (#4956)
* Studio: add API key authentication for programmatic access

External users want to hit the Studio API (chat completions with tool
calling, training, export, etc.) without going through the browser
login flow. This adds sk-unsloth- prefixed API keys that work as a
drop-in replacement for JWTs in the Authorization: Bearer header.

Backend:
- New api_keys table in SQLite (storage.py)
- create/list/revoke/validate functions with SHA-256 hashed storage
- API key detection in _get_current_subject before the JWT path
- POST/GET/DELETE /api/auth/api-keys endpoints on the auth router

Frontend:
- /api-keys page with create form, one-time key reveal, keys table
- API Keys link in desktop and mobile navbar
- Route registered with requireAuth guard

Zero changes to any existing route handler -- every endpoint that uses
Depends(get_current_subject) automatically works with API keys.

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

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

* Use actual origin in API key usage examples

The examples on /api-keys were hardcoded to localhost:8888 which is
wrong for remote users. Use window.location.origin so the examples
show the correct URL regardless of where the user is connecting from.

* Add `unsloth studio run` CLI command for one-liner model serving

Adds a `run` subcommand that starts Studio, loads a model, creates an
API key, and prints a ready-to-use curl command -- similar to
`ollama run` or `vllm serve`.

Usage: unsloth studio run -m unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL

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

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

* Add end-to-end tests for `unsloth studio run` and API key usage

Tests the 4 usage examples from the API Keys page:
1. curl basic (non-streaming) chat completions
2. curl streaming (SSE) chat completions
3. OpenAI Python SDK streaming completions
4. curl with tools (web_search + python)

Also tests --help output, invalid key rejection, and no-key rejection.
All 7 tests pass against Qwen3-1.7B-GGUF.

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

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

* Add /v1/completions, /v1/embeddings, /v1/responses endpoints and --parallel support

- llama_cpp.py: accept n_parallel param, pass to llama-server --parallel
- run.py: plumb llama_parallel_slots through to app.state
- inference.py: add /completions and /embeddings as transparent proxies to
  llama-server, add /responses as application-level endpoint that converts
  to ChatCompletionRequest; thread n_parallel through load_model
- studio.py: set llama_parallel_slots=4 for `unsloth studio run` path

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

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

* Make /v1/responses endpoint match OpenAI Responses API format

The existing /v1/responses shim returned Chat Completions format, which
broke OpenAI SDK clients using openai.responses.create(). This commit
replaces the endpoint with a proper implementation that:

- Returns `output` array with `output_text` content parts instead of
  `choices` with `message`
- Uses `input_tokens`/`output_tokens` instead of `prompt_tokens`/
  `completion_tokens` in usage
- Sets `object: "response"` and `id: "resp_..."`
- Emits named SSE events for streaming (response.created,
  response.output_text.delta, response.completed, etc.)
- Accepts all OpenAI Responses API fields (tools, store, metadata,
  previous_response_id) without erroring -- silently ignored
- Maps `developer` role to `system` and `input_text`/`input_image`
  content parts to the internal Chat format

Adds Pydantic schemas for request/response models and 23 unit tests
covering schema validation, input normalisation, and response format.

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

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

* Studio: add Anthropic-compatible /v1/messages endpoint (#4981)

* Add Anthropic-compatible /v1/messages endpoint with tool support

Translate Anthropic Messages API format to/from internal OpenAI format
and reuse the existing server-side agentic tool loop. Supports streaming
SSE (message_start, content_block_delta, etc.) and non-streaming JSON.
Includes offline unit tests and e2e tests in test_studio_run.py.

* Add enable_tools, enabled_tools, session_id to /v1/messages endpoint

Support the same shorthand as /v1/chat/completions: enable_tools=true
with an optional enabled_tools list uses built-in server tools without
requiring full Anthropic tool definitions. session_id is passed through
for sandbox isolation. max_tokens is now optional.

* Strip leaked tool-call XML from Anthropic endpoint content

Apply _TOOL_XML_RE to content events in both streaming and
non-streaming tool paths, matching the OpenAI endpoint behavior.

* Emit custom tool_result SSE event in Anthropic stream

Adds a non-standard tool_result event between the tool_use block close
and the next text block, so clients can see server-side tool execution
results. Anthropic SDKs ignore unknown event types.

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

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

* Split /v1/messages into server-side and client-side tool paths

enable_tools=true runs the existing server-side agentic loop with
built-in tools (web_search/python/terminal). A bare tools=[...] field
now triggers a client-side pass-through: client-provided tools are
forwarded to llama-server and any tool_use output is returned to the
caller with stop_reason=tool_use for client execution.

This fixes Claude Code (and any Anthropic SDK client) which sends
tools=[...] expecting client-side execution but was previously routed
through execute_tool() and failing with 'Unknown tool'.

Adds AnthropicPassthroughEmitter to convert llama-server OpenAI SSE
chunks into Anthropic SSE events, plus unit tests covering text
blocks, tool_use blocks, mixed, stop reasons, and usage.

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

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

* Fix httpcore GeneratorExit in /v1/messages passthrough stream

Explicitly aclose aiter_lines() before the surrounding async with
blocks unwind, mirroring the prior fix in external_provider.py
(a41160d3) and cc757b78's RuntimeError suppression.

* Wire stop_sequences through /v1/messages; warn on tool_choice

Plumb payload.stop_sequences to all three code paths (server-side
tool loop, no-tool plain, client-side passthrough) so Anthropic SDK
clients setting stop_sequences get the behavior they expect. The
llama_cpp backend already accepted `stop` on both generate_chat_
completion and generate_chat_completion_with_tools; the Anthropic
handler simply wasn't passing it.

tool_choice remains declared on the request model for Anthropic SDK
compatibility (the SDK often sets it by default) but is not yet
honored. Log a structured warning on each request carrying a non-
null tool_choice so the silent drop is visible to operators.

* Wire min_p / repetition_penalty / presence_penalty through /v1/messages

Align the Anthropic endpoint's sampling surface with /v1/chat/completions.
Adds the three fields as x-unsloth extensions on AnthropicMessagesRequest
and threads them through all three code paths: server-side tool loop,
no-tool plain, and client-side passthrough.

The passthrough builder emits "repeat_penalty" (not "repetition_penalty")
because that is llama-server's field name; the backend methods already
apply the same rename internally.

* Fix block ordering and prev_text reset in non-streaming tool path

_anthropic_tool_non_streaming was building the response by appending
all tool_use blocks first, then a single concatenated text block at
the end — losing generation order and merging pre-tool and post-tool
text into one block. It also never reset prev_text between synthesis
turns, so the first N characters of each post-tool turn were dropped
(where N = length of the prior turn's final cumulative text).

Rewrite to build content_blocks incrementally in generation order,
matching the streaming emitter's behavior: deltas within a turn are
merged into the trailing text block, tool_use blocks interrupt the
text sequence, and prev_text is reset on tool_end so turn N+1 diffs
against an empty baseline.

Caught by gemini-code-assist[bot] review on #4981.

* Make test_studio_run.py e2e tests pytest-compatible

Add a hybrid session-scoped studio_server fixture in conftest.py that
feeds base_url / api_key into the existing e2e test functions. Three
invocation modes are now supported:

1. Script mode (unchanged) — python tests/test_studio_run.py
2. Pytest + external server — point at a running instance via
   UNSLOTH_E2E_BASE_URL / UNSLOTH_E2E_API_KEY env vars, no per-run
   GGUF load cost
3. Pytest + fixture-managed server — pytest drives _start_server /
   _kill_server itself via --unsloth-model / --unsloth-gguf-variant,
   CI-friendly

The existing _start_server / _kill_server helpers and main() stay
untouched so the script entry point keeps working exactly as before.
Test function signatures are unchanged — the (base_url, api_key)
parameters now resolve via the new fixtures when running under
pytest.

* Rename test_studio_run.py -> test_studio_api.py

The file is entirely about HTTP API endpoint testing (OpenAI-compatible
/v1/chat/completions, Anthropic-compatible /v1/messages, API key auth,
plus a CLI --help sanity check on the command that runs the API). None
of its tests cover training, export, chat-UI, or internal-Python-API
concerns.

The old name misleadingly suggested "tests for the unsloth studio run
CLI subcommand" — the new name reflects the actual scope.

Updates:
- git mv the file (rename tracked, history preserved)
- Rewrite opening docstring to state the API surface focus and call
  out what is explicitly out of scope
- Update all 4 Usage-block path references to the new filename
- LOG_FILE renamed to test_studio_api.log
- conftest.py fixture import rewritten from test_studio_run to
  test_studio_api, plus 7 docstring/comment references updated

No functional changes to test logic, signatures, or main().

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Fix httpcore asyncgen cleanup in /v1/messages and /v1/completions

The earlier fix in 985e92a9 was incomplete: it closed aiter_lines()
explicitly but still used `async with httpx.AsyncClient()` /
`async with client.stream()` inside the generator. When the generator
is orphaned (e.g. client disconnects mid-stream and Starlette drops
the StreamingResponse iterator without explicitly calling aclose()),
Python's asyncgen finalizer runs the cleanup in a DIFFERENT task than
the one that originally entered the httpx context managers. The
`async with` exits then trigger httpcore's HTTP11ConnectionByteStream
.aclose(), which enters anyio.CancelScope.__exit__ with a mismatched
task and raises RuntimeError("Attempted to exit cancel scope in a
different task"). That error escapes any user-owned try/except
because it happens during GC finalization.

Replace `async with` with manual client/response lifecycle in both
/v1/messages passthrough and /v1/completions proxy. Close the
response and client in a finally block wrapped in
`try: ... except Exception: pass`. This suppresses RuntimeError (and
other Exception subclasses) from the anyio cleanup noise while
letting GeneratorExit (a BaseException, not Exception) propagate
cleanly so the generator terminates as Python expects.

Traceback observed in user report:
  File ".../httpcore/_async/connection_pool.py", line 404, in __aiter__
      yield part
  RuntimeError: async generator ignored GeneratorExit
...
  File ".../anyio/_backends/_asyncio.py", line 455, in __exit__
      raise RuntimeError(
  RuntimeError: Attempted to exit cancel scope in a different task

* Expand unsloth studio run banner with SDK base URL and more curl examples

Add an explicit "OpenAI / Anthropic SDK base URL" line inside the info
box so SDK users don't accidentally copy the bare server URL (without
/v1) into their OpenAI/Anthropic SDK constructors and hit 404s.

Replace the single /v1/chat/completions curl example with three
labeled blocks: chat/completions, Anthropic /messages, and OpenAI
Responses. The Anthropic example includes max_tokens (Anthropic SDKs
require it even though Studio accepts None).

All examples derived from a computed sdk_base_url so the /v1 prefix
stays in sync if the public path ever changes.

* Hash API keys with HMAC-SHA256 + persistent server secret

Stores the HMAC secret in a new app_secrets singleton table. Fixes
CodeQL py/weak-sensitive-data-hashing alert on storage.py:74-76,
394-395. Refresh tokens stay on plain SHA-256 (unchanged _hash_token)
so existing user sessions survive upgrade — API keys are new on this
branch so there is no migration.

* Use PBKDF2 for API key hashing per CodeQL recommendation

HMAC-SHA256 was still flagged by py/weak-sensitive-data-hashing.
Switch to hashlib.pbkdf2_hmac, which is in CodeQL's recommended
allowlist (Argon2/scrypt/bcrypt/PBKDF2). Persistent server-side
salt stays in app_secrets for defense-in-depth. 100k iterations to
match auth/hashing.py's password hasher.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-04-13 21:08:11 +04:00
Wasim Yousef Said
629199e3a6
fix: remove old comments (#4292)
* fix: quotation marks

* diceware passphrase generation

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

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

---------

Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-14 16:50:13 +04:00
Roland Tannous
47654cb91c Final cleanup 2026-03-12 18:28:04 +00:00
Roland Tannous
a2baf80511 Update license headers 2026-03-12 17:23:10 +00:00
Shine1i
bbb4cd0f0b feat(studio): add auth-specific paths and integrate auth database location 2026-03-11 20:19:52 +00:00
Roland Tannous
daa50d0756 Revert "Merge pull request #347 from unslothai/feature/studio-storage-roots"
This reverts commit 6b43e33ff1, reversing
changes made to 9edadaf21f.
2026-03-10 01:52:47 +00:00
Shine1i
109db14817 feat(studio): add auth-specific paths and integrate auth database location 2026-03-09 23:48:31 +00:00
Roland Tannous
d882678fe4 Add AGPL-3.0 SPDX headers to all source files 2026-03-09 20:17:45 +00:00
Leo Borcherding
a3daae1c40 fix: replace datetime.UTC with timezone.utc for Python 3.9+ compatibility
- Replace datetime.UTC with datetime.timezone.utc in authentication.py and storage.py
- Fixes ImportError on Python versions < 3.11
- timezone.utc works on Python 3.9+

Resolves #237
2026-02-24 14:37:00 -06:00
Roland Tannous
5602f7ccb4 fix: rollback auth.db user row if token generation fails during setup 2026-02-13 10:11:07 +00:00
Roland Tannous
fd3e0e5f09 chore: untrack auth.db (already in .gitignore) 2026-02-11 12:21:54 +00:00
Roland Tannous
01fcb4f713 authentication refactor - added setup token and token refresh mechanism 2026-02-11 12:09:47 +00:00
sshah229
0082627801 fixed the errors- renamed jwt to authentication, used raw jwt, and removed search route 2026-02-07 03:14:30 -07:00
sshah229
50ff5626f1 refactored the code for username/password and added pydantic models and routes for the same 2026-02-06 03:15:30 -07:00
sshah229
009e93f079 Refactored the training and model routes and added the jwt authentication 2026-02-06 03:15:30 -07:00
Roland Tannous
544d6944d1 root studio folder 2026-02-02 09:13:49 +00:00