unsloth/studio/backend/models/responses.py
Daniel Han f08aef1804 Studio (#4237)
* Rebuild Studio branch on top of main

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

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

* Fix security and code quality issues for Studio PR #4237

- Validate models_dir query param against allowed directory roots
  to prevent path traversal in /api/models/local endpoint
- Replace string startswith() with Path.is_relative_to() for
  frontend path traversal check in serve_frontend
- Sanitize SSE error messages to not leak exception details to
  clients (4 locations in inference.py)
- Bind port-discovery socket to 127.0.0.1 instead of all interfaces
  in llama_cpp backend
- Import datasets_root and resolve_output_dir in embedding training
  function to fix NameError and use managed output directory
- Remove stale .gitignore entries for package-lock.json and test
  directories so tests can be tracked in version control
- Add venv-reexecution logic to ui CLI command matching the studio
  command behavior

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

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

* Move models_dir path validation before try/except block

The HTTPException(403) was inside the try/except Exception handler,
so it would be caught and re-raised as a 500. Moving the validation
before the try block ensures the 403 is returned directly and also
makes the control flow clearer for static analysis (path is validated
before any filesystem operations).

* Use os.path.realpath + startswith for models_dir validation

CodeQL py/path-injection does not recognize Path.is_relative_to() as
a sanitizer. Switched to os.path.realpath + str.startswith which is
a recognized sanitizer pattern in CodeQL's taint analysis. The
startswith check uses root_str + os.sep to prevent prefix collisions
(e.g. /app/models_evil matching /app/models).

* Never pass user input to Path constructor in models_dir validation

CodeQL traces taint through Path(resolved) even after a startswith
barrier guard. Fix: the user-supplied models_dir is only used as a
string for comparison against allowed roots. The Path object passed
to _scan_models_dir comes from the trusted allowed_roots list, not
from user input. This fully breaks the taint chain.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-12 03:36:19 -07:00

67 lines
2.4 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Pydantic response schemas for endpoints that previously returned raw dicts.
These are small response models for training and model management routes.
"""
from pydantic import BaseModel, Field
from typing import Optional, List
# --- Training route response models ---
class TrainingStopResponse(BaseModel):
"""Response for stopping a training job"""
status: str = Field(..., description = "Current status: 'stopped' or 'idle'")
message: str = Field(..., description = "Human-readable status message")
class TrainingMetricsResponse(BaseModel):
"""Response for training metrics history"""
loss_history: List[float] = Field(
default_factory = list, description = "Loss values per step"
)
lr_history: List[float] = Field(
default_factory = list, description = "Learning rate per step"
)
step_history: List[int] = Field(default_factory = list, description = "Step numbers")
grad_norm_history: List[float] = Field(
default_factory = list, description = "Gradient norm values"
)
grad_norm_step_history: List[int] = Field(
default_factory = list, description = "Step numbers for gradient norm values"
)
current_loss: Optional[float] = Field(None, description = "Most recent loss value")
current_lr: Optional[float] = Field(None, description = "Most recent learning rate")
current_step: Optional[int] = Field(None, description = "Most recent step number")
# --- Model management route response models ---
class LoRABaseModelResponse(BaseModel):
"""Response for getting a LoRA's base model"""
lora_path: str = Field(..., description = "Path to the LoRA adapter")
base_model: str = Field(..., description = "Base model identifier")
class VisionCheckResponse(BaseModel):
"""Response for checking if a model is a vision model"""
model_name: str = Field(..., description = "Model identifier")
is_vision: bool = Field(..., description = "Whether the model is a vision model")
class EmbeddingCheckResponse(BaseModel):
"""Response for checking if a model is an embedding model"""
model_name: str = Field(..., description = "Model identifier")
is_embedding: bool = Field(
..., description = "Whether the model is an embedding/sentence-transformer model"
)