From 1e9d19126bc3ed70d6430d165e02109da2de79f8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 15 Mar 2026 07:40:39 +0000
Subject: [PATCH] studio: fix P1 issues from PR review comments

1. n_gpu_layers kwarg: accept (and ignore) in load_model signature
   so callers like llm_assist.py don't get TypeError

2. mmproj exclusion: filter out mmproj files in _find_smallest_fitting_variant
   so fallback doesn't pick a tiny vision projection as the "model"

3. Shard preservation after fallback: re-discover shards for the
   fallback variant instead of resetting to empty list, so split
   GGUFs download all shards

4. Orphan cleanup safety: only kill llama-server processes whose
   cmdline contains ".unsloth/", avoiding termination of unrelated
   llama-server instances on the same machine

5. Path expression sanitization: validate repo_id format before using
   it in cache directory lookups
---
 studio/backend/core/inference/llama_cpp.py | 43 ++++++++++++++++++----
 studio/backend/routes/models.py            |  5 +++
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index fafefbb7a7..94beefa3aa 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -307,7 +307,10 @@ class LlamaCppBackend:
             from huggingface_hub import get_paths_info, list_repo_files
 
             files = list_repo_files(hf_repo, token = hf_token)
-            gguf_files = [f for f in files if f.endswith(".gguf")]
+            gguf_files = [
+                f for f in files
+                if f.endswith(".gguf") and "mmproj" not in f.lower()
+            ]
             if not gguf_files:
                 return None
 
@@ -387,6 +390,7 @@ class LlamaCppBackend:
         is_vision: bool = False,
         n_ctx: int = 4096,
         n_threads: Optional[int] = None,
+        n_gpu_layers: Optional[int] = None,  # Accepted for caller compat, unused
     ) -> bool:
         """
         Start llama-server with a GGUF model.
@@ -516,12 +520,25 @@ class LlamaCppBackend:
                                 hf_token,
                             )
                             if smaller:
+                                fallback_file, fallback_size = smaller
                                 logger.info(
                                     f"Selected variant too large ({total_gb:.1f} GB), "
-                                    f"falling back to {smaller[0]} ({smaller[1] / (1024**3):.1f} GB)"
+                                    f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
                                 )
-                                gguf_filename = smaller[0]
-                                gguf_extra_shards = []
+                                gguf_filename = fallback_file
+                                # Re-discover shards for the fallback variant
+                                import re as _re
+                                _shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
+                                _m = _shard_pat.match(gguf_filename)
+                                _prefix = _m.group(1) if _m else None
+                                if _prefix:
+                                    gguf_extra_shards = sorted(
+                                        f for f in all_gguf_files
+                                        if f.startswith(_prefix) and f != gguf_filename
+                                        and "mmproj" not in f.lower()
+                                    )
+                                else:
+                                    gguf_extra_shards = []
                             else:
                                 raise RuntimeError(
                                     f"Not enough disk space to download any variant. "
@@ -727,13 +744,18 @@ class LlamaCppBackend:
 
     @staticmethod
     def _kill_orphaned_servers():
-        """Kill any orphaned llama-server processes from previous studio runs."""
+        """Kill orphaned llama-server processes started by studio.
+
+        Only kills processes whose binary lives under ~/.unsloth/llama.cpp/
+        to avoid terminating unrelated llama-server instances on the machine.
+        """
         import os
         import signal
 
         try:
+            # Use pgrep with full command match to identify studio-managed servers
             result = subprocess.run(
-                ["pgrep", "-f", "llama-server"],
+                ["pgrep", "-a", "-f", "llama-server"],
                 capture_output = True,
                 text = True,
                 timeout = 5,
@@ -741,9 +763,16 @@ class LlamaCppBackend:
             if result.returncode != 0:
                 return
             for line in result.stdout.strip().splitlines():
-                pid = int(line.strip())
+                parts = line.strip().split(None, 1)
+                if len(parts) < 2:
+                    continue
+                pid = int(parts[0])
+                cmdline = parts[1]
                 if pid == os.getpid():
                     continue
+                # Only kill if it's a studio-managed server (lives under .unsloth/)
+                if ".unsloth/" not in cmdline and "unsloth" not in cmdline.lower():
+                    continue
                 try:
                     os.kill(pid, signal.SIGKILL)
                     logger.info(f"Killed orphaned llama-server process (pid={pid})")
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 34069cffc6..1dee85ed82 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -561,8 +561,13 @@ async def get_gguf_variants(
         # case-insensitive match.
         cached_files: set = set()
         try:
+            import re as _re
             from huggingface_hub import constants as hf_constants
 
+            # Sanitize repo_id: must be "owner/name" with safe chars only
+            if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id):
+                raise ValueError(f"Invalid repo_id format: {repo_id}")
+
             cache_dir = Path(hf_constants.HF_HUB_CACHE)
             target = f"models--{repo_id.replace('/', '--')}".lower()
             for entry in cache_dir.iterdir():