Commit graph

7,057 commits

Author SHA1 Message Date
Daniel Han
d2e6192eb2 Honor the configured SDXL LoRA batch size on datasets smaller than the batch
The SDXL trainer drew min(train_batch_size, len(pairs)) indices, so a dataset with
fewer images than the batch trained at a smaller effective batch than configured
while the scheduler and samples-per-second still assumed the full batch. The shared
PermutationBatchSampler already refills across permutation cycles to return exactly k
indices, and the DiT trainer calls it with the full batch size, so drop the clamp and
pass train_batch_size through for parity and to honor the configured batch.
2026-07-08 01:50:23 +00:00
Daniel Han
dc65a63f16 Scope local diffusion-family detection to the leaf name, not the on-disk path
A scanned On-Device model's id is its full filesystem path, and the family-token
matcher treats any path segment as a hint, so a folder under a family-named
parent (e.g. /models/qwen-image/misc) tagged an unrelated single-file as
text-to-image. That surfaced it in the Images picker, and a load would evict the
GPU owner before from_single_file failed on the unrelated weights. Match on
Path(model.id).name so only the leaf name counts, for both the image and video
task-tag loops (the video loop is gated on the same detection).
2026-07-08 00:49:31 +00:00
Daniel Han
429ffa7392 Refuse deleting an in-use native companion repo while a GGUF is loaded
The native sd.cpp one-shot engine re-reads its companion VAE and text-encoder
files from the HF cache on every generation, but the delete-cached guard only
compared the loaded main diffusion repo_id. Deleting a companion repo such as
comfyanonymous/flux_text_encoders while a FLUX GGUF was loaded therefore
succeeded and broke the next generation. Add a loaded_repo_ids() accessor to the
native backend (mirroring loading_repo_ids(), reconstructed from the committed
family's VAE + text-encoder repos) and have the guard refuse those companions
too. Server mode and the diffusers engine hold companions in-process/VRAM, so
they are unaffected.
2026-07-08 00:08:55 +00:00
Daniel Han
f134a19b87 Validate a local video checkpoint's own suffix; unload the old engine before publishing the new one
Video: a local FILE picked for a gguf/single_file load is handed straight to the
loader (the resolver returns the file itself, ignoring gguf_filename), so its own
suffix must match the kind. A .gguf picked as single_file (or a .safetensors picked
as gguf) slipped past the gguf_filename checks, evicting the resident GPU owner
before failing in from_single_file. Reject the mismatch in validate, before the handoff.

Engine router: publish the newly selected engine only after the old one finishes
unloading. The arbiter's diffusion evictor unloads the active engine, so flipping the
active name to the new (empty) engine first let a concurrent chat/video acquire evict
that empty engine and take the GPU while the old model was still freeing VRAM.
2026-07-07 22:55:36 +00:00
Daniel Han
3050196f6b Make dense-quant candidate resolution tests independent of runner disk
The two candidate-resolution tests exercised the real cache-disk gate, so a
small CI disk (< transient + 10 GiB free) dropped the candidate and the tests
failed non-deterministically across runners. Neutralize the disk probe in the
shared selector helper; the two disk-gate tests re-patch it to cover the gate.
2026-07-07 21:47:52 +00:00
Daniel Han
cf0f5d5504 Serialize video-load GPU placement with eviction; skip LoRA adapters as base picks
Two evict/OOM fixes on the diffusion load paths:

- The video load moved a pipeline onto the GPU (apply_memory_plan) and
  committed it while holding no lock, so an unload / GPU-arbiter eviction --
  which bumps the load token and then barriers on _generate_lock before
  freeing -- could hand VIDEO to chat/images and let the new owner allocate
  concurrently with the in-flight placement, OOMing. Hold _generate_lock
  across placement + the locked commit, mirroring the image backend, so an
  evicting owner waits until this worker's placement is torn down or
  committed. Lock order stays _generate_lock -> _lock (unload takes _lock
  then releases it before the barrier), so there is no deadlock.

- resolve_local_single_file reinterpreted an On-Device folder as a base
  single_file load whenever it held exactly one .safetensors, so a PEFT LoRA
  adapter folder (adapter_config.json + adapter_model.safetensors) with a
  family-token name was picked as a base checkpoint, evicting the resident
  model before from_single_file failed on the adapter weights. Skip adapter
  folders (adapter_config.json) and the adapter_model basename so the pick
  stays a pipeline load and 400s in validation, before the GPU handoff.

Adds regression tests for both.
2026-07-07 20:57:44 +00:00
Daniel Han
26f24f1b7d Reject a local file picked as a video pipeline before the GPU handoff
The video load-request preflight gated its local-pipeline check on
root.is_dir(), so a bare local file (e.g. /models/ltx-2.safetensors) sent
with model_kind=pipeline skipped it, passed validation, and the route then
evicted the resident GPU model before from_pretrained failed on the
non-directory path. Gate on root.exists() instead (mirroring the image
loader's diffusion.validate_load_request), so a local file is rejected up
front. Add a regression test.
2026-07-07 19:34:39 +00:00
Daniel Han
a091a862df Add krea-2 to bf16-only preflight; match dataset stems case-insensitively
Two evict/corruption fixes surfaced by review of the diffusion training path:

- krea-2 sets force_bf16 in the DiT trainer spec, but the route-level
  _FORCE_BF16_FAMILIES preflight listed only qwen-image and z-image, so a
  krea-2 start with mixed_precision=fp16 passed the route check, reserved
  training and evicted resident GPU models, and only the child trainer then
  raised. Add krea-2 to the set and a drift-guard test asserting it equals
  the trainer specs whose force_bf16 is set.

- The dataset-upload same-stem duplicate check compared stems
  case-sensitively, so on case-insensitive filesystems (Windows / default
  macOS) sample.png and Sample.jpg both passed even though their caption
  sidecars sample.txt / Sample.txt resolve to the same file, silently
  sharing and corrupting one caption. Compare stems and the same-name guard
  with casefold at both the on-disk and in-batch sites.
2026-07-07 18:07:48 +00:00
Daniel Han
c5bc90c3cf Match cached repo ids on a / boundary before blocking deletion
The /delete-cached guard refused deleting a cached repo whose id is a bare
prefix of the loaded one: with Qwen/Qwen-Image-2512 loaded, deleting the
separate cached Qwen/Qwen-Image returned 400 "Unload the model before
deleting" because loaded_id.startswith(repo_id) matched across the repo
boundary. Both are real, independently cached catalog artifacts, so a
supported delete was wrongly blocked.

Add a /-boundary aware _loaded_id_matches_repo helper (mirroring
hub/services/models/deletion.py) and use it at all six guard sites (chat,
Images, Video, and their in-flight loading-id loops), so only the loaded
repo itself or a file within it blocks deletion. Add a regression test
covering the sibling-prefix pair.
2026-07-07 16:34:29 +00:00
Daniel Han
ffa8a56391 Merge image-generation bug fixes (#6872)
Fold PR #6872's image-generation fixes into the branch, deduped against the
round-12 dataset-upload and gallery integrity work already on image-generation.

Fixes carried forward from #6872:
- fp8 single-file transformer memory estimate: an fp8 checkpoint loads with no
  quantization_config and diffusers upcasts it to bf16 (~2x resident), so budget
  it accordingly in _plan_memory and estimate_safetensors_dense_mib.
- dense-quant OOM-evict preflight: when the GGUF fits resident but the dense bf16
  transformer this path materializes does not, skip the fast path up front rather
  than evict the current pipeline and OOM in finalization. Combined with the
  existing offload->resident candidate re-plan so both the family-table estimate
  and the on-disk shard measurement gate engagement (unified on the
  transformer_resident_override_mib plan override).
- ControlNet: evict the previous module and its from_pipe wrapper before loading a
  new one so swapping ControlNets within a base-model load cannot accumulate to OOM.
- ControlNet union_control_mode: raise on an unknown control type instead of
  silently defaulting to canny.
- edit-family mask rejection: raise instead of silently dropping a mask on an
  image-editing model that has no inpaint pipeline.
- companion cache: walk the snapshot dir and exclude transformer/ so the
  dense-quant prefetch's cached shards do not inflate the companion total and
  wrongly force offload.
- training: drop piecewise_constant from the LR scheduler enum and force bf16 for
  fp16-incompatible families.
- dataset upload: batch-atomic staging with the same-stem duplicate guard.
- images page: guard negative-prompt restore on guidance>0, clear stale ControlNet
  selection on restore, and revert an optimistic quant label when a pipeline load
  never starts.
- uninstall (sh + ps1): keep the owner-marker guard on sd.cpp removal.

Conflicts resolved in favour of image-generation's evolved memory system,
loadSpecFor catalog, and stop-and-save (lora_path) run detection; #6872's fp8 and
dense-preflight fixes carried forward on top. All affected backend tests pass
(test_diffusion_backend, test_diffusion_training, test_diffusion_lora_trainer,
test_video_gallery, test_diffusion_controlnet).
2026-07-07 16:14:06 +00:00
Daniel Han
598f222eca Fix video-gallery delete test to patch Path.unlink for Python 3.10
test_delete_keeps_sidecar_listable_when_mp4_unlink_fails patched os.unlink, but delete()
calls path.unlink(). On Python 3.10 Path.unlink dispatches through a cached _accessor bound
to os.unlink at import, so patching os.unlink had no effect there and the simulated mp4
unlink wrongly succeeded (the assert delete() is False failed on 3.10 while passing on
3.11+, where Path.unlink looks up os.unlink dynamically). Patch Path.unlink itself, which
delete() invokes directly on every supported Python version. Test-only; no behavior change.
2026-07-07 15:06:24 +00:00
Daniel Han
41bdc116be Reject unknown union ControlNet control types instead of defaulting to canny
union_control_mode fell back to control_mode=0 (the canny head) for ANY unmapped control
type, so a typo'd or unsupported value like 'detph' silently conditioned the map as canny
instead of failing. preprocess_control passes non-canny maps through unchanged, so that map
would be interpreted under the wrong mode with no error. Now only 'passthrough' (or an empty
type) keeps the deliberate mode-0 default; any other unknown type raises ValueError, which
the generate route maps to a 400. Known modes are unchanged.
2026-07-07 14:55:37 +00:00
Daniel Han
46a0a21d53 Drop piecewise_constant from the diffusion training API scheduler enum
Follow-up to removing piecewise_constant from the trainable scheduler allow-list: the
DiffusionTrainingStartRequest.lr_scheduler Literal still advertised it, so a client that
picked it straight from the schema passed request validation and then hit the 400 from
normalized(). Remove it from the enum too so the API only offers schedulers the trainers
can actually run, and add a test asserting the enum never advertises a scheduler outside
the validation allow-list.
2026-07-07 14:23:14 +00:00
Daniel Han
0d0a6d2b96 Drop piecewise_constant from the trainable lr_scheduler allow-list
piecewise_constant is the only diffusers scheduler that needs a step_rules string, and
neither diffusion trainer passes one (get_scheduler is called with only warmup/training
steps, and there is no config field for it). Accepting it let /diffusion/start pass
normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in
the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the
exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear
400. The remaining six schedulers all run with only warmup/training steps.
2026-07-07 13:39:34 +00:00
Daniel Han
f33a6c1b61 Harden diffusion dataset integrity: atomic example import and same-stem upload guard
Promote the example-import staging dir into the dataset folder in one atomic
same-filesystem rename instead of a per-file move loop. A hard process death between
two moves left the folder with SOME images, and the image_count>0 idempotency check
then accepted that truncated dataset as complete on retry. The folder is created
empty on this path, so a single rename is atomic; a folder holding unrelated
non-image files falls back to a per-file move rather than abort.

Reject a second uploaded image whose stem matches an existing image but differs by
extension (sample.png vs sample.jpg): both map to one <stem>.txt caption sidecar
(the kohya/diffusers convention), so keeping both would silently corrupt captions
during training. Exact-name overwrites and .txt caption sidecars stay allowed.
2026-07-07 11:57:37 +00:00
Daniel Han
fda725241d Delete gallery videos MP4-first so a locked MP4 can't orphan the record
list_videos globs *.mp4 but requires a readable sidecar, so a video whose sidecar
is gone is skipped. delete()/clear() dropped the sidecar first, so if the mp4
unlink then failed (a Windows lock from a concurrent stream/transcode), the
still-present mp4 vanished from the gallery with no way to retry the delete. Unlink
the mp4 first and only then best-effort the sidecar: the worst case is now an
orphaned sidecar, which list_videos already ignores.
2026-07-07 11:57:37 +00:00
Daniel Han
123b3ea4d6 Make the diffusion-training reservation a compare-and-set
Two overlapping /diffusion/start requests can interleave between the is_active()
check and the reservation, so reserve() itself must reject a second reservation
atomically. Otherwise both callers reserve, both free the GPU's resident chat or
image model, and the loser only 409s after the eviction -- the evict-then-fail the
reservation exists to prevent. reserve() now raises under the lock if a start is
already reserved or a job is already running.
2026-07-07 11:57:37 +00:00
Daniel Han
39ee329b26 Reserve the diffusion-training slot before freeing GPU residents
start_diffusion_training freed resident GPU models and only then called
service.start(config), which is where is_active() first flips true. During that
free-then-spawn window a concurrent /images/load or /video/load saw training as
inactive, passed its training guard, acquired the GPU, and began a background load,
so the trainer and that pipeline both allocated VRAM. Add reserve()/unreserve() to
the training service (is_active() also reports the reservation) and reserve BEFORE
the free, in a try/finally so a failed start rolls the reservation back. An
overlapping load's guard now refuses during the window. Regression tests: the route
reserves before the free (and the free sees an active service), and the service
reservation marks active then rolls back.
2026-07-07 10:58:13 +00:00
Daniel Han
318e45ee97 Carry the image loader's video preflight + speed=off guards into the video backend
Three gates the image backend has were missing on the video path:

- kind/extension mismatch: a model_kind 'single_file' with a .gguf name (or 'gguf'
  with a non-.gguf) passed the video preflight, so the route acquired VIDEO and
  evicted the resident GPU owner before the wrong single-file loader failed in the
  background. Reject the mismatch up front, mirroring the image loader.
- Windows-shaped missing local pick: the missing-path check only matched POSIX
  prefixes (/ ~ ./ ../), so a missing Windows path (C:\ / C:/ or any backslash path)
  was treated as a Hub repo and only failed after the GPU handoff. Use the image
  loader's is_absolute()/backslash path-shaped check.
- speed=off auto-quant: a full-pipeline load with Speed=off but Precision=auto still
  promoted the unset precision to auto-quant, so on a dense-capable GPU it engaged
  torchao quantization and forced the speed back to default -- silently breaking the
  user's bit-exact request. Suppress the auto promotion when speed_mode is off, as
  the image loader does.

Regression tests for each.
2026-07-07 10:58:04 +00:00
Daniel Han
6519fbfad5 Keep an unowned default-mode sd.cpp checkout on uninstall
The custom-root sd.cpp removal requires the .unsloth-studio-owned marker (written by
install_sd_cpp_prebuilt) before deleting, so a user's own stable-diffusion.cpp checkout
beside a custom root is kept. The default-mode removal of ~/.unsloth/stable-diffusion.cpp
was unconditional, so a user who keeps their own checkout at that path (or points
UNSLOTH_SD_CPP_PATH there), or a pre-marker Studio build, would have it deleted on
uninstall. Guard the default-mode removal on the same owner marker in both uninstall.sh
and uninstall.ps1. Extends the sd.cpp uninstall shell test with owned (removed) and
unowned (kept) default-mode cases.
2026-07-07 10:22:40 +00:00
Daniel Han
5eeea1407b Normalize dataset upload filenames so the UI can manage them
The dataset upload took Path(filename).name, which on POSIX does not split on a
backslash, so a Windows client sending a backslash path in the multipart filename
stored the name verbatim. The caption/thumbnail/delete endpoints then run it
through _safe_dataset_image_path, which rejects backslashes and '..', so the
labeling grid could list an image it could never preview, caption, or delete (an
orphan). Fold backslashes to forward slashes before taking the basename so the
true name is stored, and reject a basename that still contains '..' at upload
rather than persisting an unmanageable entry. Regression test covers a Windows
backslash path (stored and served under the clean basename) and a '..' rejection.
2026-07-07 09:51:51 +00:00
Daniel Han
7c9fb69d42 Scope the diffusion dataset body-cap passthrough to the exact upload route
The multipart upload passthrough was a bare /api/train/diffusion/dataset prefix,
so its JSON sub-routes (PUT .../{name}/caption/{filename}, POST .../import-example)
also bypassed the default JSON body cap and inherited the far larger upload limit.
A large caption/import body would then be buffered and parsed up to the upload cap
before the route-level length checks ran. Match the upload route by EXACT path
instead: the JSON sub-routes fall through to the normal small-JSON cap. Adds an
upload_passthrough_exact_paths param to MaxBodyMiddleware, with regression tests
that the sub-routes keep the default cap and a large sub-route body is 413'd.
2026-07-07 09:51:44 +00:00
Daniel Han
eee5a53658 Offload the gated-base HEAD preflight to a worker thread
start_diffusion_training is async but called _preflight_gated_base inline; it does
a blocking urllib urlopen HEAD to Hugging Face (up to a 5s timeout) to detect a
gated/unauthorized base repo. On a slow or unreachable network that stalled the
FastAPI event loop, freezing every concurrent status/progress/cancel request until
it returned or timed out. Wrap it in asyncio.to_thread, matching the dataset
preflight and GPU cleanup just below it. Regression test asserts it runs off the
coroutine thread.
2026-07-07 08:55:23 +00:00
Daniel Han
37185e651d Reject doomed local picks before the GPU handoff; load bare local safetensors
Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.

- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
  local path, but the base loads via from_pretrained (needs model_index.json). A
  local dir that is not a diffusers pipeline passed the trust gate, evicted the
  resident model, then failed. Add a shared _assert_local_base_is_pipeline check
  and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
  corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
  crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
  (cheap PIL header check) that the start route enables; the trainers leave it off
  since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
  .safetensors directory (no model_index.json) as a text-to-image model, but the
  picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
  such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
  so the advertised model is actually loadable.

Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.
2026-07-07 08:06:54 +00:00
Daniel Han
8476f02763 Offload the diffusion-training GPU cleanup to a worker thread
start_diffusion_training is an async route, but it called the blocking
_free_gpu_for_diffusion_training() inline. That teardown waits on generation
locks and joins the export subprocess, so it can block for seconds and freeze the
FastAPI event loop, stalling every concurrent status/progress/cancel request
until it finishes. Wrap it in asyncio.to_thread, mirroring the dataset-preflight
call just above it and the inference routes' load/unload offloading. Add a
regression test asserting the cleanup runs off the coroutine thread.
2026-07-07 06:58:07 +00:00
Daniel Han
3d482dcb88 Reset the image FBCache with the real diffusers CacheMixin hook
_reset_step_cache looked up reset_stateful_hooks on the transformer, but on a
diffusers CacheMixin transformer (Flux, QwenImage) that method lives only on the
HookRegistry; the transformer-level entry point is _reset_stateful_cache. So with
FBCache engaged on an image model the reset was a silent no-op, and the next
generation reused the previous request's first-block residual: a tensor-shape
mismatch (crash) when the resolution or batch changed, or stale cached output
otherwise. Prefer _reset_stateful_cache and fall back to reset_stateful_hooks,
matching the video backend. Update the tests to the real hook name.
2026-07-07 06:57:53 +00:00
Daniel Han
d165789462 Keep GPU ownership across an in-flight reload on unload
The images and video unload routes dropped their arbiter claim whenever the
backend was not committed-loaded, but a concurrent /load re-acquires the owner
and starts a background load that is not is_loaded for its whole download and
finalize window. Releasing during that window cleared the newer load's claim, so
a later chat/image load saw no owner, skipped eviction, and could allocate a
second heavy model on the GPU. Gate the release on loading_repo_ids() too (both
diffusion engines and the video backend expose it), not just the committed
state, so an overlapping load keeps ownership. Add regression tests for the
in-flight case on both routes.
2026-07-07 06:57:47 +00:00
Daniel Han
f98a06b72b Stub is_bf16_supported in the train-precision-modes test helper
The three train_precision_modes capability tests patched is_available and
get_device_capability but not is_bf16_supported, so on the CPU-only CI runner
(where the real probe is False) the dense modes collapsed to nf4 and the fp8 /
mxfp8 assertions failed; they only passed on a bf16 GPU dev box. An Ada or
Blackwell GPU is by definition bf16-capable, so the helper must stub it True to
exercise the capability gate the tests target.
2026-07-07 06:57:41 +00:00
Daniel Han
2e752e6900 Seed image sliders from a resident model's recipe on discovery
When the images page mounts with a diffusion model already loaded (a prior
session or another route left it resident), refreshStatus only set status; the
steps/guidance sliders kept the unrecognised-model fallback (few-step, no CFG),
so a resident full model such as flux.1-dev generated at 9 steps guidance 0 and
produced garbage until the user re-picked it. Add a one-shot effect that seeds
the sliders from defaultsFor(status.repo_id) when the page discovers a resident
model it did not load itself, and wires Reapply to it for non-GGUF pipelines.
Guarded by lastLoad and a per-repo ref so a manual slider edit is never
clobbered.
2026-07-07 06:10:36 +00:00
Daniel Han
22dad1df73 Gate mxfp8 DiT training precision before evicting resident GPU models
The start route's precision preflight folded bf16/int8/fp8 into the CUDA
requirement but omitted mxfp8, so an mxfp8 request on a GPU-less host (or an
older CUDA GPU without Blackwell) passed the preflight, evicted resident image
and chat models, then raised only in the spawned trainer child. Mirror
_resolve_base_precision: require CUDA for mxfp8 and re-check the Blackwell
(sm100+) capability up front, so a doomed run is rejected before teardown.
2026-07-07 06:10:28 +00:00
Daniel Han
eb80e66709 Studio: apply ruff-format kwarg spacing to the diffusion review changes 2026-07-07 05:53:48 +00:00
Daniel Han
c399aabd5d Merge remote-tracking branch 'origin/main' into fold-integration
# Conflicts:
#	scripts/scan_packages_baseline.json
2026-07-07 05:52:16 +00:00
Daniel Han
1c110dc867 Studio: gate the base_model card tag on image loads, fix an unload/reload arbiter race, and don't crash the offload fallback on a partial dual-DiT hook set
- The companion base for a GGUF/single-file image load is resolved from the GGUF
  repo's base_model card tag when no base_repo is passed, and that value loads via
  from_pretrained. The explicit base_repo is already trust-gated, but the card tag is
  attacker-controlled metadata on any remote repo, so it now clears the same
  unsloth/allowlist/local trust bar; an untrusted tag is dropped in favour of the
  curated family default and never reaches from_pretrained. This closes a pickle
  deserialization vector on the normal GGUF load path (a user loading an attacker's
  GGUF repo whose card points base_model at a malicious pipeline), matching the
  trust discipline the ControlNet path already applies via evaluate_file_security.
  The allowlist already contains every legitimate variant base, so variant
  resolution for the supported unsloth GGUFs is unchanged.
- The images/unload route ran the slow VRAM-freeing unload on a thread and then
  released the DIFFUSION arbiter owner unconditionally. release() is owner-guarded
  and identity-less, so a concurrent /images/load that re-acquired DIFFUSION while
  the unload ran would have its ownership cleared by the trailing release, and a
  later chat load would then see no owner, skip eviction, and OOM against the newly
  resident pipeline. The route now releases only when nothing is resident again.
- _apply_group_offload placed the resident companions before attaching the
  transformer's group-offload hooks so a companion OOM returns with no hooks and the
  whole-module fallback stays valid, but the streamed loop itself installs hooks on
  each DiT in turn. On a dual-DiT pipeline where the second tower failed after the
  first got its hooks, it returned False with hooks already installed, and the
  caller's enable_model_cpu_offload fallback then crashed (diffusers rejects it on a
  partially group-offloaded pipe). It now propagates the real failure once any hook
  is installed, so the load fails with its actual cause instead of a misleading crash.

Adds regression tests: the untrusted card tag dropped to the family default (trusted
tag still honoured, explicit base still wins), unload keeping ownership when a model
is still resident (and releasing when not), and the partial dual-DiT hook set
propagating rather than falling through to a crashing whole-module offload.
2026-07-07 05:34:30 +00:00
Daniel Han
5608081c35
Studio: apply presence_penalty on the safetensors and MLX inference paths (#6923)
* Studio: apply presence_penalty on the safetensors and MLX inference paths

The safetensors and MLX generate paths resolved the inference config and
then dropped presence_penalty before generation, so the same model applied
the configured value under GGUF and 0 under safetensors/MLX. Thread the
already-resolved presence_penalty through the orchestrator command, worker
gen_kwargs, and the safetensors/MLX generate calls, and apply it with a
small logits processor (subtract once per distinct completion token,
prompt excluded, presence not frequency, zero is a no-op, negatives raise).

Backwards compatible: presence_penalty defaults to 0.0 (byte-identical
output when unset) and the GGUF path is unchanged. Also forward min_p on
the legacy /generate/stream route and add the missing min_p field to
GenerateRequest.

* Studio: bound presence_penalty generated ids to valid vocab range on both paths

The presence-penalty logits processors index by generated token ids. The
torch path filtered only the upper bound (seen < vocab_size), so a negative
id would silently wrap to the wrong row; the MLX path had no bound at all,
and MLX out-of-bounds indexing is documented undefined behavior (crash or
memory corruption on Apple Silicon), unlike torch's harmless negative wrap.

Bound generated ids to [0, vocab) consistently on both paths:
- torch: seen[(seen >= 0) & (seen < vocab_size)] (zero-regression safety net;
  real completion tokens are always in range).
- MLX: route out-of-range/negative ids to a discarded scratch slot via
  mx.where and a (vocab + 1)-wide scatter-assign mask, then subtract. MLX has
  no boolean-mask filtering (data-dependent output shape), so this keeps a
  fixed shape, stays on-device, and preserves once-per-distinct-token
  semantics without any torch/numpy dependency.

Add torch tests for out-of-range and negative ids (only in-range distinct
ids penalized, stray ids ignored, no wrong-index wrap) and a bound-documenting
MLX test that runs on the arm64 macOS CI.

* [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-07-06 22:24:47 -07:00
Daniel Han
9674e882c2
Studio: serialize the compare-mode dispatcher lifecycle to fix a start race (#6922)
* Studio: serialize the compare-mode dispatcher lifecycle to fix a start race

_generate_dispatched (compare mode) bypasses _gen_lock so two concurrent
compare requests can both reach _start_dispatcher. The check-then-spawn there
had no lock, so both could observe no live dispatcher and each spawn one. The
extra dispatcher is orphaned (self._dispatcher_thread tracks only the last) and
during a later unload it can consume the 'unloaded' reply off _resp_queue before
unload_model's _wait_response, hanging the unload on its timeout.

Add _dispatcher_lifecycle_lock and take it around the whole body of both
_start_dispatcher and _stop_dispatcher, so start/stop cannot interleave and the
second concurrent starter sees the dispatcher alive and returns. _start_dispatcher
now returns whether it actually spawned the thread, and _generate_dispatched
derives dispatcher_preexisting from that atomic result instead of a separate
unlocked is_alive() read.

No call site holds _mailbox_lock when calling start/stop, so joining the
dispatcher (which takes _mailbox_lock) under the new lock cannot deadlock; the
lock order is always _gen_lock then _dispatcher_lifecycle_lock and is never
inverted.

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

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

* Studio: refuse dispatcher start queued behind an unload's stop

A compare request could pass the early _unload_pending check, then block in _start_dispatcher on _dispatcher_lifecycle_lock behind an unload's _stop_dispatcher. When the unload released the lock the start spawned a fresh dispatcher, which became the resp_queue reader and consumed the worker's unroutable 'unloaded' reply before unload_model's _wait_response saw it, hanging the unload for 300s.

Gate _start_dispatcher on _unload_pending under the lifecycle lock, and set _unload_pending under the same lock ahead of the stop, so any start queued behind the stop observes the unload and refuses. Ordering stays _gen_lock -> _dispatcher_lifecycle_lock. Adds a regression test forcing the queued-behind-stop interleaving.

* [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-07-06 22:09:41 -07:00
Daniel Han
f88d5c9f6b Studio: gate untrusted companion base_repo on image loads, track image/video generation for the training guard, and body-cap the whole /v1 surface
- Image load now trust-gates a client-supplied base_repo. validate_load_request
  rejects a base_repo that is not an unsloth/* repo, an allowlisted official base, or a
  local path, mirroring the repo_id gate and the video loader. The route passes
  base_repo into that pre-eviction validation, so an authenticated client can no longer
  keep model_path on a trusted GGUF while pointing base_repo at an arbitrary remote repo
  that the server would download and deserialize (a from_pretrained pickle/config path),
  and no resident model is evicted for the rejected load.
- The keepwarm middleware now tracks the image and video generation routes
  (/images/generate, /images/generations, /video/generate), so
  other_inference_request_count() sees an in-flight generation and an API-key training
  start is refused (409) before its unload would cancel that generation. endswith keeps
  the GET *-progress and */cancel variants untracked.
- The OpenAI-compatible /v1 surface is now blanket body-capped like /api/inference,
  instead of only /v1/chat/completions and /v1/completions. Every /v1 POST route
  (images/generations, audio, embeddings, responses, messages, ...) buffers a JSON body
  and none is a multipart-upload passthrough, so an unbounded ImageGenerationRequest
  prompt on /v1/images/generations can no longer be buffered outside the request limit.

Adds regression tests: the base_repo trust gate at both the backend (untrusted remote
raises, local passes) and the route (untrusted base_repo returns 400 with no load), the
keepwarm tracking of the image/video generation paths (and not the progress/cancel
variants), and the /v1 surface being body-protected.
2026-07-07 05:04:49 +00:00
Daniel Han
3506371677
Studio: keep the nudge wiring test collectable without the unsloth stack (#6924)
test_nudge_tool_calls_wiring.py imported InferenceBackend from
core.inference.inference, which pulls in unsloth (and thus unsloth_zoo)
at module scope. The dependency-light backend CI matrix job does not
install unsloth_zoo, so the import raised at collection time and aborted
the whole job (831 tests never ran). Guard that one import and fold the
safetensors InferenceBackend checks in only when the unsloth stack is
importable; the orchestrator/llama_cpp/safetensors_agentic wiring is
still asserted unconditionally, and local/full-stack runs keep the
InferenceBackend coverage.
2026-07-06 21:57:56 -07:00
Daniel Han
b6ab866c46 Studio: defer chat GPU handoff, unload video before training, stage example imports, name-aware Wan GGUF video tag, and diffusion upload body passthrough
- Chat load defers the CHAT arbiter handoff until after identifier / gpu_ids /
  training-memory validation, so a doomed chat load (bad id, unsupported gpu_ids on
  GGUF, or a training 409) no longer evicts a resident Images/Video pipeline and then
  errors. The already-loaded fast paths re-assert CHAT ownership themselves. Mirrors
  the image and video loaders, which validate before acquire_for.
- Both training-start GPU cleanups now unload a resident Video pipeline and release the
  VIDEO arbiter owner, not just Images/DIFFUSION, so starting LLM or diffusion training
  after a video generation session no longer competes with the still-resident video
  model and OOMs the run.
- Example dataset import materializes into a private staging dir and promotes into the
  dataset folder only after the whole import succeeds. A materialize that fails partway
  no longer leaves a partial dataset that a retry would treat as complete (imported=0),
  stranding the user with a truncated dataset. Hidden dirs are skipped by the dataset
  scan so the staging dir never surfaces as a dataset.
- Cached/local Wan GGUFs are now classified name-aware: arch "wan" alone is ambiguous
  between the loadable single-DiT TI2V-5B and the dual-expert A14B MoE the loader
  refuses, so _arch_to_task falls back to the repo/file name (as the loader's own
  detect_video_family does) and tags only a non-MoE match text-to-video, surfacing
  loadable Wan GGUFs in the Video picker without surfacing unloadable A14B files.
- The diffusion dataset upload route is added to the MaxBodyMiddleware upload
  passthrough, so its own get_upload_limit_bytes() cap (plus multipart overhead, and a
  raised max_upload_size_mb) applies instead of the default body limit rejecting
  near-limit batches with 413 before the handler runs.

Adds regression tests for each: the chat handoff not evicting on a doomed load, the
video unload on diffusion-training start, the atomic import leaving no partial dataset,
the name-aware Wan GGUF classification (TI2V-5B video, A14B unsupported, bare arch
unsupported), and the diffusion upload passthrough cap.
2026-07-07 04:08:41 +00:00
Daniel Han
d36aeb54f3 Studio: classify local video pipelines, harden video preflight, and gate custom-root sd.cpp removal on ownership
- _local_model_task now tags a local diffusers pipeline that resolves to a video
  family (LTX / Wan / Hunyuan) as text-to-video, mirroring the cached-repo
  _cached_repo_task, so supported local video pipelines surface in the Video
  On-Device picker instead of being routed to the Images picker where the image
  loader rejects them. Gated on _local_is_diffusers so only a real loadable
  pipeline dir reaches the video check.
- Video validate_load_request now rejects a local pipeline pick whose directory has
  no model_index.json before the GPU handoff, mirroring the image loader, so a bad
  local pipeline can no longer evict the resident model and only then fail deep in
  from_pretrained.
- The custom/env-mode uninstall now removes a sibling stable-diffusion.cpp only when
  it carries the Studio owner marker. install_sd_cpp_prebuilt writes the canonical
  .unsloth-studio-owned marker on install; uninstall.sh and uninstall.ps1 keep any
  unowned checkout (a user's own git clone of stable-diffusion.cpp beside a custom
  Studio root is no longer deleted). A pre-marker Studio build is left behind rather
  than a user file removed.

Adds regression tests: local video pipeline tagged text-to-video (and a video-named
non-pipeline dir stays untagged so it can never trigger a doomed pipeline load), the
video local-pipeline preflight rejection, the install ownership marker, and the
uninstall keeping an unowned sibling while removing an owned one.
2026-07-07 03:15:07 +00:00
Daniel Han
46ab683065
Studio: client-tool passthrough healing for safetensors and MLX (#6870)
* Studio: client-tool passthrough healing for safetensors and MLX

PR 6801 made response-side tool-call healing default-on for the client-tool
passthrough, but only on the GGUF path: the passthrough branch in
/v1/chat/completions is gated on using_gguf, and the safetensors section never
reads payload.tools, so a client-tools request against a safetensors or MLX
model silently dropped the tool schemas and returned prose with no tool_calls.

Add the missing leg. When a non-GGUF model is loaded, the request declares
client tools (or carries tool-role history), server-side tools are off, and the
template supports tools, the route now:
- renders the tools into the chat template for a single turn via the existing
  backend.generate_chat_response(..., tools=...) seam (worker templating
  already accepts role=tool and assistant.tool_calls messages, normalized with
  _openai_messages_for_passthrough);
- non-streaming: promotes text-form calls with heal_openai_message, honors the
  opt-in nudge single retry (nudge_should_retry / nudge_messages), caps healed
  calls when parallel_tool_calls=false (covers the nudge retry too), and sets
  finish_reason=tool_calls with content null on a pure tool-call turn;
- streaming: derives deltas from the worker's cumulative snapshots and feeds
  StreamToolCallHealer, emitting healed tool-call deltas and the correct
  finish chunk, guarded against repeated or shrinking snapshots.

heal_gate semantics are identical to the GGUF passthrough: default on,
auto_heal_tool_calls=false or UNSLOTH_DISABLE_TOOL_CALL_HEALING=1 relays
verbatim, tool_choice narrows promotion, undeclared names stay text. MLX rides
the same orchestrator seam, so both local backends gain the behavior.

CompletionMessage.content becomes Optional so a promoted pure tool-call turn
matches the OpenAI contract (content null when only tool_calls return).

Adds tests/test_sf_client_tools_passthrough.py (22 cases: healing, gating,
opt-outs, streaming deltas, tool-role history, dict-arguments history, forced
tool_choice, parallel cap, usage, nudge on/off/double-failure, generator error
hygiene, disconnect reset, empty output, MLX path).

* Address review: tool_choice none, developer folding, retry fallback, monitor reply

Four review follow-ups on the safetensors/MLX client-tool passthrough leg:
- tool_choice="none" keeps the tool-history templating but no longer
  advertises the tools, so a forced final-answer turn is not prompted into
  emitting markup that the (correctly disabled) healer would relay as prose.
  Mirrors the GGUF passthrough where llama-server honors tool_choice itself.
- OpenAI "developer" messages fold into a single leading system message via
  _set_or_prepend_system_message before templating; local templates reject the
  role and the fallback formatter drops it.
- A nudge retry that fails or is cancelled after the original answer exists
  falls back to the first response instead of surfacing a 500, matching the
  GGUF nudge path.
- The API monitor records the healed tool call summary instead of the raw
  markup on a promoted turn.

Adds four regression tests.

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

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

* Address review: forced tool_choice templating, content-part flattening, stream monitor parity

- A forced tool_choice function is now the only schema rendered into the
  local template, so the advertised tools and the healer allowlist can no
  longer disagree (llama-server enforces tool_choice itself on the GGUF path).
- Content-part lists are flattened to their text parts before templating.
  Remote image URLs are not decodable locally, so such requests reached this
  path with part lists that raise inside apply_chat_template on text-only
  templates; the plain non-GGUF path has always flattened them.
- The streaming monitor entry is now fed from the healed events the client
  actually receives, recording promoted calls as the [tool_calls] summary
  the non-streaming path records.

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

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

* Address review: gate passthrough on the engaged server path, deserialize templated arguments

- The client-tools gate now keys on _sf_use_tools (whether the server-side
  tool path actually claimed the request) instead of the raw mcp_enabled
  flag: with an empty MCP registry or a CLI --disable-tools policy, a client
  that sets mcp_enabled while declaring its own tools fell through to plain
  generation with the tools silently dropped. The GGUF passthrough gate has
  no mcp_enabled clause either.
- New _structured_tool_history_for_local_template deserializes assistant
  tool_calls[].function.arguments JSON strings into mappings for the
  templated copy only: spec-compliant clients send strings, but local chat
  templates iterate arguments as a mapping or raise on strings, which
  crashed or misrendered multi-turn tool history. The HTTP response and the
  GGUF wire shape keep strings.

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

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

* Tighten comments and docstrings in the client-tools passthrough

* Report first-attempt usage when a nudge retry is discarded

When nudge_should_retry fires but the retry produces no healable tool call
(or raises), the first response is still delivered to the client. The retry's
generate() had already overwritten stats_holder, so _monitor_usage recorded
the unseen retry's token counts against the request instead of the first
attempt that was actually returned. Capture the first attempt's stats before
the retry and restore them on both the no-heal and exception paths so the
monitor reports the usage of the response the caller received.

* Do not promote buffered tool markup when a stream is cancelled

The streaming client-tool heal path breaks out of the token loop when
cancel_event is set (the registry "Stop" path), but then still fell through to
healer.finalize(), which heals incomplete tool markup at EOF (allow_incomplete)
and emits a tool_calls delta plus finish_reason=tool_calls. Because the Stop
request only sets the event and leaves the SSE socket open, the client received
that promoted call and executed a tool the user had just cancelled. The disconnect
path already returns before finalize; guard finalize and the finish_reason on
cancel_event too, so a cancelled stream ends with finish_reason=stop and no tool
call. Adds a regression test driving a Stop mid-emission with buffered markup.

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

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

* Trim comments in the client-tools passthrough

* Trim client-tools passthrough comments further

* [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-07-06 19:48:36 -07:00
Daniel Han
8ba46b566a
Studio: close switch/cancel races during model load (#6918)
Fix six race conditions when a user switches or cancels a model while a
previous load or generation is still in flight, across the inference
orchestrator and the /load and /unload routes:

- Cancel an in-flight generation on a safetensors/MLX model switch and
  serialize unload with load under the inference lifecycle gate.
- Cancel an in-flight load off the lifecycle gate so a Stop-loading
  cancel does not wait out the multi-minute load; guard the dispatched
  mailbox against a racing unload.
- Recheck the loading marker after spawn and again after the load
  response before publishing, so a load cancelled mid-flight is reaped
  instead of going live.
- Discard the loading marker before tearing the subprocess down in
  cancel_load, closing a spawn-after-cancel window and an orphaned
  compare-mode dispatcher during unload.
- Match the unload target before canceling an in-flight GGUF load and
  add an off-gate fast path for the still-loading GGUF case.
- Run the Unsloth unload off the event loop so a paused SSE stream
  holding _gen_lock cannot block the loop.

Adds studio/backend/tests/test_orchestrator_unload_cancel.py covering
the unload/cancel/switch race paths.
2026-07-06 19:43:15 -07:00
Daniel Han
9dabe96786
Studio chat: tool-call nudging on by default (API stays opt-in) (#6883)
* Studio chat: tool-call nudging on by default (API stays opt-in)

Healing is already default-on everywhere and the nudge retry from the
client-tool passthrough is opt-in on the API. Studio chat had neither
signal: the frontend never sent nudge_tool_calls, and the safetensors
and MLX server-side loop lacked the GGUF loop's plan-without-action
re-prompt entirely.

Backend: the re-prompt helpers move from llama_cpp.py into
tool_call_parser.py (shared, cycle-free; the GGUF loop imports them
under its old names with zero behavior change) and
run_safetensors_tool_loop now re-prompts once at the streaming
no-tool-call exit, gated on Auto-Heal, active tools, nothing executed
yet, and short forward-looking text. Re-prompts do not consume tool
iterations.

Frontend: the chat adapter sends nudge_tool_calls from a new
nudgeToolCalls runtime setting (default true) with the same
persistence, hydration, and settings toggle plumbing as Auto-Heal.
Request-model defaults are untouched, so raw API callers stay opt-in.

* Address review: persist the nudge setting, consume the flag in the loops, skip the re-prompt after RAG autoinject

ChatSettingsPayload uses extra forbid, so a settings patch containing
nudgeToolCalls failed to persist any settings; the field is now typed
and round-trips. nudge_tool_calls now plumbs into both server-side tool
loops and gates the plan-without-action re-prompt with None meaning on,
so API callers keep today's behavior, explicit false disables it, and
Studio's default-on flag actually controls the path Studio chat runs.
The safetensors loop no longer re-prompts after RAG autoinject: the
injected retrieval bypasses the tool controller, so the nothing-executed
gate saw an empty history and re-asked after a successful retrieval.

* Safetensors loop: the plan-without-action retry requires an explicit nudge flag

The retry is new on this loop, so an omitted nudge_tool_calls must not
change existing API behavior; Studio opts in explicitly. The GGUF loop
keeps None as on because its re-prompt predates the flag.

* Suppress the plan-without-action re-prompt after a denied tool confirmation

A denial appends TOOL_REJECTED_MESSAGE but records nothing in the tool
controller history, so the nothing-executed gate re-prompted the model
to call the tool the user had just rejected, producing another
confirmation prompt. A denial now suppresses the re-prompt for the rest
of the request, mirroring the RAG autoinject handling.

* Tighten plan-without-action re-prompt comments

* Tighten plan-without-action re-prompt comments

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

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

* Studio: match unified plan-without-action nudge cap to GGUF default of 3

The shared MAX_ACT_REPROMPTS was set to 1, but GGUF's established default
(llama_cpp.py) has re-prompted a stalling model up to 3 times since #5620.
Restore the GGUF-matched cap so safetensors and MLX inherit the same
behavior, and update the safetensors cap test to assert the cap dynamically.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-06 19:41:19 -07:00
Daniel Han
c2a7b78f6b
Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load on Apple Silicon) (#6803)
* Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load)

mlx-lm 0.31.3 regressed the QK-norm archs: its strict load_weights rejects the
q_norm/k_norm tensors with "Received N parameters not in model", so gemma4 and
qwen3_5 checkpoints fail to load. Studio installs the MLX stack unpinned at
latest, which pulls 0.31.3. Verified on a real macos-14 runner: gemma4 fails to
load on 0.31.3 but loads and generates coherently on 0.31.2 and on git-main
(future 0.31.4). See mlx-lm #1242.

Exclude just that release (!=0.31.3) in the installer and the self-heal floor so
--upgrade still resolves to the newest good build, and treat an already-installed
0.31.3 as unsatisfied so the self-heal replaces it.

* Studio MLX: cover fresh-install path + robust bad-version compare

Address PR review:
- Fresh install.sh (Apple Silicon) runs the base 'uv pip install unsloth' with
  SKIP_STUDIO_BASE=1, skipping the guarded MLX-stack step, so transitive
  resolution could still pull mlx-lm 0.31.3. install.sh already exports
  UV_OVERRIDE -> overrides-darwin-arm64.txt before that install, so exclude
  mlx-lm 0.31.3 there too; this also strengthens the self-heal (same override).
- Match the known-bad version with parsed packaging.Version so 0.31.3 == 0.31.3.0
  (trailing-zero normalization) instead of raw string equality.

* Studio: exclude mlx-lm 0.31.3 on the fresh Apple Silicon install too

The overrides file only applies via UV_OVERRIDE when it exists relative to the
script, which is not true for a curl-piped install, and the guarded MLX step in
install_python_stack.py is skipped there (SKIP_STUDIO_BASE=1). So the base
install could still resolve the transitive mlx-lm to the broken 0.31.3. Append
mlx-lm!=0.31.3 to the base install on Apple Silicon (empty elsewhere), so the
fresh path pins away from 0.31.3 without waiting for the runtime self-heal.

* Studio: exclude mlx-lm 0.31.3 on the migrated install; keep the >=0.22.0 floor

The with-deps migrated install did not append ${_MLX_LM_EXCLUDE_ARG:-}, so a
curl-piped Apple Silicon migration (no repo overrides file, UV_OVERRIDE unset)
could resolve mlx-lm 0.31.3 transitively. Append the exclusion there, matching
the fresh install path. The no-torch migration is left alone since --no-deps
never resolves mlx-lm (same as the fresh no-torch path).

Also restore the >=0.22.0 floor in overrides-darwin-arm64.txt: a uv override
replaces the transitive constraint, so a bare !=0.31.3 could let the resolver
drop below the supported minimum that mlx_repair.py enforces at runtime.

* Triage huggingface_hub 1.22.0 / fastapi / multiprocess scanner false positives

The scan-packages gate red-failed on all three shards after transitive deps
bumped. Every new CRITICAL is a benign false positive, verified against upstream:

- huggingface_hub 1.22.0 added _sandbox.py for the remote HF sandbox feature.
  Its job-startup bootstrap string (fetch sbx-server into the container /tmp and
  exec it) and the SandboxPool host-reservation loop trip the staged-dropper and
  C2-loop heuristics; that script runs inside a remote HF container, not on the
  user machine. The bump also re-hashed the already-reviewed benign polling loops
  in hf_api.py and utils/_http.py. The PyPI artifact is byte-identical to the
  official v1.22.0 tag.
- fastapi 0.139.0 routing.py re-hashed the websocket keepalive while-True loop;
  byte-identical to upstream 0.139.0.
- multiprocess 0.70.19 forkserver.py and tests/__init__.py re-hashed the AF_UNIX
  fork-server IPC and fd-inheritance tests; genuine uqfoundation release, local
  IPC not network.

Added 7 reviewed allowlist entries (no blind regenerate). All three shards
(hf-stack, studio, extras) exit 0 locally.

* Tighten mlx-lm 0.31.3 exclusion comments

* Trim mlx-lm 0.31.3 exclusion comments
2026-07-06 19:40:06 -07:00
Daniel Han
fb94a79337 Studio: fix dataset upload data loss, caption over-count, and custom-root sd.cpp uninstall
- Diffusion dataset upload now streams each file into a sibling temp file and
  atomically os.replace()s it into place only after the whole file is written and
  within the size cap. A mid-batch 413 (or any abort) removes the temp, never an
  example already stored under the same name, so re-uploading a too-large batch can
  no longer truncate or delete a previously uploaded image.
- _diffusion_dataset_summary counts an image as captioned only when it resolves to a
  non-empty caption via the same sidecar-over-metadata precedence the trainer uses. An
  empty (tombstone) sidecar shadows a metadata row and makes the trainer skip the
  image, so counting it over-reported caption_count and mislabeled an effectively
  uncaptioned dataset as captioned.
- uninstall.sh/.ps1 now remove a custom/env-mode Studio's native diffusion build that
  installs beside the root as a stable-diffusion.cpp sibling (find_sd_cpp_binary
  resolves it from the Studio home's parent), guarded by the same unsafe-path check,
  and stop processes locking the default-mode stable-diffusion.cpp before removing it.

Adds regression tests for the upload data-loss and caption-count paths and a hermetic
shell test for the custom-root stable-diffusion.cpp removal.
2026-07-07 02:16:56 +00:00
Daniel Han
f109e7f0e6
Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704)
* Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes

Extends the rescue parsers in core/tool_healing.py and
core/inference/tool_call_parser.py to recognise two extra serialisations
local models commonly emit when bypassing native function calling:

* [TOOL_CALLS]name{json_args} (Devstral-Small-2, Mistral-Small-3.x).
* name[ARGS]{json_args} (reasoning-model rehearsal).

Both extractors use a brace-balance scan that honours escapes and
quoted strings so nested JSON args stay intact.

Also pre-strips <think>...</think> and [THINK]...[/THINK] blocks before
matching so calls emitted after a reasoning preamble are recognised
regardless of position.

Streaming gates (TOOL_XML_SIGNALS, llama_cpp.py _TOOL_XML_SIGNALS) and
the SSE strip regex (routes/inference.py _TOOL_XML_RE) gain the new
sentinels so the parser is actually invoked and the raw markup never
leaks to the UI.

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

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

* Strip unclosed think blocks and catch rehearsal [ARGS] mid-buffer

The pre-existing ``_THINK_TAG_RE`` only matched closed thinking
blocks (``<think>...</think>`` or ``[THINK]...[/THINK]``). During
streaming the model is still inside the open block when the parser
runs, so any tool-shaped markup the model is REHEARSING inside that
block survived the strip and could be executed as a real call.
Switch both copies of the regex (parser + healing) to accept the
trailing block being terminated by end-of-string in addition to
the explicit closer.

The ``_TOOL_XML_SIGNALS`` list on the llama_cpp streaming buffer
included ``[ARGS]`` to catch rehearsal syntax, but the gate used a
``startswith`` check against the buffer head -- rehearsal is shaped
``name[ARGS]{json}``, so the buffer never STARTS with ``[ARGS]``
and the signal had no effect. Add a substring fallback for the
bracket-style signals so the BUFFERING window can still divert the
stream into DRAINING when rehearsal markup arrives mid-buffer.

Adds three regression tests covering rehearsal inside unclosed
``<think>`` / ``[THINK]`` blocks (must yield no calls) and the
positive case after a closed think block (still parsed).

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

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

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

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

* Studio: harden bracket-tag tool-call parsing and streaming strip

Address review findings on the Mistral [TOOL_CALLS] / rehearsal [ARGS] paths:

- Accept hyphenated tool names in the bracket parsers and strip patterns.
  _MISTRAL_BRACKET_RE and _REHEARSAL_RE used \w+, which dropped or truncated
  MCP function names containing dashes (mcp__srv__list-issues). Use [\w-]+ to
  match the XML and Gemma parsers.
- Strip a partial bracket marker streamed before its opening brace. The
  trailing-unclosed patterns required the {, so a [TOOL_CALLS]web_search or
  python[ARGS] split across deltas leaked the raw marker to the UI. Match the
  bare marker to end-of-text, mirroring how the bare open tags are stripped.
  Closed pairs are unchanged so in-progress markup stays buffered until parsed.
- Strip a truncated bracket tail in the route-level display regex. _TOOL_XML_RE
  required a balanced JSON object; a tool call truncated by EOS now strips up
  to \Z, like the orphan-opening XML shapes. Complete calls still strip only
  their balanced JSON so following prose survives.

Add regression tests for hyphenated names, the streaming partial-marker strip,
and the unclosed-tail route strip.

* Studio: preserve XML parameter indentation in tool_healing

The chat template emits <parameter=k>\nVALUE\n</parameter>; the parameter-start
regex consumed the wrapping newline AND the value's first-line indentation via a
trailing \s*, then str.strip() removed the rest, corrupting code/diff arguments.
Narrow the trailing class to horizontal whitespace and trim exactly one wrapping
newline (_trim_param_value), preserving indentation. Matches SGLang's qwen3_coder
detector and the same fix on the multi-format parser. Add a regression test.

* Studio: tighten Mistral/rehearsal tool-call comments

Compress the comments in the Mistral [TOOL_CALLS] / rehearsal [ARGS] healing shim
and its callers to one or two lines, keeping the bracket-tag stripping rationale,
the thinking-block handling note, and the forge attribution intact.

Comment-only: no code or behavior change (verified with comment_tools.py check
--strip-docstrings; tests green).

* Studio: fix think-strip arg corruption and nested bracket-JSON strip

Review follow-up for the Mistral/rehearsal healing shim:

- The <think>/[THINK] strip ran unconditionally over the whole content before
  parsing, so a real tool argument that legitimately contained a <think> /
  [THINK] literal was silently corrupted. Don't delete the blocks: compute the
  reasoning-block spans and skip any tool-call candidate that STARTS inside one,
  across all parse paths (JSON, Gemma, XML, bracket, rehearsal). A rehearsed call
  inside reasoning is still ignored; a real call after </think> still parses.
- The bracket-tag display strip used a fixed one-level-nesting regex, so a call
  with two-level-nested JSON args either leaked raw markup or, in final mode, let
  the catch-all eat the trailing prose. Add a balanced-brace
  _strip_bracket_tag_calls pass (any nesting depth) used by strip_tool_call_markup
  and the route display strip.

Add regressions: <think>/[THINK] literal inside a real argument, rehearsal-inside-
think with a real call after, and two-level-nested bracket/rehearsal strip keeping
trailing prose.

* Studio: correct think-block comments to match span-skip behavior

The think-strip fix replaced the unconditional think-block strip with a
span-skip (the block is kept and any tool-call candidate starting inside it is
ignored), but two comments still described the old strip-first behavior. Update
the _THINK_TAG_RE comment and the parse_tool_calls_from_text docstring.

* Studio: parse Mistral arrays and call-ids, unify bracket parse/strip, keep it linear

- Parse the canonical Mistral array form (TOOL_CALLS followed by a JSON list of
  calls) and emit every call; parse the v11 shape that carries an opaque CALL_ID
  token between the name and ARGS (the function name is the token after
  TOOL_CALLS, never the call-id); and parse a Mistral call plus a rehearsal call
  in one message (the second was dropped yet still stripped from display).
- One shared balanced forward scan (_iter_bracket_spans) backs both the parser
  and the strip path, so they no longer diverge. It is linear: each regex is
  re-searched only once its cached match falls behind the cursor, replacing the
  per-match full-tail re-scan that was O(n^2) (O(n^3) over a stream). A length cap
  before the scan is a backstop.
- strip_tool_call_markup preserves think/reasoning blocks verbatim (the parser
  skips tool markup inside them), stripping only the visible text around them.
- _in_think uses bisect over the sorted think spans (was a linear scan per
  candidate).
- GGUF streaming strip runs the balanced bracket pre-pass before the regex
  patterns so nested-arg calls do not leak or eat trailing prose, and the
  BUFFERING ARGS detector requires the rehearsal name-ARGS shape.
- Tests: canonical array, array string-args, array strip keeps prose, Mistral
  plus rehearsal multi-call, v11 call-id name, think-rehearsal strip
  preservation, and bracket-strip linearity.

* Studio: preserve reasoning blocks in the route and streaming strip paths too

Addresses Gemini/Codex review: making strip_tool_call_markup preserve think
blocks left the route display strip and the GGUF streaming strip inconsistent,
so a rehearsed call inside a reasoning block was still deleted from the visible
text on those paths.

- Extract the think-block segmentation into one shared helper (strip_outside_think)
  and route all three strip paths through it: strip_tool_call_markup,
  _strip_tool_xml_for_display, and the GGUF _strip_tool_markup_streaming closure.
- Add a route-strip regression test that a rehearsal inside a reasoning block is
  preserved while a real call outside it is still stripped.

* Studio: fix bracket-tag strip/buffer review findings

Address the live code-review findings on the Mistral bracket-tag / rehearsal
tool-call rescue path:

- tool_healing: a literal think block inside a tool-call argument is no longer
  treated as a reasoning block. strip_outside_think now excludes think spans
  that sit inside a complete tool-call span, so the call is stripped whole
  instead of the split hiding its open/close pair and leaking the raw call.
- tool_healing: the rehearsal trailing-strip pattern requires a following brace
  or end-of-text, so prose that merely mentions name[ARGS] is not truncated as
  a phantom call. The bracket strip patterns are aligned with the parser
  regexes (whitespace, v11 [CALL_ID]/[ARGS] metadata, and the [CALL_ID]
  lookbehind).
- routes: strip a truncated canonical Mistral array ([TOOL_CALLS] [{... with no
  closing bracket) that the balanced scan cannot remove, align the display
  regex with the parser regexes, and apply the same rehearsal-prose guard.
- safetensors loop: mirror the GGUF [ARGS] rehearsal-substring check during
  BUFFERING so a rehearsal name does not stream before its [ARGS] arrives.

Adds regression tests for each; existing parser suite stays green.

* Studio: hold split rehearsal tool-name prefix in both streaming loops

A reasoning-model rehearsal call can stream the tool name and its [ARGS] arm in
separate chunks (web_search then [ARGS]{...}). The buffering detector only
recognised the rehearsal once [ARGS] was present, so the bare tool name was
emitted as visible content before the call drained and executed.

Add _is_rehearsal_prefix (mirrored in the safetensors loop and the GGUF loop):
when a no-signal buffer is a bare active-tool name -- or a partial prefix of
NAME[ARGS] -- hold it as a prefix instead of streaming it, so the next chunk's
[ARGS] flips it to a drain. A whitespace in the buffer means prose, not a split
call, so ordinary text still streams.

Adds regression tests for the split rehearsal in both loops and a guard that a
plain non-tool word still streams.

* Studio: route Anthropic tool-call cleanup through the protected display strip

The Anthropic stream, non-stream, and passthrough paths cleaned content with raw
_TOOL_XML_RE.sub instead of _strip_tool_xml_for_display, so a rehearsal call
inside <think> was deleted from the reasoning and a nested [TOOL_CALLS] call
dropped its trailing prose (the OpenAI-compatible paths already use the helper).
Route all four sites (prior-assistant cleanup, streaming content events,
non-stream aggregation, passthrough conversion) through the protected helper, and
add a source-level guard test so raw _TOOL_XML_RE.sub stays confined to the
helper itself.

* Studio: stop split rehearsal tool names leaking once streaming, uncapped, or unrestricted

The split-rehearsal guard (NAME in one chunk, [ARGS]{...} in the next) only held
the name in the initial BUFFERING state. Three gaps remained where the bare tool
name still streamed as visible content before the call drained:

- STREAMING: after prose had already streamed, both loops emitted a trailing
  active-tool-name token (and the GGUF/safetensors [ARGS] boundary was not pulled
  back over the name). Hold the trailing rehearsal token and release it on the
  next chunk, with an end-of-stream flush so a plain answer that merely ends on a
  tool-name word is never dropped.
- Buffer cap: a realistic MCP name longer than the 32-char _MAX_BUFFER_CHARS cap
  defeated the BUFFERING hold. A rehearsal prefix is self-bounding (it stops
  matching once it grows past NAME[ARGS]), so the generic cap no longer applies to
  it.
- Unrestricted mode (tools=[]): with no declared tool list, any bare identifier
  may be a NAME[ARGS] rehearsal, so the prefix check now recognises one instead of
  leaking the name and mis-parsing the call.

Regression tests cover the streaming, long-name, and unrestricted cases plus the
plain-prose paths that must not be held or corrupted.

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

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

* Studio tools: protect think blocks in safetensors streaming, hold split rehearsal on initial flush, advertise Mistral tools

Pass-3 review follow-ups on the Mistral [TOOL_CALLS] / rehearsal [ARGS] work:

- Safetensors streaming display strip now preserves think / [THINK] reasoning
  verbatim (routes through strip_outside_think like the GGUF path). A call
  rehearsed inside a reasoning block was stripped mid-stream and then restored by
  the final strip, a non-monotonic shrink/grow that corrupted append-by-length
  stream consumers and the visible reasoning.
- The first flush out of BUFFERING (safetensors and GGUF) now applies the same
  trailing-name hold the STREAMING branch uses, so a split rehearsal (prose plus a
  trailing active tool name in one chunk, [ARGS]{...} in the next) no longer leaks
  the bare name before the call drains.
- Safetensors capability gate no longer suppresses tools for Mistral [TOOL_CALLS]
  templates, which the shared bracket-tag parser now handles end to end. Llama
  python_tag stays suppressed (still unparseable).
- Route display strip applies the open-ended / bare-marker tail arms only on the
  segment after the last reasoning block (closed-only regex before it), matching
  strip_tool_call_markup, so a bare foo[ARGS] before a reasoning block is preserved
  while complete calls are still removed in every segment.

Adds regression tests for each and updates the now-stale Mistral capability test.

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

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

* Fix tool-call think-marker and bracket-wrapper edge cases

Round-1 review follow-ups on the Mistral/rehearsal tool-call healing:

- tool_healing: a reasoning marker that opens INSIDE a tool call's
  arguments is argument data, not a reasoning block. Add
  _think_spans_outside_tool_markup (start-inside test) and use it in
  both parse_tool_calls_from_text and strip_outside_think so a literal
  marker in one call's args no longer hides a later call (parse) or
  leaks the raw markup (strip) when the greedy match runs past the
  call's closer.
- tool_healing: strip the orphan Mistral v11 [/TOOL_CALLS] closer left
  behind after the balanced scan removes the call body. Add a route arm
  for the same closer in _TOOL_XML_RE / _TOOL_XML_CLOSED_RE.
- safetensors + llama_cpp streaming strip: run the open-ended (EOS
  anchored) tail patterns only on the last segment; segments before a
  reasoning block use the closed-only patterns, matching the final
  strip and the route strip. A bare foo[ARGS] before a reasoning block
  is prose, not a truncated call.
- safetensors streaming detector: validate each [ARGS] hit before
  draining. A bare foo[ARGS] in prose (no active tool name in front)
  no longer drains the rest of the turn; a later real NAME[ARGS] call
  is still found and the prose in between is preserved.

Regression tests added for each case across the parser, strip helpers,
and both streaming loops.

* Strip incomplete-XML tool markup with literal think tags; widen render-html detector

Round-2 review follow-ups.

- tool_healing: an UNCLOSED <tool_call> / <function= call that the parser still
  executes via allow_incomplete leaked its markup when an argument contained a
  literal think marker. _tool_call_markup_spans only covered closed calls, so the
  literal was treated as a reasoning block to preserve. Extend it to the
  open-ended XML tail forms (shared as _TOOL_OPEN_XML_TAIL_PATS) so a think marker
  inside an unclosed call is argument data and the call's markup is stripped. A
  complete call's opener stays bounded to its closed span, and a real reasoning
  block with no tool call is still preserved.
- safetensors render-html provisional card: _detect_render_html_tool_start was
  XML-only, so a Mistral [TOOL_CALLS]render_html or rehearsal render_html[ARGS]
  call executed but skipped the early card. Detect the earliest tool-call marker
  across every serialization the loop executes and fire when it is render_html.

Regression tests added for both.

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

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

* Studio tools: gate [ARGS] on active tools and skip think-block render_html rehearsal

Round 3 review fixes for the Mistral / rehearsal tool-call parsing path. Both are
asymmetric-fix bugs where one code path applied a guard the analogous paths did not.

- [ARGS] active-tool gating: the streaming state already validates a rehearsal
  NAME[ARGS] against the active tool list before draining, but the BUFFERING
  detection and the end-of-stream safety-net checks (safetensors and GGUF) treated
  any word[ARGS] substring as a tool boundary. An answer containing a literal
  foo[ARGS]{...} in prose, where foo is not an enabled tool, was drained, parsed into
  a disabled foo no-op, and forced an extra generation turn. Gate those checks on the
  active tool name too (unrestricted mode still accepts any name), so inactive-name
  prose is neither drained nor parsed. Adds a shared _has_genuine_tool_signal helper
  (safetensors) and _gguf_rehearsal_signal_pos / _gguf_has_genuine_tool_signal (GGUF).

- render_html provisional card vs think blocks: the parser skips tool candidates that
  start inside a <think>/[THINK] reasoning block, but the provisional render_html
  detector scanned raw content. A render_html rehearsed inside <think> followed by a
  real non-render_html call emitted a provisional render_html tool_start (reusing the
  later call's id) that the loop never executed. Drop candidates that start inside a
  think span and use the first marker of each shape outside the blocks. Also resolve
  the [TOOL_CALLS] [{...}] array shape through the parser so a nested "name" argument
  key no longer fires a false provisional card ahead of the real top-level tool name.

Adds regression tests for both loops: inactive-name foo[ARGS]{...} is not drained into
a disabled no-op or a retry turn, a think-block render_html rehearsal emits no
provisional card, and the array top-level name is read correctly.

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

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

* Gate ambiguous bare-rehearsal parse and strip on the active tool list

A bare NAME[ARGS]{json} is a genuine rehearsal call only when NAME is an
active tool; otherwise it is prose. The earlier round gated only detection
(so an inactive foo[ARGS] no longer drained the buffer or forced a retry
turn), but the parse and strip stayed unrestricted, which produced two
regressions:

1. An inactive foo[ARGS]{...} placed immediately before a real
   web_search[ARGS]{...} in the same content span made the real call fail
   to execute (parse consumed the phantom foo call).
2. An inactive foo[ARGS]{...} in a prose answer had its markup stripped
   from the visible text, corrupting the sentence to " is just syntax."

Thread enabled_tool_names through the shared parser/strip so parse and
strip apply the SAME active-tool gate as detection:

- core/tool_healing.py: _iter_bracket_spans skips an inactive rehearsal
  span; parse_tool_calls_from_text, _strip_bracket_tag_calls,
  _strip_markup_segment and strip_tool_call_markup accept and thread the
  gate; apply_tool_strip_patterns keeps an inactive rehearsal match.
- core/inference/tool_call_parser.py: wrappers forward the gate.
- core/inference/safetensors_agentic.py and core/inference/llama_cpp.py:
  compute the gate from the active tool list (None when unrestricted, to
  keep the legacy strip-all behavior) and thread it into every parse and
  streaming/final strip site.
- routes/inference.py: _strip_tool_xml_for_display accepts the gate and
  keeps an inactive rehearsal via a capture group on its rehearsal arm, so
  the display cleanup does not re-strip the already-correct loop output.
  The [TOOL_CALLS] control-token arms still strip unconditionally. Wire
  the current turn's active tool names into the GGUF and safetensors
  content-display sites.

Tests: parse and strip gate coverage in test_tool_call_parser_strict.py,
test_tool_xml_strip.py and test_safetensors_tool_loop.py; end-to-end GGUF
coverage for the real-call-after-inactive-rehearsal case and a
strengthened assertion that the inactive rehearsal prose survives intact.

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

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

* Studio: render the reasoning block for safetensors and MLX like GGUF

enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed <think>
into the generation prompt, so the model emits only the closing </think> then
the answer. The safetensors/MLX chat stream emitted that as plain content, so
the reasoning showed inline with no collapsible thinking block, while GGUF
(which surfaces reasoning via reasoning_content) rendered one. This brings
safetensors and MLX to parity.

- _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts
  inside the reasoning block and splits on the first </think>; default False
  keeps GGUF and every existing caller byte-identical. It suppresses a stray
  re-emitted <think> and holds partial markers back across chunk boundaries.
- _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the
  request, an enable_thinking or enable_thinking_effort style, and the template
  actually using the standard <think>/</think> markers. Models with a bespoke
  reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their
  answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are
  excluded too.
- sf_tool_stream and stream_chunks (the latter also serves MLX) feed text
  through the extractor, emitting reasoning_content then content deltas, with a
  per-turn reset in the tool loop and a flush before each tool_start; only the
  visible delta reaches the monitor reply. The two non-streaming drains split
  reasoning_content the same way.
- Tests: extractor prefilled mode (streaming and edge cases), the gate matrix
  including the gemma-style exclusion, and a route-replay of the tool-loop
  reasoning stream.

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

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

* studio: skip tool calls rehearsed in prefilled reasoning

Reasoning models (Qwen3.5 enable_thinking) open <think> in the prompt, so the
generated text starts inside the thought and emits only a closing </think> with
no opener. _think_spans_outside_tool_markup only found spans with an explicit
opener, so a NAME[ARGS]{...} or [TOOL_CALLS] call rehearsed in that leading
thought was parsed and executed as a real call.

Add a leading think span (offset 0 through the first close marker) when the
content opens with a bare close, so the rehearsed call is skipped and the
reasoning is preserved by strip_outside_think. Guarded by the existing call-span
check: a literal </think> inside a real call's arguments does not trigger the
span, so a genuine leading call still fires. Tests for both cases.

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

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

* studio: do not start prefilled reasoning mode when reasoning_effort is none

enable_thinking_effort models (e.g. GLM-5.2) express thinking-off via
reasoning_effort="none" rather than enable_thinking=False, but
_sf_reasoning_prefill_mode only looked at enable_thinking, so such a request
started the extractor in prefilled mode. With thinking off the model never emits
</think>, so the whole answer was captured as reasoning_content and the visible
content/stream came back empty. Thread reasoning_effort through and return False
when it is "none". Tests for none vs a real effort level.

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

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

* studio: only treat a leading bare </think> as prefilled reasoning when a real call follows

The prefilled-reasoning virtual span fired on any unmatched leading close marker,
so a non-prefilled turn that emits a real call before a stray </think> (for
example "Now web_search[ARGS]{...}</think> answer") had the call swallowed by the
span and dropped. Require that a real tool call also appear after the close (the
actual turn that follows the thought) before adding the span, so a stray close in
a normal answer no longer suppresses a genuine leading call. The rehearse-then-
call case still skips the rehearsal. Test for the stray-close case.

* Studio: trim redundant comments (comment-only, AST-verified)

* studio: keep tool_healing importable on Python 3.9

_balanced_json_span was annotated -> int | None. With no
from __future__ import annotations, that PEP 604 union is evaluated at
import time, so on Python 3.9 (which the package still supports,
requires-python >=3.9, and where external inference servers import this
module standalone) the def raises TypeError and the whole module fails
to import before any parsing runs.

Add from __future__ import annotations so annotations stay lazy strings,
matching the prevailing convention across studio/backend. No behavior
change: the module has no runtime annotation introspection.

* Studio: gate the Anthropic tool-stream display strip on declared tools

The Anthropic streaming and non-streaming tool paths called
_strip_tool_xml_for_display without enabled_tool_names, so with the default
strip-all behavior a final answer that literally contains an inactive-name
NAME[ARGS]{json} (prose, not a call) lost those bytes in the delivered text.
The GGUF and safetensors paths already pass _display_tool_name_gate(tools);
these two sites were missed when that gate was threaded through.

Compute the gate from the declared tools and pass it at both sites (threading
openai_tools into _anthropic_tool_non_streaming and its caller), so an
inactive-name rehearsal survives while an active-name one is still stripped.
Add a regression test.

* Studio: hold a split unrestricted rehearsal prefix at the bracket

In unrestricted tool mode (tools=[]) the rehearsal-prefix regex required
[A after the bracket, so a chunk boundary landing right after NAME[ (e.g.
web_search[ then ARGS]{...}) failed the prefix check and streamed the
partial tool markup web_search[ to the client before the call drained.
Restricted mode already holds this via a startswith check. Make the bracket
and each ARGS letter individually optional so NAME[ is held too, matching
the documented intent. Add a regression test.

* Studio: gate rehearsal detection and history strip on the original tool set

Two display/loop gate fixes so a spent one-shot tool is handled consistently:

- Rehearsal DETECTION (safetensors and GGUF loops) now uses the ORIGINAL tool
  list, matching the strip gate, instead of the post-removal active_tools. After a
  one-shot tool (render_html) runs it is dropped from active_tools; a repeat
  render_html[ARGS]{...} while another tool is still active was stripped from
  display yet never detected, so it was not routed to the render_html_repeat no-op
  and the turn ended as a blank continuation. Detection now fires for it.

- The GGUF assistant-history sanitiser forwards the enabled-tool-name gate (like
  the live-response strip), so a prior turn documenting an inactive foo[ARGS]{...}
  shape is preserved in the replayed prompt context instead of being deleted.

Add regression tests for both loops and the history strip.

* Studio: thread the tool-name gate through the remaining rehearsal/history sites

Follow-up to the rehearsal-detection and history-strip gate fixes, covering the
sibling sites that were missed:

- GGUF loop: the rehearsal-prefix and trailing-name hold checks now use the
  original tool list (_detect_tools) like the detection path, so a spent one-shot's
  split repeat (bare render_html then [ARGS]{...}) is held instead of flushed as
  visible text.
- The safetensors and Anthropic assistant-history sanitisers and the Anthropic
  non-streaming passthrough now forward the enabled-tool-name gate to
  _strip_tool_xml_for_display, matching the GGUF history sanitiser and the live
  strips, so a prior turn documenting an inactive foo[ARGS]{...} example is
  preserved in the replayed prompt / final text instead of deleted.

Add regression tests.

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

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

* Tile bracket-call spans per array item and include the v11 closer

Two with_spans fixes for the Mistral bracket parser, both hit through the
client-tool passthrough healers:
- A multi-call [TOOL_CALLS] array carried its whole markup span on the first
  call and zero-width spans after, so a consumer that filters promotions by
  the declared tool set either re-emitted the full raw array as text next to
  the promoted call or silently dropped a filtered call's bytes. The region is
  now tiled across the call-producing items (each call's span covers its own
  JSON object plus the separator bytes before it; the last span runs to the
  region end), so promoted markup strips exactly once and a skipped call's
  bytes stay visible.
- The v11 wrapper closer [/TOOL_CALLS] sat outside the reported span and
  leaked as stray text after promotion; the region now extends over an
  immediately-following closer.

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

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

* Address review: decouple healer signals from the loop signal set

The passthrough healer buffered on every TOOL_XML_SIGNALS entry, so the bare
[ARGS] rehearsal marker this branch adds for the loops (where it is gated on
active tool names) put legitimate prose like 'Use foo[ARGS] in templates'
into the holding state and stalled the stream until finalization. The healer
can never promote a bare rehearsal call, so it now buffers only on formats
its parser promotes: <tool_call>, <|tool_call>, <function=, [TOOL_CALLS].

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

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

* Condense comments in the Mistral tool-call rescue to contract essentials

* verify_import_hoist: exempt __future__ imports and same-diff relocations

Two false positives fired on this PR's refactor. A from __future__ import
is a compiler directive whose name never appears as a runtime load, so
HOISTED-IMPORT-UNUSED can never see it used, yet the file requires it for
PEP 604 annotations on Python 3.9. TARGET-CHANGED flagged the deliberate
move of the strip-pattern constants into core.inference.tool_call_parser
as a silent re-point even though the old module-level target was removed
and the new one added in the same diff. Both get narrow exemptions; a
re-point to a pre-existing target is still caught, and the self-test
negative controls all pass unchanged.

* Drain the whole Mistral [TOOL_CALLS] array in streaming passthrough healing

StreamToolCallHealer._drain promoted only the first parsed call per pass and
dropped the rest of the buffer past that one span. For a well-formed Mistral
parallel-tool-call array streamed through client-tool passthrough
([TOOL_CALLS][{...},{...}]), the per-item spans are contiguous, so after the
first call was promoted the residue began with ,{...}] (no leading signal) and
was flushed as raw text: every call after the first was lost.

_drain now walks the contiguous run of parsed calls (adjacent tiled spans =
one array), promoting each declared call and relaying undeclared ones as data,
and stops at the first gap (prose) or incomplete trailing block so separate
blocks still stream incrementally in document order. This mirrors the
non-streaming heal_openai_message / finalize promote-or-flush loop and the
server-side safetensors loop, which already handled multi-call arrays.

Added regression tests: 2-call array in one feed and char-by-char, an
undeclared middle call kept as text, and an array followed by trailing prose.

* Drain comma-less Mistral tool-call arrays and normalize null arguments

The array branch fed the whole body to a single json.loads, which rejects the
comma-less multi-call form the repo's own Mistral/Ollama templates render (the
range loop in ollama_template_mappers.py emits the objects with no separator) and
so dropped every call. Decode elements individually with the existing
comma-tolerant raw_decode helper, now _decode_array_items, which also returns the
objects, so all calls are recovered while the span tiling is unchanged.

Also normalize a non-object array argument such as arguments null to an empty
object, matching the wrapped tool_call path, instead of serializing None to the
string "null" that auto-heal would turn into a bogus query of "null".

* Gate safetensors reasoning prefill on the rendered generation prompt

reasoning_always_on fires on any paired <think></think> in the template,
including markup that only renders PAST assistant history (Kimi-K2-Thinking)
while the generation prompt opens no <think>. Starting the reasoning extractor
in prefilled mode there captured a normal answer entirely as reasoning_content
and returned blank visible content. Prefill only when rendering the generation
prompt actually leaves <think> open (DeepSeek-R1 / QwQ / Qwen3-Thinking);
history-only templates start the extractor in normal mode and parse the model's
own <think>...</think>. Adds a Kimi-shape regression test.

* Keep bare scalar Mistral array arguments raw instead of double-encoding

A scalar string argument in the canonical Mistral [TOOL_CALLS] array
(for example [TOOL_CALLS][{"name":"web_search","arguments":"weather"}])
was run through json.dumps, turning weather into the JSON string
"weather". The downstream argument healer then wrapped that quoted
form, so a single-string tool like web_search searched for the literal
"weather" with quotes. The <tool_call> path already keeps a scalar
argument raw; mirror it here so only a dict is serialized. Add a
regression test asserting both paths yield the same healed arguments.

* Tighten tool-call rescue and reasoning-prefill comments

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

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

* [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-07-06 18:52:13 -07:00
Daniel Han
2e75b0131c Merge remote-tracking branch 'origin/main' into fold-integration
# Conflicts:
#	scripts/scan_packages_baseline.json
2026-07-07 01:46:07 +00:00
Daniel Han
233949cc9c
scan_packages: baseline transitive-dep drift in the supply-chain scan (#6917)
The pip scan-packages gate (SCAN_ENFORCE=1) blocks on non-baselined
CRITICAL/HIGH findings. Recent upstream releases of transitive
dependencies added new files/loops that trip the pattern scanner, so all
three shards (extras, hf-stack, studio) red-failed on legitimate library
code. Add the 7 reviewed findings to scripts/scan_packages_baseline.json.

Each entry is genuine upstream code from the official PyPI archive:

- huggingface-hub huggingface_hub/_sandbox.py (staged dropper + C2 loop):
  the HF Jobs sandbox bootstrap string and its host-pool reservation
  loop. New in huggingface_hub 1.x (pulled via huggingface_hub>=0.34.0).
- huggingface-hub huggingface_hub/hf_api.py, utils/_http.py (C2 loop):
  standard polling / retry while True loops.
- fastapi fastapi/routing.py (C2 loop): websocket receive loop.
- fastmcp-slim fastmcp/cli/apps_dev.py (fs enum + network): the FastMCP
  dev CLI (PrefectHQ) making httpx/socket calls.
- cffi cffi/_cffi_gen_src.py (compile + exec): cffi generating and
  running C extension source, its core purpose.

Additive only: no existing baseline entry is changed or removed. Verified
by re-running the scanner over the full closure on Python 3.12.13 (the CI
interpreter); it now exits 0 with only MEDIUM findings remaining.
2026-07-06 18:34:18 -07:00
Daniel Han
660ef9aab7 Reconcile stale tests after stack fold: typed resolved-field roundtrip; drop dense-fbcache test superseded by video auto-quant 2026-07-07 01:32:08 +00:00
Daniel Han
de099eaecd Merge remote-tracking branch 'origin/video-hunyuan-gate' into fold-integration
# Conflicts:
#	studio/backend/core/inference/video.py
#	studio/backend/routes/models.py
#	studio/backend/tests/test_cached_gguf_routes.py
#	studio/frontend/src/features/images/images-page.tsx
#	studio/frontend/src/features/images/train/diffusion-train-panel.tsx
2026-07-07 01:15:20 +00:00
Daniel Han
c7822fa728 Merge remote-tracking branch 'origin/video-wan' into fold-integration 2026-07-07 01:08:50 +00:00