Add native GGUF intake to Studio (#5246)
* feat(studio): add Tauri native GGUF intake * feat(studio): polish native GGUF intake * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): load backend helpers during local setup * fix(studio): acquire native load lease before unload * Studio: harden native path lease verification and Tauri intake - Wrap path.resolve(strict=True) and Path.stat() in NativePathLeaseError so a deleted or unmounted GGUF returns 400 instead of leaking the full filesystem path through the generic load_model/validate_model handler. - Re-apply _reject_network_or_device_path to the resolved canonical path for defense in depth after symlink resolution. - Replace try/except ValueError pattern in the device-path guard with Path.is_relative_to; the previous shape silently swallowed NativePathLeaseError (which subclasses ValueError) so /dev,/proc,/sys were never actually rejected. - Broaden the lease redaction regex and dict-key check (Python and Rust diagnostics) to cover both native_path_lease and nativePathLease so the camelCase form emitted by Tauri/frontend payloads is also redacted. - Hoist the redact_native_paths import to module top in loggers/handlers; the recursive filter no longer pays a per-record import lookup. - Persist activeNativePathToken in the chat runtime store so the rollback branch can mint a fresh lease and reload the previous native GGUF when a new load fails after unload; clear it in clearCheckpoint and overwrite it on each successful load. - use-native-drop: read options through a ref so the Tauri onDragDropEvent listener is registered once and stays attached across option changes; reject ambiguous multi-file drops up front instead of silently registering only the first GGUF. - pick_native_model: use an async pick_file with a tokio oneshot channel instead of blocking_pick_file so the Tokio worker is not held for the duration of the OS dialog. - registerNativeModelPath: drop the duplicate sourceKind argument; the Rust command parameter is source_kind. - install_python_stack: insert the script directory (studio/) on sys.path; the previous insert pointed at studio/backend/ which does not satisfy `from backend.utils.wheel_utils import ...`. * install_python_stack: keep _BACKEND_DIR on sys.path Restore the studio/backend insertion. Although the immediately following `from backend.utils.wheel_utils import (...)` is satisfied by studio/ already being on sys.path[0] when invoked as `python studio/install_python_stack.py`, wheel_utils itself runs `from utils.native_path_leases import ...`, which requires studio/backend/ to be importable. Without the backend insertion, the existing tests/python/test_install_python_stack.py collection fails with ModuleNotFoundError: No module named 'utils'. * Studio: tighten native path lease lifecycle and Tauri intake IPC - register_native_model_path now hardcodes NativePathSourceKind::Drop on the Rust side and the frontend stops sending source_kind. The previous JS payload (source_kind only) never reached the Rust deserializer because Tauri's default ArgumentCase::Camel maps the Rust parameter source_kind to the JS key sourceKind, so drag/drop registration silently failed. Hardcoding the source kind also keeps audit metadata trustworthy on this command. - Add native_path_secret_removed_for_child_start context manager and wrap multiprocessing.Process.start() at the inference, export, training, and data-recipe job spawn sites. The previous wrapper-only scrub left UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET visible to spawn-platform import-time worker code. The wrapper run_without_native_path_secret stays as defense-in-depth inside the child. - Stop passing exc_info=True from the native-grant load/validate error logs in routes/inference.py. The structlog filter_sensitive_data processor runs before the renderer, so ConsoleRenderer formatted tracebacks bypassed redaction; the redacted str(e) preserves the message text. - Replace the os.path.normcase string equality on the resolved canonical path with Path.samefile (with a normcase fallback) so Windows leases that differ only in extended-length \\?\ prefix or short-name spelling are accepted. - Wrap consumeNativePathToken in its own try/catch in the chat runtime rollback. If the previous native-model token has aged out of TOKEN_TTL we now surface a clear modelsError instead of silently swallowing the rollback inside the outer catch. - Reject non-ASCII lease strings in _split_lease and convert UnicodeEncodeError / binascii.Error / ValueError raised by _b64decode into NativePathLeaseError so verify_native_path_lease never escapes raw exceptions to the route handler. - Tighten dropStateForPaths to mark multi-file payloads invalid so the overlay matches the post-fix drop handler that rejects the same payload. - Replace the one-shot fetch in useNativePathLeasesSupported with a delayed-retry loop so the picker/drop becomes available once the backend is up rather than staying disabled for the rest of the session after a transient failure. - Drop the unused setActiveNativePathToken setter; the value is set via setState directly in use-chat-model-runtime. - Add a toast on auto-load failure in use-native-drop so a collapsed model selector does not hide the error. - Burn the lease nonce before _validate_current_stat so a stat-failed lease is single-use even if a later state change happens to match the original size/mtime. * Studio: cache lease secret, harden native path stat checks, polish intake UX - Cache the decoded UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET on first verify and validate that it is base64-decodable and at least 32 bytes. Subsequent _decode_secret calls return from the cache and never touch os.environ, so concurrent /api/inference/load and /api/health requests no longer race with native_path_secret_removed_for_child_start scrubbing the env. native_path_leases_supported now wraps _decode_secret so the health flag matches what verify_native_path_lease actually accepts. - Replace path.is_file()/is_dir() + path.stat() with os.lstat() in _validate_current_stat and explicitly reject S_ISLNK; size and mtime checks now refer to the link itself, closing the same-size+same-mtime symlink-swap window that the prior follow-symlink stat() left open. - Add an issued_at_ms < expires_at_ms sanity check in _validate_payload to reject internally inconsistent (HMAC-protected) lease payloads. - Sort _NATIVE_PATH_REDACTIONS by length (descending) before iterating in redact_native_paths so a longer registered path is replaced before a shorter prefix path; otherwise logs containing /foo/X.gguf.bak after only /foo/X.gguf was registered would leak the .bak suffix. - classify_existing_path now re-checks the canonical path with symlink_metadata after canonicalize, so a regular file that is replaced with a symlink in the small canonicalize window is rejected at registration. - ModelSelector renders the local file picker as its own block (not in the eject ternary), so a user with an active model can still replace it via the picker rather than only via drag/drop. - useNativePathLeasesSupported caps the readiness probe at MAX_READINESS_POLLS (60 = ~5 minutes) and aborts the in-flight fetch on unmount via AbortController, so a permanently-disabled backend stops generating sustained traffic and hot-reload no longer leaks open connections. - useChooseNativeModel returns a stable useCallback closure and guards the OS dialog with a useRef so rapid double-clicks cannot open multiple dialogs and orphan Rust tokens. - Branch the multi-file drop toast: if no GGUF was present we say "Only .gguf model files can be dropped here." and otherwise "Drop a single .gguf model file." so users dropping non-GGUF attachments get an accurate explanation. * native_path_leases: lstat the signed canonical path before resolving The earlier change to lstat inside _validate_current_stat operates on grant.canonical_path, which is the post-resolve target. If the user atomically replaces the originally-signed file with a symlink to a different file of identical size and mtime, path.resolve(strict=True) follows the symlink, samefile returns True (both ends share the new inode), and the lstat in _validate_current_stat sees the regular target file rather than the symlink, so the swap goes undetected. Add an os.lstat on the signed canonical path before path.resolve(strict=True), and reject S_ISLNK there. The lstat in _validate_current_stat stays as defense-in-depth for swaps that occur strictly between resolve and stat. * Studio: scrub native lease secret before mp.Queue spawn and tighten lease lifecycle - Move _CTX.Queue / _CTX.Event / _CTX.Process construction inside native_path_secret_removed_for_child_start at the inference, export, training and data-recipe spawn sites. The first Queue creation lazily spawns Python's multiprocessing.resource_tracker child, so when it ran outside the scrub context the tracker process inherited the lease secret. Reproduced via the proc filesystem environ entry; the wrapped order keeps the tracker clean. - native_path_secret_removed_for_child_start now refcounts entries: the env var is popped on the first entry and restored only when the last context exits. Concurrent training/inference/export starts no longer serialize on the env lock across the entire proc.start yield, while still guaranteeing the env stays empty for the duration of every overlapping spawn. - run_without_native_path_secret now also nulls the module-level cached lease secret. With the existing spawn-only multiprocessing context the cache is irrelevant in practice, but a future fork caller would otherwise inherit the in-memory secret even though the env var was scrubbed. - filter_sensitive_data now applies the native lease key check on the top-level event_dict, not only on nested dicts, so a logger call that includes a lease value as a top-level keyword field actually redacts it (the bare value does not match the prefix-anchored regex). - chat-page loadNativeModelIntent now passes intent.id to clearModelIntent so a second drag-drop during an in-flight first auto-load is not wiped from the chip area when the first resolves. - Bump useNativePathLeasesSupported's MAX_READINESS_POLLS from 60 to 720 so first-run installs that compile llama.cpp from source or download large CUDA wheels (well past 5 minutes) don't permanently disable the native picker. * native_path_leases: serialize first-decode against scrub context _decode_secret used a separate _SECRET_INIT_LOCK from the env scrub's _NATIVE_PATH_ENV_LOCK, so the very first decode (before the cache is populated) could race a concurrent native_path_secret_removed_for_child_start and read os.environ during the env-empty window, raising "Native path grants require the managed desktop backend." Subsequent calls hit the cache and were already safe. Acquire _NATIVE_PATH_ENV_LOCK around the env read inside _SECRET_INIT_LOCK and fall back to _SCRUB_SAVED_SECRET when the scrub has temporarily popped the env var. Lock ordering (init then env) is consistent with no other caller, so no deadlock. * Studio: surface native model load errors and harden native path label cache - Native model load and validate now bubble up the actual exception (with paths redacted) and apply the same friendly-error rewrite the non-native path uses, so users see "CUDA OOM", "trust_remote_code required", etc. instead of a generic "Failed to load native model: <label>". - run_without_native_path_secret now also nulls _SCRUB_SAVED_SECRET so a forked grandchild that imports native_path_leases cannot recover the secret via the scrub-aware fallback in _decode_secret. - _NATIVE_PATH_LABELS now has its own 10000-entry cap independent of the 100-entry redaction list, so display_label_for_native_path no longer falls back to returning the raw canonical path after 101 native paths in one session. Redaction list keeps the 100-entry cap for log-scan performance. - _validate_payload now also rejects null bytes in display_label, which is echoed back in HTTP responses and log lines. * Studio: harden native path lease validation and chained native rollback - child_env_without_native_path_secret now copies os.environ under _NATIVE_PATH_ENV_LOCK so a concurrent scrub-context env pop cannot raise RuntimeError: dictionary changed size during iteration in a background hardware scan or other env reader. - _validate_payload and grant construction route every signed numeric field (version, issued_at_ms, expires_at_ms, size_bytes, modified_ms) through new _required_int / _optional_int helpers that wrap raw int() ValueError into NativePathLeaseError. The single upstream catcher produces 400 instead of 500 for malformed signed payloads. - verify_native_path_lease now runs _validate_current_stat before _consume_nonce, so a transient stat error on the canonical path no longer permanently burns the nonce. Concurrent verifies still serialize through _consume_nonce, so single-use is preserved. - Chained native model rollback now restores activeNativePathToken in the chat runtime store after a successful rollback loadModel. Without this, a second consecutive failed switch could not re-roll-back because the store token had been overwritten by the failed attempt. - validate_model now applies the same not_supported_hints friendly rewrite to native model errors that load_model already does, so a native .gguf that fails validation with an upstream "is not supported" message gets the same actionable wording as the non-native branch. * Studio: harden native path log redaction, status disclosure, and chip lifecycle - structlog processor chain now runs format_exc_info before filter_sensitive_data so traceback strings are produced (and then redacted) rather than passed through as untouched (type, value, tb) tuples that the JSON or console renderer formats after the redaction filter has already finished. - native_path_secret_removed_for_child_start clears _CACHED_LEASE_SECRET in addition to popping the env var, so a fork during the scrub window cannot inherit the cached bytes via the parent's heap. Parent verify calls during the window keep working through the existing scrub-aware fallback in _decode_secret. - load_model's except ValueError handler now redacts native paths and uses the native model log label when native_grant_backed is true. Previously a ValueError raised after lease verification (e.g. from ModelConfig.from_identifier or downstream GGUF parsing) returned the raw exception string in the HTTP response body. - llama_cpp_backend now records the native display label at GGUF load time, and /api/inference/status prefers it over the redaction store. After a Python backend restart the redaction store is empty; the attribute keeps the friendly label, and an absolute model_identifier with no other label source falls back to the basename so the canonical path no longer appears in active_model. - reveal_path_token uses native "reveal and select" commands on macOS (open -R) and Windows (explorer /select,) so the file is highlighted in the file manager. Linux keeps the existing parent-directory open. - Native model rollback that fails because the previous token cannot be consumed now throws a rollback-specific Error, and the outer empty catch was replaced with one that re-throws the rollback error. The rollback-specific message now reaches the user instead of being overwritten by the original load error message. - NativeModelChip tracks the Rust token's expiresAtMs on a single setTimeout, disables the Load button at expiry, and relabels it "Select again" with an explanatory tooltip so users do not click into a guaranteed-failure path after the 15-minute TTL elapses. * Studio: tighten native artifact policy, mmproj sibling check, and intake UX - is_open_safe_artifact no longer grants Open for directories. Reveal already handles directory navigation, so the change closes the attack surface where a macOS .app artifact could be launched via open_path_token + open::that_detached. - Display labels are sanitized in classify_existing_path. Control characters in filenames (newlines, tabs, NUL et al.) are replaced with spaces and the label is trimmed and capped, so a file named with embedded newlines cannot inject forged log lines or scramble the UI status panel. - validate_entry_path skips the size_bytes/modified_ms equality check when the operation is Reveal or Open. Cloud-sync agents (Dropbox, iCloud Drive, OneDrive) routinely rewrite extended-attribute metadata which bumps mtime, and the user expects Reveal/Open to remain available for files in synced folders. - llama_cpp_backend gains a _native_grant_backed flag at GGUF load success. /api/inference/status only applies the absolute-path basename fallback when that flag is true, so a non-native absolute local GGUF still reports its canonical model_identifier and unload by identifier keeps working. - Native vision GGUFs now run through _validate_native_mmproj_companion before llama-server starts: the companion mmproj must be a regular file, not a symlink, and must live in the same resolved directory as the granted GGUF. This stops a hostile sibling or symlinked mmproj from being loaded under a single-file lease. - Chained native rollback restructured: the rollback loadModel + state + refresh runs inside its own try/catch that swallows so the outer throw error surfaces the ORIGINAL load failure. The native-token consume-failure case still throws the rollback-specific message early, before the inner block runs, so its actionable guidance is preserved. - Loading-model state and the duplicate-load guard in the chat runtime hook now compare both the model id and the native path token. Two drops or picks with the same basename in different folders no longer silently dedup; the second token is honored. - chat-page loadNativeModelIntent awaits selectModel before clearing the pending intent. If selectModel returns early via dedup or throws, the chip and its token stay so the user can retry instead of losing the selection. - NativeModelChip's Reveal button is disabled when the lease has expired (Rust would reject it anyway), and the Load button label reads "Expired" instead of "Select again" so the disabled element no longer promises an action it cannot perform. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
parent
35ab5da93c
commit
e35cbfb454
46 changed files with 2596 additions and 97 deletions
|
|
@ -167,13 +167,20 @@ class JobManager:
|
|||
|
||||
run_payload = dict(run)
|
||||
run_payload["_job_id"] = job_id
|
||||
mp_q = _CTX.Queue()
|
||||
proc = _CTX.Process(
|
||||
target = run_job_process,
|
||||
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
|
||||
daemon = True,
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
proc.start()
|
||||
|
||||
with native_path_secret_removed_for_child_start():
|
||||
mp_q = _CTX.Queue()
|
||||
proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_job_process,),
|
||||
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
|
||||
daemon = True,
|
||||
)
|
||||
proc.start()
|
||||
|
||||
self._mp_q = mp_q
|
||||
self._proc = proc
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ _OXC_TOOL_DIR = Path(__file__).resolve().parent / "oxc-validator"
|
|||
_OXC_RUNNER_PATH = _OXC_TOOL_DIR / "validate.mjs"
|
||||
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -248,7 +249,7 @@ def _run_oxc_batch(
|
|||
}
|
||||
try:
|
||||
tmp_dir = ensure_dir(oxc_validator_tmp_root())
|
||||
env = dict(os.environ)
|
||||
env = child_env_without_native_path_secret()
|
||||
tmp_dir_str = str(tmp_dir)
|
||||
env["TMPDIR"] = tmp_dir_str
|
||||
env["TMP"] = tmp_dir_str
|
||||
|
|
|
|||
|
|
@ -163,21 +163,28 @@ class ExportOrchestrator:
|
|||
|
||||
def _spawn_subprocess(self, config: dict) -> None:
|
||||
"""Spawn a new export subprocess."""
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
from .worker import run_export_process
|
||||
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
|
||||
self._proc = _CTX.Process(
|
||||
target = run_export_process,
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
self._proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_export_process,),
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
|
||||
|
||||
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing import Optional, Tuple
|
|||
import numpy as np
|
||||
import torch
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -105,6 +106,7 @@ class AudioCodecManager:
|
|||
spark_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
@ -143,6 +145,7 @@ class AudioCodecManager:
|
|||
outetts_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
# Remove files that pull in heavy / incompatible dependencies
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from urllib.parse import urlparse
|
|||
|
||||
import httpx
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -592,6 +593,7 @@ class LlamaCppBackend:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0:
|
||||
|
|
@ -1776,7 +1778,7 @@ class LlamaCppBackend:
|
|||
import os
|
||||
import sys
|
||||
|
||||
env = os.environ.copy()
|
||||
env = child_env_without_native_path_secret()
|
||||
binary_dir = str(Path(binary).parent)
|
||||
|
||||
if sys.platform == "win32":
|
||||
|
|
@ -2166,6 +2168,7 @@ class LlamaCppBackend:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -166,23 +166,30 @@ class InferenceOrchestrator:
|
|||
|
||||
def _spawn_subprocess(self, config: dict) -> None:
|
||||
"""Spawn a new inference subprocess."""
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
from .worker import run_inference_process
|
||||
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
self._cancel_event = _CTX.Event()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
self._cancel_event = _CTX.Event()
|
||||
|
||||
self._proc = _CTX.Process(
|
||||
target = run_inference_process,
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"cancel_event": self._cancel_event,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
self._proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_inference_process,),
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"cancel_event": self._cancel_event,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
|
||||
|
||||
def _cancel_generation(self) -> None:
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ from utils.paths import (
|
|||
)
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -1771,6 +1772,7 @@ class UnslothTrainer:
|
|||
spark_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
@ -2005,6 +2007,7 @@ class UnslothTrainer:
|
|||
outetts_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
for fpath in [
|
||||
|
|
|
|||
|
|
@ -29,6 +29,10 @@ from typing import Optional, Tuple, Any
|
|||
|
||||
import matplotlib.pyplot as plt
|
||||
from utils.hardware import prepare_gpu_selection
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -213,20 +217,22 @@ class TrainingBackend:
|
|||
|
||||
from .worker import run_training_process
|
||||
|
||||
event_queue = _CTX.Queue()
|
||||
stop_queue = _CTX.Queue()
|
||||
|
||||
proc = _CTX.Process(
|
||||
target = run_training_process,
|
||||
kwargs = {
|
||||
"event_queue": event_queue,
|
||||
"stop_queue": stop_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
try:
|
||||
proc.start()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
event_queue = _CTX.Queue()
|
||||
stop_queue = _CTX.Queue()
|
||||
|
||||
proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_training_process,),
|
||||
kwargs = {
|
||||
"event_queue": event_queue,
|
||||
"stop_queue": stop_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
proc.start()
|
||||
except Exception:
|
||||
logger.error("Failed to start training subprocess", exc_info = True)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ from typing import Optional
|
|||
|
||||
import structlog
|
||||
|
||||
from loggers.handlers import filter_sensitive_data
|
||||
|
||||
|
||||
class LogConfig:
|
||||
"""Structured logging configuration for the application.
|
||||
|
|
@ -58,6 +60,8 @@ class LogConfig:
|
|||
structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
|
||||
structlog.processors.add_log_level, # level second
|
||||
structlog.contextvars.merge_contextvars,
|
||||
structlog.processors.format_exc_info,
|
||||
filter_sensitive_data,
|
||||
# Custom processor to flatten the extra field
|
||||
lambda logger, method_name, event_dict: {
|
||||
"timestamp": event_dict.get("timestamp"),
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ Key Components:
|
|||
- get_logger: Factory function for structured loggers
|
||||
"""
|
||||
|
||||
import re
|
||||
import time
|
||||
from typing import Callable
|
||||
|
||||
|
|
@ -22,7 +23,12 @@ import structlog
|
|||
from fastapi import Request, Response
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
_NATIVE_PATH_LEASE_RE = re.compile(
|
||||
r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
|
||||
)
|
||||
|
||||
|
||||
class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
|
|
@ -75,6 +81,12 @@ def filter_sensitive_data(logger, method_name, event_dict):
|
|||
"""Structlog processor to filter out base64 data from logs."""
|
||||
|
||||
def filter_value(value):
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = redact_native_paths(value)
|
||||
except Exception:
|
||||
pass
|
||||
value = _NATIVE_PATH_LEASE_RE.sub(r"\1<redacted native path lease>", value)
|
||||
if (
|
||||
isinstance(value, str)
|
||||
and len(value) > 100
|
||||
|
|
@ -83,12 +95,22 @@ def filter_sensitive_data(logger, method_name, event_dict):
|
|||
# Likely base64 data, truncate it
|
||||
return value[:20] + "..."
|
||||
elif isinstance(value, dict):
|
||||
return {k: filter_value(v) for k, v in value.items()}
|
||||
return {
|
||||
k: "<redacted native path lease>"
|
||||
if str(k).replace("_", "").lower() == "nativepathlease"
|
||||
else filter_value(v)
|
||||
for k, v in value.items()
|
||||
}
|
||||
elif isinstance(value, list):
|
||||
return [filter_value(item) for item in value]
|
||||
return value
|
||||
|
||||
return {k: filter_value(v) for k, v in event_dict.items()}
|
||||
return {
|
||||
k: "<redacted native path lease>"
|
||||
if str(k).replace("_", "").lower() == "nativepathlease"
|
||||
else filter_value(v)
|
||||
for k, v in event_dict.items()
|
||||
}
|
||||
|
||||
|
||||
def get_logger(name: str) -> structlog.BoundLogger:
|
||||
|
|
|
|||
|
|
@ -78,6 +78,7 @@ from utils.hardware import (
|
|||
import utils.hardware.hardware as _hw_module
|
||||
|
||||
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
||||
from utils.native_path_leases import native_path_leases_supported
|
||||
|
||||
|
||||
def get_unsloth_version() -> str:
|
||||
|
|
@ -244,6 +245,7 @@ async def health_check():
|
|||
"chat_only": _hw_module.CHAT_ONLY,
|
||||
"desktop_protocol_version": 1,
|
||||
"supports_desktop_auth": True,
|
||||
"native_path_leases_supported": native_path_leases_supported(),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ class LoadRequest(BaseModel):
|
|||
"""Request to load a model for inference"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier or local path")
|
||||
native_path_lease: Optional[str] = Field(
|
||||
None, description = "Frontend-visible signed native path grant"
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None, description = "HuggingFace token for gated models"
|
||||
)
|
||||
|
|
@ -69,6 +72,9 @@ class ValidateModelRequest(BaseModel):
|
|||
"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier or local path")
|
||||
native_path_lease: Optional[str] = Field(
|
||||
None, description = "Frontend-visible signed native path grant"
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None, description = "HuggingFace token for gated models"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -122,6 +122,13 @@ try:
|
|||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
from utils.native_path_leases import (
|
||||
NativePathLeaseError,
|
||||
display_label_for_native_path,
|
||||
is_registered_native_path_label,
|
||||
redact_native_paths,
|
||||
verify_native_path_lease,
|
||||
)
|
||||
except ImportError:
|
||||
parent_backend = backend_path.parent / "backend"
|
||||
if str(parent_backend) not in sys.path:
|
||||
|
|
@ -136,6 +143,13 @@ except ImportError:
|
|||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
from utils.native_path_leases import (
|
||||
NativePathLeaseError,
|
||||
display_label_for_native_path,
|
||||
is_registered_native_path_label,
|
||||
redact_native_paths,
|
||||
verify_native_path_lease,
|
||||
)
|
||||
|
||||
from models.inference import (
|
||||
LoadRequest,
|
||||
|
|
@ -332,6 +346,65 @@ _TOOL_XML_RE = _re.compile(
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _validate_native_mmproj_companion(
|
||||
mmproj_path: str | None, gguf_path: str | None
|
||||
) -> None:
|
||||
if not mmproj_path or not gguf_path:
|
||||
return
|
||||
import stat as _stat_module
|
||||
|
||||
mm = Path(mmproj_path)
|
||||
gguf = Path(gguf_path)
|
||||
try:
|
||||
mm_lstat = os.lstat(mm)
|
||||
except OSError as exc:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Native vision companion is no longer accessible.",
|
||||
) from exc
|
||||
if _stat_module.S_ISLNK(mm_lstat.st_mode) or not _stat_module.S_ISREG(
|
||||
mm_lstat.st_mode
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Native vision companion must be a regular file.",
|
||||
)
|
||||
try:
|
||||
if mm.resolve(strict = True).parent != gguf.resolve(strict = True).parent:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Native vision companion must live next to the selected GGUF.",
|
||||
)
|
||||
except OSError as exc:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Native vision companion is no longer accessible.",
|
||||
) from exc
|
||||
|
||||
|
||||
def _resolve_model_identifier_for_request(
|
||||
request: LoadRequest | ValidateModelRequest,
|
||||
*,
|
||||
operation: str,
|
||||
) -> tuple[str, str, bool]:
|
||||
if not request.native_path_lease:
|
||||
return request.model_path, request.model_path, False
|
||||
try:
|
||||
grant = verify_native_path_lease(
|
||||
request.native_path_lease,
|
||||
operation = operation,
|
||||
expected_kind = "model",
|
||||
expected_path_type = "file",
|
||||
allowed_suffixes = (".gguf",),
|
||||
)
|
||||
except NativePathLeaseError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc)) from exc
|
||||
display_label = (
|
||||
grant.display_label or Path(request.model_path).name or "Native model"
|
||||
)
|
||||
return str(grant.canonical_path), display_label, True
|
||||
|
||||
|
||||
# GGUF inference backend (llama-server)
|
||||
_llama_cpp_backend = LlamaCppBackend()
|
||||
|
||||
|
|
@ -355,7 +428,12 @@ async def load_model(
|
|||
|
||||
GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth.
|
||||
"""
|
||||
native_grant_backed = False
|
||||
model_log_label = request.model_path
|
||||
try:
|
||||
model_identifier, model_log_label, native_grant_backed = (
|
||||
_resolve_model_identifier_for_request(request, operation = "load-model")
|
||||
)
|
||||
# Version switching is handled automatically by the subprocess-based
|
||||
# inference backend — no need for ensure_transformers_version() here.
|
||||
|
||||
|
|
@ -369,10 +447,10 @@ async def load_model(
|
|||
and llama_backend.hf_variant
|
||||
and llama_backend.hf_variant.lower() == request.gguf_variant.lower()
|
||||
and llama_backend.model_identifier
|
||||
and llama_backend.model_identifier.lower() == request.model_path.lower()
|
||||
and llama_backend.model_identifier.lower() == model_identifier.lower()
|
||||
):
|
||||
logger.info(
|
||||
f"Model already loaded (GGUF): {request.model_path} variant={request.gguf_variant}, skipping reload"
|
||||
f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload"
|
||||
)
|
||||
inference_config = load_inference_config(llama_backend.model_identifier)
|
||||
from utils.models import is_audio_input_type
|
||||
|
|
@ -385,8 +463,12 @@ async def load_model(
|
|||
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
|
||||
return LoadResponse(
|
||||
status = "already_loaded",
|
||||
model = llama_backend.model_identifier,
|
||||
display_name = llama_backend.model_identifier,
|
||||
model = model_log_label
|
||||
if native_grant_backed
|
||||
else llama_backend.model_identifier,
|
||||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else llama_backend.model_identifier,
|
||||
is_vision = llama_backend._is_vision,
|
||||
is_lora = False,
|
||||
is_gguf = True,
|
||||
|
|
@ -412,10 +494,10 @@ async def load_model(
|
|||
else:
|
||||
if (
|
||||
backend.active_model_name
|
||||
and backend.active_model_name.lower() == request.model_path.lower()
|
||||
and backend.active_model_name.lower() == model_identifier.lower()
|
||||
):
|
||||
logger.info(
|
||||
f"Model already loaded (Unsloth): {request.model_path}, skipping reload"
|
||||
f"Model already loaded (Unsloth): {model_log_label}, skipping reload"
|
||||
)
|
||||
inference_config = load_inference_config(backend.active_model_name)
|
||||
_model_info = backend.models.get(backend.active_model_name, {})
|
||||
|
|
@ -444,8 +526,12 @@ async def load_model(
|
|||
pass
|
||||
return LoadResponse(
|
||||
status = "already_loaded",
|
||||
model = backend.active_model_name,
|
||||
display_name = backend.active_model_name,
|
||||
model = model_log_label
|
||||
if native_grant_backed
|
||||
else backend.active_model_name,
|
||||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else backend.active_model_name,
|
||||
is_vision = _model_info.get("is_vision", False),
|
||||
is_lora = _model_info.get("is_lora", False),
|
||||
is_gguf = False,
|
||||
|
|
@ -467,7 +553,7 @@ async def load_model(
|
|||
# Create config using clean factory method
|
||||
# is_lora is auto-detected from adapter_config.json on disk/HF
|
||||
config = ModelConfig.from_identifier(
|
||||
model_id = request.model_path,
|
||||
model_id = model_identifier,
|
||||
hf_token = request.hf_token,
|
||||
gguf_variant = request.gguf_variant,
|
||||
)
|
||||
|
|
@ -475,7 +561,7 @@ async def load_model(
|
|||
if not config:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Invalid model identifier: {request.model_path}",
|
||||
detail = f"Invalid model identifier: {model_log_label}",
|
||||
)
|
||||
|
||||
# Normalize gpu_ids: empty list means auto-selection, same as None
|
||||
|
|
@ -522,6 +608,10 @@ async def load_model(
|
|||
)
|
||||
else:
|
||||
# Local mode: llama-server loads via -m <path>
|
||||
if native_grant_backed and config.gguf_mmproj_file:
|
||||
_validate_native_mmproj_companion(
|
||||
config.gguf_mmproj_file, config.gguf_file
|
||||
)
|
||||
success = await asyncio.to_thread(
|
||||
llama_backend.load_model,
|
||||
gguf_path = config.gguf_file,
|
||||
|
|
@ -538,10 +628,12 @@ async def load_model(
|
|||
if not success:
|
||||
raise HTTPException(
|
||||
status_code = 500,
|
||||
detail = f"Failed to load GGUF model: {config.display_name}",
|
||||
detail = f"Failed to load GGUF model: {model_log_label if native_grant_backed else config.display_name}",
|
||||
)
|
||||
|
||||
logger.info(f"Loaded GGUF model via llama-server: {config.identifier}")
|
||||
logger.info(
|
||||
f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}"
|
||||
)
|
||||
|
||||
# Detect TTS audio by probing the loaded model's vocabulary
|
||||
from utils.models import is_audio_input_type
|
||||
|
|
@ -550,6 +642,10 @@ async def load_model(
|
|||
_gguf_is_audio = _gguf_audio in ("snac", "bicodec", "dac")
|
||||
llama_backend._is_audio = _gguf_is_audio
|
||||
llama_backend._audio_type = _gguf_audio
|
||||
llama_backend._native_display_label = (
|
||||
model_log_label if native_grant_backed else None
|
||||
)
|
||||
llama_backend._native_grant_backed = bool(native_grant_backed)
|
||||
if _gguf_is_audio:
|
||||
logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}")
|
||||
await asyncio.to_thread(llama_backend.init_audio_codec, _gguf_audio)
|
||||
|
|
@ -558,8 +654,10 @@ async def load_model(
|
|||
|
||||
return LoadResponse(
|
||||
status = "loaded",
|
||||
model = config.identifier,
|
||||
display_name = config.display_name,
|
||||
model = model_log_label if native_grant_backed else config.identifier,
|
||||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else config.display_name,
|
||||
is_vision = config.is_vision,
|
||||
is_lora = False,
|
||||
is_gguf = True,
|
||||
|
|
@ -682,10 +780,13 @@ async def load_model(
|
|||
),
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 500, detail = f"Failed to load model: {config.display_name}"
|
||||
status_code = 500,
|
||||
detail = f"Failed to load model: {model_log_label if native_grant_backed else config.display_name}",
|
||||
)
|
||||
|
||||
logger.info(f"Loaded model: {config.identifier}")
|
||||
logger.info(
|
||||
f"Loaded model: {model_log_label if native_grant_backed else config.identifier}"
|
||||
)
|
||||
|
||||
# Load inference configuration parameters
|
||||
inference_config = load_inference_config(config.identifier)
|
||||
|
|
@ -715,8 +816,10 @@ async def load_model(
|
|||
|
||||
return LoadResponse(
|
||||
status = "loaded",
|
||||
model = config.identifier,
|
||||
display_name = config.display_name,
|
||||
model = model_log_label if native_grant_backed else config.identifier,
|
||||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else config.display_name,
|
||||
is_vision = config.is_vision,
|
||||
is_lora = config.is_lora,
|
||||
is_gguf = False,
|
||||
|
|
@ -738,11 +841,17 @@ async def load_model(
|
|||
except HTTPException:
|
||||
raise
|
||||
except ValueError as e:
|
||||
if native_grant_backed:
|
||||
redacted_msg = redact_native_paths(str(e))
|
||||
logger.warning(
|
||||
"Rejected inference selection for native model %s: %s",
|
||||
model_log_label,
|
||||
redacted_msg,
|
||||
)
|
||||
raise HTTPException(status_code = 400, detail = redacted_msg)
|
||||
logger.warning("Rejected inference GPU selection: %s", e)
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading model: {e}", exc_info = True)
|
||||
msg = str(e)
|
||||
# Surface a friendlier message for models that Unsloth cannot load
|
||||
not_supported_hints = [
|
||||
"No config file found",
|
||||
|
|
@ -750,6 +859,22 @@ async def load_model(
|
|||
"is not supported",
|
||||
"does not support",
|
||||
]
|
||||
if native_grant_backed:
|
||||
redacted_msg = redact_native_paths(str(e))
|
||||
logger.error(
|
||||
"Error loading native model %s: %s",
|
||||
model_log_label,
|
||||
redacted_msg,
|
||||
)
|
||||
msg = redacted_msg
|
||||
if any(h.lower() in msg.lower() for h in not_supported_hints):
|
||||
msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
|
||||
raise HTTPException(
|
||||
status_code = 500,
|
||||
detail = f"Failed to load native model {model_log_label}: {msg}",
|
||||
)
|
||||
logger.error(f"Error loading model: {e}", exc_info = True)
|
||||
msg = str(e)
|
||||
if any(h.lower() in msg.lower() for h in not_supported_hints):
|
||||
msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
|
||||
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
|
||||
|
|
@ -766,9 +891,14 @@ async def validate_model(
|
|||
This checks that ModelConfig.from_identifier() can resolve the given
|
||||
model_path, but it does NOT actually load model weights into GPU memory.
|
||||
"""
|
||||
native_grant_backed = False
|
||||
model_log_label = request.model_path
|
||||
try:
|
||||
model_identifier, model_log_label, native_grant_backed = (
|
||||
_resolve_model_identifier_for_request(request, operation = "validate-model")
|
||||
)
|
||||
config = ModelConfig.from_identifier(
|
||||
model_id = request.model_path,
|
||||
model_id = model_identifier,
|
||||
hf_token = request.hf_token,
|
||||
gguf_variant = request.gguf_variant,
|
||||
)
|
||||
|
|
@ -776,14 +906,16 @@ async def validate_model(
|
|||
if not config:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Invalid model identifier: {request.model_path}",
|
||||
detail = f"Invalid model identifier: {model_log_label}",
|
||||
)
|
||||
|
||||
return ValidateModelResponse(
|
||||
valid = True,
|
||||
message = "Model identifier is valid.",
|
||||
identifier = config.identifier,
|
||||
display_name = getattr(config, "display_name", config.identifier),
|
||||
identifier = model_log_label if native_grant_backed else config.identifier,
|
||||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else getattr(config, "display_name", config.identifier),
|
||||
is_gguf = getattr(config, "is_gguf", False),
|
||||
is_lora = getattr(config, "is_lora", False),
|
||||
is_vision = getattr(config, "is_vision", False),
|
||||
|
|
@ -795,6 +927,26 @@ async def validate_model(
|
|||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
not_supported_hints = [
|
||||
"No config file found",
|
||||
"not yet supported",
|
||||
"is not supported",
|
||||
"does not support",
|
||||
]
|
||||
if native_grant_backed:
|
||||
redacted_msg = redact_native_paths(str(e))
|
||||
logger.error(
|
||||
"Error validating native model %s: %s",
|
||||
model_log_label,
|
||||
redacted_msg,
|
||||
)
|
||||
msg = redacted_msg
|
||||
if any(h.lower() in msg.lower() for h in not_supported_hints):
|
||||
msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Invalid native model {model_log_label}: {msg}",
|
||||
)
|
||||
logger.error(
|
||||
f"Error validating model identifier '{request.model_path}': {e}",
|
||||
exc_info = True,
|
||||
|
|
@ -819,6 +971,9 @@ async def unload_model(
|
|||
llama_backend = get_llama_cpp_backend()
|
||||
if llama_backend.is_active and (
|
||||
llama_backend.model_identifier == request.model_path
|
||||
or is_registered_native_path_label(
|
||||
llama_backend.model_identifier, request.model_path
|
||||
)
|
||||
or not llama_backend.is_loaded
|
||||
):
|
||||
llama_backend.unload_model()
|
||||
|
|
@ -966,16 +1121,27 @@ async def get_status(
|
|||
# If a GGUF model is loaded via llama-server, report that
|
||||
if llama_backend.is_loaded:
|
||||
_model_id = llama_backend.model_identifier
|
||||
_native_grant_backed = getattr(llama_backend, "_native_grant_backed", False)
|
||||
_display_model_id = getattr(
|
||||
llama_backend, "_native_display_label", None
|
||||
) or display_label_for_native_path(_model_id)
|
||||
if (
|
||||
_native_grant_backed
|
||||
and _model_id
|
||||
and _display_model_id == _model_id
|
||||
and os.path.isabs(_model_id)
|
||||
):
|
||||
_display_model_id = os.path.basename(_model_id)
|
||||
_inference_cfg = load_inference_config(_model_id) if _model_id else None
|
||||
return InferenceStatusResponse(
|
||||
active_model = _model_id,
|
||||
active_model = _display_model_id,
|
||||
is_vision = llama_backend.is_vision,
|
||||
is_gguf = True,
|
||||
gguf_variant = llama_backend.hf_variant,
|
||||
is_audio = getattr(llama_backend, "_is_audio", False),
|
||||
audio_type = getattr(llama_backend, "_audio_type", None),
|
||||
loading = [],
|
||||
loaded = [_model_id],
|
||||
loaded = [_display_model_id] if _display_model_id else [],
|
||||
inference = _inference_cfg,
|
||||
requires_trust_remote_code = bool(
|
||||
(_inference_cfg or {}).get("trust_remote_code", False)
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import subprocess
|
|||
from typing import Any, Optional
|
||||
|
||||
from loggers import get_logger
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -28,6 +29,7 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("amd-smi query failed: %s", e)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Optional
|
|||
|
||||
from loggers import get_logger
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -65,6 +66,7 @@ def get_physical_gpu_count() -> Optional[int]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
|
|
@ -90,6 +92,7 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
@ -141,6 +144,7 @@ def get_visible_gpu_utilization(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
@ -227,6 +231,7 @@ def get_backend_visible_gpu_info(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ import threading
|
|||
import yaml
|
||||
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -583,6 +584,7 @@ def _is_vision_model_subprocess(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 60,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
|
|||
406
studio/backend/utils/native_path_leases.py
Normal file
406
studio/backend/utils/native_path_leases.py
Normal file
|
|
@ -0,0 +1,406 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Verification for Tauri native path signed grants.
|
||||
|
||||
Rust signs compact ``base64url(payload_json).base64url(hmac)`` grants. The
|
||||
frontend can see and forward the grant, but cannot change it without breaking
|
||||
the HMAC. The backend verifies the original payload segment bytes, then
|
||||
re-stats the path before any native read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import stat as _stat_module
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, Iterator, Mapping
|
||||
|
||||
LEASE_SECRET_ENV = "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET"
|
||||
_MAX_NATIVE_PATH_REDACTIONS = 100
|
||||
_MAX_NATIVE_PATH_LABELS = 10_000
|
||||
_MIN_LEASE_SECRET_BYTES = 32
|
||||
|
||||
_REPLAY_LOCK = threading.Lock()
|
||||
_USED_NONCES: dict[str, int] = {}
|
||||
_REDACTION_LOCK = threading.Lock()
|
||||
_NATIVE_PATH_REDACTIONS: list[str] = []
|
||||
_NATIVE_PATH_LABELS: dict[str, str] = {}
|
||||
_NATIVE_PATH_ENV_LOCK = threading.Lock()
|
||||
_SECRET_INIT_LOCK = threading.Lock()
|
||||
_CACHED_LEASE_SECRET: bytes | None = None
|
||||
_SCRUB_REFCOUNT = 0
|
||||
_SCRUB_SAVED_SECRET: str | None = None
|
||||
|
||||
|
||||
class NativePathLeaseError(ValueError):
|
||||
"""Raised when a native path grant is missing, invalid, or unsafe."""
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class NativePathGrant:
|
||||
operation: str
|
||||
canonical_path: Path
|
||||
path_kind: str
|
||||
path_type: str
|
||||
source_kind: str
|
||||
token_id_hash: str
|
||||
display_label: str
|
||||
expires_at_ms: int
|
||||
size_bytes: int | None
|
||||
modified_ms: int | None
|
||||
|
||||
|
||||
def native_path_leases_supported() -> bool:
|
||||
try:
|
||||
_decode_secret()
|
||||
except NativePathLeaseError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def child_env_without_native_path_secret(
|
||||
env: Mapping[str, str] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Return a child-process env with the native path lease secret removed."""
|
||||
|
||||
if env is None:
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
cleaned = dict(os.environ)
|
||||
else:
|
||||
cleaned = dict(env)
|
||||
cleaned.pop(LEASE_SECRET_ENV, None)
|
||||
return cleaned
|
||||
|
||||
|
||||
def run_without_native_path_secret(
|
||||
target: Callable[..., Any],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Run a multiprocessing child target without the native path lease secret."""
|
||||
|
||||
global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET
|
||||
os.environ.pop(LEASE_SECRET_ENV, None)
|
||||
_CACHED_LEASE_SECRET = None
|
||||
_SCRUB_SAVED_SECRET = None
|
||||
return target(*args, **kwargs)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def native_path_secret_removed_for_child_start() -> Iterator[None]:
|
||||
global _SCRUB_REFCOUNT, _SCRUB_SAVED_SECRET, _CACHED_LEASE_SECRET
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
if _SCRUB_REFCOUNT == 0:
|
||||
_SCRUB_SAVED_SECRET = os.environ.pop(LEASE_SECRET_ENV, None)
|
||||
_CACHED_LEASE_SECRET = None
|
||||
_SCRUB_REFCOUNT += 1
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
_SCRUB_REFCOUNT -= 1
|
||||
if _SCRUB_REFCOUNT == 0 and _SCRUB_SAVED_SECRET is not None:
|
||||
os.environ[LEASE_SECRET_ENV] = _SCRUB_SAVED_SECRET
|
||||
_SCRUB_SAVED_SECRET = None
|
||||
|
||||
|
||||
def verify_native_path_lease(
|
||||
lease: str | None,
|
||||
*,
|
||||
operation: str,
|
||||
expected_kind: str | None = None,
|
||||
expected_path_type: str | None = None,
|
||||
allowed_suffixes: Iterable[str] | None = None,
|
||||
) -> NativePathGrant:
|
||||
if not lease:
|
||||
raise NativePathLeaseError("Native path grant is required.")
|
||||
|
||||
secret = _decode_secret()
|
||||
payload_b64, signature_b64 = _split_lease(lease)
|
||||
expected_signature = hmac.new(
|
||||
secret,
|
||||
payload_b64.encode("ascii"),
|
||||
hashlib.sha256,
|
||||
).digest()
|
||||
supplied_signature = _b64decode(signature_b64)
|
||||
if not hmac.compare_digest(expected_signature, supplied_signature):
|
||||
raise NativePathLeaseError("Native path grant signature is invalid.")
|
||||
|
||||
payload = _decode_payload(payload_b64)
|
||||
_validate_payload(payload, operation = operation, expected_kind = expected_kind)
|
||||
|
||||
path = Path(str(payload["canonical_path"]))
|
||||
_reject_network_or_device_path(path)
|
||||
try:
|
||||
signed_lstat = os.lstat(path)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
if _stat_module.S_ISLNK(signed_lstat.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
try:
|
||||
resolved = path.resolve(strict = True)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
_reject_network_or_device_path(resolved)
|
||||
if not _same_native_path(resolved, path):
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant no longer resolves to the selected path."
|
||||
)
|
||||
|
||||
grant = NativePathGrant(
|
||||
operation = str(payload["operation"]),
|
||||
canonical_path = resolved,
|
||||
path_kind = str(payload["path_kind"]),
|
||||
path_type = str(payload["path_type"]),
|
||||
source_kind = str(payload["source_kind"]),
|
||||
token_id_hash = str(payload["token_id_hash"]),
|
||||
display_label = str(payload.get("display_label") or resolved.name),
|
||||
expires_at_ms = _required_int(payload, "expires_at_ms"),
|
||||
size_bytes = _optional_int(payload.get("size_bytes")),
|
||||
modified_ms = _optional_int(payload.get("modified_ms")),
|
||||
)
|
||||
|
||||
if expected_path_type and grant.path_type != expected_path_type:
|
||||
raise NativePathLeaseError("Native path grant has the wrong path type.")
|
||||
suffixes = tuple(s.lower() for s in (allowed_suffixes or ()))
|
||||
if suffixes and resolved.suffix.lower() not in suffixes:
|
||||
raise NativePathLeaseError("Native path grant has an unsupported file type.")
|
||||
|
||||
_validate_current_stat(grant)
|
||||
_consume_nonce(str(payload["nonce"]), grant.expires_at_ms)
|
||||
_remember_native_path_for_redaction(str(resolved), grant.display_label)
|
||||
return grant
|
||||
|
||||
|
||||
def display_label_for_native_path(value: str | None) -> str | None:
|
||||
if not value:
|
||||
return value
|
||||
with _REDACTION_LOCK:
|
||||
return _NATIVE_PATH_LABELS.get(value, value)
|
||||
|
||||
|
||||
def is_registered_native_path_label(path_value: str | None, label: str | None) -> bool:
|
||||
if not path_value or not label:
|
||||
return False
|
||||
with _REDACTION_LOCK:
|
||||
return _NATIVE_PATH_LABELS.get(path_value) == label
|
||||
|
||||
|
||||
def redact_native_paths(value: str) -> str:
|
||||
with _REDACTION_LOCK:
|
||||
paths = sorted(_NATIVE_PATH_REDACTIONS, key = len, reverse = True)
|
||||
redacted = value
|
||||
for path in paths:
|
||||
for variant in {path, path.replace("/", "\\"), path.replace("\\", "/")}:
|
||||
if variant:
|
||||
redacted = redacted.replace(variant, "<native_path>")
|
||||
return redacted
|
||||
|
||||
|
||||
def _decode_secret() -> bytes:
|
||||
global _CACHED_LEASE_SECRET
|
||||
if _CACHED_LEASE_SECRET is not None:
|
||||
return _CACHED_LEASE_SECRET
|
||||
with _SECRET_INIT_LOCK:
|
||||
if _CACHED_LEASE_SECRET is not None:
|
||||
return _CACHED_LEASE_SECRET
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
encoded = os.environ.get(LEASE_SECRET_ENV)
|
||||
if encoded is None and _SCRUB_SAVED_SECRET is not None:
|
||||
encoded = _SCRUB_SAVED_SECRET
|
||||
if not encoded:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grants require the managed desktop backend."
|
||||
)
|
||||
try:
|
||||
secret = _b64decode(encoded)
|
||||
except Exception as exc:
|
||||
raise NativePathLeaseError("Native path grant secret is invalid.") from exc
|
||||
if len(secret) < _MIN_LEASE_SECRET_BYTES:
|
||||
raise NativePathLeaseError("Native path grant secret is invalid.")
|
||||
_CACHED_LEASE_SECRET = secret
|
||||
return secret
|
||||
|
||||
|
||||
def _split_lease(lease: str) -> tuple[str, str]:
|
||||
if not isinstance(lease, str):
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.")
|
||||
try:
|
||||
lease.encode("ascii")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.") from exc
|
||||
parts = lease.split(".")
|
||||
if len(parts) != 2 or not parts[0] or not parts[1]:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.")
|
||||
return parts[0], parts[1]
|
||||
|
||||
|
||||
def _decode_payload(payload_b64: str) -> dict[str, Any]:
|
||||
try:
|
||||
payload = json.loads(_b64decode(payload_b64).decode("utf-8"))
|
||||
except Exception as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.")
|
||||
return payload
|
||||
|
||||
|
||||
def _validate_payload(
|
||||
payload: dict[str, Any], *, operation: str, expected_kind: str | None
|
||||
) -> None:
|
||||
required = (
|
||||
"version",
|
||||
"operation",
|
||||
"canonical_path",
|
||||
"path_kind",
|
||||
"path_type",
|
||||
"source_kind",
|
||||
"token_id_hash",
|
||||
"issued_at_ms",
|
||||
"expires_at_ms",
|
||||
"nonce",
|
||||
)
|
||||
missing = [key for key in required if key not in payload]
|
||||
if missing:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant payload is missing required fields."
|
||||
)
|
||||
if _required_int(payload, "version") != 1:
|
||||
raise NativePathLeaseError("Native path grant version is unsupported.")
|
||||
if payload["operation"] != operation:
|
||||
raise NativePathLeaseError("Native path grant operation is invalid.")
|
||||
if expected_kind and payload["path_kind"] != expected_kind:
|
||||
raise NativePathLeaseError("Native path grant kind is invalid.")
|
||||
now_ms = int(time.time() * 1000)
|
||||
issued_at_ms = _required_int(payload, "issued_at_ms")
|
||||
expires_at_ms = _required_int(payload, "expires_at_ms")
|
||||
if issued_at_ms >= expires_at_ms:
|
||||
raise NativePathLeaseError("Native path grant timestamps are inconsistent.")
|
||||
if expires_at_ms <= now_ms:
|
||||
raise NativePathLeaseError("Native path grant has expired.")
|
||||
if issued_at_ms > now_ms + 30_000:
|
||||
raise NativePathLeaseError("Native path grant issue time is invalid.")
|
||||
for key in ("canonical_path", "nonce", "token_id_hash", "display_label"):
|
||||
raw = payload.get(key)
|
||||
if raw is None:
|
||||
continue
|
||||
if "\x00" in str(raw):
|
||||
raise NativePathLeaseError("Native path grant contains invalid characters.")
|
||||
|
||||
|
||||
def _validate_current_stat(grant: NativePathGrant) -> None:
|
||||
try:
|
||||
st = os.lstat(grant.canonical_path)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
if _stat_module.S_ISLNK(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
if grant.path_type == "file":
|
||||
if not _stat_module.S_ISREG(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
elif grant.path_type == "directory":
|
||||
if not _stat_module.S_ISDIR(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a directory.")
|
||||
else:
|
||||
raise NativePathLeaseError("Native path grant has an unsupported path type.")
|
||||
|
||||
if grant.size_bytes is not None and st.st_size != grant.size_bytes:
|
||||
raise NativePathLeaseError("Native path changed after it was selected.")
|
||||
current_modified_ms = int(st.st_mtime_ns // 1_000_000)
|
||||
if grant.modified_ms is not None and current_modified_ms != grant.modified_ms:
|
||||
raise NativePathLeaseError("Native path changed after it was selected.")
|
||||
|
||||
|
||||
def _consume_nonce(nonce: str, expires_at_ms: int) -> None:
|
||||
now_ms = int(time.time() * 1000)
|
||||
with _REPLAY_LOCK:
|
||||
for key, expiry in list(_USED_NONCES.items()):
|
||||
if expiry <= now_ms:
|
||||
_USED_NONCES.pop(key, None)
|
||||
if nonce in _USED_NONCES:
|
||||
raise NativePathLeaseError("Native path grant was already used.")
|
||||
_USED_NONCES[nonce] = expires_at_ms
|
||||
|
||||
|
||||
def _remember_native_path_for_redaction(path: str, display_label: str) -> None:
|
||||
with _REDACTION_LOCK:
|
||||
_NATIVE_PATH_LABELS[path] = display_label
|
||||
if len(_NATIVE_PATH_LABELS) > _MAX_NATIVE_PATH_LABELS:
|
||||
excess = len(_NATIVE_PATH_LABELS) - _MAX_NATIVE_PATH_LABELS
|
||||
for stale_path in list(_NATIVE_PATH_LABELS.keys())[:excess]:
|
||||
_NATIVE_PATH_LABELS.pop(stale_path, None)
|
||||
if path in _NATIVE_PATH_REDACTIONS:
|
||||
return
|
||||
_NATIVE_PATH_REDACTIONS.append(path)
|
||||
del _NATIVE_PATH_REDACTIONS[:-_MAX_NATIVE_PATH_REDACTIONS]
|
||||
|
||||
|
||||
def _reject_network_or_device_path(path: Path) -> None:
|
||||
text = str(path)
|
||||
if os.name == "nt":
|
||||
normalized = text.replace("/", "\\").lower()
|
||||
if normalized.startswith("\\\\?\\"):
|
||||
rest = normalized[4:]
|
||||
is_local_drive = len(rest) >= 3 and rest[0].isalpha() and rest[1:3] == ":\\"
|
||||
if not is_local_drive:
|
||||
raise NativePathLeaseError(
|
||||
"Network paths are not supported for native grants."
|
||||
)
|
||||
elif normalized.startswith("\\\\"):
|
||||
raise NativePathLeaseError(
|
||||
"Network paths are not supported for native grants."
|
||||
)
|
||||
if os.name != "nt":
|
||||
for root in ("/dev", "/proc", "/sys"):
|
||||
if path.is_relative_to(root):
|
||||
raise NativePathLeaseError(
|
||||
"Device and virtual filesystem paths are not supported."
|
||||
)
|
||||
if "\x00" in text:
|
||||
raise NativePathLeaseError("Native path contains invalid characters.")
|
||||
|
||||
|
||||
def _b64decode(value: str) -> bytes:
|
||||
try:
|
||||
padding = "=" * (-len(value) % 4)
|
||||
return base64.urlsafe_b64decode((value + padding).encode("ascii"))
|
||||
except (UnicodeEncodeError, binascii.Error, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.") from exc
|
||||
|
||||
|
||||
def _same_native_path(resolved: Path, signed: Path) -> bool:
|
||||
try:
|
||||
return resolved.samefile(signed)
|
||||
except OSError:
|
||||
return os.path.normcase(str(resolved)) == os.path.normcase(str(signed))
|
||||
|
||||
|
||||
def _optional_int(value: Any) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
|
||||
|
||||
def _required_int(payload: dict[str, Any], key: str) -> int:
|
||||
raw = payload.get(key)
|
||||
if raw is None:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant payload is missing required fields."
|
||||
)
|
||||
try:
|
||||
return int(raw)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
|
|
@ -36,6 +36,7 @@ import subprocess
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -504,6 +505,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool:
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0:
|
||||
|
|
@ -526,6 +528,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool:
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode != 0:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ import urllib.error
|
|||
import urllib.request
|
||||
from typing import Callable
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
FLASH_ATTN_RELEASE_BASE_URL = (
|
||||
|
|
@ -59,6 +61,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
|
|||
stderr = subprocess.PIPE,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
return None
|
||||
|
|
@ -142,6 +145,7 @@ def install_wheel(
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
attempts.append(("uv", result))
|
||||
if result.returncode == 0:
|
||||
|
|
@ -153,6 +157,7 @@ def install_wheel(
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
attempts.append(("pip", result))
|
||||
return attempts
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import {
|
|||
} from "@/components/tauri/window-titlebar";
|
||||
import { Toaster } from "@/components/ui/sonner";
|
||||
import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth";
|
||||
import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain";
|
||||
import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend";
|
||||
import { useTauriUpdate } from "@/hooks/use-tauri-update";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
|
|
@ -245,6 +246,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
|
|||
const content = showApp ? (
|
||||
<>
|
||||
<TauriUpdateLayer isExternalServer={isExternalServer} />
|
||||
<NativeIntentDrain />
|
||||
{children}
|
||||
</>
|
||||
) : (
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import { usePlatformStore } from "@/config/env";
|
|||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
ArrowDown01Icon,
|
||||
FolderSearchIcon,
|
||||
Logout01Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
|
|
@ -41,6 +42,7 @@ interface ModelSelectorProps {
|
|||
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
onFoldersChange?: () => void;
|
||||
onPickLocalModel?: () => void | Promise<void>;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
variant?: "outline" | "ghost" | "muted";
|
||||
|
|
@ -117,6 +119,7 @@ function ModelSelectorContent({
|
|||
onSelect,
|
||||
onEject,
|
||||
onFoldersChange,
|
||||
onPickLocalModel,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
className,
|
||||
|
|
@ -128,6 +131,7 @@ function ModelSelectorContent({
|
|||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
onFoldersChange?: () => void;
|
||||
onPickLocalModel?: () => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
className?: string;
|
||||
|
|
@ -170,6 +174,19 @@ function ModelSelectorContent({
|
|||
</Tabs>
|
||||
)}
|
||||
|
||||
{onPickLocalModel ? (
|
||||
<div className="mt-2 border-t border-border/70 pt-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onPickLocalModel}
|
||||
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-muted/60"
|
||||
title="Pick a model file from disk"
|
||||
>
|
||||
<HugeiconsIcon icon={FolderSearchIcon} className="size-3.5" />
|
||||
Pick a model file from disk
|
||||
</button>
|
||||
</div>
|
||||
) : null}
|
||||
{hasSelection && onEject ? (
|
||||
<div className="mt-2 border-t border-border/70 pt-2">
|
||||
<button
|
||||
|
|
@ -196,6 +213,7 @@ export function ModelSelector({
|
|||
onValueChange,
|
||||
onEject,
|
||||
onFoldersChange,
|
||||
onPickLocalModel,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
variant = "outline",
|
||||
|
|
@ -275,6 +293,11 @@ export function ModelSelector({
|
|||
setOpen(false);
|
||||
}
|
||||
|
||||
function handlePickLocalModel() {
|
||||
setOpen(false);
|
||||
void onPickLocalModel?.();
|
||||
}
|
||||
|
||||
return (
|
||||
<Popover open={open} onOpenChange={setOpen}>
|
||||
<ModelSelectorTrigger
|
||||
|
|
@ -292,6 +315,7 @@ export function ModelSelector({
|
|||
onSelect={handleSelect}
|
||||
onEject={onEject ? handleEject : undefined}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
className={contentClassName}
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ import {
|
|||
import { sentAudioNames } from "@/features/chat/api/chat-adapter";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
|
||||
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
|
|
@ -298,27 +299,41 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
|
|||
[disabled],
|
||||
);
|
||||
|
||||
const composerContent = (
|
||||
<>
|
||||
<ComposerAttachments />
|
||||
<PendingAudioChip />
|
||||
<ToolStatusDisplay />
|
||||
<ComposerPrimitive.Input
|
||||
placeholder="Send a message..."
|
||||
className="aui-composer-input mb-1 min-h-12 w-full resize-none overflow-y-auto bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm font-[450] outline-none placeholder:text-muted-foreground focus-visible:ring-0"
|
||||
minRows={1}
|
||||
maxRows={6}
|
||||
autoFocus={!disabled}
|
||||
disabled={disabled}
|
||||
aria-label="Message input"
|
||||
/>
|
||||
<ComposerAction disabled={disabled} />
|
||||
</>
|
||||
);
|
||||
|
||||
return (
|
||||
<ComposerPrimitive.Root
|
||||
className="aui-composer-root relative flex w-full flex-col"
|
||||
aria-disabled={disabled}
|
||||
onSubmit={handleSubmit}
|
||||
>
|
||||
<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:bg-accent/50">
|
||||
<ComposerAttachments />
|
||||
<PendingAudioChip />
|
||||
<ToolStatusDisplay />
|
||||
<ComposerPrimitive.Input
|
||||
placeholder="Send a message..."
|
||||
className="aui-composer-input mb-1 min-h-12 w-full resize-none overflow-y-auto bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm font-[450] outline-none placeholder:text-muted-foreground focus-visible:ring-0"
|
||||
minRows={1}
|
||||
maxRows={6}
|
||||
autoFocus={!disabled}
|
||||
disabled={disabled}
|
||||
aria-label="Message input"
|
||||
/>
|
||||
<ComposerAction disabled={disabled} />
|
||||
</ComposerPrimitive.AttachmentDropzone>
|
||||
{isTauri ? (
|
||||
// Phase 1 native model drops own Tauri local-path drops. Restore browser
|
||||
// attachment drops in Tauri when Phase 1d adds attachment-token bridging.
|
||||
<div className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow">
|
||||
{composerContent}
|
||||
</div>
|
||||
) : (
|
||||
<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:bg-accent/50">
|
||||
{composerContent}
|
||||
</ComposerPrimitive.AttachmentDropzone>
|
||||
)}
|
||||
</ComposerPrimitive.Root>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -68,7 +68,11 @@ export async function loadModel(
|
|||
const response = await authFetch("/api/inference/load", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
body: JSON.stringify({
|
||||
...payload,
|
||||
native_path_lease: payload.nativePathLease ?? null,
|
||||
nativePathLease: undefined,
|
||||
}),
|
||||
});
|
||||
return parseJsonOrThrow<LoadModelResponse>(response);
|
||||
}
|
||||
|
|
@ -81,6 +85,7 @@ export async function validateModel(
|
|||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model_path: payload.model_path,
|
||||
native_path_lease: payload.nativePathLease ?? null,
|
||||
hf_token: payload.hf_token,
|
||||
gguf_variant: payload.gguf_variant ?? null,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -8,6 +8,14 @@ import {
|
|||
ModelSelector,
|
||||
} from "@/components/assistant-ui/model-selector";
|
||||
import { Thread } from "@/components/assistant-ui/thread";
|
||||
import { NativeModelChip } from "@/features/native-intents/components/native-model-chip";
|
||||
import { NativeModelDropOverlay } from "@/features/native-intents/components/native-model-drop-overlay";
|
||||
import { useChooseNativeModel } from "@/features/native-intents/use-native-dialogs";
|
||||
import { useNativeModelDrop } from "@/features/native-intents/use-native-drop";
|
||||
import { useNativePathLeasesSupported } from "@/features/native-intents/use-native-readiness";
|
||||
import { useNativeIntentStore } from "@/features/native-intents/store";
|
||||
import type { NativeIntent } from "@/features/native-intents/types";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { GuidedTour, useGuidedTourController } from "@/features/tour";
|
||||
import { useSidebar } from "@/components/ui/sidebar";
|
||||
|
|
@ -574,6 +582,7 @@ export function ChatPage(): ReactElement {
|
|||
const modelsFromStore = useChatRuntimeStore((state) => state.models);
|
||||
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
|
||||
const modelsError = useChatRuntimeStore((state) => state.modelsError);
|
||||
const modelLoading = useChatRuntimeStore((state) => state.modelLoading);
|
||||
const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId);
|
||||
const modelOperationInProgress = useChatRuntimeStore(
|
||||
(state) => state.modelLoading,
|
||||
|
|
@ -587,6 +596,8 @@ export function ChatPage(): ReactElement {
|
|||
loadProgress,
|
||||
loadToastDismissed,
|
||||
} = useChatModelRuntime();
|
||||
const pendingNativeModelIntent = useNativeIntentStore((state) => state.pendingModelIntent);
|
||||
const nativePathLeasesSupported = useNativePathLeasesSupported();
|
||||
const refreshRef = useRef(refresh);
|
||||
const selectModelRef = useRef(selectModel);
|
||||
|
||||
|
|
@ -618,6 +629,60 @@ export function ChatPage(): ReactElement {
|
|||
return { mode: "single" };
|
||||
}, [search.thread, search.compare, search.new, activeThreadId]);
|
||||
|
||||
const hasActiveModel = Boolean(inferenceParams.checkpoint);
|
||||
const loadNativeModelIntent = useCallback(
|
||||
async (intent: NativeIntent, loadingDescription: string) => {
|
||||
const label = intent.path.displayLabel || intent.displayLabel || "Local GGUF model";
|
||||
await selectModel({
|
||||
id: label,
|
||||
nativePathToken: intent.path.token,
|
||||
isDownloaded: true,
|
||||
loadingDescription,
|
||||
forceReload: true,
|
||||
throwOnError: true,
|
||||
});
|
||||
useNativeIntentStore.getState().clearModelIntent(intent.id);
|
||||
},
|
||||
[selectModel],
|
||||
);
|
||||
const handleNativeModelDropAutoLoad = useCallback(
|
||||
(intent: NativeIntent) =>
|
||||
loadNativeModelIntent(
|
||||
intent,
|
||||
hasActiveModel
|
||||
? "Replacing with dropped local GGUF model."
|
||||
: "Loading dropped local GGUF model.",
|
||||
),
|
||||
[hasActiveModel, loadNativeModelIntent],
|
||||
);
|
||||
const handleNativeModelPickerAutoLoad = useCallback(
|
||||
(intent: NativeIntent) =>
|
||||
loadNativeModelIntent(intent, "Loading chosen local GGUF model."),
|
||||
[loadNativeModelIntent],
|
||||
);
|
||||
const canAutoLoadPickedNativeModel = useCallback(() => {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
return (
|
||||
view.mode === "single" &&
|
||||
nativePathLeasesSupported &&
|
||||
!loadingModel &&
|
||||
!modelLoading &&
|
||||
!store.modelLoading &&
|
||||
!store.params.checkpoint
|
||||
);
|
||||
}, [loadingModel, modelLoading, nativePathLeasesSupported, view.mode]);
|
||||
const chooseNativeModel = useChooseNativeModel({
|
||||
shouldAutoLoad: canAutoLoadPickedNativeModel,
|
||||
onAutoLoad: handleNativeModelPickerAutoLoad,
|
||||
});
|
||||
const nativeModelDropState = useNativeModelDrop({
|
||||
enabled: view.mode === "single",
|
||||
nativePathLeasesSupported,
|
||||
hasActiveModel,
|
||||
isModelLoading: Boolean(loadingModel) || modelLoading,
|
||||
onAutoLoad: handleNativeModelDropAutoLoad,
|
||||
});
|
||||
|
||||
const handleCheckpointChange = useCallback(
|
||||
(
|
||||
value: string,
|
||||
|
|
@ -922,6 +987,7 @@ export function ChatPage(): ReactElement {
|
|||
<div className="flex min-h-0 min-w-0 flex-1 basis-0 bg-background overflow-hidden">
|
||||
<GuidedTour {...tour.tourProps} />
|
||||
<div className="relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden">
|
||||
<NativeModelDropOverlay state={nativeModelDropState} />
|
||||
<div
|
||||
className={cn(
|
||||
"absolute top-0 left-0 right-[10px] z-30 flex h-[48px] shrink-0 items-start pt-[11px] pr-2 bg-background",
|
||||
|
|
@ -940,6 +1006,7 @@ export function ChatPage(): ReactElement {
|
|||
onValueChange={handleCheckpointChange}
|
||||
onEject={handleEject}
|
||||
onFoldersChange={refreshLocalModels}
|
||||
onPickLocalModel={isTauri ? chooseNativeModel : undefined}
|
||||
onModelsChange={refreshModelLists}
|
||||
deleteDisabled={modelOperationInProgress}
|
||||
variant="ghost"
|
||||
|
|
@ -950,6 +1017,13 @@ export function ChatPage(): ReactElement {
|
|||
className="max-w-[62vw] !pr-3 sm:max-w-none !h-[34px]"
|
||||
/>
|
||||
)}
|
||||
{pendingNativeModelIntent && view.mode !== "compare" ? (
|
||||
<NativeModelChip
|
||||
intent={pendingNativeModelIntent}
|
||||
nativeReadsDisabled={!nativePathLeasesSupported}
|
||||
onLoad={(selection) => selectModel(selection)}
|
||||
/>
|
||||
) : null}
|
||||
{loadingModel && loadToastDismissed ? (
|
||||
<ModelLoadInlineStatus
|
||||
label={
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
import { createElement, useCallback, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { consumeNativePathToken } from "@/features/native-intents/api";
|
||||
import { ModelLoadDescription } from "../components/model-load-status";
|
||||
import {
|
||||
getDownloadProgress,
|
||||
|
|
@ -47,6 +48,8 @@ type SelectedModelInput = {
|
|||
isDownloaded?: boolean;
|
||||
expectedBytes?: number;
|
||||
forceReload?: boolean;
|
||||
nativePathToken?: string;
|
||||
throwOnError?: boolean;
|
||||
};
|
||||
|
||||
const MODEL_LOAD_TOAST_CLASSNAMES = {
|
||||
|
|
@ -154,6 +157,7 @@ export function useChatModelRuntime() {
|
|||
displayName: string;
|
||||
isDownloaded?: boolean;
|
||||
isCachedLora?: boolean;
|
||||
nativePathToken?: string | null;
|
||||
} | null>(null);
|
||||
const [loadToastDismissed, setLoadToastDismissed] = useState(false);
|
||||
const [loadProgress, setLoadProgress] = useState<{
|
||||
|
|
@ -338,12 +342,20 @@ export function useChatModelRuntime() {
|
|||
typeof selection === "string" ? undefined : selection.ggufVariant;
|
||||
const forceReload =
|
||||
typeof selection === "string" ? false : selection.forceReload ?? false;
|
||||
const nativePathToken =
|
||||
typeof selection === "string" ? undefined : selection.nativePathToken;
|
||||
const throwOnError =
|
||||
typeof selection === "string" ? false : selection.throwOnError ?? false;
|
||||
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
|
||||
if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) {
|
||||
return;
|
||||
}
|
||||
// Prevent duplicate loads if already loading this model
|
||||
if (loadingModelRef.current?.id === modelId) return;
|
||||
if (
|
||||
loadingModelRef.current?.id === modelId &&
|
||||
(loadingModelRef.current?.nativePathToken ?? null) === (nativePathToken ?? null)
|
||||
)
|
||||
return;
|
||||
|
||||
const explicitIsLora =
|
||||
typeof selection === "string" ? undefined : selection.isLora;
|
||||
|
|
@ -383,7 +395,13 @@ export function useChatModelRuntime() {
|
|||
.join(" ");
|
||||
setModelsError(null);
|
||||
setLoadToastDismissedState(false);
|
||||
const loadInfo = { id: modelId, displayName, isDownloaded, isCachedLora };
|
||||
const loadInfo = {
|
||||
id: modelId,
|
||||
displayName,
|
||||
isDownloaded,
|
||||
isCachedLora,
|
||||
nativePathToken: nativePathToken ?? null,
|
||||
};
|
||||
setLoadingModel(loadInfo);
|
||||
useChatRuntimeStore.getState().setModelLoading(true);
|
||||
setLoadProgress(
|
||||
|
|
@ -413,11 +431,17 @@ export function useChatModelRuntime() {
|
|||
const hfToken = stateBeforeUnload.hfToken || null;
|
||||
const previousModelRequiresTrustRemoteCode =
|
||||
stateBeforeUnload.modelRequiresTrustRemoteCode;
|
||||
const previousActiveNativePathToken =
|
||||
stateBeforeUnload.activeNativePathToken;
|
||||
try {
|
||||
// Lightweight pre-flight validation: avoid unloading a working model
|
||||
// if the new identifier is clearly invalid (e.g. bad HF id / path).
|
||||
const validateNativePathLease = nativePathToken
|
||||
? (await consumeNativePathToken(nativePathToken, "validate-model")).nativePathLease
|
||||
: undefined;
|
||||
const validation = await validateModel({
|
||||
model_path: modelId,
|
||||
nativePathLease: validateNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: maxSeqLength,
|
||||
load_in_4bit: true,
|
||||
|
|
@ -428,6 +452,9 @@ export function useChatModelRuntime() {
|
|||
throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
|
||||
}
|
||||
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
|
||||
const loadNativePathLease = nativePathToken
|
||||
? (await consumeNativePathToken(nativePathToken, "load-model")).nativePathLease
|
||||
: undefined;
|
||||
|
||||
if (currentCheckpoint) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
|
|
@ -456,6 +483,7 @@ export function useChatModelRuntime() {
|
|||
});
|
||||
const loadResponse = await loadModel({
|
||||
model_path: modelId,
|
||||
nativePathLease: loadNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: effectiveMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
|
|
@ -528,6 +556,7 @@ export function useChatModelRuntime() {
|
|||
customContextLength: keepCustomCtx,
|
||||
defaultChatTemplate: loadResponse.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
activeNativePathToken: nativePathToken ?? null,
|
||||
});
|
||||
// Qwen3/3.5/3.6: apply thinking-mode-specific params after load
|
||||
if (modelId.toLowerCase().includes("qwen3") && (loadResponse.supports_reasoning ?? false)) {
|
||||
|
|
@ -564,9 +593,22 @@ export function useChatModelRuntime() {
|
|||
if (abortCtrl.signal.aborted) throw error;
|
||||
// If we unloaded a previous model and the new load failed, attempt a rollback.
|
||||
if (previousWasUnloaded && previousCheckpoint) {
|
||||
let rollbackNativePathLease: string | undefined;
|
||||
if (previousActiveNativePathToken) {
|
||||
try {
|
||||
rollbackNativePathLease = (
|
||||
await consumeNativePathToken(previousActiveNativePathToken, "load-model")
|
||||
).nativePathLease;
|
||||
} catch {
|
||||
throw new Error(
|
||||
"Could not reload the previous local model: please re-select the file.",
|
||||
);
|
||||
}
|
||||
}
|
||||
try {
|
||||
await loadModel({
|
||||
model_path: previousCheckpoint,
|
||||
nativePathLease: rollbackNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: rollbackMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
|
|
@ -575,9 +617,12 @@ export function useChatModelRuntime() {
|
|||
trust_remote_code:
|
||||
previousModelRequiresTrustRemoteCode || trustRemoteCode,
|
||||
});
|
||||
useChatRuntimeStore.setState({
|
||||
activeNativePathToken: previousActiveNativePathToken ?? null,
|
||||
});
|
||||
await refresh();
|
||||
} catch {
|
||||
// If rollback also fails, surface the original error.
|
||||
// Rollback also failed; surface the original load error below.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
|
|
@ -879,6 +924,9 @@ export function useChatModelRuntime() {
|
|||
const message =
|
||||
error instanceof Error ? error.message : "Failed to load model";
|
||||
setModelsError(message);
|
||||
if (throwOnError) {
|
||||
throw error instanceof Error ? error : new Error(message);
|
||||
}
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
} from "@/components/ui/dropdown-menu";
|
||||
import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
|
||||
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { useAui } from "@assistant-ui/react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
|
||||
|
|
@ -493,11 +494,15 @@ export function SharedComposer({
|
|||
<div
|
||||
className={`chat-composer-surface relative flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 transition-shadow outline-none ${dragging ? "border-ring bg-accent/50" : ""}`}
|
||||
onDragOver={(e) => {
|
||||
if (isTauri) return;
|
||||
e.preventDefault();
|
||||
setDragging(true);
|
||||
}}
|
||||
onDragLeave={() => setDragging(false)}
|
||||
onDrop={(e) => {
|
||||
// Phase 1 native model drops own Tauri local-path drops. Restore browser
|
||||
// attachment drops in Tauri when Phase 1d adds attachment-token bridging.
|
||||
if (isTauri) return;
|
||||
e.preventDefault();
|
||||
setDragging(false);
|
||||
addFiles(e.dataTransfer.files);
|
||||
|
|
|
|||
|
|
@ -226,6 +226,7 @@ type ChatRuntimeStore = {
|
|||
cachedTokens: number;
|
||||
} | null;
|
||||
modelLoading: boolean;
|
||||
activeNativePathToken: string | null;
|
||||
setModelLoading: (loading: boolean) => void;
|
||||
setModelRequiresTrustRemoteCode: (required: boolean) => void;
|
||||
setParams: (params: InferenceParams) => void;
|
||||
|
|
@ -305,6 +306,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
pendingAudioName: null,
|
||||
contextUsage: null,
|
||||
modelLoading: false,
|
||||
activeNativePathToken: null,
|
||||
setModelLoading: (loading) => set({ modelLoading: loading }),
|
||||
setModelRequiresTrustRemoteCode: (modelRequiresTrustRemoteCode) =>
|
||||
set({ modelRequiresTrustRemoteCode }),
|
||||
|
|
@ -378,6 +380,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
checkpoint: "",
|
||||
},
|
||||
activeGgufVariant: null,
|
||||
activeNativePathToken: null,
|
||||
ggufContextLength: null,
|
||||
ggufMaxContextLength: null,
|
||||
ggufNativeContextLength: null,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ export interface ListLorasResponse {
|
|||
|
||||
export interface LoadModelRequest {
|
||||
model_path: string;
|
||||
nativePathLease?: string | null;
|
||||
hf_token: string | null;
|
||||
max_seq_length: number;
|
||||
load_in_4bit: boolean;
|
||||
|
|
|
|||
46
studio/frontend/src/features/native-intents/api.ts
Normal file
46
studio/frontend/src/features/native-intents/api.ts
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
import { isTauri } from "@/lib/api-base";
|
||||
import type {
|
||||
NativeIntent,
|
||||
NativePathLeaseResponse,
|
||||
NativePathOperation,
|
||||
} from "./types";
|
||||
|
||||
async function invokeNative<T>(command: string, args?: Record<string, unknown>): Promise<T> {
|
||||
if (!isTauri) {
|
||||
throw new Error("Native desktop features are only available in the Tauri app.");
|
||||
}
|
||||
const { invoke } = await import("@tauri-apps/api/core");
|
||||
return invoke<T>(command, args);
|
||||
}
|
||||
|
||||
export async function drainNativeIntents(): Promise<NativeIntent[]> {
|
||||
if (!isTauri) return [];
|
||||
return invokeNative<NativeIntent[]>("drain_native_intents");
|
||||
}
|
||||
|
||||
export async function pickNativeModel(): Promise<NativeIntent | null> {
|
||||
if (!isTauri) return null;
|
||||
return invokeNative<NativeIntent | null>("pick_native_model");
|
||||
}
|
||||
|
||||
export async function registerNativeModelPath(path: string): Promise<NativeIntent> {
|
||||
return invokeNative<NativeIntent>("register_native_model_path", { path });
|
||||
}
|
||||
|
||||
export async function consumeNativePathToken(
|
||||
token: string,
|
||||
operation: NativePathOperation,
|
||||
): Promise<NativePathLeaseResponse> {
|
||||
return invokeNative<NativePathLeaseResponse>("consume_native_path_token", {
|
||||
token,
|
||||
operation,
|
||||
});
|
||||
}
|
||||
|
||||
export async function revealPathToken(token: string): Promise<void> {
|
||||
return invokeNative<void>("reveal_path_token", { token });
|
||||
}
|
||||
|
||||
export async function openPathToken(token: string): Promise<void> {
|
||||
return invokeNative<void>("open_path_token", { token });
|
||||
}
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
import { revealPathToken } from "../api";
|
||||
import { useNativeIntentStore } from "../store";
|
||||
import type { NativeIntent } from "../types";
|
||||
import { XIcon } from "lucide-react";
|
||||
import { useEffect, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
interface NativeModelChipProps {
|
||||
intent: NativeIntent;
|
||||
nativeReadsDisabled: boolean;
|
||||
onLoad: (selection: {
|
||||
id: string;
|
||||
nativePathToken: string;
|
||||
isDownloaded: boolean;
|
||||
loadingDescription: string;
|
||||
forceReload: boolean;
|
||||
throwOnError?: boolean;
|
||||
}) => Promise<void> | void;
|
||||
}
|
||||
|
||||
export function NativeModelChip({
|
||||
intent,
|
||||
nativeReadsDisabled,
|
||||
onLoad,
|
||||
}: NativeModelChipProps) {
|
||||
const clearModelIntent = useNativeIntentStore((state) => state.clearModelIntent);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [now, setNow] = useState(() => Date.now());
|
||||
const label = intent.path.displayLabel || intent.displayLabel || "Local GGUF model";
|
||||
const expired = intent.path.expiresAtMs <= now;
|
||||
|
||||
useEffect(() => {
|
||||
if (expired) return;
|
||||
const remaining = Math.max(0, intent.path.expiresAtMs - Date.now());
|
||||
const timer = window.setTimeout(() => setNow(Date.now()), remaining);
|
||||
return () => window.clearTimeout(timer);
|
||||
}, [expired, intent.path.expiresAtMs]);
|
||||
|
||||
async function handleLoad() {
|
||||
if (nativeReadsDisabled || expired) return;
|
||||
setLoading(true);
|
||||
try {
|
||||
await onLoad({
|
||||
id: label,
|
||||
nativePathToken: intent.path.token,
|
||||
isDownloaded: true,
|
||||
loadingDescription: "Loading selected local GGUF model.",
|
||||
forceReload: true,
|
||||
throwOnError: true,
|
||||
});
|
||||
clearModelIntent(intent.id);
|
||||
} catch {
|
||||
// selectModel reports the failure; keep the chip available for retry.
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
}
|
||||
|
||||
async function handleReveal() {
|
||||
try {
|
||||
await revealPathToken(intent.path.token);
|
||||
} catch (error) {
|
||||
toast.error("Could not reveal model", {
|
||||
description: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex min-w-0 max-w-[34rem] items-center gap-2 rounded-lg border border-border/70 bg-muted/70 px-2.5 py-1.5 text-xs">
|
||||
<span className="shrink-0 font-medium text-muted-foreground">Local GGUF</span>
|
||||
<span className="min-w-0 flex-1 truncate" title={label}>{label}</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleReveal}
|
||||
disabled={expired}
|
||||
title={expired ? "Selection expired, pick or drop the file again" : undefined}
|
||||
className="rounded-md px-2 py-1 text-muted-foreground transition-colors hover:bg-background hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50"
|
||||
>
|
||||
Reveal
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleLoad}
|
||||
disabled={nativeReadsDisabled || expired || loading}
|
||||
title={
|
||||
nativeReadsDisabled
|
||||
? "Managed desktop backend required"
|
||||
: expired
|
||||
? "Selection expired, pick or drop the file again"
|
||||
: undefined
|
||||
}
|
||||
className="rounded-md bg-foreground px-2 py-1 text-background transition-opacity disabled:cursor-not-allowed disabled:opacity-50"
|
||||
>
|
||||
{expired ? "Expired" : loading ? "Loading…" : "Load model"}
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => clearModelIntent(intent.id)}
|
||||
className="flex size-5 items-center justify-center rounded-full text-muted-foreground hover:bg-destructive hover:text-destructive-foreground"
|
||||
aria-label="Dismiss local model"
|
||||
>
|
||||
<XIcon className="size-3" />
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,70 @@
|
|||
import { cn } from "@/lib/utils";
|
||||
import { FileUpIcon } from "lucide-react";
|
||||
import type { NativeModelDropState } from "../use-native-drop";
|
||||
|
||||
function overlayCopy(state: NativeModelDropState): { title: string; description: string } {
|
||||
if (state.status === "invalid") {
|
||||
return {
|
||||
title: "GGUF models only",
|
||||
description: "Other files are not handled here yet.",
|
||||
};
|
||||
}
|
||||
if (state.status === "valid" && state.action === "replace") {
|
||||
return {
|
||||
title: "Drop to replace model",
|
||||
description: "Current model will unload first.",
|
||||
};
|
||||
}
|
||||
if (state.status === "valid" && state.action === "load") {
|
||||
return {
|
||||
title: "Drop to load model",
|
||||
description: "Adds it as the active chat model.",
|
||||
};
|
||||
}
|
||||
return {
|
||||
title: "Drop to add model chip",
|
||||
description: "Review it before loading.",
|
||||
};
|
||||
}
|
||||
|
||||
export function NativeModelDropOverlay({ state }: { state: NativeModelDropState }) {
|
||||
const isIdle = state.status === "idle";
|
||||
const isAutoLoad = state.status === "valid" && state.action !== "chip";
|
||||
const isInvalid = state.status === "invalid";
|
||||
const { title, description } = overlayCopy(state);
|
||||
|
||||
return (
|
||||
<div
|
||||
className={cn(
|
||||
"pointer-events-none absolute left-1/2 top-4 z-50 w-[clamp(16rem,28vw,22rem)] max-w-[calc(100vw-1rem)] -translate-x-1/2 transition-all duration-200 ease-out",
|
||||
isIdle ? "-translate-y-1 opacity-0" : "translate-y-0 opacity-100",
|
||||
)}
|
||||
role="status"
|
||||
aria-live="polite"
|
||||
aria-hidden={isIdle}
|
||||
>
|
||||
<div className="flex items-center gap-2.5 rounded-xl border border-border/70 bg-card/96 px-3 py-2.5 shadow-sm shadow-black/10 backdrop-blur-sm dark:shadow-black/30">
|
||||
<div
|
||||
className={cn(
|
||||
"flex size-8 shrink-0 items-center justify-center rounded-lg border transition-colors duration-200 ease-out",
|
||||
isAutoLoad
|
||||
? "border-emerald-500/30 bg-emerald-500/10 text-emerald-600 dark:text-emerald-400"
|
||||
: isInvalid
|
||||
? "border-amber-500/25 bg-amber-500/10 text-amber-700 dark:text-amber-300"
|
||||
: "border-border bg-muted text-muted-foreground",
|
||||
)}
|
||||
>
|
||||
<FileUpIcon className="size-4" aria-hidden="true" />
|
||||
</div>
|
||||
<div className="min-w-0 flex-1">
|
||||
<div className="truncate text-xs font-medium text-foreground">
|
||||
{title}
|
||||
</div>
|
||||
<div className="mt-0.5 truncate text-[11px] leading-4 text-muted-foreground">
|
||||
{description}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
import { isTauri } from "@/lib/api-base";
|
||||
import { useEffect } from "react";
|
||||
import { drainNativeIntents } from "./api";
|
||||
import { useNativeIntentStore } from "./store";
|
||||
|
||||
export function NativeIntentDrain() {
|
||||
const addIntent = useNativeIntentStore((state) => state.addIntent);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isTauri) return;
|
||||
let disposed = false;
|
||||
let unlisten: (() => void) | undefined;
|
||||
|
||||
async function drain() {
|
||||
const intents = await drainNativeIntents().catch(() => []);
|
||||
if (disposed) return;
|
||||
for (const intent of intents) addIntent(intent);
|
||||
}
|
||||
|
||||
void drain();
|
||||
void import("@tauri-apps/api/event")
|
||||
.then(({ listen }) => listen("native-intent-available", drain))
|
||||
.then((cleanup) => {
|
||||
if (disposed) {
|
||||
cleanup();
|
||||
} else {
|
||||
unlisten = cleanup;
|
||||
}
|
||||
})
|
||||
.catch(() => undefined);
|
||||
|
||||
return () => {
|
||||
disposed = true;
|
||||
unlisten?.();
|
||||
};
|
||||
}, [addIntent]);
|
||||
|
||||
return null;
|
||||
}
|
||||
23
studio/frontend/src/features/native-intents/store.ts
Normal file
23
studio/frontend/src/features/native-intents/store.ts
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
import { create } from "zustand";
|
||||
import type { NativeIntent } from "./types";
|
||||
|
||||
interface NativeIntentState {
|
||||
pendingModelIntent: NativeIntent | null;
|
||||
addIntent: (intent: NativeIntent) => void;
|
||||
clearModelIntent: (intentId?: string) => void;
|
||||
}
|
||||
|
||||
export const useNativeIntentStore = create<NativeIntentState>((set, get) => ({
|
||||
pendingModelIntent: null,
|
||||
addIntent: (intent) => {
|
||||
if (intent.kind !== "model") return;
|
||||
const current = get().pendingModelIntent;
|
||||
if (current?.path.token === intent.path.token) return;
|
||||
set({ pendingModelIntent: intent });
|
||||
},
|
||||
clearModelIntent: (intentId) => {
|
||||
const current = get().pendingModelIntent;
|
||||
if (intentId && current?.id !== intentId) return;
|
||||
set({ pendingModelIntent: null });
|
||||
},
|
||||
}));
|
||||
39
studio/frontend/src/features/native-intents/types.ts
Normal file
39
studio/frontend/src/features/native-intents/types.ts
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
export type NativePathOperation =
|
||||
| "validate-model"
|
||||
| "load-model"
|
||||
| "dataset-preview"
|
||||
| "dataset-import"
|
||||
| "attach"
|
||||
| "reveal"
|
||||
| "open";
|
||||
|
||||
export type NativePathKind = "model" | "dataset" | "attachment" | "artifact";
|
||||
|
||||
export type NativePathSourceKind =
|
||||
| "dialog"
|
||||
| "drop"
|
||||
| "deep-link"
|
||||
| "file-association"
|
||||
| "artifact";
|
||||
|
||||
export interface NativePathRef {
|
||||
token: string;
|
||||
kind: NativePathKind;
|
||||
displayLabel: string;
|
||||
allowedOperations: NativePathOperation[];
|
||||
expiresAtMs: number;
|
||||
}
|
||||
|
||||
export interface NativeIntent {
|
||||
id: string;
|
||||
kind: NativePathKind;
|
||||
sourceKind: NativePathSourceKind;
|
||||
path: NativePathRef;
|
||||
displayLabel: string;
|
||||
}
|
||||
|
||||
export interface NativePathLeaseResponse {
|
||||
nativePathLease: string;
|
||||
displayLabel: string;
|
||||
expiresAtMs: number;
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
import { useCallback, useRef } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { pickNativeModel } from "./api";
|
||||
import { useNativeIntentStore } from "./store";
|
||||
import type { NativeIntent } from "./types";
|
||||
|
||||
interface ChooseNativeModelOptions {
|
||||
shouldAutoLoad?: (intent: NativeIntent) => boolean;
|
||||
onAutoLoad?: (intent: NativeIntent) => Promise<void> | void;
|
||||
}
|
||||
|
||||
function isGgufModelIntent(intent: NativeIntent): boolean {
|
||||
const label = intent.path.displayLabel || intent.displayLabel;
|
||||
return (
|
||||
intent.kind === "model" &&
|
||||
intent.path.kind === "model" &&
|
||||
label.toLowerCase().endsWith(".gguf") &&
|
||||
intent.path.allowedOperations.includes("validate-model") &&
|
||||
intent.path.allowedOperations.includes("load-model")
|
||||
);
|
||||
}
|
||||
|
||||
export function useChooseNativeModel(options: ChooseNativeModelOptions = {}) {
|
||||
const addIntent = useNativeIntentStore((state) => state.addIntent);
|
||||
const pickingRef = useRef(false);
|
||||
const onAutoLoad = options.onAutoLoad;
|
||||
const shouldAutoLoad = options.shouldAutoLoad;
|
||||
|
||||
return useCallback(async () => {
|
||||
if (pickingRef.current) return;
|
||||
pickingRef.current = true;
|
||||
try {
|
||||
let intent: NativeIntent | null = null;
|
||||
try {
|
||||
intent = await pickNativeModel();
|
||||
} catch (error) {
|
||||
toast.error("Could not choose local model", {
|
||||
description: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (!intent) return;
|
||||
|
||||
let runAutoLoad = false;
|
||||
try {
|
||||
runAutoLoad =
|
||||
Boolean(onAutoLoad) &&
|
||||
isGgufModelIntent(intent) &&
|
||||
shouldAutoLoad?.(intent) === true;
|
||||
} catch {
|
||||
runAutoLoad = false;
|
||||
}
|
||||
if (!runAutoLoad) {
|
||||
addIntent(intent);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await onAutoLoad?.(intent);
|
||||
} catch {
|
||||
addIntent(intent);
|
||||
}
|
||||
} finally {
|
||||
pickingRef.current = false;
|
||||
}
|
||||
}, [addIntent, onAutoLoad, shouldAutoLoad]);
|
||||
}
|
||||
132
studio/frontend/src/features/native-intents/use-native-drop.ts
Normal file
132
studio/frontend/src/features/native-intents/use-native-drop.ts
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
import { isTauri } from "@/lib/api-base";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { registerNativeModelPath } from "./api";
|
||||
import { useNativeIntentStore } from "./store";
|
||||
import type { NativeIntent } from "./types";
|
||||
|
||||
export type NativeModelDropState =
|
||||
| { status: "idle" }
|
||||
| { status: "valid"; action: "load" | "replace" | "chip" }
|
||||
| { status: "invalid" };
|
||||
|
||||
interface NativeModelDropOptions {
|
||||
enabled?: boolean;
|
||||
nativePathLeasesSupported: boolean;
|
||||
hasActiveModel: boolean;
|
||||
isModelLoading: boolean;
|
||||
onAutoLoad?: (intent: NativeIntent) => Promise<void> | void;
|
||||
}
|
||||
|
||||
function ggufPaths(paths: string[]): string[] {
|
||||
return paths.filter((path) => path.toLowerCase().endsWith(".gguf"));
|
||||
}
|
||||
|
||||
function canAutoLoadPaths(paths: string[], options: NativeModelDropOptions): boolean {
|
||||
return (
|
||||
paths.length === 1 &&
|
||||
ggufPaths(paths).length === 1 &&
|
||||
options.nativePathLeasesSupported &&
|
||||
!options.isModelLoading &&
|
||||
Boolean(options.onAutoLoad)
|
||||
);
|
||||
}
|
||||
|
||||
function dropStateForPaths(
|
||||
paths: string[],
|
||||
options: NativeModelDropOptions,
|
||||
): NativeModelDropState {
|
||||
if (paths.length === 0) {
|
||||
return { status: "idle" };
|
||||
}
|
||||
const ggufs = ggufPaths(paths);
|
||||
if (paths.length !== 1 || ggufs.length !== 1) {
|
||||
return { status: "invalid" };
|
||||
}
|
||||
if (!canAutoLoadPaths(paths, options)) {
|
||||
return { status: "valid", action: "chip" };
|
||||
}
|
||||
return {
|
||||
status: "valid",
|
||||
action: options.hasActiveModel ? "replace" : "load",
|
||||
};
|
||||
}
|
||||
|
||||
export function useNativeModelDrop(options: NativeModelDropOptions): NativeModelDropState {
|
||||
const { enabled = true } = options;
|
||||
const addIntent = useNativeIntentStore((state) => state.addIntent);
|
||||
const [dropState, setDropState] = useState<NativeModelDropState>({ status: "idle" });
|
||||
const optionsRef = useRef(options);
|
||||
optionsRef.current = options;
|
||||
|
||||
useEffect(() => {
|
||||
if (!isTauri || !enabled) {
|
||||
setDropState({ status: "idle" });
|
||||
return;
|
||||
}
|
||||
let disposed = false;
|
||||
let unlisten: (() => void) | undefined;
|
||||
|
||||
void import("@tauri-apps/api/window")
|
||||
.then(({ getCurrentWindow }) => getCurrentWindow().onDragDropEvent(async (event) => {
|
||||
const currentOptions = optionsRef.current;
|
||||
if (event.payload.type === "enter") {
|
||||
setDropState(dropStateForPaths(event.payload.paths, currentOptions));
|
||||
return;
|
||||
}
|
||||
if (event.payload.type === "leave") {
|
||||
setDropState({ status: "idle" });
|
||||
return;
|
||||
}
|
||||
if (event.payload.type !== "drop") return;
|
||||
setDropState({ status: "idle" });
|
||||
const ggufs = ggufPaths(event.payload.paths);
|
||||
if (event.payload.paths.length !== 1 || ggufs.length !== 1) {
|
||||
if (event.payload.paths.length > 0) {
|
||||
toast.error(
|
||||
ggufs.length === 0
|
||||
? "Only .gguf model files can be dropped here."
|
||||
: "Drop a single .gguf model file.",
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
const ggufPath = ggufs[0];
|
||||
try {
|
||||
const intent = await registerNativeModelPath(ggufPath);
|
||||
if (disposed) return;
|
||||
if (!canAutoLoadPaths(event.payload.paths, currentOptions)) {
|
||||
addIntent(intent);
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await currentOptions.onAutoLoad?.(intent);
|
||||
} catch (error) {
|
||||
addIntent(intent);
|
||||
toast.error("Could not load dropped model", {
|
||||
description: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
toast.error("Could not use dropped model", {
|
||||
description: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
}
|
||||
}))
|
||||
.then((cleanup) => {
|
||||
if (disposed) {
|
||||
cleanup();
|
||||
} else {
|
||||
unlisten = cleanup;
|
||||
}
|
||||
})
|
||||
.catch(() => undefined);
|
||||
|
||||
return () => {
|
||||
disposed = true;
|
||||
unlisten?.();
|
||||
};
|
||||
}, [addIntent, enabled]);
|
||||
|
||||
return dropState;
|
||||
}
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
import { apiUrl, isTauri } from "@/lib/api-base";
|
||||
import { useEffect, useState } from "react";
|
||||
|
||||
const MAX_READINESS_POLLS = 720;
|
||||
|
||||
export function useNativePathLeasesSupported(): boolean {
|
||||
const [supported, setSupported] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isTauri) return;
|
||||
let disposed = false;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
let controller: AbortController | undefined;
|
||||
let polls = 0;
|
||||
|
||||
function check(delay = 0) {
|
||||
if (polls >= MAX_READINESS_POLLS) return;
|
||||
polls += 1;
|
||||
timer = setTimeout(() => {
|
||||
controller = new AbortController();
|
||||
fetch(apiUrl("/api/health"), { signal: controller.signal })
|
||||
.then((response) => response.json())
|
||||
.then((health) => {
|
||||
if (disposed) return;
|
||||
if (health?.native_path_leases_supported === true) {
|
||||
setSupported(true);
|
||||
} else {
|
||||
check(5000);
|
||||
}
|
||||
})
|
||||
.catch(() => {
|
||||
if (!disposed) check(5000);
|
||||
});
|
||||
}, delay);
|
||||
}
|
||||
|
||||
check();
|
||||
return () => {
|
||||
disposed = true;
|
||||
if (timer) clearTimeout(timer);
|
||||
controller?.abort();
|
||||
};
|
||||
}, []);
|
||||
|
||||
return supported;
|
||||
}
|
||||
|
|
@ -21,6 +21,10 @@ import tempfile
|
|||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_DIR = Path(__file__).resolve().parent / "backend"
|
||||
if str(_BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(1, str(_BACKEND_DIR))
|
||||
|
||||
from backend.utils.wheel_utils import (
|
||||
flash_attn_package_version,
|
||||
flash_attn_wheel_url,
|
||||
|
|
|
|||
80
studio/src-tauri/Cargo.lock
generated
80
studio/src-tauri/Cargo.lock
generated
|
|
@ -808,6 +808,7 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
|||
dependencies = [
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1694,6 +1695,15 @@ version = "0.4.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "home"
|
||||
version = "0.5.12"
|
||||
|
|
@ -3561,6 +3571,30 @@ dependencies = [
|
|||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rfd"
|
||||
version = "0.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a15ad77d9e70a92437d8f74c35d99b4e4691128df018833e99f90bcd36152672"
|
||||
dependencies = [
|
||||
"block2",
|
||||
"dispatch2",
|
||||
"glib-sys",
|
||||
"gobject-sys",
|
||||
"gtk-sys",
|
||||
"js-sys",
|
||||
"log",
|
||||
"objc2",
|
||||
"objc2-app-kit",
|
||||
"objc2-core-foundation",
|
||||
"objc2-foundation",
|
||||
"raw-window-handle",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
"windows-sys 0.60.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
|
|
@ -4513,6 +4547,48 @@ dependencies = [
|
|||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-dialog"
|
||||
version = "2.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a1fa4150c95ae391946cc8b8f905ab14797427caba3a8a2f79628e956da91809"
|
||||
dependencies = [
|
||||
"log",
|
||||
"raw-window-handle",
|
||||
"rfd",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tauri",
|
||||
"tauri-plugin",
|
||||
"tauri-plugin-fs",
|
||||
"thiserror 2.0.18",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-fs"
|
||||
version = "2.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "36e1ec28b79f3d0683f4507e1615c36292c0ea6716668770d4396b9b39871ed8"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"dunce",
|
||||
"glob",
|
||||
"log",
|
||||
"objc2-foundation",
|
||||
"percent-encoding",
|
||||
"schemars 0.8.22",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
"tauri",
|
||||
"tauri-plugin",
|
||||
"tauri-utils",
|
||||
"thiserror 2.0.18",
|
||||
"toml 0.9.12+spec-1.1.0",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-opener"
|
||||
version = "2.5.3"
|
||||
|
|
@ -5200,9 +5276,11 @@ checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
|||
name = "unsloth-studio"
|
||||
version = "2026.4.8"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"dirs",
|
||||
"elevated-command",
|
||||
"fix-path-env",
|
||||
"hmac",
|
||||
"libc",
|
||||
"log",
|
||||
"open",
|
||||
|
|
@ -5212,10 +5290,12 @@ dependencies = [
|
|||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"simplelog",
|
||||
"tauri",
|
||||
"tauri-build",
|
||||
"tauri-plugin-clipboard-manager",
|
||||
"tauri-plugin-dialog",
|
||||
"tauri-plugin-opener",
|
||||
"tauri-plugin-process",
|
||||
"tauri-plugin-single-instance",
|
||||
|
|
|
|||
|
|
@ -11,6 +11,9 @@ tauri-plugin-single-instance = "2"
|
|||
tauri-plugin-process = "2"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
base64 = "0.22"
|
||||
hmac = "0.12"
|
||||
sha2 = "0.10"
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
log = "0.4"
|
||||
|
|
@ -23,6 +26,7 @@ fix-path-env = { git = "https://github.com/tauri-apps/fix-path-env-rs" }
|
|||
tauri-plugin-opener = "2.5.3"
|
||||
tauri-plugin-updater = "2"
|
||||
tauri-plugin-clipboard-manager = "2"
|
||||
tauri-plugin-dialog = "2"
|
||||
rand = "0.10.0"
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
|
|
|
|||
|
|
@ -20,6 +20,12 @@ pub(crate) fn redact_text(text: &str, report: &mut RedactionReport) -> String {
|
|||
out = replace_regex(cookie_re(), &out, "$1: <redacted>", report);
|
||||
out = replace_regex(token_re(), &out, "<redacted token>", report);
|
||||
out = replace_regex(env_secret_re(), &out, "$1=<redacted>", report);
|
||||
out = replace_regex(
|
||||
native_path_lease_re(),
|
||||
&out,
|
||||
"$1<redacted native path lease>",
|
||||
report,
|
||||
);
|
||||
out = replace_known_paths(&out, report);
|
||||
out = replace_regex(windows_studio_re(), &out, "<studio_home>", report);
|
||||
out = replace_regex(windows_home_re(), &out, "%USERPROFILE%", report);
|
||||
|
|
@ -126,6 +132,16 @@ fn env_secret_re() -> &'static Regex {
|
|||
})
|
||||
}
|
||||
|
||||
fn native_path_lease_re() -> &'static Regex {
|
||||
static RE: OnceLock<Regex> = OnceLock::new();
|
||||
RE.get_or_init(|| {
|
||||
Regex::new(
|
||||
r#"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"#,
|
||||
)
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
|
||||
fn windows_studio_re() -> &'static Regex {
|
||||
static RE: OnceLock<Regex> = OnceLock::new();
|
||||
RE.get_or_init(|| {
|
||||
|
|
@ -165,6 +181,7 @@ mod tests {
|
|||
"Cookie: session=abcdef\n",
|
||||
"HF_TOKEN=hf_abcdefghijklmnopqrstuvwxyz\n",
|
||||
"API_KEY=secret123\n",
|
||||
"native_path_lease=abc.DEF_123\n",
|
||||
"url=https://user:pass@example.com/path\n",
|
||||
"email=alex@example.com\n",
|
||||
"path=/Users/alex/.unsloth/studio/logs/install.log\n",
|
||||
|
|
@ -178,6 +195,8 @@ mod tests {
|
|||
assert!(!redacted.contains("session=abcdef"));
|
||||
assert!(!redacted.contains("hf_abcdefghijklmnopqrstuvwxyz"));
|
||||
assert!(!redacted.contains("secret123"));
|
||||
assert!(!redacted.contains("abc.DEF_123"));
|
||||
assert!(redacted.contains("native_path_lease=<redacted native path lease>"));
|
||||
assert!(redacted.contains("https://<redacted>@example.com/path"));
|
||||
assert!(!redacted.contains("alex@example.com"));
|
||||
assert!(redacted.contains("<studio_home>"));
|
||||
|
|
|
|||
|
|
@ -4,6 +4,9 @@ mod commands;
|
|||
mod desktop_auth;
|
||||
mod diagnostics;
|
||||
mod install;
|
||||
mod native_backend_lease;
|
||||
mod native_intents;
|
||||
mod native_path_policy;
|
||||
mod preflight;
|
||||
mod process;
|
||||
mod update;
|
||||
|
|
@ -164,10 +167,12 @@ fn main() {
|
|||
}))
|
||||
.plugin(tauri_plugin_process::init())
|
||||
.plugin(tauri_plugin_opener::init())
|
||||
.plugin(tauri_plugin_dialog::init())
|
||||
.plugin(tauri_plugin_updater::Builder::new().build())
|
||||
.plugin(tauri_plugin_clipboard_manager::init())
|
||||
.manage(diagnostics::new_diagnostics_state())
|
||||
.manage(install::new_install_state())
|
||||
.manage(native_intents::new_native_intake_state())
|
||||
.manage(new_backend_state())
|
||||
.manage(process::new_shutdown_flag())
|
||||
.manage(update::new_update_state())
|
||||
|
|
@ -187,6 +192,13 @@ fn main() {
|
|||
commands::install_system_packages,
|
||||
desktop_auth::desktop_auth,
|
||||
diagnostics::collect_support_diagnostics,
|
||||
native_intents::drain_native_intents,
|
||||
native_intents::register_native_model_path,
|
||||
native_intents::pick_native_model,
|
||||
native_intents::consume_native_path_token,
|
||||
native_intents::register_artifact_path,
|
||||
native_intents::reveal_path_token,
|
||||
native_intents::open_path_token,
|
||||
])
|
||||
.setup(|app| {
|
||||
#[cfg(any(target_os = "windows", target_os = "linux"))]
|
||||
|
|
|
|||
199
studio/src-tauri/src/native_backend_lease.rs
Normal file
199
studio/src-tauri/src/native_backend_lease.rs
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use hmac::{Hmac, Mac};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
pub const LEASE_SECRET_ENV: &str = "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET";
|
||||
const LEASE_VERSION: u8 = 1;
|
||||
const LEASE_TTL: Duration = Duration::from_secs(2 * 60);
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum NativePathOperation {
|
||||
ValidateModel,
|
||||
LoadModel,
|
||||
DatasetPreview,
|
||||
DatasetImport,
|
||||
Attach,
|
||||
Reveal,
|
||||
Open,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum NativePathKind {
|
||||
Model,
|
||||
Dataset,
|
||||
Attachment,
|
||||
Artifact,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum NativePathSourceKind {
|
||||
Dialog,
|
||||
Drop,
|
||||
DeepLink,
|
||||
FileAssociation,
|
||||
Artifact,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum NativePathType {
|
||||
File,
|
||||
Directory,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub struct NativePathLeasePayload {
|
||||
pub version: u8,
|
||||
pub operation: NativePathOperation,
|
||||
pub canonical_path: String,
|
||||
pub path_kind: NativePathKind,
|
||||
pub path_type: NativePathType,
|
||||
pub source_kind: NativePathSourceKind,
|
||||
pub token_id_hash: String,
|
||||
pub issued_at_ms: u64,
|
||||
pub expires_at_ms: u64,
|
||||
pub nonce: String,
|
||||
pub display_label: String,
|
||||
pub size_bytes: Option<u64>,
|
||||
pub modified_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NativePathLeaseResponse {
|
||||
pub native_path_lease: String,
|
||||
pub display_label: String,
|
||||
pub expires_at_ms: u64,
|
||||
}
|
||||
|
||||
pub fn new_lease_secret() -> Vec<u8> {
|
||||
rand::random::<[u8; 32]>().to_vec()
|
||||
}
|
||||
|
||||
pub fn encode_secret_env(secret: &[u8]) -> String {
|
||||
URL_SAFE_NO_PAD.encode(secret)
|
||||
}
|
||||
|
||||
pub fn now_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_else(|_| Duration::from_secs(0))
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
pub fn token_hash(token: &str) -> String {
|
||||
hex_bytes(&Sha256::digest(token.as_bytes()))
|
||||
}
|
||||
|
||||
pub fn random_token(prefix: &str) -> String {
|
||||
format!("{}{}", prefix, hex_bytes(&rand::random::<[u8; 24]>()))
|
||||
}
|
||||
|
||||
pub fn random_nonce() -> String {
|
||||
hex_bytes(&rand::random::<[u8; 16]>())
|
||||
}
|
||||
|
||||
pub fn sign_path_lease(
|
||||
secret: &[u8],
|
||||
operation: NativePathOperation,
|
||||
canonical_path: String,
|
||||
path_kind: NativePathKind,
|
||||
path_type: NativePathType,
|
||||
source_kind: NativePathSourceKind,
|
||||
token: &str,
|
||||
display_label: String,
|
||||
size_bytes: Option<u64>,
|
||||
modified_ms: Option<u64>,
|
||||
) -> Result<NativePathLeaseResponse, String> {
|
||||
let issued_at_ms = now_ms();
|
||||
let expires_at_ms = issued_at_ms + LEASE_TTL.as_millis() as u64;
|
||||
let payload = NativePathLeasePayload {
|
||||
version: LEASE_VERSION,
|
||||
operation,
|
||||
canonical_path,
|
||||
path_kind,
|
||||
path_type,
|
||||
source_kind,
|
||||
token_id_hash: token_hash(token),
|
||||
issued_at_ms,
|
||||
expires_at_ms,
|
||||
nonce: random_nonce(),
|
||||
display_label: display_label.clone(),
|
||||
size_bytes,
|
||||
modified_ms,
|
||||
};
|
||||
sign_payload(secret, &payload).map(|native_path_lease| NativePathLeaseResponse {
|
||||
native_path_lease,
|
||||
display_label,
|
||||
expires_at_ms,
|
||||
})
|
||||
}
|
||||
|
||||
fn sign_payload(secret: &[u8], payload: &NativePathLeasePayload) -> Result<String, String> {
|
||||
let payload_json = serde_json::to_vec(payload).map_err(|e| e.to_string())?;
|
||||
let payload_b64 = URL_SAFE_NO_PAD.encode(payload_json);
|
||||
let signature = sign_bytes(secret, payload_b64.as_bytes())?;
|
||||
Ok(format!(
|
||||
"{}.{}",
|
||||
payload_b64,
|
||||
URL_SAFE_NO_PAD.encode(signature)
|
||||
))
|
||||
}
|
||||
|
||||
fn sign_bytes(secret: &[u8], bytes: &[u8]) -> Result<Vec<u8>, String> {
|
||||
let mut mac = HmacSha256::new_from_slice(secret).map_err(|e| e.to_string())?;
|
||||
mac.update(bytes);
|
||||
Ok(mac.finalize().into_bytes().to_vec())
|
||||
}
|
||||
|
||||
pub fn hex_bytes(bytes: &[u8]) -> String {
|
||||
const HEX: &[u8; 16] = b"0123456789abcdef";
|
||||
let mut out = String::with_capacity(bytes.len() * 2);
|
||||
for &byte in bytes {
|
||||
out.push(HEX[(byte >> 4) as usize] as char);
|
||||
out.push(HEX[(byte & 0x0f) as usize] as char);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn token_hash_is_stable_hex_sha256() {
|
||||
assert_eq!(
|
||||
token_hash("native-token"),
|
||||
"d0c16f641bc0a0ee6b63ff88cec29756638d19590893c340a1ae36c9fae7b07f"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_lease_has_two_base64url_parts() {
|
||||
let lease = sign_path_lease(
|
||||
b"01234567890123456789012345678901",
|
||||
NativePathOperation::ValidateModel,
|
||||
"/tmp/model.gguf".to_string(),
|
||||
NativePathKind::Model,
|
||||
NativePathType::File,
|
||||
NativePathSourceKind::Dialog,
|
||||
"token",
|
||||
"model.gguf".to_string(),
|
||||
Some(123),
|
||||
Some(456),
|
||||
)
|
||||
.unwrap();
|
||||
let parts: Vec<&str> = lease.native_path_lease.split('.').collect();
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert!(!parts[0].contains('='));
|
||||
assert!(!parts[1].contains('='));
|
||||
}
|
||||
}
|
||||
476
studio/src-tauri/src/native_intents.rs
Normal file
476
studio/src-tauri/src/native_intents.rs
Normal file
|
|
@ -0,0 +1,476 @@
|
|||
use crate::native_backend_lease::{
|
||||
encode_secret_env, now_ms, random_token, sign_path_lease, NativePathKind,
|
||||
NativePathLeaseResponse, NativePathOperation, NativePathSourceKind, NativePathType,
|
||||
};
|
||||
use crate::native_path_policy::{
|
||||
classify_artifact_path, classify_native_model_path, reveal_target, ClassifiedPath,
|
||||
NativeArtifactKind,
|
||||
};
|
||||
use serde::Serialize;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
use tauri::{AppHandle, WebviewWindow};
|
||||
use tauri_plugin_dialog::DialogExt;
|
||||
|
||||
const TOKEN_TTL: Duration = Duration::from_secs(15 * 60);
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct NativePathEntry {
|
||||
token: String,
|
||||
canonical_path: PathBuf,
|
||||
validation_policy: NativePathValidationPolicy,
|
||||
path_kind: NativePathKind,
|
||||
path_type: NativePathType,
|
||||
source_kind: NativePathSourceKind,
|
||||
allowed_operations: Vec<NativePathOperation>,
|
||||
display_label: String,
|
||||
expires_at_ms: u64,
|
||||
size_bytes: Option<u64>,
|
||||
modified_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum NativePathValidationPolicy {
|
||||
Model,
|
||||
Artifact(NativeArtifactKind),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NativePathRef {
|
||||
token: String,
|
||||
kind: NativePathKind,
|
||||
display_label: String,
|
||||
allowed_operations: Vec<NativePathOperation>,
|
||||
expires_at_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NativeIntent {
|
||||
id: String,
|
||||
kind: NativePathKind,
|
||||
source_kind: NativePathSourceKind,
|
||||
path: NativePathRef,
|
||||
display_label: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct NativeIntakeInner {
|
||||
tokens: HashMap<String, NativePathEntry>,
|
||||
queued_intents: VecDeque<NativeIntent>,
|
||||
}
|
||||
|
||||
pub struct NativeIntakeState {
|
||||
inner: Mutex<NativeIntakeInner>,
|
||||
lease_secret: Vec<u8>,
|
||||
}
|
||||
|
||||
pub fn new_native_intake_state() -> NativeIntakeState {
|
||||
NativeIntakeState {
|
||||
inner: Mutex::new(NativeIntakeInner::default()),
|
||||
lease_secret: crate::native_backend_lease::new_lease_secret(),
|
||||
}
|
||||
}
|
||||
|
||||
impl NativeIntakeState {
|
||||
pub fn lease_secret_env(&self) -> String {
|
||||
encode_secret_env(&self.lease_secret)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn enqueue_model_path(
|
||||
&self,
|
||||
path: impl AsRef<Path>,
|
||||
source_kind: NativePathSourceKind,
|
||||
) -> Result<NativeIntent, String> {
|
||||
let intent = self.register_model_path(path, source_kind)?;
|
||||
let mut inner = self.inner.lock().map_err(|e| e.to_string())?;
|
||||
inner.queued_intents.push_back(intent.clone());
|
||||
Ok(intent)
|
||||
}
|
||||
|
||||
fn register_model_path(
|
||||
&self,
|
||||
path: impl AsRef<Path>,
|
||||
source_kind: NativePathSourceKind,
|
||||
) -> Result<NativeIntent, String> {
|
||||
let classified = classify_native_model_path(path.as_ref())?;
|
||||
self.register_classified_path(classified, source_kind, NativePathValidationPolicy::Model)
|
||||
}
|
||||
|
||||
fn register_artifact(
|
||||
&self,
|
||||
kind: NativeArtifactKind,
|
||||
path: impl AsRef<Path>,
|
||||
) -> Result<NativePathRef, String> {
|
||||
let classified = classify_artifact_path(kind, path.as_ref())?;
|
||||
let entry = self.insert_entry(
|
||||
classified,
|
||||
NativePathSourceKind::Artifact,
|
||||
NativePathValidationPolicy::Artifact(kind),
|
||||
)?;
|
||||
Ok(entry.to_ref())
|
||||
}
|
||||
|
||||
fn register_classified_path(
|
||||
&self,
|
||||
classified: ClassifiedPath,
|
||||
source_kind: NativePathSourceKind,
|
||||
validation_policy: NativePathValidationPolicy,
|
||||
) -> Result<NativeIntent, String> {
|
||||
let entry = self.insert_entry(classified, source_kind, validation_policy)?;
|
||||
Ok(NativeIntent {
|
||||
id: random_token("intent_"),
|
||||
kind: entry.path_kind,
|
||||
source_kind,
|
||||
path: entry.to_ref(),
|
||||
display_label: entry.display_label.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn insert_entry(
|
||||
&self,
|
||||
classified: ClassifiedPath,
|
||||
source_kind: NativePathSourceKind,
|
||||
validation_policy: NativePathValidationPolicy,
|
||||
) -> Result<NativePathEntry, String> {
|
||||
let token = random_token("path_");
|
||||
let expires_at_ms = now_ms() + TOKEN_TTL.as_millis() as u64;
|
||||
let entry = NativePathEntry {
|
||||
token: token.clone(),
|
||||
canonical_path: classified.canonical_path,
|
||||
validation_policy,
|
||||
path_kind: classified.path_kind,
|
||||
path_type: classified.path_type,
|
||||
source_kind,
|
||||
allowed_operations: classified.allowed_operations,
|
||||
display_label: classified.display_label,
|
||||
expires_at_ms,
|
||||
size_bytes: classified.size_bytes,
|
||||
modified_ms: classified.modified_ms,
|
||||
};
|
||||
let mut inner = self.inner.lock().map_err(|e| e.to_string())?;
|
||||
inner.tokens.insert(token, entry.clone());
|
||||
Ok(entry)
|
||||
}
|
||||
|
||||
fn drain_intents(&self) -> Result<Vec<NativeIntent>, String> {
|
||||
let mut inner = self.inner.lock().map_err(|e| e.to_string())?;
|
||||
prune_expired(&mut inner);
|
||||
Ok(inner.queued_intents.drain(..).collect())
|
||||
}
|
||||
|
||||
fn entry_for_operation(
|
||||
&self,
|
||||
token: &str,
|
||||
operation: NativePathOperation,
|
||||
) -> Result<NativePathEntry, String> {
|
||||
let mut inner = self.inner.lock().map_err(|e| e.to_string())?;
|
||||
prune_expired(&mut inner);
|
||||
let entry = inner
|
||||
.tokens
|
||||
.get(token)
|
||||
.ok_or_else(|| "Native path token is unavailable or expired.".to_string())?;
|
||||
if !entry.allowed_operations.contains(&operation) {
|
||||
return Err("Native path token does not allow this operation.".to_string());
|
||||
}
|
||||
Ok(entry.clone())
|
||||
}
|
||||
|
||||
fn sign_grant(
|
||||
&self,
|
||||
token: &str,
|
||||
operation: NativePathOperation,
|
||||
) -> Result<NativePathLeaseResponse, String> {
|
||||
let entry = self.entry_for_operation(token, operation)?;
|
||||
validate_entry_path(&entry, operation)?;
|
||||
sign_path_lease(
|
||||
&self.lease_secret,
|
||||
operation,
|
||||
entry.canonical_path.to_string_lossy().to_string(),
|
||||
entry.path_kind,
|
||||
entry.path_type,
|
||||
entry.source_kind,
|
||||
&entry.token,
|
||||
entry.display_label,
|
||||
entry.size_bytes,
|
||||
entry.modified_ms,
|
||||
)
|
||||
}
|
||||
|
||||
fn path_for_operation(
|
||||
&self,
|
||||
token: &str,
|
||||
operation: NativePathOperation,
|
||||
) -> Result<NativePathEntry, String> {
|
||||
let entry = self.entry_for_operation(token, operation)?;
|
||||
validate_entry_path(&entry, operation)?;
|
||||
Ok(entry)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_entry_path(
|
||||
entry: &NativePathEntry,
|
||||
operation: NativePathOperation,
|
||||
) -> Result<(), String> {
|
||||
let classified = match entry.validation_policy {
|
||||
NativePathValidationPolicy::Model => classify_native_model_path(&entry.canonical_path)?,
|
||||
NativePathValidationPolicy::Artifact(kind) => {
|
||||
classify_artifact_path(kind, &entry.canonical_path)?
|
||||
}
|
||||
};
|
||||
let check_fingerprint = !matches!(
|
||||
operation,
|
||||
NativePathOperation::Reveal | NativePathOperation::Open
|
||||
);
|
||||
if classified.canonical_path != entry.canonical_path
|
||||
|| classified.path_kind != entry.path_kind
|
||||
|| classified.path_type != entry.path_type
|
||||
|| !classified.allowed_operations.contains(&operation)
|
||||
|| (check_fingerprint && classified.size_bytes != entry.size_bytes)
|
||||
|| (check_fingerprint && classified.modified_ms != entry.modified_ms)
|
||||
{
|
||||
return Err("Native path changed after it was selected.".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl NativePathEntry {
|
||||
fn to_ref(&self) -> NativePathRef {
|
||||
NativePathRef {
|
||||
token: self.token.clone(),
|
||||
kind: self.path_kind,
|
||||
display_label: self.display_label.clone(),
|
||||
allowed_operations: self.allowed_operations.clone(),
|
||||
expires_at_ms: self.expires_at_ms,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn prune_expired(inner: &mut NativeIntakeInner) {
|
||||
let now = now_ms();
|
||||
inner.tokens.retain(|_, entry| entry.expires_at_ms > now);
|
||||
inner
|
||||
.queued_intents
|
||||
.retain(|intent| intent.path.expires_at_ms > now);
|
||||
}
|
||||
|
||||
fn ensure_main_window(window: &WebviewWindow) -> Result<(), String> {
|
||||
if window.label() == "main" {
|
||||
Ok(())
|
||||
} else {
|
||||
Err("Native path commands are only available to the main window.".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn drain_native_intents(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
) -> Result<Vec<NativeIntent>, String> {
|
||||
ensure_main_window(&window)?;
|
||||
state.drain_intents()
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn register_native_model_path(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
path: String,
|
||||
) -> Result<NativeIntent, String> {
|
||||
ensure_main_window(&window)?;
|
||||
state.register_model_path(path, NativePathSourceKind::Drop)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn pick_native_model(
|
||||
window: WebviewWindow,
|
||||
app: AppHandle,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
) -> Result<Option<NativeIntent>, String> {
|
||||
ensure_main_window(&window)?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
app.dialog()
|
||||
.file()
|
||||
.set_title("Choose a GGUF model")
|
||||
.add_filter("GGUF models", &["gguf"])
|
||||
.pick_file(move |path| {
|
||||
let _ = tx.send(path);
|
||||
});
|
||||
let Some(file_path) = rx.await.map_err(|_| "Dialog closed".to_string())? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let path = file_path
|
||||
.into_path()
|
||||
.map_err(|_| "Only local filesystem model paths are supported.".to_string())?;
|
||||
state
|
||||
.register_model_path(path, NativePathSourceKind::Dialog)
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn consume_native_path_token(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
token: String,
|
||||
operation: NativePathOperation,
|
||||
) -> Result<NativePathLeaseResponse, String> {
|
||||
ensure_main_window(&window)?;
|
||||
match operation {
|
||||
NativePathOperation::Reveal | NativePathOperation::Open => {
|
||||
Err("Reveal/Open do not use backend path grants.".to_string())
|
||||
}
|
||||
_ => state.sign_grant(&token, operation),
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn register_artifact_path(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
kind: NativeArtifactKind,
|
||||
path: String,
|
||||
) -> Result<NativePathRef, String> {
|
||||
ensure_main_window(&window)?;
|
||||
state.register_artifact(kind, path)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn reveal_path_token(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
token: String,
|
||||
) -> Result<(), String> {
|
||||
ensure_main_window(&window)?;
|
||||
let entry = state.path_for_operation(&token, NativePathOperation::Reveal)?;
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
if entry.canonical_path.is_file() {
|
||||
return std::process::Command::new("open")
|
||||
.arg("-R")
|
||||
.arg(&entry.canonical_path)
|
||||
.spawn()
|
||||
.map(|_| ())
|
||||
.map_err(|e| format!("Failed to reveal path: {e}"));
|
||||
}
|
||||
}
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
if entry.canonical_path.is_file() {
|
||||
let mut select_arg = std::ffi::OsString::from("/select,");
|
||||
select_arg.push(entry.canonical_path.as_os_str());
|
||||
return std::process::Command::new("explorer")
|
||||
.arg(select_arg)
|
||||
.spawn()
|
||||
.map(|_| ())
|
||||
.map_err(|e| format!("Failed to reveal path: {e}"));
|
||||
}
|
||||
}
|
||||
let target = reveal_target(&entry.canonical_path);
|
||||
open::that_detached(target).map_err(|e| format!("Failed to reveal path: {e}"))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn open_path_token(
|
||||
window: WebviewWindow,
|
||||
state: tauri::State<'_, NativeIntakeState>,
|
||||
token: String,
|
||||
) -> Result<(), String> {
|
||||
ensure_main_window(&window)?;
|
||||
let entry = state.path_for_operation(&token, NativePathOperation::Open)?;
|
||||
open::that_detached(entry.canonical_path).map_err(|e| format!("Failed to open path: {e}"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_path(name: &str) -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!(
|
||||
"unsloth-native-intents-{name}-{}-{nanos}",
|
||||
std::process::id()
|
||||
))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_token_issues_distinct_validate_and_load_grants() {
|
||||
let state = new_native_intake_state();
|
||||
let path = temp_path("model").with_extension("gguf");
|
||||
fs::write(&path, b"gguf").unwrap();
|
||||
let intent = state
|
||||
.register_model_path(&path, NativePathSourceKind::Dialog)
|
||||
.unwrap();
|
||||
let validate = state
|
||||
.sign_grant(&intent.path.token, NativePathOperation::ValidateModel)
|
||||
.unwrap();
|
||||
let load = state
|
||||
.sign_grant(&intent.path.token, NativePathOperation::LoadModel)
|
||||
.unwrap();
|
||||
assert_ne!(validate.native_path_lease, load.native_path_lease);
|
||||
assert!(validate.display_label.ends_with(".gguf"));
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_token_rejects_dataset_operation() {
|
||||
let state = new_native_intake_state();
|
||||
let path = temp_path("model").with_extension("gguf");
|
||||
fs::write(&path, b"gguf").unwrap();
|
||||
let intent = state
|
||||
.register_model_path(&path, NativePathSourceKind::Dialog)
|
||||
.unwrap();
|
||||
let err = state
|
||||
.sign_grant(&intent.path.token, NativePathOperation::DatasetImport)
|
||||
.unwrap_err();
|
||||
assert!(err.contains("does not allow"));
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_token_revalidates_path_changes() {
|
||||
let state = new_native_intake_state();
|
||||
let path = temp_path("model").with_extension("gguf");
|
||||
fs::write(&path, b"gguf").unwrap();
|
||||
let intent = state
|
||||
.register_model_path(&path, NativePathSourceKind::Dialog)
|
||||
.unwrap();
|
||||
fs::write(&path, b"changed").unwrap();
|
||||
let err = state
|
||||
.sign_grant(&intent.path.token, NativePathOperation::ValidateModel)
|
||||
.unwrap_err();
|
||||
assert!(err.contains("changed"));
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn reveal_rejects_symlink_replacement() {
|
||||
use std::os::unix::fs::symlink;
|
||||
|
||||
let state = new_native_intake_state();
|
||||
let path = temp_path("model").with_extension("gguf");
|
||||
let target = temp_path("replacement").with_extension("gguf");
|
||||
fs::write(&path, b"gguf").unwrap();
|
||||
fs::write(&target, b"gguf").unwrap();
|
||||
let intent = state
|
||||
.register_model_path(&path, NativePathSourceKind::Dialog)
|
||||
.unwrap();
|
||||
fs::remove_file(&path).unwrap();
|
||||
symlink(&target, &path).unwrap();
|
||||
let err = state
|
||||
.path_for_operation(&intent.path.token, NativePathOperation::Reveal)
|
||||
.unwrap_err();
|
||||
assert!(err.contains("Symlink") || err.contains("changed"));
|
||||
let _ = fs::remove_file(path);
|
||||
let _ = fs::remove_file(target);
|
||||
}
|
||||
}
|
||||
291
studio/src-tauri/src/native_path_policy.rs
Normal file
291
studio/src-tauri/src/native_path_policy.rs
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
use crate::native_backend_lease::{NativePathKind, NativePathOperation, NativePathType};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum NativeArtifactKind {
|
||||
TrainingOutput,
|
||||
Export,
|
||||
DatasetUpload,
|
||||
RecipeArtifact,
|
||||
DiagnosticLog,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ClassifiedPath {
|
||||
pub canonical_path: PathBuf,
|
||||
pub path_kind: NativePathKind,
|
||||
pub path_type: NativePathType,
|
||||
pub allowed_operations: Vec<NativePathOperation>,
|
||||
pub display_label: String,
|
||||
pub size_bytes: Option<u64>,
|
||||
pub modified_ms: Option<u64>,
|
||||
}
|
||||
|
||||
pub fn classify_native_model_path(path: &Path) -> Result<ClassifiedPath, String> {
|
||||
let classified = classify_existing_path(path)?;
|
||||
if classified.path_type != NativePathType::File {
|
||||
return Err("Only GGUF model files are supported for native model intake.".to_string());
|
||||
}
|
||||
if !has_extension(&classified.canonical_path, "gguf") {
|
||||
return Err("Only .gguf model files are supported for native model intake.".to_string());
|
||||
}
|
||||
Ok(ClassifiedPath {
|
||||
path_kind: NativePathKind::Model,
|
||||
allowed_operations: vec![
|
||||
NativePathOperation::ValidateModel,
|
||||
NativePathOperation::LoadModel,
|
||||
NativePathOperation::Reveal,
|
||||
],
|
||||
..classified
|
||||
})
|
||||
}
|
||||
|
||||
pub fn classify_artifact_path(
|
||||
kind: NativeArtifactKind,
|
||||
path: &Path,
|
||||
) -> Result<ClassifiedPath, String> {
|
||||
let classified = classify_existing_path(path)?;
|
||||
ensure_artifact_root(kind, &classified.canonical_path)?;
|
||||
reject_sensitive_artifact(&classified.canonical_path)?;
|
||||
|
||||
let mut allowed_operations = vec![NativePathOperation::Reveal];
|
||||
if is_open_safe_artifact(&classified.canonical_path, classified.path_type) {
|
||||
allowed_operations.push(NativePathOperation::Open);
|
||||
}
|
||||
|
||||
Ok(ClassifiedPath {
|
||||
path_kind: NativePathKind::Artifact,
|
||||
allowed_operations,
|
||||
..classified
|
||||
})
|
||||
}
|
||||
|
||||
pub fn refresh_path_fingerprint(
|
||||
path: &Path,
|
||||
) -> Result<(NativePathType, Option<u64>, Option<u64>), String> {
|
||||
let metadata = fs::metadata(path).map_err(|e| format!("Path is no longer available: {e}"))?;
|
||||
let path_type = if metadata.is_file() {
|
||||
NativePathType::File
|
||||
} else if metadata.is_dir() {
|
||||
NativePathType::Directory
|
||||
} else {
|
||||
return Err("Special files are not supported.".to_string());
|
||||
};
|
||||
let size_bytes = metadata.is_file().then_some(metadata.len());
|
||||
let modified_ms = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_millis() as u64);
|
||||
Ok((path_type, size_bytes, modified_ms))
|
||||
}
|
||||
|
||||
pub fn reveal_target(path: &Path) -> PathBuf {
|
||||
if path.is_dir() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
path.parent().unwrap_or(path).to_path_buf()
|
||||
}
|
||||
}
|
||||
|
||||
fn sanitize_display_label(raw: &str) -> String {
|
||||
let cleaned: String = raw
|
||||
.chars()
|
||||
.map(|ch| if ch.is_control() { ' ' } else { ch })
|
||||
.collect();
|
||||
let trimmed = cleaned.trim();
|
||||
if trimmed.is_empty() {
|
||||
"Selected path".to_string()
|
||||
} else {
|
||||
trimmed.chars().take(160).collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_open_safe_artifact(path: &Path, path_type: NativePathType) -> bool {
|
||||
if path_type == NativePathType::Directory {
|
||||
return false;
|
||||
}
|
||||
let Some(ext) = path.extension().and_then(|ext| ext.to_str()) else {
|
||||
return false;
|
||||
};
|
||||
matches!(
|
||||
ext.to_ascii_lowercase().as_str(),
|
||||
"txt" | "log" | "json" | "jsonl" | "csv" | "tsv" | "parquet" | "md"
|
||||
)
|
||||
}
|
||||
|
||||
fn classify_existing_path(path: &Path) -> Result<ClassifiedPath, String> {
|
||||
reject_network_or_device_path(path)?;
|
||||
let symlink_metadata =
|
||||
fs::symlink_metadata(path).map_err(|e| format!("Path is not available: {e}"))?;
|
||||
if symlink_metadata.file_type().is_symlink() {
|
||||
return Err("Symlink paths are not supported for native intake.".to_string());
|
||||
}
|
||||
|
||||
let canonical_path = path
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Path could not be resolved: {e}"))?;
|
||||
reject_network_or_device_path(&canonical_path)?;
|
||||
let canonical_symlink_metadata = fs::symlink_metadata(&canonical_path)
|
||||
.map_err(|e| format!("Path is not available: {e}"))?;
|
||||
if canonical_symlink_metadata.file_type().is_symlink() {
|
||||
return Err("Symlink paths are not supported for native intake.".to_string());
|
||||
}
|
||||
let (path_type, size_bytes, modified_ms) = refresh_path_fingerprint(&canonical_path)?;
|
||||
let display_label = sanitize_display_label(
|
||||
canonical_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.unwrap_or("Selected path"),
|
||||
);
|
||||
|
||||
Ok(ClassifiedPath {
|
||||
canonical_path,
|
||||
path_kind: NativePathKind::Artifact,
|
||||
path_type,
|
||||
allowed_operations: vec![NativePathOperation::Reveal],
|
||||
display_label,
|
||||
size_bytes,
|
||||
modified_ms,
|
||||
})
|
||||
}
|
||||
|
||||
fn ensure_artifact_root(kind: NativeArtifactKind, canonical_path: &Path) -> Result<(), String> {
|
||||
let Some(home) = dirs::home_dir() else {
|
||||
return Err("Could not determine home directory.".to_string());
|
||||
};
|
||||
let studio = home.join(".unsloth").join("studio");
|
||||
let allowed_root = match kind {
|
||||
NativeArtifactKind::TrainingOutput => studio.join("outputs"),
|
||||
NativeArtifactKind::Export => studio.join("exports"),
|
||||
NativeArtifactKind::DatasetUpload => studio.join("assets").join("datasets").join("uploads"),
|
||||
NativeArtifactKind::RecipeArtifact => {
|
||||
studio.join("assets").join("datasets").join("recipes")
|
||||
}
|
||||
NativeArtifactKind::DiagnosticLog => studio.join("logs"),
|
||||
};
|
||||
let root = allowed_root
|
||||
.canonicalize()
|
||||
.map_err(|_| "Artifact root is not available.".to_string())?;
|
||||
if canonical_path == root || canonical_path.starts_with(&root) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err("Artifact path is outside the allowed artifact root.".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn reject_sensitive_artifact(path: &Path) -> Result<(), String> {
|
||||
let lowered = path.to_string_lossy().to_ascii_lowercase();
|
||||
for needle in [
|
||||
"/auth/",
|
||||
"\\auth\\",
|
||||
"/auth.db",
|
||||
"\\auth.db",
|
||||
"/studio.db",
|
||||
"\\studio.db",
|
||||
"/pid",
|
||||
"\\pid",
|
||||
] {
|
||||
if lowered.contains(needle) {
|
||||
return Err("Sensitive Studio state cannot be registered as an artifact.".to_string());
|
||||
}
|
||||
}
|
||||
if let Some(ext) = path.extension().and_then(|ext| ext.to_str()) {
|
||||
if matches!(
|
||||
ext.to_ascii_lowercase().as_str(),
|
||||
"exe" | "dll" | "dylib" | "so" | "sh" | "bash" | "zsh" | "ps1" | "bat" | "cmd"
|
||||
) {
|
||||
return Err(
|
||||
"Executable artifacts cannot be registered for native open/reveal.".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reject_network_or_device_path(path: &Path) -> Result<(), String> {
|
||||
let text = path.to_string_lossy();
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let normalized = text.replace('/', "\\").to_ascii_lowercase();
|
||||
if let Some(rest) = normalized.strip_prefix("\\\\?\\") {
|
||||
let bytes = rest.as_bytes();
|
||||
if !(bytes.len() >= 3
|
||||
&& bytes[0].is_ascii_alphabetic()
|
||||
&& bytes[1] == b':'
|
||||
&& bytes[2] == b'\\')
|
||||
{
|
||||
return Err("Network paths are not supported for native intake.".to_string());
|
||||
}
|
||||
} else if normalized.starts_with("\\\\") {
|
||||
return Err("Network paths are not supported for native intake.".to_string());
|
||||
}
|
||||
}
|
||||
#[cfg(unix)]
|
||||
{
|
||||
for root in ["/dev", "/proc", "/sys"] {
|
||||
if path.starts_with(root) {
|
||||
return Err("Device and virtual filesystem paths are not supported.".to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
if text.contains('\0') {
|
||||
return Err("Path contains invalid NUL characters.".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_extension(path: &Path, expected: &str) -> bool {
|
||||
path.extension()
|
||||
.and_then(|ext| ext.to_str())
|
||||
.map(|ext| ext.eq_ignore_ascii_case(expected))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_path(name: &str) -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!(
|
||||
"unsloth-native-policy-{name}-{}-{nanos}",
|
||||
std::process::id()
|
||||
))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gguf_model_allows_validate_load_reveal() {
|
||||
let path = temp_path("model").with_extension("gguf");
|
||||
fs::write(&path, b"gguf").unwrap();
|
||||
let classified = classify_native_model_path(&path).unwrap();
|
||||
assert_eq!(classified.path_kind, NativePathKind::Model);
|
||||
assert!(classified
|
||||
.allowed_operations
|
||||
.contains(&NativePathOperation::ValidateModel));
|
||||
assert!(classified
|
||||
.allowed_operations
|
||||
.contains(&NativePathOperation::LoadModel));
|
||||
assert!(classified
|
||||
.allowed_operations
|
||||
.contains(&NativePathOperation::Reveal));
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_gguf_model_is_rejected() {
|
||||
let path = temp_path("model").with_extension("txt");
|
||||
fs::write(&path, b"not gguf").unwrap();
|
||||
assert!(classify_native_model_path(&path).is_err());
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
}
|
||||
|
|
@ -7,7 +7,7 @@ use std::io::BufRead;
|
|||
use std::process::{Command, Stdio};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tauri::{AppHandle, Emitter};
|
||||
use tauri::{AppHandle, Emitter, Manager};
|
||||
|
||||
const MAX_LOG_LINES: usize = 1000;
|
||||
|
||||
|
|
@ -298,6 +298,13 @@ pub fn start_backend(
|
|||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
|
||||
if let Some(native_state) = app.try_state::<crate::native_intents::NativeIntakeState>() {
|
||||
cmd.env(
|
||||
crate::native_backend_lease::LEASE_SECRET_ENV,
|
||||
native_state.lease_secret_env(),
|
||||
);
|
||||
}
|
||||
|
||||
// AppImage sets LD_LIBRARY_PATH to its bundled libs, which breaks the spawned
|
||||
// Python process (wrong libpython/libz → "No module named encodings").
|
||||
// Only clear when running inside an AppImage — native .deb/.rpm installs may
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue