diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index db59d573a9..cdc28d9560 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -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 diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index afd10b02d1..44459e88c5 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -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 diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 206dbd6dbb..82de925592 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -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: diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index 895b112e85..df3bf27c16 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -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 diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index ae8224d424..6790dbc802 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 428ecaa1c3..5562820f49 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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: diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 24d06f9f62..fe8d277ac0 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -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 [ diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5437645797..5642faa189 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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 diff --git a/studio/backend/loggers/config.py b/studio/backend/loggers/config.py index 4c0d8ade28..4a27f13d38 100644 --- a/studio/backend/loggers/config.py +++ b/studio/backend/loggers/config.py @@ -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"), diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index 3add92ea1e..ddd404cdf3 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -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", 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: "" + 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: "" + if str(k).replace("_", "").lower() == "nativepathlease" + else filter_value(v) + for k, v in event_dict.items() + } def get_logger(name: str) -> structlog.BoundLogger: diff --git a/studio/backend/main.py b/studio/backend/main.py index aec335fad3..0958094ff0 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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(), } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 027f3313a1..c0bb5b53a2 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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" ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 53dd851666..592da831e4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 + 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) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 755314ca3a..fdb1ab4520 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -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) diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py index 274d9beb48..099c5fa3a5 100644 --- a/studio/backend/utils/hardware/nvidia.py +++ b/studio/backend/utils/hardware/nvidia.py @@ -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: diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 7bfc74235f..16f6d21edb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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(), ) diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py new file mode 100644 index 0000000000..a69dfab532 --- /dev/null +++ b/studio/backend/utils/native_path_leases.py @@ -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, "") + 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 diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index ebb20b85da..17af40f663 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -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: diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 00240f1e69..3ed9bda827 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -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 diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index c3a1c7dab0..62e78b809a 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -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 ? ( <> + {children} ) : ( diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 1ff946dd46..795bcb6d08 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -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; 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({ )} + {onPickLocalModel ? ( +
+ +
+ ) : null} {hasSelection && onEject ? (
+ + +
+ ); +} diff --git a/studio/frontend/src/features/native-intents/components/native-model-drop-overlay.tsx b/studio/frontend/src/features/native-intents/components/native-model-drop-overlay.tsx new file mode 100644 index 0000000000..725fb14259 --- /dev/null +++ b/studio/frontend/src/features/native-intents/components/native-model-drop-overlay.tsx @@ -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 ( +
+
+
+
+
+
+ {title} +
+
+ {description} +
+
+
+
+ ); +} diff --git a/studio/frontend/src/features/native-intents/native-intent-drain.tsx b/studio/frontend/src/features/native-intents/native-intent-drain.tsx new file mode 100644 index 0000000000..109e8e27f9 --- /dev/null +++ b/studio/frontend/src/features/native-intents/native-intent-drain.tsx @@ -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; +} diff --git a/studio/frontend/src/features/native-intents/store.ts b/studio/frontend/src/features/native-intents/store.ts new file mode 100644 index 0000000000..a9cd9e7966 --- /dev/null +++ b/studio/frontend/src/features/native-intents/store.ts @@ -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((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 }); + }, +})); diff --git a/studio/frontend/src/features/native-intents/types.ts b/studio/frontend/src/features/native-intents/types.ts new file mode 100644 index 0000000000..268638559a --- /dev/null +++ b/studio/frontend/src/features/native-intents/types.ts @@ -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; +} diff --git a/studio/frontend/src/features/native-intents/use-native-dialogs.ts b/studio/frontend/src/features/native-intents/use-native-dialogs.ts new file mode 100644 index 0000000000..45ca81343a --- /dev/null +++ b/studio/frontend/src/features/native-intents/use-native-dialogs.ts @@ -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; +} + +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]); +} diff --git a/studio/frontend/src/features/native-intents/use-native-drop.ts b/studio/frontend/src/features/native-intents/use-native-drop.ts new file mode 100644 index 0000000000..5d35511c10 --- /dev/null +++ b/studio/frontend/src/features/native-intents/use-native-drop.ts @@ -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; +} + +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({ 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; +} diff --git a/studio/frontend/src/features/native-intents/use-native-readiness.ts b/studio/frontend/src/features/native-intents/use-native-readiness.ts new file mode 100644 index 0000000000..704ae9bc96 --- /dev/null +++ b/studio/frontend/src/features/native-intents/use-native-readiness.ts @@ -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 | 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; +} diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 25d950fd3e..3fd1e6af66 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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, diff --git a/studio/src-tauri/Cargo.lock b/studio/src-tauri/Cargo.lock index e491ee534a..dd57e55dc5 100644 --- a/studio/src-tauri/Cargo.lock +++ b/studio/src-tauri/Cargo.lock @@ -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", diff --git a/studio/src-tauri/Cargo.toml b/studio/src-tauri/Cargo.toml index 269bb7a92a..d3cfe49667 100644 --- a/studio/src-tauri/Cargo.toml +++ b/studio/src-tauri/Cargo.toml @@ -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] diff --git a/studio/src-tauri/src/diagnostics/redaction.rs b/studio/src-tauri/src/diagnostics/redaction.rs index 5c5589a494..0c7a92944a 100644 --- a/studio/src-tauri/src/diagnostics/redaction.rs +++ b/studio/src-tauri/src/diagnostics/redaction.rs @@ -20,6 +20,12 @@ pub(crate) fn redact_text(text: &str, report: &mut RedactionReport) -> String { out = replace_regex(cookie_re(), &out, "$1: ", report); out = replace_regex(token_re(), &out, "", report); out = replace_regex(env_secret_re(), &out, "$1=", report); + out = replace_regex( + native_path_lease_re(), + &out, + "$1", + report, + ); out = replace_known_paths(&out, report); out = replace_regex(windows_studio_re(), &out, "", 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 = 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 = 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=")); assert!(redacted.contains("https://@example.com/path")); assert!(!redacted.contains("alex@example.com")); assert!(redacted.contains("")); diff --git a/studio/src-tauri/src/main.rs b/studio/src-tauri/src/main.rs index b4e718d867..5aeaa56497 100644 --- a/studio/src-tauri/src/main.rs +++ b/studio/src-tauri/src/main.rs @@ -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"))] diff --git a/studio/src-tauri/src/native_backend_lease.rs b/studio/src-tauri/src/native_backend_lease.rs new file mode 100644 index 0000000000..8c94a02f86 --- /dev/null +++ b/studio/src-tauri/src/native_backend_lease.rs @@ -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; + +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, + pub modified_ms: Option, +} + +#[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 { + 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, + modified_ms: Option, +) -> Result { + 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 { + 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, 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('=')); + } +} diff --git a/studio/src-tauri/src/native_intents.rs b/studio/src-tauri/src/native_intents.rs new file mode 100644 index 0000000000..30e11aa310 --- /dev/null +++ b/studio/src-tauri/src/native_intents.rs @@ -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, + display_label: String, + expires_at_ms: u64, + size_bytes: Option, + modified_ms: Option, +} + +#[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, + 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, + queued_intents: VecDeque, +} + +pub struct NativeIntakeState { + inner: Mutex, + lease_secret: Vec, +} + +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, + source_kind: NativePathSourceKind, + ) -> Result { + 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, + source_kind: NativePathSourceKind, + ) -> Result { + 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, + ) -> Result { + 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 { + 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 { + 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, 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 { + 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 { + 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 { + 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, 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 { + 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, 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 { + 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 { + 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); + } +} diff --git a/studio/src-tauri/src/native_path_policy.rs b/studio/src-tauri/src/native_path_policy.rs new file mode 100644 index 0000000000..b2ebb34621 --- /dev/null +++ b/studio/src-tauri/src/native_path_policy.rs @@ -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, + pub display_label: String, + pub size_bytes: Option, + pub modified_ms: Option, +} + +pub fn classify_native_model_path(path: &Path) -> Result { + 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 { + 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, Option), 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 { + 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); + } +} diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index ed2659c0d9..9e0cf1c17f 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -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::() { + 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