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.
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).
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
_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.
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.
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.
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.
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.
- 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.
* 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>
* 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>
- 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.
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.
- 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.
- _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.
* 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>
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.
* 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>
* 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
- 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.
* 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>
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.
Daniel Han