ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: add local speech-to-text dictation engine (#7095)

Browse source 
* Studio: add Voice settings tab (dictation, dictionary, read aloud)

New Voice tab in Settings, placed just before About:

- Dictation: microphone picker, browser STT engine, recognition language,
  and an inline mic test with a live transcript
- Dictation dictionary: entries rewrite matching speech to their exact
  spelling and casing, applied in both dictation paths
- Recent dictations: last 20 final transcripts with copy and clear, so
  text can be recovered if it lands in the wrong place
- Read aloud: optional button on assistant responses with two engines,
  curated system voices (novelty and legacy voices filtered, quality
  ranked, capped at 20) or the TTS audio model loaded in Unsloth via
  /audio/generate (e.g. Orpheus), plus speed, pitch, volume and preview

Settings persist in localStorage (unsloth_voice_settings) and are read
at call time so changes apply without reloading the runtime. Adds en
keys plus the tab label for ja, zh-CN and pt-BR.

* Studio: drop the single option STT engine select, rename TTS option

The STT engine dropdown only had one entry, so it added noise without
giving a real choice. The engine row can come back once local STT
models land. Also renames the TTS engine option Unsloth TTS model to
Load TTS model to make the action clearer.

* Studio: harden Voice settings against edge cases found in simulation

Simulated the feature across Chromium, Firefox and WebKit plus node
level unit runs and backend contract checks. Fixes from the findings:

- Dictionary rewrite used a replacement string, so entries containing
  dollar patterns corrupted transcripts (A$$AP became A$AP, $& injected
  the match). Switched to the callback form of String.replace
- Persisted voice settings now validate types on hydration: non string
  micDeviceId, dictationLanguage and ttsVoiceURI, and non boolean
  ttsEnabled fall back to defaults instead of flowing into the UI
- Dictionary entries are trimmed, capped at 120 chars and re-sanitized
  on hydration
- The Test dictation panel now falls back to the default microphone
  when the saved device is unplugged, matching the composer adapter

Test coverage: 46 unit assertions (dictionary regex edge cases across
unicode, word boundaries and injection, voice curation for simulated
macOS, Windows and Linux voice inventories, corrupt storage merge),
13 backend contract checks against /audio/generate on an isolated
instance, and 60 browser assertions across the three engines covering
rendering, degradation without SpeechRecognition, curation in a real
DOM, dictionary persistence with unicode and dollar entries, the
no-model preview error path and corrupt localStorage recovery.

* Studio: address Voice settings review feedback

Verified each review comment before acting. Confirmed and fixed:

- Editing a dictionary entry was broken in two ways: the store trimmed
  on every keystroke so spaces could not be typed, and clearing the
  field deleted the entry and unmounted the input mid edit. Updates now
  keep the raw value and a blur commit trims or removes the entry
- The unplugged mic fallback checked instanceof DOMException, but a
  cross browser probe showed Firefox and WebKit throw
  OverconstrainedError objects that are not DOMExceptions, so the
  fallback never fired there. Matching on the error name now
- When the browser ended a dictation test on its own (silence timeout),
  the mic stream stayed open. All recognition end paths now stop the
  tracks and save the transcript through a single finalize path
- The studio TTS audio element now releases its WAV data URL as soon as
  playback ends, fails or is cancelled
- Allow microphone now reports insecure contexts (no mediaDevices)
  accurately instead of claiming access was blocked
- Voice tab copy moved into i18n keys per src/i18n/AGENTS.md, so locale
  overlays can translate it; en is the baseline and parity passes
- unsloth_voice_settings added to the Reset all local preferences key
  list so voice preferences obey the reset
- Non default microphones note that the system default is used when the
  browser speech engine cannot bind a specific device, since browsers
  without the start(track) overload ignore the argument silently

Re-ran the full simulation set after the changes: 46 unit assertions,
13 backend contract checks and 60 browser assertions across Chromium,
Firefox and WebKit all pass, plus a dedicated browser probe for the
dictionary editing behavior.

* Studio: use the chat mic icon in Voice settings for consistency

The Voice tab and its buttons used the hugeicons Mic02 glyph while the
chat composer uses a custom filled mic. Extract that composer icon into
a shared lib/mic-icon component, drop the duplicate inline copies in
thread.tsx and shared-composer.tsx, and use it for the Voice tab icon
and the tab's mic buttons so the microphone looks the same everywhere.

* Studio: address second round of Voice settings review feedback

Verified each new comment against the current code first. One item was
already fixed in the previous round (recording transcripts when the
browser ends a dictation test on its own). Confirmed and fixed:

- The microphone row showed a picker with generic names when browsers
  enumerate unlabeled devices before permission, leaving no way to
  grant access from the row. It now branches on whether labels are
  visible and shows Allow microphone otherwise
- Compare chat dictation ignored the selected microphone. It now opens
  the chosen device with the same fallback rules as the main adapter,
  passes the track to recognition where supported and releases the
  stream when recognition ends
- Closing the Voice tab cancelled the shared speechSynthesis even when
  read aloud was playing a chat message. Cleanup now only cancels when
  the tab owns an active preview
- Double clicking Start test could race two recognizers and leak the
  first stream. A starting flag set before the getUserMedia await makes
  start reentrancy safe
- Turning off the read aloud setting mid playback removed the only stop
  control. The stop button now renders whenever a message is speaking
- When an engine lacks the start(track) overload, both dictation paths
  now release the selected device stream before retrying with the
  default microphone instead of holding it open
- Read aloud support no longer requires Web Speech synthesis: the
  Unsloth TTS engine only needs audio playback, so it stays available
  in WebViews without speechSynthesis, with a clear error if the system
  engine is chosen there

Not addressed here: cancelling in flight backend TTS generation on
stop. The route runs generation in a worker thread without a
cancellation path, which is shared pre existing behavior with audio
chat generation and belongs in a backend change.

All suites re-run green: 46 unit, 13 backend contract and 60 browser
matrix assertions across Chromium, Firefox and WebKit, plus probes for
the unlabeled device branch and the double click race.

* Studio: drop empty and duplicate voiceURIs so the Voice tab never renders a crashing Select item

* Studio: guard dictation mic lifecycle in Voice test and Compare composer

Release a microphone opened after the component unmounts, and stop Compare
dictation on a permission or security failure instead of silently recording
from the default device, matching the main chat adapter.

* Studio: fix dictation and read-aloud lifecycle edge cases in Voice settings

- Join final dictation chunks with a space so recorded transcripts do not merge words
- Ignore a stale recognizer onend so a quick stop then restart is not torn down
- Use previewingRef so a double click on TTS preview does not orphan the first request
- Keep the read-aloud stop control visible when a new run starts while a message is spoken
- Stop the dictionary remove button from deleting an adjacent entry on a blur then click race

* Studio: trim redundant Voice settings comments

* Studio: fix Voice preview and Compare dictation edge cases

- Only cancel the shared speechSynthesis for a system-voice preview, so stopping
  a Studio preview no longer stops an unrelated chat read-aloud
- Release the Studio preview audio and its WAV data URL on normal completion
- Iterate every finalized result in Compare dictation so batched phrases are kept
- Cap persisted recent dictations to the last 20 on hydration

* Studio: use clipboard fallback for recents and release failed preview audio

- Copy recent dictations via the copyToClipboard helper so the execCommand
  fallback works in Safari and insecure http LAN contexts
- Release the Studio preview audio when play() rejects, not just on ended/error

* Studio: add local speech-to-text dictation engine

Add an offline dictation engine that transcribes with a local faster-whisper
model, alongside the existing browser (Web Speech) engine. The browser engine
streams audio to Apple or Google speech services and needs internet; the new
engine runs on the server, works offline, and drives any chat model without
evicting it (it loads in the backend process, separate from the model
subprocess). It also gives Firefox dictation, which has no Web Speech support.

Backend: a lazily-loaded, kept-warm faster-whisper sidecar and three routes
under /api/inference/audio (stt/status, stt/load, transcribe). faster-whisper
is torch-free, so this does not disturb the existing model stack.

Frontend: a Dictation engine setting (browser or local model), a curated model
picker with sizes, and MediaRecorder capture posted to the transcribe route.
The model warms automatically when the engine is selected, with live status.

* Studio: stream local STT transcription as you speak

Local dictation showed nothing until you stopped, because the whole clip was
transcribed once on stop. Now the growing recording is re-transcribed on a
fast pass every second and emitted as live interim text, with an accurate
final pass on stop. Partial recordings decode fine, and the model refines
earlier words as more audio arrives.

Adds an interim flag to the transcribe route (beam 1, no VAD) for the fast
preview pass; the final stop uses the accurate path.

* Studio: make local dictation stop instant and reliable

Stopping local dictation waited for a final network transcription before the
session ended, so the stop button did not flip and a second click ended the
session early and dropped the text. Now stop commits the live transcript
immediately, releases the mic at once, and ignores a second stop while
finalizing. Previews run more often so the committed text is current.

* Studio: record local dictation in short clips for reliable streaming

Re-transcribing a growing buffer every second got slower as it grew, flooded
the backend, showed stale words, and could leave the stop button stuck waiting
on a backlog. Record short independent clips instead and transcribe each once,
appending the text as you speak. Work per clip is bounded, so stopping is
prompt (with a hard timeout as a safety net) and long dictations stay smooth.

* Studio: dictate then transcribe once on stop, ChatGPT style

Local STT dictation streamed by re-transcribing the growing clip, which
was quadratic and saturated the backend (multi-second lag), and stop only
halted the recorder without releasing the mic, so it kept recording. Record
the microphone continuously, release it the instant the user stops, and
transcribe the whole clip once. Stopping is immediate and the transcript
lands in about a second. Also add the tiny model for the fastest option.

* Studio: surface dictation and read-aloud failures instead of failing silently

- Compare dictation reports microphone and speech-recognition errors via toast,
  reusing the main chat adapter's describeMediaError and describeSpeechError
- Read-aloud toasts genuine model or synthesis failures while ignoring cancellations

* Studio: ChatGPT-style recording bar for dictation

Clicking the mic now drops the composer into a dedicated recording bar
with a live waveform, a discard (X) and a confirm (tick), instead of a
plain stop button. The tick stops recording and transcribes the clip;
the X throws the recording away and keeps whatever text was already in
the composer. The model adapter taps the mic with an analyser to drive
the waveform, and the router tracks the live session so the X can cancel
it without transcribing.

* Studio: transcribe dictation while speaking, ChatGPT layout

Match ChatGPT's recording layout: the bar now renders in place of the
input with the left plus button kept, the waveform in the middle, and
the discard and confirm buttons together on the right.

Cut the post-confirm delay by transcribing in the background as the user
talks. The audio is split at natural pauses (voice-activity detection off
the same analyser that drives the waveform) and each clip is transcribed
as it is cut, so confirming only has to finish the short final tail. The
model is also warmed when recording starts so the first run never pays a
cold load.

* Studio: ChatGPT waveform, hide tools while dictating, faster STT

Make the recording UI read like ChatGPT: the waveform is now a dense row
of round dots that rise into thin centered bars, and while dictating only
the plus button shows, with the mode badge and tool toggles hidden so the
bar is just the waveform and controls.

Speed up transcription: decode greedily (beam_size=1), which is several
times faster on CPU with negligible accuracy loss on short dictation
clips, and cap background segments at 6s so the final tail after confirm
stays short.

* Studio: finish ChatGPT voice bar and low-latency STT

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

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

* Studio: full-width waveform with a timer that freezes on stop

Use the full-width waveform for the recording bar: brighter, bigger bars
that advance on a fixed cadence (keeping peaks between advances) so they
glide instead of racing by, inset from the composer edges. Keep a visible
timer and the green confirm button, matching the ChatGPT reference, and
freeze the timer and waveform the moment the user confirms.

* Studio: fix multilingual local dictation

* Studio: speed up dictation and release local STT

* Studio: harden dictation finalization and STT decoding

* Studio: restore Firefox dictation fallback

* Studio: add dictation history manager

* Studio: manage speech model downloads

* Studio: remove em dash from voice model label

* Studio: move dictation history into Voice

* Studio: source local STT from Unsloth Whisper models

Point the dictation STT sidecar and its Model Hub download entries at
Unsloth's Hugging Face Whisper repos (small, large-v3-turbo, large-v3)
and run them through Transformers, so Studio only ever downloads
Unsloth-uploaded weights. Drop faster-whisper and the Systran/mobiuslabs
repos; keep the Model Hub as the only download path via local_files_only,
and keep PyAV for audio decoding.

Device selection uses float16 on CUDA and float32 on MPS and CPU, since
Whisper's decoder is unstable in float16 on MPS and repeats tokens.

Shorten the model picker labels to name plus download size and update the
STT tests for the new backend.

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

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

* Studio: smooth dictation waveform and keep pill height

* Studio: align STT model dropdown width and tidy voice copy

* Studio: guide to local engine when browser dictation is offline

* Studio: clarify voice section and STT model copy

* Studio: keep STT warm with training-aware eviction

* Harden STT lifecycle and browser compatibility

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

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

* Fix model discovery test lint

* Harden cross-browser microphone errors

* Harden cross-browser microphone errors

* Surface voice test recognition errors and fall back to Studio TTS

- Voice test now toasts non-abort speech-recognition failures instead of
  ending silently, matching the main and Compare dictation paths.
- Read-aloud routes to the backend model when the runtime lacks Web Speech
  synthesis (audio-only WebView), so it no longer errors immediately.

* Fix reviewed STT lifecycle races

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

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

* Fix read-aloud fallback controls

* Guard read-aloud stop when deleting a non-speaking message

aui.message().stopSpeaking() throws unless this message is the one being
read aloud, so calling it unconditionally rejected the delete handler before
the message was removed. Only stop speech when this message is speaking.

* Cap recent dictation transcript length before persisting

Recent dictations only limited entry count, so a long transcript stored the
full text in the persisted voice settings and a few could exceed the
localStorage quota, throwing synchronously from the uncaught dictation cleanup
path. Truncate each entry on save and on hydration, matching the dictionary cap.

* Studio: keep dictation mic clickable and guide to local model

Register the dictation adapter unconditionally so the mic stays enabled
for any engine and starts working right after switching to the local
model on an already-open thread.

When the browser engine cannot run (Firefox, Brave, non-secure origins),
clicking the mic shows a toast that points to the local speech-to-text
model instead of leaving a disabled button. The toast stacks its action
below the text with a fully rounded button.

* Studio: add bottom padding below the dictation guidance toast button

* Studio: increase bottom padding under the dictation toast button

* Studio: add bottom padding inside the dictation toast button

* Studio: add five Whisper defaults and custom model search

Add private UnslothAI Tiny and Base mirrors to the curated local STT choices while keeping Small as the default. Let users search or paste a Transformers-compatible Whisper repository and validate it end to end.

Keep short dictations in one clip to avoid repeated padded encoder work, then split longer recordings near Whisper's 30-second boundary.

Update hidden model filters and tests, including the CPU-only CI runtime stub for PyAV.

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

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

* Studio: use public Unsloth Whisper repositories

Point the Tiny and Base dictation defaults to the public unsloth repositories and remove the private mirror references from model filtering and tests.

* Studio: update Whisper download sizes

Reflect the cleaned public Tiny and Base repositories in the curated model labels.

* Studio: right-align STT model size, fix dropdown wheel scroll, refresh sizes

- Show the download size on the right of each model row so long names
  like Whisper Large v3 Turbo no longer hide it
- Update curated Whisper sizes to the safetensors weights actually
  downloaded: Tiny 151 MB, Base 290 MB, Small 967 MB
- Drive the model list scroll from a wheel handler so the mouse wheel
  scrolls it inside the Settings dialog, not just the scrollbar
- Add a search icon and shorten the placeholder to Search model

* Studio: do not search when a dictation model is picked, shrink repo label

- Treat the filled-in model text as a selection, not a query, so choosing
  a model no longer kicks off a Hugging Face search
- Make the repository line under each model name smaller

* Studio: tighten dictation model and local engine descriptions

* Studio: keep model display on pick instead of the query, shrink row text

- Guard the combobox input so selecting a model shows its name and does
  not echo the typed query back or start a search
- Map the item label to the friendly display so picks fill the field
- Reduce the model name and size text in each row

* Studio: show only the model name in the dictation field, shrink size label

- Drop the download size from the search field; the name alone is shown
  once a model is selected, with sizes kept in the dropdown list
- Reduce the size label text in each row

* Studio: clarify the dictation model description

* Studio: drop Hugging Face from the dictation model description

* Studio: move the dictation dictionary to its own Manage subpage

- Replace the inline entry list with a Manage row, matching Dictation
  history, so a long dictionary no longer crowds Voice settings
- Add a DictationDictionaryView subpage that holds the entry editor

* Studio: match STT field font, use best voice for System default

- Bump the dictation model field text to text-sm so it matches the
  engine dropdown next to it
- Resolve the System default read-aloud voice to the top curated voice
  instead of the browser default, which is a robotic legacy voice on macOS

* Studio: rerank read-aloud voices and drop duplicate voice entries

- Rank by vendor quality, then the user's locale, then a preferred list of
  natural voices, so the best voice leads instead of the first alphabetically
- Collapse voices that macOS reports twice under one name and language

* Studio: fold dictionary and recents into the dictation section

- Drop the separate Dictation dictionary and Recent dictations headings;
  their Manage rows now sit under Dictation, split by the row divider
- Shorten the custom spellings description

* Studio: add search and sort to dictation history

- Filter saved dictations by text with a search field
- Sort by newest, oldest, or A to Z; show a no-matches message
- Keep Clear all available regardless of the current filter

* Studio: settle cancelled STT loads before training and fix dictation review items

Wait for a cancelled STT load to exit and release its memory before
reporting it freed for training, so the loader cannot still be inside
from_pretrained()/.to(device) holding VRAM when the training subprocess
starts. A load that finishes before observing the cancel now gets
unloaded so the memory is actually reclaimed.

Clear the accelerator cache before the CPU fallback in load() so a failed
CUDA/MPS load does not strand reserved VRAM once the sidecar is marked
CPU-resident.

Send the saved Hugging Face token when polling STT download progress so a
gated or private repo resolves and shows the correct Load/Downloaded
state instead of reporting missing.

Mark the composer Dictate button as type="button" so clicking it does not
also submit the draft when the composer already has text or attachments.

* Studio: pin dictation settings per session and close STT startup races

Capture the STT model and language when a dictation session starts and
pass them to every queued segment and the warm-up load, so changing the
model or language mid-recording no longer transcribes the same clip with
the wrong model or a model that is not downloaded.

Check the local runtime at the top of transcribe(), before the model
cache lookup and the bounded audio decode, so a server missing PyTorch or
Transformers returns 501 up front instead of decoding a long clip first.

Treat the training startup window as active for STT device selection.
start_training frees VRAM in before_spawn but only assigns _proc later, so
a concurrent STT load could take the GPU that was just cleared. A startup
flag now reports training active from the free until the process is live,
forcing those loads to CPU; a finally clears it on every exit.

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

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

* Studio: stub the STT runtime check in transcribe orchestration tests

transcribe() now verifies the local runtime up front, so the unit tests
that exercise transcription orchestration must treat the runtime as
present to keep passing where PyTorch, Transformers, and PyAV are not
installed. Stub ensure_stt_available in the shared fixture and restore
the real check in the availability and load-rejection tests.

* Harden custom Whisper dictation models

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

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

* Add whisper.cpp dictation engine with per-engine downloads and history rework

Engines
- New GGML STT sidecar that runs a managed whisper-server subprocess with
  idle unload, plus a pinned static build script (scripts/build_whisper_cpp.sh)
- Dictation engine picker now offers Browser, Local transcription
  (whisper.cpp), and Local transcription (Transformers)
- Both local engines serve the same five curated Whisper models and download
  them directly with byte-level progress reported by /audio/stt/status
- Models auto load on selection and when their download finishes
- Unload and training admission account for both engines

Benchmarks (Apple Silicon, greedy, warm, same checkpoints)
- whisper.cpp transcribes 2.4x to 5x faster than Transformers and loads in
  about 0.45s vs 0.86s for Whisper Small
- whisper.cpp GGUF path is unchanged by the Transformers addition
  (load 0.445s -> 0.444s, short clip 0.391s -> 0.347s, long 1.197s -> 1.129s)

Voice settings UI
- Plain curated model select replaces the searchable combobox
- Single download progress bar with transfer rate for both engines
- Dictation history now stores every dictation with Show more pagination,
  a top Clear history action, and links back to the chat it was spoken into
- Archived chats dialog gets the same pagination
- Delete dialog offers deleting a dictation together with its chat

Tests: 88 backend STT tests pass, including new snapshot download coverage.
Frontend typecheck, lint, i18n parity, and production build pass.

* Merge local engines into one option and source GGML models from unslothai

Engine selection
- The dictation engine dropdown is back to two choices: Browser and Local
  transcription. The selected model decides the backend: curated ids run
  GGML checkpoints through whisper.cpp, searched Hugging Face repositories
  run safetensors through Transformers
- Model picker lists the curated models and searches Hugging Face for other
  Whisper repositories, validating them before selection. The trigger is a
  plain button so the selection never renders inside a text input
- /audio/stt/status accepts a model query param so downloaded state works
  for custom repositories; the engine param on load, transcribe, and
  download routes is derived from the model everywhere

Model source
- Curated GGML checkpoints now download from the Unsloth-hosted
  unslothai/whisper-*-GGUF repositories (one repo per model) instead of
  ggerganov/whisper.cpp; cache lookups, progress totals, and in-flight blob
  tracking are per-model

Fixes
- Voice settings and dictation history were not persisting: the quota-safe
  localStorage wrapper was declared after the store that uses it, so the
  persist storage factory failed silently. Every settings write also threw
  mid-click, which kept the model picker popover from closing on selection
- is_model_downloaded now verifies config, preprocessor config, and real
  weight files instead of trusting an offline snapshot lookup, so a partial
  download left by an aborted fetch shows the Download button instead of
  failing to load
- Removed whisper.cpp mentions from user-facing text: the ready status
  shows Loaded instead of the runtime name, picker rows show the source
  repository, and runtime error messages say local transcription runtime

Verified with automated browser sessions and live API checks: selection
closes the picker with no page errors, persisted settings hydrate on
reload, a stale partial snapshot triggers download then loads on MPS and
transcribes, and curated models download from the unslothai repos. 88
backend STT tests, typecheck, lint, i18n parity, and build pass.

* Skip the duplicate source line for custom models in the STT picker

A custom repository's display name is its id, so search results and the
appended current selection rendered the same string twice. The source
line now only renders when it differs from the name; curated rows keep
their name, unslothai source repository, and download size.

* Verify every shard of a sharded checkpoint in the downloaded check

A snapshot holding one of N shards (or a corrupt shard index) passed the
downloaded check and then failed at load. When model.safetensors.index.json
exists, every shard in its weight map must now be present. Found by
simulation; covered by a regression test.

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

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

* Rename stale _starting references in the pump resilience tests

The startup flag on TrainingBackend was renamed to _spawn_in_progress but
two tests added alongside it still asserted on the old name, failing the
Python 3.11 to 3.13 CI jobs.

* Make the selected model row clearly highlighted in the STT picker

The current selection was a faint background tint. It now uses the accent
background with a medium weight name. Two line rows use a small corner
radius; single line custom repo rows keep the pill shape.

* Address review feedback on STT snapshot checks, VRAM release, and dictation UX

Verify snapshot completeness in the load preflight so a partial download
fails before the audio is decoded, for curated and custom repos alike.
Drop the failed accelerator traceback before the CPU retry so the cache
clear can actually release that memory. Keep unloading the GGUF sidecar
after cancelling an in-flight Transformers load; both engines can hold
memory at once. Allow Auto language with English-only .en checkpoints,
matching the backend which sends no forced language. Keep the discard
button usable while a transcription is pending so a slow or hung request
cannot trap the composer in dictation mode. Stop linking Compare and
settings test dictations to the unrelated active single chat thread.

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

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

* Move the CPU retry out of the exception handler

On Python 3.10 the interpreter exception state keeps its own reference
to the traceback, so dropping it from the caught exception was not
enough to release the failed accelerator load during the retry. Leaving
the handler before clearing the cache works on every supported version.

* Address review feedback on session handoff, chat pinning, and server lifetime

Starting a dictation from a second entry point now cancels the session
it replaces, so the old recording cannot keep the microphone open or
save a transcript with no discard button pointing at it. The linked
chat is pinned when recording starts, so switching threads while a
transcription finalizes cannot relink the transcript to the newly
opened chat. whisper-server is now bound to Studio's lifetime like the
other long-lived children: PDEATHSIG on Linux, the parent job object on
Windows, and pid adoption so the shutdown sweep reaps it; before this
it survived a Ctrl+C exit as an orphan still holding the model.

* Remove the dictation mic test from Voice settings

The composer dictate button covers the same check, so the test row, its
transcript panel, the unsupported fallback row, and their strings and
search entry are gone.

* Studio STT: gate GGUF whisper-server on training and fix dictation retry and dictionary edits

GGUF (whisper.cpp) sidecar:
- Launch whisper-server with --no-gpu while training is active, mirroring the Transformers sidecar's CPU device choice, so a mid-training dictation cannot reclaim the VRAM training just freed.
- Report is_loading() during whisper-server startup so training VRAM admission accounts for the accelerator memory it is about to bind.
- Require PyAV in is_available() so /audio/stt/status reports the engine unavailable when uploads cannot be decoded, instead of loading fine and then 501ing at transcription.
- Reject a missing model before decoding audio, matching the Transformers download preflight.

Voice settings:
- The download Retry button now restarts the download; the sidecar error is sticky until a new start(), so re-polling alone never cleared it.

Dictation dictionary:
- Tabbing from an emptied entry to its remove button no longer commit-splices the row first, which shifted indices and deleted the wrong entry.

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

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

* Studio STT: fix curated GGUF whisper filenames to match hosted repos

The unslothai/whisper-*-GGUF repos host the checkpoint as whisper-<id>.bin,
not ggml-<id>.bin, so every curated dictation download and cached-path
lookup 404'd and the whisper.cpp engine could never load a model. Point
GGML_STT_MODELS at the real filenames and guard the naming with a test.

* Studio STT: validate a custom dictation repo before downloading it

The Transformers STT engine accepts an arbitrary owner/model repo, but the
download route handed it straight to snapshot_download, pulling a possibly large
non-Whisper repository into the shared HF cache. Confirm the repo is a Whisper
checkpoint first with the existing metadata-only validate_remote_model (no
weights); curated ids short-circuit and the GGUF engine (curated-only) is
unaffected. A non-Whisper repo now 422s before any download.

* Studio STT: preempt a still-loading GGUF server for training admission

A whisper-server still in its startup window binds accelerator memory but has no
loaded_model yet, so training admission could miss it and launch into an OOM.
Make the GGUF startup cancellable (cancel_pending_load signals an abort event and
terminates the starting process without the load lock; _wait_for_server observes
it and raises SttLoadCancelledError; wait_for_load_to_settle blocks on the lock
until the killed server is reaped), and always fold the GGUF sidecar into the
resident-STT summary so a resident Transformers model cannot mask a loading GGUF
server. free_stt_model_for_training now cancels an in-flight load and waits for it
to settle before training claims the memory.

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

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

* Studio STT: fall back to Transformers when whisper-server is absent

A curated dictation model (including the default small) hard-pinned the GGUF
engine, but standard installs do not ship whisper-server, so every recording
501'd instead of using the Transformers engine that serves the same checkpoint
-- the GGUF sidecar's own documented contract. Add _resolve_serving_stt_engine:
a GGUF request for a curated id (the only ids GGUF accepts, all Transformers-
servable) downgrades to Transformers when whisper-server is unavailable, applied
consistently to download, load and transcribe (not unload, which targets a
specific engine). The Voice tab likewise falls back to the Transformers status so
the model is not shown unavailable and download is not blocked.

* Studio STT: hide custom Whisper caches from the legacy model pickers

The legacy /cached-models (and /cached-gguf) routes called is_hidden_model with
only the owner/model id, which cannot reach the config-based Whisper check, so a
downloaded custom (non-curated) Whisper checkpoint was still offered as a chat
model. Pass the cached snapshot path so _path_is_whisper_model inspects the repo
config and hides it, matching the discovery route.

* Studio STT: hide GGUF dictation repos, lock-free status, unload fallback, split training eviction

- Hide the curated GGUF dictation repos (unslothai/whisper-*-GGUF) from the chat
  model inventory and pickers, backend and frontend. Only their Transformers
  safetensors companions were hidden; the GGUF repos use a different org and a
  -GGUF suffix and carry a raw .bin with no whisper config.json, so they leaked
  into chat pickers.
- Make the GGUF sidecar loaded_model/device accessors lock-free, mirroring the
  Transformers sidecar. transcribe() holds self._lock across the whole inference
  call, so /audio/stt status polls and training admission previously blocked
  behind an in-flight transcription.
- stt_unload resolves through the serving resolver: a "gguf" pick on a host
  without whisper-server is served by the Transformers fallback, so unload must
  target that engine or the resident model is never freed. Unload also attempts
  every engine even if one raises, so a failure freeing one backend no longer
  skips the other.
- free_stt_model_for_training frees the Transformers and GGUF sidecars under
  independent exception boundaries so a failure unloading one no longer skips
  the other before training claims the memory.

Adds tests/test_stt_review_fixes.py covering all four.

* Studio STT: resolve Auto dictation language for the model engine + snapshot process liveness

- The model dictation adapter sent the raw setting (the literal "auto") to the
  backend, while the browser engine resolves Auto via resolveDictationLanguage.
  A batch of non-English voice notes came back mostly English on Auto. Add
  resolveModelDictationLanguage: only the literal "auto" is resolved to a
  concrete locale, gated so it becomes a language the model AND Whisper can
  honor (mirroring the backend's known-whisper-languages set); an explicit
  language, or a locale Whisper cannot honor, stays unchanged/auto-detect. Wire
  it into both adapter call sites.
- GgmlSttSidecar._process_alive() read self._process twice; a concurrent
  unload() nulls it under the lock while loaded_model/device read lock-free, so
  a null between the two reads called None.poll(). Snapshot once. Adds a
  deterministic regression test.

* studio: tighten comments and docstrings in the dictation modules

* studio: harden dictation model downloads, GGML readiness, and recording paths

Address review findings on the STT dictation feature:

- build_whisper_cpp.sh refuses to delete a whisper.cpp tree under a custom
  Studio home unless it carries the Studio ownership marker, matching the
  setup.sh policy, and marks trees it creates
- _snapshot_is_complete validates every shard of a sharded PyTorch
  (pytorch_model.bin.index.json) checkpoint like the safetensors path, and
  requires tokenizer assets (tokenizer.json or vocab.json + merges.txt)
- custom-repo downloads pin the revision resolved at validation time and
  restrict snapshot_download to the model/tokenizer/config/preprocessor file
  classes Studio loads
- the GGML sidecar holds its port reservation until just before spawning
  whisper-server and only accepts readiness from a responder that both looks
  like whisper.cpp's server and belongs to the still-running managed child,
  probing twice, so mic audio cannot be posted to a foreign local process
- the recording adapter transcribes every non-empty segment; the RMS meter
  only shapes segment boundaries and can no longer discard quiet speech
- Compare-pane dictation can cancel a pending transcription on second click,
  with the button relabeled while finalizing
- localStorage quota recovery halves the dictation history until the save
  fits, so small histories shrink too
- the System default TTS voice resolves to the platform default voice
- new dictation UI imports go through the chat and hub feature barrels

Regression tests cover the build-script gate, sharded PyTorch and tokenizer
completeness, revision pinning and allow patterns, and the whisper-server
readiness probe.

* Fix STT download and voice picker follow-ups

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

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

* Add dictation button regression coverage

* Studio: prebuilt whisper.cpp via the shared llama.cpp install core, slim bundles paired to the llama prebuilt (#7294)

* Studio STT: add prebuilt whisper.cpp (whisper-server) installer

New install_whisper_prebuilt.py downloads a per-platform whisper-server
bundle published by the unslothai/whisper.cpp prebuilt CI into the managed
whisper.cpp dir (build/bin/whisper-server) so local dictation needs no
compiler. Mirrors install_node_prebuilt.py / install_llama_prebuilt.py:
host + backend detection, sha256 pins (whisper_prebuilt_pins.json) as the
trust anchor, staging + install lock + atomic swap, traversal-safe extract,
co-located shared libs (RUNPATH=$ORIGIN), an UNSLOTH_WHISPER_PREBUILT_INFO.json
marker with idempotent "already matches", and exit codes 0/1/2/3. Not wired
into setup yet; the pins ship empty so every asset fails closed until the
first fork release is published and its digests are reviewed in.

* Studio STT: install prebuilt whisper.cpp during setup and update

Add a fail-open whisper.cpp block to setup.sh after the llama.cpp section so
`unsloth studio update` (and a fresh install) fetch the prebuilt whisper-server
into the managed whisper.cpp dir the sidecar discovers. It skips a user-set
WHISPER_SERVER_PATH/UNSLOTH_WHISPER_CPP_PATH, honors UNSLOTH_SKIP_WHISPER_INSTALL,
forwards the resolved ROCm gfx, and never aborts setup: a busy install keeps the
existing runtime, and an unavailable prebuilt stays quiet (source build is opt-in
via UNSLOTH_WHISPER_FORCE_COMPILE) since Transformers STT and browser dictation
remain. Register UNSLOTH_WHISPER_PREBUILT_INFO.json as Studio-owned evidence.

* Studio STT: harden whisper-server child env + WSL ROCm detection

- Sidecar spawns whisper-server with a scrubbed child env that prepends the
  binary dir (co-located GPU libs) to the loader path, and on WSL2 ROCm loads
  the system HIP first (HSA_ENABLE_DXG_DETECTION=1) so a bundle's bare-metal HIP
  does not segfault on /dev/dxg. Secret-bearing vars are dropped from the child.
- find_whisper_server_binary now requires an executable, not just a file.
- Installer rocm probe passes HSA_ENABLE_DXG_DETECTION and falls back to
  /opt/rocm/bin/rocminfo so a WSL ROCm host is not misdetected as CPU-only;
  gfx parsing skips the gfx000 CPU agent and generic ISA lines.
- Tests for the child env (secret scrub, lib dir, WSL HIP precedence), the
  executable check, and the WSL rocm detection.

* Studio STT: in-app whisper.cpp prebuilt update stack + ship pins in the wheel

Mirror the llama.cpp update stack for the whisper.cpp prebuilt so Studio can
detect and install a newer whisper-server release from inside the app:
- backend/utils/whisper_cpp_freshness.py: read UNSLOTH_WHISPER_PREBUILT_INFO.json
  and compare the installed release against the newest unslothai/whisper.cpp
  release. Whisper tags are v<upstream>-unsloth.<N>, so is_behind compares a
  (major, minor, patch, serial) key with a strict downgrade guard; 24h cache;
  fail-open.
- backend/utils/whisper_cpp_update.py: run install_whisper_prebuilt.py to fetch
  and atomically swap the newest bundle, unloading the warm GGUF sidecar first.
- backend/routes/whisper.py mounted at /api/whisper (update-status + update).
- pyproject: add whisper_prebuilt_pins.json to studio package-data so the
  installer's trust anchor ships in the wheel (it is a data file, not a .py
  module, so package discovery alone does not include it; node_prebuilt_pins.json
  is listed for the same reason). Without this a pip-installed wheel had no pins
  and the prebuilt install aborted to Transformers STT.
Adds test_whisper_cpp_freshness.py (version parser, is_behind matrix + downgrade
guard, marker layouts, stale decision, fail-open).

* Studio STT: verify whisper prebuilts via the release checksum index, like llama.cpp

Re-align the whisper.cpp prebuilt installer to install_llama_prebuilt.py's trust
model: instead of a committed whisper_prebuilt_pins.json, verify every download
against the release's own whisper-prebuilt-sha256.json checksum index, fetched
from the same GitHub release.

- parse_release_checksums / fetch_release_checksums / expected_sha256_for replace
  the pins layer. The index is validated for schema/component and that its
  release_tag matches the resolved release; an asset absent from it, a release
  that does not publish it, or a manifest sha256 that disagrees with it all fail
  closed to a source build.
- resolve_release_tag now resolves the newest published release at runtime (or an
  explicit --published-release-tag), matching llama and the freshness check;
  removed the pinned-default and the UNSLOTH_WHISPER_ALLOW_UNVERIFIED opt-in.
- Delete studio/whisper_prebuilt_pins.json and drop its pyproject package-data
  entry (nothing to ship now, same as llama which has no committed pins).
- Adds test_install_whisper_prebuilt_checksums.py (index parser, fail-closed on
  uncovered asset, tampered-manifest guard, newest-release resolution).

This is a same-origin checksum (integrity, not authenticity), identical to the
llama.cpp installer; pair releases with GitHub artifact attestations for provenance.

* Resolve whisper prebuilt release via the download host (no GitHub API)

Mirror install_llama_prebuilt.py's fast path: resolve the release tag from
the releases/latest redirect and fetch the manifest + checksum index from
constructed releases/download URLs, so the common install path makes zero
api.github.com calls (unauthenticated api.github.com is capped at 60 req/hour
per IP; the download host is not). Fall back to the GitHub API only on a 404,
malformed asset, or tag mismatch.

* Studio STT: coverage-aware whisper prebuilt selection via a shared core

whisper's select_artifact returned the first os/arch/backend manifest match and
ignored the SM-coverage fields the release manifest already carries, so a
Blackwell B200 (sm_100) was served cuda12-legacy (sms 50-61) -- runnable only via
forward PTX JIT. install_llama_prebuilt.py on the same host correctly picks
cuda13-newer.

Extract the coverage-aware selection into a shared, component-agnostic core under
studio/backend/utils/prebuilt/ (selection + GPU host-capability detection), lifted
from llama's linux_cuda_choice_from_release / _artifact_covers_sms / _sm_range and
generalised over a normalised artifact. whisper's HostInfo now records the GPU
compute caps + driver CUDA version (honoring CUDA_VISIBLE_DEVICES), and
select_artifact routes CUDA/ROCm through the shared selector: every visible SM
must be covered, the tightest-covering profile wins (Blackwell-aware runtime-line
ordering), ROCm matches the gfx target exactly, and an uncovered GPU falls back to
the CPU bundle. CPU/Metal/Vulkan keep first-match. The resolver JSON, exit codes,
and "already matches" contract are unchanged.

On the B200 the installer now resolves cuda13-newer, matching llama.

* Studio STT: gate whisper CUDA selection on the on-disk runtime, like llama

The prebuilt CUDA bundles are dynamically linked and intentionally do NOT ship
libcudart/libcublas -- they load the same runtime the host already has. So the
driver's advertised CUDA version is only an upper bound: a cuda13 bundle still
needs cuda13 runtime libraries present on disk. Port llama's on-disk runtime
scan (detected_linux_runtime_lines / detected_windows_runtime_lines) into the
shared core and intersect it with the driver-compatible lines in
select_cuda_attempts. A host with a cuda13 driver but only cuda12 runtime (e.g.
torch-cuda12) now correctly gets a cuda12 bundle instead of an unloadable cuda13
one; a host with no CUDA runtime at all falls back to CPU.

Fixes a glob bug in the port (any(Path(d).glob(p) for d in dirs) tests generator
truthiness, not a match) that made every major report present; add a real
filesystem test that exercises the scan.

* studio: harden shared prebuilt core to full llama parity

Apply the review findings on the shared coverage-aware prebuilt-consumer
core so whisper.cpp selection is exactly equivalent to the llama.cpp path.

hosts.py: port llama's CUDA_VISIBLE_DEVICES handling. A GPU hidden by an
index/UUID selector now reports has_usable_nvidia False instead of staying
usable, via supports_explicit_visible_device_matching plus the physical /
explicit-match branches, and _select_visible_rows now matches rows the way
llama does (index or UUID, gpu- prefix optional) and skips unmatched tokens
rather than keeping all rows. Adds the Linux /proc/driver/nvidia/gpus
fallback and has_physical_nvidia. Adds parse_macos_version.

runtime_libs.py: the Linux on-disk scan now requires the exact libcudart /
libcublas SONAME (libcudart.so.13), not a libcudart.so.13* glob, so a bare
versioned file without the SONAME symlink no longer counts as loadable.
Hardens the ldconfig parse against an empty left-hand side.

selection.py: fix the Blackwell/torch reordering so it keys on the covering
runtime lines (falls through to the torch preference when the covering lines
were filtered out), matching linux_cuda_choice_from_release. Corrects the
compatible_runtime_lines_for_driver docstring: the bundles do not ship the
CUDA runtime, so the driver version is only an upper bound and the caller
must intersect with the on-disk scan.

install_whisper_prebuilt.py: enforce a macOS artifact's min_os (new
HostInfo.macos_version) so a bundle that cannot load on the host OS version
is dropped. Keep resolver stdout to only the JSON line by leaving logs on
stderr in --resolve-prebuilt mode, and map an unexpected probe failure to
prebuilt_available False instead of a traceback.

Tests: new host-probe suite for the visible-device logic, exact-SONAME
runtime-scan cases, macOS min_os filtering, resolver stdout-only-JSON,
exit-code mapping, and the repo key.

* studio: fix whisper prebuilt selection + launch parity gaps from review

A parallel review surfaced integration defects where the whisper path could
select or launch a bundle that cannot run on a concrete host. Each is fixed to
match install_llama_prebuilt.py.

macOS min_os: the manifest labels macOS requirements as macos-<version>
(e.g. macos-14.0), which the version parser could not read, so the guard was a
no-op and a macOS-13 host would install the macos-14 Metal bundle. Strip the
platform prefix before parsing.

ROCm gfx detection: _detect_rocm_gfx returned the first gfx token and ignored
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES. Since exact
ROCm matching treats that token as the active GPU, a mixed APU + dGPU host
(gfx1151 + gfx1100) with HIP_VISIBLE_DEVICES=1 installed the wrong archive. Route
through a shared pick_rocm_gfx_target (lifted from llama) that parses per-GPU
sections and honors the visibility vars (empty / -1 -> no AMD GPU).

--rocm-gfx override: recording the arch without setting has_rocm left the host on
its CUDA/CPU path so the ROCm bundle was never picked. --rocm-gfx now implies
has_rocm and clears NVIDIA state, like llama's _apply_host_overrides.

CUDA launch env: a CUDA bundle ships the ggml CUDA backend but not
libcudart/libcublas, and the sidecar launch env exposed only the bundle dir, so
on a host whose CUDA runtime lives only in the PyTorch wheels the selection would
gate cuda usable but the server could not load it. Add the CUDA-from-PyTorch
runtime dirs to the child loader path for CUDA bundles (bundle dir still first),
mirroring binary_env.

Also normalize a manifest artifact's supported_sms defensively (parity with
llama's parser) and document that blackwell_min_toolkit_for_caps is retained for
the Phase B llama Windows path.

Not changed (verified parity, not defects): Linux/Windows min_os is enforced
nowhere in llama (macOS only); the resolver is optimistic about the checksum
index and the install path verifies.

* studio: tighten prebuilt-core code comments

* studio: lift shared prebuilt installer core out of the whisper installer

* studio: reuse the llama.cpp prebuilt installer machinery for whisper

* studio: unify llama and whisper prebuilt installers on a shared descriptor core

* studio: consolidate prebuilt installer tests into the shared core suite

Grow tests/studio/install/test_prebuilt_core.py from 62 to 164 tests so every
component-agnostic behavior runs against both descriptors: the full seven
profile CUDA release matrix (multi-GPU, on-disk runtime gating, shuffle
stability, missing SM metadata, dotted SM normalization, no-driver fallback
policy), the ROCm gfx family matrix, macOS min_os gating and its helper,
backend resolution incl. cpu-fallback precedence and Intel-mac auto detect,
checksum-index non-object and plain-lookup cases, the tar symlink/hardlink
extraction guards moved from the llama suite, and the compute-cap, visible
device, runtime-line and Blackwell helper value tables moved verbatim from
the llama characterization suites.

Delete only tests whose exact behavior the master now asserts for the same
component: 40 pure-alias helper cases in test_selection_logic.py (replaced by
value-identical master tables plus an alias-identity pin), 6 extraction moves
and the master-absorbed zip-symlink case in the llama logic suite, 3 routing
twins in test_rocm_support.py already pinned byte-for-byte in
test_selection_logic.py, the 2 Blackwell helper tables in the backend resolve
suite, 28 whisper logic tests and 10 whisper checksum tests re-asserted by
the master whisper parameterization. Wrapper wiring pins, the llama release
plan dialect, fingerprints and every llama-only behavior stay untouched.

* studio: dedupe sidecar and update helpers into the backend prebuilt package

* studio: chain whisper.cpp prebuilt updates onto the llama.cpp update flow

* studio: consume paired slim whisper prebuilts via the llama ggml runtime

* studio: serve every whisper backend from slim prebuilts

* studio: drop the whisper fat per-accelerator selection chain

unslothai/whisper.cpp releases are slim-only from v1.9.1-unsloth.2: one
ggml-less bundle per os/arch, paired to the llama.cpp prebuilt that provides
every ggml backend. Delete the whisper-side fat CUDA/ROCm/metal/vulkan
selection glue; keep slim selection + pairing, link_ggml_runtime, and one
legacy shape, the published fat CPU bundle of an explicitly pinned pre-slim
release. Exit 2 now reads as prebuilt unavailable (whisper never source
builds); setup already treats it that way.

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

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

* Wire libomp runtime DLL alongside ggml in slim whisper installs

llama's clang-built windows-arm64 ggml-base.dll imports
libomp140.aarch64.dll, shipped in the llama bundle but not a system DLL.
Without it next to whisper-server.exe the loader fails with
STATUS_DLL_NOT_FOUND before main. MSVC x64 links vcomp140.dll from
System32 and Linux ggml uses system libgomp.so.1, so only windows-arm64
was affected. The empty-runtime guard still requires a real ggml
library; libomp alone is not a pairing.

* studio: drop whisper-side fat-selection support structure

Slim whisper bundles are selected per os/arch only; all accelerator
capability comes from the installed llama.cpp prebuilt, whose installer
already did the coverage-aware selection. Remove the machinery that only
existed to pick among fat per-accelerator whisper bundles:

- prebuilt_core: delete the generic CUDA/ROCm coverage selection
  (select_cuda_artifact, select_rocm_artifact, ArtifactView adapters,
  detected_cuda_runtime_lines, the exact-SONAME linux probe) that no
  shipped component routes through; llama keeps its own selection chain
  and whisper shadows select_artifact with the slim-only version.
  select_artifact is now a plain os/arch/backend first-match.
- install_whisper_prebuilt: drop the HostInfo CUDA fields
  (compute_caps, driver_cuda_version, torch_runtime_line) and the torch
  runtime probe that populated them; nothing reachable reads them, and
  the resolver payload sources runtime_line from the artifact.
- whisper_cpp_update: delete the standalone start_update job worker;
  whisper applies only run as the chained phase of the combined
  llama+whisper update. The status payload keeps its job field (idle).
- routes/whisper: drop the progress logger that could never fire.
- tests: remove tests of the deleted paths and tests duplicating the
  descriptor-parameterized core suite or the llama freshness suite.

Contracts unchanged: resolver JSON keys, exit codes, marker fields,
pairing logs, and the pinned pre-slim fat CPU escape hatch.

* Address review feedback on the whisper prebuilt update and install paths

- Pin the chained whisper phase to the release the freshness check
  offered, so the download-host latest pointer cannot reinstall an
  older build in a loop
- Wire the whisper prebuilt install into setup.ps1 (Windows setup
  previously skipped it entirely)
- Treat a non-executable server or missing wired ggml libraries as a
  broken install instead of reporting already matches
- Keep whisper sidecar reloads out of the job-level reload flag and
  resync chat state after a partial chained update that unloaded llama
- Repoint home and profile vars for the whisper-server subprocess at a
  managed scratch dir and drop credential-store pointers
- Clear the prebuilt marker before the opt-in source build overwrite
- Write the prebuilt marker with explicit utf-8 encoding

* Tighten comments in the whisper prebuilt consumer

* Harden the Windows whisper setup phase and the chained update edges

- setup.ps1: honor WHISPER_SERVER_PATH / UNSLOTH_WHISPER_CPP_PATH /
  UNSLOTH_SKIP_WHISPER_INSTALL, run the custom-home ownership guard
  before the atomic install, and forward the release-tag pin and ROCm
  hints like setup.sh
- sidecar: a cpu-selected install launches whisper-server with --no-gpu
  (slim wiring links every llama backend, so the flag is what keeps a
  deliberate CPU choice off the GPU)
- chained update: leave whisper unpinned on macOS (the llama phase can
  walk back there, and a newest-tag pin could be an impossible pairing
  on every retry) and treat installer exit 2 as kept-existing-runtime
  instead of failing the combined job
- job.to_tag now comes only from the llama phase, so a whisper-only
  round cannot report a llama update that never ran

* Fix slim whisper runtime follow-ups

* Address remaining whisper update reviews

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

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

* Address remaining prebuilt update reviews

* Fix remaining chained update reviews

* Fix remaining whisper runtime review edges

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

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

---------

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

---------

Co-authored-by: danielhanchen <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Michael Han 2026-07-23 01:39:03 -07:00 • committed by GitHub
commit d5cf96d628
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
81 changed files with 19343 additions and 2813 deletions
Download patch file Download diff file Expand all files Collapse all files

71
scripts/build_whisper_cpp.sh Executable file
View file

@ -0,0 +1,71 @@
#!/bin/sh
# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine.
#
# Installs into the managed Studio home so the backend's binary discovery
# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up:
# <UNSLOTH_STUDIO_HOME>/whisper.cpp/build/bin/whisper-server (custom home)
# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default)
#
# Usage:
# ./scripts/build_whisper_cpp.sh # build the pinned tag
# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh
#
# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a
# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's
# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux).
set -eu
WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}"
WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}"
STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}"
CUSTOM_STUDIO_HOME=false
if [ -n "$STUDIO_HOME" ]; then
CUSTOM_STUDIO_HOME=true
INSTALL_DIR="$STUDIO_HOME/whisper.cpp"
else
INSTALL_DIR="$HOME/.unsloth/whisper.cpp"
fi
command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; }
command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; }
# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete
# a directory under a custom Studio home unless Studio itself created it (the
# marker file below). Protects a user-managed whisper.cpp/src from rm -rf.
STUDIO_OWNED_MARKER=".unsloth-studio-owned"
if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \
[ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then
echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
exit 1
fi
echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR"
mkdir -p "$INSTALL_DIR"
: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER"
if [ ! -d "$INSTALL_DIR/src/.git" ]; then
rm -rf "$INSTALL_DIR/src"
git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src"
else
git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG"
git -C "$INSTALL_DIR/src" checkout FETCH_HEAD
fi
CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF"
if [ "${GGML_CUDA:-0}" = "1" ]; then
CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON"
fi
# shellcheck disable=SC2086
cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS
NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)"
cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU"
mkdir -p "$INSTALL_DIR/build/bin"
cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server"
echo "==> Installed $INSTALL_DIR/build/bin/whisper-server"
"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK"

876
studio/backend/core/inference/stt_ggml_sidecar.py Normal file
View file

@ -0,0 +1,876 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp (GGML/GGUF) speech-to-text sidecar for Studio dictation.
Runs the same curated Whisper checkpoints as the Transformers sidecar
(stt_sidecar.py) through whisper.cpp's `whisper-server`, ~2.5x faster at
identical quality on Apple Silicon and CPU because its Metal/CPU kernels run
the weights in f16 where PyTorch MPS requires fp32.
Owns a single `whisper-server` subprocess bound to 127.0.0.1 on an ephemeral
port; the model loads on demand, stays warm between dictations, and unloads
after the same keep-alive as the Transformers sidecar. Curated GGML checkpoints
are single files from `unslothai/whisper-*-GGUF`, downloaded directly rather
than through the Model Hub (whose variant planner only handles `.gguf` chat
layouts).
Binary discovery mirrors `_find_llama_server_binary`: env override, then managed
Studio home, then PATH. With no binary the engine is unavailable and dictation
falls back to the Transformers sidecar; `scripts/build_whisper_cpp.sh` installs
the binary.
"""
from __future__ import annotations
import io
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import threading
import time
import urllib.request
import uuid
import wave
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, Optional
from loggers import get_logger
from core.inference.stt_sidecar import (
STT_KEEP_ALIVE_SECONDS,
SttAudioDecodeError,
SttLanguageError,
SttLoadCancelledError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
_decode_audio_bounded,
_known_whisper_languages,
_TARGET_SAMPLE_RATE,
_training_active,
normalize_whisper_language,
)
from utils.prebuilt.child_env import isolate_home, scrub_env, wsl_system_rocm_lib_dirs
from utils.prebuilt.runtime_libs import dedupe_existing_dirs
from utils.prebuilt.whisper_layout import lookup_marker
from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
logger = get_logger(__name__)
# Curated GGML checkpoints, one repo per model. Keys match the Transformers
# sidecar's ids so the frontend reuses one picker; values are the single file
# inside each repo.
GGML_STT_REPOS: dict[str, str] = {
"tiny": "unslothai/whisper-tiny-GGUF",
"base": "unslothai/whisper-base-GGUF",
"small": "unslothai/whisper-small-GGUF",
"large-v3-turbo": "unslothai/whisper-large-v3-turbo-GGUF",
"large-v3": "unslothai/whisper-large-v3-GGUF",
}
GGML_STT_MODELS: dict[str, str] = {
"tiny": "whisper-tiny.bin",
"base": "whisper-base.bin",
"small": "whisper-small.bin",
"large-v3-turbo": "whisper-large-v3-turbo.bin",
"large-v3": "whisper-large-v3.bin",
}
DEFAULT_GGML_STT_MODEL = "small"
_SERVER_START_TIMEOUT_SECONDS = 120.0
_TRANSCRIBE_TIMEOUT_SECONDS = 600.0
class SttEngineUnavailableError(SttUnavailableError):
"""whisper-server is not installed; the GGUF dictation engine is off."""
def resolve_ggml_model_id(model: Optional[str]) -> str:
"""Validate a curated GGML model id. Custom repos are not supported here."""
if model is None or not str(model).strip():
return DEFAULT_GGML_STT_MODEL
normalized = str(model).strip()
if normalized in GGML_STT_MODELS:
return normalized
raise SttModelIdError(
f"STT model '{model}' is not a curated GGUF dictation model. "
f"Choose one of: {', '.join(GGML_STT_MODELS)}."
)
def _managed_whisper_cpp_dir() -> Path:
"""`<STUDIO_HOME>/whisper.cpp` in custom mode, else `~/.unsloth/whisper.cpp`.
Mirrors `managed_node_dir` / `_find_llama_server_binary` so managed runtimes
share one parent directory.
"""
legacy = Path.home() / ".unsloth" / "whisper.cpp"
try:
from utils.paths.storage_roots import studio_root
resolved = studio_root()
legacy_studio = Path.home() / ".unsloth" / "studio"
try:
is_legacy = resolved.resolve() == legacy_studio.resolve()
except (OSError, ValueError):
is_legacy = resolved == legacy_studio
return legacy if is_legacy else (resolved / "whisper.cpp")
except (ImportError, OSError, ValueError):
override = (
os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") or ""
).strip()
if override:
try:
return Path(override).expanduser().resolve() / "whisper.cpp"
except (OSError, ValueError):
return Path(override).expanduser() / "whisper.cpp"
return legacy
def find_whisper_server_binary() -> Optional[str]:
"""Locate the whisper-server binary.
Search order:
1. WHISPER_SERVER_PATH environment variable (direct path to binary)
2. UNSLOTH_WHISPER_CPP_PATH env var (custom whisper.cpp install dir)
3. managed dir: <STUDIO_HOME or ~/.unsloth>/whisper.cpp/{,build/bin/}whisper-server
4. whisper-server on PATH
"""
binary_name = "whisper-server.exe" if sys.platform == "win32" else "whisper-server"
def _layout_candidates(d: Path) -> list[Path]:
cands = [d / binary_name, d / "build" / "bin" / binary_name]
if sys.platform == "win32":
cands.append(d / "build" / "bin" / "Release" / binary_name)
return cands
env_path = os.environ.get("WHISPER_SERVER_PATH")
if env_path:
p = Path(env_path)
if _is_runnable(p):
return str(p)
custom_dir = os.environ.get("UNSLOTH_WHISPER_CPP_PATH")
if custom_dir:
for p in _layout_candidates(Path(custom_dir)):
if _is_runnable(p):
return str(p)
for p in _layout_candidates(_managed_whisper_cpp_dir()):
if _is_runnable(p):
return str(p)
return shutil.which(binary_name)
def _is_runnable(p: Path) -> bool:
"""A real whisper-server is an executable file. On Windows os.access(X_OK) is
effectively an existence check; on Unix it rejects a non-executable stub so a
half-written or wrong-mode file isn't mistaken for the server."""
return p.is_file() and (sys.platform == "win32" or os.access(p, os.X_OK))
def _whisper_install_marker(binary: str) -> Optional[dict]:
"""The prebuilt install marker above ``binary``, or None (source/custom builds)."""
return lookup_marker(binary).marker
def slim_runtime_intact(binary: str) -> bool:
"""True unless the marker says slim and the linked ggml runtime is missing
beside the server. New markers record the exact wired filenames
(linked_libraries), all of which must be present; legacy markers without the
field fall back to the per-OS core ggml name globs. A broken slim install
reads as engine-unavailable (reinstall via `unsloth studio update`), never a
crash at load."""
lookup = lookup_marker(binary)
marker = lookup.marker
if lookup.invalid or marker is None:
return not lookup.slim_collision
if not marker or marker.get("install_kind") != "slim":
return True
if lookup.authoritative:
valid = marker.get("component") == "whisper.cpp"
valid = valid and isinstance(marker.get("schema_version"), int)
valid = valid and all(
isinstance(marker.get(key), str) and marker[key]
for key in ("release_tag", "backend", "paired_llama_tag")
)
valid = valid and isinstance(marker.get("linked_libraries"), list)
valid = valid and bool(marker.get("linked_libraries"))
valid = valid and all(
isinstance(name, str) and name and Path(name).name == name
for name in marker["linked_libraries"]
)
if not valid:
return False
bin_dir = Path(binary).parent
linked = marker.get("linked_libraries")
if isinstance(linked, list) and linked and all(isinstance(name, str) for name in linked):
intact = all((bin_dir / name).is_file() for name in linked)
else:
if sys.platform == "win32":
required = ("ggml.dll", "ggml-base.dll")
elif sys.platform == "darwin":
required = ("libggml*.dylib", "libggml-base*.dylib")
else:
required = ("libggml.so*", "libggml-base.so*")
intact = all(any(p.is_file() for p in bin_dir.glob(pattern)) for pattern in required)
runtime_dirs = marker.get("linked_runtime_directories")
if intact and isinstance(runtime_dirs, list) and runtime_dirs:
intact = all(
isinstance(name, str)
and name
and (bin_dir / name).is_dir()
and any(path.is_file() for path in (bin_dir / name).rglob("*"))
for name in runtime_dirs
)
if intact and marker.get("backend") == "rocm":
expected_runtime_dirs = set() if sys.platform == "win32" else {"hipblaslt", "rocblas"}
intact = (
marker.get("runtime_wiring_version") == 2
and isinstance(runtime_dirs, list)
and set(runtime_dirs) == expected_runtime_dirs
)
if not intact:
logger.warning(
"slim whisper install is missing its linked ggml runtime at "
f"{bin_dir}; run `unsloth studio update` to reinstall it"
)
return intact
def is_available() -> bool:
binary = find_whisper_server_binary()
if binary is None:
return False
if not slim_runtime_intact(binary):
return False
try:
import av # noqa: F401
except Exception:
# No PyAV means every transcription 501s on decode.
return False
return True
def ensure_engine_available() -> str:
binary = find_whisper_server_binary()
if binary is None:
raise SttEngineUnavailableError(
"The local transcription runtime is not installed. Run "
"`unsloth studio update` to install it."
)
if not slim_runtime_intact(binary):
raise SttEngineUnavailableError(
"The local transcription runtime is missing its paired ggml "
"libraries. Run `unsloth studio update` to reinstall it."
)
return binary
# ---------------------------------------------------------------------------
# whisper-server child-process environment
# ---------------------------------------------------------------------------
# Build the whisper-server env: prepend the binary dir (co-located libs win, and
# a backstop where the loader ignores the rpath) and scrub secret-bearing vars the
# binary never needs. On WSL2 ROCm the system HIP libs go first, since a bundle's
# bare-metal HIP cannot drive /dev/dxg. A CUDA bundle ships libggml-cuda.so but not
# libcudart/libcublas (paired with the user's PyTorch), so add the
# CUDA-from-PyTorch runtime dirs the selection gated on, else the backend cannot
# resolve a runtime that lives only in wheels. Mirrors llama's binary_env(); the
# scrub/WSL/dedupe helpers live in utils.prebuilt.
# Module-level aliases keep the historical patch points for tests and callers.
_wsl_system_rocm_lib_dirs = wsl_system_rocm_lib_dirs
_dedupe_existing_dirs = dedupe_existing_dirs
def _whisper_server_child_env(binary: str) -> dict[str, str]:
"""Env for the whisper-server subprocess: secrets scrubbed, home/profile vars
repointed at a managed scratch dir (a downloaded binary must not see the real
home's token caches), co-located libs on the loader path, WSL system HIP first
on WSL2 ROCm."""
env = scrub_env(os.environ)
isolate_home(env, str(_managed_whisper_cpp_dir() / ".child_home"))
bin_dir = str(Path(binary).parent)
# A CUDA bundle needs the CUDA-from-PyTorch wheel dirs so libcudart/libcublas
# resolve at launch when they live only in site-packages/nvidia/*/lib. Placed
# after bin_dir so co-located libs still win; empty for other bundles.
cuda_runtime_dirs: list[str] = []
bundle_dir = Path(bin_dir)
has_cuda_module = any(
path.is_file()
for pattern in ("libggml-cuda.so*", "ggml-cuda*.dll")
for path in bundle_dir.glob(pattern)
)
if has_cuda_module:
try:
from utils.prebuilt.runtime_libs import python_runtime_dirs
cuda_runtime_dirs = python_runtime_dirs()
except Exception:
cuda_runtime_dirs = []
if sys.platform == "win32":
var, lead = "PATH", [bin_dir, *cuda_runtime_dirs]
elif sys.platform == "darwin":
var, lead = "DYLD_LIBRARY_PATH", [bin_dir]
else:
var, lead = "LD_LIBRARY_PATH", [bin_dir, *cuda_runtime_dirs]
wsl_rocm = _wsl_system_rocm_lib_dirs()
if wsl_rocm:
lead = [*wsl_rocm, bin_dir, *cuda_runtime_dirs]
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
existing = [p for p in env.get(var, "").split(os.pathsep) if p]
env[var] = os.pathsep.join(_dedupe_existing_dirs([*lead, *existing]))
return env
# ---------------------------------------------------------------------------
# Model file download (single files; deliberately outside the Model Hub flow)
# ---------------------------------------------------------------------------
def _cached_model_path(model_id: str) -> Optional[str]:
"""Path of a fully downloaded GGML file in the shared HF cache, else None."""
from huggingface_hub import hf_hub_download
try:
return hf_hub_download(
repo_id = GGML_STT_REPOS[model_id],
filename = GGML_STT_MODELS[model_id],
local_files_only = True,
)
except Exception:
return None
class _GgmlDownloadState:
"""Tracks one background hf_hub_download of a curated GGML file."""
def __init__(self) -> None:
self._lock = threading.Lock()
self._thread: Optional[threading.Thread] = None
self._model_id: Optional[str] = None
self._error: Optional[str] = None
self._total_bytes: Optional[int] = None
self._etag: Optional[str] = None
def status(self) -> dict:
with self._lock:
downloading = self._thread is not None and self._thread.is_alive()
return {
"downloading": downloading,
"model": self._model_id if downloading else None,
"error": self._error,
"bytes_total": self._total_bytes if downloading else None,
"bytes_done": self._incomplete_bytes() if downloading else None,
}
def _incomplete_bytes(self) -> Optional[int]:
"""Best-effort progress: size of the in-flight blob in the HF cache.
hf_hub_download writes ``blobs/<etag>.incomplete``; prefer this file's
etag, else the largest in-flight blob.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
# Caller may hold the non-reentrant self._lock; bare reads are safe.
model_id = self._model_id
if not model_id:
return None
repo_dir = (
Path(HF_HUB_CACHE)
/ f"models--{GGML_STT_REPOS[model_id].replace('/', '--')}"
/ "blobs"
)
if not repo_dir.is_dir():
return None
etag = self._etag
if etag:
target = repo_dir / f"{etag}.incomplete"
if target.is_file():
return target.stat().st_size
sizes = [p.stat().st_size for p in repo_dir.glob("*.incomplete") if p.is_file()]
return max(sizes) if sizes else None
except Exception:
return None
def start(
self,
model_id: str,
hf_token: Optional[str] = None,
) -> None:
model_id = resolve_ggml_model_id(model_id)
with self._lock:
if self._thread is not None and self._thread.is_alive():
if self._model_id == model_id:
return
raise SttModelIdError(
f"Another GGUF dictation model ('{self._model_id}') is still "
"downloading; wait for it to finish."
)
self._model_id = model_id
self._error = None
self._total_bytes = None
self._etag = None
thread = threading.Thread(target = self._run, args = (model_id, hf_token), daemon = True)
self._thread = thread
thread.start()
def _run(self, model_id: str, hf_token: Optional[str]) -> None:
repo_id = GGML_STT_REPOS[model_id]
filename = GGML_STT_MODELS[model_id]
try:
from huggingface_hub import (
get_hf_file_metadata,
hf_hub_download,
hf_hub_url,
)
try:
# One HEAD request for the total and etag.
meta = get_hf_file_metadata(hf_hub_url(repo_id, filename), token = hf_token or None)
with self._lock:
self._total_bytes = meta.size
self._etag = meta.etag
except Exception:
pass
hf_hub_download(
repo_id = repo_id,
filename = filename,
token = hf_token or None,
)
except Exception as exc:
logger.warning("GGUF STT download failed for %s: %s", model_id, exc)
with self._lock:
self._error = f"Download failed for '{model_id}'."
_download_state = _GgmlDownloadState()
def start_model_download(model: Optional[str], hf_token: Optional[str] = None) -> None:
_download_state.start(resolve_ggml_model_id(model), hf_token)
def download_status() -> dict:
return _download_state.status()
# ---------------------------------------------------------------------------
# WAV packaging
# ---------------------------------------------------------------------------
def _pcm_to_wav_bytes(decoded_audio) -> bytes:
"""Wrap decoded float32 mono 16 kHz PCM into an in-memory 16-bit WAV."""
import numpy as np
clipped = np.clip(decoded_audio, -1.0, 1.0)
pcm16 = (clipped * 32767.0).astype("<i2")
buf = io.BytesIO()
with wave.open(buf, "wb") as w:
w.setnchannels(1)
w.setsampwidth(2)
w.setframerate(_TARGET_SAMPLE_RATE)
w.writeframes(pcm16.tobytes())
return buf.getvalue()
# ---------------------------------------------------------------------------
# Sidecar
# ---------------------------------------------------------------------------
class GgmlSttSidecar:
"""Owns one whisper-server subprocess and proxies dictation to it."""
def __init__(self, keep_alive_seconds: float = STT_KEEP_ALIVE_SECONDS) -> None:
self._lock = threading.RLock()
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_id: Optional[str] = None
self._idle_timer: Optional[threading.Timer] = None
self._idle_generation = 0
self._keep_alive_seconds = keep_alive_seconds
# Set while whisper-server starts so training admission can account for
# the accelerator memory it is about to bind. Read without the lock.
self._loading = False
# A still-starting whisper-server is cancellable so training can preempt
# it before it binds accelerator memory. Assigned inside self._lock but
# acted on without it: cancel_pending_load() runs while load() holds the
# lock, so the event is the source of truth and terminating the process
# is a best-effort fast path.
self._load_cancel_event: Optional[threading.Event] = None
self._starting_process: Optional[subprocess.Popen] = None
# Set before the updater waits for _lock, then kept set while it owns
# the lock and atomically replaces the managed install tree. New loads
# fail fast instead of starting a process from files being swapped.
self._update_in_progress = False
@property
def loaded_model(self) -> Optional[str]:
# Lock-free status read (like stt_sidecar.py): transcribe() holds
# self._lock for the whole inference call (up to
# _TRANSCRIBE_TIMEOUT_SECONDS), and status polls plus training admission
# must not block behind it. _process_alive() snapshots self._process
# before poll(), which subprocess guards with _waitpid_lock, so a
# concurrent unload is safe.
return self._model_id if self._process_alive() else None
@property
def device(self) -> Optional[str]:
return "whisper.cpp" if self._process_alive() else None
def is_loading(self) -> bool:
# True only while whisper-server is starting (seconds to bind its GPU
# backend); load() sets and clears the flag around that window.
return self._loading
@property
def keep_alive_seconds(self) -> float:
return self._keep_alive_seconds
def _process_alive(self) -> bool:
# Snapshot self._process once: a concurrent unload() nulls it under the
# lock, so lock-free readers would otherwise re-read None between the
# truthiness check and .poll().
process = self._process
return process is not None and process.poll() is None
# -- idle unload ------------------------------------------------------
def _cancel_idle_unload_locked(self) -> None:
self._idle_generation += 1
if self._idle_timer is not None:
self._idle_timer.cancel()
self._idle_timer = None
def _schedule_idle_unload_locked(self) -> None:
self._cancel_idle_unload_locked()
if not self._process_alive():
return
generation = self._idle_generation
timer = threading.Timer(self._keep_alive_seconds, self._idle_unload, args = (generation,))
timer.daemon = True
self._idle_timer = timer
timer.start()
def _idle_unload(self, generation: int) -> None:
with self._lock:
if generation != self._idle_generation:
return
logger.info("Unloading idle GGUF STT model %s", self._model_id)
self._release_locked()
# -- process lifecycle -------------------------------------------------
def _release_locked(self) -> None:
self._cancel_idle_unload_locked()
process = self._process
self._process = None
self._port = None
self._model_id = None
if process is not None and process.poll() is None:
process.terminate()
try:
process.wait(timeout = 10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout = 10)
if process is not None:
forget_pid(process.pid)
def unload(self) -> None:
with self._lock:
self._release_locked()
def _raise_if_update_in_progress(self) -> None:
if self._update_in_progress:
raise SttEngineUnavailableError(
"The local transcription runtime is being updated. Try dictation again shortly."
)
@contextmanager
def update_maintenance(self) -> Iterator[bool]:
"""Block new loads while the managed whisper.cpp tree is replaced.
The flag is published before waiting for an existing transcription to
release ``_lock``. Holding that lock across the yielded installer phase
prevents Windows from relocking the executable and prevents every host
from starting a process against a partially swapped tree. The yielded
value records whether a warm model had to be unloaded.
"""
self._update_in_progress = True
try:
with self._lock:
model_was_active = self._process_alive()
self._release_locked()
yield model_was_active
finally:
self._update_in_progress = False
def cancel_pending_load(self) -> bool:
# Preempt a starting whisper-server so training does not launch while it
# binds accelerator memory. load() holds self._lock for the whole startup,
# so act without the lock: signal abort and terminate the starting
# process. _wait_for_server observes the event and raises, then load()
# reaps the process and releases the lock.
if not self._loading:
return False
event = self._load_cancel_event
if event is None:
return False
event.set()
process = self._starting_process
if process is not None and process.poll() is None:
try:
process.terminate()
except Exception:
pass
return True
def wait_for_load_to_settle(self) -> None:
# load() holds self._lock across startup and cancel cleanup, so acquiring
# it blocks until a cancelled server is killed, reaped, and its
# accelerator memory released.
with self._lock:
pass
@staticmethod
def _reserve_free_port() -> tuple[socket.socket, int]:
"""Bind an ephemeral port and keep the socket held.
The caller closes the reservation immediately before spawning
whisper-server, shrinking the window in which another local process
could bind the port. SO_REUSEADDR lets the child rebind right after.
"""
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind(("127.0.0.1", 0))
return s, s.getsockname()[1]
def _ensure_model_downloaded(self, model_id: str) -> str:
path = _cached_model_path(model_id)
if path is None:
raise SttModelNotDownloadedError(
f"STT model '{model_id}' (GGUF) is not downloaded. "
"Download it in Settings, then Voice, before loading it."
)
return path
def load(self, model: Optional[str] = None) -> None:
"""Start (or switch) whisper-server for the requested curated model."""
self._raise_if_update_in_progress()
model_id = resolve_ggml_model_id(model)
with self._lock:
self._raise_if_update_in_progress()
binary = ensure_engine_available()
if self._process_alive() and self._model_id == model_id:
self._schedule_idle_unload_locked()
return
model_path = self._ensure_model_downloaded(model_id)
self._release_locked()
reservation, port = self._reserve_free_port()
command = [binary, "-m", model_path, "--host", "127.0.0.1", "--port", str(port)]
marker = _whisper_install_marker(binary)
if _training_active():
# Keep whisper.cpp off the accelerator during training (like the
# Transformers sidecar's CPU choice) so a mid-training dictation
# cannot reclaim the VRAM training just freed.
command.append("--no-gpu")
elif marker is not None and marker.get("backend") == "cpu":
# A deliberate CPU install must stay CPU: the slim wiring links
# every llama ggml backend (including CUDA/ROCm), so without
# this flag a cpu-selected install would still grab the GPU.
command.append("--no-gpu")
logger.info(
"Starting whisper-server for STT model %s on 127.0.0.1:%s",
model_id,
port,
)
cancel_event = threading.Event()
self._load_cancel_event = cancel_event
self._loading = True
try:
# Release the reservation as late as possible: whisper-server
# binds the port moments after this close.
reservation.close()
process = subprocess.Popen(
command,
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
stdin = subprocess.DEVNULL,
# Co-located GPU libs on the loader path (WSL system HIP first),
# secrets scrubbed from the downloaded binary's env.
env = _whisper_server_child_env(binary),
# Die with Studio (Linux PDEATHSIG, Windows job) so a crash
# never orphans a server holding the model.
**child_popen_kwargs(),
)
self._starting_process = process
adopt_pid(process.pid) # terminate_all backstop for graceful exits
try:
self._wait_for_server(process, port, cancel_event)
except Exception:
if process.poll() is None:
process.kill()
process.wait(timeout = 10)
forget_pid(process.pid)
raise
self._process = process
self._port = port
self._model_id = model_id
self._schedule_idle_unload_locked()
finally:
reservation.close() # no-op when already released before spawn
self._loading = False
self._load_cancel_event = None
self._starting_process = None
@staticmethod
def _wait_for_server(
process: subprocess.Popen,
port: int,
cancel_event: Optional[threading.Event] = None,
) -> None:
deadline = time.monotonic() + _SERVER_START_TIMEOUT_SECONDS
while time.monotonic() < deadline:
if cancel_event is not None and cancel_event.is_set():
raise SttLoadCancelledError(
"GGUF STT model loading was cancelled so training could start."
)
if process.poll() is not None:
raise SttEngineUnavailableError(
"The local transcription runtime exited before becoming "
"ready; the model file may be corrupt or unsupported."
)
# Require a whisper-server-specific response twice, with the managed
# child alive around each probe. An arbitrary local process that won
# the bind race would otherwise be mistaken for the sidecar and
# receive the user's microphone audio.
if GgmlSttSidecar._probe_is_whisper_server(process, port) and (
GgmlSttSidecar._probe_is_whisper_server(process, port)
):
return
time.sleep(0.2)
raise SttEngineUnavailableError("The local transcription runtime did not start in time.")
@staticmethod
def _probe_is_whisper_server(process: subprocess.Popen, port: int) -> bool:
"""One readiness probe: our child is alive and the responder looks like
whisper.cpp's server (its index page and errors identify whisper)."""
if process.poll() is not None:
return False
try:
req = urllib.request.Request(f"http://127.0.0.1:{port}/", method = "GET")
with urllib.request.urlopen(req, timeout = 2) as response:
body = response.read(65536)
except Exception:
return False
if process.poll() is not None:
return False
return b"whisper" in body.lower()
# -- transcription ------------------------------------------------------
def transcribe(
self,
audio: bytes,
model: Optional[str] = None,
language: Optional[str] = None,
fast: bool = False,
) -> dict:
"""Transcribe encoded audio bytes via whisper-server.
Accepts any container PyAV can decode (same validation and caps as the
Transformers sidecar). Returns {text, language, duration, model}.
"""
self._raise_if_update_in_progress()
ensure_engine_available()
model_id = resolve_ggml_model_id(model)
lang = normalize_whisper_language(language)
known_languages = _known_whisper_languages()
if lang is not None and known_languages is not None and lang not in known_languages:
raise SttLanguageError(
f"Language '{language}' is not supported by STT model '{model_id}'."
)
# Reject a missing model before decoding so a long clip does not burn CPU
# only to 409 (matches the Transformers sidecar's preflight).
self._ensure_model_downloaded(model_id)
decoded_audio = _decode_audio_bounded(audio)
wav_bytes = _pcm_to_wav_bytes(decoded_audio)
with self._lock:
try:
self.load(model_id)
text = self._post_inference(wav_bytes, lang, fast)
finally:
self._schedule_idle_unload_locked()
duration = (len(decoded_audio) / _TARGET_SAMPLE_RATE) if len(decoded_audio) else None
return {
"text": text,
"language": lang,
"duration": duration,
"model": model_id,
}
def _post_inference(self, wav_bytes: bytes, lang: Optional[str], fast: bool) -> str:
boundary = uuid.uuid4().hex
fields = {
"temperature": "0.0",
"response_format": "json",
# Match the Transformers sidecar: 5-way beam search, greedy for fast.
"beam_size": "1" if fast else "5",
"language": lang or "auto",
}
parts: list[bytes] = []
for name, value in fields.items():
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
f'name="{name}"\r\n\r\n{value}\r\n'
).encode()
)
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
'name="file"; filename="dictation.wav"\r\n'
"Content-Type: audio/wav\r\n\r\n"
).encode()
+ wav_bytes
+ b"\r\n"
)
parts.append(f"--{boundary}--\r\n".encode())
body = b"".join(parts)
req = urllib.request.Request(
f"http://127.0.0.1:{self._port}/inference",
data = body,
headers = {"Content-Type": f"multipart/form-data; boundary={boundary}"},
)
try:
with urllib.request.urlopen(req, timeout = _TRANSCRIBE_TIMEOUT_SECONDS) as resp:
payload = json.load(resp)
except SttAudioDecodeError:
raise
except Exception as exc:
raise SttEngineUnavailableError(
"The local transcription runtime did not answer the request."
) from exc
text = payload.get("text")
if not isinstance(text, str):
raise SttAudioDecodeError("Could not decode the audio.")
# whisper.cpp joins segments with newlines; dictation wants one line.
return " ".join(part.strip() for part in text.splitlines() if part.strip()).strip()
_sidecar: Optional[GgmlSttSidecar] = None
def get_ggml_stt_sidecar() -> GgmlSttSidecar:
global _sidecar
if _sidecar is None:
_sidecar = GgmlSttSidecar()
return _sidecar

1130
studio/backend/core/inference/stt_sidecar.py Normal file
View file

File diff suppressed because it is too large Load diff

3
studio/backend/core/training/training.py
View file

@ -754,6 +754,9 @@ class TrainingBackend:
def __init__(self):
# Subprocess state
self._proc: Optional[mp.Process] = None
# True from the sidecar-swap handshake until the worker is recorded, so
# installs and STT loads treat the startup window as active.
self._spawn_in_progress: bool = False
self._event_queue: Any = None
self._stop_queue: Any = None
self._pump_thread: Optional[threading.Thread] = None

26
studio/backend/hub/services/models/cache_inventory.py
View file

@ -517,6 +517,19 @@ def _read_json_object(path: Path) -> dict:
return {}
def _is_whisper_model_config(config: object) -> bool:
if not isinstance(config, dict):
return False
model_type = config.get("model_type")
if isinstance(model_type, str) and model_type.strip().lower() == "whisper":
return True
architectures = config.get("architectures")
return isinstance(architectures, list) and any(
isinstance(name, str) and name == "WhisperForConditionalGeneration"
for name in architectures
)
def _read_model_card_frontmatter(path: Path) -> dict:
try:
text = path.read_text(encoding = "utf-8")
@ -547,6 +560,8 @@ def _cached_model_local_metadata(repo_path: Path) -> dict:
result: dict = {}
config = _read_json_object(snapshot / "config.json")
if _is_whisper_model_config(config):
result["_hidden_stt"] = True
quant_method = (
config.get("quantization_config", {}).get("quant_method")
if isinstance(config.get("quantization_config"), dict)
@ -581,6 +596,7 @@ def _scan_cached_models() -> list[dict]:
inspected = 0
skipped_gguf = 0
skipped_no_weights = 0
skipped_stt = 0
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
inspected += 1
@ -608,6 +624,10 @@ def _scan_cached_models() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
local_metadata = _cached_model_local_metadata(repo_path)
if local_metadata.pop("_hidden_stt", False):
skipped_stt += 1
continue
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"model",
repo_id,
@ -627,7 +647,7 @@ def _scan_cached_models() -> list[dict]:
if snapshot_partial
else None
),
**_cached_model_local_metadata(repo_path),
**local_metadata,
}
last_modified = max(
payload.last_modified,
@ -655,10 +675,12 @@ def _scan_cached_models() -> list[dict]:
continue
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
logger.info(
"Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d returned=%d",
"Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d "
"skipped_stt=%d returned=%d",
inspected,
skipped_gguf,
skipped_no_weights,
skipped_stt,
len(cached),
)
return cached

25
studio/backend/main.py
View file

@ -309,6 +309,7 @@ from routes import (
training_router,
)
from routes.llama import router as llama_router
from routes.whisper import router as whisper_router
from routes.preview import router as preview_router
from hub.routes import (
inventory_router as hub_inventory_router,
@ -755,6 +756,8 @@ app.add_middleware(SecurityHeadersMiddleware)
# headroom; non-upload routes keep the default body cap.
import json as _json_for_413 # noqa: E402
from utils.upload_limits import ( # noqa: E402
STT_AUDIO_JSON_MAX_BYTES,
STT_AUDIO_RAW_MAX_BYTES,
UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES,
default_request_body_limit_bytes,
upload_request_limit_bytes,
@ -793,6 +796,14 @@ def _get_upload_passthrough_request_max_bytes(path: str) -> int:
return default_request_body_limit_bytes()
def _get_request_body_max_bytes(path: str) -> int:
if path.startswith("/api/inference/audio/transcribe/raw"):
return STT_AUDIO_RAW_MAX_BYTES
if path.startswith("/api/inference/audio/transcribe"):
return STT_AUDIO_JSON_MAX_BYTES
return default_request_body_limit_bytes()
async def _send_411(send) -> None:
payload = _json_for_413.dumps(
{"detail": "Content-Length required for upload requests."},
@ -835,12 +846,14 @@ class MaxBodyMiddleware:
app,
max_bytes_getter,
protected_prefixes: tuple,
request_max_bytes_getter = None,
upload_passthrough_prefixes: tuple = (),
upload_passthrough_max_bytes_getter = None,
):
self.app = app
self.max_bytes_getter = max_bytes_getter
self.protected_prefixes = protected_prefixes
self.request_max_bytes_getter = request_max_bytes_getter
self.upload_passthrough_prefixes = upload_passthrough_prefixes
self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter
@ -857,6 +870,14 @@ class MaxBodyMiddleware:
except Exception:
return int(self.max_bytes_getter())
def _request_max_bytes(self, path: str) -> int:
if self.request_max_bytes_getter is None:
return int(self.max_bytes_getter())
try:
return int(self.request_max_bytes_getter(path))
except Exception:
return int(self.max_bytes_getter())
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.app(scope, receive, send)
@ -869,7 +890,7 @@ class MaxBodyMiddleware:
await self.app(scope, receive, send)
return
max_bytes = int(self.max_bytes_getter())
max_bytes = self._request_max_bytes(path)
declared = None
for name, value in scope.get("headers", []):
if name == b"content-length":
@ -934,6 +955,7 @@ app.add_middleware(
MaxBodyMiddleware,
max_bytes_getter = default_request_body_limit_bytes,
protected_prefixes = _BODY_PROTECTED_PREFIXES,
request_max_bytes_getter = _get_request_body_max_bytes,
upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES,
upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes,
)
@ -992,6 +1014,7 @@ app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"])
app.include_router(whisper_router, prefix = "/api/whisper", tags = ["whisper"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])

26
studio/backend/models/inference.py
View file

@ -187,6 +187,32 @@ class UnloadRequest(BaseModel):
model_path: str = Field(..., description = "Model identifier to unload")
class TranscribeRequest(BaseModel):
"""Speech-to-text request for the dictation STT sidecar."""
audio: str = Field(..., description = "Base64-encoded audio (any common format)")
model: Optional[str] = Field(None, description = "STT model id; defaults server-side")
language: Optional[str] = Field(None, description = "BCP-47 language, or 'auto'/None to detect")
fast: bool = Field(
False,
description = "Use low-latency single-candidate decoding for dictation",
)
engine: Optional[str] = Field(
None,
description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)",
)
class SttLoadRequest(BaseModel):
"""Warm the STT sidecar with a model without transcribing."""
model: Optional[str] = Field(None, description = "STT model id; defaults server-side")
engine: Optional[str] = Field(
None,
description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)",
)
class ValidateModelRequest(BaseModel):
"""Check whether an identifier resolves to a ModelConfig; does NOT load weights."""

1
studio/backend/requirements/extras.txt
View file

@ -10,6 +10,7 @@ omegaconf
einx
pyloudnorm
openai-whisper
av # PyAV: decode dictation audio (webm/opus/mp3/…) for the Whisper STT sidecar
uroman # 4.0 MB - used for Outetts.
MeCab # 19.9 MB - used for Outetts.
inflect # number-to-words, required by OuteTTS

344
studio/backend/routes/inference.py
View file

@ -29,6 +29,8 @@ import re as _re
from utils.models import extract_model_size_b as _extract_model_size_b
from utils.api_errors import openai_error_body, anthropic_error_body
from utils.upload_limits import STT_AUDIO_B64_MAX_CHARS, STT_AUDIO_RAW_MAX_BYTES
from hub.dependencies import get_hf_token
from core.inference.orchestrator import GenStreamError, GenStreamErrorRaised
from core.inference.llama_admission import (
LlamaAdmissionCancelled,
@ -1692,6 +1694,8 @@ async def _aiter_llama_stream_items(
from models.inference import (
LoadRequest,
UnloadRequest,
TranscribeRequest,
SttLoadRequest,
GenerateRequest,
LoadResponse,
LoadProgressResponse,
@ -6113,6 +6117,342 @@ async def generate_audio(
)
# =====================================================================
# Speech-to-text (STT) sidecar (/audio/transcribe, /audio/stt/*)
# =====================================================================
def _resolve_stt_engine(engine: Optional[str]) -> str:
"""Normalize the requested STT engine name; default is Transformers."""
normalized = (engine or "transformers").strip().lower()
if normalized in ("", "transformers", "whisper"):
return "transformers"
if normalized in ("gguf", "ggml", "whisper_cpp", "whisper.cpp"):
return "gguf"
raise HTTPException(
status_code = 422,
detail = f"Unknown STT engine '{engine}'. Use 'transformers' or 'gguf'.",
)
def _resolve_serving_stt_engine(engine: Optional[str]) -> str:
"""Resolve the engine that will actually serve a model.
whisper.cpp (gguf) only accepts curated ids, which Transformers serves too,
so when whisper-server is not installed (the common case: `unsloth studio
update` does not yet build it) fall back to Transformers instead of 501-ing
on every recording. Used for download/load/transcribe; unload targets a
specific engine via _resolve_stt_engine.
"""
resolved = _resolve_stt_engine(engine)
if resolved == "gguf":
from core.inference import stt_ggml_sidecar
if not stt_ggml_sidecar.is_available():
return "transformers"
return resolved
def _stt_sidecar_for(engine: str):
if engine == "gguf":
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
return get_ggml_stt_sidecar()
from core.inference.stt_sidecar import get_stt_sidecar
return get_stt_sidecar()
@studio_router.get("/audio/stt/status")
async def stt_status(
model: Optional[str] = None, current_subject: str = Depends(get_current_subject)
):
"""Report STT availability and which model, if any, is resident.
``model`` extends the Transformers ``downloaded_models`` check to a
custom Hugging Face repository beyond the curated defaults.
"""
from core.inference import stt_ggml_sidecar, stt_sidecar
from core.inference.stt_sidecar import (
DEFAULT_STT_MODEL,
STT_MODELS,
get_stt_sidecar,
is_available,
)
sidecar = get_stt_sidecar()
ggml = stt_ggml_sidecar.get_ggml_stt_sidecar()
transformers_downloaded = [
model_id for model_id in STT_MODELS if stt_sidecar.is_model_downloaded(model_id)
]
if model and model not in STT_MODELS and stt_sidecar.is_model_downloaded(model):
transformers_downloaded.append(model)
return JSONResponse(
content = {
"available": is_available(),
"loaded_model": sidecar.loaded_model,
"loading": sidecar.is_loading(),
"device": sidecar.device,
"keep_alive_seconds": sidecar.keep_alive_seconds,
"default_model": DEFAULT_STT_MODEL,
"models": list(STT_MODELS.keys()),
# Transformers engine, same shape as "gguf" below so clients read
# either generically. Top-level fields above kept for old clients.
"transformers": {
"available": is_available(),
"loaded_model": sidecar.loaded_model,
"loading": sidecar.is_loading(),
"device": sidecar.device,
"keep_alive_seconds": sidecar.keep_alive_seconds,
"default_model": DEFAULT_STT_MODEL,
"models": list(STT_MODELS.keys()),
"downloaded_models": transformers_downloaded,
"download": stt_sidecar.download_status(),
},
# whisper.cpp (GGUF) engine.
"gguf": {
"available": stt_ggml_sidecar.is_available(),
"loaded_model": ggml.loaded_model,
"loading": ggml.is_loading(),
"device": ggml.device,
"keep_alive_seconds": ggml.keep_alive_seconds,
"default_model": stt_ggml_sidecar.DEFAULT_GGML_STT_MODEL,
"models": list(stt_ggml_sidecar.GGML_STT_MODELS.keys()),
"downloaded_models": [
model_id
for model_id in stt_ggml_sidecar.GGML_STT_MODELS
if stt_ggml_sidecar._cached_model_path(model_id) is not None
],
"download": stt_ggml_sidecar.download_status(),
},
}
)
@studio_router.post("/audio/stt/download")
async def stt_download(
payload: SttLoadRequest,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""Start a background download of a dictation model.
Both engines download directly (a GGML checkpoint is a single file the Model
Hub's GGUF variant planner cannot express; a Transformers checkpoint is a
whole snapshot). Progress is reported by /audio/stt/status.
"""
from core.inference import stt_ggml_sidecar, stt_sidecar
from core.inference.stt_sidecar import (
SttModelCompatibilityError,
SttModelIdError,
validate_remote_model,
)
engine = _resolve_serving_stt_engine(payload.engine)
module = stt_ggml_sidecar if engine == "gguf" else stt_sidecar
try:
# Transformers accepts custom `owner/model` repos, so confirm the repo is
# a Whisper checkpoint (metadata-only) before snapshot_download pulls a
# possibly-large non-STT repo into the shared cache. Curated ids
# short-circuit; GGUF only accepts curated ids, so it needs no check.
if engine != "gguf":
validated = await asyncio.to_thread(validate_remote_model, payload.model, hf_token)
# Pin the download to the commit that was just validated so the
# repo cannot be swapped between validation and snapshot_download.
await asyncio.to_thread(
module.start_model_download,
payload.model,
hf_token,
validated.get("revision"),
)
else:
await asyncio.to_thread(module.start_model_download, payload.model, hf_token)
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
return JSONResponse(content = module.download_status())
@studio_router.post("/audio/stt/load")
async def stt_load(payload: SttLoadRequest, current_subject: str = Depends(get_current_subject)):
"""Load the selected STT model after the user starts local dictation."""
from core.inference.stt_sidecar import (
SttLoadCancelledError,
SttModelCompatibilityError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
get_stt_sidecar,
)
sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(payload.engine))
try:
await asyncio.to_thread(sidecar.load, payload.model)
except SttModelNotDownloadedError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttUnavailableError as e:
raise HTTPException(status_code = 501, detail = str(e))
except SttLoadCancelledError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
except Exception as e:
logger.error(f"STT load error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
return JSONResponse(content = {"loaded_model": sidecar.loaded_model, "device": sidecar.device})
@studio_router.post("/audio/stt/validate")
async def stt_validate(
payload: SttLoadRequest,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""Verify a Hub repository is a Whisper checkpoint before downloading it."""
from core.inference.stt_sidecar import (
SttModelCompatibilityError,
SttModelIdError,
validate_remote_model,
)
try:
result = await asyncio.to_thread(validate_remote_model, payload.model, hf_token)
except (SttModelIdError, SttModelCompatibilityError) as e:
raise HTTPException(status_code = 422, detail = str(e))
return JSONResponse(content = result)
@studio_router.post("/audio/stt/unload")
async def stt_unload(
engine: Optional[str] = None, current_subject: str = Depends(get_current_subject)
):
"""Release the local STT model when dictation is idle.
Without an engine, both sidecars unload so an engine switch in Voice
settings always frees whichever backend was resident.
"""
if engine is None:
engines = ["transformers", "gguf"]
else:
# Use the serving resolver: a "gguf" pick without whisper-server is
# actually served by the Transformers fallback, so unload must target
# that same engine or the resident model is never freed.
engines = [_resolve_serving_stt_engine(engine)]
# Attempt every engine even if one raises, so failing to unload one never
# skips freeing the other (both can be resident after a switch).
failed: list[str] = []
for name in engines:
try:
await asyncio.to_thread(_stt_sidecar_for(name).unload)
except Exception as exc: # noqa: BLE001 - report after attempting all engines
logger.warning("Failed to unload STT engine '%s': %s", name, exc)
failed.append(name)
if failed:
raise HTTPException(
status_code = 500,
detail = f"Failed to unload STT engine(s): {', '.join(failed)}",
)
return JSONResponse(content = {"loaded_model": None, "device": None})
async def _transcribe_audio_bytes(
raw: bytes,
model: Optional[str],
language: Optional[str],
fast: bool,
engine: Optional[str] = None,
) -> JSONResponse:
"""Run STT for already-decoded request bytes."""
from core.inference.stt_sidecar import (
SttAudioDecodeError,
SttAudioTooLongError,
SttLanguageError,
SttLoadCancelledError,
SttModelCompatibilityError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
)
if not raw:
raise HTTPException(status_code = 400, detail = "Audio is empty.")
if len(raw) > _MAX_AUDIO_RAW_BYTES:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(engine))
try:
result = await asyncio.to_thread(
sidecar.transcribe,
raw,
model,
language,
fast,
)
except SttUnavailableError as e:
raise HTTPException(status_code = 501, detail = str(e))
except SttLoadCancelledError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelNotDownloadedError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttLanguageError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttAudioTooLongError as e:
raise HTTPException(status_code = 413, detail = str(e))
except SttAudioDecodeError as e:
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
logger.error(f"Transcription error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
return JSONResponse(content = result)
@studio_router.post("/audio/transcribe")
async def transcribe_audio(
payload: TranscribeRequest, current_subject: str = Depends(get_current_subject)
):
"""Transcribe dictation audio to text via the STT sidecar.
Runs alongside the chat model without evicting it, so any model (including
text-only ones) can be driven by voice.
"""
b64 = payload.audio or ""
if not b64:
raise HTTPException(status_code = 400, detail = "No audio provided.")
if len(b64) > _MAX_AUDIO_B64_CHARS:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
try:
raw = base64.b64decode(b64, validate = True)
except Exception:
raise HTTPException(status_code = 400, detail = "Audio is not valid base64.")
return await _transcribe_audio_bytes(
raw, payload.model, payload.language, payload.fast, payload.engine
)
@studio_router.post("/audio/transcribe/raw")
async def transcribe_audio_raw(
request: Request,
model: Optional[str] = None,
language: Optional[str] = None,
fast: bool = False,
engine: Optional[str] = None,
current_subject: str = Depends(get_current_subject),
):
"""Transcribe a raw audio body without base64 or JSON conversion overhead."""
chunks: list[bytes] = []
size = 0
async for chunk in request.stream():
size += len(chunk)
if size > _MAX_AUDIO_RAW_BYTES:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
chunks.append(chunk)
return await _transcribe_audio_bytes(b"".join(chunks), model, language, fast, engine)
# =====================================================================
# OpenAI-Compatible Chat Completions (/chat/completions)
# =====================================================================
@ -6154,8 +6494,8 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally
# bounds the *decoded* length, since a small compressed file (opus/flac/etc.)
# can expand to a far larger PCM array than the encoded-size cap implies.
_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024
_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3
_MAX_AUDIO_RAW_BYTES = STT_AUDIO_RAW_MAX_BYTES
_MAX_AUDIO_B64_CHARS = STT_AUDIO_B64_MAX_CHARS
_MAX_AUDIO_SECONDS = 30 * 60
_WAV_HEADER_BYTES = 44
_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000

51
studio/backend/routes/llama.py
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""llama.cpp prebuilt update endpoints.
"""llama.cpp prebuilt update endpoints -- the single main update item.
GET /api/llama/update-status -> is a newer prebuilt available + job state
POST /api/llama/update -> download + atomically swap to the latest
@ -9,13 +9,19 @@ POST /api/llama/update -> download + atomically swap to the latest
Detection reuses utils.llama_cpp_freshness; the swap reuses
install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI
never blocks on a missing marker / offline GitHub.
whisper.cpp updates piggyback here: the status payload carries a whisper
sub-status (update_available is the llama OR whisper union) and the apply job
chains a whisper phase after the llama phase when whisper is behind, with a
per-phase breakdown in job.phases. All pre-existing top-level fields keep
their shape, so older clients keep working unchanged.
"""
from __future__ import annotations
import asyncio
import threading
from typing import Optional
from typing import Literal, Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
@ -38,6 +44,31 @@ class LlamaUpdateJob(BaseModel):
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
phases: Optional[dict] = Field(
None,
description = (
"Per-phase breakdown of a chained llama+whisper job "
"(name -> state/progress/to_tag/...); None for pre-chaining jobs."
),
)
class WhisperSubStatus(BaseModel):
"""The whisper piggyback inside the llama update item."""
update_available: bool = Field(
False, description = "True when the chained apply would run a whisper phase."
)
installed_tag: Optional[str] = None
latest_tag: Optional[str] = None
update_size_bytes: Optional[int] = None
skip_reason: Optional[str] = Field(
None,
description = (
"Why the whisper phase would be skipped "
"(up_to_date | local_link | source_build | not_installed | ...)."
),
)
class LlamaUpdateStatusResponse(BaseModel):
@ -46,7 +77,18 @@ class LlamaUpdateStatusResponse(BaseModel):
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
False,
description = (
"True when an update would do something: llama.cpp is behind OR the "
"whisper piggyback is behind."
),
)
llama_update_available: bool = Field(
False, description = "True when the latest llama.cpp release is newer than the install."
)
update_component: Optional[Literal["llama", "whisper"]] = Field(
None,
description = "Component whose versions the combined update banner should display.",
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
@ -62,6 +104,9 @@ class LlamaUpdateStatusResponse(BaseModel):
update_size_bytes: Optional[int] = Field(
None, description = "Download size of the prebuilt Update would fetch, in bytes."
)
whisper: Optional[WhisperSubStatus] = Field(
None, description = "Whisper piggyback sub-status; None when the probe is unavailable."
)
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)

8
studio/backend/routes/models.py
View file

@ -3026,7 +3026,9 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
if repo_info.repo_type != "model":
continue
repo_id = repo_info.repo_id
if _is_hidden_model(repo_id):
# Pass the snapshot path too so the config check also hides
# custom Whisper checkpoints, not just curated repo ids.
if _is_hidden_model(repo_id, str(repo_info.repo_path)):
continue
total_size = _repo_gguf_size_bytes(repo_info)
if total_size == 0:
@ -3083,7 +3085,9 @@ async def list_cached_models(
if repo_info.repo_type != "model":
continue
repo_id = repo_info.repo_id
if _is_hidden_model(repo_id):
# Pass the snapshot path too so the config check also hides
# custom Whisper checkpoints, not just curated repo ids.
if _is_hidden_model(repo_id, str(repo_info.repo_path)):
continue
if _repo_has_gguf_files(repo_info):
continue

55
studio/backend/routes/training.py
View file

@ -415,46 +415,29 @@ async def start_training(
try:
from routes.training_vram import (
can_keep_chat_during_training,
free_chat_models_for_training,
summarize_resident_chat,
coordinate_models_for_training,
)
resident = summarize_resident_chat()
if not resident["any"]:
return
if resident.get("loading"):
# In-flight load can't be sized -> free rather than risk OOM.
freed = free_chat_models_for_training(reason = "chat model still loading")
logger.info("Freed in-flight chat load for training: %s", freed)
return
keep, info = can_keep_chat_during_training(
model_name = training_kwargs["model_name"],
hf_token = training_kwargs["hf_token"],
training_type = training_kwargs["training_type"],
load_in_4bit = training_kwargs["load_in_4bit"],
batch_size = training_kwargs["batch_size"],
max_seq_length = training_kwargs["max_seq_length"],
lora_rank = training_kwargs["lora_r"],
target_modules = training_kwargs["target_modules"],
gradient_checkpointing = training_kwargs["gradient_checkpointing"],
optimizer = training_kwargs["optim"],
gpu_ids = training_kwargs["gpu_ids"],
)
if keep:
logger.info(
"Keeping chat model(s) loaded during training "
"(free ~%s GB, needs ~%s GB): %s",
info.get("usable_gb"),
info.get("required_gb"),
resident,
def _can_keep_resident_models():
return can_keep_chat_during_training(
model_name = training_kwargs["model_name"],
hf_token = training_kwargs["hf_token"],
training_type = training_kwargs["training_type"],
load_in_4bit = training_kwargs["load_in_4bit"],
batch_size = training_kwargs["batch_size"],
max_seq_length = training_kwargs["max_seq_length"],
lora_rank = training_kwargs["lora_r"],
target_modules = training_kwargs["target_modules"],
gradient_checkpointing = training_kwargs["gradient_checkpointing"],
optimizer = training_kwargs["optim"],
gpu_ids = training_kwargs["gpu_ids"],
)
else:
freed = free_chat_models_for_training(
reason = "insufficient VRAM to run training alongside chat",
)
logger.info("Freed chat model(s) for training: %s", freed)
freed = coordinate_models_for_training(_can_keep_resident_models)
if freed:
logger.info("Freed models for training: %s", freed)
except Exception as e:
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
logger.warning("Inference/training memory coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
from utils.transformers_version import SidecarSwapInProgress

148
studio/backend/routes/training_vram.py
View file

@ -1,15 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""VRAM coordination between chat/inference and training.
"""Memory coordination between inference and training.
Decides, from live free VRAM, whether a resident chat model can stay loaded
during training or must be unloaded, and unloads it across all backends
(HF/MLX orchestrator + llama.cpp GGUF server). In the route layer because the
GGUF accessor lives in routes/inference.py; backends are imported lazily.
Uses live free VRAM to keep resident chat and STT models when they fit. STT is
evicted before chat when training needs memory.
"""
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Callable, Dict, List, Optional, Tuple
from loggers import get_logger
@ -77,6 +75,37 @@ def summarize_resident_chat() -> Dict[str, Any]:
}
def summarize_resident_stt() -> Dict[str, Any]:
"""Report the resident dictation model (either engine). Never raises."""
try:
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
from core.inference.stt_sidecar import get_stt_sidecar
sidecar = get_stt_sidecar()
model = sidecar.loaded_model
device = sidecar.device
loading = sidecar.is_loading()
# whisper.cpp holds GPU memory via its subprocess, and both engines can be
# live at once (engine switch or direct /audio/stt/load). Always fold the
# GGUF sidecar in: a resident Transformers model must not mask a GGUF
# server still binding its backend, or admission lets training launch into
# that startup and OOM.
ggml = get_ggml_stt_sidecar()
if not model:
model = ggml.loaded_model
device = device or ggml.device
loading = loading or ggml.is_loading()
return {
"model": model,
"device": device,
"loading": loading,
"any": bool(model or loading),
}
except Exception as e:
logger.warning("Could not inspect STT sidecar: %s", e)
return {"model": None, "device": None, "loading": False, "any": False}
def can_keep_chat_during_training(
*,
model_name: str,
@ -366,3 +395,110 @@ def free_chat_models_for_training(reason: str) -> List[str]:
logger.warning("Could not unload GGUF chat model: %s", e)
return freed
def free_stt_model_for_training(reason: str) -> List[str]:
"""Unload the dictation model(s) before training. Never raises.
The Transformers and GGUF sidecars are freed under independent exception
boundaries so a failure unloading one backend never skips freeing the other
(both can hold accelerator memory at once after an engine switch).
"""
freed: List[str] = []
try:
from core.inference.stt_sidecar import get_stt_sidecar
sidecar = get_stt_sidecar()
if sidecar.is_loading() and sidecar.cancel_pending_load():
logger.info("Cancelling STT model load for training (%s)", reason)
# The loader may still be in from_pretrained()/.to(device) holding
# VRAM; wait for it to observe the cancel and release first.
sidecar.wait_for_load_to_settle()
# A load that finished before seeing the cancel leaves a resident
# model; unload it so training gets the memory back.
if sidecar.loaded_model:
sidecar.unload()
freed.append("stt:loading")
else:
model = sidecar.loaded_model
if model:
logger.info("Unloading STT model '%s' for training (%s)", model, reason)
sidecar.unload()
freed.append(f"stt:{model}")
except Exception as e:
logger.warning("Could not unload Transformers STT model: %s", e)
# Check the GGUF sidecar even after a cancelled/failed Transformers unload;
# both engines can hold memory at once (engine switch or direct load).
try:
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
ggml = get_ggml_stt_sidecar()
if ggml.is_loading() and ggml.cancel_pending_load():
logger.info("Cancelling GGUF STT model load for training (%s)", reason)
# whisper-server may still be binding its backend; wait for the
# cancelled startup to be killed and reaped before training claims
# the memory (loaded_model stays unset until it is ready).
ggml.wait_for_load_to_settle()
if ggml.loaded_model:
ggml.unload()
freed.append("stt:gguf-loading")
else:
ggml_model = ggml.loaded_model
if ggml_model:
logger.info("Unloading GGUF STT model '%s' for training (%s)", ggml_model, reason)
ggml.unload()
freed.append(f"stt:{ggml_model}")
except Exception as e:
logger.warning("Could not unload GGUF STT model: %s", e)
return freed
def coordinate_models_for_training(
can_keep: Callable[[], Tuple[bool, Dict[str, Any]]],
) -> List[str]:
"""Keep resident models when they fit, evicting STT before chat."""
resident_chat = summarize_resident_chat()
resident_stt = summarize_resident_stt()
if not resident_chat["any"] and not resident_stt["any"]:
return []
if resident_chat.get("loading"):
freed = free_stt_model_for_training(reason = "chat model still loading")
freed += free_chat_models_for_training(reason = "chat model still loading")
return freed
freed: List[str] = []
if resident_stt.get("loading"):
released_stt = free_stt_model_for_training(reason = "STT model still loading")
freed += released_stt
resident_stt = (
{"model": None, "device": None, "loading": False, "any": False}
if released_stt
else summarize_resident_stt()
)
if not resident_chat["any"] and not resident_stt["any"]:
return freed
keep, info = can_keep()
if keep:
logger.info(
"Keeping resident models loaded during training (free ~%s GB, needs ~%s GB): %s",
info.get("usable_gb"),
info.get("required_gb"),
{"chat": resident_chat, "stt": resident_stt},
)
return freed
if resident_stt["any"]:
freed += free_stt_model_for_training(reason = "insufficient training memory")
if not resident_chat["any"]:
return freed
keep, _info = can_keep()
if keep:
logger.info("Keeping chat model loaded after freeing STT: %s", resident_chat)
return freed
freed += free_chat_models_for_training(
reason = "insufficient VRAM to run training alongside chat",
)
return freed

74
studio/backend/routes/whisper.py Normal file
View file

@ -0,0 +1,74 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp prebuilt status endpoint.
GET /api/whisper/update-status -> is a newer prebuilt available + job state
Detection reuses utils.whisper_cpp_freshness and fails open so the UI never
blocks on a missing marker / offline GitHub. There is no whisper-only update
trigger: whisper updates piggyback on the single main update item
(POST /api/llama/update chains a whisper phase when whisper is behind).
"""
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject
from utils.whisper_cpp_update import get_update_status
router = APIRouter()
class WhisperUpdateJob(BaseModel):
state: str = Field("idle", description = "idle | running | success | error")
message: str = ""
from_tag: Optional[str] = None
to_tag: Optional[str] = None
reload_required: Optional[bool] = None
error: Optional[str] = None
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
class WhisperUpdateStatusResponse(BaseModel):
supported: bool = Field(
False,
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
)
installed_tag: Optional[str] = None
latest_tag: Optional[str] = None
published_repo: Optional[str] = None
installed_at_utc: Optional[str] = None
age_days: Optional[int] = None
source_build: bool = Field(
False, description = "True when there is no marker (source build) but a prebuilt is offered."
)
update_size_bytes: Optional[int] = Field(
None, description = "Download size of the prebuilt an update would fetch, in bytes."
)
job: WhisperUpdateJob = Field(default_factory = WhisperUpdateJob)
@router.get("/update-status", response_model = WhisperUpdateStatusResponse)
async def whisper_update_status(
force_refresh: bool = Query(
False, description = "Bypass the 24h release cache for an explicit check."
),
current_subject: str = Depends(get_current_subject),
) -> WhisperUpdateStatusResponse:
# Off the event loop: detection may probe the host and read GitHub.
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
return WhisperUpdateStatusResponse(**status)

66
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -191,6 +191,72 @@ def test_is_hidden_model_hides_validation_probe_everywhere():
assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF")
def test_is_hidden_model_hides_dictation_models(tmp_path):
assert models_route._is_hidden_model("unsloth/whisper-tiny")
assert models_route._is_hidden_model("unsloth/whisper-base")
assert models_route._is_hidden_model("unsloth/whisper-small")
assert models_route._is_hidden_model("unsloth/whisper-large-v3-turbo")
assert models_route._is_hidden_model(
"/hf/models--unsloth--whisper-large-v3/snapshots/abc/model.safetensors"
)
assert not models_route._is_hidden_model("user/whisper-finetune")
assert not models_route._is_hidden_model(
"C:\\cache\\models--unsloth--whisper-small-finetune\\model.safetensors"
)
custom = tmp_path / "custom-whisper"
custom.mkdir()
(custom / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
(custom / "model.safetensors").write_bytes(b"weights")
assert models_route._is_hidden_model(
"user/custom-checkpoint",
str(custom / "model.safetensors"),
)
named_only = tmp_path / "whisper-finetune"
named_only.mkdir()
(named_only / "config.json").write_text('{"model_type": "llama"}')
assert not models_route._is_hidden_model("user/whisper-finetune", str(named_only))
def test_list_cached_models_hides_custom_whisper_by_config(monkeypatch, tmp_path):
# Regression: the legacy /cached-models picker must pass the snapshot path so
# the config check hides a custom (non-curated) Whisper checkpoint; a bare
# repo id cannot ("user/whisper-finetune" is not in the curated set).
repo_path = tmp_path / "models--user--whisper-finetune"
snap = repo_path / "snapshots" / "abc"
snap.mkdir(parents = True)
(snap / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
(snap / "model.safetensors").write_bytes(b"weights")
captured: list = []
real_hidden = models_route._is_hidden_model
def spy(*values):
captured.append(values)
return real_hidden(*values)
monkeypatch.setattr(models_route, "_is_hidden_model", spy)
repo = _repo(
"user/whisper-finetune",
[SimpleNamespace(file_name = "model.safetensors", size_on_disk = 10)],
repo_path,
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
result = asyncio.run(
models_route.list_cached_models(current_subject = "test-user", hf_token = None)
)
# The route passed the snapshot path (not just the repo id) ...
assert any(str(repo_path) in values for values in captured)
# ... so the custom Whisper checkpoint is hidden from the chat picker.
assert result["cached"] == []
def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch):
"""A custom embedder with a generic basename is hidden by EXACT repo-id
match only, so unrelated cached repos that merely contain the basename stay

735
studio/backend/tests/test_combined_update.py Normal file
View file

@ -0,0 +1,735 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hermetic tests for the combined llama+whisper update item.
llama.cpp is the single main update item; whisper.cpp piggybacks on it. These
pin the union status (update_available = llama behind OR whisper behind), the
chained apply (llama phase first, whisper phase only when behind), the failure
policy (llama failure aborts; whisper failure keeps the llama partial success),
the silent whisper skips, and the backward-compatible payload shape.
"""
from __future__ import annotations
import json
import sys
import time
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import utils.llama_cpp_freshness as freshness # noqa: E402
import utils.llama_cpp_update as upd # noqa: E402
import utils.whisper_cpp_freshness as wfresh # noqa: E402
import utils.whisper_cpp_update as wupd # noqa: E402
MARKER = "UNSLOTH_PREBUILT_INFO.json"
WHISPER_MARKER = "UNSLOTH_WHISPER_PREBUILT_INFO.json"
# The top-level status and job fields that predate the whisper piggyback; the
# combined payload must stay an exact superset so current UI code keeps working.
LEGACY_STATUS_FIELDS = {
"supported",
"update_available",
"stale",
"installed_tag",
"latest_tag",
"published_repo",
"installed_at_utc",
"age_days",
"source_build",
"update_size_bytes",
"job",
}
LEGACY_JOB_FIELDS = {
"state",
"message",
"from_tag",
"to_tag",
"reload_required",
"error",
"progress",
"started_at",
"finished_at",
}
class _FakeInstallerPopen:
"""Stands in for the streamed llama installer process."""
def __init__(
self,
cmd,
*,
returncode = 0,
lines = None,
on_start = None,
**kwargs,
):
if on_start is not None:
on_start(list(cmd))
self.returncode = returncode
self.stdout = iter(lines or [])
def wait(self):
return self.returncode
def kill(self):
pass
def _patch_llama_installer(
monkeypatch,
*,
returncode = 0,
lines = None,
on_start = None,
):
# Only intercept the installer invocation: importing routes.inference inside
# the worker can Popen unrelated host probes (ldconfig etc).
def _popen(cmd, **kw):
is_installer = any("install_llama_prebuilt" in str(part) for part in cmd)
return _FakeInstallerPopen(
cmd,
returncode = returncode if is_installer else 0,
lines = lines if is_installer else None,
on_start = on_start if is_installer else None,
)
monkeypatch.setattr(upd.subprocess, "Popen", _popen)
def _write_llama_install(dir_: Path, tag: str) -> str:
"""Create a fake llama prebuilt install and return the llama-server path."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "llama-server"
binary.write_text("stub")
(dir_ / MARKER).write_text(
json.dumps(
{
"tag": tag,
"release_tag": tag,
"published_repo": "unslothai/llama.cpp",
"installed_at_utc": "2020-01-01T00:00:00Z",
}
)
)
return str(binary)
def _write_whisper_install(
dir_: Path,
tag: str,
backend: str = "cpu",
) -> str:
"""Create a fake whisper prebuilt install and return the whisper-server path."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "whisper-server"
binary.write_text("stub")
(dir_ / WHISPER_MARKER).write_text(
json.dumps(
{
"release_tag": tag,
"upstream_tag": tag.split("-")[0],
"published_repo": "unslothai/whisper.cpp",
"backend": backend,
"installed_at_utc": "2020-01-01T00:00:00Z",
}
)
)
return str(binary)
@pytest.fixture(autouse = True)
def _clean_state(monkeypatch, tmp_path):
freshness.reset_caches()
wfresh.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
wupd._resolve_memo.clear()
monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".llama_cache")
monkeypatch.setattr(wfresh, "_cache_dir", lambda: tmp_path / ".whisper_cache")
for var in (
"LLAMA_SERVER_PATH",
"UNSLOTH_LLAMA_CPP_PATH",
"WHISPER_SERVER_PATH",
"UNSLOTH_WHISPER_CPP_PATH",
):
monkeypatch.delenv(var, raising = False)
# Never hit the network in these tests.
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
yield
freshness.reset_caches()
wfresh.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
wupd._resolve_memo.clear()
def _setup_llama(
monkeypatch,
tmp_path,
*,
installed = "b9493",
latest = "b9518",
):
"""Marker-managed llama install; behind when installed != latest."""
install_dir = tmp_path / "llama.cpp"
binary = _write_llama_install(install_dir, installed)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest)
return install_dir
def _setup_whisper(
monkeypatch,
tmp_path,
*,
installed = "v1.9.1-unsloth.1",
latest = "v1.9.2-unsloth.1",
):
"""Marker-managed whisper install; behind when latest is newer."""
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, installed)
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
monkeypatch.setattr(wupd, "_installer_script", lambda: tmp_path / "install_whisper_prebuilt.py")
monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest)
return install_dir
def _patch_whisper_phase(
monkeypatch,
events,
*,
to_tag = "v1.9.2-unsloth.1",
error = None,
):
"""Record whisper phase runs without touching a real installer."""
def _run(phase, set_progress):
events.append("whisper")
if error is not None:
raise RuntimeError(error)
set_progress(0.5)
return {
"to_tag": to_tag,
"reload_required": False,
"message": f"Updated whisper.cpp to {to_tag}.",
}
monkeypatch.setattr(wupd, "run_chained_phase", _run)
def _wait_for_job():
deadline = time.time() + 10
while time.time() < deadline:
with upd._job_lock:
job = dict(upd._job)
if job["state"] in ("success", "error"):
return job
time.sleep(0.05)
with upd._job_lock:
return dict(upd._job)
# --- status: the single item folds whisper in ---
def test_status_payload_is_exact_superset_of_legacy_fields(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
st = upd.get_update_status(force_refresh = True)
assert LEGACY_STATUS_FIELDS <= set(st)
assert LEGACY_JOB_FIELDS <= set(st["job"])
# The new fields ride alongside, never replacing the legacy ones.
assert st["llama_update_available"] is True
assert st["whisper"]["update_available"] is True
assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1"
assert st["update_component"] == "llama"
def test_status_union_whisper_only_surfaces_update(monkeypatch, tmp_path):
# llama current, whisper behind: the single item still shows an update.
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
st = upd.get_update_status(force_refresh = True)
assert st["llama_update_available"] is False
assert st["whisper"]["update_available"] is True
assert st["update_available"] is True
assert st["update_component"] == "whisper"
assert st["installed_tag"] == "b9518"
assert st["latest_tag"] == "b9518"
assert st["whisper"]["installed_tag"] == "v1.9.1-unsloth.1"
assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1"
def test_status_whisper_current_does_not_flip_union(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
st = upd.get_update_status(force_refresh = True)
assert st["update_available"] is False
assert st["whisper"]["skip_reason"] == "up_to_date"
assert st["update_component"] is None
def test_status_survives_whisper_probe_failure(monkeypatch, tmp_path):
# The piggyback fails open: llama status still works without a whisper probe.
_setup_llama(monkeypatch, tmp_path)
def _boom(*, force_refresh = False):
raise RuntimeError("probe exploded")
monkeypatch.setattr(wupd, "chained_phase_plan", _boom)
st = upd.get_update_status(force_refresh = True)
assert st["update_available"] is True
assert st["whisper"] is None
# --- whisper chained_phase_plan: silent skips ---
def test_whisper_plan_skips_local_link(monkeypatch, tmp_path):
monkeypatch.setattr(wupd, "_find_binary", lambda: str(tmp_path / "whisper-server"))
monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True)
plan = wupd.chained_phase_plan()
assert plan["update_available"] is False
assert plan["skip_reason"] == "local_link"
assert plan["phase"] is None
def test_whisper_plan_skips_source_build(monkeypatch, tmp_path):
binary = tmp_path / "whisper.cpp" / "build" / "bin" / "whisper-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker
monkeypatch.setattr(wupd, "_find_binary", lambda: str(binary))
plan = wupd.chained_phase_plan()
assert plan["skip_reason"] == "source_build"
assert plan["phase"] is None
def test_whisper_update_targets_canonical_root_when_inner_marker_exists(tmp_path):
install_dir = tmp_path / "whisper.cpp"
binary = install_dir / "build" / "bin" / "whisper-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
(install_dir / WHISPER_MARKER).write_text("{}")
(binary.parent / WHISPER_MARKER).write_text("{}")
assert wupd._install_dir_for(str(binary)) == install_dir
def test_whisper_plan_skips_when_not_installed(monkeypatch):
monkeypatch.setattr(wupd, "_find_binary", lambda: None)
plan = wupd.chained_phase_plan()
assert plan["skip_reason"] == "not_installed"
assert plan["phase"] is None
def test_whisper_plan_eligible_when_behind(monkeypatch, tmp_path):
install_dir = _setup_whisper(monkeypatch, tmp_path)
script = tmp_path / "install_whisper_prebuilt.py"
script.write_text("stub")
plan = wupd.chained_phase_plan(force_refresh = True)
assert plan["update_available"] is True
assert plan["skip_reason"] is None
assert plan["phase"]["install_dir"] == install_dir
assert plan["phase"]["repo"] == "unslothai/whisper.cpp"
assert plan["phase"]["backend"] == "cpu"
# Pin to the exact release the freshness check offered: unpinned, the
# installer's download-host /releases/latest pointer can lag published_at
# and reinstall an older build in a loop.
assert plan["phase"]["pin_release_tag"] == "v1.9.2-unsloth.1"
def test_whisper_plan_requires_a_repairable_pair_for_slim_installs(monkeypatch, tmp_path):
install_dir = _setup_whisper(monkeypatch, tmp_path)
marker_path = install_dir / WHISPER_MARKER
marker = json.loads(marker_path.read_text())
marker["install_kind"] = "slim"
marker_path.write_text(json.dumps(marker))
wfresh.reset_caches()
monkeypatch.setattr(
wupd,
"_resolve_prebuilt_for_host",
lambda **kwargs: {"prebuilt_available": False},
)
plan = wupd.chained_phase_plan(force_refresh = True)
assert plan["update_available"] is False
assert plan["skip_reason"] == "paired_llama_unavailable"
repaired = wupd.chained_phase_plan(
force_refresh = True,
paired_llama_will_update = True,
)
assert repaired["update_available"] is True
assert repaired["phase"] is not None
def test_whisper_phase_pins_installer_to_checked_release(monkeypatch, tmp_path):
calls = []
monkeypatch.setattr(
wupd._flow,
"stream_installer",
lambda cmd, env, **kw: calls.append(cmd),
)
monkeypatch.setattr(wupd, "reset_caches", lambda **kw: None)
monkeypatch.setattr(wupd, "latest_published_release", lambda repo, **kw: "v9")
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v9")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": "v9",
},
lambda f: None,
)
cmd = calls[0]
assert "--published-release-tag" in cmd
assert cmd[cmd.index("--published-release-tag") + 1] == "v9"
def test_whisper_phase_exit_2_is_a_failed_phase(monkeypatch, tmp_path):
# No install occurred, so incompatibility must remain an actionable job
# error instead of producing a false success toast and hiding the banner.
def _raise_exit_2(cmd, env, **kw):
raise wupd._flow.InstallerExit(2, "installer exited 2: incompatible release")
monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_2)
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v1")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
with pytest.raises(wupd._flow.InstallerExit) as exc_info:
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": None,
},
lambda f: None,
)
assert exc_info.value.returncode == 2
def test_llama_update_survives_unavailable_whisper_module(monkeypatch, tmp_path):
import builtins
llama_dir = _setup_llama(monkeypatch, tmp_path)
monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kw: None)
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"),
)
real_import = builtins.__import__
def guarded_import(
name,
globals = None,
locals = None,
fromlist = (),
level = 0,
):
if name == "utils" and "whisper_cpp_update" in fromlist:
raise AssertionError("whisper module was re-imported after its failed probe")
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", guarded_import)
# A failed optional whisper probe must not be followed by an unconditional
# import. The valid llama phase still starts and completes.
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "unavailable"
def test_macos_status_uses_compatible_resolver_release(monkeypatch, tmp_path):
_setup_whisper(
monkeypatch,
tmp_path,
installed = "v1.9.1-unsloth.1",
latest = "v1.9.2-unsloth.1",
)
monkeypatch.setattr(wupd.sys, "platform", "darwin")
monkeypatch.setattr(
wupd,
"_resolve_prebuilt_for_host",
lambda **kw: {
"prebuilt_available": True,
"release_tag": "v1.9.1-unsloth.1",
},
)
status = wupd.get_update_status(force_refresh = True)
assert status["latest_tag"] == "v1.9.1-unsloth.1"
assert status["update_available"] is False
assert status["stale"] is False
def test_whisper_phase_integrity_failure_is_not_swallowed(monkeypatch, tmp_path):
def _raise_exit_1(cmd, env, **kw):
raise wupd._flow.InstallerExit(1, "installer exited 1: checksum mismatch")
monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_1)
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v1")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
with pytest.raises(wupd._flow.InstallerExit, match = "checksum mismatch"):
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": None,
},
lambda f: None,
)
# --- apply: the chained job ---
def test_apply_runs_llama_then_whisper(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
res = upd.start_update()
assert res["started"] is True, res
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama", "whisper"] # llama phase strictly first
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["llama"]["to_tag"] == "b9518"
assert job["phases"]["whisper"]["state"] == "success"
assert job["phases"]["whisper"]["to_tag"] == "v1.9.2-unsloth.1"
# Legacy top-level fields keep their llama meaning.
assert job["from_tag"] == "b9493"
assert job["to_tag"] == "b9518"
assert "Updated llama.cpp to b9518." in job["message"]
assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"]
assert job["progress"] == 1.0
assert LEGACY_JOB_FIELDS <= set(job)
def test_apply_llama_only_when_whisper_current(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama"]
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "up_to_date"
def test_apply_whisper_only_noops_llama(monkeypatch, tmp_path):
# llama current + whisper behind: the same single apply runs, with the llama
# phase a cheap already-matches no-op and the whisper phase doing the work.
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(monkeypatch, on_start = lambda cmd: events.append("llama"))
_patch_whisper_phase(monkeypatch, events)
res = upd.start_update()
assert res["started"] is True, res
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["whisper"] # the llama installer never ran
# The legacy job-level to_tag means "llama tag"; a whisper-only round
# leaves it unset so the UI never reports a llama update that never ran.
assert job["to_tag"] is None
assert job["phases"]["llama"]["state"] == "skipped"
assert job["phases"]["llama"]["reason"] == "up_to_date"
assert job["phases"]["whisper"]["state"] == "success"
assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"]
def test_whisper_reload_never_raises_job_reload_flag(monkeypatch, tmp_path):
# A whisper-only update that had to unload a warm sidecar reports
# reload_required on its phase, but the JOB flag stays down: the chat
# frontend resyncs (and clears the local checkpoint) off the job flag,
# which must mean "the llama server changed", not "the sidecar restarted".
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
def _whisper_phase(phase, set_progress):
return {
"to_tag": "v1.9.2-unsloth.1",
"reload_required": True,
"message": "Updated whisper.cpp to v1.9.2-unsloth.1.",
}
monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert job["phases"]["whisper"]["reload_required"] is True
assert not job["reload_required"]
def test_apply_refuses_when_both_current(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"
def test_apply_llama_failure_aborts_before_whisper(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(monkeypatch, returncode = 2, lines = ["boom: disk full\n"])
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "error", job
assert "boom" in (job["error"] or "")
assert events == [] # whisper never attempted
assert job["phases"]["llama"]["state"] == "error"
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "aborted"
assert job["message"] == "llama.cpp update failed."
def test_apply_whisper_failure_keeps_llama_partial_success(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
# An active model makes the llama phase report reload_required.
import threading
from types import ModuleType
class _FakeBackend:
def __init__(self):
self._serial_load_lock = threading.Lock()
self._llama_update_in_progress = False
self.is_active = True
def unload_model(self):
self.is_active = False
backend = _FakeBackend()
routes_pkg = ModuleType("routes")
routes_pkg.__path__ = []
inference_mod = ModuleType("routes.inference")
inference_mod.get_llama_cpp_backend = lambda: backend
monkeypatch.setitem(sys.modules, "routes", routes_pkg)
monkeypatch.setitem(sys.modules, "routes.inference", inference_mod)
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events, error = "whisper installer exploded")
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "error", job
assert events == ["llama", "whisper"]
# The message says both halves: llama landed, whisper did not.
assert "Updated llama.cpp to b9518." in job["message"]
assert "whisper.cpp update failed." in job["message"]
assert "whisper installer exploded" in (job["error"] or "")
# The llama phase's reload_required survives the whisper failure.
assert job["reload_required"] is True
assert job["to_tag"] == "b9518"
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["whisper"]["state"] == "error"
def test_apply_skips_whisper_local_link_silently(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True)
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama"]
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "local_link"
assert job["message"] == "Updated llama.cpp to b9518."
def test_chained_progress_windows(monkeypatch, tmp_path):
# The llama phase fills roughly the first 0.7 slice and whisper the rest.
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
seen = {}
def _whisper_phase(phase, set_progress):
with upd._job_lock:
seen["at_whisper_start"] = upd._job["progress"]
set_progress(0.5)
with upd._job_lock:
seen["mid_whisper"] = upd._job["progress"]
return {"to_tag": "v1.9.2-unsloth.1", "reload_required": False, "message": "ok"}
monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase)
_patch_llama_installer(
monkeypatch,
lines = ["Downloading app.tar.gz: 100.0% (35.0 MiB/35.0 MiB) at 9.0 MiB/s\n"],
on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"),
)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert seen["at_whisper_start"] == pytest.approx(0.7)
assert seen["mid_whisper"] == pytest.approx(0.7 + 0.5 * 0.3)
assert job["progress"] == 1.0

21
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -200,14 +200,9 @@ def _gpu_linux_host(caps):
)
def test_host_is_blackwell_includes_datacenter_parts():
assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100
assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103
assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120
assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper
assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins
# _host_is_blackwell / _blackwell_min_toolkit_for_host are prebuilt_core
# re-exports; their value tables moved verbatim to
# tests/studio/install/test_prebuilt_core.py.
def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile):
@ -285,16 +280,6 @@ def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter():
assert [a.name for a in kept] == [cuda13.name]
def test_blackwell_min_toolkit_is_sm_aware():
# Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it.
f = ilp._blackwell_min_toolkit_for_host
assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200
assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50
assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300
assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark
assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins
def test_sm103_host_drops_cuda128_windows_build():
# B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped.
host = _host(

231
studio/backend/tests/test_install_whisper_prebuilt_checksums.py Normal file
View file

@ -0,0 +1,231 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Trust-anchor tests for install_whisper_prebuilt.py.
Whisper verifies each download against the release's own
whisper-prebuilt-sha256.json checksum index (the same model as
install_llama_prebuilt.py), not a committed pins file. These pin the index
parser, the fail-closed behaviour when an asset is not covered, the
tampered-manifest guard, and the newest-release resolution.
"""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
iwp = importlib.import_module("install_whisper_prebuilt")
if not hasattr(iwp, "parse_release_checksums"):
pytest.skip("checksum-model symbols not present - check branch", allow_module_level = True)
_A = "0" * 64
_B = "1" * 64
_TAG = "v1.9.1-unsloth.1"
_REPO = "unslothai/whisper.cpp"
def _index(**overrides) -> dict:
payload = {
"schema_version": 1,
"component": "whisper.cpp",
"release_tag": _TAG,
"upstream_tag": "v1.9.1",
"artifacts": {
"whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz": {"sha256": _A},
"whisper-v1.9.1-unsloth.1-linux-x64-cuda12-portable.tar.gz": {"sha256": _B},
},
}
payload.update(overrides)
return payload
# parse_release_checksums / expected_sha256_for are prebuilt_core re-exports;
# their valid/fail-closed matrix is asserted against the real whisper
# descriptor in tests/studio/install/test_prebuilt_core.py. The download-host
# fast-path tests below still route through this module's parse wrapper.
# release tag resolution.
def test_resolve_release_tag_explicit_override_passthrough():
assert iwp.resolve_release_tag(_REPO, published_release_tag = "v1.9.1-unsloth.2") == (
"v1.9.1-unsloth.2"
)
def test_resolve_release_tag_resolves_newest_when_no_override(monkeypatch):
monkeypatch.setattr(iwp, "resolve_newest_release_tag", lambda repo: "v9.9.9-unsloth.9")
assert iwp.resolve_release_tag(_REPO, published_release_tag = None) == "v9.9.9-unsloth.9"
def test_resolve_newest_release_tag_picks_latest_published(monkeypatch):
releases = [
{"tag_name": "v1.9.1-unsloth.1", "published_at": "2026-01-01T00:00:00Z"},
{"tag_name": "v1.9.1-unsloth.3", "published_at": "2026-03-01T00:00:00Z"},
{"tag_name": "v1.9.1-unsloth.2", "published_at": "2026-02-01T00:00:00Z"},
{"tag_name": "draft", "published_at": "2026-09-01T00:00:00Z", "draft": True},
{"tag_name": "pre", "published_at": "2026-09-01T00:00:00Z", "prerelease": True},
]
monkeypatch.setattr(iwp, "fetch_json", lambda url: releases)
assert iwp.resolve_newest_release_tag(_REPO) == "v1.9.1-unsloth.3"
def test_resolve_newest_release_tag_none_published_fails_closed(monkeypatch):
monkeypatch.setattr(iwp, "fetch_json", lambda url: [{"tag_name": "d", "draft": True}])
with pytest.raises(iwp.PrebuiltFallback):
iwp.resolve_newest_release_tag(_REPO)
def test_pins_symbols_are_gone():
# The committed-pins trust model was removed in favour of llama's runtime index.
for gone in ("load_pins", "pins_path", "resolve_expected_sha256", "PINS_FILENAME"):
assert not hasattr(iwp, gone), f"{gone} should have been removed"
# Download-host fast path (resolve + fetch the JSON assets with no GitHub API).
_CPU_ASSET = "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz"
def _manifest() -> dict:
return {
"schema_version": 1,
"component": "whisper.cpp",
"upstream_tag": "v1.9.1",
"artifacts": [{"asset": _CPU_ASSET, "os": "linux", "arch": "x64", "backend": "cpu"}],
}
def _no_api(monkeypatch):
"""Fail loudly if any code path touches api.github.com."""
def _boom(*a, **k):
raise AssertionError("api.github.com was used on the fast path")
monkeypatch.setattr(iwp, "fetch_json", _boom)
monkeypatch.setattr(iwp, "github_release", _boom)
monkeypatch.setattr(iwp, "fetch_release_bundle", _boom)
def test_fetch_release_for_install_prefers_download_host(monkeypatch):
_no_api(monkeypatch)
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
def _dhj(url):
if url.endswith(iwp.SHA256_ASSET_NAME):
return _index()
if url.endswith(iwp.MANIFEST_ASSET_NAME):
return _manifest()
raise AssertionError(f"unexpected url {url}")
monkeypatch.setattr(iwp, "_download_host_json", _dhj)
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None)
assert bundle.release_tag == _TAG
assert checks[_CPU_ASSET] == _A
# asset_urls point at the download host (github.com), not the API.
assert bundle.asset_urls[iwp.SHA256_ASSET_NAME].startswith(
f"https://github.com/{_REPO}/releases/"
)
assert bundle.asset_urls[_CPU_ASSET].startswith(
f"https://github.com/{_REPO}/releases/download/"
)
walked = iwp._fetch_release_candidate(_REPO, _TAG)
assert iwp.SHA256_ASSET_NAME in walked.asset_urls
assert _CPU_ASSET in walked.asset_urls
def test_fetch_release_for_install_explicit_tag_skips_the_head(monkeypatch):
# An explicit tag needs no /releases/latest HEAD: resolving it must not call it.
monkeypatch.setattr(
iwp,
"_download_host_latest_release_tag",
lambda repo: (_ for _ in ()).throw(AssertionError("HEAD used for an explicit tag")),
)
monkeypatch.setattr(
iwp,
"_download_host_json",
lambda url: _index() if url.endswith(iwp.SHA256_ASSET_NAME) else _manifest(),
)
_no_api(monkeypatch)
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = _TAG)
assert bundle.release_tag == _TAG
def test_fetch_release_for_install_falls_back_to_api(monkeypatch):
# Fast path returns None (e.g. a 404) -> the API path resolves the release.
monkeypatch.setattr(iwp, "_resolve_release_via_download_host", lambda repo, tag: None)
sentinel = iwp.ReleaseBundle(repo = _REPO, release_tag = _TAG, manifest = _manifest(), asset_urls = {})
monkeypatch.setattr(iwp, "resolve_release_tag", lambda repo, *, published_release_tag: _TAG)
monkeypatch.setattr(iwp, "fetch_release_bundle", lambda repo, tag: sentinel)
monkeypatch.setattr(iwp, "fetch_release_checksums", lambda bundle: {_CPU_ASSET: _A})
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None)
assert bundle is sentinel
assert checks == {_CPU_ASSET: _A}
def test_resolve_via_download_host_sha_404_returns_none(monkeypatch):
import urllib.error
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
def _dhj(url):
raise urllib.error.HTTPError(url, 404, "not found", {}, None)
monkeypatch.setattr(iwp, "_download_host_json", _dhj)
assert iwp._resolve_release_via_download_host(_REPO, None) is None
def test_resolve_via_download_host_tag_mismatch_returns_none(monkeypatch):
# A checksum index whose self-reported release_tag disagrees is rejected (None).
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
monkeypatch.setattr(
iwp, "_download_host_json", lambda url: _index(release_tag = "v1.9.1-unsloth.2")
)
assert iwp._resolve_release_via_download_host(_REPO, None) is None
def test_download_host_latest_release_tag_parses_redirect(monkeypatch):
class _Resp:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def geturl(self):
return f"https://github.com/{_REPO}/releases/tag/{_TAG}"
class _Opener:
def open(
self,
req,
timeout = None,
):
return _Resp()
monkeypatch.setattr(iwp, "_URL_OPENER", _Opener())
assert iwp._download_host_latest_release_tag(_REPO) == _TAG
def test_download_host_latest_release_tag_404_returns_none(monkeypatch):
import urllib.error
class _Opener:
def open(
self,
req,
timeout = None,
):
raise urllib.error.HTTPError(req.full_url, 404, "nf", {}, None)
monkeypatch.setattr(iwp, "_URL_OPENER", _Opener())
assert iwp._download_host_latest_release_tag(_REPO) is None

3
studio/backend/tests/test_llama_cpp_update.py
View file

@ -119,6 +119,9 @@ def _clean_state(monkeypatch, tmp_path):
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
# Never hit the network in these tests.
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
# Keep the whisper piggyback out of the llama-only tests: no host probe, no
# whisper phase (test_combined_update.py covers the chained flow).
monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kwargs: None)
yield
freshness.reset_caches()
upd._reset_job_for_tests()

15
studio/backend/tests/test_llama_route.py
View file

@ -100,6 +100,21 @@ def test_status_response_exposes_update_size_bytes():
assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None
def test_status_response_exposes_update_component():
model = rl.LlamaUpdateStatusResponse(
supported = True,
update_available = True,
llama_update_available = False,
update_component = "whisper",
whisper = {
"update_available": True,
"installed_tag": "v1",
"latest_tag": "v2",
},
)
assert model.model_dump()["update_component"] == "whisper"
def test_status_handler_runs_off_event_loop(monkeypatch):
seen = {}

7
studio/backend/tests/test_local_llama_cpp_link.py
View file

@ -21,6 +21,13 @@ from utils import llama_cpp_update as u
from core.inference.llama_cpp import LlamaCppBackend
@pytest.fixture(autouse = True)
def _no_whisper_piggyback(monkeypatch):
# Keep the whisper piggyback probe off the host: these tests exercise the
# llama local-link contract only.
monkeypatch.setattr(u, "_whisper_chain_status", lambda **kwargs: None)
def _make_link(link: Path, target: Path) -> None:
"""Create a directory junction (Windows) / symlink (POSIX); neither needs
elevation."""

50
studio/backend/tests/test_middleware.py
View file

@ -34,6 +34,7 @@ def main_module():
def _make_protected_app(
max_bytes: int,
main_module,
request_max_bytes_getter = None,
upload_passthrough_prefixes: tuple = (),
upload_passthrough_max_bytes_getter = None,
):
@ -41,7 +42,13 @@ def _make_protected_app(
app.add_middleware(
main_module.MaxBodyMiddleware,
max_bytes_getter = lambda: max_bytes,
protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"),
protected_prefixes = (
"/v1/chat/completions",
"/api/inference",
"/api/settings",
"/api/train",
),
request_max_bytes_getter = request_max_bytes_getter,
upload_passthrough_prefixes = upload_passthrough_prefixes,
upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter,
)
@ -68,6 +75,10 @@ def _make_protected_app(
total += len(chunk)
return {"ok": True, "chunks": chunks, "total": total}
@app.post("/api/inference/audio/transcribe/raw")
async def transcribe_raw(request: Request):
return {"ok": True, "total": len(await request.body())}
@app.get("/api/train/status")
async def status_get():
return {"ok": True, "get": True}
@ -97,6 +108,43 @@ class TestMaxBodyMiddleware:
assert r.status_code == 200
assert r.json()["unprotected"] is True
def test_route_specific_cap_overrides_default(self, main_module):
app = _make_protected_app(
4096,
main_module,
request_max_bytes_getter = lambda path: (
128 if path.endswith("/transcribe/raw") else 4096
),
)
c = TestClient(app)
rejected = c.post(
"/api/inference/audio/transcribe/raw",
content = b"x" * 129,
)
accepted = c.post(
"/api/inference/audio/transcribe/raw",
content = b"x" * 128,
)
assert rejected.status_code == 413
assert accepted.status_code == 200
assert accepted.json()["total"] == 128
def test_stt_routes_use_audio_specific_caps(self, main_module):
from utils.upload_limits import (
STT_AUDIO_JSON_MAX_BYTES,
STT_AUDIO_RAW_MAX_BYTES,
)
assert (
main_module._get_request_body_max_bytes("/api/inference/audio/transcribe/raw")
== STT_AUDIO_RAW_MAX_BYTES
)
assert (
main_module._get_request_body_max_bytes("/api/inference/audio/transcribe")
== STT_AUDIO_JSON_MAX_BYTES
)
def test_settings_put_body_over_cap_rejected(self, main_module):
app = _make_protected_app(1024, main_module)
c = TestClient(app)

47
studio/backend/tests/test_model_update_robustness.py
View file

@ -400,6 +400,53 @@ def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path):
assert rows[0]["last_modified"] == 5_000.0
def test_cached_model_scan_hides_custom_whisper_repo(monkeypatch, tmp_path):
repo_path = tmp_path / "models--Org--CustomWhisper"
snapshot = repo_path / "snapshots" / ("a" * 40)
snapshot.mkdir(parents = True)
(snapshot / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
repo = SimpleNamespace(
repo_id = "Org/CustomWhisper",
repo_type = "model",
repo_path = repo_path,
revisions = [
SimpleNamespace(
files = [
SimpleNamespace(
file_name = "config.json",
size_on_disk = 10,
blob_path = None,
),
SimpleNamespace(
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
),
]
)
],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
monkeypatch.setattr(
CI,
"_cached_model_snapshot_path",
lambda _repo_path: snapshot,
)
monkeypatch.setattr(
CI.hf_cache_scan,
"is_snapshot_partial",
lambda *args, **kwargs: False,
)
assert CI._scan_cached_models() == []
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───

168
studio/backend/tests/test_stt_download_validation.py Normal file
View file

@ -0,0 +1,168 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The /audio/stt/download route must validate a custom Transformers repo before
snapshot_download pulls it into the shared HF cache.
Regression for a Codex finding: the Transformers engine accepts arbitrary
`owner/model` repos, so an authenticated caller could make Studio download a
large non-STT repository before load-time validation ever ran. Whisper-
compatibility is now enforced (metadata-only, no weights) before the background
download starts. The GGUF engine only accepts curated ids, so it is not gated.
"""
from __future__ import annotations
import asyncio
import sys
from pathlib import Path
import pytest
from fastapi import HTTPException
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
import core.inference.stt_ggml_sidecar as ggml_module # noqa: E402
import core.inference.stt_sidecar as stt_module # noqa: E402
import routes.inference as ri # noqa: E402
from core.inference.stt_sidecar import SttModelCompatibilityError # noqa: E402
from models.inference import SttLoadRequest # noqa: E402
def _run(coro):
return asyncio.run(coro)
def test_custom_non_whisper_repo_is_rejected_before_download(monkeypatch):
started: list = []
validated: list = []
def fake_validate(model, hf_token = None):
validated.append(model)
raise SttModelCompatibilityError(
f"STT model '{model}' is not a compatible Transformers Whisper model."
)
def fake_download(model, hf_token = None):
started.append(model)
monkeypatch.setattr(stt_module, "validate_remote_model", fake_validate)
monkeypatch.setattr(stt_module, "start_model_download", fake_download)
with pytest.raises(HTTPException) as excinfo:
_run(
ri.stt_download(
SttLoadRequest(model = "owner/chat-model", engine = "transformers"),
current_subject = "tester",
hf_token = None,
)
)
assert excinfo.value.status_code == 422
assert validated == ["owner/chat-model"]
# The download never starts for a repo that failed the Whisper check.
assert started == []
def test_validated_transformers_repo_downloads(monkeypatch):
started: list = []
revision = "a" * 40
monkeypatch.setattr(
stt_module,
"validate_remote_model",
lambda model, hf_token = None: {"model": model, "revision": revision},
)
monkeypatch.setattr(
stt_module,
"start_model_download",
lambda model, hf_token = None, revision = None: started.append((model, revision)),
)
monkeypatch.setattr(stt_module, "download_status", lambda: {"downloading": True})
resp = _run(
ri.stt_download(
SttLoadRequest(model = "owner/real-whisper", engine = "transformers"),
current_subject = "tester",
hf_token = None,
)
)
assert resp.status_code == 200
assert started == [("owner/real-whisper", revision)]
def test_gguf_engine_skips_the_transformers_repo_check(monkeypatch):
started: list = []
def fail_if_called(model, hf_token = None):
raise AssertionError("GGUF downloads must not run the Transformers repo check")
# whisper-server present, so the GGUF request stays on the GGUF engine.
monkeypatch.setattr(ggml_module, "is_available", lambda: True)
monkeypatch.setattr(stt_module, "validate_remote_model", fail_if_called)
monkeypatch.setattr(
ggml_module, "start_model_download", lambda model, hf_token = None: started.append(model)
)
monkeypatch.setattr(ggml_module, "download_status", lambda: {"downloading": True})
resp = _run(
ri.stt_download(
SttLoadRequest(model = "small", engine = "gguf"),
current_subject = "tester",
hf_token = None,
)
)
assert resp.status_code == 200
assert started == ["small"]
def test_resolve_serving_stt_engine_falls_back_when_whisper_server_absent(monkeypatch):
# A curated GGUF request downgrades to Transformers when whisper-server is not
# installed (both engines serve curated ids), but stays GGUF when it is.
monkeypatch.setattr(ggml_module, "is_available", lambda: False)
assert ri._resolve_serving_stt_engine("gguf") == "transformers"
monkeypatch.setattr(ggml_module, "is_available", lambda: True)
assert ri._resolve_serving_stt_engine("gguf") == "gguf"
# Transformers is unaffected by whisper-server availability.
monkeypatch.setattr(ggml_module, "is_available", lambda: False)
assert ri._resolve_serving_stt_engine("transformers") == "transformers"
def test_gguf_download_falls_back_to_transformers_when_server_absent(monkeypatch):
"""Selecting the default curated model on a host without whisper-server must
download through the Transformers engine, not 501/dead-end on GGUF."""
gguf_started: list = []
tf_started: list = []
monkeypatch.setattr(ggml_module, "is_available", lambda: False) # no whisper-server
# validate_remote_model no-ops curated ids in production; keep it a no-op here.
monkeypatch.setattr(
stt_module, "validate_remote_model", lambda model, hf_token = None: {"model": model}
)
monkeypatch.setattr(
stt_module,
"start_model_download",
lambda model, hf_token = None, revision = None: tf_started.append(model),
)
monkeypatch.setattr(stt_module, "download_status", lambda: {"downloading": True})
monkeypatch.setattr(
ggml_module,
"start_model_download",
lambda model, hf_token = None: gguf_started.append(model),
)
resp = _run(
ri.stt_download(
SttLoadRequest(model = "small", engine = "gguf"),
current_subject = "tester",
hf_token = None,
)
)
assert resp.status_code == 200
assert tf_started == ["small"] # served by Transformers instead of dead-ending on GGUF
assert gguf_started == []

780
studio/backend/tests/test_stt_ggml_sidecar.py Normal file
View file

@ -0,0 +1,780 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import http.server
import io
import json
import os
import sys
import threading
import time
import wave
from pathlib import Path
import numpy as np
import pytest
import core.inference.stt_ggml_sidecar as ggml_module
from core.inference.stt_ggml_sidecar import (
DEFAULT_GGML_STT_MODEL,
GGML_STT_MODELS,
GGML_STT_REPOS,
GgmlSttSidecar,
SttEngineUnavailableError,
find_whisper_server_binary,
resolve_ggml_model_id,
)
from core.inference.stt_sidecar import (
SttLanguageError,
SttLoadCancelledError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
)
@pytest.fixture(autouse = True)
def isolate_runtime_and_stub_audio_decoder(monkeypatch, tmp_path):
"""Unit tests exercise orchestration, not PyAV container parsing."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio"))
monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False)
monkeypatch.delenv("UNSLOTH_WHISPER_CPP_PATH", raising = False)
monkeypatch.setenv("PATH", "")
monkeypatch.setattr(
ggml_module,
"_decode_audio_bounded",
lambda audio: np.zeros(16000, dtype = np.float32),
)
# ---------------------------------------------------------------------------
# Model id resolution
# ---------------------------------------------------------------------------
def test_curated_ids_resolve():
for model_id in GGML_STT_MODELS:
assert resolve_ggml_model_id(model_id) == model_id
def test_default_model_resolves_from_none_and_blank():
assert resolve_ggml_model_id(None) == DEFAULT_GGML_STT_MODEL
assert resolve_ggml_model_id(" ") == DEFAULT_GGML_STT_MODEL
def test_custom_repo_ids_are_rejected():
with pytest.raises(SttModelIdError):
resolve_ggml_model_id("owner/model")
with pytest.raises(SttModelIdError):
resolve_ggml_model_id("large-v2")
def test_curated_ids_mirror_transformers_sidecar():
from core.inference.stt_sidecar import STT_MODELS
assert list(GGML_STT_MODELS.keys()) == list(STT_MODELS.keys())
def test_curated_filenames_match_repo_naming():
# unslothai/whisper-<id>-GGUF hosts whisper-<id>.bin; keep the download
# filename in lockstep with the repo so it resolves instead of 404ing.
for model_id, repo in GGML_STT_REPOS.items():
expected = repo.split("/", 1)[1].removesuffix("-GGUF") + ".bin"
assert GGML_STT_MODELS[model_id] == expected
# ---------------------------------------------------------------------------
# Binary discovery
# ---------------------------------------------------------------------------
def test_env_binary_override_wins(monkeypatch, tmp_path):
binary = tmp_path / "whisper-server"
binary.write_text("#!/bin/sh\n")
binary.chmod(0o755) # find_whisper_server_binary requires an executable
monkeypatch.setenv("WHISPER_SERVER_PATH", str(binary))
assert find_whisper_server_binary() == str(binary)
def test_env_dir_override_scans_layouts(monkeypatch, tmp_path):
monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False)
build_bin = tmp_path / "build" / "bin"
build_bin.mkdir(parents = True)
binary = build_bin / "whisper-server"
binary.write_text("#!/bin/sh\n")
binary.chmod(0o755) # find_whisper_server_binary requires an executable
monkeypatch.setenv("UNSLOTH_WHISPER_CPP_PATH", str(tmp_path))
assert find_whisper_server_binary() == str(binary)
def test_missing_binary_reports_unavailable(monkeypatch, tmp_path):
monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False)
monkeypatch.setenv("UNSLOTH_WHISPER_CPP_PATH", str(tmp_path / "nope"))
monkeypatch.setattr(ggml_module, "_managed_whisper_cpp_dir", lambda: tmp_path / "gone")
monkeypatch.setattr(ggml_module.shutil, "which", lambda name: None)
assert find_whisper_server_binary() is None
assert not ggml_module.is_available()
with pytest.raises(SttEngineUnavailableError):
ggml_module.ensure_engine_available()
def test_non_executable_binary_is_not_runnable(monkeypatch, tmp_path):
if sys.platform == "win32":
pytest.skip("X_OK is an existence check on Windows")
binary = tmp_path / "whisper-server"
binary.write_text("#!/bin/sh\n") # written but not chmod +x
monkeypatch.setenv("WHISPER_SERVER_PATH", str(binary))
monkeypatch.setattr(ggml_module.shutil, "which", lambda name: None)
assert find_whisper_server_binary() is None
# ---------------------------------------------------------------------------
# Slim-install launch guard
# ---------------------------------------------------------------------------
def _slim_install(
tmp_path,
*,
install_kind = "slim",
with_ggml = True,
linked_libraries = None,
backend = "cpu",
linked_runtime_directories = None,
runtime_wiring_version = None,
) -> str:
"""A managed-looking install tree: marker at the root, server in build/bin."""
install_dir = tmp_path / "whisper.cpp"
bin_dir = install_dir / "build" / "bin"
bin_dir.mkdir(parents = True)
binary = bin_dir / "whisper-server"
binary.write_text("#!/bin/sh\n")
binary.chmod(0o755)
marker: dict = {
"schema_version": 1,
"component": "whisper.cpp",
"release_tag": "v1.9.1-unsloth.1",
"backend": backend,
"paired_llama_tag": "b10069-mix-fb3d4ca",
}
if install_kind is not None:
marker["install_kind"] = install_kind
if linked_libraries is not None:
marker["linked_libraries"] = linked_libraries
if linked_runtime_directories is not None:
marker["linked_runtime_directories"] = linked_runtime_directories
for name in linked_runtime_directories:
catalog = bin_dir / name
catalog.mkdir()
(catalog / "kernel.dat").write_bytes(b"kernel")
if runtime_wiring_version is not None:
marker["runtime_wiring_version"] = runtime_wiring_version
(install_dir / "UNSLOTH_WHISPER_PREBUILT_INFO.json").write_text(json.dumps(marker))
if with_ggml:
names = (
("ggml.dll", "ggml-base.dll")
if sys.platform == "win32"
else ("libggml.so.0", "libggml-base.so.0")
)
for name in names:
(bin_dir / name).write_bytes(b"ggml")
return str(binary)
def test_slim_guard_flags_missing_ggml_links(monkeypatch, tmp_path):
# A slim marker whose linked ggml runtime is gone must read as engine
# unavailable (reinstall), never crash into a server launch.
binary = _slim_install(tmp_path, with_ggml = False)
assert ggml_module.slim_runtime_intact(binary) is False
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary)
assert not ggml_module.is_available()
with pytest.raises(SttEngineUnavailableError, match = "ggml"):
ggml_module.ensure_engine_available()
def test_slim_guard_passes_with_links_in_place(monkeypatch, tmp_path):
names = ["libggml.so.0", "libggml-base.so.0"]
binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names)
assert ggml_module.slim_runtime_intact(binary) is True
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary)
assert ggml_module.ensure_engine_available() == binary
def test_slim_guard_verifies_the_marker_linked_libraries(monkeypatch, tmp_path):
# New markers record the exact wired filenames; one missing name flips the
# install to unavailable even when the legacy core ggml names are present.
names = ["libggml.dylib", "libggml-base.dylib", "libggml-metal.dylib"]
binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names)
bin_dir = Path(binary).parent
for name in names[:-1]:
(bin_dir / name).write_bytes(b"ggml")
assert ggml_module.slim_runtime_intact(binary) is False # metal dylib absent
(bin_dir / names[-1]).write_bytes(b"ggml")
assert ggml_module.slim_runtime_intact(binary) is True
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary)
assert ggml_module.ensure_engine_available() == binary
def test_slim_guard_malformed_authoritative_marker_fails_closed(tmp_path):
for bad in ("not-a-list", [], [1, 2]):
root = tmp_path / f"case_{type(bad).__name__}_{len(str(bad))}"
root.mkdir()
binary = _slim_install(root, with_ggml = True, linked_libraries = bad)
assert ggml_module.slim_runtime_intact(binary) is False
def test_slim_guard_prefers_authoritative_root_marker(tmp_path):
names = ["libggml.so.0", "libggml-base.so.0"]
binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names)
packaging_marker = Path(binary).parent / "UNSLOTH_WHISPER_PREBUILT_INFO.json"
packaging_marker.write_text(json.dumps({"backend": "slim", "release_tag": "packaging"}))
assert ggml_module._whisper_install_marker(binary)["install_kind"] == "slim"
assert ggml_module.slim_runtime_intact(binary) is True
def test_slim_guard_rejects_invalid_root_even_with_inner_marker(tmp_path):
binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = ["libggml.so.0"])
root_marker = Path(binary).parents[2] / "UNSLOTH_WHISPER_PREBUILT_INFO.json"
root_marker.write_text("not json")
(Path(binary).parent / root_marker.name).write_text(json.dumps({"backend": "slim"}))
assert ggml_module.slim_runtime_intact(binary) is False
def test_slim_guard_rejects_missing_rocm_catalog(tmp_path):
names = ["libggml.so.0", "libggml-base.so.0", "libggml-hip.so"]
binary = _slim_install(
tmp_path,
linked_libraries = names,
backend = "rocm",
linked_runtime_directories = ["hipblaslt", "rocblas"],
runtime_wiring_version = 2,
)
bin_dir = Path(binary).parent
(bin_dir / "libggml-hip.so").write_bytes(b"ggml")
assert ggml_module.slim_runtime_intact(binary) is True
(bin_dir / "rocblas" / "kernel.dat").unlink()
assert ggml_module.slim_runtime_intact(binary) is False
def test_slim_guard_accepts_windows_rocm_dll_overlay(monkeypatch, tmp_path):
monkeypatch.setattr(ggml_module.sys, "platform", "win32")
names = ["ggml.dll", "ggml-base.dll", "ggml-hip.dll", "amdhip64.dll"]
binary = _slim_install(
tmp_path,
linked_libraries = names,
backend = "rocm",
linked_runtime_directories = [],
runtime_wiring_version = 2,
)
for name in names:
(Path(binary).parent / name).write_bytes(b"dll")
assert ggml_module.slim_runtime_intact(binary) is True
def test_slim_guard_ignores_fat_and_markerless_installs(tmp_path):
# Fat installs carry their own ggml; no marker means source/custom build.
fat = _slim_install(tmp_path / "fat", install_kind = None, with_ggml = False)
assert ggml_module.slim_runtime_intact(fat) is True
bare = tmp_path / "bare" / "whisper-server"
bare.parent.mkdir(parents = True)
bare.write_text("#!/bin/sh\n")
assert ggml_module.slim_runtime_intact(str(bare)) is True
# ---------------------------------------------------------------------------
# whisper-server child-process environment
# ---------------------------------------------------------------------------
def _loader_path_var() -> str:
return {"win32": "PATH", "darwin": "DYLD_LIBRARY_PATH"}.get(sys.platform, "LD_LIBRARY_PATH")
def test_child_env_scrubs_secrets_and_adds_lib_dir(monkeypatch, tmp_path):
monkeypatch.setenv("HF_TOKEN", "secret-token") # exact name
monkeypatch.setenv("MY_API_KEY", "nope") # marker substring
monkeypatch.setenv("HTTPS_PROXY", "http://u:p@px:8080") # url-name
monkeypatch.setenv("SOME_REMOTE", "https://u:pw@host/repo") # url-userinfo value
monkeypatch.setenv("STT_KEEPME", "keep") # benign
binary = tmp_path / "whisper-server"
binary.write_text("#!/bin/sh\n")
env = ggml_module._whisper_server_child_env(str(binary))
for scrubbed in ("HF_TOKEN", "MY_API_KEY", "HTTPS_PROXY", "SOME_REMOTE"):
assert scrubbed not in env
assert env.get("STT_KEEPME") == "keep"
assert str(tmp_path.resolve()) in env[_loader_path_var()].split(os.pathsep)
def test_child_env_isolates_home_and_cred_locations(monkeypatch, tmp_path):
# The downloaded server must not see the real home (token caches live
# there) nor explicit cred-store pointers like HF_HOME / NETRC.
monkeypatch.setenv("HOME", "/real/home")
monkeypatch.setenv("HF_HOME", "/real/hf")
monkeypatch.setenv("NETRC", "/real/.netrc")
monkeypatch.setattr(ggml_module, "_managed_whisper_cpp_dir", lambda: tmp_path / "managed")
binary = tmp_path / "whisper-server"
binary.write_text("#!/bin/sh\n")
env = ggml_module._whisper_server_child_env(str(binary))
assert env["HOME"] == str(tmp_path / "managed" / ".child_home")
assert "HF_HOME" not in env
assert "NETRC" not in env
assert (tmp_path / "managed" / ".child_home").is_dir()
def test_child_env_wsl_rocm_prepends_system_hip(monkeypatch, tmp_path):
if sys.platform != "linux":
pytest.skip("WSL ROCm library precedence is Linux-only")
rocm = tmp_path / "rocm-lib"
rocm.mkdir()
bindir = tmp_path / "bin"
bindir.mkdir()
binary = bindir / "whisper-server"
binary.write_text("#!/bin/sh\n")
monkeypatch.setattr(ggml_module, "_wsl_system_rocm_lib_dirs", lambda: [str(rocm)])
env = ggml_module._whisper_server_child_env(str(binary))
parts = env["LD_LIBRARY_PATH"].split(os.pathsep)
assert parts[0] == str(rocm.resolve()) # system HIP wins
assert str(bindir.resolve()) in parts # bundle libs still present
assert env.get("HSA_ENABLE_DXG_DETECTION") == "1"
def test_child_env_adds_cuda_runtime_dirs_for_cuda_bundle(monkeypatch, tmp_path):
# Versioned CUDA backend modules are valid too. They still need the
# CUDA-from-PyTorch wheel dirs for libcudart/libcublas at launch.
if sys.platform == "darwin":
pytest.skip("no CUDA on macOS")
import utils.prebuilt.runtime_libs as rl
bindir = tmp_path / "bin"
bindir.mkdir()
(bindir / "whisper-server").write_text("#!/bin/sh\n")
module_name = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so.0"
(bindir / module_name).write_text("")
cuda_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
cuda_dir.mkdir(parents = True)
monkeypatch.setattr(rl, "python_runtime_dirs", lambda: [str(cuda_dir)])
env = ggml_module._whisper_server_child_env(str(bindir / "whisper-server"))
parts = env[_loader_path_var()].split(os.pathsep)
assert str(bindir.resolve()) in parts
assert str(cuda_dir.resolve()) in parts
assert parts.index(str(bindir.resolve())) < parts.index(str(cuda_dir.resolve()))
def test_child_env_omits_cuda_runtime_dirs_for_cpu_bundle(monkeypatch, tmp_path):
# No libggml-cuda.so beside the binary -> a static CPU/Metal bundle -> the CUDA
# wheel discovery must not run and must not touch the loader path.
if sys.platform == "darwin":
pytest.skip("no CUDA on macOS")
import utils.prebuilt.runtime_libs as rl
bindir = tmp_path / "bin"
bindir.mkdir()
(bindir / "whisper-server").write_text("#!/bin/sh\n")
cuda_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
cuda_dir.mkdir(parents = True)
called = {"n": 0}
def _fake_dirs():
called["n"] += 1
return [str(cuda_dir)]
monkeypatch.setattr(rl, "python_runtime_dirs", _fake_dirs)
env = ggml_module._whisper_server_child_env(str(bindir / "whisper-server"))
parts = env[_loader_path_var()].split(os.pathsep)
assert str(cuda_dir.resolve()) not in parts
assert called["n"] == 0
def test_engine_unavailable_is_stt_unavailable():
# Routes map SttUnavailableError to HTTP 501; the engine error must share it.
assert issubclass(SttEngineUnavailableError, SttUnavailableError)
# ---------------------------------------------------------------------------
# WAV packaging
# ---------------------------------------------------------------------------
def test_pcm_to_wav_bytes_shape_and_rate():
pcm = np.zeros(3200, dtype = np.float32)
data = ggml_module._pcm_to_wav_bytes(pcm)
with wave.open(io.BytesIO(data)) as w:
assert w.getnchannels() == 1
assert w.getsampwidth() == 2
assert w.getframerate() == 16000
assert w.getnframes() == 3200
def test_pcm_to_wav_bytes_clips_out_of_range():
pcm = np.array([2.0, -2.0], dtype = np.float32)
data = ggml_module._pcm_to_wav_bytes(pcm)
with wave.open(io.BytesIO(data)) as w:
frames = np.frombuffer(w.readframes(2), dtype = "<i2")
assert frames[0] == 32767
assert frames[1] == -32767
# ---------------------------------------------------------------------------
# Sidecar orchestration
# ---------------------------------------------------------------------------
def _available(monkeypatch):
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: "/bin/echo")
def test_transcribe_requires_engine(monkeypatch):
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: None)
sidecar = GgmlSttSidecar()
with pytest.raises(SttEngineUnavailableError):
sidecar.transcribe(b"RIFF")
def test_transcribe_rejects_unknown_language(monkeypatch):
_available(monkeypatch)
sidecar = GgmlSttSidecar()
with pytest.raises(SttLanguageError):
sidecar.transcribe(b"RIFF", model = "small", language = "xx-QQ")
def test_load_requires_downloaded_model(monkeypatch):
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: None)
sidecar = GgmlSttSidecar()
with pytest.raises(SttModelNotDownloadedError):
sidecar.load("small")
def test_unloaded_sidecar_reports_nothing_resident():
sidecar = GgmlSttSidecar()
assert sidecar.loaded_model is None
assert sidecar.device is None
assert sidecar.is_loading() is False
sidecar.unload() # no-op, must not raise
def test_update_maintenance_unloads_and_blocks_new_loads(monkeypatch):
class FakeProcess:
pid = 4242
def __init__(self):
self.running = True
def poll(self):
return None if self.running else 0
def terminate(self):
self.running = False
def wait(self, timeout = None):
return 0
monkeypatch.setattr(ggml_module, "forget_pid", lambda _pid: None)
sidecar = GgmlSttSidecar()
sidecar._process = FakeProcess()
sidecar._model_id = "small"
with sidecar.update_maintenance() as model_was_active:
assert model_was_active is True
assert sidecar.loaded_model is None
with pytest.raises(SttEngineUnavailableError, match = "being updated"):
sidecar.load("small")
assert sidecar._update_in_progress is False
def test_server_pid_is_tracked_for_parent_lifetime(monkeypatch):
# The spawned server must be adopted for the terminate_all backstop and
# forgotten once this sidecar has reaped it.
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
class FakeProcess:
pid = 4242
def __init__(self, *args, **kwargs):
self.terminated = False
def poll(self):
return 1 if self.terminated else None
def terminate(self):
self.terminated = True
def wait(self, timeout = None):
return 0
events = []
monkeypatch.setattr(ggml_module.subprocess, "Popen", FakeProcess)
monkeypatch.setattr(ggml_module, "adopt_pid", lambda pid: events.append(("adopt", pid)))
monkeypatch.setattr(ggml_module, "forget_pid", lambda pid: events.append(("forget", pid)))
monkeypatch.setattr(
GgmlSttSidecar,
"_wait_for_server",
staticmethod(lambda process, port, cancel_event = None: None),
)
sidecar = GgmlSttSidecar()
sidecar.load("small")
assert events == [("adopt", 4242)]
sidecar.unload()
assert events == [("adopt", 4242), ("forget", 4242)]
def test_training_forces_whisper_server_off_gpu(monkeypatch):
# Mirror the Transformers sidecar: keep whisper.cpp on CPU during training
# so a mid-training dictation cannot reclaim the VRAM training just freed.
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
commands: list[list[str]] = []
class FakeProcess:
pid = 4242
def __init__(self, command, *args, **kwargs):
commands.append(command)
def poll(self):
return None
def terminate(self):
pass
def wait(self, timeout = None):
return 0
monkeypatch.setattr(ggml_module.subprocess, "Popen", FakeProcess)
monkeypatch.setattr(ggml_module, "adopt_pid", lambda pid: None)
monkeypatch.setattr(ggml_module, "forget_pid", lambda pid: None)
monkeypatch.setattr(
GgmlSttSidecar,
"_wait_for_server",
staticmethod(lambda process, port, cancel_event = None: None),
)
monkeypatch.setattr(ggml_module, "_training_active", lambda: False)
idle = GgmlSttSidecar()
idle.load("small")
assert "--no-gpu" not in commands[0]
assert idle.is_loading() is False
idle.unload()
monkeypatch.setattr(ggml_module, "_training_active", lambda: True)
training = GgmlSttSidecar()
training.load("small")
assert "--no-gpu" in commands[1]
training.unload()
def test_cpu_root_marker_forces_no_gpu_despite_inner_packaging_marker(monkeypatch, tmp_path):
names = ["libggml.so.0", "libggml-base.so.0"]
binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names)
(Path(binary).parent / "UNSLOTH_WHISPER_PREBUILT_INFO.json").write_text(
json.dumps({"backend": "slim"})
)
monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
commands: list[list[str]] = []
class FakeProcess:
pid = 4244
def __init__(self, command, *args, **kwargs):
commands.append(command)
def poll(self):
return None
def terminate(self):
pass
def wait(self, timeout = None):
return 0
monkeypatch.setattr(ggml_module.subprocess, "Popen", FakeProcess)
monkeypatch.setattr(ggml_module, "adopt_pid", lambda pid: None)
monkeypatch.setattr(ggml_module, "forget_pid", lambda pid: None)
monkeypatch.setattr(ggml_module, "_training_active", lambda: False)
monkeypatch.setattr(
GgmlSttSidecar,
"_wait_for_server",
staticmethod(lambda process, port, cancel_event = None: None),
)
sidecar = GgmlSttSidecar()
sidecar.load("small")
assert "--no-gpu" in commands[0]
sidecar.unload()
def test_startup_is_cancellable_before_training(monkeypatch):
# A whisper-server still binding its (Metal/CUDA) backend must be preemptible
# so training coordination can stop it before admitting the run, instead of
# racing an allocating subprocess.
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
class FakeProcess:
pid = 4243
def __init__(self, *args, **kwargs):
self.terminated = False
self.killed = False
def poll(self):
return -15 if (self.terminated or self.killed) else None
def terminate(self):
self.terminated = True
def kill(self):
self.killed = True
def wait(self, timeout = None):
return 0
monkeypatch.setattr(ggml_module.subprocess, "Popen", FakeProcess)
monkeypatch.setattr(ggml_module, "adopt_pid", lambda pid: None)
monkeypatch.setattr(ggml_module, "forget_pid", lambda pid: None)
# The server never reports ready, so _wait_for_server loops until cancelled.
def never_ready(req, timeout = None):
raise OSError("connection refused")
monkeypatch.setattr(ggml_module.urllib.request, "urlopen", never_ready)
sidecar = GgmlSttSidecar()
result: dict = {}
def _load():
try:
sidecar.load("small")
result["ok"] = True
except Exception as exc: # noqa: BLE001 - recorded for the assertion below
result["error"] = exc
thread = threading.Thread(target = _load)
thread.start()
try:
deadline = time.monotonic() + 5
while time.monotonic() < deadline and not sidecar.is_loading():
time.sleep(0.01)
assert sidecar.is_loading() is True
assert sidecar.cancel_pending_load() is True
# Blocks until the cancelled startup has been reaped and the lock freed.
sidecar.wait_for_load_to_settle()
finally:
thread.join(timeout = 5)
assert thread.is_alive() is False
assert isinstance(result.get("error"), SttLoadCancelledError)
assert sidecar.is_loading() is False
assert sidecar.loaded_model is None
class _FakeWhisperHandler(http.server.BaseHTTPRequestHandler):
"""Stands in for whisper-server's /inference endpoint."""
response_text = "Hello world.\n Second line."
def do_POST(self):
length = int(self.headers.get("Content-Length", "0"))
self.rfile.read(length)
body = json.dumps({"text": self.response_text}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *args):
pass
@pytest.fixture()
def fake_whisper_server():
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _FakeWhisperHandler)
thread = threading.Thread(target = server.serve_forever, daemon = True)
thread.start()
yield server.server_address[1]
server.shutdown()
def test_transcribe_joins_segments_one_line(monkeypatch, fake_whisper_server):
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
sidecar = GgmlSttSidecar()
def fake_load(model = None):
sidecar._port = fake_whisper_server
sidecar._model_id = ggml_module.resolve_ggml_model_id(model)
monkeypatch.setattr(sidecar, "load", fake_load)
result = sidecar.transcribe(b"RIFF", model = "small", language = "en", fast = True)
assert result["text"] == "Hello world. Second line."
assert result["language"] == "en"
assert result["model"] == "small"
assert result["duration"] == pytest.approx(1.0)
def test_transcribe_maps_bad_payload_to_decode_error(monkeypatch, fake_whisper_server):
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
monkeypatch.setattr(_FakeWhisperHandler, "response_text", None)
sidecar = GgmlSttSidecar()
def fake_load(model = None):
sidecar._port = fake_whisper_server
sidecar._model_id = ggml_module.resolve_ggml_model_id(model)
monkeypatch.setattr(sidecar, "load", fake_load)
from core.inference.stt_sidecar import SttAudioDecodeError
with pytest.raises(SttAudioDecodeError):
sidecar.transcribe(b"RIFF", model = "small")
def test_beam_size_matches_fast_flag(monkeypatch, fake_whisper_server):
_available(monkeypatch)
monkeypatch.setattr(ggml_module, "_cached_model_path", lambda model_id: "/tmp/ggml.bin")
seen: list[bytes] = []
orig_post = _FakeWhisperHandler.do_POST
def capture_post(handler):
length = int(handler.headers.get("Content-Length", "0"))
body = handler.rfile.read(length)
seen.append(body)
payload = json.dumps({"text": "ok"}).encode()
handler.send_response(200)
handler.send_header("Content-Type", "application/json")
handler.send_header("Content-Length", str(len(payload)))
handler.end_headers()
handler.wfile.write(payload)
monkeypatch.setattr(_FakeWhisperHandler, "do_POST", capture_post)
try:
sidecar = GgmlSttSidecar()
def fake_load(model = None):
sidecar._port = fake_whisper_server
sidecar._model_id = ggml_module.resolve_ggml_model_id(model)
monkeypatch.setattr(sidecar, "load", fake_load)
sidecar.transcribe(b"RIFF", model = "small", fast = True)
sidecar.transcribe(b"RIFF", model = "small", fast = False)
finally:
_FakeWhisperHandler.do_POST = orig_post
assert b'name="beam_size"\r\n\r\n1' in seen[0]
assert b'name="beam_size"\r\n\r\n5' in seen[1]
# Dictation defaults to deterministic decoding.
assert b'name="temperature"\r\n\r\n0.0' in seen[0]
def test_download_rejects_custom_ids():
with pytest.raises(SttModelIdError):
ggml_module.start_model_download("owner/model")
def test_download_status_idle_shape():
status = ggml_module.download_status()
assert set(status) >= {"downloading", "model", "error"}

219
studio/backend/tests/test_stt_review_fixes.py Normal file
View file

@ -0,0 +1,219 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regressions for a fresh review pass on the local STT dictation feature:
1. Curated GGUF dictation repos (unslothai/whisper-*-GGUF) must be hidden from
chat pickers, not just their Transformers safetensors companions.
2. The GGUF sidecar's loaded_model/device status accessors must be lock-free so
they never block behind an in-flight transcription (which holds self._lock).
3. A "gguf" unload on a host without whisper-server must target the Transformers
fallback that actually served it, and unload-all must attempt both backends
even if one raises.
4. free_stt_model_for_training must free the GGUF sidecar even when the
Transformers unload raises (independent exception boundaries).
"""
from __future__ import annotations
import asyncio
import sys
import threading
from pathlib import Path
import pytest
from fastapi import HTTPException
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
# 1. Hidden-model GGUF companions ------------------------------------------------
def test_curated_gguf_dictation_repos_are_hidden():
from utils.hidden_models import _HIDDEN_STT_REPO_IDS, is_hidden_model
for repo in (
"unslothai/whisper-tiny-GGUF",
"unslothai/whisper-base-GGUF",
"unslothai/whisper-small-GGUF",
"unslothai/whisper-large-v3-turbo-GGUF",
"unslothai/whisper-large-v3-GGUF",
):
assert repo in _HIDDEN_STT_REPO_IDS
assert is_hidden_model(repo) is True
# Case-insensitive, matching how the cache stores the repo id.
assert is_hidden_model(repo.lower()) is True
# A same-prefix but genuinely different repo is NOT hidden.
assert is_hidden_model("unslothai/whisper-large-v3-GGUF-finetune") is False
# 2. GGUF status accessors are lock-free ----------------------------------------
def test_gguf_status_accessors_do_not_block_on_the_inference_lock():
from core.inference.stt_ggml_sidecar import GgmlSttSidecar
sidecar = GgmlSttSidecar()
class _AliveProc:
pid = 4321
def poll(self):
return None # still running
sidecar._process = _AliveProc()
sidecar._model_id = "small"
holder_has_lock = threading.Event()
release = threading.Event()
def _hold_inference_lock():
# Mimic transcribe() holding self._lock across the whole HTTP call.
with sidecar._lock:
holder_has_lock.set()
release.wait(timeout = 5)
holder = threading.Thread(target = _hold_inference_lock)
holder.start()
assert holder_has_lock.wait(timeout = 5)
result: dict = {}
def _read_status():
result["model"] = sidecar.loaded_model
result["device"] = sidecar.device
reader = threading.Thread(target = _read_status)
reader.start()
reader.join(timeout = 2)
blocked = reader.is_alive()
release.set()
holder.join(timeout = 5)
reader.join(timeout = 5)
assert not blocked, "loaded_model/device blocked on self._lock (should be lock-free)"
assert result == {"model": "small", "device": "whisper.cpp"}
def test_process_alive_snapshots_process_against_concurrent_unload():
# _process_alive() must read self._process exactly once. The lock-free
# readers (loaded_model/device) can run while unload() nulls self._process;
# the old `self._process is not None and self._process.poll() is None` read it
# twice, so a null landing between the two reads called None.poll(). A
# property that yields the live process on the first read and None afterwards
# reproduces that interleaving deterministically.
from core.inference.stt_ggml_sidecar import GgmlSttSidecar
class _AliveProc:
def poll(self):
return None # still running
live = _AliveProc()
reads = {"n": 0}
class _RacingSidecar(GgmlSttSidecar):
@property
def _process(self):
reads["n"] += 1
return live if reads["n"] == 1 else None
@_process.setter
def _process(self, value):
pass # __init__ assigns None; the property drives the read
sidecar = GgmlSttSidecar()
sidecar.__class__ = _RacingSidecar # data descriptor wins over the instance attr
# Snapshot fix: exactly one read, no AttributeError from a second None read.
assert sidecar._process_alive() is True
assert reads["n"] == 1
# 3. Unload resolves through the serving engine + attempts every backend ---------
def test_gguf_unload_targets_transformers_fallback_without_whisper_server(monkeypatch):
import core.inference.stt_ggml_sidecar as ggml_module
import routes.inference as ri
monkeypatch.setattr(ggml_module, "is_available", lambda: False) # no whisper-server
calls: list = []
class _Sidecar:
def __init__(self, name):
self.name = name
def unload(self):
calls.append(self.name)
monkeypatch.setattr(ri, "_stt_sidecar_for", lambda name: _Sidecar(name))
resp = asyncio.run(ri.stt_unload(engine = "gguf", current_subject = "tester"))
assert resp.status_code == 200
# gguf is served by the Transformers fallback here, so that is what unloads.
assert calls == ["transformers"]
def test_unload_all_attempts_both_backends_even_when_one_fails(monkeypatch):
import routes.inference as ri
attempted: list = []
class _Sidecar:
def __init__(self, name):
self.name = name
def unload(self):
attempted.append(self.name)
if self.name == "transformers":
raise RuntimeError("boom")
monkeypatch.setattr(ri, "_stt_sidecar_for", lambda name: _Sidecar(name))
with pytest.raises(HTTPException) as excinfo:
asyncio.run(ri.stt_unload(engine = None, current_subject = "tester"))
assert excinfo.value.status_code == 500
# gguf is still attempted after the transformers unload raised.
assert attempted == ["transformers", "gguf"]
# 4. free_stt_model_for_training isolates the two backends -----------------------
def test_free_stt_frees_gguf_even_when_transformers_unload_raises(monkeypatch):
import routes.training_vram as tv
class _TransformersSidecar:
def is_loading(self):
return False
@property
def loaded_model(self):
return "whisper-small"
def unload(self):
raise RuntimeError("transformers unload failed")
class _GgmlSidecar:
def __init__(self):
self.unloaded = False
def is_loading(self):
return False
@property
def loaded_model(self):
return None if self.unloaded else "small"
def unload(self):
self.unloaded = True
ggml = _GgmlSidecar()
monkeypatch.setattr(
"core.inference.stt_sidecar.get_stt_sidecar", lambda: _TransformersSidecar()
)
monkeypatch.setattr("core.inference.stt_ggml_sidecar.get_ggml_stt_sidecar", lambda: ggml)
freed = tv.free_stt_model_for_training("test")
# The Transformers failure must not skip GGUF eviction.
assert ggml.unloaded is True
assert any("small" in entry for entry in freed)

332
studio/backend/tests/test_stt_review_fixes_2.py Normal file
View file

@ -0,0 +1,332 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regressions for the second review pass on the local STT dictation feature:
1. scripts/build_whisper_cpp.sh must not rm -rf a whisper.cpp/src tree under a
custom Studio home unless Studio itself created it (ownership marker), the
same policy studio/setup.sh applies before its destructive replacements.
2. _snapshot_is_complete must validate every shard of a sharded PyTorch
(pytorch_model.bin.index.json) checkpoint, like the safetensors path.
3. _snapshot_is_complete must require tokenizer assets (tokenizer.json or
vocab.json + merges.txt); weights + config alone decode to blank text.
4. Custom-repo downloads must pin the revision validated beforehand and
restrict snapshot_download to the model/tokenizer/config/preprocessor file
classes (TOCTOU + unbounded-download hardening).
5. The GGML sidecar's readiness probe must not treat an arbitrary local HTTP
responder as whisper-server (mic audio would be posted to it), and the port
reservation must stay held until just before spawn.
"""
from __future__ import annotations
import http.server
import json
import os
import socket
import stat
import subprocess
import sys
import threading
from pathlib import Path
from types import SimpleNamespace
import pytest
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
import core.inference.stt_ggml_sidecar as ggml_module
import core.inference.stt_sidecar as stt_sidecar_module
from core.inference.stt_ggml_sidecar import GgmlSttSidecar, SttEngineUnavailableError
from core.inference.stt_sidecar import validate_remote_model
_BUILD_SCRIPT = _BACKEND_ROOT.parents[1] / "scripts" / "build_whisper_cpp.sh"
# 1. build_whisper_cpp.sh ownership gate ----------------------------------------
def _stub_tools(tmp_path: Path) -> dict:
"""PATH with git/cmake stubs so the script never reaches a real build."""
bin_dir = tmp_path / "stub-bin"
bin_dir.mkdir(exist_ok = True)
for tool in ("git", "cmake"):
stub = bin_dir / tool
stub.write_text("#!/bin/sh\necho stub-%s-invoked >&2\nexit 1\n" % tool)
stub.chmod(stub.stat().st_mode | stat.S_IEXEC)
env = dict(os.environ)
env["PATH"] = f"{bin_dir}:{env['PATH']}"
return env
def _run_build_script(env: dict) -> subprocess.CompletedProcess:
return subprocess.run(
["sh", str(_BUILD_SCRIPT)],
env = env,
capture_output = True,
text = True,
timeout = 60,
)
def test_build_script_refuses_unowned_dir_in_custom_studio_home(tmp_path):
home = tmp_path / "studio-home"
src = home / "whisper.cpp" / "src"
src.mkdir(parents = True)
user_file = src / "user-data.txt"
user_file.write_text("precious")
env = _stub_tools(tmp_path)
env["UNSLOTH_STUDIO_HOME"] = str(home)
result = _run_build_script(env)
assert result.returncode != 0
assert "not marked as an Unsloth-owned" in result.stderr
# The unowned tree, and the user's file inside it, survived untouched.
assert user_file.read_text() == "precious"
def test_build_script_proceeds_when_marker_present(tmp_path):
home = tmp_path / "studio-home"
install = home / "whisper.cpp"
(install / "src").mkdir(parents = True)
(install / ".unsloth-studio-owned").write_text("")
env = _stub_tools(tmp_path)
env["UNSLOTH_STUDIO_HOME"] = str(home)
result = _run_build_script(env)
# Past the guard: it fails later at the stubbed git clone, not the gate.
assert "not marked as an Unsloth-owned" not in result.stderr
assert "stub-git-invoked" in result.stderr
def test_build_script_marks_fresh_custom_install_dir(tmp_path):
home = tmp_path / "studio-home"
home.mkdir()
env = _stub_tools(tmp_path)
env["UNSLOTH_STUDIO_HOME"] = str(home)
_run_build_script(env)
# A directory the script creates is marked so re-runs stay allowed.
assert (home / "whisper.cpp" / ".unsloth-studio-owned").is_file()
def test_build_script_keeps_legacy_home_behavior(tmp_path):
fake_home = tmp_path / "user-home"
src = fake_home / ".unsloth" / "whisper.cpp" / "src"
src.mkdir(parents = True)
env = _stub_tools(tmp_path)
env.pop("UNSLOTH_STUDIO_HOME", None)
env.pop("STUDIO_HOME", None)
env["HOME"] = str(fake_home)
result = _run_build_script(env)
# The legacy managed dir is always Studio-owned; no gate, straight to git.
assert "not marked as an Unsloth-owned" not in result.stderr
assert "stub-git-invoked" in result.stderr
# 2 + 3. _snapshot_is_complete --------------------------------------------------
def _base_snapshot(tmp_path: Path) -> Path:
snap = tmp_path / "snap"
snap.mkdir()
(snap / "config.json").write_text("{}")
(snap / "preprocessor_config.json").write_text("{}")
(snap / "tokenizer.json").write_text("{}")
return snap
def test_sharded_pytorch_snapshot_requires_every_shard(tmp_path):
snap = _base_snapshot(tmp_path)
index = {
"weight_map": {
"a": "pytorch_model-00001-of-00002.bin",
"b": "pytorch_model-00002-of-00002.bin",
}
}
(snap / "pytorch_model.bin.index.json").write_text(json.dumps(index))
(snap / "pytorch_model-00001-of-00002.bin").write_bytes(b"w" * 8)
# One missing .bin shard must read as incomplete, like the safetensors path.
assert stt_sidecar_module._snapshot_is_complete(snap) is False
(snap / "pytorch_model-00002-of-00002.bin").write_bytes(b"w" * 8)
assert stt_sidecar_module._snapshot_is_complete(snap) is True
def test_snapshot_without_tokenizer_assets_is_incomplete(tmp_path):
snap = _base_snapshot(tmp_path)
(snap / "model.safetensors").write_bytes(b"w" * 8)
assert stt_sidecar_module._snapshot_is_complete(snap) is True
# Weights + config but no tokenizer decodes to blank text; not complete.
(snap / "tokenizer.json").unlink()
assert stt_sidecar_module._snapshot_is_complete(snap) is False
# The slow vocab.json + merges.txt pair is an accepted alternative.
(snap / "vocab.json").write_text("{}")
assert stt_sidecar_module._snapshot_is_complete(snap) is False
(snap / "merges.txt").write_text("")
assert stt_sidecar_module._snapshot_is_complete(snap) is True
# 4. Revision pinning and allow_patterns ----------------------------------------
def test_validate_remote_model_returns_the_validated_revision(monkeypatch):
revision = "a" * 40
class _FakeApi:
def __init__(self, token = None):
pass
def model_info(
self,
repo,
expand = None,
timeout = None,
):
return SimpleNamespace(config = {"model_type": "whisper"}, sha = revision)
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "HfApi", _FakeApi)
result = validate_remote_model("someone/custom-whisper")
assert result["revision"] == revision
def test_download_pins_revision_and_limits_patterns(monkeypatch):
captured = {}
validated_revision = "a" * 40
head_revision = "b" * 40
def fake_snapshot_download(**kwargs):
captured.update(kwargs)
return "/cached"
class _FakeApi:
def __init__(self, token = None):
pass
def model_info(
self,
repo,
revision = None,
files_metadata = None,
timeout = None,
):
names = (
"config.json",
"preprocessor_config.json",
"tokenizer.json",
"model.safetensors",
)
siblings = [
SimpleNamespace(rfilename = name, size = 10, blob_id = name, lfs = None) for name in names
]
return SimpleNamespace(siblings = siblings, sha = head_revision)
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "HfApi", _FakeApi)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
state = stt_sidecar_module._SnapshotDownloadState()
# The revision resolved at validation time wins over the current head.
state._run("someone/custom-whisper", None, revision = validated_revision)
assert captured["revision"] == validated_revision
patterns = captured["allow_patterns"]
assert "model.safetensors" in patterns and "tokenizer.json" in patterns
# No wildcard that would admit arbitrary repo contents.
assert "*" not in patterns
# Without a validated revision (curated repos), pin to the metadata head.
captured.clear()
state._run("someone/custom-whisper", None)
assert captured["revision"] == head_revision
assert captured["allow_patterns"]
# 5. GGML readiness must identify whisper-server --------------------------------
class _CannedHandler(http.server.BaseHTTPRequestHandler):
body = b""
def do_GET(self): # noqa: N802
payload = type(self).body
self.send_response(200)
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *args):
pass
def _serve(body: bytes):
handler = type("Handler", (_CannedHandler,), {"body": body})
server = http.server.HTTPServer(("127.0.0.1", 0), handler)
thread = threading.Thread(target = server.serve_forever, daemon = True)
thread.start()
return server, server.server_address[1]
def _fake_alive_process():
return SimpleNamespace(poll = lambda: None, pid = 999999)
def test_wait_for_server_rejects_a_foreign_http_responder(monkeypatch):
server, port = _serve(b"<html>hello from some other local app</html>")
try:
monkeypatch.setattr(ggml_module, "_SERVER_START_TIMEOUT_SECONDS", 1.0)
with pytest.raises(SttEngineUnavailableError, match = "did not start in time"):
GgmlSttSidecar._wait_for_server(_fake_alive_process(), port)
finally:
server.shutdown()
def test_wait_for_server_accepts_the_whisper_server_page(monkeypatch):
server, port = _serve(b"<html><title>Whisper.cpp Server</title></html>")
try:
monkeypatch.setattr(ggml_module, "_SERVER_START_TIMEOUT_SECONDS", 5.0)
GgmlSttSidecar._wait_for_server(_fake_alive_process(), port)
finally:
server.shutdown()
def test_probe_requires_the_managed_child_to_be_alive():
server, port = _serve(b"whisper")
try:
dead = SimpleNamespace(poll = lambda: 0, pid = 999999)
assert GgmlSttSidecar._probe_is_whisper_server(dead, port) is False
assert GgmlSttSidecar._probe_is_whisper_server(_fake_alive_process(), port) is True
finally:
server.shutdown()
def test_port_reservation_is_held_until_released():
reservation, port = GgmlSttSidecar._reserve_free_port()
try:
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
with pytest.raises(OSError):
probe.bind(("127.0.0.1", port))
finally:
probe.close()
finally:
reservation.close()
# Released right before spawn: the port becomes bindable for the child.
child = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
child.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
try:
child.bind(("127.0.0.1", port))
finally:
child.close()

1262
studio/backend/tests/test_stt_sidecar.py Normal file
View file

File diff suppressed because it is too large Load diff

31
studio/backend/tests/test_training_pump_resilience.py
View file

@ -566,3 +566,34 @@ def test_db_run_created_before_pump_consumes_events(monkeypatch):
# The pump observed an already-created run; it would be False if the pump
# were started before the eager create.
assert seen["db_created"] is True
def test_startup_flag_reports_training_active_before_proc():
# Between freeing VRAM and _proc going live, a concurrent STT load must see
# training as active so it does not grab the just-freed GPU.
b = TrainingBackend()
b._spawn_in_progress = True
assert b.is_training_active() is True
def test_before_spawn_runs_inside_active_window(monkeypatch):
# The VRAM-freeing hook must run while training already counts as active, or
# an STT load racing it would place Whisper back on the freed GPU.
b = TrainingBackend()
_stub_spawn(monkeypatch)
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_pump_loop", lambda: setattr(b, "_pump_running", False))
active_during_free = {}
def before_spawn():
active_during_free["value"] = b.is_training_active()
assert b.start_training("job_active_window", model_name = "m", before_spawn = before_spawn) is True
if b._pump_thread is not None:
b._pump_thread.join(timeout = 2.0)
assert active_during_free["value"] is True
# The transient flag clears, but the live proc keeps training active.
assert b._spawn_in_progress is False
assert b.is_training_active() is True

246
studio/backend/tests/test_training_vram_coexistence.py
View file

@ -82,6 +82,63 @@ def _patch_backends(inf, llama):
return patch.dict(sys.modules, {"core.inference": core_inf, "routes.inference": routes_inf})
def _fake_stt_sidecar(
*,
model = None,
device = None,
loading = False,
):
sidecar = SimpleNamespace(
loaded_model = model,
device = device,
is_loading = lambda: loading,
)
sidecar.cancel_pending_load = MagicMock(return_value = loading)
sidecar.wait_for_load_to_settle = MagicMock()
sidecar.unload = MagicMock()
return sidecar
def _fake_ggml_sidecar(
*,
model = None,
device = None,
loading = False,
):
ggml = SimpleNamespace(
loaded_model = model,
device = device,
is_loading = lambda: loading,
)
ggml.cancel_pending_load = MagicMock(return_value = loading)
ggml.wait_for_load_to_settle = MagicMock()
ggml.unload = MagicMock()
return ggml
def _patch_stt(sidecar):
stt_module = types.ModuleType("core.inference.stt_sidecar")
stt_module.get_stt_sidecar = lambda: sidecar
# A fresh import of the GGUF sidecar pulls names from the fake module
# above and fails; fake it too so test ordering cannot break that import.
ggml_module = types.ModuleType("core.inference.stt_ggml_sidecar")
empty_ggml = _fake_ggml_sidecar()
ggml_module.get_ggml_stt_sidecar = lambda: empty_ggml
return patch.dict(
sys.modules,
{
"core.inference.stt_sidecar": stt_module,
"core.inference.stt_ggml_sidecar": ggml_module,
},
)
def _patch_ggml_stt(sidecar):
ggml_module = types.ModuleType("core.inference.stt_ggml_sidecar")
ggml_module.get_ggml_stt_sidecar = lambda: sidecar
return patch.dict(sys.modules, {"core.inference.stt_ggml_sidecar": ggml_module})
# ── summarize_resident_chat ──────────────────────────────────────────────────
@ -169,6 +226,49 @@ class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase):
self.assertTrue(out["any"]) # GGUF still detected
class TestSummarizeResidentStt(_GpuCacheResetMixin, unittest.TestCase):
def test_reports_resident_model(self):
sidecar = _fake_stt_sidecar(model = "small", device = "cuda")
with _patch_stt(sidecar):
out = tv.summarize_resident_stt()
self.assertEqual(out["model"], "small")
self.assertEqual(out["device"], "cuda")
self.assertTrue(out["any"])
self.assertFalse(out["loading"])
def test_reports_inflight_load(self):
sidecar = _fake_stt_sidecar(loading = True)
with _patch_stt(sidecar):
out = tv.summarize_resident_stt()
self.assertTrue(out["any"])
self.assertTrue(out["loading"])
def test_reports_empty_sidecar(self):
with _patch_stt(_fake_stt_sidecar()):
out = tv.summarize_resident_stt()
self.assertFalse(out["any"])
def test_reports_resident_gguf_when_transformers_idle(self):
ggml = _fake_ggml_sidecar(model = "small", device = "whisper.cpp")
with _patch_stt(_fake_stt_sidecar()), _patch_ggml_stt(ggml):
out = tv.summarize_resident_stt()
self.assertEqual(out["model"], "small")
self.assertEqual(out["device"], "whisper.cpp")
self.assertTrue(out["any"])
def test_resident_transformers_does_not_mask_loading_gguf(self):
# A Transformers model resident on CPU holds no VRAM, but a GGUF
# whisper-server still binding its accelerator backend does; the CPU
# model must not hide that in-flight startup from training admission.
sidecar = _fake_stt_sidecar(model = "small", device = "cpu")
ggml = _fake_ggml_sidecar(loading = True)
with _patch_stt(sidecar), _patch_ggml_stt(ggml):
out = tv.summarize_resident_stt()
self.assertEqual(out["model"], "small")
self.assertTrue(out["loading"])
self.assertTrue(out["any"])
# ── can_keep_during_training (auto mode) ─────────────────────────────────────
@ -438,5 +538,151 @@ class TestFreeChatModels(_GpuCacheResetMixin, unittest.TestCase):
self.assertEqual(freed, ["gguf:gemma.gguf"])
class TestFreeSttModel(_GpuCacheResetMixin, unittest.TestCase):
def test_unloads_resident_model(self):
sidecar = _fake_stt_sidecar(model = "small", device = "cuda")
with _patch_stt(sidecar):
freed = tv.free_stt_model_for_training(reason = "test")
sidecar.unload.assert_called_once()
self.assertEqual(freed, ["stt:small"])
def test_cancels_inflight_load_and_waits_to_settle(self):
sidecar = _fake_stt_sidecar(loading = True)
with _patch_stt(sidecar):
freed = tv.free_stt_model_for_training(reason = "test")
sidecar.cancel_pending_load.assert_called_once()
# The cancelled loader may still hold VRAM; we wait for it to release.
sidecar.wait_for_load_to_settle.assert_called_once()
# No model surfaced after the wait, so nothing to unload.
sidecar.unload.assert_not_called()
self.assertEqual(freed, ["stt:loading"])
def test_cancels_inflight_load_then_unloads_settled_model(self):
# A load that finished before observing the cancel leaves a resident
# model behind; it must be unloaded so training reclaims the memory.
sidecar = _fake_stt_sidecar(model = "small", loading = True)
with _patch_stt(sidecar):
freed = tv.free_stt_model_for_training(reason = "test")
sidecar.cancel_pending_load.assert_called_once()
sidecar.wait_for_load_to_settle.assert_called_once()
sidecar.unload.assert_called_once()
self.assertEqual(freed, ["stt:loading"])
def test_cancelled_load_still_unloads_gguf_sidecar(self):
# Cancelling a Transformers load must not skip the GGUF sidecar; both
# engines can hold memory at once (engine switch or direct load calls).
sidecar = _fake_stt_sidecar(loading = True)
ggml = _fake_ggml_sidecar(model = "small")
with _patch_stt(sidecar), _patch_ggml_stt(ggml):
freed = tv.free_stt_model_for_training(reason = "test")
sidecar.cancel_pending_load.assert_called_once()
ggml.unload.assert_called_once()
self.assertEqual(freed, ["stt:loading", "stt:small"])
def test_leaves_empty_sidecar_alone(self):
sidecar = _fake_stt_sidecar()
with _patch_stt(sidecar):
freed = tv.free_stt_model_for_training(reason = "test")
sidecar.unload.assert_not_called()
self.assertEqual(freed, [])
def test_cancels_inflight_gguf_load_and_waits_to_settle(self):
# A GGUF whisper-server still in startup has no loaded_model yet, so the
# coordinator must cancel and wait for it, not skip it, before training
# claims the accelerator memory it is binding.
sidecar = _fake_stt_sidecar() # Transformers idle
ggml = _fake_ggml_sidecar(loading = True)
with _patch_stt(sidecar), _patch_ggml_stt(ggml):
freed = tv.free_stt_model_for_training(reason = "test")
ggml.cancel_pending_load.assert_called_once()
ggml.wait_for_load_to_settle.assert_called_once()
ggml.unload.assert_not_called() # nothing surfaced after the wait
self.assertEqual(freed, ["stt:gguf-loading"])
class TestCoordinateModels(_GpuCacheResetMixin, unittest.TestCase):
def _run(self, chat, stt, keep_results):
keep = MagicMock(side_effect = keep_results)
with (
patch.object(tv, "summarize_resident_chat", return_value = chat),
patch.object(tv, "summarize_resident_stt", return_value = stt),
patch.object(
tv,
"free_stt_model_for_training",
return_value = ["stt:small"],
) as free_stt,
patch.object(
tv,
"free_chat_models_for_training",
return_value = ["hf:chat"],
) as free_chat,
):
freed = tv.coordinate_models_for_training(keep)
return freed, keep, free_stt, free_chat
def test_keeps_everything_when_training_fits(self):
chat = {"any": True, "loading": False}
stt = {"any": True, "loading": False}
freed, keep, free_stt, free_chat = self._run(
chat,
stt,
[(True, {"usable_gb": 40, "required_gb": 10})],
)
self.assertEqual(freed, [])
keep.assert_called_once()
free_stt.assert_not_called()
free_chat.assert_not_called()
def test_frees_stt_before_chat(self):
chat = {"any": True, "loading": False}
stt = {"any": True, "loading": False}
freed, keep, free_stt, free_chat = self._run(
chat,
stt,
[
(False, {"usable_gb": 8, "required_gb": 10}),
(True, {"usable_gb": 12, "required_gb": 10}),
],
)
self.assertEqual(freed, ["stt:small"])
self.assertEqual(keep.call_count, 2)
free_stt.assert_called_once()
free_chat.assert_not_called()
def test_frees_chat_when_stt_is_not_enough(self):
chat = {"any": True, "loading": False}
stt = {"any": True, "loading": False}
freed, keep, free_stt, free_chat = self._run(
chat,
stt,
[
(False, {"usable_gb": 8, "required_gb": 10}),
(False, {"usable_gb": 9, "required_gb": 10}),
],
)
self.assertEqual(freed, ["stt:small", "hf:chat"])
self.assertEqual(keep.call_count, 2)
free_stt.assert_called_once()
free_chat.assert_called_once()
def test_frees_loading_models_without_probe(self):
chat = {"any": True, "loading": True}
stt = {"any": True, "loading": True}
freed, keep, free_stt, free_chat = self._run(chat, stt, [])
self.assertEqual(freed, ["stt:small", "hf:chat"])
keep.assert_not_called()
free_stt.assert_called_once()
free_chat.assert_called_once()
def test_cancels_loading_stt_without_probe(self):
chat = {"any": False, "loading": False}
stt = {"any": True, "loading": True}
freed, keep, free_stt, free_chat = self._run(chat, stt, [])
self.assertEqual(freed, ["stt:small"])
keep.assert_not_called()
free_stt.assert_called_once()
free_chat.assert_not_called()
if __name__ == "__main__":
unittest.main()

156
studio/backend/tests/test_whisper_cpp_freshness.py Normal file
View file

@ -0,0 +1,156 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the whisper.cpp prebuilt freshness check.
Pins the whisper-specific version policy: the release-tag parser, the
is_behind decision matrix (with its downgrade guard), and one end-to-end
wiring smoke through the shared freshness flow. The shared marker-walk and
fail-open mechanics are covered by test_llama_cpp_freshness.py.
"""
from __future__ import annotations
import json
import sys
import types as _types
from datetime import datetime, timedelta, timezone
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
class _NoopLogger:
"""structlog-style logger: every method swallows positional + kwargs."""
def __getattr__(self, _name):
return lambda *a, **k: None
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("structlog", _structlog_stub)
import pytest
from utils import whisper_cpp_freshness as fr
# Helpers.
def _write_marker(install_dir: Path, **overrides) -> Path:
payload = {
"requested_tag": "latest",
"release_tag": "v1.9.1-unsloth.1",
"upstream_tag": "v1.9.1",
"published_repo": "unslothai/whisper.cpp",
"asset": "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz",
"asset_sha256": None,
"source": "published",
"installed_at_utc": (datetime.now(tz = timezone.utc) - timedelta(days = 1))
.isoformat()
.replace("+00:00", "Z"),
}
payload.update(overrides)
install_dir.mkdir(parents = True, exist_ok = True)
marker = install_dir / "UNSLOTH_WHISPER_PREBUILT_INFO.json"
marker.write_text(json.dumps(payload))
return marker
def _fake_binary(install_dir: Path) -> Path:
"""Stub whisper-server under the canonical cmake install layout."""
bin_dir = install_dir / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
bin_path = bin_dir / "whisper-server"
bin_path.write_text("stub\n")
return bin_path
@pytest.fixture(autouse = True)
def _reset(monkeypatch, tmp_path):
# Isolate disk cache per-test; never touch the real cache.
monkeypatch.setattr(fr, "_cache_dir", lambda: tmp_path / ".freshness")
fr.reset_caches()
yield
fr.reset_caches()
# parse_release_version.
def test_parse_release_version():
assert fr.parse_release_version("v1.9.1-unsloth.2") == (1, 9, 1, 2)
assert fr.parse_release_version("1.10.0") == (1, 10, 0, 0) # no v, no serial
assert fr.parse_release_version(" v2.0.0-unsloth.10 ") == (2, 0, 0, 10)
assert fr.parse_release_version("v1.9") == (1, 9, 0, 0) # padded
assert fr.parse_release_version("nightly") is None
assert fr.parse_release_version(None) is None
assert fr.parse_release_version("") is None
# is_behind decision matrix + downgrade guard.
def test_is_behind_serial_bump():
assert fr.is_behind("v1.9.1-unsloth.1", "v1.9.1-unsloth.2") is True
def test_is_behind_downgrade_guard():
# A lower serial or version is never "behind".
assert fr.is_behind("v1.9.1-unsloth.2", "v1.9.1-unsloth.1") is False
assert fr.is_behind("v1.10.0-unsloth.1", "v1.9.1-unsloth.9") is False
def test_is_behind_upstream_bump():
assert fr.is_behind("v1.9.1-unsloth.1", "v1.10.0-unsloth.1") is True
def test_is_behind_identical_is_false():
assert fr.is_behind("v1.9.1-unsloth.1", "v1.9.1-unsloth.1") is False
def test_is_behind_unparseable_differs_is_behind():
assert fr.is_behind("v1.9.1-unsloth.1", "nightly") is True
def test_is_behind_missing_side_fails_open():
assert fr.is_behind(None, "v1.9.1-unsloth.2") is False
assert fr.is_behind("v1.9.1-unsloth.1", None) is False
# check_prebuilt_freshness end-to-end.
def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, tmp_path):
_write_marker(
tmp_path,
release_tag = "v1.9.1-unsloth.1",
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 10))
.isoformat()
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(tmp_path)
monkeypatch.setattr(fr, "latest_published_release", lambda *a, **k: "v1.9.1-unsloth.3")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["has_marker"] is True
assert info["behind"] is True
assert info["stale"] is True
assert info["installed_tag"] == "v1.9.1-unsloth.1"
assert info["latest_tag"] == "v1.9.1-unsloth.3"
def test_marker_reader_prefers_install_root_over_packaging_marker(tmp_path):
root_marker = _write_marker(tmp_path, release_tag = "v1.9.1-unsloth.2")
binary = _fake_binary(tmp_path)
(binary.parent / root_marker.name).write_text(
json.dumps({"backend": "slim", "release_tag": "archive-metadata"})
)
assert fr.read_install_marker(str(binary))["release_tag"] == "v1.9.1-unsloth.2"

70
studio/backend/utils/hidden_models.py
View file

@ -9,6 +9,7 @@ which eagerly loads the model-config/checkpoint stack, and without importing
from __future__ import annotations
import json
import re
from pathlib import Path
from typing import Optional
@ -31,6 +32,61 @@ _DEFAULT_EMBEDDING_REPO_IDS = {
# fallback for Studio's static default embedder only; configured custom repos
# remain exact-match-only.
_DEFAULT_EMBEDDING_PATH_BASENAMES = {"bge-small-en-v1.5"}
# Curated Whisper dictation checkpoints (STT, never chat), hidden from the chat
# inventory and pickers: Transformers safetensors repos (unsloth/whisper-*) and
# their GGUF companions (unslothai/whisper-*-GGUF). Custom checkpoints are caught
# by config below, but the GGUF companions carry a raw .bin (no config.json), so
# they must be listed here by id or they leak into chat pickers.
_HIDDEN_STT_REPO_IDS = frozenset(
{
"unsloth/whisper-tiny",
"unsloth/whisper-base",
"unsloth/whisper-small",
"unsloth/whisper-large-v3-turbo",
"unsloth/whisper-large-v3",
"unslothai/whisper-tiny-GGUF",
"unslothai/whisper-base-GGUF",
"unslothai/whisper-small-GGUF",
"unslothai/whisper-large-v3-turbo-GGUF",
"unslothai/whisper-large-v3-GGUF",
}
)
def _config_is_whisper(path: Path) -> bool:
"""True if a config.json declares a Whisper model."""
try:
with open(path, "r", encoding = "utf-8") as file:
config = json.load(file)
except Exception:
return False
if not isinstance(config, dict):
return False
model_type = config.get("model_type")
if isinstance(model_type, str) and model_type.strip().lower() == "whisper":
return True
architectures = config.get("architectures")
return isinstance(architectures, list) and any(
isinstance(name, str) and name == "WhisperForConditionalGeneration"
for name in architectures
)
def _path_is_whisper_model(value: str) -> bool:
"""Inspect an existing local model path's config; never hides name-only matches."""
if _HF_REPO_ID_RE.fullmatch(value.strip()):
return False
path = Path(value).expanduser()
try:
if path.is_file():
path = path.parent
candidates = [path / "config.json"]
snapshots = path / "snapshots"
if snapshots.is_dir():
candidates.extend(child / "config.json" for child in snapshots.iterdir())
except OSError:
return False
return any(_config_is_whisper(candidate) for candidate in candidates)
def _safe_resolve(path: Path) -> Optional[str]:
@ -79,11 +135,11 @@ def _path_basename_is_default_embedder(value: str) -> bool:
def is_hidden_model(*values: str | None) -> bool:
"""True if any id/path is the RAG embedding model (the effective embedder
or its GGUF companion repo) or the llama.cpp install validation probe
(ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
None are usable chat models; the probe can be cached as a side effect of
installing the prebuilt llama-server and otherwise sorts smallest, so it
would be auto-selected.
or its GGUF companion repo), the llama.cpp install validation probe
(ggml-org/models / stories260K), or a curated/custom Whisper dictation
model, so pickers hide them (GGUF and non-GGUF). None are usable chat
models; the probe can be cached as a side effect of installing the prebuilt
llama-server and otherwise sorts smallest, so it would be auto-selected.
Hub repo ids are matched EXACTLY (case-insensitive full "owner/name"), so a
custom embedder with a generic basename like "org/model" cannot substring
@ -97,6 +153,7 @@ def is_hidden_model(*values: str | None) -> bool:
hidden_repo_ids = {
_PROBE_REPO_ID.lower(),
*(repo_id.lower() for repo_id in _DEFAULT_EMBEDDING_REPO_IDS),
*(repo_id.lower() for repo_id in _HIDDEN_STT_REPO_IDS),
}
exact_paths: list[str] = []
for model in {
@ -135,6 +192,9 @@ def is_hidden_model(*values: str | None) -> bool:
return True
if _path_contains_repo_id(v, hidden_repo_ids):
return True
# Custom Whisper checkpoints keep no curated repo id, so match by config.
if _path_is_whisper_model(v):
return True
if exact_paths:
resolved = _safe_resolve(Path(v).expanduser())
if resolved and resolved.lower() in exact_paths:

312
studio/backend/utils/llama_cpp_freshness.py
View file

@ -7,28 +7,28 @@ Reads UNSLOTH_PREBUILT_INFO.json (written by install_llama_prebuilt.py)
and compares the installed release tag against the latest on GitHub.
Surfaced via main.py:lifespan() and /api/inference/status. Fails open
on any missing data so we never show a misleading banner.
The mechanics (marker walk-up, GitHub fetch, memo + disk cache, report
skeleton) live in utils.prebuilt.freshness_flow; this module keeps the
llama version policy and the per-module caches its tests patch.
"""
from __future__ import annotations
import json
import os
import re
import time
from datetime import datetime, timezone
from datetime import datetime
from pathlib import Path
from typing import Optional
import structlog
from utils.prebuilt import freshness_flow as _flow
logger = structlog.get_logger(__name__)
# 3 days matches Unsloth's typical llama.cpp release cadence.
STALENESS_THRESHOLD_DAYS = 3
# 24h TTL keeps the GitHub call off the hot path and within rate limits.
_RELEASE_CACHE_TTL_SECONDS = 24 * 60 * 60
_INSTALL_MARKER_NAME = "UNSLOTH_PREBUILT_INFO.json"
_marker_cache: dict[str, Optional[dict]] = {}
@ -49,203 +49,60 @@ def _cache_dir() -> Path:
def read_install_marker(binary_path: Optional[str]) -> Optional[dict]:
"""Walk up from binary_path to find UNSLOTH_PREBUILT_INFO.json.
None = no marker (source build / custom path) or invalid JSON."""
if not binary_path:
return None
cached = _marker_cache.get(binary_path)
if cached is not None or binary_path in _marker_cache:
return cached
p = Path(binary_path)
marker: Optional[dict] = None
# Cover all _find_llama_server_binary layouts (binary is 1-4 dirs deep):
for parent in p.parents[:5]:
candidate = parent / _INSTALL_MARKER_NAME
if candidate.is_file():
try:
marker = json.loads(candidate.read_text(encoding = "utf-8"))
except (OSError, json.JSONDecodeError) as exc:
logger.debug(
"failed to parse install marker",
path = str(candidate),
error = str(exc),
)
marker = None
break
_marker_cache[binary_path] = marker
return marker
def _cache_path_for(repo: str) -> Path:
safe = repo.replace("/", "__")
return _cache_dir() / f"{safe}.json"
return _flow.read_install_marker(
binary_path,
marker_name = _INSTALL_MARKER_NAME,
cache = _marker_cache,
log_message = "failed to parse install marker",
)
def _load_disk_cache(repo: str) -> Optional[tuple[float, Optional[str]]]:
path = _cache_path_for(repo)
try:
payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, json.JSONDecodeError):
return None
ts = payload.get("fetched_at")
tag = payload.get("latest_tag")
if not isinstance(ts, (int, float)):
return None
return float(ts), tag if isinstance(tag, str) else None
return _flow.load_disk_cache(repo, _cache_dir())
def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None:
path = _cache_path_for(repo)
try:
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".tmp")
tmp.write_text(
json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}),
encoding = "utf-8",
)
tmp.replace(path)
except OSError as exc:
logger.debug("freshness cache write failed", repo = repo, error = str(exc))
_flow.save_disk_cache(
repo, latest_tag, _cache_dir(), log_message = "freshness cache write failed"
)
def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
"""Newest published release tag for `repo`, by publish time.
Resolves "latest" the way install_llama_prebuilt.py does (newest
non-draft/non-prerelease by ``published_at``), NOT via GitHub's
``/releases/latest`` pointer. That pointer sorts by commit date and can lag
behind the build the installer actually installs, so detection and apply
disagreed -- the cause of the downgrade/sticky banner. None on any failure
(offline, rate-limited, etc)."""
import urllib.error
import urllib.request
url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
headers = {
"Accept": "application/vnd.github+json",
"User-Agent": "unsloth-studio-freshness-check",
}
token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
if token:
headers["Authorization"] = f"Bearer {token}"
req = urllib.request.Request(url, headers = headers)
try:
with urllib.request.urlopen(req, timeout = timeout) as resp:
data = json.loads(resp.read().decode("utf-8"))
except (
urllib.error.URLError,
urllib.error.HTTPError,
OSError,
json.JSONDecodeError,
) as exc:
logger.debug("freshness fetch failed", repo = repo, error = str(exc))
return None
if not isinstance(data, list):
return None
published = [
r
for r in data
if isinstance(r, dict)
and not r.get("draft")
and not r.get("prerelease")
and isinstance(r.get("tag_name"), str)
and r.get("tag_name")
]
if not published:
return None
newest = max(published, key = lambda r: r.get("published_at") or "")
return newest["tag_name"]
"""Newest published release tag for `repo`, by publish time (see
freshness_flow for why this is not GitHub's /releases/latest pointer)."""
return _flow.fetch_latest_release_tag(repo, timeout, log_message = "freshness fetch failed")
def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]:
"""Latest release tag for `repo`. Memo + disk-cached (24h TTL).
None when offline and never previously cached."""
if not repo:
return None
now = time.time()
if not force_refresh:
memo = _release_memo.get(repo)
if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS:
return memo[1]
disk = _load_disk_cache(repo)
if disk and now - disk[0] < _RELEASE_CACHE_TTL_SECONDS:
_release_memo[repo] = disk
return disk[1]
latest = _fetch_latest_release_tag(repo)
if latest is None:
# Keep last-good disk value rather than poisoning with None.
disk = _load_disk_cache(repo)
if disk:
_release_memo[repo] = disk
return disk[1]
return None
_release_memo[repo] = (now, latest)
_save_disk_cache(repo, latest)
return latest
return _flow.latest_published_release(
repo,
force_refresh = force_refresh,
memo = _release_memo,
cache_dir = lambda: _cache_dir(),
fetch = lambda r: _fetch_latest_release_tag(r),
save = lambda r, tag: _save_disk_cache(r, tag),
)
def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]:
"""Asset name -> size (bytes) for the newest published release of `repo`,
selected exactly like _fetch_latest_release_tag. None on any failure."""
import urllib.error
import urllib.request
url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
headers = {
"Accept": "application/vnd.github+json",
"User-Agent": "unsloth-studio-freshness-check",
}
token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
if token:
headers["Authorization"] = f"Bearer {token}"
req = urllib.request.Request(url, headers = headers)
try:
with urllib.request.urlopen(req, timeout = timeout) as resp:
data = json.loads(resp.read().decode("utf-8"))
except (
urllib.error.URLError,
urllib.error.HTTPError,
OSError,
json.JSONDecodeError,
) as exc:
logger.debug("freshness asset fetch failed", repo = repo, error = str(exc))
return None
if not isinstance(data, list):
return None
published = [
r
for r in data
if isinstance(r, dict)
and not r.get("draft")
and not r.get("prerelease")
and isinstance(r.get("tag_name"), str)
and r.get("tag_name")
]
if not published:
return None
newest = max(published, key = lambda r: r.get("published_at") or "")
assets: dict[str, int] = {}
for a in newest.get("assets") or []:
name, size = a.get("name"), a.get("size")
if isinstance(name, str) and isinstance(size, int):
assets[name] = size
return assets
return _flow.fetch_latest_release_assets(
repo, timeout, log_message = "freshness asset fetch failed"
)
def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]:
"""Newest-release asset sizes for `repo`, memoized (24h TTL). None when
offline and never fetched. In-memory only -- a restart simply re-fetches."""
if not repo:
return None
now = time.time()
if not force_refresh:
memo = _assets_memo.get(repo)
if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS:
return memo[1]
assets = _fetch_latest_release_assets(repo)
if assets is None:
memo = _assets_memo.get(repo)
return memo[1] if memo else None
_assets_memo[repo] = (now, assets)
return assets
return _flow.latest_release_assets(
repo,
force_refresh = force_refresh,
memo = _assets_memo,
fetch = lambda r: _fetch_latest_release_assets(r),
)
def update_download_size_bytes(
@ -290,16 +147,7 @@ def update_download_size_bytes(
def _parse_installed_at(value: object) -> Optional[datetime]:
if not isinstance(value, str) or not value:
return None
s = value.replace("Z", "+00:00") if value.endswith("Z") else value
try:
dt = datetime.fromisoformat(s)
except ValueError:
return None
if dt.tzinfo is None:
dt = dt.replace(tzinfo = timezone.utc)
return dt
return _flow.parse_installed_at(value)
def parse_base_build(tag: object) -> Optional[int]:
@ -350,64 +198,27 @@ def check_prebuilt_freshness(
behind = installed genuinely older than latest (see is_behind).
stale = behind AND age >= threshold.
Fails open on missing data (behind/stale stay False)."""
out: dict = {
"has_marker": False,
"stale": False,
"behind": False,
"installed_tag": None,
"latest_tag": None,
"installed_at_utc": None,
"age_days": None,
"published_repo": None,
"threshold_days": int(threshold_days),
}
marker = read_install_marker(binary_path)
if not marker:
return out
out["has_marker"] = True
# Display prefers the normalized base ("tag"); comparison below prefers the
# full "release_tag" -- deliberately opposite fallbacks.
out["installed_tag"] = marker.get("tag") or marker.get("release_tag")
out["installed_at_utc"] = marker.get("installed_at_utc")
out["published_repo"] = marker.get("published_repo")
# The marker records both a normalized base tag ("tag", e.g. b9596) and the
# full release tag ("release_tag", e.g. b9596-mix-<sha>). Compare against the
# FULL identity, since GitHub /releases/latest returns the full tag_name --
# comparing the normalized base against the full latest is what produced the
# permanent "downgrade" banner on every mix release.
installed_full = marker.get("release_tag") or marker.get("tag")
repo = out["published_repo"]
if not repo or not installed_full:
return out
latest = latest_published_release(repo)
out["latest_tag"] = latest
out["behind"] = is_behind(installed_full, latest)
if not out["behind"]:
return out
installed_at = _parse_installed_at(out["installed_at_utc"])
if installed_at is None:
return out
now = now or datetime.now(tz = timezone.utc)
age_seconds = (now - installed_at).total_seconds()
out["age_days"] = max(0, int(age_seconds // 86400))
if age_seconds >= threshold_days * 86400:
out["stale"] = True
return out
# full release tag ("release_tag", e.g. b9596-mix-<sha>). Display prefers the
# normalized base; comparison uses the FULL identity, since GitHub
# /releases/latest returns the full tag_name -- comparing the normalized base
# against the full latest is what produced the permanent "downgrade" banner
# on every mix release. Deliberately opposite fallbacks.
return _flow.check_freshness(
binary_path,
threshold_days = threshold_days,
now = now,
read_marker = lambda p: read_install_marker(p),
latest_release = lambda repo: latest_published_release(repo),
behind = lambda installed, latest: is_behind(installed, latest),
display_tag = lambda marker: marker.get("tag") or marker.get("release_tag"),
compare_tag = lambda marker: marker.get("release_tag") or marker.get("tag"),
)
def format_stale_warning(info: dict) -> str:
"""Human-readable one-liner for stale prebuilt info."""
age = info.get("age_days")
installed = info.get("installed_tag") or "unknown"
latest = info.get("latest_tag") or "unknown"
age_str = f"{age} day{'s' if age != 1 else ''}" if age is not None else "some time"
return (
f"llama.cpp prebuilt is {age_str} behind: installed "
f"{installed}, latest {latest}. Run `unsloth studio update` "
f"to refresh."
)
return _flow.format_stale_warning(info, component = "llama.cpp")
def reset_caches(*, drop_disk: bool = False) -> None:
@ -420,13 +231,8 @@ def reset_caches(*, drop_disk: bool = False) -> None:
(see its last-good fallback) and the banner could linger. Dropping the disk
cache makes latest read as None in that offline case, so the banner fails
open (off) instead of pointing at the just-replaced build."""
_marker_cache.clear()
_release_memo.clear()
_assets_memo.clear()
if drop_disk:
import shutil
# _cache_dir() is a dedicated freshness-only subdir; it is re-created on
# the next _save_disk_cache. ignore_errors so a missing/locked dir is a
# no-op rather than breaking an otherwise successful install.
shutil.rmtree(_cache_dir(), ignore_errors = True)
_flow.reset_caches(
(_marker_cache, _release_memo, _assets_memo),
drop_disk = drop_disk,
cache_dir = lambda: _cache_dir(),
)

570
studio/backend/utils/llama_cpp_update.py
View file

@ -17,17 +17,22 @@ Design notes:
thread; callers poll get_update_status() for the job state.
- Everything fails open: a missing marker / offline GitHub / source build just
reports update_available=False and never blocks the app.
- The mechanics (managed-root resolution, local-link detection, the resolve
probe, the streamed installer run) live in utils.prebuilt.update_flow; this
module keeps the llama policy and the job dict its callers poll.
- This is the single main update item: whisper.cpp piggybacks on it. Status
folds in a whisper sub-status (update_available becomes the union) and apply
chains a whisper phase after the llama phase when whisper is behind (see
update_flow.run_chained_update and whisper_cpp_update.chained_phase_plan).
"""
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional
@ -43,7 +48,7 @@ from utils.llama_cpp_freshness import (
reset_caches,
update_download_size_bytes,
)
from utils.process_lifetime import child_popen_kwargs
from utils.prebuilt import update_flow as _flow
logger = structlog.get_logger(__name__)
@ -51,33 +56,18 @@ DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp"
_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate
# Background job state. Single in-flight update at a time, guarded by _job_lock.
_JOB_IDLE = "idle"
_JOB_RUNNING = "running"
_JOB_SUCCESS = "success"
_JOB_ERROR = "error"
_JOB_IDLE = _flow.JOB_IDLE
_JOB_RUNNING = _flow.JOB_RUNNING
_JOB_SUCCESS = _flow.JOB_SUCCESS
_JOB_ERROR = _flow.JOB_ERROR
_job_lock = threading.Lock()
_job: dict = {
"state": _JOB_IDLE,
"message": "",
"from_tag": None,
"to_tag": None,
"reload_required": None,
"error": None,
"progress": None,
"started_at": None,
"finished_at": None,
}
_job: dict = _flow.new_job()
# Matches the installer's download progress lines, e.g.
# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s".
_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(")
# The download dominates the update; extract/validate fill the last slice.
_DOWNLOAD_PROGRESS_CEILING = 0.95
def _utcnow() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
_utcnow = _flow.utcnow
_is_under = _flow.is_under
_is_external_link = _flow.is_external_link
_rocm_install_args = _flow.rocm_install_args
def _find_binary() -> Optional[str]:
@ -94,37 +84,19 @@ def _find_binary() -> Optional[str]:
def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]:
"""The directory holding UNSLOTH_PREBUILT_INFO.json -- i.e. the install root
install_llama_prebuilt.py wrote and the one we re-install into. Walks up from
the binary the same way read_install_marker() does."""
if not binary_path:
return None
p = Path(binary_path)
for parent in p.parents[:5]:
if (parent / _INSTALL_MARKER_NAME).is_file():
return parent
return None
install_llama_prebuilt.py wrote and the one we re-install into."""
return _flow.install_dir_for(binary_path, marker_name = _INSTALL_MARKER_NAME)
def _installer_script() -> Optional[Path]:
"""Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then
searches up from this file for both ``<root>/install_llama_prebuilt.py`` and
``<root>/studio/install_llama_prebuilt.py`` so it works in the dev tree and
in an installed Unsloth layout."""
env = os.environ.get("UNSLOTH_LLAMA_INSTALLER")
if env and Path(env).is_file():
return Path(env)
here = Path(__file__).resolve()
for up in here.parents:
for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"):
if cand.is_file():
return cand
return None
"""Locate install_llama_prebuilt.py (UNSLOTH_LLAMA_INSTALLER wins)."""
return _flow.find_installer_script(
env_var = "UNSLOTH_LLAMA_INSTALLER", script_name = "install_llama_prebuilt.py"
)
# Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we
# ask the installer whether an official prebuilt now exists for this host. Memo
# is 24h; only successful answers are cached so a network blip retries.
_RESOLVE_TTL_SECONDS = 24 * 60 * 60
# ask the installer whether an official prebuilt now exists for this host.
_resolve_memo: dict = {}
@ -132,39 +104,12 @@ def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict]
"""Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return
{prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or
None. Fail-open: any error -> None so a source build never blocks the app."""
now = time.time()
if not force_refresh and _resolve_memo:
if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS:
return _resolve_memo.get("value")
script = _installer_script()
if script is None:
return None
value: Optional[dict] = None
try:
proc = subprocess.run(
[
sys.executable,
str(script),
"--resolve-prebuilt",
"latest",
"--output-format",
"json",
],
capture_output = True,
text = True,
timeout = 60,
)
out = (proc.stdout or "").strip()
if proc.returncode == 0 and out:
parsed = json.loads(out.splitlines()[-1])
if isinstance(parsed, dict):
value = parsed
except Exception as exc: # pragma: no cover - subprocess/json defensive
logger.debug("llama update: resolve-prebuilt failed", error = str(exc))
value = None
if value is not None: # cache real answers; let failures retry next poll
_resolve_memo.update(at = now, value = value)
return value
return _flow.resolve_prebuilt_for_host(
force_refresh = force_refresh,
memo = _resolve_memo,
installer_script = lambda: _installer_script(),
log_message = "llama update: resolve-prebuilt failed",
)
def _installed_build_number(binary: Optional[str]) -> Optional[int]:
@ -218,38 +163,16 @@ def get_installed_llama_version() -> Optional[str]:
return f"b{n}" if n is not None else None
def _is_under(path: Path, root: Path) -> bool:
try:
p, r = path.resolve(), root.resolve()
except (OSError, ValueError):
p, r = path, root
return p == r or r in p.parents
def _llama_install_root(binary: Optional[str]) -> Optional[Path]:
"""The Unsloth-managed llama.cpp root the active binary lives under, or None
when the binary is unmanaged. Installing anywhere the active binary is not
would not replace what _find_llama_server_binary runs (which prefers a pinned
LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we
refuse rather than silently install into an inactive or foreign tree."""
marked = _install_dir_for(binary)
if marked is not None:
return marked
if not binary:
return None
# LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery;
# never auto-replace its tree (even a user's own llama.cpp checkout).
if os.environ.get("LLAMA_SERVER_PATH"):
return None
p = Path(binary)
env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
if env and _is_under(p, Path(env)):
return Path(env)
for parent in p.parents:
if parent.name == "llama.cpp":
return parent
# PATH / system / custom install: not a managed tree, so do not offer.
return None
when the binary is unmanaged (see update_flow.managed_install_root)."""
return _flow.managed_install_root(
binary,
marker_root = _install_dir_for(binary),
server_path_var = "LLAMA_SERVER_PATH",
cpp_path_var = "UNSLOTH_LLAMA_CPP_PATH",
dir_name = "llama.cpp",
)
def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
@ -324,69 +247,83 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
}
def _is_external_link(path: Optional[Path]) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
the user's own llama.cpp checkout, so Unsloth must never auto-update it."""
if path is None:
return False
try:
if os.path.islink(path):
return True
except OSError:
return False
if os.name == "nt":
try:
import stat
attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined]
return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT)
except (OSError, AttributeError):
return False
return False
def _active_install_is_local_link(binary: Optional[str]) -> bool:
"""True when the active llama-server resolves through a --with-llama-cpp-dir
local link at the canonical llama.cpp directory. An update would write
through that link into the user's own checkout (or fail), so the install is
treated as externally managed: no update is offered or applied. Checks only
up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root
above it can't trip a false positive."""
if not binary:
return False
for parent in Path(binary).parents:
if _is_external_link(parent):
return True
if parent.name == "llama.cpp":
break
return False
local link at the canonical llama.cpp directory (see
update_flow.active_install_is_local_link)."""
return _flow.active_install_is_local_link(binary, dir_name = "llama.cpp")
def _local_link_status() -> dict:
"""Status payload for a local-link install: unmanaged, no update offered."""
with _job_lock:
job = dict(_job)
return {
"supported": False,
"update_available": False,
"stale": False,
"installed_tag": None,
"latest_tag": None,
"published_repo": None,
"installed_at_utc": None,
"age_days": None,
"source_build": False,
"local_link": True,
"update_size_bytes": None,
"job": job,
return _flow.local_link_status(_job, _job_lock)
def _whisper_chain_status(
*, force_refresh: bool = False, paired_llama_will_update: bool = False
) -> Optional[dict]:
"""Whisper's piggyback plan for the combined update item (see
whisper_cpp_update.chained_phase_plan). None disables the piggyback --
fail-open so whisper can never break the llama status or apply."""
try:
from utils import whisper_cpp_update
return whisper_cpp_update.chained_phase_plan(
force_refresh = force_refresh,
paired_llama_will_update = paired_llama_will_update,
)
except Exception as exc: # pragma: no cover - defensive
logger.debug("llama update: whisper piggyback probe failed", error = str(exc))
return None
def _merge_whisper_status(status: dict, *, force_refresh: bool = False) -> dict:
"""Fold the whisper sub-status into the llama status payload: the llama
update item is the single UI surface, so update_available becomes the union
(llama behind OR whisper behind) while llama_update_available keeps the
llama-only flag. All pre-existing top-level fields are preserved."""
status["llama_update_available"] = bool(status.get("update_available"))
plan = _whisper_chain_status(
force_refresh = force_refresh,
paired_llama_will_update = status["llama_update_available"],
)
if plan is None:
status["whisper"] = None
status["update_component"] = "llama" if status["llama_update_available"] else None
return status
sub = plan.get("status") or {}
status["whisper"] = {
"update_available": bool(plan.get("update_available")),
"installed_tag": sub.get("installed_tag"),
"latest_tag": sub.get("latest_tag"),
"update_size_bytes": sub.get("update_size_bytes"),
"skip_reason": plan.get("skip_reason"),
}
whisper_update_available = bool(plan.get("update_available"))
if whisper_update_available:
status["update_available"] = True
status["update_component"] = (
"llama"
if status["llama_update_available"]
else "whisper"
if whisper_update_available
else None
)
return status
def get_update_status(*, force_refresh: bool = False) -> dict:
"""Report whether a newer prebuilt exists plus the current job state.
"""Report whether an update is available plus the current job state.
force_refresh bypasses the 24h release cache for an explicit "check now".
This is the single main update item: llama.cpp drives it and the whisper
piggyback is folded in (see _merge_whisper_status). force_refresh bypasses
the 24h release cache for an explicit "check now".
"""
status = _llama_only_status(force_refresh = force_refresh)
return _merge_whisper_status(status, force_refresh = force_refresh)
def _llama_only_status(*, force_refresh: bool = False) -> dict:
"""The llama.cpp half of get_update_status (no whisper sub-status)."""
binary = _find_binary()
# A --with-llama-cpp-dir local link is the user's own tree; never offer to
# replace it. Bail before any network/freshness work.
@ -456,33 +393,19 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
}
def _rocm_install_args(asset: Optional[str]) -> list[str]:
"""Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh.
The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx
ROCm bundles carry the family in the name (rocm-gfx110X), version-tagged
bundles only rocm/hip."""
if not asset:
return []
low = asset.lower()
if "rocm" not in low and "hip" not in low:
return []
gfx = re.search(r"-gfx[0-9a-z]+", low)
if gfx:
# _normalize_forwarded_gfx accepts the family form (gfx110x -> gfx110X).
return ["--rocm-gfx", gfx.group(0).lstrip("-")]
return ["--has-rocm"]
def _run_update(
def _run_llama_phase(
install_dir: Path,
repo: str,
asset: Optional[str],
script: Path,
pin_release_tag: Optional[str] = None,
pin_release_tag: Optional[str],
set_progress,
force_cpu: bool = False,
) -> None:
"""Worker: put the backend into a maintenance state, run the installer for
the latest prebuilt, then refresh caches so the next load uses the new build.
) -> dict:
"""The llama phase of a chained update: put the backend into a maintenance
state, run the installer for the latest prebuilt, then refresh caches so the
next load uses the new build. Returns {to_tag, reload_required, message};
raises on failure.
pin_release_tag pins the installer to that exact published release instead
of letting it re-resolve "latest" itself (see start_update for why)."""
@ -530,7 +453,6 @@ def _run_update(
if force_cpu:
cmd.append("--force-cpu")
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream progress lines into job["progress"].
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
# box would otherwise re-route and silently replace the Vulkan build.
@ -538,44 +460,12 @@ def _run_update(
# _rocm_install_args).
if asset and "vulkan" in asset.lower():
env["UNSLOTH_FORCE_VULKAN"] = "1"
proc = subprocess.Popen(
_flow.stream_installer(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = env,
**child_popen_kwargs(),
env,
set_progress = set_progress,
timeout_seconds = _INSTALL_TIMEOUT_SECONDS,
)
timed_out = threading.Event()
def _kill_on_timeout() -> None:
timed_out.set()
proc.kill()
watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout)
watchdog.daemon = True
watchdog.start()
tail_lines: list[str] = []
try:
assert proc.stdout is not None
for line in proc.stdout:
tail_lines.append(line)
if len(tail_lines) > 80:
del tail_lines[0]
m = _PROGRESS_LINE_RE.search(line)
if m is None:
continue
fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING
with _job_lock:
_job["progress"] = max(_job.get("progress") or 0.0, fraction)
returncode = proc.wait()
finally:
watchdog.cancel()
if timed_out.is_set():
raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s")
if returncode != 0:
tail = "".join(tail_lines).strip()[-1500:]
raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}")
# Drop stale caches so the banner re-checks the swapped marker.
# If GitHub is offline, latest stays unknown and the banner fails open.
@ -597,29 +487,18 @@ def _run_update(
):
raise RuntimeError(f"pinned release {pin_release_tag} but installer produced {new_tag}")
with _job_lock:
_job.update(
state = _JOB_SUCCESS,
message = (
f"Updated llama.cpp to {new_tag}."
+ (" Reload your model to use it." if model_was_active else "")
),
to_tag = new_tag,
reload_required = model_was_active,
error = None,
progress = 1.0,
finished_at = _utcnow(),
)
logger.info("llama update: success", to_tag = new_tag)
return {
"to_tag": new_tag,
"reload_required": model_was_active,
"message": (
f"Updated llama.cpp to {new_tag}."
+ (" Reload your model to use it." if model_was_active else "")
),
}
except Exception as exc:
logger.warning("llama update: failed", error = str(exc))
with _job_lock:
_job.update(
state = _JOB_ERROR,
message = "llama.cpp update failed.",
error = str(exc),
finished_at = _utcnow(),
)
raise
finally:
# Always clear maintenance state.
if backend is not None:
@ -629,50 +508,58 @@ def _run_update(
pass
def start_update() -> dict:
"""Kick off a background update. Idempotent: a second call while one is
running returns the in-flight job rather than starting another."""
# Combined-job progress split when both phases run (download sizes: the llama
# bundle dwarfs the whisper one); normalized to 0..1 when a phase is skipped.
_LLAMA_PHASE_WEIGHT = 0.7
_WHISPER_PHASE_WEIGHT = 0.3
def _plan_llama_phase() -> dict:
"""Decide how the llama phase of a combined update runs. Returns {"spec"}
when llama should install, else {"skip_reason", "refusal"}: skip_reason
marks the phase skipped inside a chained job, refusal is the started=False
response when the whisper phase has nothing to run either."""
binary = _find_binary()
# Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt
# here would write through the link into the user's own checkout (or fail)
# and silently drop the link the flag created.
if _active_install_is_local_link(binary):
return {
"started": False,
"reason": "local_link",
"message": (
"llama.cpp is a local directory linked with --with-llama-cpp-dir; "
"Unsloth won't replace it. Update your own llama.cpp checkout instead."
),
"job": get_update_status()["job"],
"skip_reason": "local_link",
"refusal": {
"started": False,
"reason": "local_link",
"message": (
"llama.cpp is a local directory linked with --with-llama-cpp-dir; "
"Unsloth won't replace it. Update your own llama.cpp checkout instead."
),
},
}
marker = read_install_marker(binary)
script = _installer_script()
if script is None:
return {
"started": False,
"reason": "installer_missing",
"message": "install_llama_prebuilt.py could not be located.",
"job": get_update_status()["job"],
"skip_reason": "installer_missing",
"refusal": {
"started": False,
"reason": "installer_missing",
"message": "install_llama_prebuilt.py could not be located.",
},
}
# A job already in flight wins over any freshness re-check below (and skips
# its network call). The final lock block re-checks to close the TOCTOU.
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
if marker:
# Mirror the detection guard: a direct POST or a stale banner must not
# start an install when the latest is not actually newer (force a fresh
# check so a stale 24h cache can't wrongly block a real update either).
status = get_update_status(force_refresh = True)
status = _llama_only_status(force_refresh = True)
if not status.get("update_available"):
return {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at the latest prebuilt.",
"job": status["job"],
"skip_reason": "up_to_date",
"refusal": {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at the latest prebuilt.",
},
}
install_dir = _install_dir_for(binary)
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
@ -693,20 +580,27 @@ def start_update() -> dict:
src = _source_build_status(binary, force_refresh = True) if binary else None
if src is None:
return {
"started": False,
"reason": "no_prebuilt_available",
"message": (
"No official llama.cpp prebuilt is available for this host, "
"so the source build cannot be swapped automatically."
),
"job": get_update_status()["job"],
"skip_reason": "no_prebuilt_available",
"refusal": {
"started": False,
"reason": "no_prebuilt_available",
"message": (
"No official llama.cpp prebuilt is available for this host, "
"so the source build cannot be swapped automatically."
),
},
}
if not src.get("update_available"):
return {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at or newer than the latest prebuilt.",
"job": get_update_status()["job"],
"skip_reason": "up_to_date",
"refusal": {
"started": False,
"reason": "up_to_date",
"message": (
"The installed llama.cpp build is already at or newer than the "
"latest prebuilt."
),
},
}
res = _resolve_prebuilt_for_host()
install_dir = _llama_install_root(binary)
@ -721,31 +615,116 @@ def start_update() -> dict:
if install_dir is None:
return {
"started": False,
"reason": "no_install_dir",
"message": "Could not determine the llama.cpp install directory.",
"job": get_update_status()["job"],
"skip_reason": "no_install_dir",
"refusal": {
"started": False,
"reason": "no_install_dir",
"message": "Could not determine the llama.cpp install directory.",
},
}
return {
"spec": {
"install_dir": install_dir,
"repo": repo,
"asset": asset,
"script": script,
"pin_release_tag": pin_release_tag,
"from_tag": from_tag,
"force_cpu": force_cpu,
}
}
def start_update() -> dict:
"""Kick off a background update job. The job chains the llama phase (the
existing flow) with a whisper phase that runs only when whisper is actually
behind; either phase no-ops cleanly when its component is current or
unmanaged. Idempotent: a second call while one is running returns the
in-flight job rather than starting another."""
# A job already in flight wins over any freshness re-check below (and skips
# its network calls). The final lock block re-checks to close the TOCTOU.
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
llama_plan = _plan_llama_phase()
llama_spec = llama_plan.get("spec")
whisper_plan = _whisper_chain_status(
force_refresh = True,
paired_llama_will_update = llama_spec is not None,
)
whisper_spec = (whisper_plan or {}).get("phase")
if llama_spec is None and whisper_spec is None:
# Nothing to run in either phase: answer with the llama refusal so the
# existing reasons (local_link / up_to_date / ...) keep their meaning.
refusal = dict(llama_plan["refusal"])
with _job_lock:
refusal["job"] = dict(_job)
return refusal
whisper_run = None
if whisper_spec is not None:
from utils import whisper_cpp_update as _whisper
whisper_run = lambda set_progress: _whisper.run_chained_phase(whisper_spec, set_progress)
phases = [
{
"name": "llama",
"weight": _LLAMA_PHASE_WEIGHT,
"failure_message": "llama.cpp update failed.",
"skip_reason": llama_plan.get("skip_reason"),
"run": (
(
lambda set_progress: _run_llama_phase(
llama_spec["install_dir"],
llama_spec["repo"],
llama_spec["asset"],
llama_spec["script"],
llama_spec["pin_release_tag"],
set_progress,
force_cpu = llama_spec.get("force_cpu", False),
)
)
if llama_spec
else None
),
},
{
"name": "whisper",
"weight": _WHISPER_PHASE_WEIGHT,
"failure_message": "whisper.cpp update failed.",
# The sidecar reload is whisper-internal; it must not trip the
# job-level reload flag the chat frontend resyncs on.
"affects_job_reload": False,
"skip_reason": (whisper_plan or {}).get("skip_reason") or "unavailable",
"run": whisper_run,
},
]
running = " + ".join(
name for name, spec in (("llama.cpp", llama_spec), ("whisper.cpp", whisper_spec)) if spec
)
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
_job.update(
state = _JOB_RUNNING,
message = "Downloading and installing the latest llama.cpp prebuilt...",
from_tag = from_tag,
message = f"Downloading and installing the latest {running} prebuilt...",
from_tag = (llama_spec or {}).get("from_tag"),
to_tag = None,
reload_required = None,
error = None,
progress = 0.0,
started_at = _utcnow(),
finished_at = None,
phases = None,
)
job_snapshot = dict(_job)
thread = threading.Thread(
target = _run_update,
args = (install_dir, repo, asset, script, pin_release_tag, force_cpu),
target = _flow.run_chained_update,
args = (phases,),
kwargs = {"job": _job, "job_lock": _job_lock},
name = "llama-cpp-update",
daemon = True,
)
@ -755,15 +734,4 @@ def start_update() -> dict:
def _reset_job_for_tests() -> None:
"""Test-only: return the job tracker to idle."""
with _job_lock:
_job.update(
state = _JOB_IDLE,
message = "",
from_tag = None,
to_tag = None,
reload_required = None,
error = None,
progress = None,
started_at = None,
finished_at = None,
)
_flow.reset_job(_job, _job_lock)

11
studio/backend/utils/prebuilt/__init__.py Normal file
View file

@ -0,0 +1,11 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Backend-importable prebuilt helpers.
The installers reuse install_llama_prebuilt.py directly; this package holds the
backend-side shapes the studio/ scripts cannot provide (the backend runs with
studio/backend as its sys.path root): runtime_libs (wheel CUDA dirs), child_env
(secret scrubbing + WSL ROCm dirs), freshness_flow and update_flow (the shared
mechanics behind the *_cpp_freshness / *_cpp_update twins).
"""

145
studio/backend/utils/prebuilt/child_env.py Normal file
View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Child-process environment hygiene for the managed ggml servers.
Secret-env scrubbing and the WSL2 ROCm library-dir probe, shared by the STT
sidecar (and any future launcher of a downloaded binary). Kept in sync with
install_llama_prebuilt.py's scrub_env / _wsl_system_rocm_lib_dirs; the backend
cannot import the studio/ installer scripts, so this copy stays importable with
only the backend root on sys.path.
"""
from __future__ import annotations
import os
import re
from typing import Mapping
SECRET_ENV_EXACT = frozenset(
{
"HF_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"GH_TOKEN",
"GITHUB_TOKEN",
"WANDB_API_KEY",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN",
"GOOGLE_APPLICATION_CREDENTIALS",
"AZURE_CLIENT_SECRET",
"KUBECONFIG",
"SSH_AUTH_SOCK",
}
)
# Case-insensitive substring markers for names we do not enumerate (no bare "KEY").
SECRET_ENV_MARKERS = (
"TOKEN",
"SECRET",
"PASSWORD",
"PASSWD",
"PASSPHRASE",
"CREDENTIAL",
"PRIVATE_KEY",
"API_KEY",
)
# Proxy / index URLs embed creds in their value; the offline server never needs them.
SECRET_ENV_URL_NAMES = frozenset(
{
"HTTP_PROXY",
"HTTPS_PROXY",
"ALL_PROXY",
"FTP_PROXY",
"RSYNC_PROXY",
"PIP_INDEX_URL",
"PIP_EXTRA_INDEX_URL",
"UV_INDEX_URL",
"UV_DEFAULT_INDEX",
"UV_EXTRA_INDEX_URL",
}
)
# Also drop values with URL userinfo creds (scheme://user:secret@host).
URL_USERINFO_RE = re.compile(r"://[^/@\s]+@")
def is_secret_env_name(name: str) -> bool:
upper = name.upper()
return (
upper in SECRET_ENV_EXACT
or upper in SECRET_ENV_URL_NAMES
or any(marker in upper for marker in SECRET_ENV_MARKERS)
)
def scrub_env(env: Mapping[str, str]) -> dict[str, str]:
"""Copy of ``env`` without secret-bearing names or URL-userinfo values."""
return {
k: v
for k, v in env.items()
if not is_secret_env_name(k) and not URL_USERINFO_RE.search(v or "")
}
# Filesystem pointers a downloaded binary could follow to on-disk credential
# stores (token caches under $HF_HOME, ~/.netrc, XDG config). Dropped, not
# repointed; the offline inference server needs none. Mirrors the cred-location
# list of the tools bypass env (core/inference/tools.py).
CRED_LOCATION_ENV_NAMES = frozenset(
{
"HF_HOME",
"HF_HUB_CACHE",
"HUGGINGFACE_HUB_CACHE",
"HF_XET_CACHE",
"TRANSFORMERS_CACHE",
"HF_DATASETS_CACHE",
"XDG_CONFIG_HOME",
"XDG_CACHE_HOME",
"XDG_DATA_HOME",
"NETRC",
"BASH_ENV",
"GIT_CONFIG_GLOBAL",
"GIT_CONFIG_SYSTEM",
"GIT_ASKPASS",
"SSH_ASKPASS",
"HOMEDRIVE",
"HOMEPATH",
}
)
# Home dirs are repointed (not dropped): loaders and SDKs expect them present,
# but they must not resolve to the user's real profile with its token caches.
HOME_ENV_NAMES = ("HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA")
def isolate_home(env: dict[str, str], scratch_dir: str) -> dict[str, str]:
"""Repoint home/profile vars at ``scratch_dir`` and drop credential-store
pointers so a compromised downloaded server cannot read token caches or cred
files through the environment. Mutates and returns ``env``."""
os.makedirs(scratch_dir, exist_ok = True)
for name in HOME_ENV_NAMES:
if name in env:
env[name] = scratch_dir
for name in CRED_LOCATION_ENV_NAMES:
env.pop(name, None)
return env
def wsl_system_rocm_lib_dirs() -> list[str]:
"""System ROCm lib dir(s) to load before a bundle's HIP on WSL2. Strict no-op
off WSL (needs /dev/dxg, a "microsoft" /proc/version, and a librocdxg)."""
try:
if not os.path.exists("/dev/dxg"):
return []
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
except OSError:
return []
dirs: list[str] = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists(
os.path.join(d, "librocdxg.so.1")
):
dirs.append(d)
return dirs

325
studio/backend/utils/prebuilt/freshness_flow.py Normal file
View file

@ -0,0 +1,325 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared mechanics of the llama.cpp / whisper.cpp prebuilt freshness checks.
The component modules (utils.llama_cpp_freshness / utils.whisper_cpp_freshness)
keep their public names, per-module caches, and version-comparison policy;
everything mechanical (marker walk-up, GitHub release fetch, memo + disk cache,
the freshness report skeleton) lives here, parameterized by call-time callables
so the modules' monkeypatch seams keep working.
"""
from __future__ import annotations
import json
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Optional
import structlog
logger = structlog.get_logger(__name__)
# 24h TTL keeps the GitHub call off the hot path and within rate limits.
RELEASE_CACHE_TTL_SECONDS = 24 * 60 * 60
def read_install_marker(
binary_path: Optional[str],
*,
marker_name: str,
cache: dict[str, Optional[dict]],
log_message: str,
) -> Optional[dict]:
"""Walk up from binary_path to find the install marker JSON.
None = no marker (source build / custom path) or invalid JSON."""
if not binary_path:
return None
cached = cache.get(binary_path)
if cached is not None or binary_path in cache:
return cached
p = Path(binary_path)
marker: Optional[dict] = None
# Cover all managed binary layouts (binary is 1-4 dirs deep).
for parent in p.parents[:5]:
candidate = parent / marker_name
if candidate.is_file():
try:
marker = json.loads(candidate.read_text(encoding = "utf-8"))
except (OSError, json.JSONDecodeError) as exc:
logger.debug(log_message, path = str(candidate), error = str(exc))
marker = None
break
cache[binary_path] = marker
return marker
def cache_path_for(repo: str, cache_dir: Path) -> Path:
safe = repo.replace("/", "__")
return cache_dir / f"{safe}.json"
def load_disk_cache(repo: str, cache_dir: Path) -> Optional[tuple[float, Optional[str]]]:
path = cache_path_for(repo, cache_dir)
try:
payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, json.JSONDecodeError):
return None
ts = payload.get("fetched_at")
tag = payload.get("latest_tag")
if not isinstance(ts, (int, float)):
return None
return float(ts), tag if isinstance(tag, str) else None
def save_disk_cache(
repo: str, latest_tag: Optional[str], cache_dir: Path, *, log_message: str
) -> None:
path = cache_path_for(repo, cache_dir)
try:
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".tmp")
tmp.write_text(
json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}),
encoding = "utf-8",
)
tmp.replace(path)
except OSError as exc:
logger.debug(log_message, repo = repo, error = str(exc))
def _fetch_newest_published_release(
repo: str, timeout: float, *, log_message: str
) -> Optional[dict]:
"""Newest published (non-draft/non-prerelease) release object for `repo`, by
``published_at``.
Resolves "latest" the way the installers do, NOT via GitHub's
``/releases/latest`` pointer, which sorts by commit date and can lag the
build the installer installs (detection and apply then disagree -- the
downgrade/sticky-banner bug). None on any failure (offline, rate-limited)."""
import os
import urllib.error
import urllib.request
url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
headers = {
"Accept": "application/vnd.github+json",
"User-Agent": "unsloth-studio-freshness-check",
}
token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
if token:
headers["Authorization"] = f"Bearer {token}"
req = urllib.request.Request(url, headers = headers)
try:
with urllib.request.urlopen(req, timeout = timeout) as resp:
data = json.loads(resp.read().decode("utf-8"))
except (
urllib.error.URLError,
urllib.error.HTTPError,
OSError,
json.JSONDecodeError,
) as exc:
logger.debug(log_message, repo = repo, error = str(exc))
return None
if not isinstance(data, list):
return None
published = [
r
for r in data
if isinstance(r, dict)
and not r.get("draft")
and not r.get("prerelease")
and isinstance(r.get("tag_name"), str)
and r.get("tag_name")
]
if not published:
return None
return max(published, key = lambda r: r.get("published_at") or "")
def fetch_latest_release_tag(
repo: str,
timeout: float = 5.0,
*,
log_message: str,
) -> Optional[str]:
"""Newest published release tag for `repo`, by publish time. None on failure."""
newest = _fetch_newest_published_release(repo, timeout, log_message = log_message)
return newest["tag_name"] if newest else None
def fetch_latest_release_assets(
repo: str,
timeout: float = 5.0,
*,
log_message: str,
) -> Optional[dict[str, int]]:
"""Asset name -> size (bytes) for the newest published release of `repo`,
selected exactly like fetch_latest_release_tag. None on any failure."""
newest = _fetch_newest_published_release(repo, timeout, log_message = log_message)
if newest is None:
return None
assets: dict[str, int] = {}
for a in newest.get("assets") or []:
name, size = a.get("name"), a.get("size")
if isinstance(name, str) and isinstance(size, int):
assets[name] = size
return assets
def latest_published_release(
repo: str,
*,
force_refresh: bool,
memo: dict[str, tuple[float, Optional[str]]],
cache_dir: Callable[[], Path],
fetch: Callable[[str], Optional[str]],
save: Callable[[str, Optional[str]], None],
) -> Optional[str]:
"""Latest release tag for `repo`. Memo + disk-cached (24h TTL).
None when offline and never previously cached."""
if not repo:
return None
now = time.time()
if not force_refresh:
cached = memo.get(repo)
if cached and now - cached[0] < RELEASE_CACHE_TTL_SECONDS:
return cached[1]
disk = load_disk_cache(repo, cache_dir())
if disk and now - disk[0] < RELEASE_CACHE_TTL_SECONDS:
memo[repo] = disk
return disk[1]
latest = fetch(repo)
if latest is None:
# Keep the last-good disk value rather than poison it with None.
disk = load_disk_cache(repo, cache_dir())
if disk:
memo[repo] = disk
return disk[1]
return None
memo[repo] = (now, latest)
save(repo, latest)
return latest
def latest_release_assets(
repo: str,
*,
force_refresh: bool,
memo: dict[str, tuple[float, dict[str, int]]],
fetch: Callable[[str], Optional[dict[str, int]]],
) -> Optional[dict[str, int]]:
"""Newest-release asset sizes for `repo`, memoized (24h TTL). None when
offline and never fetched. In-memory only -- a restart re-fetches."""
if not repo:
return None
now = time.time()
if not force_refresh:
cached = memo.get(repo)
if cached and now - cached[0] < RELEASE_CACHE_TTL_SECONDS:
return cached[1]
assets = fetch(repo)
if assets is None:
cached = memo.get(repo)
return cached[1] if cached else None
memo[repo] = (now, assets)
return assets
def parse_installed_at(value: object) -> Optional[datetime]:
if not isinstance(value, str) or not value:
return None
s = value.replace("Z", "+00:00") if value.endswith("Z") else value
try:
dt = datetime.fromisoformat(s)
except ValueError:
return None
if dt.tzinfo is None:
dt = dt.replace(tzinfo = timezone.utc)
return dt
def check_freshness(
binary_path: Optional[str],
*,
threshold_days: int,
now: Optional[datetime],
read_marker: Callable[[Optional[str]], Optional[dict]],
latest_release: Callable[[str], Optional[str]],
behind: Callable[[Optional[str], Optional[str]], bool],
display_tag: Callable[[dict], Any],
compare_tag: Callable[[dict], Any],
) -> dict:
"""Freshness report skeleton shared by both components; the component's
marker-tag choice and is_behind policy come in as callables. Fails open on
missing data (behind/stale stay False)."""
out: dict = {
"has_marker": False,
"stale": False,
"behind": False,
"installed_tag": None,
"latest_tag": None,
"installed_at_utc": None,
"age_days": None,
"published_repo": None,
"threshold_days": int(threshold_days),
}
marker = read_marker(binary_path)
if not marker:
return out
out["has_marker"] = True
out["installed_tag"] = display_tag(marker)
out["installed_at_utc"] = marker.get("installed_at_utc")
out["published_repo"] = marker.get("published_repo")
installed_full = compare_tag(marker)
repo = out["published_repo"]
if not repo or not installed_full:
return out
latest = latest_release(repo)
out["latest_tag"] = latest
out["behind"] = behind(installed_full, latest)
if not out["behind"]:
return out
installed_at = parse_installed_at(out["installed_at_utc"])
if installed_at is None:
return out
now = now or datetime.now(tz = timezone.utc)
age_seconds = (now - installed_at).total_seconds()
out["age_days"] = max(0, int(age_seconds // 86400))
if age_seconds >= threshold_days * 86400:
out["stale"] = True
return out
def format_stale_warning(info: dict, *, component: str) -> str:
"""Human-readable one-liner for stale prebuilt info."""
age = info.get("age_days")
installed = info.get("installed_tag") or "unknown"
latest = info.get("latest_tag") or "unknown"
age_str = f"{age} day{'s' if age != 1 else ''}" if age is not None else "some time"
return (
f"{component} prebuilt is {age_str} behind: installed "
f"{installed}, latest {latest}. Run `unsloth studio update` "
f"to refresh."
)
def reset_caches(
caches: tuple[dict, ...], *, drop_disk: bool, cache_dir: Callable[[], Path]
) -> None:
"""Drop the in-memory freshness caches; with drop_disk also the on-disk 24h
release cache (see the component modules for why)."""
for cache in caches:
cache.clear()
if drop_disk:
import shutil
# cache_dir() is a dedicated freshness-only subdir, re-created on the next
# save_disk_cache. ignore_errors so a missing/locked dir is a no-op rather
# than breaking an otherwise successful install.
shutil.rmtree(cache_dir(), ignore_errors = True)

65
studio/backend/utils/prebuilt/runtime_libs.py Normal file
View file

@ -0,0 +1,65 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""CUDA runtime dirs shipped inside Python wheels, for the STT sidecar's child env.
Kept in sync with install_llama_prebuilt.py's python_runtime_dirs; the backend
cannot import the studio/ installer scripts, so this small copy stays importable
with only the backend root on sys.path.
"""
from __future__ import annotations
import site
import sys
from pathlib import Path
from typing import Iterable
def dedupe_existing_dirs(paths: Iterable[str | Path]) -> list[str]:
unique: list[str] = []
seen: set[str] = set()
for raw in paths:
if not raw:
continue
try:
path = Path(raw).expanduser()
if not path.is_dir():
continue
resolved = str(path.resolve())
except (OSError, ValueError):
continue
if resolved in seen:
continue
seen.add(resolved)
unique.append(resolved)
return unique
def python_runtime_dirs() -> list[str]:
"""CUDA runtime dirs shipped inside Python wheels (torch + nvidia-* wheels)."""
candidates: list[Path] = []
search_roots = [Path(entry) for entry in sys.path if entry]
try:
search_roots.extend(Path(path) for path in site.getsitepackages())
except Exception:
pass
try:
user_site = site.getusersitepackages()
if user_site:
search_roots.append(Path(user_site))
except Exception:
pass
for root in search_roots:
if not root.is_dir():
continue
candidates.extend(root.glob("nvidia/*/lib")) # Linux convention
candidates.extend(root.glob("nvidia/*/bin")) # legacy modular Windows wheels
candidates.extend(root.glob("nvidia/*/bin/x86_64")) # CUDA 13 Windows wheel layout
candidates.extend(root.glob("nvidia/*/bin/x64"))
candidates.extend(root.glob("nvidia/*/Library/bin")) # conda-style repacks
candidates.extend(root.glob("nvidia/*/Library/bin/x86_64"))
candidates.extend(root.glob("nvidia/*/Library/bin/x64"))
candidates.extend(root.glob("torch/lib"))
return dedupe_existing_dirs(candidates)

447
studio/backend/utils/prebuilt/update_flow.py Normal file
View file

@ -0,0 +1,447 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared mechanics of the llama.cpp / whisper.cpp in-app prebuilt updates.
The component modules (utils.llama_cpp_update / utils.whisper_cpp_update) keep
their public names, job dicts, and update policy (version comparison, pinning,
pre/post install steps); everything mechanical (managed-root resolution,
local-link detection, the resolve probe, the streamed installer run) lives here,
parameterized so the modules' monkeypatch seams keep working.
"""
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Callable, Optional
import structlog
from utils.process_lifetime import child_popen_kwargs
logger = structlog.get_logger(__name__)
# Markerless (source-build) resolve answers are memoized for 24h; only
# successful answers are cached so a network blip retries.
RESOLVE_TTL_SECONDS = 24 * 60 * 60
# Matches the installer's download progress lines, e.g.
# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s".
PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(")
# The download dominates the update; extract/validate fill the last slice.
DOWNLOAD_PROGRESS_CEILING = 0.95
class InstallerExit(RuntimeError):
"""Installer subprocess exited nonzero; carries the exit code so phase
runners can special-case contractual codes (whisper's 2 = unavailable)."""
def __init__(self, returncode: int, message: str) -> None:
super().__init__(message)
self.returncode = returncode
JOB_IDLE = "idle"
JOB_RUNNING = "running"
JOB_SUCCESS = "success"
JOB_ERROR = "error"
# Per-phase states inside a chained job's "phases" breakdown.
PHASE_PENDING = "pending"
PHASE_RUNNING = "running"
PHASE_SUCCESS = "success"
PHASE_ERROR = "error"
PHASE_SKIPPED = "skipped"
_IDLE_JOB_FIELDS = dict(
state = JOB_IDLE,
message = "",
from_tag = None,
to_tag = None,
reload_required = None,
error = None,
progress = None,
started_at = None,
finished_at = None,
phases = None,
)
def new_job() -> dict:
"""A fresh idle job-state dict (one per component module)."""
return dict(_IDLE_JOB_FIELDS)
def reset_job(job: dict, job_lock: threading.Lock) -> None:
"""Return a job tracker to idle (test seam)."""
with job_lock:
job.update(_IDLE_JOB_FIELDS)
def utcnow() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
def is_under(path: Path, root: Path) -> bool:
try:
p, r = path.resolve(), root.resolve()
except (OSError, ValueError):
p, r = path, root
return p == r or r in p.parents
def install_dir_for(binary_path: Optional[str], *, marker_name: str) -> Optional[Path]:
"""The directory holding the install marker: the install root the installer
wrote and the one we re-install into. Walks up from the binary like the
freshness marker reader does."""
if not binary_path:
return None
p = Path(binary_path)
for parent in p.parents[:5]:
if (parent / marker_name).is_file():
return parent
return None
def find_installer_script(*, env_var: str, script_name: str) -> Optional[Path]:
"""Locate the installer script. Honours the env override, then searches up
from this file for both ``<root>/<script>`` and ``<root>/studio/<script>`` so
it works in the dev tree and in an installed Unsloth layout."""
env = os.environ.get(env_var)
if env and Path(env).is_file():
return Path(env)
here = Path(__file__).resolve()
for up in here.parents:
for cand in (up / script_name, up / "studio" / script_name):
if cand.is_file():
return cand
return None
def resolve_prebuilt_for_host(
*,
force_refresh: bool,
memo: dict,
installer_script: Callable[[], Optional[Path]],
log_message: str,
extra_args: tuple[str, ...] = (),
) -> Optional[dict]:
"""Run ``<installer> --resolve-prebuilt latest --output-format json`` (no
download); return the parsed payload or None. Fail-open: any error -> None so
a source build never blocks the app."""
now = time.time()
cache_key = tuple(extra_args)
if not force_refresh and memo.get("key") == cache_key:
if now - memo.get("at", 0.0) < RESOLVE_TTL_SECONDS:
return memo.get("value")
script = installer_script()
if script is None:
return None
value: Optional[dict] = None
try:
cmd = [
sys.executable,
str(script),
"--resolve-prebuilt",
"latest",
"--output-format",
"json",
*extra_args,
]
proc = subprocess.run(
cmd,
capture_output = True,
text = True,
timeout = 60,
)
out = (proc.stdout or "").strip()
if proc.returncode == 0 and out:
parsed = json.loads(out.splitlines()[-1])
if isinstance(parsed, dict):
value = parsed
except Exception as exc: # pragma: no cover - subprocess/json defensive
logger.debug(log_message, error = str(exc))
value = None
if value is not None: # cache real answers; let failures retry next poll
memo.update(at = now, key = cache_key, value = value)
return value
def is_external_link(path: Optional[Path]) -> bool:
"""True when ``path`` is a locally-linked component dir: a POSIX symlink or a
Windows junction / reparse point. Such a link resolves into the user's own
checkout, so Unsloth must never auto-update it."""
if path is None:
return False
try:
if os.path.islink(path):
return True
except OSError:
return False
if os.name == "nt":
try:
import stat
attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined]
return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT)
except (OSError, AttributeError):
return False
return False
def active_install_is_local_link(binary: Optional[str], *, dir_name: str) -> bool:
"""True when the active server binary resolves through a locally-linked
component directory. An update would write through that link into the user's
checkout (or fail), so the install is treated as externally managed: none is
offered or applied. Checks only up to and including the component dir so a
symlinked HOME / studio root above it can't trip a false positive."""
if not binary:
return False
for parent in Path(binary).parents:
if is_external_link(parent):
return True
if parent.name == dir_name:
break
return False
def managed_install_root(
binary: Optional[str],
*,
marker_root: Optional[Path],
server_path_var: str,
cpp_path_var: str,
dir_name: str,
) -> Optional[Path]:
"""The Unsloth-managed component root the active binary lives under, or None
when unmanaged. Installing where the active binary is not would not replace
what discovery runs (a pinned server path, then the custom dir, then a
component tree), so we refuse rather than install into an inactive or foreign
tree."""
if marker_root is not None:
return marker_root
if not binary:
return None
# The server-path pin is an explicit user choice that wins in discovery; never
# auto-replace its tree (even the user's own checkout).
if os.environ.get(server_path_var):
return None
p = Path(binary)
env = os.environ.get(cpp_path_var)
if env and is_under(p, Path(env)):
return Path(env)
for parent in p.parents:
if parent.name == dir_name:
return parent
# PATH / system / custom install: not a managed tree, so do not offer.
return None
def local_link_status(job: dict, job_lock: threading.Lock) -> dict:
"""Status payload for a local-link install: unmanaged, no update offered."""
with job_lock:
snapshot = dict(job)
return {
"supported": False,
"update_available": False,
"stale": False,
"installed_tag": None,
"latest_tag": None,
"published_repo": None,
"installed_at_utc": None,
"age_days": None,
"source_build": False,
"local_link": True,
"update_size_bytes": None,
"job": snapshot,
}
def rocm_install_args(asset: Optional[str]) -> list[str]:
"""Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh.
The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx ROCm
bundles carry the family in the name (rocm-gfx110X), version-tagged bundles
only rocm/hip."""
if not asset:
return []
low = asset.lower()
if "rocm" not in low and "hip" not in low:
return []
gfx = re.search(r"-gfx[0-9a-z]+", low)
if gfx:
return ["--rocm-gfx", gfx.group(0).lstrip("-")]
return ["--has-rocm"]
def stream_installer(
cmd: list[str],
env: dict[str, str],
*,
timeout_seconds: int,
job: Optional[dict] = None,
job_lock: Optional[threading.Lock] = None,
set_progress: Optional[Callable[[float], None]] = None,
) -> None:
"""Run the installer, streaming its progress lines into job["progress"]
(or through set_progress when given, e.g. a chained-phase progress window).
Raises RuntimeError on timeout or a nonzero exit (with an output tail)."""
if set_progress is None:
assert job is not None and job_lock is not None
def set_progress(fraction: float) -> None:
with job_lock:
job["progress"] = max(job.get("progress") or 0.0, fraction)
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = env,
**child_popen_kwargs(),
)
timed_out = threading.Event()
def _kill_on_timeout() -> None:
timed_out.set()
proc.kill()
watchdog = threading.Timer(timeout_seconds, _kill_on_timeout)
watchdog.daemon = True
watchdog.start()
tail_lines: list[str] = []
try:
assert proc.stdout is not None
for line in proc.stdout:
tail_lines.append(line)
if len(tail_lines) > 80:
del tail_lines[0]
m = PROGRESS_LINE_RE.search(line)
if m is None:
continue
set_progress(min(float(m.group(1)) / 100.0, 1.0) * DOWNLOAD_PROGRESS_CEILING)
returncode = proc.wait()
finally:
watchdog.cancel()
if timed_out.is_set():
raise RuntimeError(f"installer timed out after {timeout_seconds}s")
if returncode != 0:
tail = "".join(tail_lines).strip()[-1500:]
raise InstallerExit(returncode, f"installer exited {returncode}: {tail or 'no output'}")
def _new_phase_record(spec: dict) -> dict:
"""Initial breakdown entry for one phase of a chained job."""
runnable = spec.get("run") is not None
return {
"state": PHASE_PENDING if runnable else PHASE_SKIPPED,
"reason": None if runnable else spec.get("skip_reason"),
"progress": None,
"to_tag": None,
"reload_required": None,
"message": "",
"error": None,
}
def run_chained_update(phases: list[dict], *, job: dict, job_lock: threading.Lock) -> None:
"""Run update phases in order into one shared job dict (the worker of a
combined llama+whisper apply).
Each phase spec: ``name`` (breakdown key), ``weight`` (progress slice,
normalized over runnable phases), ``run`` (callable(set_progress) -> result
dict with to_tag/reload_required/message, raises on failure; None = skipped)
and ``skip_reason`` / ``failure_message``. A failing phase aborts the chain:
later phases are marked skipped (reason "aborted") and the job goes to error,
keeping the reload_required and messages of already-succeeded phases so a
partial success stays visible."""
runnable = [p for p in phases if p.get("run") is not None]
total_weight = sum(float(p.get("weight") or 1.0) for p in runnable) or 1.0
with job_lock:
job["phases"] = {p["name"]: _new_phase_record(p) for p in phases}
offset = 0.0
done_messages: list[str] = []
reload_required = False
primary_to_tag: Optional[str] = None
for index, phase in enumerate(phases):
if phase.get("run") is None:
continue
name = phase["name"]
weight = float(phase.get("weight") or 1.0) / total_weight
with job_lock:
job["phases"][name].update(state = PHASE_RUNNING, progress = 0.0)
def set_progress(
fraction: float,
*,
_name: str = name,
_base: float = offset,
_slice: float = weight,
) -> None:
f = max(0.0, min(float(fraction), 1.0))
with job_lock:
record = job["phases"][_name]
record["progress"] = max(record.get("progress") or 0.0, f)
job["progress"] = max(job.get("progress") or 0.0, _base + f * _slice)
try:
result = phase["run"](set_progress) or {}
except Exception as exc:
failure = phase.get("failure_message") or f"{name} update failed."
with job_lock:
job["phases"][name].update(state = PHASE_ERROR, error = str(exc))
for later in phases[index + 1 :]:
if later.get("run") is not None:
job["phases"][later["name"]].update(state = PHASE_SKIPPED, reason = "aborted")
# A partial success keeps its messages and reload_required so the
# caller sees the earlier phase did land.
job.update(
state = JOB_ERROR,
message = " ".join(done_messages + [failure]),
to_tag = primary_to_tag,
error = str(exc),
finished_at = utcnow(),
)
if done_messages:
job["reload_required"] = reload_required
return
set_progress(1.0)
offset += weight
with job_lock:
job["phases"][name].update(
state = PHASE_SUCCESS,
to_tag = result.get("to_tag"),
reload_required = result.get("reload_required"),
message = result.get("message") or "",
)
if result.get("message"):
done_messages.append(result["message"])
# Only phases affecting the primary (llama) server may raise the job-level
# reload flag: the frontend resyncs chat model state off it, and a
# whisper-only sidecar reload must not clear the chat checkpoint. Per-phase
# reload_required stays visible under job["phases"].
if phase.get("affects_job_reload", True):
reload_required = reload_required or bool(result.get("reload_required"))
# The legacy job-level to_tag means "the llama build now installed";
# a whisper-only round must leave it unset or the UI reports a llama
# update that never ran (per-phase to_tag remains under phases).
if primary_to_tag is None:
primary_to_tag = result.get("to_tag")
with job_lock:
job.update(
state = JOB_SUCCESS,
message = " ".join(done_messages) or "Already up to date.",
to_tag = primary_to_tag,
reload_required = reload_required,
error = None,
progress = 1.0,
finished_at = utcnow(),
)

74
studio/backend/utils/prebuilt/whisper_layout.py Normal file
View file

@ -0,0 +1,74 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Canonical whisper.cpp install-root and marker lookup helpers."""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
MARKER_NAME = "UNSLOTH_WHISPER_PREBUILT_INFO.json"
@dataclass(frozen = True)
class MarkerLookup:
marker: Optional[dict]
root: Optional[Path]
authoritative: bool
invalid: bool = False
slim_collision: bool = False
def canonical_install_root(binary_path: Optional[str]) -> Optional[Path]:
"""Return the managed root for the two supported CMake binary layouts."""
if not binary_path:
return None
parent = Path(binary_path).parent
if parent.name == "bin" and parent.parent.name == "build":
return parent.parent.parent
if (
parent.name == "Release"
and parent.parent.name == "bin"
and parent.parent.parent.name == "build"
):
return parent.parent.parent.parent
return None
def _parse_marker(path: Path) -> tuple[Optional[dict], bool]:
try:
payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, json.JSONDecodeError):
return None, True
return (payload, False) if isinstance(payload, dict) else (None, True)
def lookup_marker(binary_path: Optional[str]) -> MarkerLookup:
"""Prefer the install-root marker over archive metadata beside the binary."""
if not binary_path:
return MarkerLookup(None, None, False)
binary = Path(binary_path)
root = canonical_install_root(binary_path)
inner_marker = binary.parent / MARKER_NAME
inner_payload, _ = _parse_marker(inner_marker) if inner_marker.is_file() else (None, False)
slim_collision = bool(
inner_payload
and (inner_payload.get("install_kind") == "slim" or inner_payload.get("backend") == "slim")
)
if root is not None:
root_marker = root / MARKER_NAME
if root_marker.is_file():
marker, invalid = _parse_marker(root_marker)
return MarkerLookup(marker, root, True, invalid, slim_collision)
if slim_collision:
return MarkerLookup(None, root, False, True, True)
for parent in binary.parents[:5]:
candidate = parent / MARKER_NAME
if candidate.is_file():
marker, invalid = _parse_marker(candidate)
return MarkerLookup(marker, parent, False, invalid, slim_collision)
return MarkerLookup(None, None, False)

3
studio/backend/utils/upload_limits.py
View file

@ -14,6 +14,9 @@ MIN_UPLOAD_LIMIT_MB = 1
MAX_UPLOAD_LIMIT_MB = 8192
_BYTES_PER_MB = 1024 * 1024
MULTIPART_OVERHEAD_BYTES = 10 * _BYTES_PER_MB
STT_AUDIO_RAW_MAX_BYTES = 25 * _BYTES_PER_MB
STT_AUDIO_B64_MAX_CHARS = ((STT_AUDIO_RAW_MAX_BYTES + 2) // 3) * 4
STT_AUDIO_JSON_MAX_BYTES = STT_AUDIO_B64_MAX_CHARS + 64 * 1024
LOCAL_SEED_UPLOAD_MAX_BYTES = 100 * _BYTES_PER_MB
LOCAL_SEED_UPLOAD_MAX_LABEL = "100MB"

254
studio/backend/utils/whisper_cpp_freshness.py Normal file
View file

@ -0,0 +1,254 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp prebuilt freshness check.
Reads UNSLOTH_WHISPER_PREBUILT_INFO.json (written by install_whisper_prebuilt.py)
and compares the installed release tag against the latest on GitHub. Surfaced via
utils.whisper_cpp_update (GET /api/whisper/update-status and the combined
llama+whisper update status). Fails open on any missing data so we never show a
misleading banner.
The mechanics (marker walk-up, GitHub fetch, memo + disk cache, report skeleton)
live in utils.prebuilt.freshness_flow; this module keeps the whisper version
policy and the per-module caches its tests patch.
"""
from __future__ import annotations
import re
from datetime import datetime
from pathlib import Path
from typing import Optional
import structlog
from utils.prebuilt import freshness_flow as _flow
from utils.prebuilt.whisper_layout import lookup_marker
logger = structlog.get_logger(__name__)
# 3 days matches Unsloth's typical whisper.cpp release cadence.
STALENESS_THRESHOLD_DAYS = 3
_INSTALL_MARKER_NAME = "UNSLOTH_WHISPER_PREBUILT_INFO.json"
_marker_cache: dict[str, Optional[dict]] = {}
_release_memo: dict[str, tuple[float, Optional[str]]] = {}
# Newest-release asset sizes (name -> bytes), memoized like the tag (24h TTL).
_assets_memo: dict[str, tuple[float, dict[str, int]]] = {}
def _cache_dir() -> Path:
"""Lazy import so tests can stub storage_roots."""
try:
from utils.paths.storage_roots import cache_root
return cache_root() / "whisper_cpp_freshness"
except Exception:
return Path.home() / ".unsloth" / "studio" / "cache" / "whisper_cpp_freshness"
def read_install_marker(binary_path: Optional[str]) -> Optional[dict]:
"""Walk up from binary_path to find UNSLOTH_WHISPER_PREBUILT_INFO.json.
None = no marker (source build / custom path) or invalid JSON."""
if not binary_path:
return None
cached = _marker_cache.get(binary_path)
if cached is not None or binary_path in _marker_cache:
return cached
marker = lookup_marker(binary_path).marker
_marker_cache[binary_path] = marker
return marker
def _load_disk_cache(repo: str) -> Optional[tuple[float, Optional[str]]]:
return _flow.load_disk_cache(repo, _cache_dir())
def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None:
_flow.save_disk_cache(
repo, latest_tag, _cache_dir(), log_message = "whisper freshness cache write failed"
)
def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
"""Newest published release tag for `repo`, by publish time (see
freshness_flow for why this is not GitHub's /releases/latest pointer)."""
return _flow.fetch_latest_release_tag(
repo, timeout, log_message = "whisper freshness fetch failed"
)
def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]:
"""Latest release tag for `repo`. Memo + disk-cached (24h TTL).
None when offline and never previously cached."""
return _flow.latest_published_release(
repo,
force_refresh = force_refresh,
memo = _release_memo,
cache_dir = lambda: _cache_dir(),
fetch = lambda r: _fetch_latest_release_tag(r),
save = lambda r, tag: _save_disk_cache(r, tag),
)
def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]:
"""Asset name -> size (bytes) for the newest published release of `repo`,
selected exactly like _fetch_latest_release_tag. None on any failure."""
return _flow.fetch_latest_release_assets(
repo, timeout, log_message = "whisper freshness asset fetch failed"
)
def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]:
"""Newest-release asset sizes for `repo`, memoized (24h TTL). None when
offline and never fetched. In-memory only -- a restart simply re-fetches."""
return _flow.latest_release_assets(
repo,
force_refresh = force_refresh,
memo = _assets_memo,
fetch = lambda r: _fetch_latest_release_assets(r),
)
def _asset_platform_suffix(asset: str, installed_tag: Optional[str]) -> Optional[str]:
"""The tag-independent ``<os>-<arch>-<accel>.<ext>`` suffix identifying this
host's bundle, from stripping the ``whisper-<tag>-`` prefix off the installed
asset name. Falls back to anchoring on the platform token when the marker tag
does not line up with the asset's embedded tag."""
if isinstance(installed_tag, str) and installed_tag:
tag = installed_tag if installed_tag.startswith("v") else f"v{installed_tag}"
prefix = f"whisper-{tag}-"
if asset.startswith(prefix):
return asset[len(prefix) :]
m = re.search(r"-((?:linux|macos|windows)-.*)$", asset)
return m.group(1) if m else None
def update_download_size_bytes(
marker: Optional[dict],
latest_tag: Optional[str],
repo: Optional[str],
*,
force_refresh: bool = False,
) -> Optional[int]:
"""Download size of the latest-release asset matching this host's installed
bundle (same platform/arch/accel suffix). None when there is no marker asset,
the latest assets can't be read, or no match.
Assets are named ``whisper-<tag>-<os>-<arch>-<accel>.<ext>`` and the fork
builds every slice itself, so only the publish repo is consulted (no
upstream/binary_repo passthrough)."""
if not marker or not latest_tag or not repo:
return None
installed_asset = marker.get("asset")
if not isinstance(installed_asset, str) or not installed_asset:
return None
suffix = _asset_platform_suffix(installed_asset, marker.get("release_tag"))
if not suffix:
return None
assets = latest_release_assets(repo, force_refresh = force_refresh)
if not assets:
return None
latest_v = latest_tag if latest_tag.startswith("v") else f"v{latest_tag}"
want = f"whisper-{latest_v}-{suffix}"
if want in assets:
return assets[want]
# Tag formatting can vary; fall back to the platform+accel suffix.
for name, size in assets.items():
if name.endswith(suffix):
return size
return None
def parse_release_version(tag: object) -> Optional[tuple]:
"""Comparable key ``(upstream_major, upstream_minor, upstream_patch,
unsloth_serial)`` for a tag like ``v1.9.1-unsloth.2``.
Tolerant of a leading ``v`` and a missing ``-unsloth.N`` suffix (serial then
0); the upstream version is padded to major/minor/patch. None when the version
component is not purely numeric."""
if not isinstance(tag, str):
return None
s = tag.strip()
if not s:
return None
if s[0] in ("v", "V"):
s = s[1:]
serial = 0
m = re.search(r"-unsloth\.(\d+)$", s)
if m:
serial = int(m.group(1))
s = s[: m.start()]
parts = s.split(".")
nums: list[int] = []
for part in parts:
if not part.isdigit():
return None
nums.append(int(part))
if not nums:
return None
while len(nums) < 3:
nums.append(0)
return (nums[0], nums[1], nums[2], serial)
def is_behind(installed: Optional[str], latest: Optional[str]) -> bool:
"""Whether `installed` is genuinely behind `latest`.
- identical tags -> not behind (clears the sticky banner post-update)
- both parse -> behind only when the latest version key is strictly greater
(a lower/equal version is never "behind" -> downgrade guard)
- either fails to parse and they differ -> behind (plain inequality)
"""
if not installed or not latest:
return False
installed, latest = installed.strip(), latest.strip()
if installed == latest:
return False
installed_key, latest_key = parse_release_version(installed), parse_release_version(latest)
if installed_key is not None and latest_key is not None:
return latest_key > installed_key
# One side is unparseable and the tags already differ: treat as behind.
return True
def check_prebuilt_freshness(
binary_path: Optional[str],
*,
threshold_days: int = STALENESS_THRESHOLD_DAYS,
now: Optional[datetime] = None,
) -> dict:
"""Returns {has_marker, stale, behind, installed_tag, latest_tag,
installed_at_utc, age_days, published_repo, threshold_days}.
behind = installed genuinely older than latest (see is_behind).
stale = behind AND age >= threshold.
Fails open on missing data (behind/stale stay False)."""
# The whisper marker records a single ``release_tag`` (e.g. v1.9.1-unsloth.2);
# both display and comparison use it directly.
return _flow.check_freshness(
binary_path,
threshold_days = threshold_days,
now = now,
read_marker = lambda p: read_install_marker(p),
latest_release = lambda repo: latest_published_release(repo),
behind = lambda installed, latest: is_behind(installed, latest),
display_tag = lambda marker: marker.get("release_tag"),
compare_tag = lambda marker: marker.get("release_tag"),
)
def reset_caches(*, drop_disk: bool = False) -> None:
"""Drop the in-memory freshness caches. The no-arg form is test-only.
With ``drop_disk = True`` also delete the on-disk 24h release cache. Used by
the post-install/update path: clearing memory alone leaves the stale
same-version value on disk, so an offline post-install GitHub refresh would
replay it (latest_published_release's last-good fallback) and the banner could
linger. Dropping the disk cache makes latest read None in that offline case,
so the banner fails open (off) rather than point at the just-replaced build."""
_flow.reset_caches(
(_marker_cache, _release_memo, _assets_memo),
drop_disk = drop_disk,
cache_dir = lambda: _cache_dir(),
)

490
studio/backend/utils/whisper_cpp_update.py Normal file
View file

@ -0,0 +1,490 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""In-app whisper.cpp prebuilt update.
Builds on utils.whisper_cpp_freshness (which detects whether a newer prebuilt
release exists) and adds the *apply* half: run install_whisper_prebuilt.py to
download the newest bundle for this host and atomically swap it in, so the next
model load uses it. Applies only run as the whisper phase of the combined
llama+whisper update (utils.llama_cpp_update.start_update chains
run_chained_phase); there is no standalone whisper update trigger.
Design notes:
- Detection is delegated to check_prebuilt_freshness(). We surface an
``update_available`` flag (installed_tag != latest_tag), laxer than freshness'
``stale`` (which also requires the install to be >= 3 days old). The UI shows
the single main update item on update_available.
- Everything fails open: a missing marker / offline GitHub / source build just
reports update_available=False and never blocks the app.
- The mechanics (managed-root resolution, local-link detection, the resolve
probe, the streamed installer run) live in utils.prebuilt.update_flow; this
module keeps the whisper policy. The ``job`` dict in the status payload stays
idle (kept for response-shape stability): chained applies report progress
through the llama job.
"""
from __future__ import annotations
import os
import re
import subprocess
import sys
import threading
from pathlib import Path
from typing import Optional
import structlog
from utils.prebuilt import update_flow as _flow
from utils.prebuilt.whisper_layout import canonical_install_root
from utils.whisper_cpp_freshness import (
_INSTALL_MARKER_NAME,
check_prebuilt_freshness,
is_behind,
latest_published_release,
latest_release_assets,
parse_release_version,
read_install_marker,
reset_caches,
update_download_size_bytes,
)
logger = structlog.get_logger(__name__)
DEFAULT_PUBLISHED_REPO = "unslothai/whisper.cpp"
_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + extract/validate
# Always-idle job payload: whisper applies run inside the chained llama job, so
# nothing flips this to running. Kept so status payload shapes stay stable.
_job_lock = threading.Lock()
_job: dict = _flow.new_job()
_rocm_install_args = _flow.rocm_install_args
def _find_binary() -> Optional[str]:
"""Locate the active whisper-server binary via the STT sidecar's own resolver
so update targets exactly what Unsloth runs. Lazy import keeps the heavy
inference module off this module's import path."""
try:
from core.inference.stt_ggml_sidecar import find_whisper_server_binary
return find_whisper_server_binary()
except Exception as exc: # pragma: no cover - defensive
logger.debug("whisper update: binary discovery failed", error = str(exc))
return None
def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]:
"""The directory holding UNSLOTH_WHISPER_PREBUILT_INFO.json: the install root
install_whisper_prebuilt.py wrote (``<install-dir>`` whose canonical server is
``build/bin/whisper-server``) and the one we re-install into."""
root = canonical_install_root(binary_path)
if root is not None and (root / _INSTALL_MARKER_NAME).is_file():
return root
return _flow.install_dir_for(binary_path, marker_name = _INSTALL_MARKER_NAME)
def _installer_script() -> Optional[Path]:
"""Locate install_whisper_prebuilt.py (UNSLOTH_WHISPER_INSTALLER wins)."""
return _flow.find_installer_script(
env_var = "UNSLOTH_WHISPER_INSTALLER", script_name = "install_whisper_prebuilt.py"
)
# Markerless (source-build) installs have no UNSLOTH_WHISPER_PREBUILT_INFO.json,
# so we ask the installer whether an official prebuilt now exists for this host.
_resolve_memo: dict = {}
def _resolve_prebuilt_for_host(
*, force_refresh: bool = False, backend: Optional[str] = None
) -> Optional[dict]:
"""Run install_whisper_prebuilt.py --resolve-prebuilt (no download); return
{prebuilt_available, repo, release_tag, upstream_tag, backend, asset, os,
arch, ...} or None. Fail-open: any error -> None so a source build never
blocks the app."""
extra_args = ("--backend", backend) if backend else ()
return _flow.resolve_prebuilt_for_host(
force_refresh = force_refresh,
memo = _resolve_memo,
installer_script = lambda: _installer_script(),
log_message = "whisper update: resolve-prebuilt failed",
extra_args = extra_args,
)
def _installed_whisper_version(binary: Optional[str]) -> Optional[str]:
"""Best-effort ``v<A.B.C>`` from ``whisper-server --version``. None when the
binary is missing, reports no version, or cannot run. Used only for the
markerless source-build downgrade guard, so it fails open to None."""
if not binary:
return None
try:
proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20)
except Exception: # pragma: no cover - defensive
return None
m = re.search(r"v?(\d+\.\d+\.\d+)", (proc.stderr or "") + (proc.stdout or ""))
if not m:
return None
return f"v{m.group(1)}"
def _whisper_install_root(binary: Optional[str]) -> Optional[Path]:
"""The Unsloth-managed whisper.cpp root the active binary lives under, or None
when the binary is unmanaged (see update_flow.managed_install_root)."""
return _flow.managed_install_root(
binary,
marker_root = _install_dir_for(binary),
server_path_var = "WHISPER_SERVER_PATH",
cpp_path_var = "UNSLOTH_WHISPER_CPP_PATH",
dir_name = "whisper.cpp",
)
def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
"""Update status for a markerless (source-build) install: offer the official
prebuilt when one exists for this host and is newer than the installed binary.
None -> caller falls through to the no-marker default (unsupported)."""
res = _resolve_prebuilt_for_host(force_refresh = force_refresh)
if not res or not res.get("prebuilt_available"):
return None
release_tag = res.get("release_tag")
if not release_tag:
return None
# No resolvable install root (e.g. a pinned WHISPER_SERVER_PATH we cannot
# manage) means an apply would not take effect, so do not offer it.
if _whisper_install_root(binary) is None:
return None
installed_tag = _installed_whisper_version(binary)
installed_key = parse_release_version(installed_tag) if installed_tag else None
latest_key = parse_release_version(release_tag)
if installed_key is None or latest_key is None:
# Unknown installed/latest version (involuntary source-build case): treat
# as behind so we still offer the prebuilt.
update_available = True
else:
# Same downgrade guard as is_behind: only a strictly newer key is behind.
update_available = latest_key > installed_key
latest = release_tag
# Size of the resolved prebuilt, so source builds show it like the marker
# path. Fails open to None (offline / asset absent from release).
update_size_bytes = None
if update_available:
asset_name = res.get("asset")
if isinstance(asset_name, str) and asset_name:
try:
assets = latest_release_assets(res.get("repo"), force_refresh = force_refresh)
if assets:
update_size_bytes = assets.get(asset_name)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("whisper update: source-build size lookup failed", error = str(exc))
with _job_lock:
job = dict(_job)
return {
"supported": True,
"update_available": update_available,
"stale": False,
"installed_tag": installed_tag,
"latest_tag": latest,
"published_repo": res.get("repo"),
"installed_at_utc": None,
"age_days": None,
"source_build": True,
"update_size_bytes": update_size_bytes,
"job": job,
}
def _active_install_is_local_link(binary: Optional[str]) -> bool:
"""True when the active whisper-server resolves through a locally-linked
whisper.cpp directory (see update_flow.active_install_is_local_link)."""
return _flow.active_install_is_local_link(binary, dir_name = "whisper.cpp")
def _local_link_status() -> dict:
"""Status payload for a local-link install: unmanaged, no update offered."""
return _flow.local_link_status(_job, _job_lock)
def get_update_status(*, force_refresh: bool = False) -> dict:
"""Report whether a newer prebuilt exists plus the current job state.
force_refresh bypasses the 24h release cache for an explicit "check now".
"""
binary = _find_binary()
# A locally-linked whisper.cpp dir is the user's own tree; never offer to
# replace it. Bail before any network/freshness work.
if _active_install_is_local_link(binary):
return _local_link_status()
marker = read_install_marker(binary)
# No marker = source build / custom path. Offer the official prebuilt if one
# exists for this host.
if marker is None and binary is not None:
src = _source_build_status(binary, force_refresh = force_refresh)
if src is not None:
return src
repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO
if force_refresh and repo:
# Prime the cache so the freshness read below sees the newest tag.
try:
latest_published_release(repo, force_refresh = True)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("whisper update: force refresh failed", error = str(exc))
freshness = check_prebuilt_freshness(binary)
installed = freshness.get("installed_tag")
latest = freshness.get("latest_tag")
compatible_override = False
if sys.platform == "darwin" and marker is not None:
# The newest published release may require a newer macOS. Ask the same
# host-aware resolver the installer uses so the banner compares against
# the newest release this host can actually install, avoiding a repeated
# offer of an incompatible release after walkback.
resolved = _resolve_prebuilt_for_host(
force_refresh = force_refresh,
backend = marker.get("backend") if isinstance(marker.get("backend"), str) else None,
)
compatible_latest = (resolved or {}).get("release_tag")
if (resolved or {}).get("prebuilt_available") and isinstance(compatible_latest, str):
compatible_override = compatible_latest != latest
latest = compatible_latest
# `behind` compares the release version with a downgrade guard, so a lagging
# /releases/latest or lower published tag can't show a false update
# (whisper_cpp_freshness.is_behind).
update_available = bool(
freshness.get("has_marker")
and (
is_behind(installed, latest)
if sys.platform == "darwin" and marker is not None
else freshness.get("behind")
)
)
# Size of the prebuilt Update would download, for the banner. Only when an
# update is offered; fails open to None (offline / no matching asset).
update_size_bytes = None
if update_available and not compatible_override:
try:
update_size_bytes = update_download_size_bytes(
marker,
latest,
freshness.get("published_repo") or repo,
force_refresh = force_refresh,
)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("whisper update: size lookup failed", error = str(exc))
with _job_lock:
job = dict(_job)
return {
"supported": bool(freshness.get("has_marker")),
"update_available": update_available,
"stale": bool(update_available and freshness.get("stale")),
"installed_tag": installed,
"latest_tag": latest,
"published_repo": freshness.get("published_repo") or repo,
"installed_at_utc": freshness.get("installed_at_utc"),
"age_days": freshness.get("age_days"),
"source_build": False,
"update_size_bytes": update_size_bytes,
"job": job,
}
def _install_latest(
install_dir: Path,
repo: str,
asset: Optional[str],
backend: Optional[str],
script: Path,
set_progress,
pin_release_tag: Optional[str] = None,
) -> dict:
"""Replace whisper.cpp while the sidecar blocks every new load."""
try:
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
sidecar = get_ggml_stt_sidecar()
except Exception as exc:
# Replacing the tree without the singleton's maintenance barrier would
# reopen the Windows executable-lock and stale-process races. Fail closed.
raise RuntimeError("could not coordinate the whisper.cpp sidecar update") from exc
# update_maintenance publishes its guard before waiting for an existing
# transcription, unloads the warm server, and holds the sidecar lock across
# the complete atomic install. No new process can relock or outlive the tree.
with sidecar.update_maintenance() as model_was_active:
return _install_latest_while_blocked(
install_dir,
repo,
asset,
backend,
script,
set_progress,
pin_release_tag = pin_release_tag,
model_was_active = model_was_active,
)
def _install_latest_while_blocked(
install_dir: Path,
repo: str,
asset: Optional[str],
backend: Optional[str],
script: Path,
set_progress,
*,
pin_release_tag: Optional[str],
model_was_active: bool,
) -> dict:
"""Run the installer with the sidecar already in update maintenance."""
cmd = [
sys.executable,
str(script),
"--install-dir",
str(install_dir),
"--whisper-tag",
"latest",
"--published-repo",
repo,
]
# Preserve the installed accelerator across updates. Left unpinned the
# installer re-detects the host, fine on unchanged hardware but able to
# reroute a deliberate choice (e.g. cpu on a GPU box); forwarding the marker's
# backend keeps the same slice.
if isinstance(backend, str) and backend:
cmd.extend(["--backend", backend])
if pin_release_tag:
cmd.extend(["--published-release-tag", pin_release_tag])
cmd.extend(_rocm_install_args(asset))
logger.info("whisper update: installing", cmd = " ".join(cmd))
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
# Every nonzero exit is a failed phase. In particular, exit 2 means the
# requested release was incompatible and no install happened, so reporting
# success would hide the banner and toast an update that never landed.
_flow.stream_installer(
cmd,
env,
set_progress = set_progress,
timeout_seconds = _INSTALL_TIMEOUT_SECONDS,
)
# Drop stale caches so the banner re-checks the swapped marker. If GitHub is
# offline, latest stays unknown and the banner fails open.
reset_caches(drop_disk = True)
try:
latest_published_release(repo, force_refresh = True)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("whisper update: post-install freshness refresh failed", error = str(exc))
new_marker = read_install_marker(_find_binary())
new_tag = (new_marker or {}).get("release_tag")
logger.info("whisper update: success", to_tag = new_tag)
return {
"to_tag": new_tag,
"reload_required": model_was_active,
"message": (
f"Updated whisper.cpp to {new_tag}."
+ (" Reload your model to use it." if model_was_active else "")
),
}
def chained_phase_plan(
*, force_refresh: bool = False, paired_llama_will_update: bool = False
) -> dict:
"""Whisper's side of the combined llama+whisper update item.
Returns {status, update_available, skip_reason, phase}: `status` is the
marker-path status dict (or a minimal one when whisper is skipped),
`update_available` says the chained apply would run a whisper phase, and
`phase` carries what run_chained_phase needs. Only marker-managed installs are
chained: local links, source builds and unmanaged/pinned paths are silently
skipped so whisper can never block a llama update. Never raises; failures
degrade to a skip."""
binary = _find_binary()
if _active_install_is_local_link(binary):
return {
"status": _local_link_status(),
"update_available": False,
"skip_reason": "local_link",
"phase": None,
}
marker = read_install_marker(binary)
if marker is None:
# No marker: whisper is absent or a source/custom build. The standalone
# utils API can update source builds; the chain does not.
return {
"status": None,
"update_available": False,
"skip_reason": "source_build" if binary else "not_installed",
"phase": None,
}
status = get_update_status(force_refresh = force_refresh)
plan: dict = {"status": status, "update_available": False, "skip_reason": None, "phase": None}
if not status.get("update_available"):
# Skew note: when llama just updated but whisper is already latest, a slim
# install keeps hardlinks to the OLD llama ggml inodes -- still the exact
# build whisper was installed against, so skipping is correct and needs no
# re-wiring. A whisper phase that does run re-wires via the installer
# (prepare_runtime_payload).
plan["skip_reason"] = "up_to_date"
return plan
if marker.get("install_kind") == "slim" and not paired_llama_will_update:
# A slim install can only be refreshed from a completed managed llama
# prebuilt. Ask the installer through its read-only resolver so local
# links, markerless/current source builds, and incomplete managed trees
# never produce an Update button that can only fail. When the llama
# phase will run first, it supplies the repaired pairing instead.
resolved = _resolve_prebuilt_for_host(
force_refresh = force_refresh,
backend = marker.get("backend") if isinstance(marker.get("backend"), str) else None,
)
if not (resolved or {}).get("prebuilt_available"):
plan["skip_reason"] = "paired_llama_unavailable"
return plan
script = _installer_script()
if script is None:
plan["skip_reason"] = "installer_missing"
return plan
install_dir = _install_dir_for(binary)
if install_dir is None:
plan["skip_reason"] = "no_install_dir"
return plan
plan["update_available"] = True
plan["phase"] = {
"install_dir": install_dir,
"repo": marker.get("published_repo") or DEFAULT_PUBLISHED_REPO,
"asset": marker.get("asset"),
"backend": marker.get("backend"),
"script": script,
# Install exactly the release the check offered: the installer's unpinned
# "latest" prefers the download-host /releases/latest pointer, which sorts
# by commit date and can lag the published_at pick the freshness check
# used, reinstalling an older build in a loop (the #6219 class the llama
# phase pins against). Not on macOS: the llama phase is unpinned there
# (walk-back to an os-compatible release), so pinning whisper to the
# newest tag could be an impossible pairing (min_os / requires_llama_tag)
# on every retry.
"pin_release_tag": None if sys.platform == "darwin" else status.get("latest_tag"),
}
return plan
def run_chained_phase(phase: dict, set_progress) -> dict:
"""Run the whisper phase of a combined update (spec from chained_phase_plan):
same unload/install/cache-refresh path as the standalone job, reporting
progress through the chained job's window rather than whisper's own."""
return _install_latest(
phase["install_dir"],
phase["repo"],
phase["asset"],
phase["backend"],
phase["script"],
set_progress,
pin_release_tag = phase.get("pin_release_tag"),
)

231
studio/frontend/src/components/assistant-ui/chat-dictation-bar.tsx Normal file
View file

@ -0,0 +1,231 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Spinner } from "@/components/ui/spinner";
import {
cancelActiveStudioDictation,
subscribeDictationLevel,
} from "@/features/chat";
import { useAui, useAuiState } from "@assistant-ui/react";
import { CheckIcon, XIcon } from "lucide-react";
import { type FC, useEffect, useRef, useState } from "react";
import { TooltipIconButton } from "./tooltip-icon-button";
// Dense row of dots that rise into centered recording bars.
const BAR_COUNT = 84;
// Peak height multiple for the loudest audio (dot is 4px). Under the 40px pill
// so bars don't touch the edges.
const MAX_SCALE = 8;
// Time for a sample to slide one bar-width left (drift speed). Bars interpolate
// between samples each frame, so a slower interval stays smooth.
const PUSH_INTERVAL_MS = 165;
// If no real mic level arrives for this long (e.g. Web Audio unavailable), fall
// back to a gentle idle shimmer so the bar stays alive.
const IDLE_AFTER_MS = 450;
const WAVE_BAR_IDS = Array.from(
{ length: BAR_COUNT },
(_, index) => `wave-bar-${index}`,
);
function formatElapsed(ms: number): string {
const total = Math.floor(ms / 1000);
const m = Math.floor(total / 60);
const s = total % 60;
return `${m}:${s.toString().padStart(2, "0")}`;
}
/**
* Recording UI shown in place of the composer input: a live waveform with
* discard and confirm on the right. Confirm transcribes; discard keeps the
* existing composer text.
*/
export const ChatDictationBar: FC = () => {
const aui = useAui();
const isDictating = useAuiState((s) => s.composer.dictation != null);
const [transcribing, setTranscribing] = useState(false);
const [elapsed, setElapsed] = useState(0);
const transcribingRef = useRef(false);
// Extra slot: newest sample lands here; the last visible bar slides toward it.
const barsRef = useRef<number[]>(new Array(BAR_COUNT + 1).fill(0));
const rowRef = useRef<HTMLDivElement | null>(null);
useEffect(() => {
if (!isDictating) {
return;
}
const startedAt = Date.now();
let peak = 0; // loudest level seen since the last waveform advance
let smoothed = 0;
let lastLevelAt = 0;
const barEls = rowRef.current
? Array.from(rowRef.current.children).filter(
(el): el is HTMLElement => el instanceof HTMLElement,
)
: [];
// Paint the bars imperatively (not via React state) so the waveform does not
// thrash renders. The spans carry no style prop, so React never overwrites
// these transforms on an elapsed-timer re-render. At rest each is a round dot
// (scaleY 1); louder audio scales it into a thin centered bar.
for (const el of barEls) {
el.style.transform = "scaleY(1)";
el.style.opacity = "0.62";
}
// Push one new sample, dropping the oldest.
const commitSample = () => {
const bars = barsRef.current;
let level = peak;
peak = 0;
if (Date.now() - lastLevelAt > IDLE_AFTER_MS) {
level = 0.075 + 0.055 * (1 + Math.sin(Date.now() / 360));
}
// Quiet speech needs a perceptual lift. Fast attack, slower release:
// removes twitching while preserving clear peaks.
const visual = Math.min(1, Math.max(0, level) ** 0.62 * 1.45);
smoothed =
visual >= smoothed
? smoothed * 0.2 + visual * 0.8
: smoothed * 0.78 + visual * 0.22;
bars.push(smoothed);
while (bars.length > BAR_COUNT + 1) {
bars.shift();
}
};
// Keep the loudest mic level between advances so downsampling to
// PUSH_INTERVAL_MS doesn't swallow peaks in quiet gaps.
const unsub = subscribeDictationLevel((level) => {
if (level > peak) {
peak = level;
}
lastLevelAt = Date.now();
});
// Repaint every frame; bars interpolate between samples so the wave glides.
// Timer updates at most once per second.
let lastPushAt = performance.now();
let shownSecond = -1;
let raf = 0;
const frame = () => {
raf = requestAnimationFrame(frame);
if (transcribingRef.current) {
return;
}
const now = performance.now();
const bars = barsRef.current;
let steps = 0;
while (now - lastPushAt >= PUSH_INTERVAL_MS && steps < BAR_COUNT + 1) {
commitSample();
lastPushAt += PUSH_INTERVAL_MS;
steps++;
}
// Drop stale backlog if rAF was paused (tab backgrounded).
if (now - lastPushAt >= PUSH_INTERVAL_MS) {
lastPushAt = now;
}
const phase = Math.min(1, (now - lastPushAt) / PUSH_INTERVAL_MS);
for (let i = 0; i < barEls.length; i++) {
// Bar i drifts toward its right neighbour as phase goes 0 to 1.
const a = bars[i] ?? 0;
const b = bars[i + 1] ?? a;
const v = a + (b - a) * phase;
barEls[i].style.transform = `scaleY(${1 + v * (MAX_SCALE - 1)})`;
barEls[i].style.opacity = `${0.62 + v * 0.38}`;
}
const elapsedMs = Date.now() - startedAt;
const second = Math.floor(elapsedMs / 1000);
if (second !== shownSecond) {
shownSecond = second;
setElapsed(elapsedMs);
}
};
raf = requestAnimationFrame(frame);
// Reset in cleanup (dictation end or unmount) so the next session starts
// fresh, without a synchronous setState in the effect body.
return () => {
unsub();
cancelAnimationFrame(raf);
transcribingRef.current = false;
setTranscribing(false);
setElapsed(0);
barsRef.current = new Array(BAR_COUNT + 1).fill(0);
};
}, [isDictating]);
if (!isDictating) {
return null;
}
const discard = () => {
cancelActiveStudioDictation();
};
const confirm = () => {
// Freeze the timer + waveform before the transcription round trip completes
// and the session ends.
transcribingRef.current = true;
setTranscribing(true);
aui.composer().stopDictation();
};
return (
<fieldset
// order-2 places the bar in the input's slot after the left "+" tools.
className="unsloth-dictation-bar order-2 m-0 flex min-w-0 flex-1 items-center gap-2 border-0 p-0"
aria-label="Voice recording"
>
<div
ref={rowRef}
aria-hidden="true"
className="unsloth-dictation-wave grid h-10 min-w-0 flex-1 items-center overflow-hidden px-2"
style={{
gridTemplateColumns: `repeat(${BAR_COUNT}, minmax(1px, 3px))`,
justifyContent: "space-between",
}}
>
{WAVE_BAR_IDS.map((barId) => (
<span
key={barId}
className="h-1 w-full origin-center rounded-full bg-foreground opacity-[0.62] will-change-transform"
/>
))}
</div>
<span className="shrink-0 tabular-nums text-sm text-muted-foreground">
{formatElapsed(elapsed)}
</span>
<div className="flex shrink-0 items-center gap-1">
<TooltipIconButton
type="button"
tooltip="Discard recording"
aria-label="Discard recording"
variant="ghost"
onClick={discard}
className="size-8 rounded-full text-muted-foreground hover:text-foreground"
>
<XIcon className="size-5" />
</TooltipIconButton>
<TooltipIconButton
type="button"
tooltip={transcribing ? "Transcribing…" : "Stop and transcribe"}
aria-label="Stop and transcribe"
variant="default"
onClick={confirm}
disabled={transcribing}
className="size-8 rounded-full"
>
{transcribing ? (
<Spinner className="size-4" />
) : (
<CheckIcon className="size-5" />
)}
</TooltipIconButton>
</div>
</fieldset>
);
};

200
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -34,6 +34,11 @@ import { RenderHtmlToolUI } from "@/components/assistant-ui/tool-ui-render-html"
import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python";
import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
import { ChatDictationBar } from "@/components/assistant-ui/chat-dictation-bar";
import {
isStudioDictationAvailable,
notifyStudioDictationUnavailable,
} from "@/features/chat";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import {
IntentAwareScrollProvider,
@ -1435,6 +1440,7 @@ const Composer: FC<{
disableQueue?: boolean;
}> = ({ disabled, threadId, menuSide, disableQueue }) => {
const aui = useAui();
const isDictating = useAuiState((s) => s.composer.dictation != null);
const pageDragging = useContext(PageDragContext);
const { overlay, closeOverlay } = useGeneratedImageOverlay();
const setImageToolsEnabled = useChatRuntimeStore(
@ -1856,79 +1862,101 @@ const Composer: FC<{
const composerContent = (
<>
<ComposerAttachments />
<PendingAudioChip />
<ThreadDocumentsBar
threadId={referenceThreadId}
onIndexingChange={handleIndexingChange}
/>
<ToolStatusDisplay />
{!isDictating ? (
<>
<ComposerAttachments />
<PendingAudioChip />
</>
) : null}
{/* Keep indexing state subscribed while dictating, but hide its chips so
the waveform stays the composer's only status indicator. */}
<div className={isDictating ? "hidden" : "contents"}>
<ThreadDocumentsBar
threadId={referenceThreadId}
onIndexingChange={handleIndexingChange}
/>
</div>
{!isDictating ? <ToolStatusDisplay /> : null}
<div
className="unsloth-composer-line"
// The permission pill is always visible, so keep the two-row layout
// expanded and leave the primary tool toggles accessible in every mode.
data-expanded="true"
// expanded whenever not dictating; dictation collapses to the bar.
data-expanded={!isDictating ? "true" : "false"}
data-dictating={isDictating ? "true" : undefined}
>
<div
className="unsloth-composer-left"
data-pill-compact={pillsCompact ? "true" : undefined}
>
<ComposerToolsMenu side={effectiveMenuSide} />
{/* Permission-level pill: always visible and opens the permission
level dropdown. */}
<PermissionModeComposerPill side={effectiveMenuSide} />
<WebSearchToggle />
<CodeToolsToggle />
<ImagesToggle />
<KnowledgeBaseComposerButton side={effectiveMenuSide} />
{artifactsEnabled ? <ArtifactsToggle /> : null}
{mcpEnabledForChat ? (
<McpComposerButton side={effectiveMenuSide} />
{/* While dictating, show only the "+"; hide the pill and tool toggles
so the waveform is the sole status indicator. */}
{!isDictating ? (
<>
{/* Permission-level pill: always visible, opens the level dropdown. */}
<PermissionModeComposerPill side={effectiveMenuSide} />
<WebSearchToggle />
<CodeToolsToggle />
<ImagesToggle />
<KnowledgeBaseComposerButton side={effectiveMenuSide} />
{artifactsEnabled ? <ArtifactsToggle /> : null}
{mcpEnabledForChat ? (
<McpComposerButton side={effectiveMenuSide} />
) : null}
</>
) : null}
</div>
<ComposerPrimitive.Input
placeholder={
overlay ? "Type your edits for your image" : "Ask anything"
}
ref={inputRef}
className="aui-composer-input unsloth-composer-input"
minRows={1}
maxRows={12}
autoFocus={!disabled}
disabled={disabled}
aria-label={overlay ? "Image edit instructions" : "Message input"}
// dir="auto": browser picks LTR/RTL from the first strong char;
// no effect on Latin / CJK / Devanagari.
dir="auto"
{...inputProps}
/>
<ComposerRightControls
disabled={
disabled ||
!hasSendableContent ||
isComposing ||
hasPendingAttachments
}
// disableQueue (project new-chat composer) also blocks the queue
// button, so a running thread shows Stop instead of Queue.
queueDisabled={disableQueue || !canQueueCurrentPrompt}
onQueueClick={() => {
if (disableQueue) return;
const queuedPrompt = composerText.trim();
if (queuedPrompt.length === 0) {
return;
}
flushResourcesSync(() => {
aui.composer().setText("");
});
startPromptQueue([queuedPrompt], createPromptQueueTarget(), true);
}}
onSendClick={interceptSend}
onStopClick={stopQueue}
pendingSend={pendingSend}
menuSide={effectiveMenuSide}
queueThreadIds={promptQueueThreadIds}
/>
{isDictating ? (
// The recording UI replaces the input and send controls; only the
// left plus stays visible alongside it.
<ChatDictationBar />
) : (
<>
<ComposerPrimitive.Input
placeholder={
overlay ? "Type your edits for your image" : "Ask anything"
}
ref={inputRef}
className="aui-composer-input unsloth-composer-input"
minRows={1}
maxRows={12}
autoFocus={!disabled}
disabled={disabled}
aria-label={overlay ? "Image edit instructions" : "Message input"}
// dir="auto": browser picks LTR/RTL from the first strong char;
// no effect on Latin / CJK / Devanagari.
dir="auto"
{...inputProps}
/>
<ComposerRightControls
disabled={
disabled ||
!hasSendableContent ||
isComposing ||
hasPendingAttachments
}
// disableQueue (project new-chat composer) also blocks the queue
// button, so a running thread shows Stop instead of Queue.
queueDisabled={disableQueue || !canQueueCurrentPrompt}
onQueueClick={() => {
if (disableQueue) return;
const queuedPrompt = composerText.trim();
if (queuedPrompt.length === 0) {
return;
}
flushResourcesSync(() => {
aui.composer().setText("");
});
startPromptQueue([queuedPrompt], createPromptQueueTarget(), true);
}}
onSendClick={interceptSend}
onStopClick={stopQueue}
pendingSend={pendingSend}
menuSide={effectiveMenuSide}
queueThreadIds={promptQueueThreadIds}
/>
</>
)}
</div>
</>
);
@ -3362,32 +3390,36 @@ const ComposerRightControls: FC<{
findPromptQueueEntry(s, queueThreadIds),
);
const isQueueRunning = Boolean(queueEntry);
const aui = useAui();
// Keep the mic clickable: if the engine can't run here, explain and point to
// the local model instead of disabling the button.
const startDictation = () => {
if (!isStudioDictationAvailable()) {
notifyStudioDictationUnavailable();
return;
}
try {
aui.composer().startDictation();
} catch {
notifyStudioDictationUnavailable();
}
};
return (
<div className="aui-composer-action-wrapper flex shrink-0 items-center gap-1.5">
<ReasoningToggle side={menuSide} />
{/* Starts dictation; the recording bar then covers the input row and owns
the stop and discard actions. */}
<ComposerPrimitive.If dictation={false}>
<ComposerPrimitive.Dictate asChild={true}>
<TooltipIconButton
tooltip="Dictate"
aria-label="Dictate"
variant="ghost"
className="size-8 rounded-full text-foreground"
>
<MicIcon className="size-5" />
</TooltipIconButton>
</ComposerPrimitive.Dictate>
</ComposerPrimitive.If>
<ComposerPrimitive.If dictation={true}>
<ComposerPrimitive.StopDictation asChild={true}>
<TooltipIconButton
tooltip="Stop dictation"
aria-label="Stop dictation"
variant="ghost"
className="size-8 rounded-full text-destructive"
>
<SquareIcon className="size-3 animate-pulse fill-current" />
</TooltipIconButton>
</ComposerPrimitive.StopDictation>
<TooltipIconButton
tooltip="Dictate"
aria-label="Dictate"
type="button"
variant="ghost"
className="size-8 rounded-full text-foreground"
onClick={startDictation}
>
<MicIcon className="size-5" />
</TooltipIconButton>
</ComposerPrimitive.If>
<AuiIf condition={({ thread }) => !thread.isRunning && !isQueueRunning}>
<ComposerPrimitive.Send asChild={true}>

12
studio/frontend/src/components/llama-update-banner.tsx
View file

@ -93,16 +93,17 @@ export function LlamaUpdateBanner({
});
async function handleUpdate() {
const component = status?.component ?? "llama.cpp";
const result = await apply();
if (result?.ok) {
const updatedTag = result.tag ?? status?.latest_tag ?? "the latest build";
const reloadHint = result.reloadRequired
? " Reload your model to use it."
: "";
toast.success(`llama.cpp updated to ${updatedTag}.${reloadHint}`);
toast.success(`${component} updated to ${updatedTag}.${reloadHint}`);
} else if (result) {
toast.error(
`llama.cpp update failed: ${result.error ?? "unknown error"}`,
`${component} update failed: ${result.error ?? "unknown error"}`,
);
}
}
@ -113,6 +114,7 @@ export function LlamaUpdateBanner({
status != null &&
(status.update_available || applying);
const sizeBytes = status?.update_size_bytes ?? null;
const component = status?.component ?? "llama.cpp";
const sizeLabel =
sizeBytes && sizeBytes > 0
? `${Math.round(sizeBytes / (1024 * 1024))} MB`
@ -142,7 +144,7 @@ export function LlamaUpdateBanner({
type="button"
onClick={dismiss}
className="absolute top-2.5 right-3 flex size-6 items-center justify-center rounded-full text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground"
aria-label="Dismiss llama.cpp update notification"
aria-label={`Dismiss ${component} update notification`}
>
<svg
aria-hidden="true"
@ -170,7 +172,7 @@ export function LlamaUpdateBanner({
/>
<div className="min-w-0">
<p className="font-heading text-base font-medium text-foreground">
{applying ? "Updating llama.cpp..." : "New llama.cpp update"}
{applying ? `Updating ${component}...` : `New ${component} update`}
</p>
<p className="mt-0.5 text-xs text-muted-foreground">
{status?.installed_tag ?? "unknown"} &rarr;{" "}
@ -189,7 +191,7 @@ export function LlamaUpdateBanner({
<div
className="mb-1.5 mt-4 h-1 overflow-hidden rounded-full bg-muted"
role="progressbar"
aria-label="Updating llama.cpp"
aria-label={`Updating ${component}`}
aria-valuemin={0}
aria-valuemax={100}
aria-valuenow={

10
studio/frontend/src/components/ui/combobox.tsx
View file

@ -7,6 +7,7 @@
import { Combobox as ComboboxPrimitive } from "@base-ui/react";
import * as React from "react";
import { useWheelScrollRef } from "@/hooks";
import { createContext, useState } from "react";
import { Button } from "@/components/ui/button";
@ -100,13 +101,19 @@ function ComboboxInput({
disabled = false,
showTrigger = true,
showClear = false,
startAddon,
...props
}: ComboboxPrimitive.Input.Props & {
showTrigger?: boolean;
showClear?: boolean;
/** Optional leading content (e.g. a search icon) rendered before the input. */
startAddon?: React.ReactNode;
}): React.ReactElement {
return (
<InputGroup className={cn("w-auto", className)}>
{startAddon && (
<InputGroupAddon align="inline-start">{startAddon}</InputGroupAddon>
)}
<ComboboxPrimitive.Input
render={<InputGroupInput disabled={disabled} />}
{...props}
@ -195,8 +202,11 @@ function ComboboxList({
className,
...props
}: ComboboxPrimitive.List.Props): React.ReactElement {
const listRef = useWheelScrollRef<HTMLDivElement>();
return (
<ComboboxPrimitive.List
ref={listRef}
data-slot="combobox-list"
className={cn(
"no-scrollbar max-h-[min(calc(--spacing(72)---spacing(9)),calc(var(--available-height)---spacing(9)))] scroll-py-1 overflow-y-auto p-1 data-empty:p-0 overflow-y-auto overscroll-contain",

91
studio/frontend/src/features/chat/adapters/dictation-level.ts Normal file
View file

@ -0,0 +1,91 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
type LevelListener = (level: number) => void;
type FrameListener = (rawRms: number, now: number) => void;
const levelListeners = new Set<LevelListener>();
/** Subscribe to the live microphone level (0..1) during dictation. */
export function subscribeDictationLevel(listener: LevelListener): () => void {
levelListeners.add(listener);
return () => {
levelListeners.delete(listener);
};
}
function publishLevel(level: number): void {
for (const listener of levelListeners) {
listener(level);
}
}
/**
* Start a lightweight Web Audio meter for the shared recording waveform and
* optional voice-activity detection. Returns an idempotent cleanup function.
*/
export function startDictationLevelMeter(
source: MediaStream,
onFrame?: FrameListener,
): () => void {
let audioContext: AudioContext | null = null;
let levelRaf = 0;
let stopped = false;
const stop = () => {
if (stopped) {
return;
}
stopped = true;
if (levelRaf) {
cancelAnimationFrame(levelRaf);
}
levelRaf = 0;
audioContext?.close().catch(() => {
// A closing or already-closed context is harmless.
});
audioContext = null;
publishLevel(0);
};
try {
const Ctx =
window.AudioContext ||
(window as unknown as { webkitAudioContext?: typeof AudioContext })
.webkitAudioContext;
if (!Ctx) {
return stop;
}
audioContext = new Ctx();
audioContext.resume().catch(() => {
// Some browsers resume automatically after microphone permission.
});
const node = audioContext.createMediaStreamSource(source);
const analyser = audioContext.createAnalyser();
analyser.fftSize = 512;
node.connect(analyser);
const data = new Uint8Array(analyser.frequencyBinCount);
const tick = () => {
if (stopped) {
return;
}
analyser.getByteTimeDomainData(data);
let sum = 0;
for (const sample of data) {
const value = (sample - 128) / 128;
sum += value * value;
}
const rms = Math.sqrt(sum / data.length);
onFrame?.(rms, performance.now());
// Perceptual boost so normal speech remains clearly visible.
publishLevel(Math.min(1, rms * 3.2));
levelRaf = requestAnimationFrame(tick);
};
levelRaf = requestAnimationFrame(tick);
} catch {
// The waveform is cosmetic; recording can continue without Web Audio.
stop();
}
return stop;
}

139
studio/frontend/src/features/chat/adapters/studio-dictation-adapter.tsx Normal file
View file

@ -0,0 +1,139 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useSettingsDialogStore } from "@/features/settings/stores/settings-dialog-store";
import {
type DictationEngine,
useVoiceSettingsStore,
} from "@/features/settings/stores/voice-settings-store";
import { toast } from "@/lib/toast";
import type { DictationAdapter } from "@assistant-ui/react";
import { StudioModelDictationAdapter } from "./studio-model-dictation-adapter";
import {
type StudioDictationSession,
StudioWebSpeechDictationAdapter,
} from "./studio-web-speech-dictation-adapter";
// The one live dictation session, so the recording bar's discard (X) can cancel
// it without going through assistant-ui (which only exposes stop, i.e.
// transcribe). Cancelling emits no transcript, so composer text is untouched.
let activeSession: StudioDictationSession | null = null;
/** Discard the current dictation without transcribing. Safe to call when idle. */
export function cancelActiveStudioDictation(): void {
const session = activeSession;
activeSession = null;
session?.cancel();
}
/**
* Routes dictation to the engine chosen in Voice settings, resolved at listen()
* time so switching engines applies without reloading the chat runtime.
*/
/** Both local engines (Transformers and GGUF) record via MediaRecorder. */
function usesModelRecording(dictationEngine: DictationEngine): boolean {
return dictationEngine === "model";
}
export class StudioDictationAdapter implements DictationAdapter {
// Chat linked in Recent dictations. undefined follows the active single chat;
// null records no chat (composers outside it, e.g. Compare).
private readonly chatId: string | null | undefined;
constructor(options: { chatId?: string | null } = {}) {
this.chatId = options.chatId;
}
static isSupported(
dictationEngine: DictationEngine = useVoiceSettingsStore.getState()
.dictationEngine,
): boolean {
return usesModelRecording(dictationEngine)
? StudioModelDictationAdapter.isSupported()
: StudioWebSpeechDictationAdapter.isSupported();
}
listen(): StudioDictationSession {
const session = this.createSession();
// A second entry point (chat, Compare, settings test) replaces the active
// session; cancel the old one so it cannot keep the mic open or save a
// transcript with no discard button pointing at it.
cancelActiveStudioDictation();
activeSession = session;
// Forget the session once it ends so a later cancel is a no-op.
const clear = () => {
if (activeSession === session) {
activeSession = null;
}
};
session.onSpeechEnd(clear);
session.onEnd?.(clear);
return session;
}
private createSession(): StudioDictationSession {
const { dictationEngine } = useVoiceSettingsStore.getState();
if (usesModelRecording(dictationEngine)) {
if (StudioModelDictationAdapter.isSupported()) {
return new StudioModelDictationAdapter({ chatId: this.chatId }).listen();
}
throw new Error(
"Local model dictation is not supported in this browser.",
);
}
if (StudioWebSpeechDictationAdapter.isSupported()) {
return new StudioWebSpeechDictationAdapter({ chatId: this.chatId }).listen();
}
throw new Error("Browser dictation is not supported in this browser.");
}
}
/** Whether dictation can run now for the chosen engine. */
export function isStudioDictationAvailable(
dictationEngine: DictationEngine = useVoiceSettingsStore.getState()
.dictationEngine,
): boolean {
return StudioDictationAdapter.isSupported(dictationEngine);
}
/** Explain why dictation can't start and point the user to the local model. */
export function notifyStudioDictationUnavailable(
dictationEngine: DictationEngine = useVoiceSettingsStore.getState()
.dictationEngine,
): void {
// Both engines need a secure context (localhost or HTTPS).
if (typeof window !== "undefined" && !window.isSecureContext) {
toast.error("Voice typing needs a secure connection.", {
description:
"Open Studio at http://127.0.0.1 (localhost) or over HTTPS to dictate.",
});
return;
}
if (usesModelRecording(dictationEngine)) {
// Defensive: MediaRecorder is effectively always present here.
toast.error("Voice recording isn't available in this browser.");
return;
}
// Browser Web Speech is missing (e.g. Firefox). Stack text and button so the
// action sits below, not squeezed into a side column.
const toastId = toast.error("Voice typing isn't available in this browser.", {
description: (
<div className="mt-0.5 flex flex-col items-start gap-2 pb-1.5">
<span>
Choose the local speech-to-text model in Voice settings to dictate
here.
</span>
<button
type="button"
onClick={() => {
useSettingsDialogStore.getState().openDialog("voice");
toast.dismiss(toastId);
}}
className="rounded-full bg-foreground px-2.5 pt-1 pb-1.5 text-xs font-medium text-background transition-colors hover:bg-foreground/90"
>
Open Voice settings
</button>
</div>
),
});
}

637
studio/frontend/src/features/chat/adapters/studio-model-dictation-adapter.ts Normal file
View file

@ -0,0 +1,637 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { hubTokenHeader } from "@/features/hub/lib/hub-token-header";
import {
applyDictationDictionary,
isCuratedSttModel,
recordRecentDictation,
resolveModelDictationLanguage,
useVoiceSettingsStore,
} from "@/features/settings/stores/voice-settings-store";
import type { DictationAdapter } from "@assistant-ui/react";
import { toast } from "sonner";
import { startDictationLevelMeter } from "./dictation-level";
import {
type StudioDictationSession,
isMissingDeviceError,
resolveDictationChatId,
} from "./studio-web-speech-dictation-adapter";
// Fine timeslice so the buffer is ready the moment a segment is cut or stopped.
const SEGMENT_TIMESLICE_MS = 250;
// Whisper pads input to 30s. Short dictation stays one clip; long dictation cuts
// at the first pause after 20s or before the 30s boundary.
const MIN_SEGMENT_MS = 20_000;
const MAX_SEGMENT_MS = 28_000;
const SILENCE_CUT_MS = 280;
// Raw RMS (0..1) above which a frame counts as speech (well above the room floor
// after noise suppression).
const VOICE_RMS = 0.015;
// Prefer Opus (small, widely supported); fall back to whatever the browser
// records. The backend decodes any of these with PyAV.
const PREFERRED_MIME_TYPES = [
"audio/webm;codecs=opus",
"audio/webm",
"audio/ogg;codecs=opus",
"audio/mp4",
];
function pickMimeType(): string | undefined {
if (typeof MediaRecorder === "undefined") return undefined;
for (const type of PREFERRED_MIME_TYPES) {
if (MediaRecorder.isTypeSupported(type)) return type;
}
return undefined;
}
const stopStream = (stream: MediaStream | null) => {
for (const track of stream?.getTracks() ?? []) {
track.stop();
}
};
/** Backend STT engine, decided by the model: curated ids run GGML through
* whisper.cpp; a custom HF repo is safetensors and runs through Transformers. */
export type SttEngine = "transformers" | "gguf";
export function sttEngineFor(model: string): SttEngine {
return isCuratedSttModel(model) ? "gguf" : "transformers";
}
/** POST audio to the STT sidecar and return the transcript. */
export async function transcribeAudioBlob(
blob: Blob,
options: {
model?: string;
language?: string;
engine?: SttEngine;
signal?: AbortSignal;
} = {},
): Promise<string> {
const settings = useVoiceSettingsStore.getState();
const model = options.model ?? settings.sttModel;
const language = resolveModelDictationLanguage(
model,
options.language ?? settings.dictationLanguage,
);
const engine = options.engine ?? sttEngineFor(model);
const params = new URLSearchParams({ model, fast: "true", engine });
if (language) params.set("language", language);
const response = await authFetch(
`/api/inference/audio/transcribe/raw?${params.toString()}`,
{
method: "POST",
headers: { "Content-Type": blob.type || "application/octet-stream" },
body: blob,
signal: options.signal,
},
);
if (!response.ok) {
const body = (await response.json().catch(() => null)) as {
detail?: string;
} | null;
const detail = body?.detail ?? `HTTP ${response.status}`;
if (response.status === 501) {
throw new Error(
"Speech-to-text is not available on this server. Run `unsloth studio update` to install it.",
);
}
throw new Error(detail);
}
const data = (await response.json()) as { text?: string };
return (data.text ?? "").trim();
}
export interface SttDownloadStatus {
downloading: boolean;
model: string | null;
error: string | null;
bytes_total: number | null;
bytes_done: number | null;
}
export interface SttEngineStatus {
available: boolean;
loaded_model: string | null;
loading: boolean;
device: string | null;
keep_alive_seconds: number;
default_model: string;
models: string[];
downloaded_models: string[];
download: SttDownloadStatus;
}
export interface SttStatus {
available: boolean;
loaded_model: string | null;
loading: boolean;
device: string | null;
keep_alive_seconds: number;
default_model: string;
models: string[];
/** Per-engine state; absent on servers predating the engine split. */
transformers?: SttEngineStatus;
gguf?: SttEngineStatus;
}
// Keep load/unload requests ordered so a new recording cannot race an unload
// still finishing for the previous one.
let sttLifecycle: Promise<void> = Promise.resolve();
function queueSttLifecycle(operation: () => Promise<void>): Promise<void> {
const result = sttLifecycle.catch(() => {}).then(operation);
sttLifecycle = result.catch(() => {});
return result;
}
/** Report whether STT is installed and which model, if any, is resident.
* Passing a model extends the downloaded check to custom repos. */
export async function fetchSttStatus(
refreshKey?: number,
model?: string,
): Promise<SttStatus> {
const params = new URLSearchParams();
if (refreshKey !== undefined) params.set("refresh", String(refreshKey));
if (model) params.set("model", model);
const query = params.toString();
const response = await authFetch(
`/api/inference/audio/stt/status${query ? `?${query}` : ""}`,
);
if (!response.ok) throw new Error(`HTTP ${response.status}`);
return (await response.json()) as SttStatus;
}
/** Verify a custom Hub repository is a Transformers Whisper checkpoint. */
export async function validateSttModel(
model: string,
hfToken?: string,
): Promise<void> {
const response = await authFetch("/api/inference/audio/stt/validate", {
method: "POST",
headers: {
"Content-Type": "application/json",
...hubTokenHeader(hfToken),
},
body: JSON.stringify({ model }),
});
if (!response.ok) {
const body = (await response.json().catch(() => null)) as {
detail?: string;
} | null;
throw new Error(body?.detail ?? `HTTP ${response.status}`);
}
}
/** Load a selected model that is already downloaded. */
export function loadSttModel(model: string, engine?: SttEngine): Promise<void> {
const resolvedEngine = engine ?? sttEngineFor(model);
return queueSttLifecycle(async () => {
const response = await authFetch("/api/inference/audio/stt/load", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ model, engine: resolvedEngine }),
});
if (!response.ok) {
const body = (await response.json().catch(() => null)) as {
detail?: string;
} | null;
throw new Error(body?.detail ?? `HTTP ${response.status}`);
}
});
}
/** Start a background download of a dictation model. */
export async function startSttDownload(
model: string,
hfToken?: string,
): Promise<void> {
const response = await authFetch("/api/inference/audio/stt/download", {
method: "POST",
headers: {
"Content-Type": "application/json",
...hubTokenHeader(hfToken),
},
body: JSON.stringify({ model, engine: sttEngineFor(model) }),
});
if (!response.ok) {
const body = (await response.json().catch(() => null)) as {
detail?: string;
} | null;
throw new Error(body?.detail ?? `HTTP ${response.status}`);
}
}
/** Release the local STT model and its RAM/VRAM allocations. */
export function unloadSttModel(): Promise<void> {
return queueSttLifecycle(async () => {
const response = await authFetch("/api/inference/audio/stt/unload", {
method: "POST",
});
if (!response.ok) {
const body = (await response.json().catch(() => null)) as {
detail?: string;
} | null;
throw new Error(body?.detail ?? `HTTP ${response.status}`);
}
});
}
/**
* Local model dictation. Short recordings use one pass; long ones split near
* Whisper's 30s window. Confirm keeps text, discard removes it, and either
* releases the microphone immediately.
*/
export class StudioModelDictationAdapter implements DictationAdapter {
private readonly chatId: string | null | undefined;
constructor(options: { chatId?: string | null } = {}) {
this.chatId = options.chatId;
}
static isSupported(): boolean {
return (
typeof window !== "undefined" &&
window.isSecureContext &&
typeof MediaRecorder !== "undefined" &&
navigator.mediaDevices?.getUserMedia !== undefined
);
}
listen(): DictationAdapter.Session {
if (!StudioModelDictationAdapter.isSupported()) {
throw new Error("Recording is not supported in this browser.");
}
// Pin the model, language, and linked chat chosen when recording began, so a
// mid-session settings change or thread switch cannot affect later segments
// or relink the saved transcript.
const { sttModel: sessionModel, dictationLanguage } =
useVoiceSettingsStore.getState();
const sessionLanguage = resolveModelDictationLanguage(
sessionModel,
dictationLanguage,
);
const sessionEngine = sttEngineFor(sessionModel);
const sessionChatId = resolveDictationChatId(this.chatId);
const speechStartCallbacks = new Set<() => void>();
const speechEndCallbacks = new Set<
(result: DictationAdapter.Result) => void
>();
const speechCallbacks = new Set<
(result: DictationAdapter.Result) => void
>();
const endCallbacks = new Set<() => void>();
let stream: MediaStream | null = null;
let ended = false;
let cancelled = false;
let finalizing = false;
const abortController = new AbortController();
const mimeType = pickMimeType();
// Shared waveform meter also feeds this adapter's pause detector.
let stopLevelMeter = () => {
// Replaced after microphone access succeeds.
};
let onAudioFrame: (rawRms: number, now: number) => void = () => {};
let resolveEnded: (() => void) | null = null;
const endedPromise = new Promise<void>((resolve) => {
resolveEnded = resolve;
});
// --- Background transcription pipeline ---------------------------------
// Each segment is a self-contained clip transcribed on its own; results are
// stored by index so the final text keeps its order.
type Segment = {
index: number;
chunks: Blob[];
startedAt: number;
voiced: boolean;
recorder: MediaRecorder;
};
const results: string[] = [];
const queue: { index: number; blob: Blob }[] = [];
let worker = false;
let currentSeg: Segment | null = null;
let segCounter = 0;
let pendingRecorders = 0;
let silenceMs = 0;
let lastFrameAt = 0;
let cutting = false;
let finalCutDone = false;
let reportedTranscriptionError = false;
const reportTranscriptionError = (error: unknown) => {
if (reportedTranscriptionError || cancelled || ended) return;
reportedTranscriptionError = true;
const message =
error instanceof Error && error.message
? error.message
: "A recorded segment could not be transcribed.";
console.error("STT transcription error:", error);
toast.error(message);
};
const buildTranscript = () =>
results
.filter((part) => part?.trim())
.join(" ")
.trim();
const finishSession = (
reason: "stopped" | "cancelled" | "error",
transcript?: string,
) => {
if (ended) return;
ended = true;
stopLevelMeter();
if (currentSeg && currentSeg.recorder.state !== "inactive") {
try {
currentSeg.recorder.stop();
} catch {
// ignore
}
}
session.status = { type: "ended", reason };
stopStream(stream);
stream = null;
const corrected = transcript ? applyDictationDictionary(transcript) : "";
if (reason !== "cancelled" && corrected) {
for (const callback of speechCallbacks) {
callback({ transcript: corrected, isFinal: true });
}
recordRecentDictation(corrected, sessionChatId);
}
for (const callback of speechEndCallbacks) {
callback({ transcript: corrected });
}
for (const callback of endCallbacks) callback();
resolveEnded?.();
};
// Finish once the final segment has been cut and the queue has drained.
const maybeComplete = () => {
if (ended || cancelled || !finalizing || !finalCutDone) return;
if (pendingRecorders === 0 && queue.length === 0 && !worker) {
finishSession("stopped", buildTranscript());
}
};
// Transcribe queued segments one at a time so the backend is never flooded.
const processQueue = () => {
if (worker || cancelled || ended) return;
const item = queue.shift();
if (!item) {
maybeComplete();
return;
}
worker = true;
void (async () => {
try {
const text = await transcribeAudioBlob(item.blob, {
model: sessionModel,
language: sessionLanguage,
engine: sessionEngine,
signal: abortController.signal,
});
if (!cancelled) results[item.index] = text;
} catch (error) {
if (!cancelled && !abortController.signal.aborted) {
// Keep transcribed segments, but never hide that part was lost.
reportTranscriptionError(error);
}
} finally {
worker = false;
processQueue();
}
})();
};
// Every non-empty recording is transcribed. The RMS meter only shapes
// segment boundaries; a quiet microphone or suspended AudioContext can keep
// it below VOICE_RMS for real speech, so it must never discard audio.
// Whisper returns an empty transcript for genuine silence.
const enqueueSegment = (index: number, blob: Blob) => {
if (blob.size > 0) {
queue.push({ index, blob });
processQueue();
} else {
results[index] = "";
maybeComplete();
}
};
// Start recording a fresh segment on the shared mic stream.
const startSegment = () => {
if (ended || cancelled || !stream) return;
const seg: Segment = {
index: segCounter++,
chunks: [],
startedAt: performance.now(),
voiced: false,
recorder: new MediaRecorder(
stream,
mimeType ? { mimeType } : undefined,
),
};
currentSeg = seg;
silenceMs = 0;
seg.recorder.addEventListener("dataavailable", (event) => {
if (event.data.size > 0) seg.chunks.push(event.data);
});
seg.recorder.addEventListener("stop", () => {
pendingRecorders = Math.max(0, pendingRecorders - 1);
if (cancelled || ended) {
maybeComplete();
return;
}
const blob = new Blob(seg.chunks, {
type: seg.recorder.mimeType || "audio/webm",
});
enqueueSegment(seg.index, blob);
});
pendingRecorders += 1;
try {
seg.recorder.start(SEGMENT_TIMESLICE_MS);
} catch (error) {
pendingRecorders = Math.max(0, pendingRecorders - 1);
if (currentSeg === seg) currentSeg = null;
throw error;
}
};
// Close the current segment at a pause and open the next, so recording stays
// continuous while each clip is independently decodable.
const cutSegment = () => {
const seg = currentSeg;
if (cutting || !seg || finalizing) return;
cutting = true;
const rec = seg.recorder;
if (rec.state !== "inactive") {
rec.addEventListener(
"stop",
() => {
cutting = false;
},
{ once: true },
);
try {
rec.stop();
} catch {
cutting = false;
}
} else {
cutting = false;
}
startSegment();
};
// Pause detector: mark voiced frames. Short dictations stay one segment; long
// ones cut at a pause after the target duration, or at the hard limit.
onAudioFrame = (rawRms, now) => {
const seg = currentSeg;
if (!seg || finalizing) {
lastFrameAt = now;
return;
}
if (rawRms > VOICE_RMS) {
seg.voiced = true;
silenceMs = 0;
} else if (lastFrameAt) {
silenceMs += now - lastFrameAt;
}
lastFrameAt = now;
const duration = now - seg.startedAt;
const pauseBreak =
seg.voiced && duration > MIN_SEGMENT_MS && silenceMs > SILENCE_CUT_MS;
if (!cutting && (pauseBreak || duration > MAX_SEGMENT_MS)) {
cutSegment();
}
};
const session: StudioDictationSession = {
status: { type: "starting" },
stop: async () => {
if (!ended && !finalizing) {
finalizing = true;
// Stop publishing zero-valued frames at once so the UI can switch to
// its transcription shimmer.
stopLevelMeter();
const seg = currentSeg;
// Cut the final segment (its buffer survives) so only the short tail is
// left to transcribe, then release the mic immediately.
if (seg && seg.recorder.state !== "inactive") {
seg.recorder.addEventListener(
"stop",
() => {
finalCutDone = true;
maybeComplete();
},
{ once: true },
);
try {
seg.recorder.stop();
} catch {
finalCutDone = true;
maybeComplete();
}
} else {
finalCutDone = true;
maybeComplete();
}
stopStream(stream);
stream = null;
}
await endedPromise;
},
cancel: () => {
if (ended) return;
cancelled = true;
finalizing = true;
abortController.abort();
finishSession("cancelled");
},
onSpeechStart: (callback) => {
speechStartCallbacks.add(callback);
return () => {
speechStartCallbacks.delete(callback);
};
},
onSpeechEnd: (callback) => {
speechEndCallbacks.add(callback);
return () => {
speechEndCallbacks.delete(callback);
};
},
onSpeech: (callback) => {
speechCallbacks.add(callback);
return () => {
speechCallbacks.delete(callback);
};
},
onEnd: (callback: () => void) => {
endCallbacks.add(callback);
return () => {
endCallbacks.delete(callback);
};
},
};
void (async () => {
try {
const { micDeviceId } = useVoiceSettingsStore.getState();
const baseAudio: MediaTrackConstraints = {
echoCancellation: true,
noiseSuppression: true,
};
try {
stream = await navigator.mediaDevices.getUserMedia({
audio:
micDeviceId && micDeviceId !== "default"
? { ...baseAudio, deviceId: { exact: micDeviceId } }
: baseAudio,
});
} catch (error) {
// Saved mic may be unplugged; fall back to the default.
if (micDeviceId !== "default" && isMissingDeviceError(error)) {
stream = await navigator.mediaDevices.getUserMedia({
audio: baseAudio,
});
} else {
throw error;
}
}
if (ended || cancelled) {
stopStream(stream);
stream = null;
return;
}
// Warm the model only after mic access. The backend loads cache-only
// and never downloads here.
void loadSttModel(sessionModel, sessionEngine).catch(
reportTranscriptionError,
);
stopLevelMeter = startDictationLevelMeter(stream, (rawRms, now) => {
onAudioFrame(rawRms, now);
});
startSegment();
session.status = { type: "running" };
for (const callback of speechStartCallbacks) callback();
} catch (error) {
const message = isMissingDeviceError(error)
? "No microphone was found for dictation."
: "Dictation could not access the microphone.";
console.error("STT microphone error:", error);
toast.error(message);
finishSession("error");
}
})();
return session;
}
}

118
studio/frontend/src/features/chat/adapters/studio-speech-synthesis-adapter.ts
View file

@ -6,17 +6,19 @@ import { useVoiceSettingsStore } from "@/features/settings/stores/voice-settings
import { toast } from "@/lib/toast";
import type { SpeechSynthesisAdapter } from "@assistant-ui/react";
/** Voice for a stored voiceURI; undefined lets the browser pick. */
/** Voice for a stored voiceURI. "default" resolves to the voice the platform
* marks as its default, so the "System default" choice means what it says
* instead of falling back to a curated pick. Undefined lets the browser pick. */
export function findTtsVoice(
voiceURI: string,
): SpeechSynthesisVoice | undefined {
if (typeof window === "undefined" || !window.speechSynthesis) {
return undefined;
}
if (!voiceURI || voiceURI === "default") return undefined;
return window.speechSynthesis
.getVoices()
.find((voice) => voice.voiceURI === voiceURI);
if (!voiceURI) return undefined;
const voices = window.speechSynthesis.getVoices();
if (voiceURI === "default") return voices.find((voice) => voice.default);
return voices.find((voice) => voice.voiceURI === voiceURI);
}
// macOS novelty and legacy Eloquence voices that sound robotic and flood the picker.
@ -59,19 +61,54 @@ function voiceBaseName(voice: SpeechSynthesisVoice): string {
return name;
}
// Well-known natural English voices, best first. Breaks ties when the name
// carries no quality hint, so basic voices are not just alphabetical.
const PREFERRED_VOICE_NAMES = [
"samantha",
"alex",
"ava",
"allison",
"susan",
"tom",
"daniel",
"serena",
"karen",
"moira",
"tessa",
"fiona",
];
// Quality tier from vendor hints in the voice name.
function voiceQualityScore(voice: SpeechSynthesisVoice): number {
const name = voice.name.toLowerCase();
let score = 0;
if (name.includes("premium")) score += 8;
if (name.includes("enhanced")) score += 7;
if (name.includes("natural") || name.includes("neural")) score += 6;
if (name.includes("siri")) score += 6;
if (name.includes("google")) score += 5;
if (name.includes("microsoft")) score += 4;
if (voice.default) score += 3;
if (name.includes("premium")) score += 100;
if (name.includes("siri")) score += 90;
if (name.includes("enhanced")) score += 80;
if (name.includes("natural") || name.includes("neural")) score += 70;
if (name.includes("google")) score += 40;
if (name.includes("microsoft")) score += 30;
return score;
}
// Higher for voices in the user's exact region, then the same language.
function voiceLocaleScore(voice: SpeechSynthesisVoice): number {
const navLang =
typeof navigator !== "undefined" && navigator.language
? navigator.language.toLowerCase()
: "en-us";
const lang = voice.lang.toLowerCase().replace("_", "-");
if (lang === navLang) return 2;
if (langBase(lang) === langBase(navLang)) return 1;
return 0;
}
// Rank in the preferred list, best first; 0 when the voice is not listed.
function voicePreferredRank(voice: SpeechSynthesisVoice): number {
const index = PREFERRED_VOICE_NAMES.indexOf(voiceBaseName(voice));
return index === -1 ? 0 : PREFERRED_VOICE_NAMES.length - index;
}
function langBase(tag: string): string {
return tag.toLowerCase().split(/[-_]/)[0] ?? "";
}
@ -86,8 +123,8 @@ const MAX_CURATED_VOICES = 20;
export function curateSystemVoices(
voices: SpeechSynthesisVoice[],
selectedVoiceURI?: string,
dictationLanguage = useVoiceSettingsStore.getState().dictationLanguage,
): SpeechSynthesisVoice[] {
const { dictationLanguage } = useVoiceSettingsStore.getState();
const wantedLangs = new Set<string>(["en"]);
if (typeof navigator !== "undefined" && navigator.language) {
wantedLangs.add(langBase(navigator.language));
@ -107,12 +144,33 @@ export function curateSystemVoices(
});
kept.sort((a, b) => {
const scoreDiff = voiceQualityScore(b) - voiceQualityScore(a);
if (scoreDiff !== 0) return scoreDiff;
const quality = voiceQualityScore(b) - voiceQualityScore(a);
if (quality !== 0) return quality;
const locale = voiceLocaleScore(b) - voiceLocaleScore(a);
if (locale !== 0) return locale;
const preferred = voicePreferredRank(b) - voicePreferredRank(a);
if (preferred !== 0) return preferred;
const byDefault = Number(b.default) - Number(a.default);
if (byDefault !== 0) return byDefault;
return a.name.localeCompare(b.name);
});
const curated = kept.slice(0, MAX_CURATED_VOICES);
// macOS reports some voices twice (compact + enhanced) under one name. Keep
// one per name and language, preferring the selected voice then the best
// ranked, so no duplicates show.
const keyOf = (voice: SpeechSynthesisVoice) =>
`${voiceBaseName(voice)}|${voice.lang.toLowerCase()}`;
const winners = new Map<string, string>();
for (const voice of kept) {
const key = keyOf(voice);
if (!winners.has(key)) winners.set(key, voice.voiceURI);
if (voice.voiceURI === selectedVoiceURI) winners.set(key, voice.voiceURI);
}
const deduped = kept.filter(
(voice) => winners.get(keyOf(voice)) === voice.voiceURI,
);
const curated = deduped.slice(0, MAX_CURATED_VOICES);
if (
selectedVoiceURI &&
selectedVoiceURI !== "default" &&
@ -126,6 +184,22 @@ export function curateSystemVoices(
return curated;
}
/**
* Best voice when none is chosen. The browser default on macOS is often a
* robotic legacy voice, so fall back to the top curated voice instead.
*/
function defaultTtsVoice(): SpeechSynthesisVoice | undefined {
if (typeof window === "undefined" || !window.speechSynthesis) {
return undefined;
}
const voices = window.speechSynthesis.getVoices();
return (
curateSystemVoices(voices)[0] ??
voices.find((voice) => voice.default) ??
voices[0]
);
}
/** Build an utterance from the current Voice settings. */
export function createConfiguredUtterance(
text: string,
@ -133,7 +207,7 @@ export function createConfiguredUtterance(
const { ttsVoiceURI, ttsRate, ttsPitch, ttsVolume } =
useVoiceSettingsStore.getState();
const utterance = new SpeechSynthesisUtterance(text);
const voice = findTtsVoice(ttsVoiceURI);
const voice = findTtsVoice(ttsVoiceURI) ?? defaultTtsVoice();
if (voice) {
utterance.voice = voice;
utterance.lang = voice.lang;
@ -272,8 +346,14 @@ export class StudioSpeechSynthesisAdapter implements SpeechSynthesisAdapter {
if (res.status.type === "ended") return;
// Surface genuine read-aloud failures; a cancelled/interrupted utterance
// is a normal stop, not an error, and must not toast.
if (reason === "error" && error !== "interrupted" && error !== "canceled") {
toast.error(error instanceof Error ? error.message : "Read aloud failed.");
if (
reason === "error" &&
error !== "interrupted" &&
error !== "canceled"
) {
toast.error(
error instanceof Error ? error.message : "Read aloud failed.",
);
}
res.status = { type: "ended", reason, error };
for (const handler of subscribers) handler();

185
studio/frontend/src/features/chat/adapters/studio-web-speech-dictation-adapter.ts
View file

@ -1,6 +1,7 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useSettingsDialogStore } from "@/features/settings/stores/settings-dialog-store";
import {
applyDictationDictionary,
recordRecentDictation,
@ -9,6 +10,39 @@ import {
} from "@/features/settings/stores/voice-settings-store";
import type { DictationAdapter } from "@assistant-ui/react";
import { toast } from "sonner";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import { startDictationLevelMeter } from "./dictation-level";
/** Chat open while dictating, so the saved dictation can link back to it. */
export function activeDictationChatId(): string | undefined {
return useChatRuntimeStore.getState().activeThreadId ?? undefined;
}
/**
* Resolve the chat a saved dictation links to. undefined falls back to the
* active single chat; null (composers without one, e.g. Compare) means none.
*/
export function resolveDictationChatId(
chatId: string | null | undefined,
): string | undefined {
if (chatId === undefined) return activeDictationChatId();
return chatId ?? undefined;
}
// Reused id so repeated network failures replace, not stack, the same toast.
const NETWORK_TOAST_ID = "dictation-network-offline";
// Short grace for the browser to finalize its latest interim hypothesis. If it
// doesn't, promote that text instead of stalling on a finalization spinner.
const STOP_FINALIZATION_GRACE_MS = 350;
/**
* A dictation session with an extra onEnd hook (not part of the assistant-ui
* interface) so non-runtime callers can reset when the session ends by itself.
*/
export type StudioDictationSession = DictationAdapter.Session & {
onEnd?: (callback: () => void) => () => void;
};
const getSpeechRecognitionAPI = ():
| SpeechRecognitionConstructor
@ -18,7 +52,9 @@ const getSpeechRecognitionAPI = ():
};
const stopStream = (stream: MediaStream | null) => {
stream?.getTracks().forEach((track) => track.stop());
for (const track of stream?.getTracks() ?? []) {
track.stop();
}
};
const mediaErrorName = (error: unknown): unknown =>
@ -48,7 +84,10 @@ export const describeMediaError = (error: unknown): string => {
: "Dictation could not access the microphone.";
};
export const describeSpeechError = (error: string, message?: string): string => {
export const describeSpeechError = (
error: string,
message?: string,
): string => {
if (error === "not-allowed") {
return "Speech recognition was blocked by the browser. Check microphone permissions for this Unsloth page.";
}
@ -68,18 +107,21 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
private readonly language: string | undefined;
private readonly continuous: boolean;
private readonly interimResults: boolean;
private readonly chatId: string | null | undefined;
constructor(
options: {
language?: string;
continuous?: boolean;
interimResults?: boolean;
chatId?: string | null;
} = {},
) {
// Resolved from Voice settings at listen() time unless overridden.
this.language = options.language;
this.continuous = options.continuous ?? true;
this.interimResults = options.interimResults ?? true;
this.chatId = options.chatId;
}
static isSupported(): boolean {
@ -101,6 +143,9 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
recognition.lang = this.language ?? resolveDictationLanguage();
recognition.continuous = this.continuous;
recognition.interimResults = this.interimResults;
// Pin the linked chat now so a thread switch during finalization cannot
// relink the transcript to the newly opened chat.
const sessionChatId = resolveDictationChatId(this.chatId);
const speechStartCallbacks = new Set<() => void>();
const speechEndCallbacks = new Set<
@ -109,22 +154,38 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
const speechCallbacks = new Set<
(result: DictationAdapter.Result) => void
>();
const endCallbacks = new Set<() => void>();
let stream: MediaStream | null = null;
let finalTranscript = "";
let ended = false;
let started = false;
let stopping = false;
let stopRequestedAt = 0;
let stopFallbackTimer = 0;
let stopLevelMeter = () => {
// Replaced after microphone access succeeds.
};
const interimParts = new Map<number, string>();
let resolveEnded: (() => void) | null = null;
const endedPromise = new Promise<void>((resolve) => {
resolveEnded = resolve;
});
const session: DictationAdapter.Session = {
const session: StudioDictationSession = {
status: { type: "starting" },
stop: async () => {
if (!ended && started) {
stopping = true;
stopRequestedAt = performance.now();
recognition.stop();
// Ending the track now gives the browser an audio endpoint to finalize
// and releases the mic without waiting on its remote speech service.
stopLevelMeter();
stopStream(stream);
stream = null;
scheduleStopFallback();
} else if (!ended) {
finish("stopped");
}
@ -132,11 +193,9 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
},
cancel: () => {
if (!ended && started) {
recognition.abort();
} else if (!ended) {
finish("cancelled");
}
if (ended) return;
if (started) recognition.abort();
finish("cancelled");
},
onSpeechStart: (callback) => {
@ -159,23 +218,75 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
speechCallbacks.delete(callback);
};
},
// Beyond DictationAdapter: lets callers reset UI when the session ends on
// its own (silence, error), not just via stop().
onEnd: (callback: () => void) => {
endCallbacks.add(callback);
return () => {
endCallbacks.delete(callback);
};
},
};
const currentInterimTranscript = () =>
[...interimParts.entries()]
.sort(([a], [b]) => a - b)
.map(([, transcript]) => transcript.trim())
.filter(Boolean)
.join(" ");
const promoteInterim = () => {
const interim = currentInterimTranscript();
interimParts.clear();
if (!interim) return false;
const corrected = applyDictationDictionary(interim).trim();
if (!corrected) return false;
finalTranscript = finalTranscript
? `${finalTranscript} ${corrected}`
: corrected;
return true;
};
function scheduleStopFallback(): void {
if (!stopping || ended || stopFallbackTimer) return;
const elapsed = performance.now() - stopRequestedAt;
const delay = Math.max(0, STOP_FINALIZATION_GRACE_MS - elapsed);
stopFallbackTimer = window.setTimeout(() => {
stopFallbackTimer = 0;
if (!ended) {
promoteInterim();
finish("stopped");
recognition.abort();
}
}, delay);
}
const finish = (reason: "stopped" | "cancelled" | "error") => {
if (ended) return;
ended = true;
stopping = false;
if (stopFallbackTimer) window.clearTimeout(stopFallbackTimer);
stopFallbackTimer = 0;
session.status = { type: "ended", reason };
stopLevelMeter();
stopStream(stream);
stream = null;
if (finalTranscript) {
for (const callback of speechEndCallbacks) {
callback({ transcript: finalTranscript });
const transcript = reason === "cancelled" ? "" : finalTranscript;
if (transcript) {
for (const callback of speechCallbacks) {
callback({ transcript, isFinal: true });
}
if (reason !== "cancelled") {
recordRecentDictation(finalTranscript);
}
finalTranscript = "";
recordRecentDictation(transcript, sessionChatId);
}
// assistant-ui uses this callback to leave dictation mode; required even
// for silence and cancelled recordings.
for (const callback of speechEndCallbacks) {
callback({ transcript });
}
finalTranscript = "";
interimParts.clear();
for (const callback of endCallbacks) callback();
resolveEnded?.();
};
@ -188,6 +299,7 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
});
recognition.addEventListener("result", (event) => {
if (ended) return;
const speechEvent = event as SpeechRecognitionEvent;
for (
let i = speechEvent.resultIndex;
@ -198,6 +310,7 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
if (!result) continue;
const transcript = result[0]?.transcript ?? "";
if (result.isFinal) {
interimParts.delete(i);
const corrected = applyDictationDictionary(transcript);
// Join final chunks with a single space so recorded transcripts do
// not merge words when a browser omits leading whitespace.
@ -207,25 +320,34 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
? `${finalTranscript} ${trimmed}`
: trimmed;
}
for (const callback of speechCallbacks) {
callback({ transcript: corrected, isFinal: true });
}
} else {
for (const callback of speechCallbacks) {
callback({ transcript, isFinal: false });
}
interimParts.set(i, transcript);
}
}
const interim = currentInterimTranscript();
if (interim) {
scheduleStopFallback();
}
});
recognition.addEventListener("end", () => {
if (ended) {
return;
}
promoteInterim();
finish("stopped");
});
recognition.addEventListener("error", (event) => {
if (ended) return;
const errorEvent = event as SpeechRecognitionErrorEvent;
if (errorEvent.error === "aborted") {
finish("cancelled");
if (stopping) {
promoteInterim();
finish("stopped");
} else {
finish("cancelled");
}
return;
}
const description = describeSpeechError(
@ -233,7 +355,22 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
errorEvent.message,
);
console.error("Dictation error:", errorEvent.error, errorEvent.message);
toast.error(description);
if (errorEvent.error === "network") {
// Online speech service unreachable; point the user to the offline
// local engine (the toast opens Voice settings).
toast.error("No internet connection", {
id: NETWORK_TOAST_ID,
description:
"Browser dictation needs the online speech service. Switch to Local in Voice settings and pick a local STT model to dictate offline.",
action: {
label: "Open Voice settings",
onClick: () =>
useSettingsDialogStore.getState().openDialog("voice"),
},
});
} else {
toast.error(description);
}
finish("error");
});
@ -275,6 +412,7 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
"NotFoundError",
);
}
stopLevelMeter = startDictationLevelMeter(stream);
try {
recognition.start(audioTrack);
} catch (error) {
@ -285,6 +423,7 @@ export class StudioWebSpeechDictationAdapter implements DictationAdapter {
"Dictation start(audioTrack) failed; retrying start().",
error,
);
stopLevelMeter();
stopStream(stream);
stream = null;
recognition.start();

22
studio/frontend/src/features/chat/index.ts
View file

@ -119,3 +119,25 @@ export {
updateChatProjectInstructions,
useChatProjects,
} from "./hooks/use-chat-projects";
export { subscribeDictationLevel } from "./adapters/dictation-level";
export {
StudioDictationAdapter,
cancelActiveStudioDictation,
isStudioDictationAvailable,
notifyStudioDictationUnavailable,
} from "./adapters/studio-dictation-adapter";
export {
StudioModelDictationAdapter,
fetchSttStatus,
loadSttModel,
startSttDownload,
unloadSttModel,
validateSttModel,
type SttDownloadStatus,
} from "./adapters/studio-model-dictation-adapter";
export {
StudioSpeechSynthesisAdapter,
createConfiguredUtterance,
curateSystemVoices,
generateStudioTtsAudio,
} from "./adapters/studio-speech-synthesis-adapter";

13
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -33,8 +33,8 @@ import {
useRef,
} from "react";
import { toast } from "sonner";
import { StudioDictationAdapter } from "./adapters/studio-dictation-adapter";
import { StudioSpeechSynthesisAdapter } from "./adapters/studio-speech-synthesis-adapter";
import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter";
import {
ThreadAutosaveHandle,
createOpenAIStreamAdapter,
@ -1195,13 +1195,10 @@ function useStudioRuntimeAdapters(
[aui, modelType, pairId],
);
const dictation = useMemo(
() =>
StudioWebSpeechDictationAdapter.isSupported()
? new StudioWebSpeechDictationAdapter()
: undefined,
[],
);
// Always register the adapter so the mic stays clickable for any engine. The
// engine is resolved at listen() time and the composer shows guidance when it
// cannot run, so engine switches also work on an already-mounted thread.
const dictation = useMemo(() => new StudioDictationAdapter(), []);
const speech = useMemo(
() =>
StudioSpeechSynthesisAdapter.isSupported()

247
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -24,16 +24,12 @@ import {
} from "@/components/ui/dropdown-menu";
import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
import {
describeMediaError,
describeSpeechError,
isMissingDeviceError,
} from "@/features/chat/adapters/studio-web-speech-dictation-adapter";
import {
applyDictationDictionary,
recordRecentDictation,
resolveDictationLanguage,
useVoiceSettingsStore,
} from "@/features/settings/stores/voice-settings-store";
StudioDictationAdapter,
isStudioDictationAvailable,
notifyStudioDictationUnavailable,
} from "@/features/chat/adapters/studio-dictation-adapter";
import type { StudioDictationSession } from "@/features/chat/adapters/studio-web-speech-dictation-adapter";
import { useVoiceSettingsStore } from "@/features/settings/stores/voice-settings-store";
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
import { isTauri } from "@/lib/api-base";
import { isDownloadCancelled } from "@/lib/native-files";
@ -224,173 +220,95 @@ function formatReasoningDisabledLabel(
function useDictation(
setText: (value: string | ((prev: string) => string)) => void,
) {
// Re-render support state when the user switches recognition engines.
const dictationEngine = useVoiceSettingsStore((s) => s.dictationEngine);
const [isDictating, setIsDictating] = useState(false);
const recognitionRef = useRef<SpeechRecognition | null>(null);
const streamRef = useRef<MediaStream | null>(null);
// True while a stopped recording's final audio is still transcribing; a
// second click then cancels the pending transcription instead of re-stopping.
const [isFinalizing, setIsFinalizing] = useState(false);
const sessionRef = useRef<StudioDictationSession | null>(null);
const startingRef = useRef(false);
// Guards the getUserMedia await so a mic opened after unmount is released.
const disposedRef = useRef(false);
const releaseStream = useCallback(() => {
streamRef.current?.getTracks().forEach((track) => track.stop());
streamRef.current = null;
}, []);
const finalizingRef = useRef(false);
const start = useCallback(async () => {
const SpeechRecognitionAPI =
typeof window !== "undefined" &&
(window.SpeechRecognition ??
(
window as unknown as {
webkitSpeechRecognition?: typeof SpeechRecognition;
}
).webkitSpeechRecognition);
if (!SpeechRecognitionAPI) {
if (startingRef.current || sessionRef.current) return;
// Unsupported engine (e.g. Firefox): explain and steer to the local model.
if (!isStudioDictationAvailable()) {
notifyStudioDictationUnavailable();
return;
}
if (startingRef.current || recognitionRef.current) return;
startingRef.current = true;
// Open the microphone chosen in Voice settings, matching the main chat
// adapter, so Compare dictation honors the same device selection.
let audioTrack: MediaStreamTrack | undefined;
const { micDeviceId } = useVoiceSettingsStore.getState();
if (navigator.mediaDevices?.getUserMedia) {
try {
let stream: MediaStream;
try {
stream = await navigator.mediaDevices.getUserMedia({
audio:
micDeviceId && micDeviceId !== "default"
? { deviceId: { exact: micDeviceId } }
: true,
});
} catch (error) {
// Saved mic may be unplugged; fall back to the default device.
if (micDeviceId !== "default" && isMissingDeviceError(error)) {
stream = await navigator.mediaDevices.getUserMedia({
audio: true,
});
} else {
throw error;
}
}
streamRef.current = stream;
audioTrack = stream.getAudioTracks()[0];
} catch (error) {
// Permission/security failure: report it and stop instead of silently
// recording from a different default device, matching the main adapter.
startingRef.current = false;
releaseStream();
setIsDictating(false);
toast.error(describeMediaError(error));
return;
}
}
if (disposedRef.current) {
releaseStream();
startingRef.current = false;
return;
}
const recognition = new SpeechRecognitionAPI() as SpeechRecognition;
recognition.continuous = true;
recognition.interimResults = true;
recognition.lang = resolveDictationLanguage();
let sessionTranscript = "";
recognition.onresult = (event: SpeechRecognitionEvent) => {
// Iterate every result from resultIndex; a single event can carry more
// than one finalized phrase and dropping the rest loses dictated words.
for (let i = event.resultIndex; i < event.results.length; i++) {
const result = event.results[i];
if (!result?.isFinal) continue;
const transcript = applyDictationDictionary(
result[0]?.transcript?.trim() ?? "",
);
if (!transcript) continue;
sessionTranscript = sessionTranscript
? `${sessionTranscript} ${transcript}`
: transcript;
setText((prev) => (prev ? `${prev} ${transcript}` : transcript));
}
};
recognition.onerror = (event) => {
// Report speech-service failures like the main adapter; aborted is a
// normal stop, not an error.
const errorEvent = event as SpeechRecognitionErrorEvent;
if (errorEvent.error !== "aborted") {
toast.error(describeSpeechError(errorEvent.error, errorEvent.message));
}
setIsDictating(false);
};
recognition.onend = () => {
// A stop()+immediate restart can install a new recognizer before this
// old one ends; only tear down shared refs when we are still current.
if (recognitionRef.current === recognition) {
releaseStream();
recognitionRef.current = null;
setIsDictating(false);
}
if (sessionTranscript) {
recordRecentDictation(sessionTranscript);
sessionTranscript = "";
}
};
let session: StudioDictationSession;
try {
if (audioTrack) {
try {
recognition.start(audioTrack);
} catch {
// No start(track) overload: recognition captures from the default
// device, so release the selected-device stream.
releaseStream();
recognition.start();
}
} else {
recognition.start();
}
// Routes to the engine chosen in Voice settings (browser or STT model),
// honoring the selected microphone, language, and dictionary. Compare
// feeds two panes, so recent dictations must not link the unrelated
// single-chat active thread.
session = new StudioDictationAdapter({ chatId: null }).listen();
} catch {
startingRef.current = false;
releaseStream();
notifyStudioDictationUnavailable();
return;
}
recognitionRef.current = recognition;
startingRef.current = false;
sessionRef.current = session;
setIsDictating(true);
}, [setText, releaseStream]);
// Append final transcripts; the adapter has already applied the dictionary
// and records the session in Recent dictations.
session.onSpeech((result) => {
if (!result.isFinal) return;
const transcript = result.transcript?.trim() ?? "";
if (transcript) {
setText((prev) => (prev ? `${prev} ${transcript}` : transcript));
}
});
session.onEnd?.(() => {
if (sessionRef.current === session) sessionRef.current = null;
finalizingRef.current = false;
setIsFinalizing(false);
setIsDictating(false);
});
startingRef.current = false;
}, [setText]);
const stop = useCallback(() => {
if (recognitionRef.current) {
recognitionRef.current.stop();
recognitionRef.current = null;
const session = sessionRef.current;
if (!session) return;
// A second click while the final segment is transcribing discards the
// pending transcription instead of leaving the pane stuck until timeout.
if (finalizingRef.current) {
session.cancel();
if (sessionRef.current === session) sessionRef.current = null;
finalizingRef.current = false;
setIsFinalizing(false);
setIsDictating(false);
return;
}
releaseStream();
setIsDictating(false);
}, [releaseStream]);
finalizingRef.current = true;
setIsFinalizing(true);
// Keep the session and dictation state alive while its final audio segment
// is transcribed. onEnd clears both after the transcript callbacks run.
void session.stop().catch((error) => {
console.error("Could not stop dictation:", error);
session.cancel();
if (sessionRef.current === session) sessionRef.current = null;
finalizingRef.current = false;
setIsFinalizing(false);
setIsDictating(false);
});
}, []);
useEffect(() => {
disposedRef.current = false;
return () => {
disposedRef.current = true;
if (recognitionRef.current) {
recognitionRef.current.abort();
}
streamRef.current?.getTracks().forEach((track) => track.stop());
streamRef.current = null;
sessionRef.current?.cancel();
sessionRef.current = null;
};
}, []);
const supported =
typeof window !== "undefined" &&
!!(
window.SpeechRecognition ??
(window as unknown as { webkitSpeechRecognition?: unknown })
.webkitSpeechRecognition
);
const supported = StudioDictationAdapter.isSupported(dictationEngine);
return { isDictating, start, stop, supported };
return { isDictating, isFinalizing, start, stop, supported };
}
export type CompareHandles = MutableRefObject<Record<string, CompareHandle>>;
@ -835,9 +753,9 @@ export function SharedComposer({
const {
isDictating,
isFinalizing: isDictationFinalizing,
start: startDictation,
stop: stopDictation,
supported: dictationSupported,
} = useDictation(setText);
useEffect(() => {
@ -1483,7 +1401,7 @@ export function SharedComposer({
}
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
if (!busy) {
if (!busy && !isDictating) {
send();
}
}
@ -1494,7 +1412,8 @@ export function SharedComposer({
pendingImages.length > 0 ||
pendingAudio !== null) &&
!busy &&
!isComposing;
!isComposing &&
!isDictating;
// Adjustable "+" menu items, keyed by id. Pinned ones render at the top
// level; the rest fall into the "More" overflow submenu. Core items (photos,
@ -2269,7 +2188,7 @@ export function SharedComposer({
</button>
)
) : null}
{dictationSupported && (
{
<>
{!isDictating ? (
<TooltipIconButton
@ -2285,19 +2204,27 @@ export function SharedComposer({
</TooltipIconButton>
) : (
<TooltipIconButton
tooltip="Stop dictation"
tooltip={
isDictationFinalizing
? "Cancel transcription"
: "Stop dictation"
}
side="bottom"
variant="ghost"
size="icon"
className="size-8 rounded-full text-destructive"
onClick={stopDictation}
aria-label="Stop dictation"
aria-label={
isDictationFinalizing
? "Cancel transcription"
: "Stop dictation"
}
>
<SquareIcon className="size-3 animate-pulse fill-current" />
</TooltipIconButton>
)}
</>
)}
}
{isQueueRunning ? (
<button
type="button"

1
studio/frontend/src/features/hub/index.ts
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export {
DownloadProgressBar,
downloadManager,
jobKeyOf,
subscribeJobListeners,

38
studio/frontend/src/features/hub/lib/hidden-models.ts
View file

@ -6,18 +6,36 @@ import { getInventoryVersion } from "../stores/inventory-events";
// Infra models hidden from browse/preview lists (Hub Discover, the chat model
// selector, and local on-device rows). Mirrors the backend
// `utils.hidden_models`: the RAG embedding model and the llama.cpp validation
// probe are not usable chat models. Server-confirmed cache rows are trusted
// because the backend applies variant-aware filtering. Optimistic cache rows
// still use these needles until the server confirms them. The dynamic matchers
// fetched from `/api/hub/hidden-models` add the user's configured embedder as
// exact repo ids and exact resolved paths, never substring needles. Per-repo
// views are not filtered, so reinstall flows still show downloaded files.
// `utils.hidden_models`: the RAG embedding model, STT dictation models, and the
// llama.cpp validation probe are not usable chat models. Server-confirmed cache
// rows are trusted because the backend applies variant-aware filtering.
// Optimistic cache rows still use these needles until the server confirms them.
// The dynamic matchers fetched from `/api/hub/hidden-models` add the user's
// configured embedder as exact repo ids and exact resolved paths, never
// substring needles. Per-repo views are not filtered, so reinstall flows still
// show downloaded files.
const HIDDEN_NEEDLES = [
"bge-small-en-v1.5", // RAG embedder: unsloth/bge-small-en-v1.5[-GGUF]
"ggml-org/models", // llama.cpp validation probe repo
"stories260k.gguf", // probe filename (carries .gguf so it stays specific)
];
const HIDDEN_STT_REPOS = new Set([
// Transformers safetensors repos and their whisper.cpp GGUF companions
// (unslothai/whisper-*-GGUF): STT-only, never chat models.
"unsloth/whisper-tiny",
"unsloth/whisper-base",
"unsloth/whisper-small",
"unsloth/whisper-large-v3-turbo",
"unsloth/whisper-large-v3",
"unslothai/whisper-tiny-gguf",
"unslothai/whisper-base-gguf",
"unslothai/whisper-small-gguf",
"unslothai/whisper-large-v3-turbo-gguf",
"unslothai/whisper-large-v3-gguf",
]);
const HIDDEN_STT_CACHE_NAMES = [...HIDDEN_STT_REPOS].map((repo) =>
repo.replace("/", "--"),
);
let dynamicNeedles: readonly string[] = [];
let dynamicExactIds: readonly string[] = [];
@ -81,7 +99,13 @@ export function isHiddenModelId(
return false;
}
const lower = v.toLowerCase();
const normalized = lower.trim().replace(/^\/+|\/+$/g, "");
const pathParts = lower.split(/[\\/]/);
return (
HIDDEN_STT_REPOS.has(normalized) ||
HIDDEN_STT_CACHE_NAMES.some((name) =>
pathParts.includes(`models--${name}`),
) ||
HIDDEN_NEEDLES.some((needle) => lower.includes(needle)) ||
dynamicNeedles.some((needle) => lower.includes(needle)) ||
dynamicExactIds.includes(lower) ||

22
studio/frontend/src/features/settings/components/archived-chats-dialog.tsx
View file

@ -22,10 +22,14 @@ import {
import { toast } from "@/lib/toast";
import { ArchiveRestoreIcon, Delete02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Button } from "@/components/ui/button";
import { useNavigate, useRouterState } from "@tanstack/react-router";
import { useState } from "react";
import { useSettingsDialogStore } from "../stores/settings-dialog-store";
/** Archived chats shown per page; "Show more" reveals the next page. */
const ARCHIVED_PAGE_SIZE = 20;
function formatCreatedAt(ms: number): string {
return new Date(ms).toLocaleDateString(undefined, {
year: "numeric",
@ -54,6 +58,9 @@ export function ArchivedChatsView() {
const [confirmingDelete, setConfirmingDelete] = useState<SidebarItem | null>(
null,
);
// Pagination: the view remounts with its settings tab, so plain state
// restarts from the first page on each visit.
const [visibleCount, setVisibleCount] = useState(ARCHIVED_PAGE_SIZE);
// Open an archived chat: leave it archived, just navigate to it.
function openChat(item: SidebarItem) {
@ -113,7 +120,7 @@ export function ArchivedChatsView() {
<span className="w-32 shrink-0">Date created</span>
<span className="w-16 shrink-0" />
</div>
{archivedItems.map((item) => (
{archivedItems.slice(0, visibleCount).map((item) => (
<div
key={item.id}
className="group flex items-center gap-4 border-b border-border/40 px-1 py-2.5 text-sm last:border-0"
@ -159,6 +166,19 @@ export function ArchivedChatsView() {
</span>
</div>
))}
{archivedItems.length > visibleCount ? (
<div className="flex justify-center pt-3">
<Button
variant="outline"
size="sm"
onClick={() =>
setVisibleCount(visibleCount + ARCHIVED_PAGE_SIZE)
}
>
Show more ({archivedItems.length - visibleCount})
</Button>
</div>
) : null}
</div>
)}

132
studio/frontend/src/features/settings/components/dictation-dictionary-view.tsx Normal file
View file

@ -0,0 +1,132 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { useT } from "@/i18n";
import {
ArrowLeft01Icon,
Delete02Icon,
PlusSignIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useState } from "react";
import { useVoiceSettingsStore } from "../stores/voice-settings-store";
// Full-page editor for the dictation dictionary. Kept on its own subpage so a
// long list of entries does not crowd the main Voice settings.
export function DictationDictionaryView({ onBack }: { onBack: () => void }) {
const t = useT();
const dictionary = useVoiceSettingsStore((s) => s.dictionary);
const addDictionaryEntry = useVoiceSettingsStore((s) => s.addDictionaryEntry);
const updateDictionaryEntry = useVoiceSettingsStore(
(s) => s.updateDictionaryEntry,
);
const commitDictionaryEntry = useVoiceSettingsStore(
(s) => s.commitDictionaryEntry,
);
const removeDictionaryEntry = useVoiceSettingsStore(
(s) => s.removeDictionaryEntry,
);
const [newEntry, setNewEntry] = useState("");
const handleAddEntry = () => {
const trimmed = newEntry.trim();
if (!trimmed) return;
addDictionaryEntry(trimmed);
setNewEntry("");
};
return (
<div className="flex flex-col gap-6">
<header className="flex items-center gap-2">
<button
type="button"
onClick={onBack}
aria-label={t("settings.voice.dictionary.backToVoice")}
className="inline-flex size-7 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
</button>
<h1 className="font-heading text-xl font-semibold">
{t("settings.voice.title")}
</h1>
</header>
<div className="flex flex-col gap-1">
<h2 className="text-sm font-semibold">
{t("settings.voice.dictionary.sectionTitle")}
</h2>
<p className="text-xs text-muted-foreground">
{t("settings.voice.dictionary.sectionDescription")}
</p>
</div>
<div className="flex flex-col">
{dictionary.map((entry, index) => (
<div
// biome-ignore lint/suspicious/noArrayIndexKey: entries are editable in place
key={index}
className="flex items-center gap-2 py-1.5"
>
<Input
value={entry}
onChange={(e) => updateDictionaryEntry(index, e.target.value)}
// Tabbing to this row's remove button must not commit-splice the
// row first, which shifts indices and deletes the wrong entry.
onBlur={(e) => {
if (
e.relatedTarget instanceof HTMLElement &&
e.relatedTarget.dataset.removeIndex === String(index)
) {
return;
}
commitDictionaryEntry(index);
}}
className="h-8 flex-1 text-sm"
aria-label={`Dictionary entry ${index + 1}`}
/>
<Button
variant="ghost"
size="icon"
className="size-8 shrink-0 text-muted-foreground hover:text-destructive"
data-remove-index={index}
// Keep the click from blurring an empty input first, which would
// commit-splice this row and make onClick delete the next one.
onMouseDown={(e) => e.preventDefault()}
onClick={() => removeDictionaryEntry(index)}
aria-label={`Remove dictionary entry ${index + 1}`}
>
<HugeiconsIcon icon={Delete02Icon} className="size-3.5" />
</Button>
</div>
))}
<div className="flex items-center gap-2 py-1.5">
<Input
value={newEntry}
onChange={(e) => setNewEntry(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter") {
e.preventDefault();
handleAddEntry();
}
}}
placeholder="Jane Doe"
className="h-8 flex-1 text-sm"
aria-label="New dictionary entry"
/>
<Button
variant="outline"
size="sm"
className="shrink-0"
onClick={handleAddEntry}
disabled={!newEntry.trim()}
>
<HugeiconsIcon icon={PlusSignIcon} className="mr-1.5 size-3.5" />
{t("settings.voice.dictionary.addEntry")}
</Button>
</div>
</div>
</div>
);
}

515
studio/frontend/src/features/settings/components/recent-dictations-view.tsx Normal file
View file

@ -0,0 +1,515 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
InputGroup,
InputGroupAddon,
InputGroupInput,
} from "@/components/ui/input-group";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import {
deleteChatItem,
useChatRuntimeStore,
type SidebarItem,
} from "@/features/chat";
import { useT } from "@/i18n";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { toast } from "@/lib/toast";
import {
ArrowLeft01Icon,
Copy01Icon,
Delete02Icon,
Message01Icon,
Search01Icon,
ViewIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate } from "@tanstack/react-router";
import { useMemo, useState } from "react";
import {
type RecentDictation,
useVoiceSettingsStore,
} from "../stores/voice-settings-store";
import { useSettingsDialogStore } from "../stores/settings-dialog-store";
/** Dictations shown per page; "Show more" reveals the next page. */
const PAGE_SIZE = 20;
type PendingDelete =
| { kind: "one"; dictation: RecentDictation }
| { kind: "all" }
| null;
type SortOrder = "newest" | "oldest" | "az";
function formatDictationDate(ms: number): string {
return new Date(ms).toLocaleString(undefined, {
year: "numeric",
month: "short",
day: "numeric",
hour: "numeric",
minute: "2-digit",
});
}
export function RecentDictationsView({
selectedId,
onSelect,
onBack,
}: {
selectedId: string | null;
onSelect: (id: string | null) => void;
onBack: () => void;
}) {
const t = useT();
const recentDictations = useVoiceSettingsStore((s) => s.recentDictations);
const removeRecentDictation = useVoiceSettingsStore(
(s) => s.removeRecentDictation,
);
const clearRecentDictations = useVoiceSettingsStore(
(s) => s.clearRecentDictations,
);
const [pendingDelete, setPendingDelete] = useState<PendingDelete>(null);
const [search, setSearch] = useState("");
const [sortOrder, setSortOrder] = useState<SortOrder>("newest");
// Pagination is keyed to the search/sort inputs so changing either restarts
// from the first page without an effect.
const pageKey = `${sortOrder}::${search}`;
const [page, setPage] = useState({ key: pageKey, count: PAGE_SIZE });
const visibleCount = page.key === pageKey ? page.count : PAGE_SIZE;
const navigate = useNavigate();
const closeSettings = useSettingsDialogStore((s) => s.closeDialog);
const selected =
recentDictations.find((dictation) => dictation.id === selectedId) ?? null;
function openChat(chatId: string) {
closeSettings();
void navigate({ to: "/chat", search: { thread: chatId } });
}
const visible = useMemo(() => {
const query = search.trim().toLowerCase();
const filtered = query
? recentDictations.filter((d) => d.text.toLowerCase().includes(query))
: recentDictations;
const sorted = [...filtered];
if (sortOrder === "oldest") {
sorted.sort((a, b) => a.at - b.at);
} else if (sortOrder === "az") {
sorted.sort((a, b) => a.text.localeCompare(b.text));
} else {
sorted.sort((a, b) => b.at - a.at);
}
return sorted;
}, [recentDictations, search, sortOrder]);
async function handleCopy(text: string) {
if (await copyToClipboard(text)) {
toast.success(t("settings.voice.recents.copied"));
} else {
toast.error(t("settings.voice.recents.copyFailed"));
}
}
function confirmDelete() {
try {
if (pendingDelete?.kind === "one") {
removeRecentDictation(pendingDelete.dictation.id);
if (pendingDelete.dictation.id === selectedId) {
onSelect(null);
}
} else if (pendingDelete?.kind === "all") {
clearRecentDictations();
onSelect(null);
}
} finally {
setPendingDelete(null);
}
}
// Delete a linked dictation together with the chat it was used in.
async function confirmDeleteWithChat() {
if (pendingDelete?.kind !== "one") {
return;
}
const dictation = pendingDelete.dictation;
setPendingDelete(null);
if (dictation.chatId) {
const item: SidebarItem = {
type: "single",
id: dictation.chatId,
title: "",
createdAt: 0,
updatedAt: 0,
};
try {
await deleteChatItem(
item,
useChatRuntimeStore.getState().activeThreadId ?? undefined,
() => {
// The deleted chat was open; leave the user on a fresh chat.
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
});
},
);
} catch (error) {
toast.error(t("settings.voice.recents.deleteWithChatFailed"), {
description: error instanceof Error ? error.message : undefined,
});
return;
}
}
removeRecentDictation(dictation.id);
if (dictation.id === selectedId) {
onSelect(null);
}
}
return (
<div className="flex flex-col gap-6">
<header className="flex items-center gap-2">
<button
type="button"
onClick={() => (selected ? onSelect(null) : onBack())}
aria-label={
selected
? t("settings.voice.recents.backToRecents")
: t("settings.voice.recents.backToVoice")
}
className="inline-flex size-7 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
</button>
<h1 className="font-heading text-xl font-semibold">
{t("settings.voice.title")}
</h1>
</header>
<div className="flex flex-col gap-1">
<h2 className="text-sm font-semibold">
{selected
? t("settings.voice.recents.detailTitle")
: t("settings.voice.recents.sectionTitle")}
</h2>
<p className="text-xs text-muted-foreground">
{selected
? formatDictationDate(selected.at)
: t("settings.voice.recents.pageDescription")}
</p>
</div>
{selected ? (
<div className="flex min-h-0 flex-col gap-4">
<article className="rounded-lg border border-border/60 bg-muted/20 p-4">
<p className="whitespace-pre-wrap break-words text-sm leading-relaxed text-foreground select-text">
{selected.text}
</p>
</article>
<div className="flex justify-end gap-2">
{selected.chatId ? (
<Button
variant="outline"
size="sm"
onClick={() => openChat(selected.chatId as string)}
>
<HugeiconsIcon
icon={Message01Icon}
className="mr-1.5 size-3.5"
/>
{t("settings.voice.recents.openChat")}
</Button>
) : null}
<Button
variant="outline"
size="sm"
onClick={async () => {
await handleCopy(selected.text);
}}
>
<HugeiconsIcon icon={Copy01Icon} className="mr-1.5 size-3.5" />
{t("settings.voice.recents.copy")}
</Button>
<Button
variant="outline"
size="sm"
onClick={() =>
setPendingDelete({ kind: "one", dictation: selected })
}
className="text-destructive hover:border-destructive/60 hover:text-destructive"
>
<HugeiconsIcon icon={Delete02Icon} className="mr-1.5 size-3.5" />
{t("settings.voice.recents.delete")}
</Button>
</div>
</div>
) : recentDictations.length === 0 ? (
<p className="py-8 text-center text-sm text-muted-foreground">
{t("settings.voice.recents.empty")}
</p>
) : (
<div className="flex flex-col gap-3">
<div className="flex items-center gap-2">
<InputGroup className="h-8 flex-1">
<InputGroupAddon align="inline-start">
<HugeiconsIcon
icon={Search01Icon}
strokeWidth={2}
className="size-3.5 text-muted-foreground"
/>
</InputGroupAddon>
<InputGroupInput
value={search}
onChange={(e) => setSearch(e.target.value)}
placeholder={t("settings.voice.recents.searchPlaceholder")}
aria-label={t("settings.voice.recents.searchPlaceholder")}
className="text-sm"
/>
</InputGroup>
<Select
value={sortOrder}
onValueChange={(value) => setSortOrder(value as SortOrder)}
>
<SelectTrigger
size="sm"
className="w-36 shrink-0"
aria-label={t("settings.voice.recents.sortLabel")}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="newest">
{t("settings.voice.recents.sortNewest")}
</SelectItem>
<SelectItem value="oldest">
{t("settings.voice.recents.sortOldest")}
</SelectItem>
<SelectItem value="az">
{t("settings.voice.recents.sortAlpha")}
</SelectItem>
</SelectContent>
</Select>
<Button
variant="outline"
size="sm"
onClick={() => setPendingDelete({ kind: "all" })}
className="shrink-0 text-destructive hover:border-destructive/60 hover:text-destructive"
>
<HugeiconsIcon icon={Delete02Icon} className="mr-1.5 size-3.5" />
{t("settings.voice.recents.clear")}
</Button>
</div>
{visible.length === 0 ? (
<p className="py-8 text-center text-sm text-muted-foreground">
{t("settings.voice.recents.noMatches")}
</p>
) : (
<div className="flex flex-col">
<div className="flex items-center gap-4 border-b border-border/60 px-1 pb-2 text-xs font-semibold text-foreground">
<span className="min-w-0 flex-1">
{t("settings.voice.recents.dictationColumn")}
</span>
<span className="hidden w-40 shrink-0 items-center gap-1.5 sm:flex">
{/* Spacer matching the row's chat-link icon slot so the
header starts exactly where the dates start. */}
<span className="size-3.5 shrink-0" />
{t("settings.voice.recents.dateColumn")}
</span>
<span className="w-24 shrink-0" />
</div>
{visible.slice(0, visibleCount).map((dictation) => (
<div
key={dictation.id}
className="group flex items-stretch gap-2 border-b border-border/40 last:border-0"
>
<button
type="button"
onClick={() =>
dictation.chatId
? openChat(dictation.chatId)
: onSelect(dictation.id)
}
aria-label={
dictation.chatId
? t("settings.voice.recents.openChat")
: t("settings.voice.recents.view")
}
className="flex min-w-0 flex-1 items-start gap-4 rounded-md px-1 py-3 text-left transition-colors hover:bg-muted/50 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<span className="min-w-0 flex-1">
<span className="line-clamp-3 whitespace-pre-wrap break-words text-sm text-foreground">
{dictation.text}
</span>
<span className="mt-1 flex items-center gap-1.5 text-xs text-muted-foreground sm:hidden">
{dictation.chatId ? (
<HugeiconsIcon
icon={Message01Icon}
className="size-3"
/>
) : null}
{formatDictationDate(dictation.at)}
</span>
</span>
<span className="hidden w-40 shrink-0 items-center gap-1.5 text-left text-xs text-muted-foreground tabular-nums sm:flex">
{/* Fixed icon slot keeps every date starting at the
same column whether or not a chat is linked. */}
<span className="flex size-3.5 shrink-0 items-center justify-center">
{dictation.chatId ? (
<HugeiconsIcon
icon={Message01Icon}
className="size-3.5"
aria-label={t("settings.voice.recents.openChat")}
/>
) : null}
</span>
<span className="whitespace-nowrap">
{formatDictationDate(dictation.at)}
</span>
</span>
</button>
<span className="flex w-24 shrink-0 items-center justify-end gap-1">
{dictation.chatId ? (
<button
type="button"
onClick={() => onSelect(dictation.id)}
aria-label={t("settings.voice.recents.view")}
title={t("settings.voice.recents.view")}
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={ViewIcon}
strokeWidth={1.75}
className="size-4"
/>
</button>
) : null}
<button
type="button"
onClick={async () => {
await handleCopy(dictation.text);
}}
aria-label={t("settings.voice.recents.copy")}
title={t("settings.voice.recents.copy")}
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={Copy01Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
<button
type="button"
onClick={() =>
setPendingDelete({ kind: "one", dictation })
}
aria-label={t("settings.voice.recents.delete")}
title={t("settings.voice.recents.delete")}
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={Delete02Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
</span>
</div>
))}
</div>
)}
{visible.length > visibleCount ? (
<div className="flex justify-center pt-1">
<Button
variant="outline"
size="sm"
onClick={() =>
setPage({ key: pageKey, count: visibleCount + PAGE_SIZE })
}
>
{t("settings.voice.recents.showMore", {
count: visible.length - visibleCount,
})}
</Button>
</div>
) : null}
</div>
)}
<AlertDialog
open={pendingDelete !== null}
onOpenChange={(open) => {
if (!open) {
setPendingDelete(null);
}
}}
>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>
{pendingDelete?.kind === "all"
? t("settings.voice.recents.clearTitle")
: t("settings.voice.recents.deleteTitle")}
</AlertDialogTitle>
<AlertDialogDescription>
{pendingDelete?.kind === "all"
? t("settings.voice.recents.clearDescription")
: pendingDelete?.kind === "one" &&
pendingDelete.dictation.chatId
? t("settings.voice.recents.deleteLinkedDescription")
: t("settings.voice.recents.deleteDescription")}
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>{t("common.cancel")}</AlertDialogCancel>
{pendingDelete?.kind === "one" && pendingDelete.dictation.chatId ? (
<AlertDialogAction
variant="outline"
className="text-destructive hover:border-destructive/60 hover:text-destructive"
onClick={(event) => {
// Close only through state so the dismissal can't race a
// click on whatever ends up under the pointer.
event.preventDefault();
void confirmDeleteWithChat();
}}
>
{t("settings.voice.recents.deleteWithChat")}
</AlertDialogAction>
) : null}
<AlertDialogAction
variant="destructive"
onClick={(event) => {
event.preventDefault();
confirmDelete();
}}
>
{pendingDelete?.kind === "all"
? t("settings.voice.recents.clearConfirm")
: t("common.delete")}
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

1
studio/frontend/src/features/settings/settings-search.ts
View file

@ -108,7 +108,6 @@ export const SETTINGS_SEARCH_INDEX: Record<SettingsTab, TranslationKey[]> = {
"settings.voice.dictation.sectionTitle",
"settings.voice.dictation.microphoneLabel",
"settings.voice.dictation.languageLabel",
"settings.voice.dictation.testLabel",
"settings.voice.dictionary.sectionTitle",
"settings.voice.recents.sectionTitle",
"settings.voice.readAloud.sectionTitle",

300
studio/frontend/src/features/settings/stores/voice-settings-store.ts
View file

@ -2,28 +2,118 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { create } from "zustand";
import { persist } from "zustand/middleware";
import { createJSONStorage, persist } from "zustand/middleware";
// Voice preferences in localStorage. Adapters read them at call time so
// changes apply without reloading the chat runtime.
export interface RecentDictation {
id: string;
text: string;
at: number;
/** Chat the dictation was spoken into, when one was open at the time. */
chatId?: string;
}
const MAX_RECENT_DICTATIONS = 20;
// Dictation history is kept in full; the list view paginates. QUOTA_TRIM_KEEP is
// the emergency floor if localStorage runs out of room (see persist wrapper).
const QUOTA_TRIM_KEEP = 200;
// Cap stored transcript length so a few long dictations cannot bloat the
// persisted blob and trip a synchronous localStorage quota error on save.
const MAX_RECENT_DICTATION_LENGTH = 2000;
const MAX_DICTIONARY_ENTRIES = 100;
const MAX_DICTIONARY_ENTRY_LENGTH = 120;
/** Five curated Whisper choices, mirrored by the backend (stt_sidecar.py). */
export const STT_MODELS = [
"tiny",
"base",
"small",
"large-v3-turbo",
"large-v3",
] as const;
export type DefaultSttModel = (typeof STT_MODELS)[number];
/** A curated id or a user-selected Hugging Face `owner/model` repository. */
export type SttModel = string;
/** Whisper repos downloaded through Studio's existing Model Hub manager. */
export const STT_MODEL_REPOS: Record<DefaultSttModel, string> = {
tiny: "unsloth/whisper-tiny",
base: "unsloth/whisper-base",
small: "unsloth/whisper-small",
"large-v3-turbo": "unsloth/whisper-large-v3-turbo",
"large-v3": "unsloth/whisper-large-v3",
};
export const DEFAULT_STT_MODEL: DefaultSttModel = "small";
const HF_REPO_ID =
/^[A-Za-z0-9][A-Za-z0-9._-]{0,95}\/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$/;
export function isSttModelId(value: string): boolean {
const normalized = value.trim();
return (
(STT_MODELS as readonly string[]).includes(normalized) ||
HF_REPO_ID.test(normalized)
);
}
export function normalizeSttModel(value: unknown): SttModel {
if (typeof value !== "string") {
return DEFAULT_STT_MODEL;
}
const normalized = value.trim();
return isSttModelId(normalized) ? normalized : DEFAULT_STT_MODEL;
}
export function getSttModelRepo(model: SttModel): string {
return STT_MODEL_REPOS[model as DefaultSttModel] ?? normalizeSttModel(model);
}
// All curated models are multilingual. Custom `.en` checkpoints are treated as
// English-only so a later language change falls back safely.
export const ENGLISH_ONLY_STT_MODELS: ReadonlySet<SttModel> = new Set([]);
/** Whether a model can honor the selected dictation language. */
export function isSttModelLanguageCompatible(
model: SttModel,
language: string,
): boolean {
const isEnglishOnly =
ENGLISH_ONLY_STT_MODELS.has(model) ||
getSttModelRepo(model).toLowerCase().endsWith(".en");
if (!isEnglishOnly) {
return true;
}
const normalized = language.trim().replaceAll("_", "-").toLowerCase();
// Auto sends no forced language, which English-only checkpoints accept.
return normalized === "auto" || normalized.split("-", 1)[0] === "en";
}
export type DictationEngine = "browser" | "model";
/**
* Whether a model id is one of the five curated Whisper choices. Curated models
* run GGML through whisper.cpp; custom repos are safetensors and run through
* Transformers.
*/
export function isCuratedSttModel(model: SttModel): boolean {
return (STT_MODELS as readonly string[]).includes(model.trim());
}
export interface VoiceSettingsState {
/** Input device for dictation. "default" = system default microphone. */
micDeviceId: string;
setMicDeviceId: (value: string) => void;
/**
* "browser": Web Speech API. "model": local transcription; the model decides
* the backend (whisper.cpp for curated GGML, Transformers for custom repos).
*/
dictationEngine: DictationEngine;
setDictationEngine: (value: DictationEngine) => void;
/** STT model to use when dictationEngine is "model". */
sttModel: SttModel;
setSttModel: (value: SttModel) => void;
/** BCP 47 tag for speech recognition, or "auto" for the browser locale. */
dictationLanguage: string;
setDictationLanguage: (value: string) => void;
@ -38,7 +128,8 @@ export interface VoiceSettingsState {
/** Final transcripts, newest first, so text can be recovered. */
recentDictations: RecentDictation[];
addRecentDictation: (text: string) => void;
addRecentDictation: (text: string, chatId?: string) => void;
removeRecentDictation: (id: string) => void;
clearRecentDictations: () => void;
/** Show the read-aloud button on assistant responses. */
@ -61,14 +152,81 @@ export interface VoiceSettingsState {
setTtsVolume: (value: number) => void;
}
/**
* localStorage wrapper that keeps the full dictation history until the browser's
* quota is hit, then drops the oldest entries instead of losing the whole save.
*/
const quotaSafeLocalStorage = {
getItem: (key: string) => localStorage.getItem(key),
removeItem: (key: string) => localStorage.removeItem(key),
setItem: (key: string, value: string) => {
try {
localStorage.setItem(key, value);
return;
} catch {
// Quota exceeded: trim dictation history, oldest first, and retry.
}
let parsed: { state?: { recentDictations?: RecentDictation[] } };
try {
parsed = JSON.parse(value) as typeof parsed;
} catch {
return;
}
const state = parsed.state;
const recents = state?.recentDictations;
if (!state || !Array.isArray(recents)) return;
// Halve the history until the save fits, down to an empty history, so even
// a small history shrinks when another store already consumed the quota.
let keep = Math.min(QUOTA_TRIM_KEEP, recents.length);
for (;;) {
keep = keep >= recents.length ? Math.floor(recents.length / 2) : keep;
try {
state.recentDictations = recents.slice(0, keep);
localStorage.setItem(key, JSON.stringify(parsed));
return;
} catch {
// Still over quota: trim harder.
}
if (keep === 0) return;
keep = Math.floor(keep / 2);
}
},
};
export const useVoiceSettingsStore = create<VoiceSettingsState>()(
persist(
(set) => ({
micDeviceId: "default",
setMicDeviceId: (micDeviceId) => set({ micDeviceId }),
dictationEngine: "browser",
setDictationEngine: (dictationEngine) => set({ dictationEngine }),
sttModel: DEFAULT_STT_MODEL,
setSttModel: (value) =>
set((state) => {
const sttModel = normalizeSttModel(value);
return {
sttModel: isSttModelLanguageCompatible(
sttModel,
state.dictationLanguage,
)
? sttModel
: DEFAULT_STT_MODEL,
};
}),
dictationLanguage: "auto",
setDictationLanguage: (dictationLanguage) => set({ dictationLanguage }),
setDictationLanguage: (dictationLanguage) =>
set((state) => ({
dictationLanguage,
sttModel: isSttModelLanguageCompatible(
state.sttModel,
dictationLanguage,
)
? state.sttModel
: DEFAULT_STT_MODEL,
})),
dictionary: [],
addDictionaryEntry: (value) =>
@ -111,17 +269,31 @@ export const useVoiceSettingsStore = create<VoiceSettingsState>()(
})),
recentDictations: [],
addRecentDictation: (text) =>
addRecentDictation: (text, chatId) =>
set((state) => {
const trimmed = text.trim().slice(0, MAX_RECENT_DICTATION_LENGTH);
if (!trimmed) return state;
if (!trimmed) {
return state;
}
const at = Date.now();
return {
recentDictations: [
{ text: trimmed, at: Date.now() },
{
id: `${at}-${Math.random().toString(36).slice(2, 10)}`,
text: trimmed,
at,
...(chatId ? { chatId } : {}),
},
...state.recentDictations,
].slice(0, MAX_RECENT_DICTATIONS),
],
};
}),
removeRecentDictation: (id) =>
set((state) => ({
recentDictations: state.recentDictations.filter(
(dictation) => dictation.id !== id,
),
})),
clearRecentDictations: () => set({ recentDictations: [] }),
ttsEnabled: true,
@ -142,30 +314,37 @@ export const useVoiceSettingsStore = create<VoiceSettingsState>()(
}),
{
name: "unsloth_voice_settings",
storage: createJSONStorage(() => quotaSafeLocalStorage),
merge: (persisted, current) => {
const saved = persisted as Partial<VoiceSettingsState> | undefined;
const dictationLanguage = asString(saved?.dictationLanguage, "auto");
// "gguf" was a short-lived separate engine choice; both local
// backends now live under "model".
const savedEngine = saved?.dictationEngine as string | undefined;
const dictationEngine: DictationEngine =
savedEngine === "model" || savedEngine === "gguf"
? "model"
: "browser";
const savedSttModel = normalizeSttModel(saved?.sttModel);
const sttModel = isSttModelLanguageCompatible(
savedSttModel,
dictationLanguage,
)
? savedSttModel
: DEFAULT_STT_MODEL;
return {
...current,
micDeviceId: asString(saved?.micDeviceId, "default"),
dictationLanguage: asString(saved?.dictationLanguage, "auto"),
dictationEngine,
sttModel,
dictationLanguage,
dictionary: Array.isArray(saved?.dictionary)
? saved.dictionary
.filter((v): v is string => typeof v === "string" && !!v.trim())
.map((v) => v.trim().slice(0, MAX_DICTIONARY_ENTRY_LENGTH))
.slice(0, MAX_DICTIONARY_ENTRIES)
: [],
recentDictations: Array.isArray(saved?.recentDictations)
? saved.recentDictations
.filter(
(v): v is RecentDictation =>
typeof v?.text === "string" && typeof v?.at === "number",
)
.slice(0, MAX_RECENT_DICTATIONS)
.map((v) => ({
text: v.text.slice(0, MAX_RECENT_DICTATION_LENGTH),
at: v.at,
}))
: [],
recentDictations: normalizeRecentDictations(saved?.recentDictations),
ttsEnabled:
typeof saved?.ttsEnabled === "boolean" ? saved.ttsEnabled : true,
ttsEngine: saved?.ttsEngine === "studio" ? "studio" : "system",
@ -183,6 +362,40 @@ function asString(value: unknown, fallback: string): string {
return typeof value === "string" && value ? value : fallback;
}
function normalizeRecentDictations(value: unknown): RecentDictation[] {
if (!Array.isArray(value)) {
return [];
}
const normalized: RecentDictation[] = [];
for (const [index, entry] of value.entries()) {
if (!entry || typeof entry !== "object") {
continue;
}
const candidate = entry as Partial<RecentDictation>;
if (
typeof candidate.text !== "string" ||
!candidate.text.trim() ||
typeof candidate.at !== "number" ||
!Number.isFinite(candidate.at)
) {
continue;
}
normalized.push({
id:
typeof candidate.id === "string" && candidate.id
? candidate.id
: `legacy-${candidate.at}-${index}`,
text: candidate.text.trim().slice(0, MAX_RECENT_DICTATION_LENGTH),
at: candidate.at,
...(typeof candidate.chatId === "string" && candidate.chatId
? { chatId: candidate.chatId }
: {}),
});
}
return normalized;
}
function clampNumber(
value: unknown,
min: number,
@ -202,6 +415,47 @@ export function resolveDictationLanguage(setting?: string): string {
: "en-US";
}
// Whisper's language codes: transformers `LANGUAGES` keys plus the backend's
// BCP-47 aliases (stt_sidecar.py `_WHISPER_LANGUAGE_ALIASES`). Keep in sync with
// `_known_whisper_languages()`; Auto only resolves to a code in this set, so a
// UI locale Whisper cannot honor stays on auto-detect.
const WHISPER_DICTATION_LANGUAGES = new Set([
"af", "am", "ar", "as", "az", "ba", "be", "bg", "bn", "bo", "br", "bs",
"ca", "cmn", "cs", "cy", "da", "de", "el", "en", "es", "et", "eu", "fa",
"fi", "fil", "fo", "fr", "gl", "gu", "ha", "haw", "he", "hi", "hr", "ht",
"hu", "hy", "id", "in", "is", "it", "iw", "ja", "ji", "jw", "ka", "kk",
"km", "kn", "ko", "la", "lb", "ln", "lo", "lt", "lv", "mg", "mi", "mk",
"ml", "mn", "mr", "ms", "mt", "my", "nb", "ne", "nl", "nn", "no", "oc",
"pa", "pl", "ps", "pt", "ro", "ru", "sa", "sd", "si", "sk", "sl", "sn",
"so", "sq", "sr", "su", "sv", "sw", "ta", "te", "tg", "th", "tk", "tl",
"tr", "tt", "uk", "ur", "uz", "vi", "yi", "yo", "yue", "zh",
]);
/**
* Resolve Auto for the model STT engine (the browser engine resolves it via
* `resolveDictationLanguage`). Only the literal "auto" resolves to a concrete
* locale; an explicit or malformed setting passes through unchanged. Gated so
* Auto only becomes a language the model AND Whisper can honor, else it stays
* auto-detect rather than forcing a locale Whisper cannot handle (e.g. Irish)
* or 422ing every dictation.
*/
export function resolveModelDictationLanguage(
model: SttModel,
requested: string,
): string {
if (requested !== "auto") return requested;
const resolved = resolveDictationLanguage(requested);
const primary = resolved
.trim()
.replaceAll("_", "-")
.toLowerCase()
.split("-", 1)[0];
return isSttModelLanguageCompatible(model, resolved) &&
WHISPER_DICTATION_LANGUAGES.has(primary)
? resolved
: requested;
}
function escapeRegExp(value: string): string {
return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}
@ -240,6 +494,6 @@ export function applyDictationDictionary(
}
/** Record a finished dictation so it can be recovered from settings. */
export function recordRecentDictation(text: string): void {
useVoiceSettingsStore.getState().addRecentDictation(text);
export function recordRecentDictation(text: string, chatId?: string): void {
useVoiceSettingsStore.getState().addRecentDictation(text, chatId);
}

1163
studio/frontend/src/features/settings/tabs/voice-tab.tsx
View file

File diff suppressed because it is too large Load diff

2
studio/frontend/src/hooks/index.ts
View file

@ -1,7 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { useDebouncedValue } from "./use-debounced-value";
export { useGpuInfo } from "./use-gpu-info";
export { useGpuUtilization } from "./use-gpu-utilization";
@ -11,3 +10,4 @@ export { useHfTokenValidation } from "./use-hf-token-validation";
export { useTauriBackend } from "./use-tauri-backend";
export { useCollapseScrollLock } from "./use-collapse-scroll-lock";
export { useSystemInfo } from "./use-system";
export { useWheelScrollRef } from "./use-wheel-scroll-ref";

34
studio/frontend/src/hooks/use-llama-update-check.ts
View file

@ -30,6 +30,7 @@ export interface LlamaUpdateJob {
export interface LlamaUpdateStatus {
supported: boolean;
update_available: boolean;
component: "llama.cpp" | "whisper.cpp";
installed_tag: string | null;
latest_tag: string | null;
// Prebuilt download size in bytes, if known.
@ -55,13 +56,30 @@ function parseJob(value: unknown): LlamaUpdateJob {
function parseStatus(value: unknown): LlamaUpdateStatus | null {
if (!value || typeof value !== "object") return null;
const s = value as Record<string, unknown>;
const component =
s.update_component === "whisper" ? "whisper.cpp" : "llama.cpp";
const whisper =
s.whisper && typeof s.whisper === "object"
? (s.whisper as Record<string, unknown>)
: null;
// Legacy top-level version fields intentionally retain their llama meaning.
// A whisper-only update must display the nested whisper release instead of
// presenting equal llama tags as a new llama update.
const details = component === "whisper.cpp" && whisper ? whisper : s;
return {
supported: s.supported === true,
update_available: s.update_available === true,
installed_tag: typeof s.installed_tag === "string" ? s.installed_tag : null,
latest_tag: typeof s.latest_tag === "string" ? s.latest_tag : null,
component,
installed_tag:
typeof details.installed_tag === "string"
? details.installed_tag
: null,
latest_tag:
typeof details.latest_tag === "string" ? details.latest_tag : null,
update_size_bytes:
typeof s.update_size_bytes === "number" ? s.update_size_bytes : null,
typeof details.update_size_bytes === "number"
? details.update_size_bytes
: null,
job: parseJob(s.job),
};
}
@ -163,8 +181,12 @@ export function useLlamaUpdateCheck({
// can drop or double-fire the notification.
const notifyReloadIfNeeded = useCallback(
(job: Pick<LlamaUpdateJob, "state" | "reload_required" | "finished_at">) => {
// "error" is included for partial chained updates: the llama phase can
// land (and unload the server) before a later phase fails, and the
// backend keeps reload_required set in exactly that case. Without the
// resync the chat UI would keep pointing at the unloaded model.
if (
job.state === "success" &&
(job.state === "success" || job.state === "error") &&
job.reload_required &&
job.finished_at !== reloadNotifiedForRef.current
) {
@ -202,7 +224,9 @@ export function useLlamaUpdateCheck({
reloadRequired: s.job.reload_required,
});
} else if (s.job.state === "error") {
// Keep the banner visible so retry is available.
// Keep the banner visible so retry is available. A partial chained
// update can still have unloaded the llama server before failing.
notifyReloadIfNeeded(s.job);
onDone?.({ ok: false, error: s.job.error });
} else {
onDone?.({ ok: false, error: "update did not complete" });

43
studio/frontend/src/hooks/use-wheel-scroll-ref.ts Normal file
View file

@ -0,0 +1,43 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useRef } from "react";
/**
* Callback ref for scroll containers inside modal scroll locks. The lock may
* cancel native wheel scrolling, so apply the delta before it reaches the
* portaled dialog boundary.
*/
export function useWheelScrollRef<T extends HTMLElement>() {
const detachRef = useRef<(() => void) | null>(null);
return useCallback((node: T | null) => {
detachRef.current?.();
detachRef.current = null;
if (!node) {
return;
}
const onWheel = (event: WheelEvent) => {
if (event.ctrlKey || node.scrollHeight <= node.clientHeight) {
return;
}
const multiplier =
event.deltaMode === WheelEvent.DOM_DELTA_LINE
? 16
: event.deltaMode === WheelEvent.DOM_DELTA_PAGE
? node.clientHeight
: 1;
const delta = event.deltaY * multiplier;
if (delta === 0) {
return;
}
node.scrollTop += delta;
event.preventDefault();
event.stopPropagation();
};
node.addEventListener("wheel", onWheel, { passive: false });
detachRef.current = () => node.removeEventListener("wheel", onWheel);
}, []);
}

89
studio/frontend/src/i18n/locales/en.ts
View file

@ -105,9 +105,45 @@ export const en = {
},
voice: {
title: "Voice",
description: "Microphone, dictation, and read-aloud",
description: "Microphone, dictation, speech-to-text and read-aloud",
dictation: {
sectionTitle: "Dictation",
engineLabel: "Dictation engine",
engineBrowser: "Browser",
engineBrowserDescription:
"Transcribes audio using your browser speech service. Select 'Local transcription' to use a STT model.",
engineModel: "Local transcription",
engineModelDescription:
"Runs a speech-to-text (STT) model locally and works offline. Download, load, then unloads after inactivity.",
sttModelLabel: "Speech recognition model",
sttModelDescription: "Choose or search a STT model to run locally.",
sttModelSearchPlaceholder: "Search model",
sttModelSearching: "Searching Hugging Face…",
sttModelValidating: "Checking Whisper compatibility…",
sttModelNoResults: "No Whisper models found",
sttModelInvalid: "This repository cannot be used for dictation",
sttModelFailed: "Could not load the STT model",
sttModelUnsupported: "Recording is not supported in this browser",
sttChecking: "Checking…",
sttOnDemand: "Downloaded",
sttLoadingModel: "Loading model…",
sttReady: "Loaded on {device}",
sttLoaded: "Loaded",
sttUnavailable:
"Not installed on this server. Run `unsloth studio update` to enable local dictation.",
sttRetry: "Retry",
sttDownloadChecking: "Checking download status…",
sttNotDownloaded: "Not downloaded",
sttDownloadStatusFailed: "Could not check download status",
sttDownload: "Download",
sttDownloading: "Downloading… {progress}%",
sttCancelDownload: "Cancel",
sttCancellingDownload: "Cancelling…",
sttDownloadComplete: "Speech recognition model downloaded",
sttDownloadFailed: "Could not download the speech recognition model",
sttLoad: "Load",
sttUnload: "Unload",
sttUnloading: "Unloading…",
microphoneLabel: "Microphone",
microphoneDescription: "Used for dictation",
microphoneFallbackHint:
@ -118,35 +154,58 @@ export const en = {
"Microphone access was blocked. Allow microphone access for this Unsloth page, then try again.",
micAccessUnsupported:
"Microphone access is not supported in this browser or context.",
micOpenFailed:
"Could not open the selected microphone. Check permissions or pick another device.",
systemDefault: "System default",
savedMicDisconnected: "Saved microphone (not connected)",
languageLabel: "Dictation language",
languageDescription: "Language to recognize",
languageAuto: "Auto (browser language)",
testLabel: "Test dictation",
testDescription: "Speak to check your mic and settings",
startTest: "Start test",
stopTest: "Stop test",
listening: "Listening…",
testSaved: "Saved to recent dictations",
notSupported: "Not supported in this browser",
},
dictionary: {
sectionTitle: "Dictation dictionary",
sectionDescription:
"Apply the spelling entered here when dictation recognizes the same words or phrase",
sectionDescription: "Set how dictation spells specific words or phrases",
manageLabel: "Custom spellings",
manage: "Manage",
backToVoice: "Back to Voice",
addEntry: "Add entry",
},
recents: {
sectionTitle: "Recent dictations",
sectionTitle: "Dictation history",
sectionDescription:
"Your recent dictations will appear here so you can recover text",
"Every dictation is saved here so you can recover text",
manageLabel: "Dictation history",
manage: "Manage",
pageDescription:
"Every dictation is saved. View, copy, or delete them, or open the chat a dictation was used in.",
searchPlaceholder: "Search dictations",
sortLabel: "Sort dictations",
sortNewest: "Newest",
sortOldest: "Oldest",
sortAlpha: "A to Z",
noMatches: "No dictations match your search",
detailTitle: "Saved dictation",
backToVoice: "Back to Voice",
backToRecents: "Back to recent dictations",
view: "View full dictation",
empty: "No dictations yet",
dictationColumn: "Dictation",
dateColumn: "Date created",
copy: "Copy dictation",
copied: "Copied to clipboard",
copyFailed: "Could not copy to clipboard",
clear: "Clear recent dictations",
delete: "Delete dictation",
deleteTitle: "Delete dictation",
deleteDescription:
"Delete this saved dictation? This cannot be undone.",
deleteLinkedDescription:
"Delete this saved dictation? You can also delete the chat it was used in. This cannot be undone.",
deleteWithChat: "Delete chat and dictation",
deleteWithChatFailed: "Could not delete the chat",
clear: "Clear history",
clearTitle: "Clear dictation history",
clearDescription: "Delete all saved dictations? This cannot be undone.",
clearConfirm: "Clear all",
showMore: "Show more ({count})",
openChat: "Open chat",
},
readAloud: {
sectionTitle: "Read aloud",

6
studio/frontend/src/index.css
View file

@ -1674,6 +1674,12 @@ html[data-chat-font] .aui-root {
@apply flex w-full flex-wrap items-center gap-0.5 px-1;
}
/* Recording is a dedicated one-line mode. Active tool pills must not force
the composer back into the normal two-row layout. */
.unsloth-composer-line[data-dictating="true"] {
@apply flex-nowrap;
}
.unsloth-composer-left {
@apply flex min-w-0 flex-wrap items-center gap-0.5;
order: 1;

1258
studio/install_llama_prebuilt.py
View file

File diff suppressed because it is too large Load diff

1377
studio/install_whisper_prebuilt.py Normal file
View file

File diff suppressed because it is too large Load diff

2427
studio/prebuilt_core.py Normal file
View file

File diff suppressed because it is too large Load diff

84
studio/setup.ps1
View file

@ -2655,12 +2655,13 @@ if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) {
}
$StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome)
# Directory-local evidence that Unsloth created $Path, used to adopt a custom-home
# llama.cpp predating the .unsloth-studio-owned marker (see setup.sh). Only the
# prebuilt UNSLOTH_PREBUILT_INFO.json counts; source builds are indistinguishable
# from a user clone on Windows and stay under the strict guard.
# llama.cpp or whisper.cpp predating the .unsloth-studio-owned marker (see
# setup.sh). Only Unsloth prebuilt markers count; source builds are
# indistinguishable from a user clone on Windows and stay under the strict guard.
function Test-StudioOwnedAdoptable {
param([Parameter(Mandatory = $true)][string]$Path)
if (Test-Path -LiteralPath (Join-Path $Path "UNSLOTH_PREBUILT_INFO.json") -PathType Leaf) { return $true }
if (Test-Path -LiteralPath (Join-Path $Path "UNSLOTH_WHISPER_PREBUILT_INFO.json") -PathType Leaf) { return $true }
return $false
}
function Assert-StudioOwnedOrAbsent {
@ -3767,6 +3768,83 @@ if ($LocalLlamaCppLinked) {
}
}
# ==========================================================================
# PHASE 3.4: Install the whisper.cpp prebuilt (dictation runtime)
# ==========================================================================
# Mirrors the llama.cpp prebuilt install above; current whisper releases are
# slim bundles that reuse the llama install's ggml runtime, so this runs after
# llama. Failure is never fatal: local dictation falls back to Transformers STT.
$WhisperCppDir = Join-Path $UnslothHome "whisper.cpp"
$WhisperInstaller = Join-Path $PSScriptRoot "install_whisper_prebuilt.py"
# Same opt-outs as setup.sh: a user-configured binary/dir or an explicit skip
# disables the managed install entirely.
if ($env:WHISPER_SERVER_PATH -or $env:UNSLOTH_WHISPER_CPP_PATH) {
substep "whisper.cpp: using a user-configured binary/dir; skipping managed install"
} elseif ($env:UNSLOTH_SKIP_WHISPER_INSTALL -eq "1") {
substep "whisper.cpp: install skipped (UNSLOTH_SKIP_WHISPER_INSTALL=1)"
} elseif (Test-Path -LiteralPath $WhisperInstaller) {
# The installer's atomic activation replaces the whole directory, so the
# custom-home ownership guard must run first (mirrors the llama block).
if ($StudioHomeIsCustom) {
Assert-StudioOwnedOrAbsent -Path $WhisperCppDir -Label "whisper.cpp install"
}
$whisperArgs = @($WhisperInstaller, "--install-dir", $WhisperCppDir)
if ($env:UNSLOTH_WHISPER_RELEASE_TAG) {
$whisperArgs += @("--published-release-tag", $env:UNSLOTH_WHISPER_RELEASE_TAG)
}
if ($script:ROCmGfxArch) {
$whisperArgs += @("--rocm-gfx", $script:ROCmGfxArch)
} elseif ($HasROCm) {
$whisperArgs += "--has-rocm"
}
$prevEAPWhisper = $ErrorActionPreference
$ErrorActionPreference = "Continue"
$previousNativeErrorPreferenceW = $null
$restoreNativeErrorPreferenceW = $false
if ($PSVersionTable.PSVersion.Major -ge 7) {
$previousNativeErrorPreferenceW = $PSNativeCommandUseErrorActionPreference
$PSNativeCommandUseErrorActionPreference = $false
$restoreNativeErrorPreferenceW = $true
}
try {
$whisperOutput = & python @whisperArgs 2>&1 | Out-String
$whisperExit = $LASTEXITCODE
} finally {
if ($restoreNativeErrorPreferenceW) {
$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreferenceW
}
}
$ErrorActionPreference = $prevEAPWhisper
if ($whisperExit -eq 0) {
if ($whisperOutput -match "already matches") {
step "whisper.cpp" "prebuilt up to date"
} else {
step "whisper.cpp" "prebuilt installed"
}
if ($StudioHomeIsCustom -and (Test-Path -LiteralPath $WhisperCppDir -PathType Container)) {
Mark-StudioOwned -Path $WhisperCppDir
}
} elseif ($whisperExit -eq 3) {
step "whisper.cpp" "install busy; keeping existing runtime" "Yellow"
} elseif ($whisperExit -eq 2) {
$requiredWhisperLlamaTag = "unknown"
if ($whisperOutput -match "slim bundle requires llama\.cpp ([^;\s]+)") {
$requiredWhisperLlamaTag = $Matches[1]
}
$installedWhisperLlamaTag = "unknown"
$llamaMarker = Join-Path $LlamaCppDir "UNSLOTH_PREBUILT_INFO.json"
if (Test-Path -LiteralPath $llamaMarker -PathType Leaf) {
try {
$markerPayload = Get-Content -LiteralPath $llamaMarker -Raw | ConvertFrom-Json
if ($markerPayload.release_tag) { $installedWhisperLlamaTag = $markerPayload.release_tag }
} catch {}
}
step "whisper.cpp" "no compatible prebuilt (installed llama.cpp $installedWhisperLlamaTag; whisper requires $requiredWhisperLlamaTag); curated whisper.cpp dictation is unavailable; publish paired releases in llama.cpp then whisper.cpp order; browser and Transformers dictation remain available" "Yellow"
} else {
step "whisper.cpp" "prebuilt install failed; curated whisper.cpp dictation is unavailable; retry setup or inspect verbose output; browser and Transformers dictation remain available" "Yellow"
}
}
# ==========================================================================
# PHASE 3.5: Install OpenSSL dev (for HTTPS support in llama-server)
# ==========================================================================

99
studio/setup.sh
View file

@ -537,11 +537,13 @@ if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then
fi
# Directory-local evidence Unsloth created "$1": only prebuilt-installer metadata
# counts (UNSLOTH_PREBUILT_INFO.json for llama.cpp, UNSLOTH_NODE_PREBUILT_INFO.json
# for Node), both written only by our installers. Mirrors the setup.ps1 Node guard.
# A markerless source build stays strict since this runs right before an rm -rf.
# for Node, UNSLOTH_WHISPER_PREBUILT_INFO.json for whisper.cpp), all written only
# by our installers. Mirrors the setup.ps1 Node guard. A markerless source build
# stays strict since this runs right before an rm -rf.
_studio_owned_adoptable() {
[ -f "$1/UNSLOTH_PREBUILT_INFO.json" ] && return 0
[ -f "$1/UNSLOTH_NODE_PREBUILT_INFO.json" ] && return 0
[ -f "$1/UNSLOTH_WHISPER_PREBUILT_INFO.json" ] && return 0
return 1
}
_assert_studio_owned_or_absent() {
@ -1965,6 +1967,99 @@ if [ ! -L "$LLAMA_CPP_DIR" ] && {
_remove_agent_instruction_files "$LLAMA_CPP_DIR"
fi
# ── whisper.cpp (local speech-to-text dictation engine) ──
# Optional runtime for local dictation. Fail-open: any failure leaves the
# Transformers STT engine and browser dictation working, so it never aborts
# setup (unlike llama.cpp). Runs in 'unsloth studio update' too so the runtime
# installs/refreshes without a compiler. Installs beside llama.cpp under the
# same managed home the sidecar's _managed_whisper_cpp_dir() resolves.
WHISPER_CPP_DIR="$UNSLOTH_HOME/whisper.cpp"
if [ -n "${WHISPER_SERVER_PATH:-}" ] || [ -n "${UNSLOTH_WHISPER_CPP_PATH:-}" ]; then
verbose_substep "whisper.cpp: using a user-configured binary/dir; skipping managed install"
elif [ "${UNSLOTH_SKIP_WHISPER_INSTALL:-0}" = "1" ]; then
verbose_substep "whisper.cpp: install skipped (UNSLOTH_SKIP_WHISPER_INSTALL=1)"
else
if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then
_assert_studio_owned_or_absent "$WHISPER_CPP_DIR" "whisper.cpp install"
fi
_WHISPER_CMD=(python "$SCRIPT_DIR/install_whisper_prebuilt.py" --install-dir "$WHISPER_CPP_DIR")
if [ -n "${UNSLOTH_WHISPER_RELEASE_TAG:-}" ]; then
_WHISPER_CMD+=(--published-release-tag "$UNSLOTH_WHISPER_RELEASE_TAG")
fi
if [ -n "${_setup_gfx:-}" ]; then
_WHISPER_CMD+=(--rocm-gfx "$_setup_gfx")
elif [ "$_setup_amd_detected" = true ]; then
_WHISPER_CMD+=(--has-rocm)
fi
_WHISPER_LOG="$(mktemp)"
set +e
if _is_verbose; then
"${_WHISPER_CMD[@]}" 2>&1 | tee "$_WHISPER_LOG"
_WHISPER_STATUS=${PIPESTATUS[0]}
else
"${_WHISPER_CMD[@]}" >"$_WHISPER_LOG" 2>&1
_WHISPER_STATUS=$?
fi
set -e
if [ "$_WHISPER_STATUS" -eq 0 ]; then
if grep -Fq "already matches" "$_WHISPER_LOG"; then
step "whisper.cpp" "prebuilt up to date"
else
step "whisper.cpp" "prebuilt installed"
fi
if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then
: > "$WHISPER_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
fi
rm -f "$_WHISPER_LOG"
elif [ "$_WHISPER_STATUS" -eq 3 ]; then
# A warm dictation server holds the binary; keep the old install.
step "whisper.cpp" "install busy; keeping existing runtime" "$C_WARN"
rm -f "$_WHISPER_LOG"
else
# A source build is opt-in. Keep the installer log until fallback has
# finished so setup can distinguish release skew from an operational
# installer failure and report the exact pairing when available.
_WHISPER_RECOVERED=false
_WHISPER_BUILD="$SCRIPT_DIR/../scripts/build_whisper_cpp.sh"
if [ "${UNSLOTH_WHISPER_FORCE_COMPILE:-0}" = "1" ] && [ -f "$_WHISPER_BUILD" ] \
&& command -v cmake >/dev/null 2>&1 && command -v git >/dev/null 2>&1; then
substep "whisper.cpp prebuilt unavailable; building from source (UNSLOTH_WHISPER_FORCE_COMPILE=1)..."
# The source build overwrites whisper-server in the managed dir but
# knows nothing about the prebuilt marker; a stale marker would make
# a later setup run report "already matches" and skip repairing the
# prebuilt over the source binary. Drop it before building.
rm -f "$WHISPER_CPP_DIR/UNSLOTH_WHISPER_PREBUILT_INFO.json" 2>/dev/null || true
if run_quiet_no_exit "whisper.cpp source build" sh "$_WHISPER_BUILD"; then
_WHISPER_RECOVERED=true
step "whisper.cpp" "source build installed"
if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then
: > "$WHISPER_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
fi
else
:
fi
fi
if [ "$_WHISPER_RECOVERED" != true ]; then
if [ "$_WHISPER_STATUS" -eq 2 ]; then
_WHISPER_REQUIRED_TAG="$(sed -n 's/.*slim bundle requires llama\.cpp \([^; ]*\).*/\1/p' "$_WHISPER_LOG" | tail -n 1)"
_WHISPER_INSTALLED_TAG="$(python - "$UNSLOTH_HOME/llama.cpp/UNSLOTH_PREBUILT_INFO.json" <<'PY' 2>/dev/null || true
import json, sys
try:
print(json.load(open(sys.argv[1], encoding="utf-8")).get("release_tag", ""))
except Exception:
pass
PY
)"
_WHISPER_PAIRING="installed llama.cpp ${_WHISPER_INSTALLED_TAG:-unknown}; whisper requires ${_WHISPER_REQUIRED_TAG:-unknown}"
step "whisper.cpp" "no compatible prebuilt ($_WHISPER_PAIRING); curated whisper.cpp dictation is unavailable; publish the paired releases in llama.cpp then whisper.cpp order; browser and Transformers dictation remain available" "$C_WARN"
else
step "whisper.cpp" "prebuilt install failed; curated whisper.cpp dictation is unavailable; retry setup or inspect verbose output; browser and Transformers dictation remain available" "$C_WARN"
fi
fi
rm -f "$_WHISPER_LOG"
fi
fi
# ── Footer ──
if [ "$_LLAMA_ONLY" = "1" ]; then
echo ""

104
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -22,7 +22,6 @@ sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback
extract_archive = INSTALL_LLAMA_PREBUILT.extract_archive
binary_env = INSTALL_LLAMA_PREBUILT.binary_env
is_secret_env_name = INSTALL_LLAMA_PREBUILT.is_secret_env_name
scrub_env = INSTALL_LLAMA_PREBUILT.scrub_env
@ -105,105 +104,10 @@ def approved_checksums_for(
)
def test_extract_archive_allows_safe_tar_symlink_chain(tmp_path: Path):
archive_path = tmp_path / "bundle.tar.gz"
payload = b"shared-object"
with tarfile.open(archive_path, "w:gz") as archive:
versioned = tarfile.TarInfo("libllama.so.0.0.1")
versioned.size = len(payload)
archive.addfile(versioned, io_bytes(payload))
soname = tarfile.TarInfo("libllama.so.0")
soname.type = tarfile.SYMTYPE
soname.linkname = "libllama.so.0.0.1"
archive.addfile(soname)
linker_name = tarfile.TarInfo("libllama.so")
linker_name.type = tarfile.SYMTYPE
linker_name.linkname = "libllama.so.0"
archive.addfile(linker_name)
destination = tmp_path / "extract"
extract_archive(archive_path, destination)
assert (destination / "libllama.so.0.0.1").read_bytes() == payload
assert (destination / "libllama.so.0").is_symlink()
assert (destination / "libllama.so").is_symlink()
assert (destination / "libllama.so").resolve().read_bytes() == payload
def test_extract_archive_allows_safe_tar_hardlink(tmp_path: Path):
archive_path = tmp_path / "bundle.tar.gz"
payload = b"quantize"
with tarfile.open(archive_path, "w:gz") as archive:
target = tarfile.TarInfo("llama-quantize")
target.size = len(payload)
archive.addfile(target, io_bytes(payload))
hardlink = tarfile.TarInfo("llama-quantize-copy")
hardlink.type = tarfile.LNKTYPE
hardlink.linkname = "llama-quantize"
archive.addfile(hardlink)
destination = tmp_path / "extract"
extract_archive(archive_path, destination)
assert (destination / "llama-quantize-copy").read_bytes() == payload
assert not (destination / "llama-quantize-copy").is_symlink()
def test_extract_archive_rejects_absolute_tar_symlink_target(tmp_path: Path):
archive_path = tmp_path / "bundle.tar.gz"
with tarfile.open(archive_path, "w:gz") as archive:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "/tmp/libllama.so.0"
archive.addfile(entry)
with pytest.raises(PrebuiltFallback, match = "archive link used an absolute target"):
extract_archive(archive_path, tmp_path / "extract")
def test_extract_archive_rejects_escaping_tar_symlink_target(tmp_path: Path):
archive_path = tmp_path / "bundle.tar.gz"
with tarfile.open(archive_path, "w:gz") as archive:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "../outside/libllama.so.0"
archive.addfile(entry)
with pytest.raises(PrebuiltFallback, match = "archive link escaped destination"):
extract_archive(archive_path, tmp_path / "extract")
def test_extract_archive_rejects_unresolved_tar_symlink_target(tmp_path: Path):
archive_path = tmp_path / "bundle.tar.gz"
with tarfile.open(archive_path, "w:gz") as archive:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "libllama.so.0"
archive.addfile(entry)
with pytest.raises(PrebuiltFallback, match = "unresolved link entries"):
extract_archive(archive_path, tmp_path / "extract")
def test_extract_archive_rejects_zip_symlink_entry(tmp_path: Path):
archive_path = tmp_path / "bundle.zip"
with zipfile.ZipFile(archive_path, "w") as archive:
info = zipfile.ZipInfo("libllama.so")
info.create_system = 3
info.external_attr = 0o120777 << 16
archive.writestr(info, "libllama.so.0")
with pytest.raises(PrebuiltFallback, match = "zip archive contained a symlink entry"):
extract_archive(archive_path, tmp_path / "extract")
# The extract_archive guard tests (safe symlink chain / hardlink, absolute or
# escaping or unresolved symlink targets, zip symlink entries) moved verbatim
# to tests/studio/install/test_prebuilt_core.py: extract_archive is the shared
# prebuilt_core implementation, re-exported by this installer.
def test_remove_agent_instruction_files_does_not_follow_links(tmp_path: Path):

1357
tests/studio/install/test_install_whisper_prebuilt_logic.py Normal file
View file

File diff suppressed because it is too large Load diff

918
tests/studio/install/test_prebuilt_core.py Normal file
View file

@ -0,0 +1,918 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared descriptor-parameterized tests for studio/prebuilt_core.py.
Runs the component-agnostic core against BOTH shipped descriptors -- the real
whisper descriptor exported by install_whisper_prebuilt and a llama-flavored
descriptor built here the way a hypothetical third ggml-family component would
plug in (descriptor only, no installer module). Covers os/arch selection,
checksum fail-closed behavior, extraction guards, the resolver payload, and the
ops monkeypatch seam, so a new component gets this coverage for free.
The llama installer's shipped release-plan machinery is intentionally NOT
routed through the generic flow (its characterization suites pin it); the
llama descriptor here exercises the canonical dialect a future migration
would use, including the "no fallback backend -> report no prebuilt" policy.
"""
import importlib.util
import io
import json
import sys
import tarfile
import zipfile
from pathlib import Path
from types import SimpleNamespace
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
STUDIO_DIR = PACKAGE_ROOT / "studio"
if str(STUDIO_DIR) not in sys.path:
sys.path.insert(0, str(STUDIO_DIR))
SPEC = importlib.util.spec_from_file_location(
"studio_prebuilt_core", STUDIO_DIR / "prebuilt_core.py"
)
assert SPEC is not None and SPEC.loader is not None
core = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = core
SPEC.loader.exec_module(core)
WSPEC = importlib.util.spec_from_file_location(
"studio_install_whisper_prebuilt_for_core", STUDIO_DIR / "install_whisper_prebuilt.py"
)
assert WSPEC is not None and WSPEC.loader is not None
iwp = importlib.util.module_from_spec(WSPEC)
sys.modules[WSPEC.name] = iwp
WSPEC.loader.exec_module(iwp)
LLAMA_DESCRIPTOR = core.ComponentDescriptor(
component = "llama.cpp",
log_prefix = "llama-prebuilt",
published_repo = "unslothai/llama.cpp",
manifest_asset_name = "llama-prebuilt-manifest.json",
sha256_asset_name = "llama-prebuilt-sha256.json",
metadata_filename = "UNSLOTH_LLAMA_PREBUILT_INFO.json",
user_agent = "unsloth-studio-llama-prebuilt",
# A GPU-selection miss reports "no prebuilt" so the caller can fall back to
# a source build instead of silently degrading to CPU.
fallback_backend = None,
server_binary_name = lambda host: "llama-server",
runtime_bin_dir = lambda install_dir, host: install_dir / "build" / "bin",
)
class Component:
"""One descriptor under test plus the mutable namespace behind its ops."""
def __init__(self, descriptor):
self.descriptor = descriptor
self.namespace = core.component_namespace(descriptor)
self.namespace["log"] = lambda message: None # keep test output quiet
self.ops = core.ModuleOps(self.namespace)
@property
def falls_back_to_cpu(self):
return self.descriptor.fallback_backend == "cpu"
@pytest.fixture(params = ["whisper", "llama"])
def component(request):
if request.param == "whisper":
return Component(iwp.DESCRIPTOR)
return Component(LLAMA_DESCRIPTOR)
def make_host(
component,
*,
os_token = "linux",
arch_token = "x64",
is_windows = False,
is_macos = False,
is_apple_silicon = False,
has_usable_nvidia = False,
has_rocm = False,
rocm_gfx = None,
macos_version = None,
):
if component.descriptor is iwp.DESCRIPTOR:
return iwp.HostInfo(
system = {"linux": "Linux", "macos": "Darwin", "windows": "Windows"}[os_token],
machine = "x86_64" if arch_token == "x64" else "arm64",
whisper_os = os_token,
whisper_arch = arch_token,
archive_ext = ".zip" if is_windows else ".tar.gz",
is_windows = is_windows,
is_macos = is_macos,
is_apple_silicon = is_apple_silicon,
has_usable_nvidia = has_usable_nvidia,
has_rocm = has_rocm,
rocm_gfx = rocm_gfx,
macos_version = macos_version,
)
# Descriptor-only component: the core default host_platform_tokens hook
# reads .os_token/.arch_token off a plain host object.
return SimpleNamespace(
os_token = os_token,
arch_token = arch_token,
is_windows = is_windows,
is_macos = is_macos,
is_apple_silicon = is_apple_silicon,
has_usable_nvidia = has_usable_nvidia,
has_rocm = has_rocm,
rocm_gfx = rocm_gfx,
macos_version = macos_version,
)
def artifact(
os_ = "linux",
arch = "x64",
backend = "cpu",
asset = "bundle.tar.gz",
**extra,
):
payload = {"os": os_, "arch": arch, "backend": backend, "asset": asset}
payload.update(extra)
return payload
def manifest_for(component, artifacts, **extra):
payload = {
"schema_version": 1,
"component": extra.pop("component_name", component.descriptor.component),
"upstream_tag": "v1.0.0",
"source_commit": "a" * 40,
"artifacts": artifacts,
}
payload.update(extra)
return payload
# ── Manifest parsing ──
def test_parse_manifest_normalizes(component):
manifest = component.ops.parse_manifest(
manifest_for(component, [artifact(), "not-a-dict", {"os": "linux"}]), label = "m"
)
assert manifest["component"] == component.descriptor.component
assert manifest["upstream_tag"] == "v1.0.0"
# Non-dict entries and entries without an asset name are dropped.
assert [a["asset"] for a in manifest["artifacts"]] == ["bundle.tar.gz"]
def test_parse_manifest_rejects_wrong_component(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_manifest(
manifest_for(component, [artifact()], component_name = "other.cpp"), label = "m"
)
def test_parse_manifest_rejects_unknown_schema(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_manifest(
manifest_for(component, [artifact()], schema_version = 99), label = "m"
)
def test_parse_manifest_rejects_non_object(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_manifest(["nope"], label = "m")
# An object without an 'artifacts' list is rejected too.
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_manifest(
{"schema_version": 1, "component": component.descriptor.component}, label = "m"
)
# ── Selection matrix ──
def test_select_cpu_first_match(component):
manifest = component.ops.parse_manifest(
manifest_for(
component,
[
artifact(backend = "cpu", asset = "first-cpu.tar.gz"),
artifact(backend = "cpu", asset = "second-cpu.tar.gz"),
],
),
label = "m",
)
host = make_host(component)
chosen = component.ops.select_artifact(manifest, host, "cpu")
assert chosen["asset"] == "first-cpu.tar.gz"
def test_select_respects_os_arch(component):
manifest = component.ops.parse_manifest(
manifest_for(component, [artifact(os_ = "windows", backend = "cpu")]), label = "m"
)
host = make_host(component)
assert component.ops.select_artifact(manifest, host, "cpu") is None
def test_fallback_policy_differs_per_descriptor(component):
# No asset for the requested backend: whisper degrades to the CPU asset of
# the same release, the llama-flavored descriptor reports no prebuilt
# (source-build fallback).
manifest = component.ops.parse_manifest(
manifest_for(component, [artifact(backend = "cpu", asset = "cpu.tar.gz")]), label = "m"
)
host = make_host(component, has_usable_nvidia = True)
assert component.ops.select_artifact(manifest, host, "cuda") is None
if component.falls_back_to_cpu:
chosen, backend, used_fallback = component.ops.select_artifact_with_fallback(
manifest, host, "cuda"
)
assert (chosen["asset"], backend, used_fallback) == ("cpu.tar.gz", "cpu", True)
else:
with pytest.raises(core.PrebuiltFallback):
component.ops.select_artifact_with_fallback(manifest, host, "cuda")
def test_macos_min_os_gate(component):
manifest = component.ops.parse_manifest(
manifest_for(
component,
[
artifact(
os_ = "macos",
arch = "arm64",
backend = "metal",
asset = "metal-new.tar.gz",
min_os = "macos-15.0",
)
],
),
label = "m",
)
old_host = make_host(
component,
os_token = "macos",
arch_token = "arm64",
is_macos = True,
is_apple_silicon = True,
macos_version = (14, 7),
)
new_host = make_host(
component,
os_token = "macos",
arch_token = "arm64",
is_macos = True,
is_apple_silicon = True,
macos_version = (15, 1),
)
assert component.ops.select_artifact(manifest, old_host, "metal") is None
chosen = component.ops.select_artifact(manifest, new_host, "metal")
assert chosen["asset"] == "metal-new.tar.gz"
def _metal_artifact(asset, min_os):
return artifact(
os_ = "macos",
arch = "arm64",
backend = "metal",
asset = asset,
min_os = min_os,
)
def _arm_mac_host(component, macos_version):
return make_host(
component,
os_token = "macos",
arch_token = "arm64",
is_macos = True,
is_apple_silicon = True,
macos_version = macos_version,
)
def test_macos_min_os_filters_to_compatible_bundle(component):
# host 14.0 can't load the 15.0 bundle; the 13.0 bundle is picked instead.
manifest = component.ops.parse_manifest(
manifest_for(
component,
[
_metal_artifact("metal-new.tar.gz", "macos-15.0"),
_metal_artifact("metal.tar.gz", "macos-13.0"),
],
),
label = "m",
)
host = _arm_mac_host(component, (14, 0))
assert component.ops.select_artifact(manifest, host, "metal")["asset"] == "metal.tar.gz"
def test_macos_min_os_unknown_host_version_keeps_artifact(component):
# Unknown host macOS version -> defer to runtime validation, don't reject.
manifest = component.ops.parse_manifest(
manifest_for(component, [_metal_artifact("metal-new.tar.gz", "macos-15.0")]), label = "m"
)
host = _arm_mac_host(component, None)
assert component.ops.select_artifact(manifest, host, "metal")["asset"] == "metal-new.tar.gz"
def test_macos_min_os_accepts_bare_version_format(component):
# A bare "14.0" (no 'macos-' prefix) must still parse, for forward-compat.
manifest = component.ops.parse_manifest(
manifest_for(component, [_metal_artifact("metal.tar.gz", "14.0")]), label = "m"
)
host = _arm_mac_host(component, (13, 0))
assert component.ops.select_artifact(manifest, host, "metal") is None # 13.0 < 14.0
def test_macos_min_os_ok_helper_handles_prefix_and_bare(component):
host14 = _arm_mac_host(component, (14, 0))
# The live manifest format is 'macos-<ver>'; the prefix must be stripped.
assert component.ops.macos_min_os_ok(host14, "macos-14.0") is True
assert component.ops.macos_min_os_ok(host14, "macos-15.0") is False
assert component.ops.macos_min_os_ok(host14, "13.3") is True # bare also parses
assert component.ops.macos_min_os_ok(host14, None) is True # unknown -> defer
host_unknown = _arm_mac_host(component, None)
assert component.ops.macos_min_os_ok(host_unknown, "macos-15.0") is True
# ── Backend resolution ──
def test_resolve_backend_auto_and_validation(component):
gpu_host = make_host(component, has_usable_nvidia = True)
assert component.ops.resolve_backend(gpu_host, "auto", cpu_fallback = False) == "cuda"
assert component.ops.resolve_backend(gpu_host, "auto", cpu_fallback = True) == "cpu"
# --cpu-fallback wins over an explicit backend too.
assert component.ops.resolve_backend(gpu_host, "cuda", cpu_fallback = True) == "cpu"
# An explicit supported backend passes through untouched.
assert component.ops.resolve_backend(gpu_host, "vulkan", cpu_fallback = False) == "vulkan"
mac_host = make_host(
component, os_token = "macos", arch_token = "arm64", is_macos = True, is_apple_silicon = True
)
assert component.ops.resolve_backend(mac_host, None, cpu_fallback = False) == "metal"
# Intel mac has no Metal bundle in the P0 matrix -> cpu.
intel_mac = make_host(component, os_token = "macos", arch_token = "x64", is_macos = True)
assert component.ops.resolve_backend(intel_mac, "auto", cpu_fallback = False) == "cpu"
rocm_host = make_host(component, has_rocm = True, rocm_gfx = "gfx1100")
assert component.ops.resolve_backend(rocm_host, "auto", cpu_fallback = False) == "rocm"
bare_host = make_host(component)
assert component.ops.resolve_backend(bare_host, None, cpu_fallback = False) == "cpu"
with pytest.raises(core.PrebuiltFallback):
component.ops.resolve_backend(gpu_host, "tpu", cpu_fallback = False)
# ── Checksum index: fail closed ──
def _index_for(
component,
tag = "v1",
artifacts = None,
):
return {
"schema_version": 1,
"component": component.descriptor.component,
"release_tag": tag,
"artifacts": artifacts
if artifacts is not None
else {"bundle.tar.gz": {"sha256": "0" * 64}},
}
def test_parse_release_checksums_valid(component):
out = component.ops.parse_release_checksums("r", "v1", _index_for(component))
assert out == {"bundle.tar.gz": "0" * 64}
@pytest.mark.parametrize(
"mutation",
[
{"component": "other.cpp"},
{"schema_version": 99},
{"release_tag": "v2"},
{"artifacts": {"bundle.tar.gz": {"sha256": "nope"}}},
{"artifacts": "not-a-map"},
],
)
def test_parse_release_checksums_fails_closed(component, mutation):
payload = _index_for(component)
payload.update(mutation)
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_release_checksums("r", "v1", payload)
def test_parse_release_checksums_rejects_non_object(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.parse_release_checksums("r", "v1", ["not", "a", "dict"])
def test_expected_sha256_covered_asset_plain_lookup(component):
assert component.ops.expected_sha256_for({"a.tar.gz": "0" * 64}, "a.tar.gz") == "0" * 64
def test_expected_sha256_missing_asset_fails_closed(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.expected_sha256_for({"a.tar.gz": "0" * 64}, "b.tar.gz")
def test_expected_sha256_manifest_disagreement_fails_closed(component):
with pytest.raises(core.PrebuiltFallback):
component.ops.expected_sha256_for(
{"a.tar.gz": "0" * 64}, "a.tar.gz", manifest_sha256 = "1" * 64
)
assert (
component.ops.expected_sha256_for(
{"a.tar.gz": "0" * 64}, "a.tar.gz", manifest_sha256 = "0" * 64
)
== "0" * 64
)
# ── Extraction guards ──
def test_extract_archive_rejects_traversal(tmp_path):
archive = tmp_path / "evil.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
data = b"x"
info = tarfile.TarInfo("../escape.txt")
info.size = len(data)
tar.addfile(info, io.BytesIO(data))
with pytest.raises(core.PrebuiltFallback):
core.extract_archive(archive, tmp_path / "out")
def test_extract_archive_rejects_absolute_member(tmp_path):
archive = tmp_path / "abs.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
data = b"x"
info = tarfile.TarInfo("/abs.txt")
info.size = len(data)
tar.addfile(info, io.BytesIO(data))
with pytest.raises(core.PrebuiltFallback):
core.extract_archive(archive, tmp_path / "out")
def test_extract_archive_rejects_zip_symlink(tmp_path):
archive = tmp_path / "link.zip"
with zipfile.ZipFile(archive, "w") as zf:
info = zipfile.ZipInfo("link")
info.create_system = 3
info.external_attr = 0o120777 << 16
zf.writestr(info, "target")
with pytest.raises(core.PrebuiltFallback, match = "zip archive contained a symlink entry"):
core.extract_archive(archive, tmp_path / "out")
def test_extract_archive_rejects_absolute_zip_member(tmp_path):
archive = tmp_path / "abs.zip"
with zipfile.ZipFile(archive, "w") as zf:
zf.writestr("/etc/passwd", b"pwn")
with pytest.raises(core.PrebuiltFallback):
core.extract_archive(archive, tmp_path / "out")
def test_extract_archive_allows_safe_tar_symlink_chain(tmp_path):
archive = tmp_path / "bundle.tar.gz"
payload = b"shared-object"
with tarfile.open(archive, "w:gz") as tar:
versioned = tarfile.TarInfo("libllama.so.0.0.1")
versioned.size = len(payload)
tar.addfile(versioned, io.BytesIO(payload))
soname = tarfile.TarInfo("libllama.so.0")
soname.type = tarfile.SYMTYPE
soname.linkname = "libllama.so.0.0.1"
tar.addfile(soname)
linker_name = tarfile.TarInfo("libllama.so")
linker_name.type = tarfile.SYMTYPE
linker_name.linkname = "libllama.so.0"
tar.addfile(linker_name)
destination = tmp_path / "extract"
core.extract_archive(archive, destination)
assert (destination / "libllama.so.0.0.1").read_bytes() == payload
assert (destination / "libllama.so.0").is_symlink()
assert (destination / "libllama.so").is_symlink()
assert (destination / "libllama.so").resolve().read_bytes() == payload
def test_extract_archive_allows_safe_tar_hardlink(tmp_path):
archive = tmp_path / "bundle.tar.gz"
payload = b"quantize"
with tarfile.open(archive, "w:gz") as tar:
target = tarfile.TarInfo("llama-quantize")
target.size = len(payload)
tar.addfile(target, io.BytesIO(payload))
hardlink = tarfile.TarInfo("llama-quantize-copy")
hardlink.type = tarfile.LNKTYPE
hardlink.linkname = "llama-quantize"
tar.addfile(hardlink)
destination = tmp_path / "extract"
core.extract_archive(archive, destination)
assert (destination / "llama-quantize-copy").read_bytes() == payload
assert not (destination / "llama-quantize-copy").is_symlink()
def test_extract_archive_rejects_absolute_tar_symlink_target(tmp_path):
archive = tmp_path / "bundle.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "/tmp/libllama.so.0"
tar.addfile(entry)
with pytest.raises(core.PrebuiltFallback, match = "archive link used an absolute target"):
core.extract_archive(archive, tmp_path / "extract")
def test_extract_archive_rejects_escaping_tar_symlink_target(tmp_path):
archive = tmp_path / "bundle.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "../outside/libllama.so.0"
tar.addfile(entry)
with pytest.raises(core.PrebuiltFallback, match = "archive link escaped destination"):
core.extract_archive(archive, tmp_path / "extract")
def test_extract_archive_rejects_unresolved_tar_symlink_target(tmp_path):
archive = tmp_path / "bundle.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
entry = tarfile.TarInfo("libllama.so")
entry.type = tarfile.SYMTYPE
entry.linkname = "libllama.so.0"
tar.addfile(entry)
with pytest.raises(core.PrebuiltFallback, match = "unresolved link entries"):
core.extract_archive(archive, tmp_path / "extract")
def test_extract_archive_rejects_unknown_format(tmp_path):
archive = tmp_path / "blob.xz"
archive.write_bytes(b"data")
with pytest.raises(core.PrebuiltFallback):
core.extract_archive(archive, tmp_path / "out")
def test_restore_tar_exec_bits(tmp_path):
payload = tmp_path / "server"
payload.write_bytes(b"#!/bin/sh\n")
payload.chmod(0o755)
archive = tmp_path / "bundle.tar.gz"
with tarfile.open(archive, "w:gz") as tar:
tar.add(payload, arcname = "bundle/server")
out = tmp_path / "out"
core.extract_archive(archive, out)
extracted = out / "bundle" / "server"
assert extracted.is_file()
core.restore_tar_exec_bits(archive, out)
assert extracted.stat().st_mode & 0o111
# ── Resolver payload ──
def _fake_release(component, artifacts):
ns = component.namespace
manifest = component.ops.parse_manifest(manifest_for(component, artifacts), label = "m")
bundle = core.ReleaseBundle(
repo = component.descriptor.published_repo,
release_tag = "v1",
manifest = manifest,
asset_urls = {},
)
checksums = {str(a["asset"]): "0" * 64 for a in artifacts}
ns["fetch_release_for_install"] = lambda repo, *, published_release_tag: (bundle, checksums)
return bundle
def test_resolve_prebuilt_payload_keys(component):
_fake_release(component, [artifact(backend = "cpu", asset = "cpu.tar.gz")])
host = make_host(component)
payload = component.ops.resolve_prebuilt(
host,
published_repo = component.descriptor.published_repo,
published_release_tag = None,
backend = "cpu",
cpu_fallback = True,
)
assert payload == {
"prebuilt_available": True,
"repo": component.descriptor.published_repo,
"release_tag": "v1",
"upstream_tag": "v1.0.0",
"backend": "cpu",
"requested_backend": "cpu",
"cpu_fallback": False,
"asset": "cpu.tar.gz",
"os": "linux",
"arch": "x64",
"runtime_line": None,
}
def test_resolve_prebuilt_unavailable_payload(component):
ns = component.namespace
def boom(repo, *, published_release_tag):
raise core.PrebuiltFallback("no release")
ns["fetch_release_for_install"] = boom
host = make_host(component)
payload = component.ops.resolve_prebuilt(
host,
published_repo = component.descriptor.published_repo,
published_release_tag = None,
backend = "cpu",
cpu_fallback = True,
)
assert payload == {"prebuilt_available": False, "repo": component.descriptor.published_repo}
def test_emit_resolver_output_formats(capsys):
payload = {"prebuilt_available": True, "asset": "a.tar.gz"}
core.emit_resolver_output(payload, output_format = "json")
assert json.loads(capsys.readouterr().out) == payload
core.emit_resolver_output(payload, output_format = "plain")
assert capsys.readouterr().out.strip() == "a.tar.gz"
core.emit_resolver_output({"prebuilt_available": False}, output_format = "plain")
assert json.loads(capsys.readouterr().out) == {"prebuilt_available": False}
# ── Marker / fingerprint ──
def test_install_fingerprint_is_stable_and_sensitive(component):
kwargs = dict(
published_repo = component.descriptor.published_repo,
release_tag = "v1",
upstream_tag = "v1.0.0",
source_commit = "a" * 40,
asset = "cpu.tar.gz",
asset_sha256 = "0" * 64,
backend = "cpu",
runtime_line = None,
coverage = {},
)
first = core.compute_install_fingerprint(**kwargs)
assert first == core.compute_install_fingerprint(**kwargs)
changed = core.compute_install_fingerprint(**{**kwargs, "asset_sha256": "1" * 64})
assert changed != first
def test_write_and_match_marker(component, tmp_path):
host = make_host(component)
install_dir = tmp_path / "install"
manifest = component.ops.parse_manifest(
manifest_for(component, [artifact(backend = "cpu", asset = "cpu.tar.gz")]), label = "m"
)
selection = component.ops.selection_from_artifact(
published_repo = component.descriptor.published_repo,
release_tag = "v1",
manifest = manifest,
artifact = manifest["artifacts"][0],
backend = "cpu",
asset_sha256 = "0" * 64,
)
# No marker on disk yet -> not a match.
assert not component.ops.existing_install_matches(install_dir, host, selection)
bin_dir = component.ops.runtime_bin_dir(install_dir, host)
bin_dir.mkdir(parents = True)
component.ops.write_prebuilt_metadata(install_dir, selection)
# Marker alone is not enough: the server binary must exist too.
assert not component.ops.existing_install_matches(install_dir, host, selection)
(bin_dir / component.ops.server_binary_name(host)).write_bytes(b"bin")
marker = json.loads((install_dir / component.descriptor.metadata_filename).read_text())
assert marker["component"] == component.descriptor.component
assert marker["install_fingerprint"] == selection.fingerprint()
assert component.ops.existing_install_matches(install_dir, host, selection)
# A different selection (new sha) must force a reinstall.
other = core.InstallSelection(
**{
**selection.__dict__,
"asset_sha256": "1" * 64,
}
)
assert not component.ops.existing_install_matches(install_dir, host, other)
def test_slim_selection_fields_are_additive(component, tmp_path):
"""The slim pairing identity rides InstallSelection additively: it never
enters the fingerprint, and only a slim selection writes marker fields."""
manifest = component.ops.parse_manifest(
manifest_for(component, [artifact(backend = "cpu", asset = "cpu.tar.gz")]), label = "m"
)
selection = component.ops.selection_from_artifact(
published_repo = component.descriptor.published_repo,
release_tag = "v1",
manifest = manifest,
artifact = manifest["artifacts"][0],
backend = "cpu",
asset_sha256 = "0" * 64,
)
import dataclasses
slim = dataclasses.replace(
selection,
install_kind = "slim",
paired_llama_tag = "b10069-mix-fb3d4ca",
linked_from = "/llama/build/bin",
linked_libraries = ("libggml.so.0", "libggml-base.so.0"),
)
assert slim.fingerprint() == selection.fingerprint() # no change to the computation
fat_dir, slim_dir = tmp_path / "fat", tmp_path / "slim"
fat_dir.mkdir(), slim_dir.mkdir()
component.ops.write_prebuilt_metadata(fat_dir, selection)
fat_marker = json.loads((fat_dir / component.descriptor.metadata_filename).read_text())
for key in ("install_kind", "paired_llama_tag", "linked_from", "linked_libraries"):
assert key not in fat_marker
component.ops.write_prebuilt_metadata(slim_dir, slim)
slim_marker = json.loads((slim_dir / component.descriptor.metadata_filename).read_text())
assert slim_marker["install_kind"] == "slim"
assert slim_marker["paired_llama_tag"] == "b10069-mix-fb3d4ca"
assert slim_marker["linked_from"] == "/llama/build/bin"
assert slim_marker["linked_libraries"] == ["libggml.so.0", "libggml-base.so.0"]
assert set(slim_marker) == set(fat_marker) | {
"install_kind",
"paired_llama_tag",
"linked_from",
"linked_libraries",
}
def test_core_slim_hooks_default_inert(component, tmp_path):
# A component without its own hooks stages nothing extra and adds no
# resolver fields (llama's probe output must stay byte-identical).
assert component.ops.resolver_payload_extra({"install_kind": "slim"}) == {}
host = make_host(component)
selection = object()
assert component.ops.prepare_runtime_payload(tmp_path, host, selection) is None
def test_busy_activation_restores_previous_install(monkeypatch, tmp_path):
install_dir = tmp_path / "whisper.cpp"
staged_root = tmp_path / "staged"
install_dir.mkdir()
staged_root.mkdir()
(install_dir / "version").write_text("old")
(staged_root / "version").write_text("new")
real_replace = core.os.replace
def locked_activation(source, destination):
if Path(source) == staged_root:
raise PermissionError(13, "Permission denied")
return real_replace(source, destination)
monkeypatch.setattr(core.os, "replace", locked_activation)
with pytest.raises(core.BusyInstallConflict):
core.swap_into_place(staged_root, install_dir)
assert (install_dir / "version").read_text() == "old"
assert staged_root.is_dir()
assert not list(tmp_path.glob(".whisper.cpp.old-*"))
# ── Host/GPU token helpers (component-independent core functions) ──
# Value tables moved verbatim from the llama characterization suite; these are
# pure functions with no descriptor sensitivity, so they run unparameterized.
@pytest.mark.parametrize(
"value,expected",
[
("8.6", "86"),
("07.05", "75"),
("75", "75"),
(86, "86"),
("", None),
(" ", None),
("x.y", None),
("8.6.0", None),
("9.0", "90"),
],
)
def test_normalize_compute_cap(value, expected):
assert core.normalize_compute_cap(value) == expected
@pytest.mark.parametrize(
"values,expected",
[
(["8.6", "86", "8.6"], ["86"]), # deduplication
(["9.0", "7.5", "8.6"], ["75", "86", "90"]), # numeric sort
(["8.6", "bad", "", "7.5"], ["75", "86"]), # drops invalid
([], []),
],
)
def test_normalize_compute_caps(values, expected):
assert core.normalize_compute_caps(values) == expected
@pytest.mark.parametrize(
"value,expected",
[
(None, None),
("", []),
("-1", []),
("0", ["0"]),
("0,1,2", ["0", "1", "2"]),
(" 0 , 1 ", ["0", "1"]),
],
)
def test_parse_cuda_visible_devices(value, expected):
assert core.parse_cuda_visible_devices(value) == expected
@pytest.mark.parametrize(
"visible,expected",
[
(["0", "1", "2"], True),
(["GPU-abc123"], True),
(None, False),
([], False),
(["0", "MIG-device"], False),
],
)
def test_supports_explicit_visible_device_matching(visible, expected):
assert core.supports_explicit_visible_device_matching(visible) is expected
_GPU_ROWS = [
("0", "GPU-aaa", "8.6"),
("1", "GPU-bbb", "7.5"),
("2", "GPU-ccc", "8.9"),
]
@pytest.mark.parametrize(
"visible,expected_indices",
[
(None, [0, 1, 2]), # no filter returns all
([], []),
(["0", "2"], [0, 2]), # filter by index
(["gpu-bbb"], [1]), # UUID match is case insensitive
(["0", "0"], [0]), # same device requested twice is deduplicated
(["99"], []), # unknown token matches nothing
],
)
def test_select_visible_gpu_rows(visible, expected_indices):
expected = [_GPU_ROWS[i] for i in expected_indices]
assert core.select_visible_gpu_rows(_GPU_ROWS, visible) == expected
@pytest.mark.parametrize(
"driver,expected",
[
(None, []),
((11, 8), []),
((12, 4), ["cuda12"]),
((13, 0), ["cuda13", "cuda12"]),
((14, 0), ["cuda14", "cuda13", "cuda12"]), # future major derives lines
],
)
def test_compatible_linux_runtime_lines(driver, expected):
host = SimpleNamespace(driver_cuda_version = driver)
assert core.compatible_linux_runtime_lines(host) == expected
@pytest.mark.parametrize(
"value,expected",
[
("12.6", "cuda12"),
("13.0", "cuda13"),
("11.8", None),
(None, None),
("", None),
],
)
def test_runtime_line_from_cuda_version(value, expected):
assert core.runtime_line_from_cuda_version(value) == expected
def _caps_host(caps):
return SimpleNamespace(compute_caps = list(caps))
def test_host_is_blackwell_includes_datacenter_parts():
assert core.host_is_blackwell(_caps_host(["10.0"])) is True # B200 sm_100
assert core.host_is_blackwell(_caps_host(["10.3"])) is True # B300 sm_103
assert core.host_is_blackwell(_caps_host(["12.0"])) is True # RTX 50 sm_120
assert core.host_is_blackwell(_caps_host(["12.1"])) is True # DGX Spark sm_121
assert core.host_is_blackwell(_caps_host(["9.0"])) is False # Hopper
assert core.host_is_blackwell(_caps_host(["8.0"])) is False # Ampere
assert core.host_is_blackwell(_caps_host(["9.0", "10.0"])) is True # highest cap wins
def test_blackwell_min_toolkit_is_sm_aware():
# Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it.
f = core.blackwell_min_toolkit_for_host
assert f(_caps_host(["10.0"])) == (12, 8) # B200
assert f(_caps_host(["12.0"])) == (12, 8) # RTX 50
assert f(_caps_host(["10.3"])) == (12, 9) # B300
assert f(_caps_host(["12.1"])) == (12, 9) # DGX Spark
assert f(_caps_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins
# ── The ops seam ──
def test_module_ops_prefers_module_globals_over_core_defaults(component):
ns = dict(component.namespace)
calls = []
def fake_download_file(url, destination):
calls.append(url)
destination.write_bytes(b"data")
ns["download_file"] = fake_download_file
ops = core.ModuleOps(ns)
assert ops.download_file is fake_download_file
# Core defaults still resolve (and come back bound) for everything else.
assert callable(ops.fetch_json)
with pytest.raises(AttributeError):
_ = ops.does_not_exist_anywhere

28
tests/studio/install/test_rocm_support.py
View file

@ -232,6 +232,10 @@ UPSTREAM_ASSETS = {
class TestResolveUpstreamAssetChoice:
"""Verify that the asset selection logic picks the right binary for each platform."""
# The plain cpu-linux / windows-cpu / macos-arm64 routing cases live in
# test_selection_logic.py::TestResolveUpstreamAssetChoice (exact-name pins);
# this class keeps the ROCm/NVIDIA-precedence dialect only.
@patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
def test_nvidia_linux_gets_cpu_asset(self, mock_assets):
"""NVIDIA host should NOT hit the ROCm path -- gets CPU asset (CUDA handled elsewhere)."""
@ -249,30 +253,6 @@ class TestResolveUpstreamAssetChoice:
assert choice.install_kind == "linux-rocm"
assert "rocm" in choice.name
@patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
def test_cpu_linux_gets_cpu_asset(self, mock_assets):
"""CPU-only Linux host should get CPU asset."""
host = cpu_host()
choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
assert choice.install_kind == "linux-cpu"
assert "ubuntu-x64" in choice.name
@patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
def test_macos_arm64_gets_macos_asset(self, mock_assets):
"""macOS arm64 host should get macOS asset."""
host = macos_host()
choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
assert choice.install_kind == "macos-arm64"
assert "macos-arm64" in choice.name
@patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
def test_windows_cpu_gets_cpu_asset(self, mock_assets):
"""Windows CPU-only host should get Windows CPU asset."""
host = windows_host()
choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
assert choice.install_kind == "windows-cpu"
assert "win-cpu" in choice.name
@patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
def test_windows_rocm_gets_hip_asset(self, mock_assets):
"""Windows ROCm host should get Windows HIP asset."""

203
tests/studio/install/test_selection_logic.py
View file

@ -31,17 +31,8 @@ PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback
LinuxCudaSelection = INSTALL_LLAMA_PREBUILT.LinuxCudaSelection
UPSTREAM_REPO = INSTALL_LLAMA_PREBUILT.UPSTREAM_REPO
normalize_compute_cap = INSTALL_LLAMA_PREBUILT.normalize_compute_cap
normalize_compute_caps = INSTALL_LLAMA_PREBUILT.normalize_compute_caps
parse_cuda_visible_devices = INSTALL_LLAMA_PREBUILT.parse_cuda_visible_devices
supports_explicit_visible_device_matching = (
INSTALL_LLAMA_PREBUILT.supports_explicit_visible_device_matching
)
select_visible_gpu_rows = INSTALL_LLAMA_PREBUILT.select_visible_gpu_rows
compatible_linux_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_linux_runtime_lines
pick_windows_cuda_runtime = INSTALL_LLAMA_PREBUILT.pick_windows_cuda_runtime
compatible_windows_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines
runtime_line_from_cuda_version = INSTALL_LLAMA_PREBUILT.runtime_line_from_cuda_version
apply_approved_hashes = INSTALL_LLAMA_PREBUILT.apply_approved_hashes
linux_cuda_choice_from_release = INSTALL_LLAMA_PREBUILT.linux_cuda_choice_from_release
windows_cuda_attempts = INSTALL_LLAMA_PREBUILT.windows_cuda_attempts
@ -452,165 +443,25 @@ class TestStudioLocalhostIpv6Warning:
# ===========================================================================
# A. normalize_compute_cap
# A-F, H. Component-independent GPU/token helpers: the behavior tables live in
# tests/studio/install/test_prebuilt_core.py (they are pure prebuilt_core
# functions). This pin proves the installer still re-exports them from core, so
# the master tables keep covering the names this module and its callers use.
# ===========================================================================
class TestNormalizeComputeCap:
def test_dotted_86(self):
assert normalize_compute_cap("8.6") == "86"
def test_dotted_leading_zero(self):
assert normalize_compute_cap("07.05") == "75"
def test_already_normalized(self):
assert normalize_compute_cap("75") == "75"
def test_int_input(self):
assert normalize_compute_cap(86) == "86"
def test_empty_string(self):
assert normalize_compute_cap("") is None
def test_whitespace(self):
assert normalize_compute_cap(" ") is None
def test_non_numeric(self):
assert normalize_compute_cap("x.y") is None
def test_triple_part(self):
assert normalize_compute_cap("8.6.0") is None
def test_zero_minor(self):
assert normalize_compute_cap("9.0") == "90"
# ===========================================================================
# B. normalize_compute_caps
# ===========================================================================
class TestNormalizeComputeCaps:
def test_deduplication(self):
assert normalize_compute_caps(["8.6", "86", "8.6"]) == ["86"]
def test_numeric_sort(self):
assert normalize_compute_caps(["9.0", "7.5", "8.6"]) == ["75", "86", "90"]
def test_drops_invalid(self):
assert normalize_compute_caps(["8.6", "bad", "", "7.5"]) == ["75", "86"]
def test_empty_input(self):
assert normalize_compute_caps([]) == []
# ===========================================================================
# C. parse_cuda_visible_devices
# ===========================================================================
class TestParseCudaVisibleDevices:
def test_none(self):
assert parse_cuda_visible_devices(None) is None
def test_empty(self):
assert parse_cuda_visible_devices("") == []
def test_minus_one(self):
assert parse_cuda_visible_devices("-1") == []
def test_single(self):
assert parse_cuda_visible_devices("0") == ["0"]
def test_multi(self):
assert parse_cuda_visible_devices("0,1,2") == ["0", "1", "2"]
def test_whitespace_stripped(self):
assert parse_cuda_visible_devices(" 0 , 1 ") == ["0", "1"]
# ===========================================================================
# D. supports_explicit_visible_device_matching
# ===========================================================================
class TestSupportsExplicitVisibleDeviceMatching:
def test_all_digits(self):
assert supports_explicit_visible_device_matching(["0", "1", "2"]) is True
def test_gpu_prefix(self):
assert supports_explicit_visible_device_matching(["GPU-abc123"]) is True
def test_none(self):
assert supports_explicit_visible_device_matching(None) is False
def test_empty(self):
assert supports_explicit_visible_device_matching([]) is False
def test_mixed_invalid(self):
assert supports_explicit_visible_device_matching(["0", "MIG-device"]) is False
# ===========================================================================
# E. select_visible_gpu_rows
# ===========================================================================
class TestSelectVisibleGpuRows:
ROWS = [
("0", "GPU-aaa", "8.6"),
("1", "GPU-bbb", "7.5"),
("2", "GPU-ccc", "8.9"),
]
def test_none_returns_all(self):
assert select_visible_gpu_rows(self.ROWS, None) == list(self.ROWS)
def test_empty_returns_empty(self):
assert select_visible_gpu_rows(self.ROWS, []) == []
def test_filter_by_index(self):
result = select_visible_gpu_rows(self.ROWS, ["0", "2"])
assert result == [("0", "GPU-aaa", "8.6"), ("2", "GPU-ccc", "8.9")]
def test_filter_by_uuid_case_insensitive(self):
result = select_visible_gpu_rows(self.ROWS, ["gpu-bbb"])
assert result == [("1", "GPU-bbb", "7.5")]
def test_dedup_same_device(self):
result = select_visible_gpu_rows(self.ROWS, ["0", "0"])
assert result == [("0", "GPU-aaa", "8.6")]
def test_missing_token(self):
result = select_visible_gpu_rows(self.ROWS, ["99"])
assert result == []
# ===========================================================================
# F. compatible_linux_runtime_lines
# ===========================================================================
class TestCompatibleLinuxRuntimeLines:
def test_no_driver(self):
host = make_host(driver_cuda_version = None)
assert compatible_linux_runtime_lines(host) == []
def test_driver_11_8(self):
host = make_host(driver_cuda_version = (11, 8))
assert compatible_linux_runtime_lines(host) == []
def test_driver_12_4(self):
host = make_host(driver_cuda_version = (12, 4))
assert compatible_linux_runtime_lines(host) == ["cuda12"]
def test_driver_13_0(self):
host = make_host(driver_cuda_version = (13, 0))
assert compatible_linux_runtime_lines(host) == ["cuda13", "cuda12"]
def test_future_major_derives_lines(self):
host = make_host(driver_cuda_version = (14, 0))
assert compatible_linux_runtime_lines(host) == ["cuda14", "cuda13", "cuda12"]
def test_core_helper_aliases_bound_to_prebuilt_core():
import prebuilt_core as _core
for name in (
"normalize_compute_cap",
"normalize_compute_caps",
"parse_cuda_visible_devices",
"supports_explicit_visible_device_matching",
"select_visible_gpu_rows",
"compatible_linux_runtime_lines",
"runtime_line_from_cuda_version",
):
assert getattr(INSTALL_LLAMA_PREBUILT, name) is getattr(_core, name), name
# ===========================================================================
@ -669,28 +520,6 @@ class TestCompatibleWindowsRuntimeLines:
assert compatible_windows_runtime_lines(host) == ["cuda14", "cuda13", "cuda12"]
# ===========================================================================
# H. runtime_line_from_cuda_version
# ===========================================================================
class TestRuntimeLineFromCudaVersion:
def test_cuda_12(self):
assert runtime_line_from_cuda_version("12.6") == "cuda12"
def test_cuda_13(self):
assert runtime_line_from_cuda_version("13.0") == "cuda13"
def test_cuda_11(self):
assert runtime_line_from_cuda_version("11.8") is None
def test_none(self):
assert runtime_line_from_cuda_version(None) is None
def test_empty(self):
assert runtime_line_from_cuda_version("") is None
# ===========================================================================
# I. apply_approved_hashes
# ===========================================================================

29
tests/studio/install/test_setup_whisper_status.py Normal file
View file

@ -0,0 +1,29 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""Contract checks for actionable whisper prebuilt setup outcomes."""
from pathlib import Path
ROOT = Path(__file__).resolve().parents[3]
def test_shell_setup_distinguishes_release_skew_from_install_failure():
script = (ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8")
assert 'if [ "$_WHISPER_STATUS" -eq 2 ]' in script
assert "installed llama.cpp ${_WHISPER_INSTALLED_TAG:-unknown}" in script
assert "whisper requires ${_WHISPER_REQUIRED_TAG:-unknown}" in script
assert "prebuilt install failed" in script
assert "retry setup or inspect verbose output" in script
assert "curated whisper.cpp dictation is unavailable" in script
assert "browser and Transformers dictation remain available" in script
def test_powershell_setup_distinguishes_release_skew_from_install_failure():
script = (ROOT / "studio" / "setup.ps1").read_text(encoding = "utf-8")
assert "elseif ($whisperExit -eq 2)" in script
assert "installed llama.cpp $installedWhisperLlamaTag" in script
assert "whisper requires $requiredWhisperLlamaTag" in script
assert "prebuilt install failed" in script
assert "retry setup or inspect verbose output" in script
assert "curated whisper.cpp dictation is unavailable" in script
assert "browser and Transformers dictation remain available" in script

42
tests/studio/playwright_extra_ui.py
View file

@ -482,6 +482,14 @@ with sync_playwright() as p:
step("Settings dialog: cycle through tabs")
page.goto(f"{BASE}/chat")
composer.wait_for(state = "visible", timeout = 60_000)
dictate = page.get_by_role("button", name = "Dictate").first
if dictate.count() == 0:
fail("Chat Dictate button not found")
elif dictate.get_attribute("type") != "button":
fail("Chat Dictate control must use type=button, not submit the composer")
else:
info("OK Chat Dictate control is type=button")
page.keyboard.press("Control+,")
page.wait_for_timeout(800)
settings = page.get_by_role("dialog").first
@ -501,6 +509,7 @@ with sync_playwright() as p:
"Appearance",
"Chat",
"Developer",
"Voice",
"About",
)
seen_tabs = []
@ -528,6 +537,39 @@ with sync_playwright() as p:
soft_fail(f"Settings tab '{tab_name}' body suspiciously short: {body_text}")
except Exception as exc:
soft_fail(f"Settings tab '{tab_name}' click failed: {exc!r}")
step("Voice model picker: real mouse-wheel scrolling")
voice_tab = page.get_by_role(
"button", name = re.compile(r"^\s*Voice(?:\s+New)?\s*$", re.I)
).first
if voice_tab.count() == 0:
fail("Voice settings tab not found")
else:
voice_tab.click()
page.get_by_label("Dictation engine").click()
page.get_by_role("option", name = "Local transcription").click()
page.get_by_label("Speech recognition model").click()
page.get_by_placeholder("Search model").fill("whisper")
results = page.get_by_test_id("stt-model-results")
try:
page.wait_for_function(
"""() => {
const node = document.querySelector('[data-testid="stt-model-results"]');
return !!node && node.scrollHeight > node.clientHeight;
}""",
timeout = 30_000,
)
results.hover()
page.mouse.wheel(0, 700)
page.wait_for_function(
"""() => {
const node = document.querySelector('[data-testid="stt-model-results"]');
return !!node && node.scrollTop > 0;
}""",
timeout = 5_000,
)
info("OK Voice model picker mouse wheel changed scrollTop")
except Exception as exc:
fail(f"Voice model picker did not wheel-scroll: {exc!r}")
shoot("10-settings-tabs-visited")
page.keyboard.press("Escape")
page.wait_for_timeout(300)

8
tests/test_studio_install_workspace_guard.py
View file

@ -254,6 +254,14 @@ def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard():
), "Assert-StudioOwnedOrAbsent must precede the install_llama_prebuilt.py call"
def test_setup_ps1_adopts_existing_whisper_prebuilt_marker():
text = SETUP_PS1.read_text(encoding = "utf-8")
helper_start = text.index("function Test-StudioOwnedAdoptable")
helper_end = text.index("function Assert-StudioOwnedOrAbsent", helper_start)
helper = text[helper_start:helper_end]
assert "UNSLOTH_WHISPER_PREBUILT_INFO.json" in helper
def test_env_mode_passes_when_venv_marker_present(tmp_path):
"""install.sh env-mode guard must accept the in-VENV .unsloth-studio-owned marker as a sentinel."""
studio_home = tmp_path / "ws"