From a932294627176b742ff425817f320c60c28d284d Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 14 May 2026 07:24:20 -0500 Subject: [PATCH 01/50] MLX training support for Studio on Apple Silicon (#5340) * mlx fixes * Fix studio integration, local dataset files, chat templates without the torch gpu imports * pass grad norm in mlx worker * fix(studio): pass MLX grad clipping settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * mlx: update grad value * fix(mlx): address ci and clipping review * fix backward compatibility and CI tests * unsloth local is mlx function * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * dont reference runtime * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio mlx: hardcode value clipping, drop max_grad_value from frontend Simplifies the MLX grad-clipping plumbing now that we are standardising on elementwise value clipping at [-5, 5] for the compiled MLX path and norm clipping disabled. The MLX worker no longer reads max_grad_norm / max_grad_value from the request; both are pinned in one place. Frontend stops sending the field at all, and the TypeScript request type drops it to match. Non-MLX (CUDA/AMD/Intel) is untouched and continues to pick up HF TrainingArguments' default max_grad_norm = 1.0. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../backend/core/inference/mlx_inference.py | 4 +- studio/backend/core/training/training.py | 2 + studio/backend/core/training/worker.py | 87 +++++++++++++++++-- studio/backend/models/training.py | 13 +++ studio/backend/routes/training.py | 2 + .../tests/test_mlx_inference_backend.py | 9 +- .../tests/test_training_raw_support.py | 44 ++++++++++ .../backend/utils/datasets/chat_templates.py | 25 +++++- .../src/features/training/api/mappers.ts | 1 + .../src/features/training/types/api.ts | 1 + tests/studio/run_real_mlx_smoke.py | 6 +- tests/studio/test_is_mlx_dispatch_gate.py | 66 +++++++++----- .../test_mlx_training_worker_behaviors.py | 8 +- unsloth/__init__.py | 37 +++++--- unsloth/chat_templates.py | 14 ++- unsloth/device_type.py | 23 ++++- unsloth/kernels/utils.py | 9 ++ 17 files changed, 295 insertions(+), 56 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index ee7c9bbd51..e7bce2d33e 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -116,11 +116,11 @@ class MLXInferenceBackend: ) try: - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as e: raise ImportError( "Unsloth: MLX inference requires unsloth-zoo with the MLX modules " - "(unsloth_zoo.mlx_loader). Reinstall via install.sh on Apple Silicon." + "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e model, tokenizer_or_processor = FastMLXModel.from_pretrained( diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 549d733252..556545e680 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -214,6 +214,8 @@ class TrainingBackend: "max_steps": kwargs.get("max_steps", 0), "save_steps": kwargs.get("save_steps", 0), "weight_decay": kwargs.get("weight_decay", 0.001), + "max_grad_norm": kwargs.get("max_grad_norm", 0.0), + "max_grad_value": kwargs.get("max_grad_value"), "random_seed": kwargs.get("random_seed", 3407), "packing": kwargs.get("packing", False), "optim": kwargs.get("optim", "adamw_8bit"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ef5cafb175..8cd1d1f9c8 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -417,6 +417,55 @@ def _normalize_mlx_studio_scheduler(value): return raw +def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: + """Resolve Studio local dataset uploads without importing the GPU trainer.""" + from utils.paths import resolve_dataset_path + + all_files: list[str] = [] + for dataset_file in file_paths or []: + file_path = ( + dataset_file + if os.path.isabs(dataset_file) + else str(resolve_dataset_path(dataset_file)) + ) + file_path_obj = Path(file_path) + + if file_path_obj.is_dir(): + parquet_dir = ( + file_path_obj / "parquet-files" + if (file_path_obj / "parquet-files").exists() + else file_path_obj + ) + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if parquet_files: + all_files.extend(str(p) for p in parquet_files) + continue + + candidates: list[Path] = [] + for ext in (".json", ".jsonl", ".csv", ".parquet"): + candidates.extend(sorted(file_path_obj.glob(f"*{ext}"))) + if candidates: + all_files.extend(str(c) for c in candidates) + continue + + raise ValueError(f"No supported data files in directory: {file_path_obj}") + + all_files.append(str(file_path_obj)) + + return all_files + + +def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: + first_ext = Path(files[0]).suffix.lower() + if first_ext in (".json", ".jsonl"): + return "json" + if first_ext == ".csv": + return "csv" + if first_ext == ".parquet": + return "parquet" + raise ValueError(f"Unsupported dataset format: {files[0]}") + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -442,8 +491,8 @@ def _run_mlx_training(event_queue, stop_queue, config): import mlx.core as mx try: - from unsloth_zoo.mlx_loader import FastMLXModel - from unsloth_zoo.mlx_trainer import ( + from unsloth_zoo.mlx.loader import FastMLXModel + from unsloth_zoo.mlx.trainer import ( MLXTrainer, MLXTrainingConfig, train_on_responses_only, @@ -451,7 +500,7 @@ def _run_mlx_training(event_queue, stop_queue, config): except ImportError as e: raise ImportError( "Unsloth: MLX training requires unsloth-zoo with the MLX modules " - "(unsloth_zoo.mlx_loader / unsloth_zoo.mlx_trainer). Reinstall via " + "(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via " "install.sh on Apple Silicon." ) from e from datasets import load_dataset @@ -572,7 +621,6 @@ def _run_mlx_training(event_queue, stop_queue, config): return ds def _load_local(file_paths): - from core.training.trainer import UnslothTrainer from datasets import load_from_disk if len(file_paths) == 1: @@ -581,10 +629,10 @@ def _run_mlx_training(event_queue, stop_queue, config): (p / "dataset_info.json").exists() or (p / "state.json").exists() ): return load_from_disk(str(p)) - all_files = UnslothTrainer._resolve_local_files(file_paths) + all_files = _resolve_mlx_local_dataset_files(file_paths) if not all_files: raise ValueError("No local dataset files found") - loader = UnslothTrainer._loader_for_files(all_files) + loader = _mlx_local_dataset_loader_for_files(all_files) return load_dataset(loader, data_files = all_files, split = "train") if hf_dataset: @@ -718,6 +766,10 @@ def _run_mlx_training(event_queue, stop_queue, config): else: eval_steps_val = int(eval_steps_val) + # MLX: value-clip grads to [-5, 5]; norm clipping disabled for compile-friendliness. + max_grad_norm = 0.0 + max_grad_value = 5.0 # TODO: expose MLX grad-clip in Studio UI for power users + trainer = MLXTrainer( model = model, tokenizer = tokenizer, @@ -732,6 +784,8 @@ def _run_mlx_training(event_queue, stop_queue, config): lr_scheduler_type = lr_scheduler_type, optim = optim_name, weight_decay = float(config.get("weight_decay", 0.001) or 0.001), + max_grad_norm = max_grad_norm, + max_grad_value = max_grad_value, logging_steps = 1, max_seq_length = max_seq_length, seed = config.get("random_seed", 3407), @@ -820,7 +874,17 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 9. Real-time progress callback ── _send("status", status_message = f"Training {model_name}...") - def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens): + def _on_step( + step, + total, + loss, + lr, + tok_s, + peak_gb, + elapsed, + num_tokens, + grad_norm = None, + ): eta = (elapsed / step * (total - step)) if step > 0 else 0 _send( "progress", @@ -831,7 +895,7 @@ def _run_mlx_training(event_queue, stop_queue, config): total_steps = total, elapsed_seconds = elapsed, eta_seconds = max(0, eta), - grad_norm = None, + grad_norm = grad_norm, num_tokens = num_tokens, eval_loss = None, status_message = None, @@ -846,6 +910,11 @@ def _run_mlx_training(event_queue, stop_queue, config): "train/tokens_per_sec": tok_s, "train/peak_gb": peak_gb, "train/num_tokens": num_tokens, + **( + {"train/grad_norm": grad_norm} + if grad_norm is not None + else {} + ), }, step = step, ) @@ -857,6 +926,8 @@ def _run_mlx_training(event_queue, stop_queue, config): tb_writer.add_scalar("train/learning_rate", lr, step) tb_writer.add_scalar("train/tokens_per_sec", tok_s, step) tb_writer.add_scalar("train/peak_gb", peak_gb, step) + if grad_norm is not None: + tb_writer.add_scalar("train/grad_norm", grad_norm, step) except Exception: pass diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 31f1d575d7..34ae9cfe72 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -262,6 +262,19 @@ class TrainingStartRequest(BaseModel): max_steps: Optional[int] = Field(None, description = "Maximum training steps") save_steps: int = Field(100, description = "Steps between checkpoints") weight_decay: float = Field(0.001, description = "Weight decay") + max_grad_norm: float = Field( + 0.0, + ge = 0, + description = "Global gradient norm clipping threshold. Set 0 to disable.", + ) + max_grad_value: Optional[float] = Field( + None, + ge = 0, + description = ( + "Elementwise gradient value clipping threshold. Set 0 to disable. " + "If omitted, MLX defaults to 1 unless max_grad_norm is set." + ), + ) random_seed: int = Field(42, description = "Random seed") packing: bool = Field(False, description = "Enable sequence packing") optim: str = Field("adamw_8bit", description = "Optimizer") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 19202f3883..6875a13206 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -215,6 +215,8 @@ async def start_training( "max_steps": request.max_steps, "save_steps": request.save_steps, "weight_decay": request.weight_decay, + "max_grad_norm": request.max_grad_norm, + "max_grad_value": request.max_grad_value, "random_seed": request.random_seed, "packing": request.packing, "optim": request.optim, diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 868e537372..ce447bdd1f 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -56,11 +56,14 @@ def _install_fake_fast_mlx(monkeypatch, calls): return _DummyModel(), _DummyTokenizer() unsloth_zoo_pkg = types.ModuleType("unsloth_zoo") - mlx_loader = types.ModuleType("unsloth_zoo.mlx_loader") + mlx_pkg = types.ModuleType("unsloth_zoo.mlx") + mlx_loader = types.ModuleType("unsloth_zoo.mlx.loader") mlx_loader.FastMLXModel = _FastMLXModel - unsloth_zoo_pkg.mlx_loader = mlx_loader + unsloth_zoo_pkg.mlx = mlx_pkg + mlx_pkg.loader = mlx_loader monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg) - monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx_loader", mlx_loader) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader) def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index 876ee34686..c4aaaf3298 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -70,6 +70,50 @@ class TestTrainingRawSupport(unittest.TestCase): self.assertTrue(config["load_in_4bit"]) self.assertEqual(config["embedding_learning_rate"], 1e-5) + def test_training_backend_forwards_grad_clipping_controls(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch( + "core.training.training.prepare_gpu_selection", + return_value = ([0], {"selection_mode": "auto"}), + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", + return_value = DummyThread(), + ), + ): + backend.start_training( + job_id = "test-grad-clip", + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + max_grad_norm = 0.7, + max_grad_value = 0.0, + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertEqual(config["max_grad_norm"], 0.7) + self.assertEqual(config["max_grad_value"], 0.0) + def test_training_route_forwards_embedding_learning_rate(self): training_route = _load_route_module( "training_route_module_raw_support", diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 35fbaba8f0..cfdd811853 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -28,6 +28,23 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair {}""" +def _is_mlx_runtime() -> bool: + try: + from unsloth_zoo.mlx import is_mlx_available + except ImportError: + return False + return is_mlx_available() + + +def _chat_template_kwargs() -> dict: + if not _is_mlx_runtime(): + return {} + return { + "patch_saving": False, + "use_zoo_tokenizer_patch": True, + } + + def get_tokenizer_chat_template(tokenizer, model_name): """ Gets appropriate chat template for tokenizer based on model. @@ -60,6 +77,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): tokenizer = get_chat_template( tokenizer, chat_template = matched_template, + **_chat_template_kwargs(), ) except Exception as e: logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}") @@ -79,6 +97,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): tokenizer = get_chat_template( tokenizer, chat_template = "chatml", + **_chat_template_kwargs(), ) except Exception as e: logger.info(f"⚠️ Failed to apply default ChatML template: {e}") @@ -255,7 +274,11 @@ def apply_chat_template_to_dataset( if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template): try: from unsloth.chat_templates import get_chat_template - tokenizer = get_chat_template(tokenizer, chat_template = "alpaca") + tokenizer = get_chat_template( + tokenizer, + chat_template = "alpaca", + **_chat_template_kwargs(), + ) logger.info(f"📝 Set alpaca chat template on tokenizer for model saving") except Exception as e: logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}") diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 5e68ccd72c..5d81f0df9c 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -83,6 +83,7 @@ export function buildTrainingStartPayload( save_steps: config.saveSteps, eval_steps: config.evalSteps, weight_decay: config.weightDecay, + max_grad_norm: 0.0, random_seed: config.randomSeed, packing: isEmbedding ? false : config.packing, optim: config.optimizerType, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index fb8a2f899e..0cb881e634 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -31,6 +31,7 @@ export interface TrainingStartRequest { save_steps: number; eval_steps: number; weight_decay: number; + max_grad_norm: number; random_seed: number; packing: boolean; optim: string; diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 168e9ad329..f0c90dd9c6 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -211,8 +211,8 @@ def cmd_train(args) -> int: workdir.mkdir(parents = True, exist_ok = True) import mlx.core as mx - from unsloth_zoo.mlx_loader import FastMLXModel - from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.loader import FastMLXModel + from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig hf_token = os.environ.get("HF_TOKEN") or None @@ -440,7 +440,7 @@ def cmd_reload(args) -> int: return _reload_gguf(save_dir, metrics) import mlx.core as mx - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.loader import FastMLXModel from mlx_lm import generate hf_token = os.environ.get("HF_TOKEN") or None diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py index fc07a497e7..043f2796aa 100644 --- a/tests/studio/test_is_mlx_dispatch_gate.py +++ b/tests/studio/test_is_mlx_dispatch_gate.py @@ -7,8 +7,9 @@ Two gates drive every dispatch decision in Studio's MLX path: 1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated once at import time and read by Studio worker code to choose between - the GPU and MLX trainer / inference / export paths. Defined as - ``Darwin AND arm64 AND find_spec("mlx") is not None``. + the GPU and MLX trainer / inference / export paths. It delegates to + the shared zoo MLX runtime gate, with a local import barrier while the + paired unsloth-zoo runtime rollout is in flight. 2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is @@ -18,8 +19,8 @@ Two gates drive every dispatch decision in Studio's MLX path: These gates are the canaries for "MLX support accidentally hijacks CUDA/AMD/Intel users". The tests here: - * verify the source-level structure of the ``_IS_MLX`` expression so an - accidental rewrite (e.g. dropping the ``arm64`` check) is caught, + * verify the source-level structure of the ``_IS_MLX`` helper so an + accidental rewrite importing zoo before the local MLX precheck is caught, * exercise the runtime gate logic under a spoofed Darwin+arm64 platform with a fake ``mlx`` module in ``sys.modules`` to confirm both gates flip True together, @@ -64,20 +65,36 @@ def test_is_mlx_gate_uses_three_required_predicates(): target = node.value break assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py" - assert isinstance(target, ast.BoolOp) and isinstance( - target.op, ast.And - ), "_IS_MLX must be a BoolOp(And) of platform + mlx checks" - + assert isinstance(target, ast.Call), "_IS_MLX must call the shared MLX helper" expr_src = ast.unparse(target) assert ( - "platform.system()" in expr_src and "Darwin" in expr_src - ), "_IS_MLX must check platform.system() == 'Darwin'" + expr_src == "_is_mlx_available()" + ), "_IS_MLX must delegate to the shared MLX runtime gate" + + helper = None + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "_is_mlx_available": + helper = node + break + assert helper is not None, "_is_mlx_available helper not found" + + helper_src = ast.unparse(helper) assert ( - "platform.machine()" in expr_src and "arm64" in expr_src - ), "_IS_MLX must check platform.machine() == 'arm64'" + "platform.system()" in helper_src + and "'Darwin'" in helper_src + and "platform.machine()" in helper_src + and "'arm64'" in helper_src + and "find_spec" in helper_src + and "'mlx'" in helper_src + and "from unsloth_zoo.mlx import is_mlx_available" in helper_src + ), "_IS_MLX helper must precheck local MLX predicates before importing zoo" assert ( - "find_spec" in expr_src and "'mlx'" in expr_src - ), "_IS_MLX must check importlib.util.find_spec('mlx')" + "from unsloth_zoo.mlx import is_mlx_available" in helper_src + and "return is_mlx_available()" in helper_src + ), "_IS_MLX helper must delegate final detection to the shared zoo MLX runtime gate" + assert helper_src.index("UNSLOTH_FORCE_GPU_PATH") < helper_src.index( + "from unsloth_zoo.mlx import is_mlx_available" + ), "_IS_MLX helper must run the local MLX precheck before importing zoo" # --------------------------------------------------------------------------- @@ -87,13 +104,14 @@ def test_is_mlx_gate_uses_three_required_predicates(): # --------------------------------------------------------------------------- -def _evaluate_is_mlx_gate(platform_module, importlib_util): - """Re-evaluate the _IS_MLX expression using injected dependencies. +def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module): + """Re-evaluate the local _is_mlx_available precheck using injected dependencies. - Mirrors the assignment in unsloth/__init__.py exactly. + Mirrors only the cheap import barrier before unsloth imports unsloth_zoo. """ return ( - platform_module.system() == "Darwin" + os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" + and platform_module.system() == "Darwin" and platform_module.machine() == "arm64" and importlib_util.find_spec("mlx") is not None ) @@ -112,7 +130,9 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch): monkeypatch.setattr(platform, "system", lambda: "Darwin") monkeypatch.setattr(platform, "machine", lambda: "arm64") - assert _evaluate_is_mlx_gate(platform, importlib.util) is True + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is True def test_is_mlx_gate_false_when_mlx_missing(monkeypatch): @@ -133,7 +153,9 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch): monkeypatch.setattr(importlib.util, "find_spec", _no_mlx) - assert _evaluate_is_mlx_gate(platform, importlib.util) is False + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False def test_is_mlx_gate_false_on_non_apple_silicon(): @@ -147,7 +169,9 @@ def test_is_mlx_gate_false_on_non_apple_silicon(): pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.") - assert _evaluate_is_mlx_gate(platform, importlib.util) is False + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False # --------------------------------------------------------------------------- diff --git a/tests/studio/test_mlx_training_worker_behaviors.py b/tests/studio/test_mlx_training_worker_behaviors.py index 6c067ea00b..78b229d6e9 100644 --- a/tests/studio/test_mlx_training_worker_behaviors.py +++ b/tests/studio/test_mlx_training_worker_behaviors.py @@ -46,8 +46,8 @@ def test_wandb_init_strips_secret_keys(): def test_local_dataset_loader_uses_load_dataset_path(): src = WORKER.read_text() - assert "_resolve_local_files" in src - assert "_loader_for_files" in src + assert "_resolve_mlx_local_dataset_files" in src + assert "_mlx_local_dataset_loader_for_files" in src assert "data_files = all_files" in src or "data_files=all_files" in src @@ -84,7 +84,7 @@ def test_poll_stop_returns_on_broken_pipe(): def test_unsloth_zoo_mlx_imports_have_friendly_error(): src = WORKER.read_text() - assert "from unsloth_zoo.mlx_loader import FastMLXModel" in src - assert "from unsloth_zoo.mlx_trainer import" in src + assert "from unsloth_zoo.mlx.loader import FastMLXModel" in src + assert "from unsloth_zoo.mlx.trainer import" in src assert "raise ImportError" in src assert "install.sh" in src diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 9b620a5c76..75fb169e83 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -12,16 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os, platform, importlib.util +import os, importlib.util, platform os.environ["UNSLOTH_IS_PRESENT"] = "1" + +def _is_mlx_available(): + # Transitional import barrier: while the paired unsloth-zoo MLX runtime + # rollout is in flight, keep non-Apple-Silicon imports from touching + # unsloth_zoo here. After both PRs are released together and + # unsloth_zoo.mlx is guaranteed to be import-safe on GPU hosts, + # this helper can collapse back to the centralized zoo runtime call below. + if ( + os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1" + or platform.system() != "Darwin" + or platform.machine() != "arm64" + or importlib.util.find_spec("mlx") is None + ): + return False + try: + from unsloth_zoo.mlx import is_mlx_available + except ImportError: + return False + return is_mlx_available() + + # Detect Apple Silicon + MLX before any torch/numpy imports -_IS_MLX = ( - platform.system() == "Darwin" - and platform.machine() == "arm64" - and importlib.util.find_spec("mlx") is not None -) +_IS_MLX = _is_mlx_available() if _IS_MLX: try: @@ -31,18 +48,18 @@ if _IS_MLX: "Unsloth: MLX support requires `unsloth-zoo` with MLX modules. " "Reinstall with `pip install unsloth-zoo` or rerun install.sh." ) from _e - # The mlx_trainer / mlx_loader submodules ship with unsloth-zoo's MLX + # The mlx.trainer / mlx.loader submodules ship with unsloth-zoo's MLX # support. An older installed unsloth-zoo (e.g. from PyPI before the # MLX release lands) will satisfy `import unsloth_zoo` but be missing # these submodules. Surface the same friendly install hint instead of # a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( "Unsloth: MLX support requires an unsloth-zoo build that includes " - "`unsloth_zoo.mlx_trainer` and `unsloth_zoo.mlx_loader`. Upgrade with " + "`unsloth_zoo.mlx.trainer` and `unsloth_zoo.mlx.loader`. Upgrade with " "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 8376dc7e39..1c94e10f9f 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -30,11 +30,9 @@ __all__ = [ from transformers import StoppingCriteria, StoppingCriteriaList from torch import LongTensor, FloatTensor from transformers.models.llama.modeling_llama import logger -from .save import patch_saving_functions import os import shutil from .tokenizer_utils import * -from .models._utils import patch_tokenizer import re from .ollama_template_mappers import OLLAMA_TEMPLATES from unsloth_zoo.dataset_utils import ( @@ -1844,6 +1842,8 @@ def get_chat_template( mapping = {"role" : "role", "content" : "content", "user" : "user", "assistant" : "assistant"}, map_eos_token = True, system_message = None, + patch_saving = True, + use_zoo_tokenizer_patch = False, ): assert(type(map_eos_token) is bool) old_tokenizer = tokenizer @@ -2026,6 +2026,12 @@ def get_chat_template( .replace("'user'", "'" + mapping["user"] + "'")\ .replace("'assistant'", "'" + mapping["assistant"] + "'") + if use_zoo_tokenizer_patch: + # Studio MLX avoids the model-utils tokenizer wrapper because that + # import path pulls in Torch/GPU-specific modules before MLX training. + from unsloth_zoo.tokenizer_utils import patch_tokenizer + else: + from .models._utils import patch_tokenizer _, tokenizer = patch_tokenizer(model = None, tokenizer = tokenizer) tokenizer.padding_side = old_padding_side @@ -2059,7 +2065,9 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - tokenizer = patch_saving_functions(tokenizer) + if patch_saving: + from .save import patch_saving_functions + tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile diff --git a/unsloth/device_type.py b/unsloth/device_type.py index a42d2b9fab..9bad9be0e4 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -20,21 +20,40 @@ __all__ = [ "DEVICE_COUNT", "ALLOW_PREQUANTIZED_MODELS", "ALLOW_BITSANDBYTES", + "is_mlx_available", ] -import torch import functools import inspect +import os from unsloth_zoo.utils import Version +def is_mlx_available(): + try: + from unsloth_zoo.mlx import is_mlx_available as _is_mlx_available + except ImportError: + return False + return _is_mlx_available() + + +_IS_MLX = is_mlx_available() + +if not _IS_MLX: + import torch + + @functools.cache def is_hip(): + if _IS_MLX: + return False return bool(getattr(getattr(torch, "version", None), "hip", None)) @functools.cache def get_device_type(): + if _IS_MLX: + return "mlx" if hasattr(torch, "cuda") and torch.cuda.is_available(): if is_hip(): return "hip" @@ -64,6 +83,8 @@ DEVICE_TYPE: str = get_device_type() DEVICE_TYPE_TORCH = DEVICE_TYPE if DEVICE_TYPE_TORCH == "hip": DEVICE_TYPE_TORCH = "cuda" +elif DEVICE_TYPE_TORCH == "mlx": + DEVICE_TYPE_TORCH = "mps" @functools.cache diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index dd5a9cbf0e..f77baea281 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -160,6 +160,10 @@ else: # INTEL GPU Specific Logic if DEVICE_TYPE == "xpu": _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream +elif DEVICE_TYPE == "mlx": + + def _gpu_getCurrentRawStream(_index = 0): + return 0 # NVIDIA GPU Default Logic elif hasattr(torch._C, "_cuda_getCurrentRawStream"): _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream @@ -206,6 +210,11 @@ if DEVICE_TYPE == "xpu": XPU_STREAMS = () WEIGHT_BUFFERS = [] ABSMAX_BUFFERS = [] +elif DEVICE_TYPE == "mlx": + CUDA_STREAMS = () + XPU_STREAMS = () + WEIGHT_BUFFERS = [] + ABSMAX_BUFFERS = [] else: # NVIDIA GPU Default Logic if DEVICE_COUNT > 0: From 4192fe6ebea2a89fde4daf410577ff30bab54bdd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 05:43:58 -0700 Subject: [PATCH 02/50] studio: drop unused max_grad_value schema + route plumbing (#5424) * studio: drop unused max_grad_value schema + route plumbing The MLX worker hardcodes max_grad_value to 5.0 after PR #5340. The schema field, frontend payload type, route forwarder, and start_training kwarg threading were all left in place as a transitional buffer for old clients. The field is now genuinely unused everywhere except inside the MLX worker, so the schema, route forwarder, and config-build entries can go. Pydantic still tolerates older clients that send max_grad_value because TrainingStartRequest's model_config defaults to extra=ignore. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- scripts/verify_comment_only_diff.py | 8 +++++--- studio/backend/core/training/training.py | 1 - studio/backend/models/training.py | 8 -------- studio/backend/routes/training.py | 1 - studio/backend/tests/test_training_raw_support.py | 2 -- 5 files changed, 5 insertions(+), 15 deletions(-) diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index 068d3244df..90eafb7f8f 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -35,6 +35,7 @@ Example: git diff --name-only origin/main..HEAD \\ | xargs python scripts/verify_comment_only_diff.py --base origin/main """ + from __future__ import annotations import argparse @@ -49,7 +50,9 @@ import yaml def _git_show(rev: str, path: str) -> str: return subprocess.check_output( - ["git", "show", f"{rev}:{path}"], text = True, stderr = subprocess.DEVNULL, + ["git", "show", f"{rev}:{path}"], + text = True, + stderr = subprocess.DEVNULL, ) @@ -145,8 +148,7 @@ def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None: elif isinstance(b, list): if len(b) != len(a): print( - f" list len at {prefix or '/'}: " - f"{len(b)} -> {len(a)}", + f" list len at {prefix or '/'}: " f"{len(b)} -> {len(a)}", ) for i, (bi, ai) in enumerate(zip(b, a)): _walk_yaml_diff(bi, ai, f"{prefix}[{i}]") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 556545e680..e4abb64b8b 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -215,7 +215,6 @@ class TrainingBackend: "save_steps": kwargs.get("save_steps", 0), "weight_decay": kwargs.get("weight_decay", 0.001), "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": kwargs.get("max_grad_value"), "random_seed": kwargs.get("random_seed", 3407), "packing": kwargs.get("packing", False), "optim": kwargs.get("optim", "adamw_8bit"), diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 34ae9cfe72..7c53b0fee5 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -267,14 +267,6 @@ class TrainingStartRequest(BaseModel): ge = 0, description = "Global gradient norm clipping threshold. Set 0 to disable.", ) - max_grad_value: Optional[float] = Field( - None, - ge = 0, - description = ( - "Elementwise gradient value clipping threshold. Set 0 to disable. " - "If omitted, MLX defaults to 1 unless max_grad_norm is set." - ), - ) random_seed: int = Field(42, description = "Random seed") packing: bool = Field(False, description = "Enable sequence packing") optim: str = Field("adamw_8bit", description = "Optimizer") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6875a13206..6e2413b3e9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -216,7 +216,6 @@ async def start_training( "save_steps": request.save_steps, "weight_decay": request.weight_decay, "max_grad_norm": request.max_grad_norm, - "max_grad_value": request.max_grad_value, "random_seed": request.random_seed, "packing": request.packing, "optim": request.optim, diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index c4aaaf3298..384247a191 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -107,12 +107,10 @@ class TestTrainingRawSupport(unittest.TestCase): model_name = "unsloth/test", training_type = "LoRA/QLoRA", max_grad_norm = 0.7, - max_grad_value = 0.0, ) config = mock_process.call_args.kwargs["kwargs"]["config"] self.assertEqual(config["max_grad_norm"], 0.7) - self.assertEqual(config["max_grad_value"], 0.0) def test_training_route_forwards_embedding_learning_rate(self): training_route = _load_route_module( From 000ca89301ba8fb9a58ab718b885d567cfbb4ca1 Mon Sep 17 00:00:00 2001 From: "U. I. I. Derbashi" <32719218+uderbashi@users.noreply.github.com> Date: Thu, 14 May 2026 15:48:28 +0200 Subject: [PATCH 03/50] Studio: Passing batch size for eval (#5168) * add eval batch size * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index a3f063694f..62f1e23e60 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3208,6 +3208,9 @@ class UnslothTrainer: if eval_steps_val > 0: config_args["eval_strategy"] = "steps" config_args["eval_steps"] = eval_steps_val + config_args["per_device_eval_batch_size"] = config_args[ + "per_device_train_batch_size" + ] logger.info( f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n" ) From 79adfd9c71ac57fbbc96b0798af83d8cd7949c35 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Thu, 14 May 2026 18:13:50 +0400 Subject: [PATCH 04/50] studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420) * studio: skip flash-attn install on Blackwell GPUs (sm_100+) Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120, or sm_121, and the older-arch wheels fail to load on Blackwell. Add a shared has_blackwell_gpu() helper and gate both the install-time (install_python_stack._ensure_flash_attn) and runtime (worker._ensure_flash_attn_for_long_context) paths on it. Detection uses nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows. * test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests prefers_prebuilt_wheel and falls_back_to_pypi exercise the install paths that the Blackwell guard now short-circuits. Make them explicit about non-Blackwell so they pass on real Blackwell hosts. * studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH - Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a single process avoid redundant nvidia-smi spawns. Tests clear the cache via setup_method/teardown_method. - In _ensure_flash_attn, run the NO_TORCH short-circuit before the Blackwell check so GGUF-only users (who never install torch anyway) do not see a Blackwell warning. Blackwell check still runs above the IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see the explicit reason rather than a silent OS skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: add has_blackwell_gpu to mlx worker test wheel_utils stub test_mlx_training_worker_config loads worker.py against a hand-rolled utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol list so worker's import line resolves. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/worker.py | 7 + .../tests/test_mlx_training_worker_config.py | 1 + .../tests/test_training_worker_flash_attn.py | 25 +++ studio/backend/utils/wheel_utils.py | 44 ++++ studio/install_python_stack.py | 14 +- .../test_flash_attn_install_python_stack.py | 195 ++++++++++++++++++ 6 files changed, 284 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 8cd1d1f9c8..6b3b3b6609 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -30,6 +30,7 @@ from utils.hardware import apply_gpu_ids from utils.wheel_utils import ( direct_wheel_url, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -313,6 +314,12 @@ def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None: if not _should_try_runtime_flash_attn_install(max_seq_length): return + if has_blackwell_gpu(): + _send_status( + event_queue, + "Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + ) + return installed = _install_package_wheel_first( event_queue = event_queue, diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 5900af4e3d..98c7bdaa55 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,6 +37,7 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", + "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 41a7c87df1..0737bdc82f 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -37,6 +37,7 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -65,6 +66,7 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -112,6 +114,29 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() +def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): + statuses: list[str] = [] + install_mock = mock.Mock() + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr( + worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True + ) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) + monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) + + install_mock.assert_not_called() + assert len(statuses) == 1 + assert "Blackwell" in statuses[0] + + def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 3ed9bda827..5c42e890d1 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations +import functools import json import logging import platform @@ -22,6 +23,49 @@ FLASH_ATTN_RELEASE_BASE_URL = ( ) +@functools.lru_cache(maxsize = 1) +def has_blackwell_gpu() -> bool: + """Return True if any visible NVIDIA GPU has compute capability >= 10.0 + (Blackwell: sm_100, sm_120, sm_121, ...). + + Dao-AILab does not publish prebuilt flash-attention wheels for these + architectures, and the older-arch wheels fail to load on Blackwell, so + callers use this gate to skip the flash-attn install/upgrade path. + + Result is cached for the process lifetime since GPU hardware does not + change. Tests that mock subprocess/nvidia-smi must call + ``has_blackwell_gpu.cache_clear()`` before each invocation. + """ + exe = shutil.which("nvidia-smi") + if not exe: + return False + try: + result = subprocess.run( + [exe, "--query-gpu=compute_cap", "--format=csv,noheader"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + env = child_env_without_native_path_secret(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + if result.returncode != 0: + return False + for line in result.stdout.splitlines(): + cap = line.strip() + if not cap: + continue + major_part = cap.split(".", 1)[0] + try: + major = int(major_part) + except ValueError: + continue + if major >= 10: + return True + return False + + def linux_wheel_platform_tag() -> str | None: machine = platform.machine().lower() if sys.platform.startswith("linux"): diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3fd1e6af66..ab234ad566 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -28,6 +28,7 @@ if str(_BACKEND_DIR) not in sys.path: from backend.utils.wheel_utils import ( flash_attn_package_version, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -628,10 +629,19 @@ def _flash_attn_install_disabled() -> bool: def _ensure_flash_attn() -> None: - if NO_TORCH or IS_WINDOWS or IS_MACOS: - return if _flash_attn_install_disabled(): return + if NO_TORCH: + return + if has_blackwell_gpu(): + _step( + "warning", + "Skipping flash-attn: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + _cyan, + ) + return + if IS_WINDOWS or IS_MACOS: + return if ( subprocess.run( [sys.executable, "-c", "import flash_attn"], diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 9881f2258e..49f4350a7b 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -10,8 +10,133 @@ from unittest import mock STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio" sys.path.insert(0, str(STUDIO_DIR)) +sys.path.insert(0, str(STUDIO_DIR / "backend")) import install_python_stack as ips +from backend.utils import wheel_utils + + +def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess: + return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "") + + +class TestHasBlackwellGpu: + def setup_method(self): + wheel_utils.has_blackwell_gpu.cache_clear() + + def teardown_method(self): + wheel_utils.has_blackwell_gpu.cache_clear() + + def test_returns_false_when_nvidia_smi_missing(self): + with mock.patch.object(wheel_utils.shutil, "which", return_value = None): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_true_for_sm_100(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_true_for_sm_120(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_true_for_sm_121(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_false_for_sm_90(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_for_sm_89(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_mixed_gpus_with_one_blackwell_returns_true(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("8.0\n10.0\n"), + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_false_when_nvidia_smi_fails(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("", returncode = 1), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_on_subprocess_timeout(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_on_malformed_output(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("not-a-number\n\n"), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False class TestFlashAttnWheelSelection: @@ -234,6 +359,76 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() + def test_blackwell_gpu_skips_install_with_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", False), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = True), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) + + def test_blackwell_gpu_on_windows_emits_blackwell_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", True), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = True), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) + + def test_non_blackwell_windows_does_not_emit_blackwell_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", True), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = False), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert not any("Blackwell" in msg for _, msg in step_messages) + class TestInstallPythonStackFlashAttnIntegration: def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int: From ba833b4ade401115618b8d2c2347ee58e2feafcf Mon Sep 17 00:00:00 2001 From: Tenith Hasintha <102300949+Tenith01@users.noreply.github.com> Date: Thu, 14 May 2026 19:45:27 +0530 Subject: [PATCH 05/50] Fix: Add missing utf-8 encoding to text-mode file operations (#5356) Fixes #2795 by explicitly adding encoding='utf-8' to open() calls. This prevents UnicodeDecodeError on Windows with non-UTF-8 system locales when processing files containing UTF-8 characters. Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- unsloth/kernels/moe/autotune_cache.py | 4 ++-- unsloth/kernels/moe/benchmark/utils.py | 2 +- unsloth/models/sentence_transformer.py | 8 ++++---- unsloth/save.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/unsloth/kernels/moe/autotune_cache.py b/unsloth/kernels/moe/autotune_cache.py index f23d9688ea..eac6b02b08 100644 --- a/unsloth/kernels/moe/autotune_cache.py +++ b/unsloth/kernels/moe/autotune_cache.py @@ -71,7 +71,7 @@ def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]: return None try: - with open(cache_file, "r") as f: + with open(cache_file, "r", encoding = "utf-8") as f: cached_data = json.load(f) # Verify cache is still valid (same device, etc.) @@ -118,7 +118,7 @@ def save_cached_config( } try: - with open(cache_file, "w") as f: + with open(cache_file, "w", encoding = "utf-8") as f: json.dump(cache_data, f, indent = 2) logger.info(f"Saved MoE kernel config cache: {cache_key}") except Exception as e: diff --git a/unsloth/kernels/moe/benchmark/utils.py b/unsloth/kernels/moe/benchmark/utils.py index 21905d8df1..56396687e4 100644 --- a/unsloth/kernels/moe/benchmark/utils.py +++ b/unsloth/kernels/moe/benchmark/utils.py @@ -183,7 +183,7 @@ def save_autotune_results(autotune_cache, mode, ref_time, fused_time, results_di filename = "_".join(key) save_path = f"{save_dir}/{filename}.json" print(f"Saving autotune results to {save_path}") - with open(save_path, "w") as f: + with open(save_path, "w", encoding = "utf-8") as f: result = { **config.all_kwargs(), "ref_time": ref_time, diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 52ff268314..c53e3a7a81 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -70,7 +70,7 @@ def _save_pretrained_torchao( modules_path = os.path.join(save_directory, "modules.json") if os.path.exists(modules_path): try: - with open(modules_path, "r") as f: + with open(modules_path, "r", encoding = "utf-8") as f: modules = json.load(f) for m in modules: if m.get("type", "").endswith("Transformer"): @@ -177,7 +177,7 @@ def _save_pretrained_gguf( modules_path = os.path.join(save_directory, "modules.json") if os.path.exists(modules_path): try: - with open(modules_path, "r") as f: + with open(modules_path, "r", encoding = "utf-8") as f: modules = json.load(f) for m in modules: if m.get("type", "").endswith("Transformer"): @@ -542,7 +542,7 @@ class FastSentenceTransformer(FastModel): model_name, "modules.json", token = token ) - with open(modules_json_path, "r") as f: + with open(modules_json_path, "r", encoding = "utf-8") as f: modules_config = json.load(f) pooling_config_path = None @@ -566,7 +566,7 @@ class FastSentenceTransformer(FastModel): break if pooling_config_path: - with open(pooling_config_path, "r") as f: + with open(pooling_config_path, "r", encoding = "utf-8") as f: pooling_config = json.load(f) # from here: # https://github.com/huggingface/sentence-transformers/blob/main/sentence_transformers/models/Pooling.py#L43 diff --git a/unsloth/save.py b/unsloth/save.py index 472eb14933..d67a1a0550 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2641,7 +2641,7 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi ) readme_path = os.path.join(actual_save_directory, "README.md") - with open(readme_path, "w") as f: + with open(readme_path, "w", encoding = "utf-8") as f: f.write(readme_content) api.upload_file( From 739ebeea8241d125c1bb8fa401207e62dcacb326 Mon Sep 17 00:00:00 2001 From: Tenith Hasintha <102300949+Tenith01@users.noreply.github.com> Date: Thu, 14 May 2026 20:19:47 +0530 Subject: [PATCH 06/50] Fix/issue 3667 vicuna template (#5357) * Fix: Add missing utf-8 encoding to text-mode file operations Fixes #2795 by explicitly adding encoding='utf-8' to open() calls. This prevents UnicodeDecodeError on Windows with non-UTF-8 system locales when processing files containing UTF-8 characters. * fix(chat_templates): escape apostrophe in vicuna default system message The unescaped apostrophe in 'user's' broke Jinja2 template parsing when _change_system_message() substituted the default message into the template string via re.sub. Escape it with \' to match the existing vicuna_old pattern. Fixes #3667 --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- unsloth/chat_templates.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 1c94e10f9f..e8a34cbc60 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -211,7 +211,7 @@ vicuna_ollama = _ollama_template("vicuna") vicuna_eos_token = "eos_token" CHAT_TEMPLATES["vicuna"] = (vicuna_template, vicuna_eos_token, False, vicuna_ollama,) -DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions." +DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user\\'s questions." # =========================================== Vicuna Old # https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template From ab21dc25b41a1d7fd142d6beeea572ed4295ae7d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 19:56:21 -0700 Subject: [PATCH 07/50] tests: public-api surface drift detector (companion to test_import_fixes_drift.py) (#5428) * tests: ship public-api surface drift detector + wire into Core matrix Companion to tests/test_import_fixes_drift.py (PR #5414): that file catches drift in THIRD-PARTY libs (transformers / trl / triton / peft / vllm / torchcodec / xformers); this file catches drift in unsloth's OWN public-surface API -- the top-9 classmethods + symbols that unslothai/notebooks calls at ~2000 cumulative sites. Closes the gap where a refactor on this repo (e.g. renaming FastLanguageModel.from_pretrained -> .load) would pass unsloth CI green and surface only on the next unslothai/notebooks CI run, or worse, on a user's Colab crash report. Coverage (call-site counts measured against unslothai/notebooks main): test_fast_language_model_class_present test_fast_language_model_from_pretrained_kwargs 506 sites test_fast_language_model_get_peft_model_kwargs 304 sites test_fast_language_model_for_inference_callable 370 sites test_fast_vision_model_class_and_methods (4 methods) test_fast_vision_model_get_peft_model_vision_kwargs (4 kwargs) test_fast_model_class_and_methods (2 methods) test_fast_model_from_pretrained_kwargs 103 sites test_is_bf16_supported_or_alias_callable 48 + 8 sites Each test asserts the healthy public shape via inspect.signature; on regression fires pytest.fail("DRIFT DETECTED: ...") -- never pytest.skip -- so the Core matrix cell goes red. Mirrors the same skeleton used by tests/test_import_fixes_drift.py. Wired as a new step in consolidated-tests-ci.yml right after the import_fixes drift step, inside every Core matrix cell. Local verification on transformers 4.57.6 + unsloth main: pytest tests/test_public_api_surface.py -v -> 9 passed in 0.02s * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 11 + tests/test_public_api_surface.py | 216 ++++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 tests/test_public_api_surface.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index abceb91567..218d467f04 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -304,6 +304,17 @@ jobs: run: | python -m pytest -v --tb=short tests/test_import_fixes_drift.py + - name: public-api surface drift detectors (9 tests, HARD GATE) + # Companion to test_import_fixes_drift.py: that file catches + # third-party drift; this one catches drift in unsloth's OWN + # public surface (FastLanguageModel / FastVisionModel / + # FastModel + their classmethods + is_bf16_supported). A + # rename here would silently break the unslothai/notebooks tree + # one PR cycle later -- this gate catches it BEFORE the + # breakage reaches users. + run: | + python -m pytest -v --tb=short tests/test_public_api_surface.py + - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) # 16 tests across 5 files. They live inside tests/saving/ and # tests/utils/, both of which Repo tests (CPU) excludes via --ignore diff --git a/tests/test_public_api_surface.py b/tests/test_public_api_surface.py new file mode 100644 index 0000000000..50cc504749 --- /dev/null +++ b/tests/test_public_api_surface.py @@ -0,0 +1,216 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Public-API surface drift detectors for unsloth itself. + +Companion to tests/test_import_fixes_drift.py: that file catches drift +in THIRD-PARTY libraries (transformers / trl / triton / peft / etc.) +that unsloth's import_fixes patches around. This file catches drift in +unsloth's OWN public-surface API -- the top-10 symbols and classmethods +that the unslothai/notebooks tree (and therefore every user on Colab) +calls. If a refactor on this repo renames FastLanguageModel.from_pretrained +or drops one of the documented kwargs, the test fires DRIFT DETECTED +here BEFORE the breakage reaches users. + +Call-site counts measured against unslothai/notebooks @ main: + FastLanguageModel.from_pretrained 506 + FastLanguageModel.for_inference 370 + FastLanguageModel.get_peft_model 304 + FastVisionModel.for_inference 183 + FastVisionModel.from_pretrained 176 + FastVisionModel.get_peft_model 99 + FastVisionModel.for_training 60 + FastModel.from_pretrained 103 + FastModel.get_peft_model 67 + +Mirrors the unsloth-zoo / unsloth drift-detector skeleton: +``pytest.importorskip("unsloth")`` to gate, assert the healthy upstream +shape, ``pytest.fail("DRIFT DETECTED: ...")`` (never ``pytest.skip``) on +regression so the matrix cell goes red. +""" + +from __future__ import annotations + +import inspect + +import pytest + + +def _signature_param_names(callable_obj) -> set[str]: + try: + sig = inspect.signature(callable_obj) + except (TypeError, ValueError): + return set() + return set(sig.parameters) + + +def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]: + """True if every name in ``kwargs`` is either a named parameter on + ``callable_obj`` OR the callable's signature has a ``**kwargs`` + catch-all. Returns (ok, missing_set).""" + try: + sig = inspect.signature(callable_obj) + except (TypeError, ValueError): + return True, set() + params = sig.parameters + has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + if has_var_kw: + return True, set() + missing = kwargs - set(params) + return (not missing), missing + + +# =========================================================================== +# FastLanguageModel: the headline class. 506 from_pretrained + 370 +# for_inference + 304 get_peft_model call sites across the notebooks. +# =========================================================================== + + +def test_fast_language_model_class_present(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastLanguageModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastLanguageModel is missing; every " + "LoRA notebook fails at the first import cell." + ) + + +def test_fast_language_model_from_pretrained_kwargs(): + """from_pretrained must accept the canonical kwargs the notebooks pass.""" + unsloth = pytest.importorskip("unsloth") + required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"} + ok, missing = _accepts(unsloth.FastLanguageModel.from_pretrained, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastLanguageModel.from_pretrained dropped " + f"kwargs {sorted(missing)}; 506 notebook call sites would " + f"crash with TypeError." + ) + + +def test_fast_language_model_get_peft_model_kwargs(): + unsloth = pytest.importorskip("unsloth") + required = { + "r", + "lora_alpha", + "lora_dropout", + "target_modules", + "bias", + "use_gradient_checkpointing", + "random_state", + } + ok, missing = _accepts(unsloth.FastLanguageModel.get_peft_model, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastLanguageModel.get_peft_model dropped " + f"kwargs {sorted(missing)}; 304 notebook call sites would crash." + ) + + +def test_fast_language_model_for_inference_callable(): + unsloth = pytest.importorskip("unsloth") + if not callable(getattr(unsloth.FastLanguageModel, "for_inference", None)): + pytest.fail( + "DRIFT DETECTED: FastLanguageModel.for_inference is missing; " + "370 inference-cell call sites would crash." + ) + + +# =========================================================================== +# FastVisionModel: 183 + 176 + 99 + 60 call sites across vision notebooks. +# =========================================================================== + + +def test_fast_vision_model_class_and_methods(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastVisionModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastVisionModel is missing; every " + "vision fine-tuning notebook fails at import." + ) + cls = unsloth.FastVisionModel + missing = [ + m + for m in ("from_pretrained", "get_peft_model", "for_inference", "for_training") + if not callable(getattr(cls, m, None)) + ] + if missing: + pytest.fail(f"DRIFT DETECTED: FastVisionModel is missing methods {missing}.") + + +def test_fast_vision_model_get_peft_model_vision_kwargs(): + """Vision-specific kwargs the notebooks pass on the vision LoRA path.""" + unsloth = pytest.importorskip("unsloth") + required = { + "finetune_vision_layers", + "finetune_language_layers", + "finetune_attention_modules", + "finetune_mlp_modules", + } + ok, missing = _accepts(unsloth.FastVisionModel.get_peft_model, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastVisionModel.get_peft_model dropped " + f"vision kwargs {sorted(missing)}." + ) + + +# =========================================================================== +# FastModel: the modern unified entry point. 103 + 67 call sites. +# =========================================================================== + + +def test_fast_model_class_and_methods(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastModel is missing; the modern " + "unified entry point used by 100+ notebooks would crash." + ) + missing = [ + m + for m in ("from_pretrained", "get_peft_model") + if not callable(getattr(unsloth.FastModel, m, None)) + ] + if missing: + pytest.fail(f"DRIFT DETECTED: FastModel is missing methods {missing}.") + + +def test_fast_model_from_pretrained_kwargs(): + unsloth = pytest.importorskip("unsloth") + required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"} + ok, missing = _accepts(unsloth.FastModel.from_pretrained, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastModel.from_pretrained dropped kwargs " + f"{sorted(missing)}; 103 notebook call sites would crash." + ) + + +# =========================================================================== +# Bf16 helper alias (renamed once already; keep both accepted). +# =========================================================================== + + +def test_is_bf16_supported_or_alias_callable(): + """48 notebook import sites for is_bf16_supported plus 8 for the + legacy is_bfloat16_supported alias. Either must remain importable.""" + unsloth = pytest.importorskip("unsloth") + has_new = callable(getattr(unsloth, "is_bf16_supported", None)) + has_old = callable(getattr(unsloth, "is_bfloat16_supported", None)) + if not (has_new or has_old): + pytest.fail( + "DRIFT DETECTED: neither unsloth.is_bf16_supported nor " + "unsloth.is_bfloat16_supported is callable; dtype probing " + "in 50+ notebooks fails." + ) From cb15a7a5b66bec7ae8d4beea45bbda008cadfb9f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 20:27:14 -0700 Subject: [PATCH 08/50] add UNSLOTH_ALLOW_CPU=1 path for CPU-only CI (#5429) Lets `import unsloth.trainer` succeed on hosts without a CUDA/XPU/HIP accelerator (typical of zoo's source-inspection test matrix). The env var is read exactly once per process via @functools.cache on `get_device_type()`, so production hosts pay no runtime cost. Three edits beyond the device_type fallback: * `_gpu_init.py:212/247` -- the bf16 + libcuda/bnb setup blocks call `torch.cuda.get_device_capability()` and `libcuda_dirs()`/`bnb.functional.lib.*` unconditionally when DEVICE_TYPE == "cuda". Guard with `and torch.cuda.is_available()` so the new CPU-CI sentinel doesn't fault those. * `_gpu_init.py:353` -- gate `_patch_trl_trainer()` (the `_backwards_compatible_trainer.__init__` wrapper). Under UNSLOTH_ALLOW_CPU we want pristine upstream TRL classes for downstream `inspect.getsource(SFTTrainer)` drift detectors. * `models/_utils.py:1196` -- same `and torch.cuda.is_available()` guard for `get_device_capability()` at import time. * `models/rl.py:PatchFastRL` -- early-return under UNSLOTH_ALLOW_CPU=1 so the heavier `patch_trl_rl_trainers()` (which replaces `trl.SFTTrainer` with the compiled `UnslothSFTTrainer` class) doesn't fire either. Without this gate the drift detectors that do `inspect.getsource(SFTTrainer)` see the wrapper source and spurious fail. Local sanity: `UNSLOTH_ALLOW_CPU=1 python -c "import unsloth.trainer"` succeeds on a CPU-only venv, `trl.SFTTrainer.__init__.__qualname__` stays `SFTTrainer.__init__` (not `UnslothSFTTrainer.__init__`), and `inspect.getsource(SFTTrainer)` still contains `self._signature_columns`. Without the env var on a CUDA host, TRL is still patched normally (verified `UnslothSFTTrainer.__init__`). --- unsloth/_gpu_init.py | 19 +++++++++++++++---- unsloth/device_type.py | 6 ++++++ unsloth/models/_utils.py | 2 +- unsloth/models/rl.py | 5 +++++ 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index aa94c4a568..a30111b529 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -209,7 +209,7 @@ del fix_peft_transformers_weight_conversion_import del patch_peft_weight_converter_compatibility # Torch 2.4 has including_emulation -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): major_version, minor_version = torch.cuda.get_device_capability() SUPPORTS_BFLOAT16 = major_version >= 8 @@ -233,12 +233,18 @@ elif DEVICE_TYPE == "xpu": # torch.xpu.is_bf16_supported() does not have including_emulation # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported() +else: + # CPU-only CI under UNSLOTH_ALLOW_CPU=1. We can't probe device + # capability, so assume no bf16 -- training won't run on this host + # anyway, this branch only exists to let `import unsloth.trainer` + # succeed for source-inspection tests. + SUPPORTS_BFLOAT16 = False # For Gradio HF Spaces? # if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ: import triton -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): libcuda_dirs = lambda: None if Version(triton.__version__) >= Version("3.0.0"): try: @@ -349,5 +355,10 @@ from unsloth_zoo.rl_environments import ( launch_openenv, ) -# Patch TRL trainers for backwards compatibility -_patch_trl_trainer() +# Patch TRL trainers for backwards compatibility. +# Skipped under UNSLOTH_ALLOW_CPU=1 (CPU-only CI) because rebinding +# trl.SFTTrainer.__init__ to a generic wrapper changes +# inspect.getsource(SFTTrainer.__init__) and corrupts downstream +# drift detectors that anchor on the pristine upstream source. +if os.environ.get("UNSLOTH_ALLOW_CPU", "0") != "1": + _patch_trl_trainer() diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 9bad9be0e4..6a82e42e8c 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -63,6 +63,10 @@ def get_device_type(): # Check torch.accelerator if hasattr(torch, "accelerator"): if not torch.accelerator.is_available(): + # Test-only CPU fallback. The env var is read exactly once per + # process because get_device_type is @functools.cache'd. + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return "cuda" raise NotImplementedError( "Unsloth cannot find any torch accelerator? You need a GPU." ) @@ -73,6 +77,8 @@ def get_device_type(): f"But `torch.accelerator.current_accelerator()` works with it being = `{accelerator}`\n" f"Please reinstall torch - it's most likely broken :(" ) + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return "cuda" raise NotImplementedError( "Unsloth currently only works on NVIDIA, AMD and Intel GPUs." ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ccd75aa000..df498e89fb 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1193,7 +1193,7 @@ SUPPORTS_BFLOAT16 = False HAS_FLASH_ATTENTION = False HAS_FLASH_ATTENTION_SOFTCAPPING = False -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): major_version, minor_version = torch.cuda.get_device_capability() torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ee9bdda26a..31a498eada 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -2270,6 +2270,11 @@ def patch_trl_vllm_generation(): def PatchFastRL(algorithm = None, FastLanguageModel = None): if FastLanguageModel is not None: PatchRL(FastLanguageModel) + # Under UNSLOTH_ALLOW_CPU=1 (CPU-only CI), skip TRL trainer rewriting so + # downstream `inspect.getsource(trl.SFTTrainer)` drift detectors see the + # pristine upstream class, not the compiled Unsloth* wrappers. + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return # Install the disable_gradient_checkpointing noop BEFORE # patch_trl_rl_trainers. patch_trl_rl_trainers imports extra trl.* trainer # submodules while generating the compiled cache; any new trl.* modules From 63c6750532bb1c8657b09be87bac31f4f4468802 Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 14 May 2026 20:31:20 -0700 Subject: [PATCH 09/50] fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) (#5350) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) When a flat local GGUF directory holds several unrelated models with their own mmproj siblings, detect_mmproj_file() returned the first projector it walked into. For the layout reported in #5347 (Qwen weights + a Gemma mmproj in the same dir) that meant llama-server was launched with --mmproj pointing at the Gemma projector, which fails to load and surfaces as a confusing crash. Disambiguation rules: - Drop candidates whose family token (qwen/gemma/llama/mistral/phi/...) disagrees with the model's family. Candidates with no recognised family token (e.g. the HF-convention 'mmproj-F16.gguf') are kept. - Among same-family candidates, prefer the one whose stem shares the longest prefix with the model (Qwen3.5-9B mmproj beats Qwen3.5-35B mmproj for a Qwen3.5-9B model). - If every candidate is dropped, return None — better than attaching a wrong projector and getting a server-launch failure. Tests cover the cross-family block, multi-candidate prefix tie-break, HF-convention 'mmproj-F16.gguf', unrecognised families, and the existing search_root walk. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: word-bounded family match, expanded token list, launcher guard Tighten the family-token detector to match only on word boundaries so substring collisions stop tagging false families: phi no longer matches sapphire, yi no longer matches yip, mimo no longer matches mimosa, and mistral does not bleed into ministral/magistral/devstral. Pick the token whose first occurrence is leftmost in the filename rather than the first hit in tuple order, so merge models disambiguate predictably (llama-phi tags llama; phi-llama tags phi). Expand _MODEL_FAMILY_TOKENS with the families an audit of the unsloth HF org turned up that the previous list missed: devstral, ministral, magistral (Mistral-derivative naming), nemotron, kimi, nanonets, cosmos, mimo, apriel, lfm. Without these, a flat local GGUF directory containing one of these weights plus an unrelated renamed projector still hit the original #5347 failure. Add mmproj_matches_model_family() and call it at the llama-server launch site in core/inference/llama_cpp.py. detect_mmproj_file already drops cross-family candidates at discovery time, but mmproj_path can also reach the launcher via config injection or future overrides; this guard keeps those paths from silently loading a known-wrong projector. Tests: 12 new cases covering substring rejection, leftmost-position selection, new family tokens, a new flat-dir Nemotron + Gemma rejection case, and the launcher-level guard. All 21 detect_mmproj_file tests and the existing 106 llama_cpp tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: pair via GGUF general.* metadata, not just filenames Real Unsloth vision GGUFs carry rich identity metadata that has been ignored by the discovery path. Every projector under the unsloth org has general.type='mmproj' plus general.base_model.0.repo_url pointing at the same upstream HF repo as its weight, and the equivalent basename, base_model.0.name, and base_model.0.organization fields. A flat-dir mismatch is therefore decidable from the headers alone, no matter how the user has renamed the files. Add utils/models/gguf_metadata.py with read_gguf_general_metadata(): a fast (~30 ms) header walk that pulls only the general.* string fields and skips everything else, cached by (resolved path, mtime_ns, size). Mirrors the parser shape already used by LlamaCppBackend._read_gguf_metadata so the format handling is consistent. is_mmproj_by_metadata() returns True/False/None from general.type, and pairing_score() returns 100 for an exact base_model URL match, 80 for basename plus organization match, 60 for basename only, -1 for definitive metadata disagreement, and 0 when neither side has enough metadata to decide. Rewire detect_mmproj_file() to a two-stage selector: 1. Detect projectors via metadata (general.type) when present, else fall back to the filename substring heuristic. This recovers headerless projectors AND projectors whose name does not contain 'mmproj' but whose header advertises one. 2. Score each candidate against the weight via pairing_score. Drop candidates with score -1 (definitive metadata disagreement). For candidates with score 0 (no usable metadata) fall back to the existing filename family-token check, dropping recognised-family mismatches. Pick the survivor with the highest (score, longest_prefix, -len(stem)) tuple, so a metadata URL match always wins over a filename-prefix match. Tests: 16 new cases. tests/test_gguf_metadata.py covers the parser (missing file, non-GGUF, string extraction, walking past arrays and uint32s, cache invalidation by mtime/size) and the score helpers. tests/test_detect_mmproj_file.py adds end-to-end cases that synthesise real on-disk GGUF headers: URL match wins over a longer-prefix sibling, URL mismatch returns None even when filenames match, a projector named 'vision-projector.gguf' is still discovered via general.type, and a 100-score header match outranks a near-perfect filename prefix on a headerless candidate. All 75 tests across detect_mmproj_file, gguf_metadata, llama_cpp load progress, cached gguf routes, trained model scan, and vision cache pass. * studio/mmproj: shorten comments and docstrings across the #5347 changes Trim verbose explanations to one-line statements of intent. The behaviour is unchanged: 161 tests across detect_mmproj_file, gguf_metadata, llama_cpp_load_progress (+ matrix), llama_server_args, llama_cpp_cache_aware_disk_check, trained_model_scan, and vision_cache all pass. * studio/mmproj: shorten remaining detect_mmproj_file body comments Trim the docstring and the dir-walking block comments inside detect_mmproj_file to one-liners. Behaviour unchanged; 44 mmproj + gguf_metadata + llama_cpp_load_progress tests pass. * studio/mmproj: cap gguf_metadata cache below ceiling on every insert The eviction branch popped exactly one entry when len >= max, so the cache size could only converge to the cap when entries were added slowly enough for natural growth. After a sandbox sim that reduced the cap mid-run, len stayed above the cap because each insert popped one and added one. Switch to a while loop so we evict until len is strictly below the cap before inserting. Steady-state behaviour at the default 4096 ceiling is unchanged. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 18 +- .../backend/tests/test_detect_mmproj_file.py | 326 ++++++++++++++++++ studio/backend/tests/test_gguf_metadata.py | 216 ++++++++++++ studio/backend/utils/models/gguf_metadata.py | 233 +++++++++++++ studio/backend/utils/models/model_config.py | 198 ++++++++--- 5 files changed, 949 insertions(+), 42 deletions(-) create mode 100644 studio/backend/tests/test_detect_mmproj_file.py create mode 100644 studio/backend/tests/test_gguf_metadata.py create mode 100644 studio/backend/utils/models/gguf_metadata.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 35933e6685..7ef687035c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2367,8 +2367,20 @@ class LlamaCppBackend: if not Path(mmproj_path).is_file(): logger.warning(f"mmproj file not found: {mmproj_path}") else: - cmd.extend(["--mmproj", mmproj_path]) - logger.info(f"Using mmproj for vision: {mmproj_path}") + # #5347 guard for paths that bypass detect_mmproj_file. + from utils.models.model_config import ( + mmproj_matches_model_family, + ) + + if not mmproj_matches_model_family(model_path, mmproj_path): + logger.warning( + f"Skipping mmproj with mismatched family: " + f"model={Path(model_path).name}, " + f"mmproj={Path(mmproj_path).name}" + ) + else: + cmd.extend(["--mmproj", mmproj_path]) + logger.info(f"Using mmproj for vision: {mmproj_path}") # Option C: add --api-key for direct client access when enabled import os as _os @@ -3747,7 +3759,7 @@ class LlamaCppBackend: except json.JSONDecodeError: logger.debug( - f"Skipping malformed SSE line: " f"{line[:100]}" + f"Skipping malformed SSE line: {line[:100]}" ) if _stream_done: break # exit outer for diff --git a/studio/backend/tests/test_detect_mmproj_file.py b/studio/backend/tests/test_detect_mmproj_file.py new file mode 100644 index 0000000000..cdb73448be --- /dev/null +++ b/studio/backend/tests/test_detect_mmproj_file.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for :func:`utils.models.model_config.detect_mmproj_file` (#5347).""" + +from __future__ import annotations + +from pathlib import Path + +import struct + +from utils.models.model_config import ( + _detect_family_token, + detect_mmproj_file, + mmproj_matches_model_family, +) + + +_GGUF_MAGIC = 0x46554747 + + +def _gguf_with_general(path: Path, fields: dict) -> Path: + """Write a minimal GGUF with only ``general.*`` string KVs.""" + body = b"" + for k, v in fields.items(): + kb = k.encode("utf-8") + vb = v.encode("utf-8") + body += struct.pack(" Path: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"") + return path + + +def test_returns_none_when_no_mmproj(tmp_path: Path): + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_single_matching_family_mmproj_picked(tmp_path: Path): + """Single same-family projector: returned (historical behaviour).""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + mmproj = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_hf_style_unprefixed_mmproj_still_works(tmp_path: Path): + """HF convention: weight + ``mmproj-F16.gguf`` sibling.""" + model = _touch(tmp_path / "model.gguf") + mmproj = _touch(tmp_path / "mmproj-F16.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_blocks_single_cross_family_projector(tmp_path: Path): + """#5347 core: Qwen weight + lone Gemma mmproj returns None.""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_picks_matching_family_among_mixed_candidates(tmp_path: Path): + """Mixed Qwen + Gemma projectors: pick Qwen, drop Gemma.""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + qwen_mm = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) == str(qwen_mm.resolve()) + + +def test_prefers_longest_prefix_within_same_family(tmp_path: Path): + """Same family, different sizes: longest shared stem prefix wins.""" + model = _touch(tmp_path / "Qwen3.5-35B-A3B-UD-Q4_K_L.gguf") + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + big_mm = _touch(tmp_path / "Qwen3.5-35B-A3B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(big_mm.resolve()) + + +def test_unrecognised_family_does_not_break_detection(tmp_path: Path): + """Unknown model family must not return None on a sole candidate.""" + model = _touch(tmp_path / "MyCustomBrand-7B-Q4_K_M.gguf") + mmproj = _touch(tmp_path / "MyCustomBrand-7B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_directory_path_returns_first_candidate(tmp_path: Path): + """Directory path: no model stem to compare; legacy first-candidate.""" + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + result = detect_mmproj_file(str(tmp_path)) + assert result is not None + assert "mmproj" in Path(result).name.lower() + + +def test_search_root_walk_still_works(tmp_path: Path): + """Snapshot layout: weight in quant subdir, mmproj at snapshot root.""" + snapshot = tmp_path / "snapshot" + weight = _touch(snapshot / "BF16" / "Qwen3.5-9B-BF16.gguf") + mmproj = _touch(snapshot / "Qwen3.5-9B-BF16-mmproj.gguf") + result = detect_mmproj_file(str(weight), search_root = str(snapshot)) + assert result == str(mmproj.resolve()) + + +# -- Family token detection: word-bounded matching ---------------------- + + +def test_family_token_phi_does_not_match_sapphire(): + """``phi`` substring inside ``sapphire`` must not tag Phi.""" + assert _detect_family_token("sapphire-7b-q4_k_m.gguf") is None + + +def test_family_token_yi_does_not_match_tinyish_names(): + """``yi`` must not cross letter boundaries (``yip``).""" + assert _detect_family_token("yip-7b.gguf") is None + assert _detect_family_token("yi-vl-6b.gguf") == "yi" + + +def test_family_token_mimo_does_not_match_mimosa(): + """``mimo`` must not tag ``mimosa``.""" + assert _detect_family_token("mimosa-rosa-7b.gguf") is None + assert _detect_family_token("MiMo-VL-7B-RL-BF16.gguf") == "mimo" + + +def test_family_token_mistral_does_not_match_ministral(): + """Pin Mistral-derivative tagging.""" + assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral" + assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral" + assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral" + assert ( + _detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf") + == "devstral" + ) + + +def test_family_token_picks_leftmost_when_multiple_present(): + """Leftmost family token wins, not tuple order.""" + assert _detect_family_token("llama-phi-merge.gguf") == "llama" + assert _detect_family_token("phi-llama-merge.gguf") == "phi" + assert _detect_family_token("llama3-3b-instruct.gguf") == "llama" + + +def test_family_token_new_families_recognised(): + """Catalogue-audit additions tag correctly.""" + assert _detect_family_token("NVIDIA-Nemotron-3-Nano-Omni-30B.gguf") == "nemotron" + assert _detect_family_token("Kimi-K2.6-BF16.gguf") == "kimi" + assert _detect_family_token("Nanonets-OCR-s-BF16.gguf") == "nanonets" + assert _detect_family_token("Cosmos-Reason1-7B-BF16.gguf") == "cosmos" + assert _detect_family_token("Apriel-1.5-15b-Thinker-BF16.gguf") == "apriel" + assert _detect_family_token("LFM2.5-VL-1.6B-BF16.gguf") == "lfm" + + +# -- Cross-family rejection with the expanded token list ---------------- + + +def test_blocks_cross_family_for_new_token_pair(tmp_path: Path): + """Nemotron weight + lone Gemma projector returns None.""" + model = _touch( + tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf" + ) + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_picks_devstral_mmproj_in_mixed_dir(tmp_path: Path): + """Devstral weight + Devstral mmproj + a Qwen mmproj: pick Devstral.""" + model = _touch(tmp_path / "Devstral-Small-2-24B-Instruct-2512-BF16.gguf") + dev_mm = _touch(tmp_path / "Devstral-Small-2-mmproj-bf16.gguf") + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(dev_mm.resolve()) + + +# -- Launcher-level family guard ---------------------------------------- + + +def test_mmproj_family_guard_blocks_cross_family(): + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/gemma-4-26B-A4B-it.mmproj-q8_0.gguf", + ) + is False + ) + + +def test_mmproj_family_guard_allows_same_family(): + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/Qwen3.5-9B-BF16-mmproj.gguf", + ) + is True + ) + + +def test_mmproj_family_guard_allows_generic_hf_mmproj(): + """No family token on the projector: wildcard.""" + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/mmproj-F16.gguf", + ) + is True + ) + + +def test_mmproj_family_guard_allows_unrecognised_model_family(): + """No family token on the model: wildcard.""" + assert ( + mmproj_matches_model_family( + "/models/Apriel-1.5-15b-Thinker-BF16.gguf", + "/models/mmproj-F16.gguf", + ) + is True + ) + + +# -- Metadata-primary pairing in detect_mmproj_file --------------------- + + +def test_metadata_url_match_picked_over_filename_lookalike(tmp_path: Path): + """URL match beats a longer-prefix sibling.""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B-Q4_K_M.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + # Closer filename prefix, wrong upstream. + _gguf_with_general( + tmp_path / "Qwen3.5-9B-mmproj-bf16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-1.5B", + }, + ) + # Matching upstream. + correct = _gguf_with_general( + tmp_path / "mmproj-BF16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(correct.resolve()) + + +def test_metadata_url_mismatch_dropped(tmp_path: Path): + """Filenames match family but metadata disagrees: returns None.""" + weight = _gguf_with_general( + tmp_path / "qwen-9b.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + _gguf_with_general( + tmp_path / "qwen-9b-mmproj.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B", + }, + ) + assert detect_mmproj_file(str(weight)) is None + + +def test_metadata_identifies_mmproj_without_filename_hint(tmp_path: Path): + """Projector named ``vision-projector.gguf`` discovered via header.""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + projector = _gguf_with_general( + tmp_path / "vision-projector.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(projector.resolve()) + + +def test_metadata_score_outranks_filename_prefix(tmp_path: Path): + """Score 100 (URL match) beats score 0 (long filename prefix).""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B-Q4_K_M.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + # Headerless: long shared stem, score 0. + _touch(tmp_path / "Qwen3.5-9B-Q4_K_M-mmproj.gguf") + # Headered: generic name, score 100. + correct = _gguf_with_general( + tmp_path / "mmproj-BF16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(correct.resolve()) diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py new file mode 100644 index 0000000000..cf1a17347f --- /dev/null +++ b/studio/backend/tests/test_gguf_metadata.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for :mod:`utils.models.gguf_metadata`. Synthesise small GGUF +headers in tmp dirs so we never depend on real model files.""" + +from __future__ import annotations + +import struct +from pathlib import Path +from typing import Iterable, Mapping + +from utils.models.gguf_metadata import ( + is_mmproj_by_metadata, + pairing_score, + read_gguf_general_metadata, +) + + +_GGUF_MAGIC = 0x46554747 +_VTYPE_STRING = 8 +_VTYPE_UINT32 = 4 +_VTYPE_ARRAY = 9 + + +def _enc_string(s: str) -> bytes: + b = s.encode("utf-8") + return struct.pack(" bytes: + return _enc_string(key) + struct.pack(" bytes: + return ( + _enc_string(key) + struct.pack(" bytes: + vals = list(values) + out = _enc_string(key) + struct.pack(" Path: + """Minimal GGUF: header + KV body, no tensors.""" + extra_uint32 = extra_uint32 or {} + extra_string_arrays = extra_string_arrays or {} + kv_count = len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + body = b"" + for k, v in general_strings.items(): + body += _enc_kv_string(k, v) + for k, v in extra_uint32.items(): + body += _enc_kv_uint32(k, v) + for k, v in extra_string_arrays.items(): + body += _enc_kv_string_array(k, v) + header = struct.pack( + " Optional[_CacheKey]: + try: + st = os.stat(path) + except OSError: + return None + try: + resolved = str(Path(path).resolve()) + except OSError: + resolved = str(path) + return (resolved, st.st_mtime_ns, st.st_size) + + +def read_gguf_general_metadata(path: str) -> Optional[Dict[str, str]]: + """Return ``general.*`` strings from a GGUF header, or ``None`` if + the file is missing, unreadable, or not a GGUF. ``{}`` means the + file is valid but carries none of the wanted keys.""" + key = _cache_key(path) + if key is None: + return None + with _CACHE_LOCK: + if key in _METADATA_CACHE: + return _METADATA_CACHE[key] + result = _parse_gguf_header(path) + with _CACHE_LOCK: + # Arbitrary eviction; header reads are cheap so true LRU is overkill. + while len(_METADATA_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _METADATA_CACHE.pop(next(iter(_METADATA_CACHE))) + except StopIteration: + break + _METADATA_CACHE[key] = result + return result + + +def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: + out: Dict[str, str] = {} + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: # 1 MB sanity bound + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: # 4 MB sanity bound + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + out[key] = sbytes.decode("utf-8", "replace") + else: + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"read_gguf_general_metadata: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"read_gguf_general_metadata: parse failure on {path}: {e}") + return None + return out + + +# Strings (8) and arrays (9) are handled inline. +_FIXED_VTYPE_SIZES: Dict[int, int] = { + 0: 1, # uint8 + 1: 1, # int8 + 2: 2, # uint16 + 3: 2, # int16 + 4: 4, # uint32 + 5: 4, # int32 + 6: 4, # float32 + 7: 1, # bool + 10: 8, # uint64 + 11: 8, # int64 + 12: 8, # float64 +} + + +def _skip_gguf_value(f, vtype: int) -> bool: + """Advance past one GGUF value. False on truncation or unknown type.""" + if vtype == 8: # STRING + slen_bytes = f.read(8) + if len(slen_bytes) < 8: + return False + slen = struct.unpack(" 1 << 30: # 1 GB sanity bound + return False + return len(f.read(slen)) == slen + if vtype == 9: # ARRAY + head = f.read(12) + if len(head) < 12: + return False + atype, alen = struct.unpack(" 1 << 30: + return False + if atype == 8: + for _ in range(alen): + slen_bytes = f.read(8) + if len(slen_bytes) < 8: + return False + slen = struct.unpack(" 1 << 30: + return False + if len(f.read(slen)) != slen: + return False + return True + sz = _FIXED_VTYPE_SIZES.get(atype) + if sz is None: + return False + total = sz * alen + return len(f.read(total)) == total + sz = _FIXED_VTYPE_SIZES.get(vtype) + if sz is None: + return False + return len(f.read(sz)) == sz + + +def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]: + """True/False from ``general.type``; None means fall back to filename.""" + if not meta: + return None + t = meta.get("general.type") + if t is None: + return None + return t.lower() == "mmproj" + + +def pairing_score( + weight_meta: Optional[Dict[str, str]], + mmproj_meta: Optional[Dict[str, str]], +) -> int: + """Pairing confidence: 100 = base_model URL match, 80 = basename + org, + 60 = basename, -1 = definitive mismatch, 0 = decide from filename.""" + if not weight_meta or not mmproj_meta: + return 0 + + w_url = weight_meta.get("general.base_model.0.repo_url") + p_url = mmproj_meta.get("general.base_model.0.repo_url") + if w_url and p_url: + return 100 if w_url.strip().rstrip("/") == p_url.strip().rstrip("/") else -1 + + w_base = weight_meta.get("general.basename") + p_base = mmproj_meta.get("general.basename") + w_org = weight_meta.get("general.base_model.0.organization") or weight_meta.get( + "general.organization" + ) + p_org = mmproj_meta.get("general.base_model.0.organization") or mmproj_meta.get( + "general.organization" + ) + if w_base and p_base and w_org and p_org: + if w_base.lower() == p_base.lower() and w_org.lower() == p_org.lower(): + return 80 + return -1 + + if w_base and p_base: + return 60 if w_base.lower() == p_base.lower() else -1 + + return 0 diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index ebf85c5320..bf7f7a009b 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -19,6 +19,11 @@ from utils.paths import ( resolve_export_dir, ) from utils.utils import without_hf_auth +from utils.models.gguf_metadata import ( + is_mmproj_by_metadata, + pairing_score, + read_gguf_general_metadata, +) import structlog from loggers import get_logger import os @@ -801,12 +806,15 @@ _AUDIO_TOKEN_PATTERNS = { "whisper": lambda tokens: "<|startoftranscript|>" in tokens, "audio_vlm": lambda tokens: "" in tokens, "bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens), - "dac": lambda tokens: "<|audio_start|>" in tokens - and "<|audio_end|>" in tokens - and "<|text_start|>" in tokens - and "<|text_end|>" in tokens, - "snac": lambda tokens: sum(1 for t in tokens if t.startswith(" 10000, + "dac": lambda tokens: ( + "<|audio_start|>" in tokens + and "<|audio_end|>" in tokens + and "<|text_start|>" in tokens + and "<|text_end|>" in tokens + ), + "snac": lambda tokens: ( + sum(1 for t in tokens if t.startswith(" 10000 + ), } @@ -913,6 +921,85 @@ def _is_mmproj(filename: str) -> bool: return "mmproj" in filename.lower() +# Family tokens for #5347's filename fallback. Lowercase. Order does not +# matter (see ``_detect_family_token``). +_MODEL_FAMILY_TOKENS: tuple[str, ...] = ( + "qwen", + "gemma", + "llama", + "mistral", + "ministral", + "magistral", + "devstral", + "phi", + "deepseek", + "internvl", + "minicpm", + "llava", + "glm", + "yi", + "command-r", + "molmo", + "pixtral", + "smolvlm", + "moondream", + "granite", + "ovis", + "nemotron", + "kimi", + "nanonets", + "cosmos", + "mimo", + "apriel", + "lfm", +) + + +# Word-bounded match: any letter on either side disqualifies. Stops +# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc. +_FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {} + + +def _family_token_re(token: str) -> "_re.Pattern[str]": + pat = _FAMILY_TOKEN_RE_CACHE.get(token) + if pat is None: + pat = _re.compile(rf"(?:^|[^a-z])({_re.escape(token)})(?:[^a-z]|$)") + _FAMILY_TOKEN_RE_CACHE[token] = pat + return pat + + +def _detect_family_token(filename: str) -> Optional[str]: + """Leftmost-position match; ties prefer the longer token.""" + name = filename.lower() + best: Optional[tuple[int, int, str]] = None # (start, -len, token) + for token in _MODEL_FAMILY_TOKENS: + m = _family_token_re(token).search(name) + if m is None: + continue + key = (m.start(1), -len(token), token) + if best is None or key < best: + best = key + return None if best is None else best[2] + + +def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool: + """Defense-in-depth guard for the launcher: True unless both filenames + carry recognised family tokens that disagree.""" + model_fam = _detect_family_token(Path(model_path).name) + mmproj_fam = _detect_family_token(Path(mmproj_path).name) + if model_fam is None or mmproj_fam is None: + return True + return model_fam == mmproj_fam + + +def _shared_prefix_len(a: str, b: str) -> int: + n = min(len(a), len(b)) + for i in range(n): + if a[i] != b[i]: + return i + return n + + def _is_gguf_filename(filename: str) -> bool: return filename.lower().endswith(".gguf") @@ -927,33 +1014,18 @@ def _iter_gguf_files(directory: Path, recursive: bool = False): def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: - """ - Find the mmproj (vision projection) GGUF file for a given model. + """Find the mmproj GGUF for a model. - Args: - path: Directory to search — or a .gguf file (uses its parent dir - as the starting point). - search_root: Optional outer directory that should also be scanned - (and any directory between it and ``path``). This handles - local layouts where the model weights live in a quant-named - subdir (``snapshot/BF16/foo.gguf``) but the mmproj sits at - the snapshot root (``snapshot/mmproj-BF16.gguf``). When - ``None``, only the immediate parent dir is scanned, matching - the historical behavior. - - Returns: - Full path to the mmproj .gguf file, or None if not found. - """ + ``path``: directory or a .gguf file. ``search_root``: optional ancestor + to also walk (snapshot layouts where the weight is in ``snapshot/BF16/`` + but the projector sits at ``snapshot/``). Returns the projector path or + ``None``.""" p = Path(path) start_dir = p.parent if p.is_file() else p if not start_dir.is_dir(): return None - # Build the list of dirs to scan: immediate dir first, then walk up - # to (and including) ``search_root`` if it is an ancestor. We walk - # incrementally rather than recursing into ``search_root`` so we - # don't accidentally pick up an mmproj from a sibling subdir - # belonging to a different model variant. + # Walk incrementally so a sibling subdir's mmproj cannot leak in. seen: set[Path] = set() scan_order: list[Path] = [] @@ -969,12 +1041,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional _add(start_dir) - # When ``path`` is a symlink (e.g. Ollama's ``.studio_links/...gguf`` - # -> ``blobs/sha256-...``), the symlink's parent directory rarely - # contains the mmproj sibling; the real mmproj file lives next to - # the symlink target. Add the target's parent to the scan so vision - # GGUFs that are surfaced via symlinks are still recognised as - # vision models. + # Ollama's .studio_links/foo.gguf -> blobs/sha256-...: also scan target dir. try: if p.is_symlink() and p.is_file(): target_parent = p.resolve().parent @@ -986,14 +1053,12 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional try: root_resolved = Path(search_root).resolve() start_resolved = start_dir.resolve() - # Only walk if start_dir is inside (or equal to) search_root. if root_resolved == start_resolved or ( start_resolved.is_relative_to(root_resolved) if hasattr(start_resolved, "is_relative_to") else str(start_resolved).startswith(str(root_resolved) + "/") ): cur = start_resolved - # Walk up from start_dir to (and including) root_resolved. while cur != root_resolved and cur.parent != cur: cur = cur.parent _add(cur) @@ -1002,11 +1067,66 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional except OSError: pass + candidates: list[Path] = [] + seen_resolved: set[Path] = set() for d in scan_order: for f in _iter_gguf_files(d): - if _is_mmproj(f.name): - return str(f.resolve()) - return None + try: + resolved = f.resolve() + except OSError: + continue + if resolved in seen_resolved: + continue + # Prefer ``general.type=='mmproj'``; fall back to filename. + meta = read_gguf_general_metadata(str(resolved)) + by_meta = is_mmproj_by_metadata(meta) + if by_meta is True or (by_meta is None and _is_mmproj(f.name)): + seen_resolved.add(resolved) + candidates.append(resolved) + + if not candidates: + return None + + # Directory path: no model name to compare against; legacy behaviour. + if not p.is_file(): + return str(candidates[0]) + + # Stage 1: GGUF metadata. Stage 2: filename family token (#5347). + model_stem = p.stem.lower() + model_family = _detect_family_token(p.name) + weight_meta = read_gguf_general_metadata(str(p)) + + scored: list[tuple[int, Path]] = [] + for c in candidates: + cand_meta = read_gguf_general_metadata(str(c)) + meta_score = pairing_score(weight_meta, cand_meta) + if meta_score == -1: + logger.info(f"detect_mmproj_file: dropped {c.name} (metadata mismatch)") + continue + if meta_score == 0 and model_family is not None: + # Unrecognised candidate family is a wildcard (``mmproj-F16.gguf``). + cand_family = _detect_family_token(c.name) + if cand_family is not None and cand_family != model_family: + logger.info( + f"detect_mmproj_file: dropped {c.name} " + f"(filename family {cand_family!r} vs model {model_family!r})" + ) + continue + scored.append((meta_score, c)) + + if not scored: + return None + + # Score first, then longest shared prefix, then shorter stem. + best = max( + scored, + key = lambda sc: ( + sc[0], + _shared_prefix_len(model_stem, sc[1].stem.lower()), + -len(sc[1].stem), + ), + ) + return str(best[1]) def detect_gguf_model(path: str) -> Optional[str]: @@ -1360,7 +1480,7 @@ def detect_gguf_model_remote( if attempt < 2: time.sleep(2**attempt) logger.warning( - f"Could not check GGUF files for '{repo_id}' after 3 attempts: " f"{last_err}" + f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}" ) return None From bbd0ba0c259353fc5147a3255af1af34185b9ed2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 21:57:04 -0700 Subject: [PATCH 10/50] studio/mmproj: skip unwanted GGUF values via seek instead of read (#5431) The previous _skip_gguf_value walked past discarded values with f.read(n), which allocates and immediately drops a Python bytes object. For weight GGUFs that carry tokenizer.ggml.tokens (~150K unicode strings) this wasted ~10 MB of allocation per cold call. Switch the discard path to f.seek(n, 1). The kernel never has to copy the bytes into userspace and Python never allocates. Truncation is now detected on the next read attempt rather than inline (an out-of-range seek on a regular file is legal and the next read returns short). Measured on real downloaded GGUFs (Qwen3.5-4B IQ2_XXS 1.52 GB, bartowski Qwen3.5-4B IQ2_M 1.70 GB, Qwen3.5-4B-MTP IQ2_M 1.94 GB): before: 142 ms cold per weight, ~11 MB read after: 90 ms cold per weight, ~4 MB read Mmproj reads are unaffected (no tokenizer to skip). Cached re-reads remain ~50 microseconds. All 161 in-tree backend tests + 85 isolated sandbox tests pass. --- studio/backend/utils/models/gguf_metadata.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index 9aab88329d..5629bac58b 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -152,7 +152,9 @@ _FIXED_VTYPE_SIZES: Dict[int, int] = { def _skip_gguf_value(f, vtype: int) -> bool: - """Advance past one GGUF value. False on truncation or unknown type.""" + """Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal + on a regular file so truncation is detected on the next read; we + only return False for unknown types or sanity-bound overflow.""" if vtype == 8: # STRING slen_bytes = f.read(8) if len(slen_bytes) < 8: @@ -160,7 +162,8 @@ def _skip_gguf_value(f, vtype: int) -> bool: slen = struct.unpack(" 1 << 30: # 1 GB sanity bound return False - return len(f.read(slen)) == slen + f.seek(slen, 1) + return True if vtype == 9: # ARRAY head = f.read(12) if len(head) < 12: @@ -176,18 +179,18 @@ def _skip_gguf_value(f, vtype: int) -> bool: slen = struct.unpack(" 1 << 30: return False - if len(f.read(slen)) != slen: - return False + f.seek(slen, 1) return True sz = _FIXED_VTYPE_SIZES.get(atype) if sz is None: return False - total = sz * alen - return len(f.read(total)) == total + f.seek(sz * alen, 1) + return True sz = _FIXED_VTYPE_SIZES.get(vtype) if sz is None: return False - return len(f.read(sz)) == sz + f.seek(sz, 1) + return True def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]: From 5345b10b6ad4034a8dbe633fce9ce4c509942801 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 01:25:23 -0700 Subject: [PATCH 11/50] ci: install ipython so transformers.utils.notebook imports cleanly in zoo pytest (#5437) unsloth_zoo's drift-detector tests/test_zoo_source_upstream_refs.py:: test_logging_utils_utils_notebook resolves transformers.utils.notebook, which executes ``import IPython.display as disp`` at module scope. The Core matrix install list did not include IPython, so the import raised ModuleNotFoundError and the test failed with: DRIFT DETECTED: transformers.utils.notebook exists but its imports fail on this install (ModuleNotFoundError: No module named 'IPython') The test message itself states the resolution: "Either install the dep in CI or remove the zoo reference." Installing keeps the upstream-refs detector functional. Add ipython to the matrix install list. --- .github/workflows/consolidated-tests-ci.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 218d467f04..2b0ce43a7c 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -206,7 +206,8 @@ jobs: 'numpy<3' pytest==9.0.3 pytest-asyncio httpx \ protobuf sentencepiece triton \ psutil packaging tqdm safetensors datasets \ - 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' + 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \ + ipython # torchvision: unsloth_zoo.vision_utils imports it at module scope. pip install --index-url https://download.pytorch.org/whl/cpu \ 'torch>=2.4,<2.11' 'torchvision<0.26' From 762657afd24673323dd78ef7fc842260406b9793 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 03:51:55 -0700 Subject: [PATCH 12/50] studio/mlx: lower per-element grad clip default from 5.0 to 1.0 (#5440) Studio's MLX training worker explicitly pinned ``max_grad_value=5.0`` into the ``MLXTrainingConfig`` so it would override the zoo default regardless. The 5.0 threshold was effectively no protection -- per- element transformer gradients in steady state are 1e-3..1e-1, so |g_i| > 5 basically never fires even on spike batches, mixed-precision overflow, or RL gradient bursts. Switch to 1.0: - matches the universal LLM clip_grad_norm=1.0 baseline (HF Trainer / TRL / PEFT / AutoTrain) while staying on MLX's fast per-element ``tree_map(mx.clip)`` path (no global reduction) - actually catches outliers without distorting Adam's normalised updates (typical post-warmup |g_i| << 1.0) - lines up with the new MLXTrainingConfig default in unslothai/unsloth-zoo so Studio doesn't silently disagree with what zoo ships No UI change; the TODO to expose grad clipping in Studio settings remains. Existing trained runs are unaffected: only newly-spawned training workers pick up the tighter clip. --- studio/backend/core/training/worker.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 6b3b3b6609..4434436ca3 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -773,9 +773,11 @@ def _run_mlx_training(event_queue, stop_queue, config): else: eval_steps_val = int(eval_steps_val) - # MLX: value-clip grads to [-5, 5]; norm clipping disabled for compile-friendliness. + # MLX: per-element clip to [-1, 1]; norm clip disabled (it needs a + # global reduction that breaks MLX's eager pipeline). 1.0 (not 5.0): + # |g_i| > 5 rarely fires, so the historical 5.0 was effectively no-op. max_grad_norm = 0.0 - max_grad_value = 5.0 # TODO: expose MLX grad-clip in Studio UI for power users + max_grad_value = 1.0 # TODO: expose MLX grad-clip in Studio UI for power users trainer = MLXTrainer( model = model, From 30f6280835670a8d6c553cb05d803af82ab81e91 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 03:53:48 -0700 Subject: [PATCH 13/50] studio/frontend: drop unused next dependency (#5438) The frontend is a Vite SPA wrapped by Tauri and served by FastAPI's StaticFiles in web mode. Nothing in src imports from next/, no next.config exists, and no script invokes the Next.js server. The package was dead weight in node_modules and was being flagged by SCA scanners under CVE-2026-44578 (Next.js SSRF via WebSocket upgrade) despite the vulnerable code path never being reachable. next-themes is unrelated and stays; its only peers are react and react-dom. Verified with npm install + npm run build (tsc -b && vite build), clean exit, dist/ produced as before. --- studio/frontend/package-lock.json | 798 +----------------------------- studio/frontend/package.json | 1 - 2 files changed, 1 insertion(+), 798 deletions(-) diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index 464f47c09c..ee0d8a7832 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -33,7 +33,7 @@ "@streamdown/math": "1.0.2", "@streamdown/mermaid": "1.0.2", "@tailwindcss/vite": "^4.2.2", - "@tanstack/react-router": "^1.159.10", + "@tanstack/react-router": "1.169.2", "@tanstack/react-table": "^8.21.3", "@tauri-apps/api": "^2.10.1", "@tauri-apps/plugin-clipboard-manager": "^2.3.2", @@ -56,7 +56,6 @@ "lucide-react": "^1.7.0", "mammoth": "^1.11.0", "motion": "^12.34.0", - "next": "^16.1.6", "next-themes": "^0.4.6", "node-forge": "^1.4.0", "radix-ui": "^1.4.3", @@ -1542,472 +1541,6 @@ "mlly": "^1.8.2" } }, - "node_modules/@img/colour": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@img/colour/-/colour-1.1.0.tgz", - "integrity": "sha512-Td76q7j57o/tLVdgS746cYARfSyxk8iEfRxewL9h4OMzYhbW4TAcppl0mT4eyqXddh6L/jwoM75mo7ixa/pCeQ==", - "license": "MIT", - "optional": true, - "engines": { - "node": ">=18" - } - }, - "node_modules/@img/sharp-darwin-arm64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-darwin-arm64/-/sharp-darwin-arm64-0.34.5.tgz", - "integrity": "sha512-imtQ3WMJXbMY4fxb/Ndp6HBTNVtWCUI0WdobyheGf5+ad6xX8VIDO8u2xE4qc/fr08CKG/7dDseFtn6M6g/r3w==", - "cpu": [ - "arm64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-darwin-arm64": "1.2.4" - } - }, - "node_modules/@img/sharp-darwin-x64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-darwin-x64/-/sharp-darwin-x64-0.34.5.tgz", - "integrity": "sha512-YNEFAF/4KQ/PeW0N+r+aVVsoIY0/qxxikF2SWdp+NRkmMB7y9LBZAVqQ4yhGCm/H3H270OSykqmQMKLBhBJDEw==", - "cpu": [ - "x64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-darwin-x64": "1.2.4" - } - }, - "node_modules/@img/sharp-libvips-darwin-arm64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-arm64/-/sharp-libvips-darwin-arm64-1.2.4.tgz", - "integrity": "sha512-zqjjo7RatFfFoP0MkQ51jfuFZBnVE2pRiaydKJ1G/rHZvnsrHAOcQALIi9sA5co5xenQdTugCvtb1cuf78Vf4g==", - "cpu": [ - "arm64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "darwin" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-darwin-x64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-x64/-/sharp-libvips-darwin-x64-1.2.4.tgz", - "integrity": "sha512-1IOd5xfVhlGwX+zXv2N93k0yMONvUlANylbJw1eTah8K/Jtpi15KC+WSiaX/nBmbm2HxRM1gZ0nSdjSsrZbGKg==", - "cpu": [ - "x64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "darwin" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-arm": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm/-/sharp-libvips-linux-arm-1.2.4.tgz", - "integrity": "sha512-bFI7xcKFELdiNCVov8e44Ia4u2byA+l3XtsAj+Q8tfCwO6BQ8iDojYdvoPMqsKDkuoOo+X6HZA0s0q11ANMQ8A==", - "cpu": [ - "arm" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-arm64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm64/-/sharp-libvips-linux-arm64-1.2.4.tgz", - "integrity": "sha512-excjX8DfsIcJ10x1Kzr4RcWe1edC9PquDRRPx3YVCvQv+U5p7Yin2s32ftzikXojb1PIFc/9Mt28/y+iRklkrw==", - "cpu": [ - "arm64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-ppc64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-ppc64/-/sharp-libvips-linux-ppc64-1.2.4.tgz", - "integrity": "sha512-FMuvGijLDYG6lW+b/UvyilUWu5Ayu+3r2d1S8notiGCIyYU/76eig1UfMmkZ7vwgOrzKzlQbFSuQfgm7GYUPpA==", - "cpu": [ - "ppc64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-riscv64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-riscv64/-/sharp-libvips-linux-riscv64-1.2.4.tgz", - "integrity": "sha512-oVDbcR4zUC0ce82teubSm+x6ETixtKZBh/qbREIOcI3cULzDyb18Sr/Wcyx7NRQeQzOiHTNbZFF1UwPS2scyGA==", - "cpu": [ - "riscv64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-s390x": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-s390x/-/sharp-libvips-linux-s390x-1.2.4.tgz", - "integrity": "sha512-qmp9VrzgPgMoGZyPvrQHqk02uyjA0/QrTO26Tqk6l4ZV0MPWIW6LTkqOIov+J1yEu7MbFQaDpwdwJKhbJvuRxQ==", - "cpu": [ - "s390x" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linux-x64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-x64/-/sharp-libvips-linux-x64-1.2.4.tgz", - "integrity": "sha512-tJxiiLsmHc9Ax1bz3oaOYBURTXGIRDODBqhveVHonrHJ9/+k89qbLl0bcJns+e4t4rvaNBxaEZsFtSfAdquPrw==", - "cpu": [ - "x64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linuxmusl-arm64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-arm64/-/sharp-libvips-linuxmusl-arm64-1.2.4.tgz", - "integrity": "sha512-FVQHuwx1IIuNow9QAbYUzJ+En8KcVm9Lk5+uGUQJHaZmMECZmOlix9HnH7n1TRkXMS0pGxIJokIVB9SuqZGGXw==", - "cpu": [ - "arm64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-libvips-linuxmusl-x64": { - "version": "1.2.4", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-x64/-/sharp-libvips-linuxmusl-x64-1.2.4.tgz", - "integrity": "sha512-+LpyBk7L44ZIXwz/VYfglaX/okxezESc6UxDSoyo2Ks6Jxc4Y7sGjpgU9s4PMgqgjj1gZCylTieNamqA1MF7Dg==", - "cpu": [ - "x64" - ], - "license": "LGPL-3.0-or-later", - "optional": true, - "os": [ - "linux" - ], - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-linux-arm": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm/-/sharp-linux-arm-0.34.5.tgz", - "integrity": "sha512-9dLqsvwtg1uuXBGZKsxem9595+ujv0sJ6Vi8wcTANSFpwV/GONat5eCkzQo/1O6zRIkh0m/8+5BjrRr7jDUSZw==", - "cpu": [ - "arm" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-arm": "1.2.4" - } - }, - "node_modules/@img/sharp-linux-arm64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm64/-/sharp-linux-arm64-0.34.5.tgz", - "integrity": "sha512-bKQzaJRY/bkPOXyKx5EVup7qkaojECG6NLYswgktOZjaXecSAeCWiZwwiFf3/Y+O1HrauiE3FVsGxFg8c24rZg==", - "cpu": [ - "arm64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-arm64": "1.2.4" - } - }, - "node_modules/@img/sharp-linux-ppc64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-ppc64/-/sharp-linux-ppc64-0.34.5.tgz", - "integrity": "sha512-7zznwNaqW6YtsfrGGDA6BRkISKAAE1Jo0QdpNYXNMHu2+0dTrPflTLNkpc8l7MUP5M16ZJcUvysVWWrMefZquA==", - "cpu": [ - "ppc64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-ppc64": "1.2.4" - } - }, - "node_modules/@img/sharp-linux-riscv64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-riscv64/-/sharp-linux-riscv64-0.34.5.tgz", - "integrity": "sha512-51gJuLPTKa7piYPaVs8GmByo7/U7/7TZOq+cnXJIHZKavIRHAP77e3N2HEl3dgiqdD/w0yUfiJnII77PuDDFdw==", - "cpu": [ - "riscv64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-riscv64": "1.2.4" - } - }, - "node_modules/@img/sharp-linux-s390x": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-s390x/-/sharp-linux-s390x-0.34.5.tgz", - "integrity": "sha512-nQtCk0PdKfho3eC5MrbQoigJ2gd1CgddUMkabUj+rBevs8tZ2cULOx46E7oyX+04WGfABgIwmMC0VqieTiR4jg==", - "cpu": [ - "s390x" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-s390x": "1.2.4" - } - }, - "node_modules/@img/sharp-linux-x64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-x64/-/sharp-linux-x64-0.34.5.tgz", - "integrity": "sha512-MEzd8HPKxVxVenwAa+JRPwEC7QFjoPWuS5NZnBt6B3pu7EG2Ge0id1oLHZpPJdn3OQK+BQDiw9zStiHBTJQQQQ==", - "cpu": [ - "x64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linux-x64": "1.2.4" - } - }, - "node_modules/@img/sharp-linuxmusl-arm64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-arm64/-/sharp-linuxmusl-arm64-0.34.5.tgz", - "integrity": "sha512-fprJR6GtRsMt6Kyfq44IsChVZeGN97gTD331weR1ex1c1rypDEABN6Tm2xa1wE6lYb5DdEnk03NZPqA7Id21yg==", - "cpu": [ - "arm64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linuxmusl-arm64": "1.2.4" - } - }, - "node_modules/@img/sharp-linuxmusl-x64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-x64/-/sharp-linuxmusl-x64-0.34.5.tgz", - "integrity": "sha512-Jg8wNT1MUzIvhBFxViqrEhWDGzqymo3sV7z7ZsaWbZNDLXRJZoRGrjulp60YYtV4wfY8VIKcWidjojlLcWrd8Q==", - "cpu": [ - "x64" - ], - "license": "Apache-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-libvips-linuxmusl-x64": "1.2.4" - } - }, - "node_modules/@img/sharp-wasm32": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-wasm32/-/sharp-wasm32-0.34.5.tgz", - "integrity": "sha512-OdWTEiVkY2PHwqkbBI8frFxQQFekHaSSkUIJkwzclWZe64O1X4UlUjqqqLaPbUpMOQk6FBu/HtlGXNblIs0huw==", - "cpu": [ - "wasm32" - ], - "license": "Apache-2.0 AND LGPL-3.0-or-later AND MIT", - "optional": true, - "dependencies": { - "@emnapi/runtime": "^1.7.0" - }, - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-win32-arm64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-arm64/-/sharp-win32-arm64-0.34.5.tgz", - "integrity": "sha512-WQ3AgWCWYSb2yt+IG8mnC6Jdk9Whs7O0gxphblsLvdhSpSTtmu69ZG1Gkb6NuvxsNACwiPV6cNSZNzt0KPsw7g==", - "cpu": [ - "arm64" - ], - "license": "Apache-2.0 AND LGPL-3.0-or-later", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-win32-ia32": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-ia32/-/sharp-win32-ia32-0.34.5.tgz", - "integrity": "sha512-FV9m/7NmeCmSHDD5j4+4pNI8Cp3aW+JvLoXcTUo0IqyjSfAZJ8dIUmijx1qaJsIiU+Hosw6xM5KijAWRJCSgNg==", - "cpu": [ - "ia32" - ], - "license": "Apache-2.0 AND LGPL-3.0-or-later", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - } - }, - "node_modules/@img/sharp-win32-x64": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-x64/-/sharp-win32-x64-0.34.5.tgz", - "integrity": "sha512-+29YMsqY2/9eFEiW93eqWnuLcWcufowXewwSNIT6UwZdUUCrM3oFjMWH/Z6/TMmb4hlFenmfAVbpWeup2jryCw==", - "cpu": [ - "x64" - ], - "license": "Apache-2.0 AND LGPL-3.0-or-later", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - } - }, "node_modules/@inquirer/ansi": { "version": "2.0.5", "resolved": "https://registry.npmjs.org/@inquirer/ansi/-/ansi-2.0.5.tgz", @@ -2268,140 +1801,6 @@ "@emnapi/runtime": "^1.7.1" } }, - "node_modules/@next/env": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.4.tgz", - "integrity": "sha512-dKkkOzOSwFYe5RX6y26fZgkSpVAlIOJKQHIiydQcrWH6y/97+RceSOAdjZ14Qa3zLduVUy0TXcn+EiM6t4rPgw==", - "license": "MIT" - }, - "node_modules/@next/swc-darwin-arm64": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.4.tgz", - "integrity": "sha512-OXTFFox5EKN1Ym08vfrz+OXxmCcEjT4SFMbNRsWZE99dMqt2Kcusl5MqPXcW232RYkMLQTy0hqgAMEsfEd/l2A==", - "cpu": [ - "arm64" - ], - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-darwin-x64": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.4.tgz", - "integrity": "sha512-XhpVnUfmYWvD3YrXu55XdcAkQtOnvaI6wtQa8fuF5fGoKoxIUZ0kWPtcOfqJEWngFF/lOS9l3+O9CcownhiQxQ==", - "cpu": [ - "x64" - ], - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-arm64-gnu": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.4.tgz", - "integrity": "sha512-Mx/tjlNA3G8kg14QvuGAJ4xBwPk1tUHq56JxZ8CXnZwz1Etz714soCEzGQQzVMz4bEnGPowzkV6Xrp6wAkEWOQ==", - "cpu": [ - "arm64" - ], - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-arm64-musl": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.4.tgz", - "integrity": "sha512-iVMMp14514u7Nup2umQS03nT/bN9HurK8ufylC3FZNykrwjtx7V1A7+4kvhbDSCeonTVqV3Txnv0Lu+m2oDXNg==", - "cpu": [ - "arm64" - ], - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-x64-gnu": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.4.tgz", - "integrity": "sha512-EZOvm1aQWgnI/N/xcWOlnS3RQBk0VtVav5Zo7n4p0A7UKyTDx047k8opDbXgBpHl4CulRqRfbw3QrX2w5UOXMQ==", - "cpu": [ - "x64" - ], - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-x64-musl": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.4.tgz", - "integrity": "sha512-h9FxsngCm9cTBf71AR4fGznDEDx1hS7+kSEiIRjq5kO1oXWm07DxVGZjCvk0SGx7TSjlUqhI8oOyz7NfwAdPoA==", - "cpu": [ - "x64" - ], - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-win32-arm64-msvc": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.4.tgz", - "integrity": "sha512-3NdJV5OXMSOeJYijX+bjaLge3mJBlh4ybydbT4GFoB/2hAojWHtMhl3CYlYoMrjPuodp0nzFVi4Tj2+WaMg+Ow==", - "cpu": [ - "arm64" - ], - "license": "MIT", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-win32-x64-msvc": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.4.tgz", - "integrity": "sha512-kMVGgsqhO5YTYODD9IPGGhA6iprWidQckK3LmPeW08PIFENRmgfb4MjXHO+p//d+ts2rpjvK5gXWzXSMrPl9cw==", - "cpu": [ - "x64" - ], - "license": "MIT", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } - }, "node_modules/@noble/ciphers": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/@noble/ciphers/-/ciphers-1.3.0.tgz", @@ -6452,15 +5851,6 @@ "react": "^18.0.0 || ^19.0.0" } }, - "node_modules/@swc/helpers": { - "version": "0.5.15", - "resolved": "https://registry.npmjs.org/@swc/helpers/-/helpers-0.5.15.tgz", - "integrity": "sha512-JQ5TuMi45Owi4/BIMAJBoSQoOJu12oOk/gADqlcUL9JEdHB8vyjUSsxqeNXnmXHjYKMi2WcYtezGEEhqUI/E2g==", - "license": "Apache-2.0", - "dependencies": { - "tslib": "^2.8.0" - } - }, "node_modules/@tabby_ai/hijri-converter": { "version": "1.0.5", "resolved": "https://registry.npmjs.org/@tabby_ai/hijri-converter/-/hijri-converter-1.0.5.tgz", @@ -8467,12 +7857,6 @@ "node": ">= 12" } }, - "node_modules/client-only": { - "version": "0.0.1", - "resolved": "https://registry.npmjs.org/client-only/-/client-only-0.0.1.tgz", - "integrity": "sha512-IV3Ou0jSMzZrd3pZ48nLkT9DA7Ag1pnPzaiQhpW7c3RbcqqzvzzVu+L8gfqMp/8IM2MQtSiqaCxrrcfu8I8rMA==", - "license": "MIT" - }, "node_modules/cliui": { "version": "8.0.1", "resolved": "https://registry.npmjs.org/cliui/-/cliui-8.0.1.tgz", @@ -13196,59 +12580,6 @@ "node": ">= 0.6" } }, - "node_modules/next": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/next/-/next-16.2.4.tgz", - "integrity": "sha512-kPvz56wF5frc+FxlHI5qnklCzbq53HTwORaWBGdT0vNoKh1Aya9XC8aPauH4NJxqtzbWsS5mAbctm4cr+EkQ2Q==", - "license": "MIT", - "dependencies": { - "@next/env": "16.2.4", - "@swc/helpers": "0.5.15", - "baseline-browser-mapping": "^2.9.19", - "caniuse-lite": "^1.0.30001579", - "postcss": "8.4.31", - "styled-jsx": "5.1.6" - }, - "bin": { - "next": "dist/bin/next" - }, - "engines": { - "node": ">=20.9.0" - }, - "optionalDependencies": { - "@next/swc-darwin-arm64": "16.2.4", - "@next/swc-darwin-x64": "16.2.4", - "@next/swc-linux-arm64-gnu": "16.2.4", - "@next/swc-linux-arm64-musl": "16.2.4", - "@next/swc-linux-x64-gnu": "16.2.4", - "@next/swc-linux-x64-musl": "16.2.4", - "@next/swc-win32-arm64-msvc": "16.2.4", - "@next/swc-win32-x64-msvc": "16.2.4", - "sharp": "^0.34.5" - }, - "peerDependencies": { - "@opentelemetry/api": "^1.1.0", - "@playwright/test": "^1.51.1", - "babel-plugin-react-compiler": "*", - "react": "^18.2.0 || 19.0.0-rc-de68d2f4-20241204 || ^19.0.0", - "react-dom": "^18.2.0 || 19.0.0-rc-de68d2f4-20241204 || ^19.0.0", - "sass": "^1.3.0" - }, - "peerDependenciesMeta": { - "@opentelemetry/api": { - "optional": true - }, - "@playwright/test": { - "optional": true - }, - "babel-plugin-react-compiler": { - "optional": true - }, - "sass": { - "optional": true - } - } - }, "node_modules/next-themes": { "version": "0.4.6", "resolved": "https://registry.npmjs.org/next-themes/-/next-themes-0.4.6.tgz", @@ -13838,34 +13169,6 @@ "points-on-curve": "0.2.0" } }, - "node_modules/postcss": { - "version": "8.4.31", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.4.31.tgz", - "integrity": "sha512-PS08Iboia9mts/2ygV3eLpY5ghnUcfLV/EXTOW1E2qYxJKGGBUtNjN76FYHnMs36RmARn41bC0AZmn+rR0OVpQ==", - "funding": [ - { - "type": "opencollective", - "url": "https://opencollective.com/postcss/" - }, - { - "type": "tidelift", - "url": "https://tidelift.com/funding/github/npm/postcss" - }, - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "dependencies": { - "nanoid": "^3.3.6", - "picocolors": "^1.0.0", - "source-map-js": "^1.0.2" - }, - "engines": { - "node": "^10 || ^12 || >=14" - } - }, "node_modules/postcss-selector-parser": { "version": "7.1.1", "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.1.tgz", @@ -13879,24 +13182,6 @@ "node": ">=4" } }, - "node_modules/postcss/node_modules/nanoid": { - "version": "3.3.12", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz", - "integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==", - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "bin": { - "nanoid": "bin/nanoid.cjs" - }, - "engines": { - "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" - } - }, "node_modules/powershell-utils": { "version": "0.1.0", "resolved": "https://registry.npmjs.org/powershell-utils/-/powershell-utils-0.1.0.tgz", @@ -15163,64 +14448,6 @@ "url": "https://github.com/sponsors/colinhacks" } }, - "node_modules/sharp": { - "version": "0.34.5", - "resolved": "https://registry.npmjs.org/sharp/-/sharp-0.34.5.tgz", - "integrity": "sha512-Ou9I5Ft9WNcCbXrU9cMgPBcCK8LiwLqcbywW3t4oDV37n1pzpuNLsYiAV8eODnjbtQlSDwZ2cUEeQz4E54Hltg==", - "hasInstallScript": true, - "license": "Apache-2.0", - "optional": true, - "dependencies": { - "@img/colour": "^1.0.0", - "detect-libc": "^2.1.2", - "semver": "^7.7.3" - }, - "engines": { - "node": "^18.17.0 || ^20.3.0 || >=21.0.0" - }, - "funding": { - "url": "https://opencollective.com/libvips" - }, - "optionalDependencies": { - "@img/sharp-darwin-arm64": "0.34.5", - "@img/sharp-darwin-x64": "0.34.5", - "@img/sharp-libvips-darwin-arm64": "1.2.4", - "@img/sharp-libvips-darwin-x64": "1.2.4", - "@img/sharp-libvips-linux-arm": "1.2.4", - "@img/sharp-libvips-linux-arm64": "1.2.4", - "@img/sharp-libvips-linux-ppc64": "1.2.4", - "@img/sharp-libvips-linux-riscv64": "1.2.4", - "@img/sharp-libvips-linux-s390x": "1.2.4", - "@img/sharp-libvips-linux-x64": "1.2.4", - "@img/sharp-libvips-linuxmusl-arm64": "1.2.4", - "@img/sharp-libvips-linuxmusl-x64": "1.2.4", - "@img/sharp-linux-arm": "0.34.5", - "@img/sharp-linux-arm64": "0.34.5", - "@img/sharp-linux-ppc64": "0.34.5", - "@img/sharp-linux-riscv64": "0.34.5", - "@img/sharp-linux-s390x": "0.34.5", - "@img/sharp-linux-x64": "0.34.5", - "@img/sharp-linuxmusl-arm64": "0.34.5", - "@img/sharp-linuxmusl-x64": "0.34.5", - "@img/sharp-wasm32": "0.34.5", - "@img/sharp-win32-arm64": "0.34.5", - "@img/sharp-win32-ia32": "0.34.5", - "@img/sharp-win32-x64": "0.34.5" - } - }, - "node_modules/sharp/node_modules/semver": { - "version": "7.7.4", - "resolved": "https://registry.npmjs.org/semver/-/semver-7.7.4.tgz", - "integrity": "sha512-vFKC2IEtQnVhpT78h1Yp8wzwrf8CM+MzKMHGJZfBtzhZNycRFnXsHk6E5TxIkkMsgNS7mdX3AGB7x2QM2di4lA==", - "license": "ISC", - "optional": true, - "bin": { - "semver": "bin/semver.js" - }, - "engines": { - "node": ">=10" - } - }, "node_modules/shebang-command": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/shebang-command/-/shebang-command-2.0.0.tgz", @@ -15601,29 +14828,6 @@ "inline-style-parser": "0.2.7" } }, - "node_modules/styled-jsx": { - "version": "5.1.6", - "resolved": "https://registry.npmjs.org/styled-jsx/-/styled-jsx-5.1.6.tgz", - "integrity": "sha512-qSVyDTeMotdvQYoHWLNGwRFJHC+i+ZvdBRYosOFgC+Wg1vx4frN2/RG/NA7SYqqvKNLf39P2LSRA2pu6n0XYZA==", - "license": "MIT", - "dependencies": { - "client-only": "0.0.1" - }, - "engines": { - "node": ">= 12.0.0" - }, - "peerDependencies": { - "react": ">= 16.8.0 || 17.x.x || ^18.0.0-0 || ^19.0.0-0" - }, - "peerDependenciesMeta": { - "@babel/core": { - "optional": true - }, - "babel-plugin-macros": { - "optional": true - } - } - }, "node_modules/stylis": { "version": "4.4.0", "resolved": "https://registry.npmjs.org/stylis/-/stylis-4.4.0.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index c69b2fdf3e..e8ede65526 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -64,7 +64,6 @@ "lucide-react": "^1.7.0", "mammoth": "^1.11.0", "motion": "^12.34.0", - "next": "^16.1.6", "next-themes": "^0.4.6", "node-forge": "^1.4.0", "radix-ui": "^1.4.3", From 9a81a5e8e7d049a4894dbe409a95831b8d095f06 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Fri, 15 May 2026 15:49:08 +0400 Subject: [PATCH 14/50] Update version-compat-ci.yml (#5445) --- .github/workflows/version-compat-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index 1ebea81066..2fbdd15747 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -214,7 +214,7 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - with: { path: unsloth } + path: unsloth - name: Clone unsloth-zoo @ main run: | # github.com occasionally 500s on the git fetch; retry so a From e81b942d2698e5d8a977f7412d338e1b1a72de67 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Fri, 15 May 2026 16:05:14 +0400 Subject: [PATCH 15/50] ci: merge duplicate `with:` keys in workflow checkout steps (#5447) Two `with:` mapping keys on the same step caused GitHub's workflow loader to reject the file (silently dropping persist-credentials: false under YAML "last key wins"). Merge into a single `with:` block in notebooks-ci.yml (3 sites) and version-compat-ci.yml (1 site). --- .github/workflows/notebooks-ci.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 0881c5ef3a..673b2f3cc5 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -200,7 +200,7 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - with: { path: unsloth } + path: unsloth - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: repository: unslothai/notebooks @@ -246,7 +246,7 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - with: { path: unsloth } + path: unsloth - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: repository: unslothai/notebooks @@ -352,7 +352,7 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - with: { path: unsloth } + path: unsloth - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: repository: unslothai/notebooks From 3f8c67263638281d91ac304f5ac44040992ae68a Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Fri, 15 May 2026 16:34:14 +0400 Subject: [PATCH 16/50] studio/chat: built-in web search for OpenAI, Anthropic, OpenRouter, Kimi (#5443) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio: wire the chat Search button to OpenAI's built-in web_search tool When the active model is an OpenAI external provider and the user clicks the existing Search pill in the composer, the chat-completion request now carries the unified enable_tools shorthand: enable_tools: true enabled_tools: ["web_search"] The backend's stream_chat_completion threads enabled_tools through to _stream_openai_responses, which translates it into the Responses API tool schema: body["tools"] = [{"type": "web_search"}] per the OpenAI Responses tool spec (https://developers.openai.com/api/docs/guides/tools). OpenAI then runs the search server-side before the model replies; the search- informed answer streams back through the existing response.output_text.delta path. web_search_call lifecycle events are silently ignored for now — sources / status indicators are follow-up scope. Frontend: - provider-capabilities.ts: new providerSupportsBuiltinWebSearch() helper. Returns true only for `openai` today; Anthropic (web_search_20250305), Gemini grounded-search, and OpenRouter variants can be added later with matching backend translation. - chat-page.tsx: both model-switch paths (the onChange handler and the inferenceParams.checkpoint useEffect) set supportsTools to match the new helper, and force toolsEnabled=false on every external switch so the Search toggle is opt-in by default. - chat-adapter.ts: external branch adds enable_tools + enabled_tools=["web_search"] to the request body when the toggle is on AND the active provider supports built-in web-search. Local-model branch is unchanged — it continues to route the same shorthand through our local tool runtime. Backend: - routes/inference.py: forwards payload.enabled_tools to stream_chat_completion at the proxy site (line 1599). - external_provider.py: stream_chat_completion gains an enabled_tools parameter; _stream_openai_responses appends {"type": "web_search"} to body["tools"] when the list contains "web_search". Other tools (file_search, code_interpreter, image_generation, computer_use_preview) are easy follow-ups in the same block. Reuses the existing pydantic ChatCompletionRequest.enabled_tools field, so no schema migrations. * studio/backend: surface OpenAI server-side web_search in the chat UI When the user has the chat Search button toggled on and OpenAI's /v1/responses invokes the built-in web_search tool, _stream_openai_responses now translates the tool's lifecycle events and citation annotations into the same _toolEvent shape that local-tool calls use. The result: the chat UI shows a web_search tool-call card mid-stream, then lists the cited sources at the end of the message — identical to how local web_search renders. SSE event translation: - response.output_item.added with item.type=web_search_call -> emit _toolEvent tool_start. Carries item.action.query as args when OpenAI ships it on the added event. - response.output_item.done with item.type=web_search_call -> backfill the query if it only arrives on the done variant. The existing reasoning branch on the same event is preserved as an if/elif under a shared isinstance guard. - response.output_text.annotation.added with type=url_citation -> collect into the most-recent web_search_call.citations list. - response.output_text.delta with inline annotations[] (older API variant) -> same collection path, so both wire shapes work. - response.completed -> emit _toolEvent tool_end per call with citations formatted as Title: \nURL: <url>\nSnippet: <snippet> blocks joined by `\n---\n`. The frontend's parseSourcesFromResult already lifts this format into source content parts at end-of-stream. - response.incomplete -> close out web_search cards with whatever citations had landed, so a truncated response does not leave a perpetually "running" tool card in the UI. Both reasoning and web_search work simultaneously on the same turn — the body sends `reasoning: {effort, summary}` and `tools: [{type: "web_search"}]` independently, and the SSE handler tracks them through separate channels. Diagnostic: finally-block logger now reports per stream web_search_requested - whether the client asked for it web_search_invocations - how many calls OpenAI actually made citations - total URLs cited queries - the search queries the model issued reasoning_emitted - whether <think> content was streamed so reports of "I clicked Search and nothing happened" can be triaged from the backend log without browser devtools. * studio/backend: fix empty query + per-card '(no sources cited)' on OpenAI web_search Two display bugs on the OpenAI Responses web_search → chat-UI bridge: 1. Tool cards showed "Searching for ''" — query missing. OpenAI's response.output_item.added for web_search_call does not reliably populate action.query across API versions; the canonical place is output_item.done. The previous code emitted tool_start at added with empty args and tried to backfill at done, but the frontend's _toolEvent: tool_start is a one-shot push (no update mechanism), so the args stayed empty. Fix: defer both tool_start *and* a placeholder tool_end emission to output_item.done, where action.query is guaranteed populated. added now just initialises tracking. Frontend then renders one card per call with the right "Searching for: <query>" label. 2. Every card showed "(no sources cited)". The previous code tried to attribute url_citation annotations to individual web_search_call invocations, but OpenAI's annotations carry no link back to a specific search call — they're just URLs the model cited from the aggregated search pool. With N invocations and M annotations, the previous logic bucketed all M into the last call and stamped "(no sources cited)" on the rest. Fix: collect citations into a single shared all_url_citations list, dedup by URL. At response.completed (and response.incomplete) overwrite the *last* web_search_call's tool_end result with the aggregated Title:/URL:/Snippet: blocks. The frontend's parseSourcesFromResult already flatMaps every web_search result, so one non-empty result is enough to surface the full source-pill set at the message tail. Other tool cards get an empty result string (no '(no sources)' text). Diagnostic log unchanged in shape; total_citations now reads len(all_url_citations) directly. * studio/chat: split Code and Search pill gates so external models cannot enable Code The previous wire-up set supportsTools=true for OpenAI external models to light up the Search pill, but supportsTools also gates the Code pill, so Code became clickable for OpenAI even though external providers have no local code execution. Separate the two gates so each pill reflects what's actually available: - chat-runtime-store: new `supportsBuiltinWebSearch: boolean` flag. Distinct from supportsTools — that one still means "runtime has a local tool sandbox" (Code, python, our DuckDuckGo web_search). This one means "the active external provider exposes a server-side web_search tool we can opt into" (OpenAI's /v1/responses today). - chat-page model-switch (both code paths): for external models, supportsTools is now forced to false (no local Code path) and supportsBuiltinWebSearch follows providerSupportsBuiltinWebSearch. Local-model paths are unaffected — they only set supportsTools. - shared-composer: Search pill gates on `searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch)`. Code pill gates on `codeDisabled = !modelLoaded || !supportsTools` — strictly the local runtime, so external models keep Code greyed out. A `toolsDisabled = codeDisabled` alias is left in place for any later-touched call site that may still reference the old name. No backend changes — chat-adapter already calls providerSupportsBuiltinWebSearch directly, independent of the store flags, so the request shape and the backend translation are unchanged. * studio/chat: default external reasoning effort to medium, not the carry-over When switching to an external model with reasoning support, the effort dropdown was inheriting whatever value the user had set on a prior model — frequently "xhigh" left over from a previous Opus/gpt-5 session. That meant every fresh OpenAI/Anthropic selection started at Extra High, burning tokens unintentionally. Both model-switch sites in chat-page (the useEffect on inferenceParams.checkpoint and the onChange callback) now pick "medium" whenever the new model's level list contains it, instead of the clamped carry-over. The clamp still fires as a fallback for the narrow case where a model doesn't expose medium (e.g. gpt-5.3-chat- latest which only has medium anyway — no change there). Users can still pick another level explicitly via the Think dropdown. * studio/chat: also light the Search pill in the welcome-screen composer There are two composers in the chat feature. shared-composer.tsx renders inside an active thread, and assistant-ui/thread.tsx has its own WebSearchToggle / CodeToolsToggle that ship the welcome-screen "Send a message…" composer (visible before the first user message). The previous fix split supportsTools and supportsBuiltinWebSearch in shared-composer but never touched the welcome-screen toggles in thread.tsx — they both still gated on supportsTools alone, so the Search pill stayed greyed on the welcome screen even for OpenAI external models that legitimately support web_search server-side. Mirror the shared-composer rule in WebSearchToggle: disabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch) CodeToolsToggle is left as-is — its current `disabled = !(modelLoaded && supportsTools)` is correct: external models have no local code-execution sandbox, so Code stays greyed when supportsTools=false (which is what chat-page now writes for external selections). * studio/backend: wire Anthropic server-side web_search end-to-end Mirrors the OpenAI web_search integration for Anthropic's web_search_20250305 tool. When the user toggles Search on with an Anthropic model selected, the request now carries the documented tool entry: tools: [{type: "web_search_20250305", name: "web_search", max_uses: 5}] on /v1/messages, and the SSE translation surfaces tool cards + source pills in the chat UI exactly the same way as OpenAI. stream_chat_completion now forwards enabled_tools into the Anthropic branch (was only doing this for the OpenAI Responses branch). _stream_anthropic gains an enabled_tools parameter and the web_search request-body block plus three additional event handlers: - content_block_start with type=server_tool_use, name=web_search: start tracking a new call. id becomes the tool_call_id. - content_block_delta with type=input_json_delta inside a server_tool_use block: buffer the partial_json so we can read out the search query when the block closes. - content_block_start with type=web_search_tool_result: capture the per-call result list (urls + titles) that Anthropic ships inline. - content_block_stop: closes whichever block we're inside — * server_tool_use -> emit _toolEvent: tool_start with the parsed query as args. * web_search_tool_result -> emit _toolEvent: tool_end with Title:/URL: blocks the frontend's parseSourcesFromResult lifts into source pills. * thinking block -> existing </think> close. Unlike OpenAI we get per-call results directly, so no aggregated- last-call fallback is needed — each tool card carries its own citations. Diagnostic log on stream completion now reports web_search_requested / invocations / total_results / queries, matching the OpenAI shape. Frontend providerSupportsBuiltinWebSearch returns true for 'anthropic' as well, so the Search pill lights up on Claude models the same way it does on OpenAI. The existing chat-adapter external branch already sends enabled_tools=['web_search'] based on this helper — no adapter changes needed. * studio: wire OpenRouter built-in web search via :online model suffix OpenRouter exposes a universal "add web search to any model" shortcut: append `:online` to the model id and the gateway runs the search server-side, streaming citations back as annotations on text deltas. Documented at https://openrouter.ai/docs/features/web-search Hook the existing Search toggle into that path: Backend (external_provider.py, default OAI-compat branch): - When provider_type == 'openrouter' and enabled_tools contains 'web_search', rewrite body['model']: openai/gpt-4o -> openai/gpt-4o:online anthropic/claude-sonnet-4-5:free -> anthropic/claude-sonnet-4-5:online Any existing `:variant` (`:free`, `:nitro`, etc.) is replaced — OpenRouter variants are mutually exclusive. - `openrouter/free` is skipped: it's a meta-router and `:online` is not a valid suffix on it (the gateway 400s). - A one-line INFO log fires whenever the rewrite happens so the diagnostic backend log shows exactly which model id the request was promoted to. Frontend (provider-capabilities.ts): - providerSupportsBuiltinWebSearch now returns true for 'openrouter' alongside 'openai' and 'anthropic'. The Search pill lights up and the existing chat-adapter external branch already forwards enabled_tools=['web_search'] based on this helper — no adapter changes needed. No new SSE event handling: OpenRouter does not emit a separate web_search_call event the way OpenAI/Anthropic do. Citations come back as text annotations via the existing reasoning_details path the adapter already parses, so source data flows through without extra translation. A per-call tool-card UX ("Searching for: …") would require synthesizing one client-side; deferred to a follow-up if the bare-citation flow feels too minimal. * studio: wire Mistral built-in web search connector Same shape as OpenAI's web_search tool, lives on /v1/chat/completions instead of /v1/responses. When the chat Search pill is toggled on with a Mistral model selected, the backend now appends {"type": "web_search"} to body["tools"] before the request goes out. Idempotent — won't double-append if a future call site adds it first. Models in the registry allowlist that don't support the connector (codestral, devstral, ministral, mistral-tiny) will surface a 400 from upstream; the existing default-path error log captures it. Mistral's docs: https://docs.mistral.ai/capabilities/agents/connectors/websearch Frontend providerSupportsBuiltinWebSearch returns true for 'mistral' now, alongside openai / anthropic / openrouter. The Search pill lights up for Mistral models and the existing adapter branch already sends enabled_tools=['web_search'] off this helper — no adapter changes. No SSE translation yet — Mistral streams citations inline as text annotations or `references` in the final assistant content, not as a separate web_search_call event. Citations flow through to the message body as text; a per-call tool-card UX with "Searching for: …" indicators is a follow-up if needed. * studio/backend: fix OpenRouter web_search to use plugins shape + synthesize tool card Two changes against the actual OpenRouter docs at https://openrouter.ai/docs/guides/features/plugins/web-search: Request shape: The previous commit appended :online to the model id, which works on concrete model ids but rejects on meta-routers like openrouter/free — and that's exactly the model the user was testing with, so neither the request rewrite nor the diagnostic log fired. Switch to the universal plugins shape: body["plugins"] = [{"id": "web"}] Per the docs this is "exactly equivalent" to :online but works on every model id including openrouter/free and openrouter/auto. No model suffix manipulation, idempotent if added twice. Tool-card synthesis: OpenRouter doesn't emit a structured web_search_call event the way OpenAI/Anthropic do — citations come back only as `annotations` of type=url_citation on delta/message objects. To match the chat-UI tool-card UX the user expects ("Searching for: …" indicator, source pills at message tail), synthesize the events client-side in the default OAI-compat stream loop: - On stream open (after the 200 status check): yield a synthetic _toolEvent: tool_start with tool_name=web_search, fixed id "openrouter_web_search". The chat-UI then renders the running tool card before any text streams. - During the SSE loop: scan every chunk's choices[].delta and choices[].message for `annotations: [{type: "url_citation", url_citation: {url, title, content}}]` entries. Dedup by URL into a citations list. Handles both the nested-url_citation shape OpenRouter documents and the flat-on-annotation shape some upstreams ship. - On [DONE] (or stream-close without [DONE]): emit synthetic tool_end carrying the citations as Title: …\nURL: …\nSnippet: …\n---\n… blocks the existing parseSourcesFromResult lifts into source pills at message tail. Diagnostic log on completion now also reports web_search_requested + citation count alongside the existing chosen-model / event-count telemetry. * studio: drop Mistral built-in web_search — connector lives on Agents API only Mistral's web_search is exclusively on /v1/agents + /v1/conversations; sending it on /v1/chat/completions returns "WebSearchTool connector is not supported". Wiring it would require a dedicated Agents streaming path. Remove from the frontend capability map and revert the chat-completions tool injection. * studio: wire Kimi $web_search builtin via two-call round-trip Kimi's $web_search lives on /v1/chat/completions but requires a client round-trip per https://platform.kimi.ai/docs/guide/use-web-search: the first call returns tool_calls with function.arguments populated; the caller echoes those arguments back as a role=tool message; the second call streams the final answer with search results incorporated. The docs also mandate thinking=disabled while the builtin is active. Backend: new _stream_kimi_web_search helper dispatched from stream_chat_completion when provider_type=='kimi' and 'web_search' in enabled_tools. Buffers tool_calls across deltas, falls back to a plain stream if the model declines to search, and synthesizes tool_start (with parsed query) / tool_end (with any url_citation annotations) so the chat UI's web-search card behaves the same as other providers. Frontend: kimi added to providerSupportsBuiltinWebSearch so the Search pill lights up in the composer. * studio/chat: mutual exclusion of Think + Search on Kimi composer Kimi's $web_search builtin requires thinking=disabled per https://platform.kimi.ai/docs/guide/use-web-search, so the two states cannot coexist. Make the pills mutually exclusive in both composers (shared and welcome-screen): clicking Search turns Think off; clicking Think back on turns Search off. Default Think to on when a Kimi model is selected — k2.6/k2.5 ship with thinking enabled out of the box. * studio/chat: fix wrong provider var name in onChange branch selectedProvider, not provider — TS2304 in tsc -b. * studio/backend: add diagnostics to Kimi $web_search round-trip Log the actual function.arguments from the first call (so we can see the model's search query) and the second call's usage.prompt_tokens + any annotation type names that came through. prompt_tokens spiking above the input message length is direct proof the server injected search results into context. annotation_types lets us learn the shape Kimi uses for citations if/when they emit any. * studio: per-provider defaults — Anthropic xhigh + Search on, OpenAI high + Search on, Opus 4.7 gains max Anthropic: Think effort defaults to the highest level the model supports (xhigh on 4.6/4.7, high on 4.5) and Search starts on, since the web_search_20250305 tool returns structured citations end-to-end. OpenAI: Think effort defaults to 'high' (the gpt-5.x reasoning sweet spot for /v1/responses + web_search) and Search starts on. Opus 4.7: 'max' added as an effort level above 'xhigh' in both backend (_ANTHROPIC_THINKING_SPECS) and frontend (ANTHROPIC_REASONING_MODELS). Kimi diagnostics: emit tool_end immediately after tool_start so the web-search card transitions to 'complete' before the second-call answer streams, log first-call args + second-call usage/prompt_tokens + any annotation type names, request stream_options.include_usage so the second call exposes usage in SSE. * studio/backend: harden Kimi fallback path with HTTPError handler + manual aiter_lines loop Addresses PR review feedback (#5443): the no-search fallback streaming path was using `async for response.aiter_lines()` and had no `httpx.HTTPError` guard around the POST. Switch to the manual __anext__ loop pattern used elsewhere in this module (avoids the Python 3.13 + httpcore 1.0.x GeneratorExit propagation issue) and wrap the whole request in a try/except so network failures surface as a proper SSE error frame instead of a raw traceback. --- .../core/inference/external_provider.py | 1007 ++++++++++++++++- studio/backend/routes/inference.py | 1 + .../src/components/assistant-ui/thread.tsx | 45 +- .../src/features/chat/api/chat-adapter.ts | 15 + .../frontend/src/features/chat/chat-page.tsx | 113 +- .../features/chat/provider-capabilities.ts | 40 +- .../src/features/chat/shared-composer.tsx | 56 +- .../chat/stores/chat-runtime-store.ts | 12 + 8 files changed, 1258 insertions(+), 31 deletions(-) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index f5b67eef70..6937518b5d 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -41,7 +41,7 @@ _ANTHROPIC_THINKING_SPECS = ( _AnthropicThinkingSpec( prefixes = ("claude-opus-4-7",), kind = "adaptive", - efforts = ("none", "low", "medium", "high", "xhigh"), + efforts = ("none", "low", "medium", "high", "xhigh", "max"), ), _AnthropicThinkingSpec( prefixes = ("claude-opus-4-6", "claude-sonnet-4-6"), @@ -141,6 +141,34 @@ def _apply_mistral_reasoning_controls( _http_client = httpx.AsyncClient() +def _build_kimi_tool_end( + synthetic_chunk_fn: Any, + tool_call_id: str, + citations: list[dict[str, str]], +) -> str: + """Format Kimi web_search citations into the tool_end payload. + + Same shape parseSourcesFromResult on the frontend expects for the + other built-in web_search providers: `Title: ...\\nURL: ...\\n + Snippet: ...\\n---\\n...`. If no citations were emitted, fall back + to a generic "(search complete)" string so the UI still shows the + tool card transitioning to a completed state. + """ + blocks: list[str] = [] + for cit in citations: + line = f"Title: {cit['title']}\nURL: {cit['url']}" + if cit.get("snippet"): + line += f"\nSnippet: {cit['snippet']}" + blocks.append(line) + return synthetic_chunk_fn( + { + "type": "tool_end", + "tool_call_id": tool_call_id, + "result": "\n---\n".join(blocks) if blocks else "(search complete)", + } + ) + + class ExternalProviderClient: """Async proxy for OpenAI-compatible external LLM APIs.""" @@ -199,6 +227,7 @@ class ExternalProviderClient: top_k: Optional[int] = None, enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, + enabled_tools: Optional[list[str]] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: """ @@ -222,6 +251,7 @@ class ExternalProviderClient: top_k, enable_thinking, reasoning_effort, + enabled_tools, ): yield line return @@ -240,6 +270,28 @@ class ExternalProviderClient: max_tokens, enable_thinking, reasoning_effort, + enabled_tools, + ): + yield line + return + + # Kimi's $web_search is a builtin_function that requires a client + # round-trip: the first call returns a tool_calls envelope with + # function.arguments populated; the caller echoes those arguments + # back as a role=tool message; the second call streams the final + # answer with the search incorporated. The doc also mandates + # disabling thinking while $web_search is active. Route to a + # dedicated helper so the default OAI-compat path stays single-pass. + # https://platform.kimi.ai/docs/guide/use-web-search + if ( + self.provider_type == "kimi" + and enabled_tools + and "web_search" in enabled_tools + ): + async for line in self._stream_kimi_web_search( + messages, + model, + max_tokens, ): yield line return @@ -317,6 +369,29 @@ class ExternalProviderClient: else: body["reasoning"] = {"enabled": False} + # OpenRouter web-search plugin — universal shape that works + # for every model id, including the `openrouter/free` and + # `openrouter/auto` meta-routers. Documented at + # https://openrouter.ai/docs/guides/features/plugins/web-search + # The `:online` model-suffix shortcut is "exactly equivalent + # to" this plugin per the same doc, but only works on + # concrete model ids — meta-routers reject the suffix. + # `plugins: [{id: "web"}]` works everywhere, no model id + # rewrite needed, and idempotent if some future call site + # adds the entry first. + if enabled_tools and "web_search" in enabled_tools: + plugins = list(body.get("plugins") or []) + if not any( + isinstance(p, dict) and p.get("id") == "web" for p in plugins + ): + plugins.append({"id": "web"}) + body["plugins"] = plugins + logger.info( + "OpenRouter web_search: attached plugins=[{id: 'web'}] " + "(model=%s)", + body.get("model"), + ) + url = f"{self.base_url}/chat/completions" logger.info( "Proxying chat completion to %s (provider=%s, model=%s)", @@ -362,6 +437,94 @@ class ExternalProviderClient: # error" in the UI with no trail on the server side. event_counts: dict[str, int] = {} chosen_model: Optional[str] = None + # Web-search tool-card synthesis for OpenRouter. The gateway + # doesn't emit structured web_search_call events — citations + # come back as `annotations` of type=url_citation on delta / + # message objects. Mirror the OpenAI/Anthropic UX by yielding + # a synthetic tool_start at stream open and tool_end at + # stream close with the collected citation list. + web_search_active = ( + self.provider_type == "openrouter" + and bool(enabled_tools) + and "web_search" in (enabled_tools or []) + ) + web_search_tool_id = "openrouter_web_search" + web_search_citations: list[dict[str, str]] = [] + web_search_tool_started = False + web_search_tool_ended = False + + def _emit_synthetic_tool_event(payload: dict[str, Any]) -> str: + chunk = { + "id": f"chatcmpl-{self.provider_type}-synthetic", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + def _record_or_url_citation(payload: Any) -> None: + if not isinstance(payload, dict): + return + if payload.get("type") != "url_citation": + return + # OpenRouter (and OpenAI Chat Completions web_search) + # nest the citation under url_citation; some variants + # ship the fields flat on the annotation itself. Accept + # both. + cit = payload.get("url_citation") + if not isinstance(cit, dict): + cit = payload + url = cit.get("url", "") if isinstance(cit, dict) else "" + if not url or not isinstance(url, str): + return + if any(c["url"] == url for c in web_search_citations): + return + title = cit.get("title") or url + snippet = cit.get("content") or cit.get("snippet") or "" + web_search_citations.append( + { + "url": url, + "title": title, + "snippet": snippet if isinstance(snippet, str) else "", + } + ) + + def _build_web_search_tool_end() -> str: + blocks: list[str] = [] + for cit in web_search_citations: + line = f"Title: {cit['title']}\nURL: {cit['url']}" + if cit.get("snippet"): + line += f"\nSnippet: {cit['snippet']}" + blocks.append(line) + return _emit_synthetic_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": ( + "\n---\n".join(blocks) + if blocks + else "(search complete)" + ), + } + ) + + if web_search_active: + yield _emit_synthetic_tool_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": web_search_tool_id, + "arguments": {}, + } + ) + web_search_tool_started = True + try: while True: try: @@ -374,6 +537,17 @@ class ExternalProviderClient: data_str = line[len("data:") :].strip() if data_str == "[DONE]": event_counts["done"] = event_counts.get("done", 0) + 1 + # Emit synthetic tool_end with collected + # citations BEFORE forwarding [DONE], so the + # tool-card transitions to "complete" in the + # UI before the stream closes. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + yield _build_web_search_tool_end() + web_search_tool_ended = True elif data_str: try: parsed = _json.loads(data_str) @@ -406,17 +580,52 @@ class ExternalProviderClient: parsed.get("model"), str ): chosen_model = parsed["model"] + # When the user has web_search on, scan + # every chunk's delta and message + # objects for url_citation annotations. + # Different OpenRouter upstreams place + # them in different spots. + if web_search_active: + choices = parsed.get("choices") or [] + if isinstance(choices, list): + for choice in choices: + if not isinstance(choice, dict): + continue + for envelope in ( + choice.get("delta"), + choice.get("message"), + ): + if not isinstance(envelope, dict): + continue + for ann in ( + envelope.get("annotations") + or [] + ): + _record_or_url_citation(ann) yield line + # Stream ended without [DONE] (some upstreams just close + # the connection). Emit tool_end so the card doesn't + # stay in "running" forever. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + yield _build_web_search_tool_end() + web_search_tool_ended = True except GeneratorExit: await response.aclose() # set PoolByteStream._closed=True FIRST await lines_gen.aclose() # now safe — aclose() is a no-op raise finally: logger.info( - "%s stream complete (model=%s, chosen=%s, events=%s)", + "%s stream complete (model=%s, chosen=%s, " + "web_search_requested=%s, citations=%s, events=%s)", self.provider_type, model, chosen_model, + web_search_active, + len(web_search_citations), event_counts, ) await response.aclose() @@ -444,6 +653,384 @@ class ExternalProviderClient: self.provider_type, ) + async def _stream_kimi_web_search( + self, + messages: list[dict[str, Any]], + model: str, + max_tokens: Optional[int], + ) -> AsyncGenerator[str, None]: + """ + Kimi $web_search round-trip. + + Wire flow (per https://platform.kimi.ai/docs/guide/use-web-search): + 1. POST messages with tools=[{type: "builtin_function", + function: {name: "$web_search"}}] and thinking=disabled. + 2. Stream the first response — accumulate function.arguments + across tool_call deltas until finish_reason="tool_calls". + Do NOT forward those tool_call chunks to the client (they + are an internal protocol step, not user-visible output). + 3. Build a second request: original messages + the assistant + message carrying the tool_calls + a role=tool message that + echoes the same arguments back verbatim (per Kimi docs, + the caller "just needs to submit tool_call.function.arguments + to Kimi as they are" — the server actually runs the search). + 4. Stream the second response — that is the final answer the + user sees, with search results already incorporated. + + We synthesize tool_start (with the parsed query) when step (2) + completes, and tool_end (with any url_citation annotations the + second stream emits) before [DONE], so the chat UI shows the + same web-search tool card as the other providers. + """ + url = f"{self.base_url}/chat/completions" + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": True, + # $web_search forbids thinking; sending the toggle silently + # would have the server reject the request with 400. + "thinking": {"type": "disabled"}, + "tools": [ + {"type": "builtin_function", "function": {"name": "$web_search"}} + ], + } + if max_tokens is not None: + body["max_tokens"] = max_tokens + + # Strip body fields the Kimi registry declares unusable + # (temperature/top_p — see body_omit in providers.py). + from core.inference.providers import get_provider_info + + provider_info = get_provider_info(self.provider_type) or {} + for field in provider_info.get("body_omit", ()): + body.pop(field, None) + + tool_call_id = "kimi_web_search" + synthetic_id = f"chatcmpl-{self.provider_type}-synthetic" + + def _synthetic_chunk(payload: dict[str, Any]) -> str: + chunk = { + "id": synthetic_id, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + logger.info( + "Kimi $web_search round-trip starting (model=%s, url=%s)", + model, + url, + ) + + # ---- First call: collect the model's $web_search tool_call ---- + tool_calls_acc: dict[int, dict[str, Any]] = {} + try: + async with _http_client.stream( + "POST", + url, + json = body, + headers = self._auth_headers(), + timeout = self._stream_timeout, + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors = "replace") + logger.error( + "Kimi first-call returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + lines_gen = response.aiter_lines().__aiter__() + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line.strip() or not line.startswith("data:"): + continue + data_str = line[len("data:") :].strip() + if data_str == "[DONE]": + break + try: + parsed = _json.loads(data_str) + except Exception: + continue + for choice in parsed.get("choices") or []: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") or {} + for tc in delta.get("tool_calls") or []: + if not isinstance(tc, dict): + continue + idx = tc.get("index", 0) + slot = tool_calls_acc.setdefault( + idx, + { + "id": tc.get("id") or f"call_{idx}", + "type": "function", + "function": {"name": "", "arguments": ""}, + }, + ) + if tc.get("id"): + slot["id"] = tc["id"] + fn = tc.get("function") or {} + if fn.get("name"): + slot["function"]["name"] = fn["name"] + if fn.get("arguments"): + slot["function"]["arguments"] += fn["arguments"] + if choice.get("finish_reason") == "tool_calls": + break + except GeneratorExit: + await response.aclose() + await lines_gen.aclose() + raise + finally: + await response.aclose() + await lines_gen.aclose() + except httpx.HTTPError as exc: + logger.error("Kimi first-call HTTP error: %s", exc) + yield _error_sse_line( + 502, + f"Error communicating with kimi: {exc}", + self.provider_type, + ) + return + + # If the model decided not to search, fall back to a plain + # streaming call without the builtin tool. That mirrors the UX + # of every other provider when web_search is on but the model + # didn't actually need it. + search_calls = [ + tc + for tc in tool_calls_acc.values() + if tc["function"]["name"] == "$web_search" + ] + if not search_calls: + logger.info( + "Kimi $web_search: model did not invoke search; " + "falling back to plain stream" + ) + fallback_body = dict(body) + fallback_body.pop("tools", None) + try: + async with _http_client.stream( + "POST", + url, + json = fallback_body, + headers = self._auth_headers(), + timeout = self._stream_timeout, + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors = "replace") + logger.error( + "Kimi fallback returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + # Manual __anext__ loop instead of `async for` — see the + # comment in stream_chat_completion for the Python 3.13 + + # httpcore 1.0.x GeneratorExit interaction this avoids. + lines_gen = response.aiter_lines().__aiter__() + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if line.strip(): + yield line + except GeneratorExit: + await response.aclose() + await lines_gen.aclose() + raise + finally: + await response.aclose() + await lines_gen.aclose() + except httpx.HTTPError as exc: + logger.error("Kimi fallback HTTP error: %s", exc) + yield _error_sse_line( + 502, + f"Error communicating with kimi: {exc}", + self.provider_type, + ) + return + + # Synthesize tool_start with the parsed search query so the + # chat UI's web-search card shows "Searching for: ...". + first_args_raw = search_calls[0]["function"]["arguments"] or "{}" + try: + first_args = _json.loads(first_args_raw) + except Exception: + first_args = {} + # Log the raw arguments so we can confirm the server actually + # ran the search. The shape is documented loosely but in practice + # the model emits `{"search_result":{"search_id":...}, + # "usage":{"total_tokens":N}}` — an opaque receipt where N is the + # token cost of the injected search context. The query string is + # NOT present; Kimi runs the search server-side during the first + # call and bakes the results straight into the model's context. + logger.info( + "Kimi $web_search: %d tool_call(s), args[0]=%s", + len(search_calls), + first_args_raw[:500], + ) + first_args_search_tokens: Optional[int] = None + if isinstance(first_args, dict): + usage_block = first_args.get("usage") + if isinstance(usage_block, dict): + tok = usage_block.get("total_tokens") + if isinstance(tok, int): + first_args_search_tokens = tok + yield _synthetic_chunk( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": tool_call_id, + "arguments": first_args if isinstance(first_args, dict) else {}, + } + ) + # Kimi's search has already executed server-side by the time the + # first call returns (the tool_call envelope encodes the search + # result reference, not a query for us to dispatch). Emit + # tool_end NOW so the UI's web-search card transitions to + # "complete" before the second call starts streaming the + # answer, instead of after — otherwise the card sits in + # "running" all the way through the answer streaming and the + # user perceives the model answering before search finishes. + yield _build_kimi_tool_end(_synthetic_chunk, tool_call_id, []) + + # ---- Second call: echo the tool_calls back and stream answer ---- + assistant_msg = { + "role": "assistant", + "content": "", + "tool_calls": list(tool_calls_acc.values()), + } + tool_msgs = [ + { + "role": "tool", + "tool_call_id": tc["id"], + "name": tc["function"]["name"], + "content": tc["function"]["arguments"], + } + for tc in tool_calls_acc.values() + ] + followup_body = dict(body) + followup_body["messages"] = list(messages) + [assistant_msg] + tool_msgs + # Ask the SSE stream to include a final `usage` block so we can + # see prompt_tokens (which jumps to thousands when the server + # injects search context). Without this, OpenAI-compat streams + # omit usage entirely. Kimi follows the same convention. + followup_body["stream_options"] = {"include_usage": True} + # Keep the tool definition on the second call so the model can + # decide to search again mid-turn if needed. Kimi's doc shows + # the same tools array on every step. + + try: + async with _http_client.stream( + "POST", + url, + json = followup_body, + headers = self._auth_headers(), + timeout = self._stream_timeout, + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors = "replace") + logger.error( + "Kimi second-call returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + lines_gen = response.aiter_lines().__aiter__() + # Diagnostics: latch usage.prompt_tokens from the final + # chunk. The Kimi docs say search results count toward + # prompt_tokens, so a big value here is direct evidence + # the server actually injected results into context. + last_usage: Optional[dict[str, Any]] = None + annotation_shapes: set[str] = set() + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line.strip(): + continue + if line.startswith("data:"): + data_str = line[len("data:") :].strip() + if data_str and data_str != "[DONE]": + try: + parsed = _json.loads(data_str) + except Exception: + parsed = None + if isinstance(parsed, dict): + usage = parsed.get("usage") + if isinstance(usage, dict): + last_usage = usage + # Scan annotations only for diagnostics — + # Kimi today doesn't emit url_citation, but + # if a future model version starts to we'll + # see the type name in the final log line + # and can wire it into the tool_end payload. + for choice in parsed.get("choices") or []: + if not isinstance(choice, dict): + continue + for envelope in ( + choice.get("delta"), + choice.get("message"), + ): + if not isinstance(envelope, dict): + continue + for ann in ( + envelope.get("annotations") or [] + ): + if isinstance(ann, dict): + annotation_shapes.add( + str(ann.get("type") or "?") + ) + yield line + except GeneratorExit: + await response.aclose() + await lines_gen.aclose() + raise + finally: + logger.info( + "Kimi $web_search complete (model=%s, " + "search_ctx_tokens=%s, annotation_types=%s, " + "prompt_tokens=%s, completion_tokens=%s)", + model, + first_args_search_tokens, + sorted(annotation_shapes) or None, + (last_usage or {}).get("prompt_tokens"), + (last_usage or {}).get("completion_tokens"), + ) + await response.aclose() + await lines_gen.aclose() + except httpx.HTTPError as exc: + logger.error("Kimi second-call HTTP error: %s", exc) + yield _error_sse_line( + 502, + f"Error communicating with kimi: {exc}", + self.provider_type, + ) + async def _stream_anthropic( self, messages: list[dict[str, Any]], @@ -454,6 +1041,7 @@ class ExternalProviderClient: top_k: Optional[int] = None, enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, + enabled_tools: Optional[list[str]] = None, ) -> AsyncGenerator[str, None]: """ Call the Anthropic Messages API and translate its SSE to OpenAI format. @@ -602,6 +1190,25 @@ class ExternalProviderClient: if body.get("max_tokens", 0) <= budget_tokens: body["max_tokens"] = budget_tokens + 1024 + # Anthropic server-side web_search — see + # https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/web-search-tool + # The tool type is date-pinned (web_search_20250305 today) and + # Anthropic dispatches search calls server-side, returning + # server_tool_use + web_search_tool_result blocks in the SSE + # stream, plus url-citation annotations on text deltas. We + # translate all of that into our local _toolEvent shape so the + # chat UI renders web_search exactly like OpenAI's path. + if enabled_tools and "web_search" in enabled_tools: + anthropic_tools = list(body.get("tools") or []) + anthropic_tools.append( + { + "type": "web_search_20250305", + "name": "web_search", + "max_uses": 5, + } + ) + body["tools"] = anthropic_tools + url = f"{self.base_url}/messages" completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}" @@ -659,6 +1266,17 @@ class ExternalProviderClient: # "no thinking content" — distinguishes "Anthropic never sent # thinking_delta" from "frontend didn't render the chunks". event_counts: dict[str, int] = {} + # web_search state. Anthropic emits the query inside an + # `input_json_delta` stream on a `server_tool_use` content + # block, then a separate `web_search_tool_result` block + # with the URL list. Unlike OpenAI we get per-call results + # directly, so each tool card carries its own citations. + # `current_server_tool_use`: {id, name, partial_json_buffer} + # `current_result_block`: {tool_use_id, results} + # Both go to None when the matching content_block_stop fires. + current_server_tool_use: Optional[dict[str, Any]] = None + current_result_block: Optional[dict[str, Any]] = None + web_search_calls: dict[str, dict[str, Any]] = {} def _content_chunk(text: str) -> str: chunk = { @@ -674,6 +1292,37 @@ class ExternalProviderClient: } return f"data: {_json.dumps(chunk)}" + def _emit_tool_event(payload: dict[str, Any]) -> str: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + def _format_web_search_results( + results: list[Any], + ) -> str: + blocks: list[str] = [] + for r in results: + if not isinstance(r, dict): + continue + if r.get("type") != "web_search_result": + continue + url = r.get("url", "") + title = r.get("title") or url + if not url: + continue + blocks.append(f"Title: {title}\nURL: {url}") + return "\n---\n".join(blocks) + try: while True: try: @@ -702,7 +1351,39 @@ class ExternalProviderClient: key = event_type or "<unknown>" event_counts[key] = event_counts.get(key, 0) + 1 - if event_type == "content_block_delta": + if event_type == "content_block_start": + content_block = event.get("content_block") or {} + block_type = content_block.get("type") + if ( + block_type == "server_tool_use" + and content_block.get("name") == "web_search" + ): + tool_use_id = content_block.get("id", "") or ( + f"ws_{len(web_search_calls)}" + ) + current_server_tool_use = { + "id": tool_use_id, + "buffer": "", + } + web_search_calls[tool_use_id] = { + "query": "", + "results": [], + } + elif block_type == "web_search_tool_result": + tool_use_id = content_block.get("tool_use_id", "") + # Anthropic sometimes ships the full results + # list on the start event; sometimes deltas + # follow. Capture whatever is present and + # finalize on content_block_stop. + content = content_block.get("content") or [] + current_result_block = { + "tool_use_id": tool_use_id, + "results": list(content) + if isinstance(content, list) + else [], + } + + elif event_type == "content_block_delta": delta = event.get("delta", {}) delta_type = delta.get("type") if delta_type == "thinking_delta": @@ -730,16 +1411,79 @@ class ExternalProviderClient: text = delta.get("text", "") if text: yield _content_chunk(text) + # Citations on text deltas are attached + # per-call by Anthropic via the + # `web_search_tool_result` block; we don't + # need to scrape them off the text events. + elif ( + delta_type == "input_json_delta" + and current_server_tool_use is not None + ): + # Streamed partial_json carrying the search + # query. Buffer until content_block_stop. + current_server_tool_use["buffer"] += delta.get( + "partial_json", "" + ) # signature_delta and any other delta types are # intentionally skipped — they carry trust / # verification metadata, not user-visible content. elif event_type == "content_block_stop": - # Close the <think> tag when the thinking block - # ends, in case no text_delta follows (e.g. - # display=omitted on Claude 4.7, or thinking-only - # turns). - if thinking_open: + if current_server_tool_use is not None: + # End of the server_tool_use block — parse the + # accumulated input_json into a query and + # emit tool_start. The matching tool_end fires + # later when the web_search_tool_result block + # closes with the actual results. + buffer = current_server_tool_use["buffer"] + query = "" + if buffer: + try: + parsed = _json.loads(buffer) + if isinstance(parsed, dict): + q = parsed.get("query", "") + if isinstance(q, str): + query = q + except Exception: + query = "" + tool_use_id = current_server_tool_use["id"] + if tool_use_id in web_search_calls: + web_search_calls[tool_use_id]["query"] = query + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": tool_use_id, + "arguments": ( + {"query": query} if query else {} + ), + } + ) + current_server_tool_use = None + elif current_result_block is not None: + # End of a web_search_tool_result — emit + # tool_end carrying the search results as + # Title:/URL: blocks. parseSourcesFromResult + # on the frontend lifts these into source + # pills at message tail. + tool_use_id = current_result_block["tool_use_id"] + results = current_result_block["results"] + if tool_use_id in web_search_calls: + web_search_calls[tool_use_id]["results"] = results + result_text = _format_web_search_results(results) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": tool_use_id, + "result": (result_text or "(search complete)"), + } + ) + current_result_block = None + elif thinking_open: + # Close the <think> tag when the thinking block + # ends, in case no text_delta follows (e.g. + # display=omitted on Claude 4.7, or thinking- + # only turns). yield _content_chunk("</think>") thinking_open = False @@ -778,16 +1522,30 @@ class ExternalProviderClient: await lines_gen.aclose() # now safe — aclose() is a no-op raise finally: - # Surface per-event-type counts so reports of "no - # reasoning panel content" can be triaged at a glance: - # zero `content_block_delta:thinking_delta` entries - # means Anthropic skipped thinking for this prompt - # (adaptive can choose to); non-zero means thinking - # arrived and we wrapped it — any visual gap is then - # on the frontend. + # Surface per-event-type counts + web_search summary so + # reports of "no reasoning panel content" / "Search + # didn't do anything" can be triaged at a glance. + web_search_requested = bool( + enabled_tools and "web_search" in enabled_tools + ) + web_search_invocations = len(web_search_calls) + total_results = sum( + len(sc.get("results") or []) for sc in web_search_calls.values() + ) + queries = [ + sc["query"] + for sc in web_search_calls.values() + if sc.get("query") + ] logger.info( - "Anthropic stream event counts (model=%s): %s", + "Anthropic stream complete (model=%s, " + "web_search_requested=%s, web_search_invocations=%s, " + "results=%s, queries=%s, events=%s)", model, + web_search_requested, + web_search_invocations, + total_results, + queries, event_counts, ) await response.aclose() @@ -824,6 +1582,7 @@ class ExternalProviderClient: max_tokens: Optional[int], enable_thinking: Optional[bool], reasoning_effort: Optional[str], + enabled_tools: Optional[list[str]] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -918,6 +1677,20 @@ class ExternalProviderClient: if max_tokens is not None: body["max_output_tokens"] = max_tokens + # OpenAI server-side tools — see + # https://developers.openai.com/api/docs/guides/tools + # The frontend's Search button maps to the unified + # enabled_tools=["web_search"] shorthand; translate that into the + # Responses-API tool schema. Other built-in tools (file_search, + # code_interpreter, image_generation, computer_use_preview) can be + # added with the same pattern when we surface their toggles. + if enabled_tools: + tools_array: list[dict[str, Any]] = [] + if "web_search" in enabled_tools: + tools_array.append({"type": "web_search"}) + if tools_array: + body["tools"] = tools_array + url = f"{self.base_url}/responses" completion_id = f"chatcmpl-openai-{model.replace('/', '-')}" @@ -950,6 +1723,64 @@ class ExternalProviderClient: done_emitted = False reasoning_open = False reasoning_emitted = False + # Per-call state for OpenAI's server-side web_search tool. Mapped + # back into our local _toolEvent shape so the existing chat-UI + # renderer surfaces web_search the same way it does for local + # tool calls: a "Searching…" tool-call card, then a `tool_end` + # carrying citations formatted as + # Title: …\nURL: …\nSnippet: …\n---\n… + # blocks (which the frontend's parseSourcesFromResult lifts + # into source content parts at end of stream). + # web_search_calls preserves insertion order so we can apply + # the aggregated citation list onto the *last* call's + # tool_end — that's the one the frontend's source-pill + # extraction reads (parseSourcesFromResult flatMaps every + # web_search result, so a single non-empty result is enough + # to surface all sources at message tail). + # OpenAI emits url_citation annotations on text deltas, not + # per call — there's no wire field linking a citation back + # to a specific search invocation. Hence the shared list. + # web_search_calls: { item_id -> {query} } + web_search_calls: dict[str, dict[str, Any]] = {} + all_url_citations: list[dict[str, str]] = [] + + def _emit_tool_event(payload: dict[str, Any]) -> str: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + def _record_url_citation(payload: dict[str, Any]) -> None: + """Append a url_citation onto the shared all_url_citations + list. Dedup by URL — the same source can be cited multiple + times across deltas. We do NOT try to attribute citations + to individual web_search_call invocations because OpenAI's + annotation events don't carry that linkage.""" + if payload.get("type") != "url_citation": + return + url = payload.get("url", "") + if not url: + return + if any(c["url"] == url for c in all_url_citations): + return + title = payload.get("title") or url + snippet = payload.get("snippet") or payload.get("quote") or "" + all_url_citations.append( + { + "url": url, + "title": title, + "snippet": snippet, + } + ) def _extract_reasoning_text(payload: Any) -> str: if payload is None: @@ -1023,13 +1854,39 @@ class ExternalProviderClient: yield _chunk_with_text("</think>") reasoning_open = False yield _chunk_with_text(delta_text) + # Some API versions inline url citations on the + # delta event itself rather than as a separate + # response.output_text.annotation.added event. + for ann in event.get("annotations") or []: + if isinstance(ann, dict): + _record_url_citation(ann) - elif event_type == "response.output_item.done": + elif event_type == "response.output_text.annotation.added": + ann = event.get("annotation") + if isinstance(ann, dict): + _record_url_citation(ann) + + elif event_type == "response.output_item.added": + # Track the call early but do NOT emit tool_start + # yet — action.query is not reliably populated on + # added across OpenAI API versions, and the + # frontend's tool_start is a one-shot push (no + # update mechanism). Wait for output_item.done. item = event.get("item", {}) if ( isinstance(item, dict) - and item.get("type") == "reasoning" + and item.get("type") == "web_search_call" ): + item_id = item.get("id", "") or ( + f"ws_{len(web_search_calls)}" + ) + web_search_calls.setdefault(item_id, {"query": ""}) + + elif event_type == "response.output_item.done": + item = event.get("item", {}) + if not isinstance(item, dict): + continue + if item.get("type") == "reasoning": summary_text = _extract_reasoning_text( item.get("summary") ) @@ -1039,6 +1896,46 @@ class ExternalProviderClient: reasoning_open = True yield _chunk_with_text(summary_text) reasoning_emitted = True + elif item.get("type") == "web_search_call": + # done is the canonical place to read the + # query, so emit both tool_start and tool_end + # here. Frontend then renders a card per call + # with the proper "Searching: <query>" label. + # Citations are aggregated separately and the + # *last* call's result is overwritten at + # response.completed with the citation list + # (so the source-pill extraction at message + # tail surfaces them once). + item_id = item.get("id", "") or ( + f"ws_{len(web_search_calls)}" + ) + action = item.get("action") + query = ( + action.get("query", "") + if isinstance(action, dict) + else "" + ) + web_search_calls[item_id] = {"query": query} + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": item_id, + "arguments": ( + {"query": query} if query else {} + ), + } + ) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": item_id, + # Empty result — the last call gets + # overwritten with citations at + # response.completed. + "result": "", + } + ) elif isinstance(event_type, str) and "reasoning" in event_type: reasoning_delta = _extract_reasoning_text(event) @@ -1053,6 +1950,32 @@ class ExternalProviderClient: if reasoning_open: yield _chunk_with_text("</think>") reasoning_open = False + # Apply the aggregated citation list onto the + # *last* web_search call by overwriting its + # tool_end result. The frontend's + # parseSourcesFromResult flatMaps every + # web_search tool-call result, so a single + # non-empty result is enough to surface the + # whole source-pill set at the message tail — + # no need to fan out across every card (which + # would just duplicate the same pills). + if web_search_calls and all_url_citations: + last_id = list(web_search_calls.keys())[-1] + blocks: list[str] = [] + for cit in all_url_citations: + line = ( + f"Title: {cit['title']}\n" f"URL: {cit['url']}" + ) + if cit.get("snippet"): + line += f"\nSnippet: {cit['snippet']}" + blocks.append(line) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": last_id, + "result": "\n---\n".join(blocks), + } + ) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -1070,6 +1993,29 @@ class ExternalProviderClient: if reasoning_open: yield _chunk_with_text("</think>") reasoning_open = False + # Same backfill as response.completed — apply + # whatever citations we managed to gather + # before truncation onto the last call. All + # earlier tool cards already have their proper + # query + empty placeholder result from the + # output_item.done emissions above. + if web_search_calls and all_url_citations: + last_id = list(web_search_calls.keys())[-1] + blocks = [] + for cit in all_url_citations: + line = ( + f"Title: {cit['title']}\n" f"URL: {cit['url']}" + ) + if cit.get("snippet"): + line += f"\nSnippet: {cit['snippet']}" + blocks.append(line) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": last_id, + "result": "\n---\n".join(blocks), + } + ) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -1103,6 +2049,31 @@ class ExternalProviderClient: await lines_gen.aclose() raise finally: + # Summarise what the model actually did this turn so + # support reports of "I clicked Search and got nothing" + # can be triaged at a glance: was the tool requested, + # did OpenAI invoke it, and how many sources came back? + web_search_requested = bool( + enabled_tools and "web_search" in enabled_tools + ) + web_search_invocations = len(web_search_calls) + total_citations = len(all_url_citations) + queries = [ + sc["query"] + for sc in web_search_calls.values() + if sc.get("query") + ] + logger.info( + "OpenAI Responses stream complete (model=%s, " + "web_search_requested=%s, web_search_invocations=%s, " + "citations=%s, queries=%s, reasoning_emitted=%s)", + model, + web_search_requested, + web_search_invocations, + total_citations, + queries, + reasoning_emitted, + ) await response.aclose() await lines_gen.aclose() diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 59928be3cf..55d74d1cbe 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1595,6 +1595,7 @@ async def _proxy_to_external_provider( top_k = payload.top_k, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, + enabled_tools = payload.enabled_tools, stream = payload.stream, ) try: diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index d75a8cce10..4988cd5a46 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -496,6 +496,9 @@ const ReasoningToggle: FC = () => { externalSelection != null ? externalProviders.find((p) => p.id === externalSelection.providerId) : undefined; + const isKimiExternal = selectedExternalProvider?.providerType === "kimi"; + const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled); + const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled); const effectiveExternalModelId = selectedExternalProvider?.providerType === "openrouter" && externalSelection?.modelId === "openrouter/free" && @@ -587,6 +590,11 @@ const ReasoningToggle: FC = () => { setReasoningEffort(level); setReasoningEnabled(true); applyQwenThinkingParams(true); + // Kimi's $web_search builtin forbids thinking, so + // enabling thinking flips the Search pill off. + if (isKimiExternal && toolsEnabled) { + setToolsEnabled(false); + } }} > {formatEffortLabel(level)} @@ -613,6 +621,11 @@ const ReasoningToggle: FC = () => { const next = !reasoningEnabled; setReasoningEnabled(next); applyQwenThinkingParams(next); + // Mutual exclusion with the Search pill on Kimi — see the + // dropdown branch above and shared-composer for the same rule. + if (isKimiExternal && next && toolsEnabled) { + setToolsEnabled(false); + } }} className="composer-pill-btn" data-active={ @@ -680,16 +693,44 @@ const WebSearchToggle: FC = () => { const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, ); + const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const supportsTools = useChatRuntimeStore((s) => s.supportsTools); + // External providers (OpenAI today) expose a server-side web_search tool + // even when the local tool runtime is unavailable — gate the Search pill + // on either source so it lights up on external models too. Mirror of + // shared-composer's searchDisabled. + const supportsBuiltinWebSearch = useChatRuntimeStore( + (s) => s.supportsBuiltinWebSearch, + ); const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled); const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled); - const disabled = !(modelLoaded && supportsTools); + const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled); + const externalProviders = useExternalProvidersStore((s) => s.providers); + const externalSelection = parseExternalModelId(checkpoint); + const selectedExternalProvider = + externalSelection != null + ? externalProviders.find((p) => p.id === externalSelection.providerId) + : undefined; + const isKimiExternal = selectedExternalProvider?.providerType === "kimi"; + const disabled = + !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); return ( <button type="button" disabled={disabled} - onClick={() => setToolsEnabled(!toolsEnabled)} + onClick={() => { + const next = !toolsEnabled; + setToolsEnabled(next); + // Kimi's $web_search builtin requires thinking=disabled (see + // https://platform.kimi.ai/docs/guide/use-web-search). Keep + // the two pills mutually exclusive so the visible state always + // matches what the backend ends up sending. + if (isKimiExternal) { + setReasoningEnabled(!next); + applyQwenThinkingParams(!next); + } + }} className="composer-pill-btn" data-active={toolsEnabled && !disabled ? "true" : "false"} aria-label={toolsEnabled ? "Disable web search" : "Enable web search"} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 3e8cec0443..30b256f8ba 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -35,6 +35,7 @@ import { getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, + providerSupportsBuiltinWebSearch, } from "../provider-capabilities"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import { isMultimodalResponse } from "../types/api"; @@ -1005,6 +1006,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(externalCapabilities?.presencePenalty ? { presence_penalty: params.presencePenalty } : {}), + // Built-in web search: when the user has the Search toggle + // on AND the active provider supports a server-side + // web_search tool (currently OpenAI's /v1/responses), pass + // the enable_tools shorthand. Backend translates + // enabled_tools=["web_search"] into the provider's tool + // schema — for OpenAI that's `tools: [{type:"web_search"}]` + // on the Responses body, see _stream_openai_responses. + ...(toolsEnabled && + providerSupportsBuiltinWebSearch(externalProvider.providerType) + ? { + enable_tools: true, + enabled_tools: ["web_search"], + } + : {}), provider_id: externalProvider.id, provider_type: externalBackendProviderType, external_model: externalSelection.modelId, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 29428b2ef9..3fa294d465 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -50,6 +50,7 @@ import { clampReasoningEffortToLevels, getExternalReasoningCapabilities, getProviderCapabilities, + providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; import { @@ -678,9 +679,56 @@ export function ChatPage(): ReactElement { preferredEffort, effortLevels, ); + // Per-provider default effort. Anthropic gets the highest available + // level (xhigh on 4.6/4.7, high on 4.5) since Claude's adaptive + // thinking adjusts cost per turn — sitting at the top of the dial + // gives users the strongest answers and the model can still skip + // thinking when the turn is trivial. OpenAI gets "high" by default + // — the gpt-5.x reasoning models accept high across the board and + // it's the right cost/quality sweet spot for Responses-API tools + // (web search included). Everyone else gets "medium" as a balanced + // default. Users can pick another level via the Think dropdown. + const isAnthropic = provider?.providerType === "anthropic"; + const isOpenAI = provider?.providerType === "openai"; + const anthropicTopEffort = effortLevels.includes("xhigh") + ? "xhigh" + : effortLevels.includes("high") + ? "high" + : clampedEffort; + const openaiDefaultEffort = effortLevels.includes("high") + ? "high" + : effortLevels.includes("medium") + ? "medium" + : clampedEffort; const nextReasoningEffort = reasoningCaps.supportsReasoning - ? clampedEffort + ? isAnthropic + ? anthropicTopEffort + : isOpenAI + ? openaiDefaultEffort + : effortLevels.includes("medium") + ? "medium" + : clampedEffort : state.reasoningEffort; + const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( + provider?.providerType, + ); + // Kimi's k2.6/k2.5 default to thinking enabled on the server side + // (per https://platform.kimi.ai/docs/models). Mirror that default + // in the UI so the Think pill comes up clicked when the user picks + // a Kimi model. The Search pill stays off by default; the mutual- + // exclusion handlers in the composer flip the two when needed. + const isKimi = provider?.providerType === "kimi"; + // Web search is on by default for the two providers we trust most + // for it: Anthropic (web_search_20250305 server tool, structured + // citations) and OpenAI (/v1/responses web_search, structured + // citations). Other providers stay off-by-default — OpenRouter's + // plugins shape and Kimi's $web_search builtin still work when the + // user opts in via the pill, but they're a notch less reliable so + // we don't pre-enable them. + const searchOnByDefault = + supportsBuiltinWebSearch && + (provider?.providerType === "anthropic" || + provider?.providerType === "openai"); useChatRuntimeStore.setState({ supportsReasoning: reasoningCaps.supportsReasoning, reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn, @@ -690,10 +738,22 @@ export function ChatPage(): ReactElement { reasoningEffort: nextReasoningEffort, reasoningEnabled: reasoningCaps.supportsReasoning ? reasoningCaps.supportsReasoningOff - ? state.reasoningEnabled + ? isKimi + ? true + : state.reasoningEnabled : true : state.reasoningEnabled, supportsPreserveThinking: false, + // External models never give us a local tool runtime (no Code + // execution, no python sandbox), so `supportsTools` must be + // false — that's what gates the Code pill in the composer. + // `supportsBuiltinWebSearch` is the separate flag that lets the + // Search pill light up for providers (currently just OpenAI) who + // run web_search server-side. + supportsTools: false, + supportsBuiltinWebSearch, + toolsEnabled: searchOnByDefault, + codeToolsEnabled: false, }); }, [externalProviders, inferenceParams.checkpoint]); const canCompare = useMemo(() => { @@ -810,8 +870,29 @@ export function ChatPage(): ReactElement { preferredEffort, effortLevels, ); + // Same per-provider default policy as the useEffect path above: + // Anthropic picks the highest available level, OpenAI picks + // "high", everyone else picks "medium". + const isAnthropic = selectedProvider?.providerType === "anthropic"; + const isOpenAI = selectedProvider?.providerType === "openai"; + const anthropicTopEffort = effortLevels.includes("xhigh") + ? "xhigh" + : effortLevels.includes("high") + ? "high" + : clampedEffort; + const openaiDefaultEffort = effortLevels.includes("high") + ? "high" + : effortLevels.includes("medium") + ? "medium" + : clampedEffort; const nextReasoningEffort = reasoningCaps.supportsReasoning - ? clampedEffort + ? isAnthropic + ? anthropicTopEffort + : isOpenAI + ? openaiDefaultEffort + : effortLevels.includes("medium") + ? "medium" + : clampedEffort : store.reasoningEffort; // Clear any cached router-picked openrouter/free model unless the // user is staying on openrouter/free — otherwise the chip would @@ -823,6 +904,20 @@ export function ChatPage(): ReactElement { ...store.params, checkpoint: value, }); + const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( + selectedProvider?.providerType, + ); + // See sibling useEffect above: Kimi's k2.x default to thinking + // enabled, so the Think pill comes up clicked. Search pill stays + // off by default; mutual exclusion flips them via the composer. + const isKimi = selectedProvider?.providerType === "kimi"; + // Mirror of sibling useEffect: Anthropic and OpenAI get Search + // on-by-default since their server tools emit structured + // citations end-to-end. OpenRouter and Kimi stay off-by-default. + const searchOnByDefault = + supportsBuiltinWebSearch && + (selectedProvider?.providerType === "anthropic" || + selectedProvider?.providerType === "openai"); useChatRuntimeStore.setState({ activeGgufVariant: null, ggufContextLength: null, @@ -837,10 +932,20 @@ export function ChatPage(): ReactElement { reasoningEffort: nextReasoningEffort, reasoningEnabled: reasoningCaps.supportsReasoning ? reasoningCaps.supportsReasoningOff - ? store.reasoningEnabled + ? isKimi + ? true + : store.reasoningEnabled : true : store.reasoningEnabled, supportsPreserveThinking: false, + // External models have no local tool runtime → supportsTools=false + // keeps the Code pill greyed out. supportsBuiltinWebSearch is the + // separate flag the composer reads to light up the Search pill + // when the provider offers a server-side web_search tool. + supportsTools: false, + supportsBuiltinWebSearch, + toolsEnabled: searchOnByDefault, + codeToolsEnabled: false, ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }), }); return; diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 0688d2d673..44da42e27b 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -83,6 +83,44 @@ export function clampReasoningEffortToLevels( */ export const EXTERNAL_MAX_OUTPUT_TOKENS = 32768; +/** + * Whether the external provider offers a built-in web-search tool that the + * model invokes server-side. When `true`, the chat composer's Search button + * is available for that provider and the chat-adapter forwards + * `enable_tools: true, enabled_tools: ["web_search"]` on the request — the + * backend routes the call through the provider's tool schema: + * - OpenAI: `tools: [{type: "web_search"}]` on /v1/responses + * - Anthropic: `tools: [{type: "web_search_20250305", name: "web_search", + * max_uses: 5}]` on /v1/messages + * - OpenRouter: `plugins: [{id: "web"}]` on /v1/chat/completions (the + * router's universal web-search shape; works for every + * underlying model including the `openrouter/free` router). + * - Kimi: `tools: [{type: "builtin_function", function: {name: + * "$web_search"}}]` with `thinking: {type: + * "disabled"}`. Requires a client round-trip: + * the first call returns the search args; the backend + * echoes them back as a role=tool message; the second + * call streams the answer. Handled in + * _stream_kimi_web_search on the backend. + * + * Mistral is intentionally excluded: their `web_search` connector lives on + * the Agents API (`/v1/agents` + `/v1/conversations`), not chat completions, + * and returns `"WebSearchTool connector is not supported"` if injected into + * /v1/chat/completions. Wiring it would require a dedicated Agents streaming + * path. Gemini's grounded-search can be added with the same pattern when + * matching backend translation lands. + */ +export function providerSupportsBuiltinWebSearch( + providerType: string | null | undefined, +): boolean { + return ( + providerType === "openai" || + providerType === "anthropic" || + providerType === "openrouter" || + providerType === "kimi" + ); +} + /** * Per-provider minimum on the outbound max_tokens. Kimi's docs require * `max_tokens >= 16000` whenever a thinking model is in use so the @@ -239,7 +277,7 @@ const NO_REASONING_CAPS: ReasoningCaps = { const ANTHROPIC_REASONING_MODELS = [ { prefixes: ["claude-opus-4-7"], - levels: ["none", "low", "medium", "high", "xhigh"], + levels: ["none", "low", "medium", "high", "xhigh", "max"], }, { prefixes: ["claude-opus-4-6", "claude-sonnet-4-6"], diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index ceef9f501b..18f0673433 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -304,6 +304,9 @@ export function SharedComposer({ const preserveThinking = useChatRuntimeStore((s) => s.preserveThinking); const setPreserveThinking = useChatRuntimeStore((s) => s.setPreserveThinking); const supportsTools = useChatRuntimeStore((s) => s.supportsTools); + const supportsBuiltinWebSearch = useChatRuntimeStore( + (s) => s.supportsBuiltinWebSearch, + ); const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled); const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled); const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled); @@ -345,13 +348,31 @@ export function SharedComposer({ const reasoningLockedOn = effectiveSupportsReasoning && (effectiveReasoningAlwaysOn || !effectiveSupportsReasoningOff); + // Kimi's $web_search builtin mandates thinking=disabled per the docs at + // https://platform.kimi.ai/docs/guide/use-web-search. Both pills stay + // clickable for Kimi, but turning one on flips the other off — the + // click handlers below enforce this mutual exclusion so the visible + // state always matches what the backend actually sends. + const isKimiExternal = selectedExternalProvider?.providerType === "kimi"; const effectiveReasoningEnabled = reasoningLockedOn ? true : reasoningEnabled; const effectiveReasoningVisualEnabled = effectiveReasoningEnabled && reasoningEffort !== "none"; const reasoningDisabled = !modelLoaded || !effectiveSupportsReasoning; const showReasoningControl = effectiveSupportsReasoning || effectiveReasoningAlwaysOn; - const toolsDisabled = !modelLoaded || !supportsTools; + // Two-pill gating: Search pill lights up when the runtime has either + // a local tool runtime (supportsTools, gives us our Code/python + local + // web_search) OR a server-side web_search the provider runs for us + // (supportsBuiltinWebSearch, currently just OpenAI's /v1/responses). + // Code pill is gated on `supportsTools` only — external providers + // never give us code execution, so the pill must stay disabled even + // when Search is available. + const searchDisabled = + !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); + const codeDisabled = !modelLoaded || !supportsTools; + // Backwards-compatible alias for any other call site that may still + // reference `toolsDisabled` (rare; both pills used it before). + const toolsDisabled = codeDisabled; const setPendingAudioStore = useChatRuntimeStore((s) => s.setPendingAudio); const clearPendingAudioStore = useChatRuntimeStore((s) => s.clearPendingAudio); @@ -766,6 +787,11 @@ export function SharedComposer({ setReasoningEffort(level); setReasoningEnabled(true); applyQwenThinkingParams(true); + // Mutual exclusion: turning thinking on for a + // Kimi model forces the web_search builtin off. + if (isKimiExternal && toolsEnabled) { + setToolsEnabled(false); + } }} > {formatReasoningEffortLabel(level, externalSelection?.modelId)} @@ -789,6 +815,12 @@ export function SharedComposer({ const next = !reasoningEnabled; setReasoningEnabled(next); applyQwenThinkingParams(next); + // Mutual exclusion: Kimi's $web_search builtin + // requires thinking off, so turning thinking on flips + // the Search pill off (and vice versa). + if (isKimiExternal && next && toolsEnabled) { + setToolsEnabled(false); + } }} className={cn( "flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors", @@ -845,10 +877,22 @@ export function SharedComposer({ )} <button type="button" - disabled={toolsDisabled} - onClick={() => setToolsEnabled(!toolsEnabled)} + disabled={searchDisabled} + onClick={() => { + const next = !toolsEnabled; + setToolsEnabled(next); + // Kimi's $web_search builtin requires thinking=disabled + // (https://platform.kimi.ai/docs/guide/use-web-search). + // Toggle the Think pill off when Search comes on, and + // back on when Search goes off — mutual exclusion that + // mirrors what the backend enforces. + if (isKimiExternal) { + setReasoningEnabled(!next); + applyQwenThinkingParams(!next); + } + }} className="composer-pill-btn" - data-active={toolsEnabled && !toolsDisabled ? "true" : "false"} + data-active={toolsEnabled && !searchDisabled ? "true" : "false"} aria-label={toolsEnabled ? "Disable web search" : "Enable web search"} > <GlobeIcon className="size-3.5" /> @@ -856,10 +900,10 @@ export function SharedComposer({ </button> <button type="button" - disabled={toolsDisabled} + disabled={codeDisabled} onClick={() => setCodeToolsEnabled(!codeToolsEnabled)} className="composer-pill-btn" - data-active={codeToolsEnabled && !toolsDisabled ? "true" : "false"} + data-active={codeToolsEnabled && !codeDisabled ? "true" : "false"} aria-label={codeToolsEnabled ? "Disable code execution" : "Enable code execution"} > <CodeToggleIcon className="size-3.5" /> diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 92ce1441ae..4cbed5899b 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -229,6 +229,16 @@ type ChatRuntimeStore = { supportsPreserveThinking: boolean; preserveThinking: boolean; supportsTools: boolean; + /** + * Whether the active external provider exposes a server-side + * web_search tool (OpenAI's /v1/responses today). Distinct from + * `supportsTools` — that flag governs the local tool runtime (Code, + * python sandbox, our DuckDuckGo web_search). This one only enables + * the chat composer's Search pill for external models and leaves + * the Code pill disabled, because external providers do not give + * us code execution. Local models keep `supportsTools` only. + */ + supportsBuiltinWebSearch: boolean; toolsEnabled: boolean; codeToolsEnabled: boolean; toolStatus: string | null; @@ -320,6 +330,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({ supportsPreserveThinking: false, preserveThinking: loadBool(PRESERVE_THINKING_KEY, false), supportsTools: false, + supportsBuiltinWebSearch: false, toolsEnabled: false, codeToolsEnabled: false, toolStatus: null, @@ -430,6 +441,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({ reasoningEffortLevels: ["low", "medium", "high"], supportsPreserveThinking: false, supportsTools: false, + supportsBuiltinWebSearch: false, toolsEnabled: false, codeToolsEnabled: false, toolStatus: null, From e0e606a24a96d8053dbce3adbf9bf71ce2d2f70a Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 05:35:19 -0700 Subject: [PATCH 17/50] ci: make compiler-cache shim test order-independent (#5449) The shim test_compile_real_modeling_module[*] was failing on all three RMSNorm families (llama / qwen3 / gemma3) on the Core 4.57.6 matrix cell because the preceding test_compile_every_transformers_ model_type sweep already invokes unsloth_compile_transformers for every model_type, which sets modeling.__UNSLOTH_PATCHED__ = True. unsloth_zoo.compiler.unsloth_compile_transformers (zoo compiler.py :3318-3324) early-returns when that marker is already set, without re-emitting the cache file. The targeted shim test then asserts the file exists and fails with "compiler did not write" against the temp cache path. Drop the unsloth-added marker (and any leftover cache file from the sweep) before invoking the compile so the test exercises a fresh emit regardless of collection order. Marker-only fix -- transformers version-agnostic (works on 4.57.6 + 5.x); does not touch zoo internals. --- .github/workflows/consolidated-tests-ci.yml | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 2b0ce43a7c..897367de0c 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -1032,16 +1032,28 @@ jobs: """Spot-check on the three production-relevant families that the compile_every sweep also covers; this case verifies the emitted cache file has the model-specific RMSNorm class - attribute, not just that the file parses + imports.""" + attribute, not just that the file parses + imports. + + Note on test isolation: ``unsloth_compile_transformers`` + early-returns when ``modeling.__UNSLOTH_PATCHED__`` is set, + so once an earlier test in the same collection patches the + module the next call won't re-emit the cache file. Drop the + marker (and any stale cache file) before invoking so this + test is order-independent.""" import importlib as _il try: - _il.import_module( + modeling = _il.import_module( f"transformers.models.{model_type}.modeling_{model_type}" ) except ModuleNotFoundError: pytest.skip( f"transformers build lacks model_type={model_type}" ) + if hasattr(modeling, "__UNSLOTH_PATCHED__"): + delattr(modeling, "__UNSLOTH_PATCHED__") + combined = _CACHE / f"unsloth_compiled_module_{model_type}.py" + if combined.exists(): + combined.unlink() unsloth_compile_transformers( model_type=model_type, fast_lora_forwards=False, ) @@ -1049,7 +1061,6 @@ jobs: f"transformers.models.{model_type}.modeling_{model_type}" ) assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True - combined = _CACHE / f"unsloth_compiled_module_{model_type}.py" _verify_file(combined, must_expose=[rms_class]) From 4999753514160f7d88f90172800d8c1be2c64018 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 15 May 2026 14:13:28 +0100 Subject: [PATCH 18/50] Studio: o3 reasoning summary payload (#5426) * fix: o3 reasoning summary payload * fix: omit reasoning.summary for o3 in enable_thinking branch --------- Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- .../core/inference/external_provider.py | 12 +++- .../test_openai_responses_translation.py | 62 +++++++++++++++++++ 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 6937518b5d..bd14e73d16 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -29,6 +29,7 @@ logger = structlog.get_logger(__name__) # still accept it. Match the 4-7 line specifically so we keep the knob # live on every other Claude generation. _ANTHROPIC_TOP_K_DEPRECATED = re.compile(r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)") +_OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)") class _AnthropicThinkingSpec(NamedTuple): @@ -1659,6 +1660,9 @@ class ExternalProviderClient: # to wrap, and the chat reasoning panel stays blank. Always pair # an explicit effort with summary except for the explicit "off" # case (effort: "none"), where summaries are pointless. + summary_unsupported = bool( + _OPENAI_REASONING_SUMMARY_UNSUPPORTED.match(model.strip().lower()) + ) if reasoning_effort in ( "minimal", "low", @@ -1667,11 +1671,15 @@ class ExternalProviderClient: "max", "xhigh", ): - body["reasoning"] = {"effort": reasoning_effort, "summary": "auto"} + body["reasoning"] = {"effort": reasoning_effort} + if not summary_unsupported: + body["reasoning"]["summary"] = "auto" elif reasoning_effort == "none" or enable_thinking is False: body["reasoning"] = {"effort": "none"} elif enable_thinking is True: - body["reasoning"] = {"effort": "medium", "summary": "auto"} + body["reasoning"] = {"effort": "medium"} + if not summary_unsupported: + body["reasoning"]["summary"] = "auto" if instructions_parts: body["instructions"] = "\n\n".join(instructions_parts) if max_tokens is not None: diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py index 4ad6a19ea9..22ccba7058 100644 --- a/studio/backend/tests/test_openai_responses_translation.py +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -286,6 +286,68 @@ def test_responses_reasoning_effort_included_when_requested(monkeypatch): assert captured["body"]["reasoning"] == {"effort": "high", "summary": "auto"} +def test_responses_reasoning_summary_omitted_for_o3(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _responses_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "o3", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = "high", + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "high"} + + +def test_responses_reasoning_summary_omitted_for_o3_with_enable_thinking(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _responses_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "o3", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = True, + reasoning_effort = None, + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "medium"} + + def test_responses_reasoning_effort_none_omits_summary(monkeypatch): captured: dict = {} From 7e90cae345ac51d81f61012e9ecba813e286e450 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 07:46:36 -0700 Subject: [PATCH 19/50] ci: compiler-cache-shim must mutate live module globals + skip rerun (#5452) The shim test pinned UNSLOTH_COMPILE_LOCATION via env before importing unsloth_zoo.compiler, but tests/conftest.py runs `import unsloth` first, which transitively imports unsloth_zoo.compiler with the default cache path. The shim's later env-set never took effect on the captured module global, so the compiler silently wrote artefacts to the default cache and the per-model file assertion failed under Core (HF=4.57.6 + TRL<1). Two fixes: 1) After import, mutate the live module globals directly (UNSLOTH_COMPILE_LOCATION, UNSLOTH_COMPILE_USE_TEMP) so they reflect the hermetic tmp dir regardless of who imported the module first. The same pattern is already used in _compiler_cache_invariants_shim._isolate_cache. 2) test_compile_real_modeling_module no longer re-runs unsloth_compile_transformers after a sweep already patched the module. The compile is not idempotent in-process: re-running on a module whose class forwards were already rewritten corrupts the inspect source/line cache and the second-pass emitted file raises IndentationError / OSError "lineno is out of bounds" on import. The sweep already emitted a valid cache file for every non-KNOWN_BROKEN model_type, so verify that artefact directly; trigger a compile only when running this test in isolation. Verified locally: pytest -q tests/_zoo_compiler_cache_shim.py (5 passed, 1 skipped) pytest -q tests/.._real_modeling_module (3 passed) --- .github/workflows/consolidated-tests-ci.yml | 44 ++++++++++++--------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 897367de0c..de928e83ee 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -887,14 +887,23 @@ jobs: import _zoo_aggressive_cuda_spoof as _spoof _spoof.apply() - # Hermetic cache dir + force compile path BEFORE importing - # unsloth_zoo.compiler (its globals capture env at module load). + # Hermetic cache dir + force compile path. The compiler's + # globals (UNSLOTH_COMPILE_LOCATION, UNSLOTH_COMPILE_USE_TEMP) + # are captured at module load; an earlier conftest `import + # unsloth` may have already imported unsloth_zoo.compiler with + # the default "unsloth_compiled_cache" path. Mutate the live + # module globals after import so this shim is robust to that + # ordering. Otherwise the compiler silently writes to the + # default cache and the per-model file assertion fails. _CACHE = pathlib.Path(tempfile.mkdtemp(prefix="unsloth_cache_")) os.environ["UNSLOTH_COMPILE_LOCATION"] = str(_CACHE) os.environ["UNSLOTH_COMPILE_OVERWRITE"] = "1" os.environ.pop("UNSLOTH_COMPILE_DISABLE", None) import pytest + import unsloth_zoo.compiler as _zoo_compiler + _zoo_compiler.UNSLOTH_COMPILE_LOCATION = str(_CACHE) + _zoo_compiler.UNSLOTH_COMPILE_USE_TEMP = False from unsloth_zoo.compiler import unsloth_compile_transformers @@ -1034,12 +1043,14 @@ jobs: emitted cache file has the model-specific RMSNorm class attribute, not just that the file parses + imports. - Note on test isolation: ``unsloth_compile_transformers`` - early-returns when ``modeling.__UNSLOTH_PATCHED__`` is set, - so once an earlier test in the same collection patches the - module the next call won't re-emit the cache file. Drop the - marker (and any stale cache file) before invoking so this - test is order-independent.""" + ``unsloth_compile_transformers`` is not idempotent in- + process: calling it twice on the same modeling module + after rewriting class attributes corrupts the inspect + source/line cache and the second emitted file is malformed + Python. The sweep above already produced a valid cache + file for every non-KNOWN_BROKEN model_type, so just verify + that artefact here. Trigger a compile only when running + this test in isolation (no sweep preceded).""" import importlib as _il try: modeling = _il.import_module( @@ -1049,17 +1060,14 @@ jobs: pytest.skip( f"transformers build lacks model_type={model_type}" ) - if hasattr(modeling, "__UNSLOTH_PATCHED__"): - delattr(modeling, "__UNSLOTH_PATCHED__") combined = _CACHE / f"unsloth_compiled_module_{model_type}.py" - if combined.exists(): - combined.unlink() - unsloth_compile_transformers( - model_type=model_type, fast_lora_forwards=False, - ) - modeling = _il.import_module( - f"transformers.models.{model_type}.modeling_{model_type}" - ) + if not combined.exists(): + unsloth_compile_transformers( + model_type=model_type, fast_lora_forwards=False, + ) + modeling = _il.import_module( + f"transformers.models.{model_type}.modeling_{model_type}" + ) assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True _verify_file(combined, must_expose=[rms_class]) From 920920592e6c5245d315573ff98e82dbce6e755c Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 15 May 2026 16:29:21 +0100 Subject: [PATCH 20/50] Polish/cloud to providers (#5450) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * polish: update provider dropdown and rename cloud * fix: tighten custom provider fallback handling * fix: external provider fallback typing * studio: wire the chat Search button to OpenAI's built-in web_search tool When the active model is an OpenAI external provider and the user clicks the existing Search pill in the composer, the chat-completion request now carries the unified enable_tools shorthand: enable_tools: true enabled_tools: ["web_search"] The backend's stream_chat_completion threads enabled_tools through to _stream_openai_responses, which translates it into the Responses API tool schema: body["tools"] = [{"type": "web_search"}] per the OpenAI Responses tool spec (https://developers.openai.com/api/docs/guides/tools). OpenAI then runs the search server-side before the model replies; the search- informed answer streams back through the existing response.output_text.delta path. web_search_call lifecycle events are silently ignored for now — sources / status indicators are follow-up scope. Frontend: - provider-capabilities.ts: new providerSupportsBuiltinWebSearch() helper. Returns true only for `openai` today; Anthropic (web_search_20250305), Gemini grounded-search, and OpenRouter variants can be added later with matching backend translation. - chat-page.tsx: both model-switch paths (the onChange handler and the inferenceParams.checkpoint useEffect) set supportsTools to match the new helper, and force toolsEnabled=false on every external switch so the Search toggle is opt-in by default. - chat-adapter.ts: external branch adds enable_tools + enabled_tools=["web_search"] to the request body when the toggle is on AND the active provider supports built-in web-search. Local-model branch is unchanged — it continues to route the same shorthand through our local tool runtime. Backend: - routes/inference.py: forwards payload.enabled_tools to stream_chat_completion at the proxy site (line 1599). - external_provider.py: stream_chat_completion gains an enabled_tools parameter; _stream_openai_responses appends {"type": "web_search"} to body["tools"] when the list contains "web_search". Other tools (file_search, code_interpreter, image_generation, computer_use_preview) are easy follow-ups in the same block. Reuses the existing pydantic ChatCompletionRequest.enabled_tools field, so no schema migrations. * studio/backend: surface OpenAI server-side web_search in the chat UI When the user has the chat Search button toggled on and OpenAI's /v1/responses invokes the built-in web_search tool, _stream_openai_responses now translates the tool's lifecycle events and citation annotations into the same _toolEvent shape that local-tool calls use. The result: the chat UI shows a web_search tool-call card mid-stream, then lists the cited sources at the end of the message — identical to how local web_search renders. SSE event translation: - response.output_item.added with item.type=web_search_call -> emit _toolEvent tool_start. Carries item.action.query as args when OpenAI ships it on the added event. - response.output_item.done with item.type=web_search_call -> backfill the query if it only arrives on the done variant. The existing reasoning branch on the same event is preserved as an if/elif under a shared isinstance guard. - response.output_text.annotation.added with type=url_citation -> collect into the most-recent web_search_call.citations list. - response.output_text.delta with inline annotations[] (older API variant) -> same collection path, so both wire shapes work. - response.completed -> emit _toolEvent tool_end per call with citations formatted as Title: <title>\nURL: <url>\nSnippet: <snippet> blocks joined by `\n---\n`. The frontend's parseSourcesFromResult already lifts this format into source content parts at end-of-stream. - response.incomplete -> close out web_search cards with whatever citations had landed, so a truncated response does not leave a perpetually "running" tool card in the UI. Both reasoning and web_search work simultaneously on the same turn — the body sends `reasoning: {effort, summary}` and `tools: [{type: "web_search"}]` independently, and the SSE handler tracks them through separate channels. Diagnostic: finally-block logger now reports per stream web_search_requested - whether the client asked for it web_search_invocations - how many calls OpenAI actually made citations - total URLs cited queries - the search queries the model issued reasoning_emitted - whether <think> content was streamed so reports of "I clicked Search and nothing happened" can be triaged from the backend log without browser devtools. * studio/backend: fix empty query + per-card '(no sources cited)' on OpenAI web_search Two display bugs on the OpenAI Responses web_search → chat-UI bridge: 1. Tool cards showed "Searching for ''" — query missing. OpenAI's response.output_item.added for web_search_call does not reliably populate action.query across API versions; the canonical place is output_item.done. The previous code emitted tool_start at added with empty args and tried to backfill at done, but the frontend's _toolEvent: tool_start is a one-shot push (no update mechanism), so the args stayed empty. Fix: defer both tool_start *and* a placeholder tool_end emission to output_item.done, where action.query is guaranteed populated. added now just initialises tracking. Frontend then renders one card per call with the right "Searching for: <query>" label. 2. Every card showed "(no sources cited)". The previous code tried to attribute url_citation annotations to individual web_search_call invocations, but OpenAI's annotations carry no link back to a specific search call — they're just URLs the model cited from the aggregated search pool. With N invocations and M annotations, the previous logic bucketed all M into the last call and stamped "(no sources cited)" on the rest. Fix: collect citations into a single shared all_url_citations list, dedup by URL. At response.completed (and response.incomplete) overwrite the *last* web_search_call's tool_end result with the aggregated Title:/URL:/Snippet: blocks. The frontend's parseSourcesFromResult already flatMaps every web_search result, so one non-empty result is enough to surface the full source-pill set at the message tail. Other tool cards get an empty result string (no '(no sources)' text). Diagnostic log unchanged in shape; total_citations now reads len(all_url_citations) directly. * studio/chat: split Code and Search pill gates so external models cannot enable Code The previous wire-up set supportsTools=true for OpenAI external models to light up the Search pill, but supportsTools also gates the Code pill, so Code became clickable for OpenAI even though external providers have no local code execution. Separate the two gates so each pill reflects what's actually available: - chat-runtime-store: new `supportsBuiltinWebSearch: boolean` flag. Distinct from supportsTools — that one still means "runtime has a local tool sandbox" (Code, python, our DuckDuckGo web_search). This one means "the active external provider exposes a server-side web_search tool we can opt into" (OpenAI's /v1/responses today). - chat-page model-switch (both code paths): for external models, supportsTools is now forced to false (no local Code path) and supportsBuiltinWebSearch follows providerSupportsBuiltinWebSearch. Local-model paths are unaffected — they only set supportsTools. - shared-composer: Search pill gates on `searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch)`. Code pill gates on `codeDisabled = !modelLoaded || !supportsTools` — strictly the local runtime, so external models keep Code greyed out. A `toolsDisabled = codeDisabled` alias is left in place for any later-touched call site that may still reference the old name. No backend changes — chat-adapter already calls providerSupportsBuiltinWebSearch directly, independent of the store flags, so the request shape and the backend translation are unchanged. * studio/chat: default external reasoning effort to medium, not the carry-over When switching to an external model with reasoning support, the effort dropdown was inheriting whatever value the user had set on a prior model — frequently "xhigh" left over from a previous Opus/gpt-5 session. That meant every fresh OpenAI/Anthropic selection started at Extra High, burning tokens unintentionally. Both model-switch sites in chat-page (the useEffect on inferenceParams.checkpoint and the onChange callback) now pick "medium" whenever the new model's level list contains it, instead of the clamped carry-over. The clamp still fires as a fallback for the narrow case where a model doesn't expose medium (e.g. gpt-5.3-chat- latest which only has medium anyway — no change there). Users can still pick another level explicitly via the Think dropdown. * studio/chat: also light the Search pill in the welcome-screen composer There are two composers in the chat feature. shared-composer.tsx renders inside an active thread, and assistant-ui/thread.tsx has its own WebSearchToggle / CodeToolsToggle that ship the welcome-screen "Send a message…" composer (visible before the first user message). The previous fix split supportsTools and supportsBuiltinWebSearch in shared-composer but never touched the welcome-screen toggles in thread.tsx — they both still gated on supportsTools alone, so the Search pill stayed greyed on the welcome screen even for OpenAI external models that legitimately support web_search server-side. Mirror the shared-composer rule in WebSearchToggle: disabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch) CodeToolsToggle is left as-is — its current `disabled = !(modelLoaded && supportsTools)` is correct: external models have no local code-execution sandbox, so Code stays greyed when supportsTools=false (which is what chat-page now writes for external selections). * studio/backend: wire Anthropic server-side web_search end-to-end Mirrors the OpenAI web_search integration for Anthropic's web_search_20250305 tool. When the user toggles Search on with an Anthropic model selected, the request now carries the documented tool entry: tools: [{type: "web_search_20250305", name: "web_search", max_uses: 5}] on /v1/messages, and the SSE translation surfaces tool cards + source pills in the chat UI exactly the same way as OpenAI. stream_chat_completion now forwards enabled_tools into the Anthropic branch (was only doing this for the OpenAI Responses branch). _stream_anthropic gains an enabled_tools parameter and the web_search request-body block plus three additional event handlers: - content_block_start with type=server_tool_use, name=web_search: start tracking a new call. id becomes the tool_call_id. - content_block_delta with type=input_json_delta inside a server_tool_use block: buffer the partial_json so we can read out the search query when the block closes. - content_block_start with type=web_search_tool_result: capture the per-call result list (urls + titles) that Anthropic ships inline. - content_block_stop: closes whichever block we're inside — * server_tool_use -> emit _toolEvent: tool_start with the parsed query as args. * web_search_tool_result -> emit _toolEvent: tool_end with Title:/URL: blocks the frontend's parseSourcesFromResult lifts into source pills. * thinking block -> existing </think> close. Unlike OpenAI we get per-call results directly, so no aggregated- last-call fallback is needed — each tool card carries its own citations. Diagnostic log on stream completion now reports web_search_requested / invocations / total_results / queries, matching the OpenAI shape. Frontend providerSupportsBuiltinWebSearch returns true for 'anthropic' as well, so the Search pill lights up on Claude models the same way it does on OpenAI. The existing chat-adapter external branch already sends enabled_tools=['web_search'] based on this helper — no adapter changes needed. * studio: wire OpenRouter built-in web search via :online model suffix OpenRouter exposes a universal "add web search to any model" shortcut: append `:online` to the model id and the gateway runs the search server-side, streaming citations back as annotations on text deltas. Documented at https://openrouter.ai/docs/features/web-search Hook the existing Search toggle into that path: Backend (external_provider.py, default OAI-compat branch): - When provider_type == 'openrouter' and enabled_tools contains 'web_search', rewrite body['model']: openai/gpt-4o -> openai/gpt-4o:online anthropic/claude-sonnet-4-5:free -> anthropic/claude-sonnet-4-5:online Any existing `:variant` (`:free`, `:nitro`, etc.) is replaced — OpenRouter variants are mutually exclusive. - `openrouter/free` is skipped: it's a meta-router and `:online` is not a valid suffix on it (the gateway 400s). - A one-line INFO log fires whenever the rewrite happens so the diagnostic backend log shows exactly which model id the request was promoted to. Frontend (provider-capabilities.ts): - providerSupportsBuiltinWebSearch now returns true for 'openrouter' alongside 'openai' and 'anthropic'. The Search pill lights up and the existing chat-adapter external branch already forwards enabled_tools=['web_search'] based on this helper — no adapter changes needed. No new SSE event handling: OpenRouter does not emit a separate web_search_call event the way OpenAI/Anthropic do. Citations come back as text annotations via the existing reasoning_details path the adapter already parses, so source data flows through without extra translation. A per-call tool-card UX ("Searching for: …") would require synthesizing one client-side; deferred to a follow-up if the bare-citation flow feels too minimal. * studio: wire Mistral built-in web search connector Same shape as OpenAI's web_search tool, lives on /v1/chat/completions instead of /v1/responses. When the chat Search pill is toggled on with a Mistral model selected, the backend now appends {"type": "web_search"} to body["tools"] before the request goes out. Idempotent — won't double-append if a future call site adds it first. Models in the registry allowlist that don't support the connector (codestral, devstral, ministral, mistral-tiny) will surface a 400 from upstream; the existing default-path error log captures it. Mistral's docs: https://docs.mistral.ai/capabilities/agents/connectors/websearch Frontend providerSupportsBuiltinWebSearch returns true for 'mistral' now, alongside openai / anthropic / openrouter. The Search pill lights up for Mistral models and the existing adapter branch already sends enabled_tools=['web_search'] off this helper — no adapter changes. No SSE translation yet — Mistral streams citations inline as text annotations or `references` in the final assistant content, not as a separate web_search_call event. Citations flow through to the message body as text; a per-call tool-card UX with "Searching for: …" indicators is a follow-up if needed. * studio/backend: fix OpenRouter web_search to use plugins shape + synthesize tool card Two changes against the actual OpenRouter docs at https://openrouter.ai/docs/guides/features/plugins/web-search: Request shape: The previous commit appended :online to the model id, which works on concrete model ids but rejects on meta-routers like openrouter/free — and that's exactly the model the user was testing with, so neither the request rewrite nor the diagnostic log fired. Switch to the universal plugins shape: body["plugins"] = [{"id": "web"}] Per the docs this is "exactly equivalent" to :online but works on every model id including openrouter/free and openrouter/auto. No model suffix manipulation, idempotent if added twice. Tool-card synthesis: OpenRouter doesn't emit a structured web_search_call event the way OpenAI/Anthropic do — citations come back only as `annotations` of type=url_citation on delta/message objects. To match the chat-UI tool-card UX the user expects ("Searching for: …" indicator, source pills at message tail), synthesize the events client-side in the default OAI-compat stream loop: - On stream open (after the 200 status check): yield a synthetic _toolEvent: tool_start with tool_name=web_search, fixed id "openrouter_web_search". The chat-UI then renders the running tool card before any text streams. - During the SSE loop: scan every chunk's choices[].delta and choices[].message for `annotations: [{type: "url_citation", url_citation: {url, title, content}}]` entries. Dedup by URL into a citations list. Handles both the nested-url_citation shape OpenRouter documents and the flat-on-annotation shape some upstreams ship. - On [DONE] (or stream-close without [DONE]): emit synthetic tool_end carrying the citations as Title: …\nURL: …\nSnippet: …\n---\n… blocks the existing parseSourcesFromResult lifts into source pills at message tail. Diagnostic log on completion now also reports web_search_requested + citation count alongside the existing chosen-model / event-count telemetry. * studio: drop Mistral built-in web_search — connector lives on Agents API only Mistral's web_search is exclusively on /v1/agents + /v1/conversations; sending it on /v1/chat/completions returns "WebSearchTool connector is not supported". Wiring it would require a dedicated Agents streaming path. Remove from the frontend capability map and revert the chat-completions tool injection. * studio: wire Kimi $web_search builtin via two-call round-trip Kimi's $web_search lives on /v1/chat/completions but requires a client round-trip per https://platform.kimi.ai/docs/guide/use-web-search: the first call returns tool_calls with function.arguments populated; the caller echoes those arguments back as a role=tool message; the second call streams the final answer with search results incorporated. The docs also mandate thinking=disabled while the builtin is active. Backend: new _stream_kimi_web_search helper dispatched from stream_chat_completion when provider_type=='kimi' and 'web_search' in enabled_tools. Buffers tool_calls across deltas, falls back to a plain stream if the model declines to search, and synthesizes tool_start (with parsed query) / tool_end (with any url_citation annotations) so the chat UI's web-search card behaves the same as other providers. Frontend: kimi added to providerSupportsBuiltinWebSearch so the Search pill lights up in the composer. * studio/chat: mutual exclusion of Think + Search on Kimi composer Kimi's $web_search builtin requires thinking=disabled per https://platform.kimi.ai/docs/guide/use-web-search, so the two states cannot coexist. Make the pills mutually exclusive in both composers (shared and welcome-screen): clicking Search turns Think off; clicking Think back on turns Search off. Default Think to on when a Kimi model is selected — k2.6/k2.5 ship with thinking enabled out of the box. * studio/chat: fix wrong provider var name in onChange branch selectedProvider, not provider — TS2304 in tsc -b. * studio/backend: add diagnostics to Kimi $web_search round-trip Log the actual function.arguments from the first call (so we can see the model's search query) and the second call's usage.prompt_tokens + any annotation type names that came through. prompt_tokens spiking above the input message length is direct proof the server injected search results into context. annotation_types lets us learn the shape Kimi uses for citations if/when they emit any. * studio: per-provider defaults — Anthropic xhigh + Search on, OpenAI high + Search on, Opus 4.7 gains max Anthropic: Think effort defaults to the highest level the model supports (xhigh on 4.6/4.7, high on 4.5) and Search starts on, since the web_search_20250305 tool returns structured citations end-to-end. OpenAI: Think effort defaults to 'high' (the gpt-5.x reasoning sweet spot for /v1/responses + web_search) and Search starts on. Opus 4.7: 'max' added as an effort level above 'xhigh' in both backend (_ANTHROPIC_THINKING_SPECS) and frontend (ANTHROPIC_REASONING_MODELS). Kimi diagnostics: emit tool_end immediately after tool_start so the web-search card transitions to 'complete' before the second-call answer streams, log first-call args + second-call usage/prompt_tokens + any annotation type names, request stream_options.include_usage so the second call exposes usage in SSE. * studio/backend: harden Kimi fallback path with HTTPError handler + manual aiter_lines loop Addresses PR review feedback (#5443): the no-search fallback streaming path was using `async for response.aiter_lines()` and had no `httpx.HTTPError` guard around the POST. Switch to the manual __anext__ loop pattern used elsewhere in this module (avoids the Python 3.13 + httpcore 1.0.x GeneratorExit propagation issue) and wrap the whole request in a try/except so network failures surface as a proper SSE error frame instead of a raw traceback. * feat: prompt caching frontend for openai/anthropic * studio/chat: route vLLM provider to /v1/chat/completions, not /v1/responses vLLM's /v1/responses rebuilds messages through the loaded model's chat template, which 400s on strict-alternation templates like Gemma 3 ("Conversation roles must alternate user/assistant/..."). Stop collapsing vllm -> openai in the frontend so the backend sees the real provider type and falls through to the standard chat-completions path. Register vllm as a hidden entry in PROVIDER_REGISTRY so supports_vision and provider-create validation work without surfacing it in the cloud-provider dropdown. * studio/chat: wire prompt caching for OpenAI and Anthropic external providers Backend half of the prompt_caching toggle that already exists in the chat settings panel. Scoped to OpenAI cloud (/v1/responses) and Anthropic (/v1/messages); every other provider plumbs the flag as a no-op. - Anthropic: attach cache_control={type:ephemeral} to the system block so the static prefix is reused across turns. Without the marker Anthropic caches nothing, so this is the only way to make the toggle do real work on /v1/messages. - OpenAI: opt into prompt_cache_retention="24h" — same price as the default in_memory policy per the OpenAI docs, but the cache survives ~24 hours of idle instead of ~5-10 minutes. The model picker is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which accept the parameter (gpt-5.5+ already defaults to "24h" so it's a no-op there). - Treats `enable_prompt_caching=None` as enabled to match the frontend default for both providers; pass `false` explicitly to opt out. * studio/chat: log cache token counts on OpenAI and Anthropic stream completion Surface cache usage in the existing "stream complete" info logs so prompt-caching behavior can be verified by tailing the studio backend log instead of opening the provider dashboard. - Anthropic: latch usage from message_start (input + cache_creation + cache_read counts) and message_delta (output_tokens), then include in the per-request summary. cache_read_input_tokens > 0 confirms the cache_control marker on the system block is doing its job. - OpenAI Responses: latch usage from response.completed and response.incomplete, extract usage.input_tokens_details.cached_tokens (the /v1/responses field name, not prompt_tokens_details). A non-zero value on turn N proves prompt_cache_retention="24h" let the prefix hit the cache instead of being recomputed. * studio/backend: strip temperature/top_p for Claude 4.7 family Anthropic Opus 4.7 removed temperature, top_p, and top_k as a launch breaking change ("Sampling parameters removed" in the 4.7 release notes at https://platform.claude.com/docs/en/about-claude/models/whats-new-claude-4-7). Setting any of them to a non-default value returns 400 "<param> is deprecated for this model". The existing guard only handled top_k; temperature was still being sent unconditionally and is now breaking opus-4-7 requests. Rename _ANTHROPIC_TOP_K_DEPRECATED to _ANTHROPIC_4_7_SAMPLING_REMOVED to reflect the broader scope, omit temperature from the base body on 4.7, and skip the thinking-mode temperature=1 override on 4.7 (still applied on 4.5/4.6 where it's required). Existing thinking_translation tests target 4.5/4.6 / mock the wire so they're unaffected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/chat: anchor Anthropic prompt cache on the latest message too A system-only cache_control marker is a no-op when the system prompt is empty or shorter than Anthropic's ~1024-token cache floor — caching silently does nothing (both cache_creation and cache_read return 0). Add a second cache_control breakpoint on the final block of the latest conversation message so the entire prefix (system + prior turns + new user turn) becomes eligible for caching. On turn N+1, Anthropic rehydrates everything up through turn N's marker instead of recomputing it. Up to 4 breakpoints are allowed per request; we use at most 2 (system + tail). Tail rebuild avoids mutating the caller's content list so an image-bearing turn still slots cleanly into the cached prefix. * studio/chat: gate vLLM reasoning toggle on provider config Add a "This server runs a reasoning model" checkbox on the vLLM provider config. When off (default), the chat Think pill stays hidden and no enable_thinking ever reaches vLLM. When on, the pill renders, per-turn state flows through the existing enable_thinking plumbing, and the backend proxy lifts it onto chat_template_kwargs.enable_thinking so vLLM's Jinja template honours it. * chore: clean vLLM reasoning-toggle comments * studio/chat: gate prompt_cache_retention to actual OpenAI cloud requests Addresses Codex P1 review on _stream_openai_responses. The frontend only sends enable_prompt_caching for the openai/anthropic UI provider types, so ollama/llama.cpp/"custom" requests reach this helper with the flag as None. The previous `is not False` check treated None as enabled and injected prompt_cache_retention="24h" into every request including those bound for non-OpenAI servers, which would 400 on servers that implement /v1/responses but not the retention parameter. Match the public OpenAI host (api.openai.com) on the client base_url before adding the field so it only lands on actual OpenAI cloud requests. Studio's openai picker is already registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which accept the parameter. --------- Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 204 ++++++++++++++-- studio/backend/core/inference/providers.py | 32 ++- studio/backend/models/inference.py | 11 + studio/backend/routes/inference.py | 1 + .../public/provider-logos/llama_cpp.svg | 1 + .../frontend/public/provider-logos/ollama.svg | 14 ++ .../frontend/public/provider-logos/vllm.svg | 1 + .../assistant-ui/model-selector.tsx | 16 ++ .../src/components/assistant-ui/thread.tsx | 4 + .../src/features/chat/api-provider-logo.tsx | 8 +- .../src/features/chat/api/chat-adapter.ts | 23 +- .../frontend/src/features/chat/chat-page.tsx | 24 +- .../features/chat/chat-providers-dialog.tsx | 225 ++++++++++++------ .../src/features/chat/chat-settings-sheet.tsx | 39 +++ .../src/features/chat/external-providers.ts | 146 +++++++++++- .../features/chat/provider-capabilities.ts | 36 ++- .../src/features/chat/shared-composer.tsx | 4 + .../frontend/src/features/chat/types/api.ts | 1 + .../src/features/settings/settings-dialog.tsx | 2 +- 19 files changed, 681 insertions(+), 111 deletions(-) create mode 100644 studio/frontend/public/provider-logos/llama_cpp.svg create mode 100644 studio/frontend/public/provider-logos/ollama.svg create mode 100644 studio/frontend/public/provider-logos/vllm.svg diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index bd14e73d16..48a5f0e879 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -24,11 +24,17 @@ import structlog # sites use printf-style positional args, which structlog accepts. logger = structlog.get_logger(__name__) -# Claude 4.7 (Opus/Sonnet/Haiku) deprecated top_k and returns 400 -# "top_k is deprecated for this model" when it is set. 3.x and 4.5/4.6 -# still accept it. Match the 4-7 line specifically so we keep the knob -# live on every other Claude generation. -_ANTHROPIC_TOP_K_DEPRECATED = re.compile(r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)") +# Claude 4.7 (Opus/Sonnet/Haiku) removed temperature, top_p, and top_k — +# the API returns 400 "<param> is deprecated for this model" if any of +# them is set to a non-default value. The "Sampling parameters removed" +# section of the 4.7 release notes is the authoritative reference: +# https://platform.claude.com/docs/en/about-claude/models/whats-new-claude-4-7 +# 3.x and 4.5/4.6 still accept all three; match the 4-7 line strictly so +# the knobs keep working on earlier families. The trailing -4-7[-.]/EOL +# anchor keeps future versions (e.g. claude-opus-5) unaffected. +_ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile( + r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)" +) _OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)") @@ -229,6 +235,7 @@ class ExternalProviderClient: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, enabled_tools: Optional[list[str]] = None, + enable_prompt_caching: Optional[bool] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: """ @@ -253,6 +260,7 @@ class ExternalProviderClient: enable_thinking, reasoning_effort, enabled_tools, + enable_prompt_caching, ): yield line return @@ -272,6 +280,7 @@ class ExternalProviderClient: enable_thinking, reasoning_effort, enabled_tools, + enable_prompt_caching, ): yield line return @@ -346,6 +355,13 @@ class ExternalProviderClient: _apply_mistral_reasoning_controls( body, model, enable_thinking, reasoning_effort ) + elif self.provider_type == "vllm" and enable_thinking is not None: + # vLLM gates thinking via chat_template_kwargs.enable_thinking. + tpl_kw = body.get("chat_template_kwargs") + if not isinstance(tpl_kw, dict): + tpl_kw = {} + tpl_kw["enable_thinking"] = bool(enable_thinking) + body["chat_template_kwargs"] = tpl_kw # OpenRouter exposes a unified `reasoning` parameter on every # chat-completion request — the gateway routes it to whichever @@ -1043,6 +1059,7 @@ class ExternalProviderClient: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, enabled_tools: Optional[list[str]] = None, + enable_prompt_caching: Optional[bool] = None, ) -> AsyncGenerator[str, None]: """ Call the Anthropic Messages API and translate its SSE to OpenAI format. @@ -1112,24 +1129,79 @@ class ExternalProviderClient: else: filtered.append(msg) + # Claude 4.7 family removed temperature / top_p / top_k entirely. + # The earlier guard only handled top_k; temperature is now also + # rejected with 400 "temperature is deprecated for this model". + # Latch the match once and reuse it everywhere temperature or + # top_k would otherwise be set — including the thinking-mode + # override below, which used to force temperature=1. + sampling_removed = bool(_ANTHROPIC_4_7_SAMPLING_REMOVED.match(model)) + body: dict[str, Any] = { "model": model, "messages": filtered, "max_tokens": max_tokens or 1024, # required by Anthropic - "temperature": temperature, "stream": True, } - # top_k is deprecated on Claude 4.7 (Opus/Sonnet/Haiku) — the API - # returns 400 "top_k is deprecated for this model" when it is set. - # 3.x and 4.5/4.6 still accept it, so gate strictly on the 4.7 ids. - if ( - top_k is not None - and top_k > 0 - and not _ANTHROPIC_TOP_K_DEPRECATED.match(model) - ): + if not sampling_removed: + body["temperature"] = temperature + if top_k is not None and top_k > 0 and not sampling_removed: body["top_k"] = top_k + # Anthropic only caches a prefix when at least one cache_control + # marker is attached to it — the frontend defaults + # enable_prompt_caching to True for Anthropic, so treat `None` the + # same as True here (callers that don't set the flag still get + # caching). Pass False explicitly to opt out. + prompt_caching_enabled = enable_prompt_caching is not False + if system: - body["system"] = system + if prompt_caching_enabled: + # System block is the most stable prefix across turns, so + # it gets its own breakpoint. Skipped when system is + # empty — there's nothing to cache, and an empty marker + # is a no-op. + body["system"] = [ + { + "type": "text", + "text": system, + "cache_control": {"type": "ephemeral"}, + } + ] + else: + body["system"] = system + + if prompt_caching_enabled and filtered: + # Second breakpoint at the end of the conversation. Anthropic + # caches the longest matching prefix up to a cache_control + # marker; placing one on the latest message means turn N+1 + # rehydrates everything up through turn N from cache instead + # of recomputing it. This is what makes caching actually work + # when the system prompt is empty or shorter than Anthropic's + # ~1024-token cache floor — the conversation history carries + # the bulk of the input tokens. Anthropic allows up to 4 + # breakpoints per request; we use at most 2 (system + tail). + last_msg = filtered[-1] + content = last_msg.get("content") + if isinstance(content, str): + last_msg["content"] = [ + { + "type": "text", + "text": content, + "cache_control": {"type": "ephemeral"}, + } + ] + elif isinstance(content, list) and content: + # Don't mutate the caller's list. Rebuild the tail with + # cache_control attached to the final block so an + # upstream image-bearing turn still cleanly slots into + # the cache as part of the conversational prefix. + head = list(content[:-1]) + tail = content[-1] + if isinstance(tail, dict): + head.append({**tail, "cache_control": {"type": "ephemeral"}}) + else: + head.append(tail) + last_msg["content"] = head thinking_spec = _anthropic_thinking_spec(model) allowed_efforts = ( thinking_spec.efforts @@ -1155,13 +1227,15 @@ class ExternalProviderClient: if effort and effort != "none": # Anthropic rejects top_k whenever thinking is enabled. body.pop("top_k", None) - # Anthropic requires temperature=1 whenever thinking is enabled, - # AND forbids top_p in the same request: setting both produces + # Earlier families (4.5/4.6) require temperature=1 when + # thinking is enabled and forbid top_p in the same request: # "temperature and top_p cannot both be specified for this # model. Please use only one." - # The base body never sets top_p, but pop defensively in case - # an upstream edit ever adds it before this branch runs. - body["temperature"] = 1 + # On Claude 4.7, temperature was removed entirely — sending + # any value (including 1) returns 400 — so skip the override + # there and let the model use its default sampling. + if not sampling_removed: + body["temperature"] = 1 body.pop("top_p", None) if thinking_spec and thinking_spec.kind == "adaptive": # `display` defaults to "omitted" on Claude Opus 4.7 (per the @@ -1278,6 +1352,13 @@ class ExternalProviderClient: current_server_tool_use: Optional[dict[str, Any]] = None current_result_block: Optional[dict[str, Any]] = None web_search_calls: dict[str, dict[str, Any]] = {} + # Cache usage tracking. message_start carries the input + # accounting (incl. cache_creation_input_tokens and + # cache_read_input_tokens); message_delta carries cumulative + # output_tokens. Both are surfaced in the "stream complete" + # log so prompt caching can be verified per-request without + # opening the Anthropic dashboard. + last_usage: dict[str, Any] = {} def _content_chunk(text: str) -> str: chunk = { @@ -1352,6 +1433,16 @@ class ExternalProviderClient: key = event_type or "<unknown>" event_counts[key] = event_counts.get(key, 0) + 1 + # message_start carries the input-side usage block + # including cache_creation_input_tokens and + # cache_read_input_tokens. message_delta updates + # output_tokens (and may overwrite the input fields + # with final values). Merge both into last_usage. + if event_type == "message_start": + start_usage = (event.get("message") or {}).get("usage") + if isinstance(start_usage, dict): + last_usage.update(start_usage) + if event_type == "content_block_start": content_block = event.get("content_block") or {} block_type = content_block.get("type") @@ -1489,6 +1580,9 @@ class ExternalProviderClient: thinking_open = False elif event_type == "message_delta": + delta_usage = event.get("usage") + if isinstance(delta_usage, dict): + last_usage.update(delta_usage) stop_reason = event.get("delta", {}).get("stop_reason") if stop_reason: if thinking_open: @@ -1538,15 +1632,28 @@ class ExternalProviderClient: for sc in web_search_calls.values() if sc.get("query") ] + # cache_read_input_tokens > 0 on turn N proves the + # cache_control marker on the system block is doing + # its job — turn 1 will show cache_creation > 0 + # instead. cache_creation tokens are billed at a + # small premium; cache_read tokens are billed at a + # discount. logger.info( "Anthropic stream complete (model=%s, " "web_search_requested=%s, web_search_invocations=%s, " - "results=%s, queries=%s, events=%s)", + "results=%s, queries=%s, " + "input_tokens=%s, output_tokens=%s, " + "cache_creation_input_tokens=%s, " + "cache_read_input_tokens=%s, events=%s)", model, web_search_requested, web_search_invocations, total_results, queries, + last_usage.get("input_tokens"), + last_usage.get("output_tokens"), + last_usage.get("cache_creation_input_tokens"), + last_usage.get("cache_read_input_tokens"), event_counts, ) await response.aclose() @@ -1584,6 +1691,7 @@ class ExternalProviderClient: enable_thinking: Optional[bool], reasoning_effort: Optional[str], enabled_tools: Optional[list[str]] = None, + enable_prompt_caching: Optional[bool] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -1685,6 +1793,27 @@ class ExternalProviderClient: if max_tokens is not None: body["max_output_tokens"] = max_tokens + # Prompt caching on /v1/responses is automatic and free, but the + # default in-memory policy only survives ~5-10 min of inactivity + # (up to ~1 hr). Opt into the 24-hour retention policy so a chat + # left idle overnight still hits the cache on the next turn. + # Pricing is identical to in_memory per OpenAI's docs. + # + # Gated on the base URL because ollama / llama.cpp / "custom" + # presets all collapse to provider_type="openai" in + # toExternalBackendProviderType, so they also land in this + # helper. Those servers expose /v1/responses-shaped routes in + # some configurations but don't implement + # prompt_cache_retention — sending the field unconditionally + # would 400 them. Match the public OpenAI host strictly so the + # field only goes to OpenAI cloud. Studio's openai model picker + # is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which + # accept this parameter (gpt-5.5+ already defaults to "24h" and + # rejects "in_memory", so it's a safe no-op there). + is_openai_cloud = "api.openai.com" in (self.base_url or "") + if is_openai_cloud and enable_prompt_caching is not False: + body["prompt_cache_retention"] = "24h" + # OpenAI server-side tools — see # https://developers.openai.com/api/docs/guides/tools # The frontend's Search button maps to the unified @@ -1731,6 +1860,12 @@ class ExternalProviderClient: done_emitted = False reasoning_open = False reasoning_emitted = False + # Latched from response.completed / response.incomplete so + # the final log can surface input_tokens_details.cached_tokens — + # the field that proves prompt_cache_retention="24h" is + # actually hitting OpenAI's cache instead of recomputing + # the prefix every turn. + last_usage: Optional[dict[str, Any]] = None # Per-call state for OpenAI's server-side web_search tool. Mapped # back into our local _toolEvent shape so the existing chat-UI # renderer surfaces web_search the same way it does for local @@ -1955,6 +2090,9 @@ class ExternalProviderClient: reasoning_emitted = True elif event_type == "response.completed": + completed_usage = (event.get("response") or {}).get("usage") + if isinstance(completed_usage, dict): + last_usage = completed_usage if reasoning_open: yield _chunk_with_text("</think>") reasoning_open = False @@ -1998,6 +2136,11 @@ class ExternalProviderClient: yield f"data: {_json.dumps(chunk)}" elif event_type == "response.incomplete": + incomplete_usage = (event.get("response") or {}).get( + "usage" + ) + if isinstance(incomplete_usage, dict): + last_usage = incomplete_usage if reasoning_open: yield _chunk_with_text("</think>") reasoning_open = False @@ -2071,16 +2214,33 @@ class ExternalProviderClient: for sc in web_search_calls.values() if sc.get("query") ] + # cached_input_tokens > 0 on turn N proves + # prompt_cache_retention="24h" is letting the previous + # turn's prefix hit the cache instead of being + # recomputed. On /v1/responses the field is nested as + # usage.input_tokens_details.cached_tokens (not + # prompt_tokens_details, which is the /v1/chat/completions + # shape). + cached_input_tokens = None + if isinstance(last_usage, dict): + details = last_usage.get("input_tokens_details") + if isinstance(details, dict): + cached_input_tokens = details.get("cached_tokens") logger.info( "OpenAI Responses stream complete (model=%s, " "web_search_requested=%s, web_search_invocations=%s, " - "citations=%s, queries=%s, reasoning_emitted=%s)", + "citations=%s, queries=%s, reasoning_emitted=%s, " + "input_tokens=%s, output_tokens=%s, " + "cached_input_tokens=%s)", model, web_search_requested, web_search_invocations, total_citations, queries, reasoning_emitted, + (last_usage or {}).get("input_tokens"), + (last_usage or {}).get("output_tokens"), + cached_input_tokens, ) await response.aclose() await lines_gen.aclose() diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index 4b6d7d6b17..143ced95f1 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -218,6 +218,28 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { # are always among the top regardless of the API's order. "model_id_limit": 15, }, + "vllm": { + "display_name": "vLLM", + # User-supplied via provider_base_url; the route layer already falls + # back to the payload's base_url when the registry entry has none. + "base_url": "", + "default_models": [], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + # Force /v1/chat/completions in stream_chat_completion — vLLM's + # /v1/responses rebuilds messages and runs them through the loaded + # model's chat template, which 400s on strict-alternation templates + # (Gemma 3 raises "Conversation roles must alternate user/assistant + # /user/assistant/..."). The chat-completions path takes messages + # verbatim and avoids that template gauntlet. + "notes": "Self-hosted vLLM server. Always routed to /v1/chat/completions.", + # Surfaced through the frontend's CUSTOM_PROVIDER_PRESETS, not the + # /api/providers/registry dropdown — see list_available_providers. + "hidden": True, + }, "openrouter": { "display_name": "OpenRouter", "base_url": "https://openrouter.ai/api/v1", @@ -269,9 +291,17 @@ def get_base_url(provider_type: str) -> str | None: def list_available_providers() -> list[dict[str, Any]]: - """Return all registered providers (for the /registry endpoint).""" + """Return all registered providers (for the /registry endpoint). + + Hidden entries (``"hidden": True``) are filtered out — they exist in the + registry only for backend lookups (e.g. ``supports_vision`` for vLLM) and + are surfaced in the frontend via ``CUSTOM_PROVIDER_PRESETS`` instead of + the cloud-provider dropdown. + """ result = [] for provider_type, info in PROVIDER_REGISTRY.items(): + if info.get("hidden"): + continue result.append( { "provider_type": provider_type, diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 013328f6c6..f2eed314ee 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -593,6 +593,17 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Override base URL for the external provider.", ) + enable_prompt_caching: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, " + "attaches cache_control={type:ephemeral} to the system block so the " + "static prefix is reused across turns. On OpenAI cloud, caching is " + "automatic for prompts >=1024 tokens and this flag is informational. " + "Ignored for every other provider (mistral, gemini, kimi, openrouter, " + "vllm, local, etc.). Treated as enabled when omitted." + ), + ) # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 55d74d1cbe..b223bcf981 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1596,6 +1596,7 @@ async def _proxy_to_external_provider( enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, enabled_tools = payload.enabled_tools, + enable_prompt_caching = payload.enable_prompt_caching, stream = payload.stream, ) try: diff --git a/studio/frontend/public/provider-logos/llama_cpp.svg b/studio/frontend/public/provider-logos/llama_cpp.svg new file mode 100644 index 0000000000..218cc1de88 --- /dev/null +++ b/studio/frontend/public/provider-logos/llama_cpp.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" xml:space="preserve" viewBox="0 0 512 512"><path d="m356.4 201.3-32.8 58.3c-43.3-33.3-107.4-38.2-150.7-2.4-69.8 57.6-64.9 190.8 43.7 191.6 30.4 0 56.2-14.3 83.9-23.8l14.6 58.1c-24.6 11.4-49.6 23.1-76.6 26.7-246 33.5-231.9-321.6-9.5-340.1 46.7-3.9 87.8 8.3 127.6 31.6zm-169.9-55.9c-37.4 11.2-72.2 31.8-98.5 60.8-4.9-58.8 8.3-177.7 73.7-201 9.7-3.4 43-11.9 42.1 5.3-1 17.3-24.1 46.9-29.7 63-9.7 28.2-.7 47.6 12.6 72.2zm92.4 252.8h-36.5v-41.3h-41.3v-34h37.7l3.6-3.6v-40.1h36.5V323h38.9v34h-38.9zm133.7-41.3v41.3h-36.5v-41.3h-38.9v-34h38.9v-43.8h36.5v40.1l3.6 3.6h37.7v34h-41.3zM305.4 31.4c4.9 7.3-22.6 38.7-27 46.7-12.6 23.8-4.1 37.4 5.3 60-27.5-4.1-53-.7-80.2 2.4C209.6 88.3 239 12.2 305.4 31.4" style="fill:#ff8236"/></svg> diff --git a/studio/frontend/public/provider-logos/ollama.svg b/studio/frontend/public/provider-logos/ollama.svg new file mode 100644 index 0000000000..d3b6a42dd7 --- /dev/null +++ b/studio/frontend/public/provider-logos/ollama.svg @@ -0,0 +1,14 @@ +<svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="294 159 1405.09 1857.06"> +<g clip-path="url(#clip0_1758_1066)"> +<path d="M599.877 159.522C582.544 162.322 561.744 171.388 547.077 182.588C502.677 216.322 468.277 287.922 453.744 377.122C448.277 410.855 444.544 457.655 444.544 493.388C444.544 535.522 449.477 589.388 456.544 626.589C458.144 634.855 458.944 642.188 458.277 642.722C457.744 643.255 451.211 648.588 443.877 654.455C418.811 674.455 390.144 705.255 370.411 733.388C332.544 787.122 308.011 848.188 297.744 914.322C293.744 940.455 292.677 993.255 295.877 1019.39C302.944 1079.66 321.077 1130.59 352.144 1177.26L362.277 1192.32L359.344 1197.26C338.544 1232.19 320.811 1282.72 312.544 1331.26C306.011 1369.66 305.211 1379.92 305.211 1431.39C305.211 1483.26 305.877 1493.52 312.011 1529.39C319.344 1572.32 334.277 1617.79 350.944 1648.06C356.411 1657.92 369.744 1678.46 371.344 1679.52C371.877 1679.79 370.277 1684.72 367.744 1690.46C348.544 1732.46 332.144 1788.32 325.344 1835.39C320.544 1867.66 319.877 1878.06 319.877 1912.06C319.877 1955.39 322.277 1976.46 331.344 2010.99L332.677 2016.06H389.744H446.944L443.211 2008.99C420.144 1966.32 418.011 1887.12 437.877 1808.06C446.944 1771.52 457.211 1744.72 476.411 1707.79L487.877 1685.39V1671.66C487.877 1658.86 487.611 1657.39 483.477 1648.99C480.277 1642.59 476.011 1637.12 468.411 1629.66C455.477 1617.12 446.144 1603.92 438.677 1587.66C405.877 1516.46 399.477 1410.72 422.544 1320.59C432.144 1282.99 448.011 1249.52 464.677 1231.26C476.011 1218.72 481.877 1204.72 481.877 1190.19C481.877 1175.12 476.544 1162.72 464.544 1149.79C430.144 1112.99 408.944 1068.19 401.344 1016.06C390.544 941.788 410.144 860.855 454.677 796.722C498.277 733.788 559.477 693.388 627.877 682.589C643.211 680.055 671.877 680.455 687.877 683.388C705.344 686.455 716.277 685.522 727.477 680.188C741.344 673.655 748.277 665.522 756.411 646.855C763.611 630.188 769.211 621.122 784.277 602.322C802.411 579.788 819.877 564.455 847.877 545.922C879.877 524.988 916.277 509.788 952.544 502.455C965.744 499.788 971.877 499.388 996.544 499.388C1021.21 499.388 1027.34 499.788 1040.54 502.455C1093.74 513.255 1146.54 540.722 1188.68 579.655C1197.74 588.055 1219.48 614.988 1226.41 626.188C1229.08 630.588 1233.74 639.922 1236.68 646.855C1244.81 665.522 1251.74 673.655 1265.61 680.188C1276.41 685.388 1287.74 686.455 1304.54 683.655C1331.08 679.122 1351.48 679.522 1377.48 684.855C1466.01 702.722 1543.08 775.655 1577.21 873.388C1606.94 959.122 1598.54 1048.86 1554.28 1117.39C1546.81 1128.99 1539.34 1138.32 1528.54 1149.79C1505.21 1174.72 1505.21 1205.66 1528.41 1231.26C1566.54 1272.99 1590.41 1375.66 1583.21 1466.19C1578.41 1525.92 1563.08 1579.39 1542.01 1609.66C1538.28 1614.99 1530.54 1624.06 1524.68 1629.66C1517.08 1637.12 1512.81 1642.59 1509.61 1648.99C1505.48 1657.39 1505.21 1658.86 1505.21 1671.66V1685.39L1516.68 1707.79C1535.88 1744.72 1546.14 1771.52 1555.21 1808.06C1574.81 1886.06 1573.08 1963.66 1550.68 2007.79C1548.81 2011.52 1547.21 2014.99 1547.21 2015.39C1547.21 2015.79 1572.68 2016.06 1603.88 2016.06H1660.41L1661.88 2010.32C1662.68 2007.26 1664.01 2002.59 1664.68 1999.92C1666.14 1994.06 1669.08 1976.72 1671.48 1960.06C1673.74 1943.26 1673.74 1881.39 1671.48 1862.72C1662.94 1794.99 1648.68 1741.26 1625.34 1690.46C1622.81 1684.72 1621.21 1679.79 1621.74 1679.52C1622.41 1679.12 1626.14 1673.79 1630.14 1667.79C1659.21 1623.79 1677.08 1568.46 1686.14 1495.39C1688.54 1475.26 1688.54 1388.72 1686.14 1369.39C1679.74 1319.52 1672.01 1285.66 1659.21 1251.39C1653.88 1237.12 1639.74 1206.99 1633.74 1197.26L1630.81 1192.32L1640.94 1177.26C1672.01 1130.59 1690.14 1079.66 1697.21 1019.39C1700.41 993.255 1699.34 940.455 1695.34 914.322C1684.94 848.055 1660.54 787.255 1622.68 733.388C1602.94 705.255 1574.28 674.455 1549.21 654.455C1541.88 648.588 1535.34 643.255 1534.81 642.722C1534.14 642.188 1534.94 634.855 1536.54 626.589C1552.68 542.455 1552.14 437.522 1535.21 355.522C1520.54 284.055 1493.88 227.255 1459.48 194.455C1432.01 168.322 1404.01 157.122 1370.41 159.255C1293.34 163.788 1231.21 252.455 1206.68 392.188C1202.68 414.722 1199.21 441.122 1199.21 448.322C1199.21 451.122 1198.68 453.388 1198.01 453.388C1197.34 453.388 1192.14 450.722 1186.54 447.388C1127.08 412.188 1060.94 393.388 996.544 393.388C932.144 393.388 866.011 412.188 806.544 447.388C800.944 450.722 795.744 453.388 795.077 453.388C794.411 453.388 793.877 451.122 793.877 448.322C793.877 440.855 790.277 413.655 786.411 392.188C764.144 266.722 713.077 183.655 645.211 162.722C635.877 159.922 609.344 158.055 599.877 159.522ZM622.544 268.055C641.744 283.255 663.077 326.722 675.344 375.388C677.611 384.188 680.011 394.322 680.677 398.055C681.211 401.655 682.677 409.788 683.877 416.055C689.077 444.322 691.477 474.855 691.744 512.055L691.877 548.722L682.677 562.322L673.477 576.055H652.011C626.944 576.055 602.011 579.255 578.144 585.655C569.611 587.788 561.344 589.922 559.744 590.322C557.211 590.855 556.811 590.055 555.344 579.122C547.477 519.788 547.877 454.055 556.544 399.388C566.144 338.455 588.544 283.255 610.411 266.988C615.611 263.122 616.544 263.255 622.544 268.055ZM1382.81 267.122C1396.01 276.855 1410.54 302.722 1421.34 335.788C1443.08 401.922 1449.21 492.722 1437.74 579.122C1436.28 590.055 1435.88 590.855 1433.34 590.322C1431.74 589.922 1423.48 587.788 1414.94 585.655C1391.08 579.255 1366.14 576.055 1341.08 576.055H1319.61L1310.41 562.322L1301.21 548.722L1301.34 512.055C1301.61 460.322 1306.41 419.922 1317.88 374.988C1330.01 326.722 1351.48 283.255 1370.54 268.055C1376.54 263.255 1377.48 263.122 1382.81 267.122Z" fill="black"/> +<path d="M975.877 938.189C946.944 940.989 939.077 942.055 925.21 944.855C902.677 949.522 872.544 959.922 851.61 970.189C778.81 1005.79 728.677 1065.12 713.344 1133.79C710.277 1147.39 709.877 1151.92 709.877 1174.86C709.877 1197.52 710.277 1202.46 713.21 1215.39C733.61 1305.12 816.277 1371.39 923.21 1383.52C946.41 1386.06 1046.68 1386.06 1069.88 1383.52C1155.74 1373.79 1229.61 1327.26 1262.81 1261.92C1271.61 1244.46 1275.88 1233.12 1279.88 1215.39C1282.81 1202.46 1283.21 1197.52 1283.21 1174.86C1283.21 1151.92 1282.81 1147.39 1279.74 1133.79C1257.48 1034.06 1160.68 955.522 1042.01 940.589C1026.54 938.722 986.01 937.122 975.877 938.189ZM1025.74 1010.72C1065.34 1014.99 1105.21 1029.12 1137.21 1050.46C1154.41 1061.92 1178.68 1085.92 1189.08 1101.66C1201.88 1121.12 1209.21 1140.99 1212.54 1165.12C1214.01 1176.19 1213.21 1184.59 1209.21 1202.46C1202.94 1229.12 1183.48 1256.99 1157.21 1276.46C1144.94 1285.39 1119.48 1298.32 1103.88 1303.39C1074.28 1312.86 1054.94 1314.59 985.877 1314.06C940.81 1313.66 932.81 1313.26 919.877 1310.86C875.744 1302.59 840.81 1284.99 815.477 1258.19C794.944 1236.59 785.61 1216.86 780.544 1184.99C778.277 1170.19 782.544 1145.66 791.21 1124.99C801.744 1099.79 828.944 1068.46 855.877 1050.46C887.077 1029.66 928.144 1014.86 965.877 1010.86C980.41 1009.26 1011.21 1009.26 1025.74 1010.72Z" fill="black"/> +<path d="M945.61 1108.06C935.477 1113.52 928.41 1127.39 930.543 1137.66C932.943 1148.72 942.677 1159.92 957.877 1169.12C966.01 1174.06 966.543 1174.72 966.943 1179.66C967.21 1182.59 966.143 1190.99 964.677 1198.46C963.077 1205.79 961.877 1213.52 961.877 1215.66C962.01 1221.39 967.343 1230.72 972.943 1235.26C977.877 1239.26 978.81 1239.39 992.677 1239.79C1005.34 1240.19 1008.01 1239.92 1013.08 1237.52C1026.14 1231.12 1029.48 1219.39 1024.68 1196.86C1020.68 1178.06 1021.48 1175.12 1031.48 1169.39C1042.01 1163.26 1053.21 1152.46 1056.54 1145.12C1062.94 1131.12 1057.08 1115.26 1042.94 1107.92C1039.48 1106.19 1035.21 1105.39 1028.94 1105.39C1019.21 1105.39 1012.94 1107.66 1001.48 1114.99L994.943 1119.12L990.81 1116.59C973.877 1106.59 970.81 1105.39 960.543 1105.52C953.21 1105.52 949.21 1106.19 945.61 1108.06Z" fill="black"/> +<path d="M621.878 953.255C598.278 960.722 580.678 978.055 571.611 1002.72C567.211 1014.46 565.078 1032.99 566.945 1042.99C571.345 1066.86 590.945 1088.59 613.211 1094.59C641.211 1101.92 662.145 1097.12 680.678 1078.72C691.478 1068.19 697.345 1058.99 703.211 1044.06C707.478 1033.52 707.745 1031.66 707.745 1016.72L707.878 1000.72L702.278 989.255C693.345 971.122 677.211 957.655 658.545 952.722C648.011 950.055 631.078 950.189 621.878 953.255Z" fill="black"/> +<path d="M1334.01 952.855C1315.74 957.789 1299.48 971.389 1290.81 989.255L1285.21 1000.72L1285.34 1016.72C1285.34 1031.66 1285.61 1033.52 1289.88 1044.06C1295.74 1058.99 1301.61 1068.19 1312.41 1078.72C1330.94 1097.12 1351.88 1101.92 1379.88 1094.59C1396.01 1090.32 1412.14 1076.72 1419.88 1060.86C1426.54 1047.39 1428.14 1037.66 1426.01 1022.32C1421.08 987.255 1400.54 961.789 1370.01 952.855C1361.08 950.189 1343.74 950.189 1334.01 952.855Z" fill="black"/> +</g> +<defs> +<clipPath id="clip0_1758_1066"> +<rect width="5849.33" height="2016" fill="transparent"/> +</clipPath> +</defs> +</svg> diff --git a/studio/frontend/public/provider-logos/vllm.svg b/studio/frontend/public/provider-logos/vllm.svg new file mode 100644 index 0000000000..0c8a13de01 --- /dev/null +++ b/studio/frontend/public/provider-logos/vllm.svg @@ -0,0 +1 @@ +<svg version="1.1" viewBox="0.0 0.0 96.0 96.0" fill="none" stroke="none" stroke-linecap="square" stroke-miterlimit="10" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns="http://www.w3.org/2000/svg"><clipPath id="g31e21232314_0_33.0"><path d="m0 0l96.0 0l0 96.0l-96.0 0l0 -96.0z" clip-rule="nonzero"/></clipPath><g clip-path="url(#g31e21232314_0_33.0)"><path fill="#d9d9d9" d="m41.04961 80.271324l1.8897629 0l0 2.3307114l-1.8897629 0z" fill-rule="evenodd"/><path fill="#d9d9d9" d="m42.221855 81.45145l1.8897629 0l0 2.3307037l-1.8897629 0z" fill-rule="evenodd"/><g filter="url(#shadowFilter-g31e21232314_0_33.1)"><use xlink:href="#g31e21232314_0_33.1" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.1" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.1"><path fill="#d9d9d9" d="m42.22417 28.470434l0 55.307083l-27.653543 -55.307083z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.2)"><use xlink:href="#g31e21232314_0_33.2" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.2" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.2"><path fill="#d9d9d9" d="m42.223038 83.77752l21.729656 0l18.653545 -70.385826l-25.574802 13.461943z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.3)"><use xlink:href="#g31e21232314_0_33.3" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.3" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.3"><path fill="#fdb515" d="m41.0477 27.293962l0 55.30709l-27.653542 -55.30709z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.4)"><use xlink:href="#g31e21232314_0_33.4" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.4" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.4"><path fill="#30a2ff" d="m41.046566 82.60105l21.72966 0l18.653545 -70.385826l-25.574806 13.461943z" fill-rule="evenodd"/></g></g></svg> diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index b4f7dd08d2..dc8bffb2b7 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -10,10 +10,12 @@ import { } from "@/components/ui/popover"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { usePlatformStore } from "@/config/env"; +import { isCustomProviderType } from "@/features/chat/external-providers"; import { cn } from "@/lib/utils"; import { ArrowDown01Icon, CloudIcon, + DashboardSquare01Icon, FolderSearchIcon, Logout01Icon, Search01Icon, @@ -40,6 +42,9 @@ const PROVIDER_LOGO_EXT: Record<string, "svg" | "png" | "jpg"> = { kimi: "jpg", qwen: "png", openrouter: "svg", + vllm: "svg", + ollama: "svg", + llama_cpp: "svg", }; function providerLogoSrc(providerType: string | undefined): string | undefined { @@ -59,6 +64,17 @@ function ExternalProviderLogo({ title?: string; }) { const src = providerLogoSrc(providerType); + if (!src && isCustomProviderType(providerType)) { + return ( + <span title={title} aria-hidden={true} className="inline-flex shrink-0"> + <HugeiconsIcon + icon={DashboardSquare01Icon} + className={cn("shrink-0", className)} + /> + </span> + ); + } + if (!src) return null; return ( <img diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 4988cd5a46..0579f84896 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -510,6 +510,10 @@ const ReasoningToggle: FC = () => { ? getExternalReasoningCapabilities( selectedExternalProvider?.providerType, effectiveExternalModelId, + { + isReasoningProvider: + selectedExternalProvider?.isReasoningModel === true, + }, ) : null; const effectiveReasoningStyle = diff --git a/studio/frontend/src/features/chat/api-provider-logo.tsx b/studio/frontend/src/features/chat/api-provider-logo.tsx index bd4f05b2ff..09794e3acf 100644 --- a/studio/frontend/src/features/chat/api-provider-logo.tsx +++ b/studio/frontend/src/features/chat/api-provider-logo.tsx @@ -4,6 +4,7 @@ import { cn } from "@/lib/utils"; import { DashboardSquare01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { isCustomProviderType } from "./external-providers"; /** * Registry logos live at `public/provider-logos/{provider_type}.{ext}` where `provider_type` @@ -19,6 +20,9 @@ const PROVIDER_LOGO_EXT: Record<string, "svg" | "png" | "jpg"> = { kimi: "jpg", qwen: "png", openrouter: "svg", + vllm: "svg", + ollama: "svg", + llama_cpp: "svg", }; export function apiProviderLogoSrc( @@ -42,7 +46,8 @@ interface ApiProviderLogoProps { * OpenAI's asset is black-on-transparent; it is inverted in dark mode for contrast. */ export function ApiProviderLogo({ providerType, className, title }: ApiProviderLogoProps) { - if (providerType === "custom") { + const src = apiProviderLogoSrc(providerType); + if (!src && isCustomProviderType(providerType)) { return ( <span title={title} aria-hidden className="inline-flex shrink-0"> <HugeiconsIcon icon={DashboardSquare01Icon} className={cn("shrink-0", className)} /> @@ -50,7 +55,6 @@ export function ApiProviderLogo({ providerType, className, title }: ApiProviderL ); } - const src = apiProviderLogoSrc(providerType); if (!src) return null; return ( <img diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 30b256f8ba..29c999678d 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -28,6 +28,8 @@ import { getExternalProviderApiKey, loadExternalProviders, parseExternalModelId, + supportsProviderPromptCaching, + toExternalBackendProviderType, } from "../external-providers"; import { EXTERNAL_MAX_OUTPUT_TOKENS, @@ -742,13 +744,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (isExternalRequest && !externalProvider) { toast.error("External provider not found.", { - description: "Open API Providers and re-add this provider.", + description: "Open Connections and re-add this provider.", }); throw new Error("External provider not found."); } if (isExternalRequest && !externalApiKey) { toast.error("Missing API key for selected external provider.", { - description: "Open API Providers and set the API key again.", + description: "Open Connections and set the API key again.", }); throw new Error("Missing external provider API key."); } @@ -929,10 +931,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { supportsPreserveThinking, preserveThinking, } = runtime; - const externalBackendProviderType = - externalProvider?.providerType === "custom" - ? "openai" - : externalProvider?.providerType; + const externalBackendProviderType = toExternalBackendProviderType( + externalProvider?.providerType, + ); const externalCapabilities = getProviderCapabilities( externalProvider?.providerType, ); @@ -943,6 +944,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ? getExternalReasoningCapabilities( externalProvider.providerType, externalSelection.modelId, + { + isReasoningProvider: + externalProvider.isReasoningModel === true, + }, ) : { supportsReasoning, @@ -1028,6 +1033,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { forceRefreshPublicKey, ), provider_base_url: externalProvider.baseUrl || null, + ...(supportsProviderPromptCaching(externalProvider.providerType) + ? { + enable_prompt_caching: + externalProvider.enablePromptCaching ?? true, + } + : {}), ...(externalReasoningCaps.supportsReasoning ? externalReasoningCaps.reasoningStyle === "reasoning_effort" ? externalReasoningEnabled diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 3fa294d465..2ebd8f45f0 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -550,6 +550,7 @@ export function ChatPage(): ReactElement { const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen); const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen); const externalProviders = useExternalProvidersStore((s) => s.providers); + const setExternalProviders = useExternalProvidersStore((s) => s.setProviders); useEffect(() => { const threadId = search.thread; @@ -629,14 +630,16 @@ export function ChatPage(): ReactElement { const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle); const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort); const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff); - const activeExternalProviderType = useMemo(() => { + const activeExternalProvider = useMemo(() => { const selection = parseExternalModelId(inferenceParams.checkpoint); if (!selection) return null; - const provider = externalProviders.find( - (p) => p.id === selection.providerId, + return ( + externalProviders.find( + (p) => p.id === selection.providerId, + ) ?? null ); - return provider?.providerType ?? null; }, [externalProviders, inferenceParams.checkpoint]); + const activeExternalProviderType = activeExternalProvider?.providerType ?? null; const activeProviderCapabilities = useMemo(() => { const selection = parseExternalModelId(inferenceParams.checkpoint); if (!selection) return null; @@ -671,6 +674,7 @@ export function ChatPage(): ReactElement { const reasoningCaps = getExternalReasoningCapabilities( provider?.providerType, selection.modelId, + { isReasoningProvider: provider?.isReasoningModel === true }, ); const state = useChatRuntimeStore.getState(); const preferredEffort = state.reasoningEffort; @@ -863,6 +867,10 @@ export function ChatPage(): ReactElement { const reasoningCaps = getExternalReasoningCapabilities( selectedProvider?.providerType, selectedExternal?.modelId, + { + isReasoningProvider: + selectedProvider?.isReasoningModel === true, + }, ); const preferredEffort = store.reasoningEffort; const effortLevels = reasoningCaps.reasoningEffortLevels; @@ -1426,6 +1434,14 @@ export function ChatPage(): ReactElement { onParamsChange={setInferenceParams} isExternalModel={isExternalModel} providerCapabilities={activeProviderCapabilities} + activeExternalProvider={activeExternalProvider} + onExternalProviderChange={(updatedProvider) => { + setExternalProviders( + externalProviders.map((provider) => + provider.id === updatedProvider.id ? updatedProvider : provider, + ), + ); + }} externalProviderType={activeExternalProviderType} onReloadModel={() => { const state = useChatRuntimeStore.getState(); diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 3297e6a992..0127369eb4 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -15,7 +15,9 @@ import { Label } from "@/components/ui/label"; import { Select, SelectContent, + SelectGroup, SelectItem, + SelectSeparator, SelectTrigger, SelectValue, } from "@/components/ui/select"; @@ -23,7 +25,6 @@ import { Spinner } from "@/components/ui/spinner"; import { Textarea } from "@/components/ui/textarea"; import { ArrowLeft02Icon, - DashboardSquare01Icon, Delete02Icon, Edit03Icon, PlusSignIcon, @@ -47,9 +48,19 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { + CUSTOM_BACKEND_PROVIDER_TYPE, + CUSTOM_PROVIDER_PRESETS, + customProviderBaseUrlPlaceholder, + customProviderDisplayName, + customProviderModelIdsPlaceholder, getExternalProviderApiKey, + isCustomProviderType, + LEGACY_CUSTOM_PROVIDER_TYPE, removeExternalProviderApiKey, setExternalProviderApiKey, + supportsProviderPromptCaching, + supportsProviderReasoningToggle, + toExternalBackendProviderType, } from "./external-providers"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ @@ -57,12 +68,11 @@ const PROVIDER_FORM_EASE: [number, number, number, number] = [ 0.165, 0.84, 0.44, 1, ]; const PROVIDER_FORM_DURATION = 0.2; -const CUSTOM_PROVIDER_TYPE = "custom"; -const CUSTOM_BACKEND_PROVIDER_TYPE = "openai"; const CUSTOM_PROVIDER_MISSING_KEY_MESSAGE = "No API key found, please make sure API key is added and valid for this provider."; const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/; const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]); +const HIDDEN_PROVIDER_TYPES = new Set(["qwen"]); const OPENROUTER_EXCLUDED_MODELS = new Set([ "google/chirp-3", "kwaivgi/kling-v3.0-pro", @@ -82,37 +92,37 @@ function resolveUiProviderTypeFromConfig( registryRows: ProviderRegistryEntry[], existingProviderType: string | undefined, ): string { - if (existingProviderType === CUSTOM_PROVIDER_TYPE) { - return CUSTOM_PROVIDER_TYPE; + if (existingProviderType && isCustomProviderType(existingProviderType)) { + return existingProviderType; } if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) { return configProviderType; } + const displayName = (configDisplayName ?? "").trim().toLowerCase(); + const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find( + (preset) => preset.displayName.toLowerCase() === displayName, + ); + if (matchingCustomPreset) { + return matchingCustomPreset.providerType; + } const openAiRegistry = registryRows.find( (entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE, ); if (!openAiRegistry) { return configProviderType; } - const displayName = (configDisplayName ?? "").trim().toLowerCase(); const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase(); if (displayName.length > 0 && displayName !== openAiDisplayName) { - return CUSTOM_PROVIDER_TYPE; + return LEGACY_CUSTOM_PROVIDER_TYPE; } const configUrl = normalizeUrl(configBaseUrl ?? ""); const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? ""); if (configUrl.length > 0 && configUrl !== defaultUrl) { - return CUSTOM_PROVIDER_TYPE; + return LEGACY_CUSTOM_PROVIDER_TYPE; } return configProviderType; } -function toBackendProviderType(uiProviderType: string): string { - return uiProviderType === CUSTOM_PROVIDER_TYPE - ? CUSTOM_BACKEND_PROVIDER_TYPE - : uiProviderType; -} - function parseManualModelIds(text: string): string[] { const seen = new Set<string>(); const out: string[] = []; @@ -175,22 +185,22 @@ export function ChatProvidersSettings({ const [manualModelIds, setManualModelIds] = useState(""); const [modelSearchQuery, setModelSearchQuery] = useState(""); const [customProviderName, setCustomProviderName] = useState("Custom"); + const [isReasoningModel, setIsReasoningModel] = useState(false); const reduceMotion = useReducedMotion(); - const isCustomProvider = providerType === CUSTOM_PROVIDER_TYPE; + const isCustomProvider = isCustomProviderType(providerType); + const showReasoningToggle = supportsProviderReasoningToggle(providerType); const registryByType = useMemo( () => new Map(registry.map((entry) => [entry.provider_type, entry])), [registry], ); - const hasCustomInRegistry = registryByType.has(CUSTOM_PROVIDER_TYPE); - const isCuratedModelList = useMemo(() => { return registryByType.get(providerType)?.model_list_mode === "curated"; }, [registryByType, providerType]); const isManualModelList = isCustomProvider || isCuratedModelList; const modelsPanelKey = isCustomProvider - ? "custom" + ? providerType || "custom" : isCuratedModelList ? "curated" : "remote"; @@ -225,7 +235,12 @@ export function ChatProvidersSettings({ useEffect(() => { if (!providerType || editingProviderId) return; const entry = registryByType.get(providerType); - if (!entry) return; + if (!entry) { + if (isCustomProviderType(providerType)) { + setCustomProviderName(customProviderDisplayName(providerType)); + } + return; + } // Seed the registry's default_models for every provider — curated and // remote alike. For remote-mode providers, loadModels() will replace // this with the union of defaults + the live /models response once the @@ -297,6 +312,12 @@ export function ChatProvidersSettings({ baseUrl: config.base_url ?? "", models: existingModels, availableModels: existing?.availableModels ?? [], + enablePromptCaching: supportsProviderPromptCaching(uiProviderType) + ? (existing?.enablePromptCaching ?? true) + : undefined, + isReasoningModel: supportsProviderReasoningToggle(uiProviderType) + ? existing?.isReasoningModel === true + : undefined, createdAt: existing?.createdAt ?? createdAt, updatedAt, }; @@ -328,7 +349,8 @@ export function ChatProvidersSettings({ setSelectedModelIds([]); setManualModelIds(""); setModelSearchQuery(""); - setCustomProviderName("Custom"); + setCustomProviderName(customProviderDisplayName(providerType)); + setIsReasoningModel(false); } function openAddProvider() { @@ -384,7 +406,7 @@ export function ChatProvidersSettings({ const trimmed = input.trim(); if (!trimmed) { if (required) { - throw new Error("Base URL is required for custom providers."); + throw new Error("Base URL is required for this connection."); } return null; } @@ -397,7 +419,7 @@ export function ChatProvidersSettings({ return; } if (isCustomProvider) { - toast.info("Custom providers use manual model IDs."); + toast.info("This connection uses manual model IDs."); return; } if (isCuratedModelList) { @@ -458,10 +480,10 @@ export function ChatProvidersSettings({ toast.error("Choose a provider first."); return; } - const backendProviderType = toBackendProviderType(providerType); + const backendProviderType = toExternalBackendProviderType(providerType); const selectedRegistryEntry = registryByType.get(backendProviderType); const displayName = isCustomProvider - ? customProviderName.trim() || "Custom" + ? customProviderName.trim() || customProviderDisplayName(providerType) : (selectedRegistryEntry?.display_name ?? providerType); if (!isCustomProvider && !apiKey.trim()) { toast.error("API key is required."); @@ -511,17 +533,21 @@ export function ChatProvidersSettings({ const updatedAt = Number.isFinite(Date.parse(created.updated_at)) ? Date.parse(created.updated_at) : Date.now(); + const uiProviderType = isCustomProvider + ? providerType + : created.provider_type; const provider: ExternalProviderConfig = { id: created.id, - providerType: isCustomProvider - ? CUSTOM_PROVIDER_TYPE - : created.provider_type, + providerType: uiProviderType, name: created.display_name, baseUrl: created.base_url ?? "", models: modelsToSave, availableModels: manualModels ? [] : pruneProviderModelIds(providerType, availableModels), + isReasoningModel: supportsProviderReasoningToggle(uiProviderType) + ? isReasoningModel + : undefined, createdAt, updatedAt, }; @@ -553,7 +579,7 @@ export function ChatProvidersSettings({ return; } const isEditingCustomProvider = - existing.providerType === CUSTOM_PROVIDER_TYPE; + isCustomProviderType(existing.providerType); if (!isEditingCustomProvider && !apiKey.trim()) { toast.error("API key is required."); return; @@ -597,7 +623,8 @@ export function ChatProvidersSettings({ ); const updated = await updateProviderConfig(editingProviderId, { displayName: isEditingCustomProvider - ? customProviderName.trim() || "Custom" + ? customProviderName.trim() || + customProviderDisplayName(existing.providerType) : existing.name, baseUrl, }); @@ -620,6 +647,11 @@ export function ChatProvidersSettings({ availableModels: manualModels ? [] : pruneProviderModelIds(existing.providerType, availableModels), + isReasoningModel: supportsProviderReasoningToggle( + existing.providerType, + ) + ? isReasoningModel + : undefined, updatedAt, } : provider, @@ -640,12 +672,19 @@ export function ChatProvidersSettings({ setEditingProviderId(provider.id); setPage("form"); setProviderType(provider.providerType); - setCustomProviderName(provider.name || "Custom"); + setCustomProviderName( + provider.name || customProviderDisplayName(provider.providerType), + ); setApiKey(getExternalProviderApiKey(provider.id)); setShowApiKey(false); setBaseUrlDraft(provider.baseUrl); setModelSearchQuery(""); - if (provider.providerType === CUSTOM_PROVIDER_TYPE) { + setIsReasoningModel( + supportsProviderReasoningToggle(provider.providerType) + ? provider.isReasoningModel === true + : false, + ); + if (isCustomProviderType(provider.providerType)) { setAvailableModels([]); setSelectedModelIds([]); setManualModelIds(provider.models.join("\n")); @@ -696,7 +735,7 @@ export function ChatProvidersSettings({ async function testProvider(provider: ExternalProviderConfig) { const savedKey = getExternalProviderApiKey(provider.id).trim(); if (!savedKey) { - if (provider.providerType === CUSTOM_PROVIDER_TYPE) { + if (isCustomProviderType(provider.providerType)) { await editProvider(provider); toast.info(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE); return; @@ -707,7 +746,9 @@ export function ChatProvidersSettings({ } try { const result = await testProviderConnection({ - providerType: toBackendProviderType(provider.providerType), + providerType: + toExternalBackendProviderType(provider.providerType) ?? + provider.providerType, apiKey: savedKey, baseUrl: provider.baseUrl || null, }); @@ -715,7 +756,7 @@ export function ChatProvidersSettings({ toast.success(result.message); } else { if ( - provider.providerType === CUSTOM_PROVIDER_TYPE && + isCustomProviderType(provider.providerType) && result.message.includes("Illegal header value b'Bearer '") ) { toast.error(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE); @@ -726,7 +767,7 @@ export function ChatProvidersSettings({ } catch (error) { const message = error instanceof Error ? error.message : "Unknown error"; if ( - provider.providerType === CUSTOM_PROVIDER_TYPE && + isCustomProviderType(provider.providerType) && message.includes("Illegal header value b'Bearer '") ) { toast.error(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE); @@ -753,7 +794,7 @@ export function ChatProvidersSettings({ </Button> <div className="flex min-w-0 items-center gap-2 leading-none"> <span className="text-xs font-medium text-muted-foreground"> - Cloud + Connections </span> <span className="size-1 rounded-full bg-muted-foreground/35" /> <span className="truncate text-xs font-medium text-muted-foreground"> @@ -774,7 +815,7 @@ export function ChatProvidersSettings({ Provider </Label> <p className="text-xs leading-snug text-muted-foreground"> - Supported registry or Custom. + Supported registry or local OpenAI-compatible connection. </p> </div> <Select @@ -786,6 +827,9 @@ export function ChatProvidersSettings({ setSelectedModelIds([]); setManualModelIds(""); setModelSearchQuery(""); + if (isCustomProviderType(value)) { + setCustomProviderName(customProviderDisplayName(value)); + } }} > <SelectTrigger @@ -796,32 +840,46 @@ export function ChatProvidersSettings({ <SelectValue placeholder="Choose a provider" /> </SelectTrigger> <SelectContent> - {registry.map((entry) => ( - <SelectItem - key={entry.provider_type} - value={entry.provider_type} - > - <span className="flex items-center gap-2"> - <ApiProviderLogo - providerType={entry.provider_type} - className="size-4" - title={entry.display_name} - /> - {entry.display_name} - </span> - </SelectItem> - ))} - {hasCustomInRegistry ? null : ( - <SelectItem value={CUSTOM_PROVIDER_TYPE}> - <span className="flex items-center gap-2"> - <HugeiconsIcon - icon={DashboardSquare01Icon} - className="size-4" - /> - Custom - </span> - </SelectItem> - )} + <SelectGroup> + {CUSTOM_PROVIDER_PRESETS.map((preset) => ( + <SelectItem + key={preset.providerType} + value={preset.providerType} + > + <span className="flex items-center gap-2"> + <ApiProviderLogo + providerType={preset.providerType} + className="size-4" + title={preset.displayName} + /> + {preset.displayName} + </span> + </SelectItem> + ))} + </SelectGroup> + <SelectSeparator /> + <SelectGroup> + {registry + .filter( + (entry) => + !HIDDEN_PROVIDER_TYPES.has(entry.provider_type), + ) + .map((entry) => ( + <SelectItem + key={entry.provider_type} + value={entry.provider_type} + > + <span className="flex items-center gap-2"> + <ApiProviderLogo + providerType={entry.provider_type} + className="size-4" + title={entry.display_name} + /> + {entry.display_name} + </span> + </SelectItem> + ))} + </SelectGroup> </SelectContent> </Select> </div> @@ -902,11 +960,35 @@ export function ChatProvidersSettings({ type="text" value={baseUrlDraft} onChange={(event) => setBaseUrlDraft(event.target.value)} - placeholder="https://my-vllm-server.com/v1" + placeholder={customProviderBaseUrlPlaceholder(providerType)} className="h-9 text-sm" /> </div> ) : null} + + {showReasoningToggle ? ( + <div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1"> + <Label + htmlFor="provider-is-reasoning" + className="text-sm font-medium" + > + Reasoning model + </Label> + <label + htmlFor="provider-is-reasoning" + className="flex cursor-pointer items-center gap-2 text-sm" + > + <Checkbox + id="provider-is-reasoning" + checked={isReasoningModel} + onCheckedChange={(checked) => + setIsReasoningModel(checked === true) + } + /> + This server runs a reasoning model + </label> + </div> + ) : null} </div> </section> @@ -952,7 +1034,7 @@ export function ChatProvidersSettings({ } title={ isCustomProvider - ? "Custom providers use manual model IDs" + ? "This connection uses manual model IDs" : isCuratedModelList ? "Full catalog is not fetched for this provider" : undefined @@ -986,7 +1068,7 @@ export function ChatProvidersSettings({ onChange={(event) => setManualModelIds(event.target.value) } - placeholder={"gpt-4o-mini\nQwen/Qwen3-14B"} + placeholder={customProviderModelIdsPlaceholder(providerType)} rows={5} className="min-h-[100px] resize-y font-mono text-sm" /> @@ -1197,9 +1279,9 @@ export function ChatProvidersSettings({ <div className="flex min-h-0 flex-col gap-6"> <header className="flex flex-col gap-1 pr-8"> <div className="flex min-w-0 flex-col gap-1"> - <h1 className="font-heading text-lg font-semibold">Cloud</h1> + <h1 className="font-heading text-lg font-semibold">Connections</h1> <p className="text-xs leading-relaxed text-muted-foreground"> - Manage cloud provider connections for chat through the Studio proxy. + Manage model provider connections for chat through the Studio proxy. </p> </div> </header> @@ -1237,7 +1319,8 @@ export function ChatProvidersSettings({ const detail = provider.baseUrl || registryEntry?.base_url || ""; const providerLabel = - registryEntry?.display_name ?? provider.providerType; + registryEntry?.display_name ?? + customProviderDisplayName(provider.providerType); const modelSummary = formatModelSummary(provider.models); return ( <div @@ -1346,9 +1429,9 @@ export function ChatProvidersDialog({ className="flex max-h-[90dvh] w-[96vw] flex-col gap-0 overflow-y-auto p-8 sm:max-w-none md:max-w-[44rem]" > <DialogHeader className="sr-only"> - <DialogTitle>Cloud</DialogTitle> + <DialogTitle>Connections</DialogTitle> <DialogDescription> - Manage external model providers for chat. + Manage external model connections for chat. </DialogDescription> </DialogHeader> <ChatProvidersSettings diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 1eb6fc35d7..74e4943b3c 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -64,6 +64,10 @@ import { Fragment, type ReactNode } from "react"; import { useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; +import { + type ExternalProviderConfig, + supportsProviderPromptCaching, +} from "./external-providers"; import { applyPresetParams, BUILTIN_PRESET_NAMES, @@ -517,6 +521,8 @@ interface ChatSettingsPanelProps { * per-param visibility in the sampling section. */ providerCapabilities?: ProviderCapabilities | null; + activeExternalProvider?: ExternalProviderConfig | null; + onExternalProviderChange?: (provider: ExternalProviderConfig) => void; /** * Backend provider type for the active external model (e.g. "kimi", * "anthropic", "openai"), or `null` for local models. Drives the @@ -533,6 +539,8 @@ export function ChatSettingsPanel({ onParamsChange, isExternalModel = false, providerCapabilities = null, + activeExternalProvider = null, + onExternalProviderChange, externalProviderType = null, onReloadModel, }: ChatSettingsPanelProps) { @@ -662,6 +670,11 @@ export function ChatSettingsPanel({ Boolean(currentCheckpoint) && modelRequiresTrustRemoteCode && !(params.trustRemoteCode ?? false); + const showPromptCachingControl = + activeExternalProvider != null && + supportsProviderPromptCaching(activeExternalProvider.providerType); + const promptCachingEnabled = + activeExternalProvider?.enablePromptCaching !== false; function set<K extends keyof InferenceParams>(key: K) { return (v: InferenceParams[K]) => { @@ -1145,6 +1158,32 @@ export function ChatSettingsPanel({ </div> </CollapsibleSection> + {showPromptCachingControl && activeExternalProvider ? ( + <CollapsibleSection label="Provider" defaultOpen={true}> + <div className="flex items-center justify-between gap-3 pt-1"> + <div className="flex min-w-0 items-center gap-1.5"> + <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg"> + Prompt caching + </span> + <InfoHint> + Reuse compatible prompt prefixes for lower latency and cost. + </InfoHint> + </div> + <Switch + className="panel-switch shrink-0" + checked={promptCachingEnabled} + onCheckedChange={(checked) => { + onExternalProviderChange?.({ + ...activeExternalProvider, + enablePromptCaching: checked, + }); + }} + aria-label="Enable prompt caching" + /> + </div> + </CollapsibleSection> + ) : null} + <CollapsibleSection label="System Prompt" defaultOpen={true}> <button type="button" diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 867d0461cc..5f645042bf 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -14,10 +14,147 @@ export interface ExternalProviderConfig { models: string[]; /** Cached available model ids from the provider's /models response. */ availableModels?: string[]; + /** Whether to ask supported hosted providers to use prompt caching. */ + enablePromptCaching?: boolean; + /** User-pinned: the loaded vLLM model supports `enable_thinking`. */ + isReasoningModel?: boolean; createdAt: number; updatedAt: number; } +const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]); + +export function supportsProviderPromptCaching( + providerType: string | null | undefined, +): boolean { + return providerType != null && PROMPT_CACHING_PROVIDER_TYPES.has(providerType); +} + +// Provider types that expose the connection-level "reasoning model" +// toggle. vLLM's OpenAI-compat endpoint doesn't advertise this per model. +const REASONING_TOGGLE_PROVIDER_TYPES = new Set(["vllm"]); + +export function supportsProviderReasoningToggle( + providerType: string | null | undefined, +): boolean { + return ( + providerType != null && REASONING_TOGGLE_PROVIDER_TYPES.has(providerType) + ); +} + +export const CUSTOM_BACKEND_PROVIDER_TYPE = "openai"; +export const LEGACY_CUSTOM_PROVIDER_TYPE = "custom"; + +export const CUSTOM_PROVIDER_PRESETS = [ + { + providerType: "llama_cpp", + displayName: "llama.cpp", + baseUrlPlaceholder: "http://localhost:8080/v1", + modelIdsPlaceholder: "gpt-oss-20b\nqwen3-14b", + }, + { + providerType: "vllm", + displayName: "vLLM", + baseUrlPlaceholder: "https://my-vllm-server.com/v1", + modelIdsPlaceholder: "openai/gpt-oss-20b\nQwen/Qwen3-14B", + }, + { + providerType: "ollama", + displayName: "Ollama", + baseUrlPlaceholder: "http://localhost:11434/v1", + modelIdsPlaceholder: "gpt-oss:20b\nqwen3:14b", + }, +] as const; + +const CUSTOM_PROVIDER_LABELS: Record<string, string> = { + [LEGACY_CUSTOM_PROVIDER_TYPE]: "Custom", + ...Object.fromEntries( + CUSTOM_PROVIDER_PRESETS.map((preset) => [ + preset.providerType, + preset.displayName, + ]), + ), +}; + +const CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS: Record<string, string> = { + [LEGACY_CUSTOM_PROVIDER_TYPE]: "https://my-vllm-server.com/v1", + ...Object.fromEntries( + CUSTOM_PROVIDER_PRESETS.map((preset) => [ + preset.providerType, + preset.baseUrlPlaceholder, + ]), + ), +}; + +const CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS: Record<string, string> = { + [LEGACY_CUSTOM_PROVIDER_TYPE]: "openai/gpt-oss-20b\nQwen/Qwen3-14B", + ...Object.fromEntries( + CUSTOM_PROVIDER_PRESETS.map((preset) => [ + preset.providerType, + preset.modelIdsPlaceholder, + ]), + ), +}; + +export function isCustomProviderType( + providerType: string | null | undefined, +): boolean { + if (!providerType) return false; + return providerType in CUSTOM_PROVIDER_LABELS; +} + +export function customProviderDisplayName( + providerType: string | null | undefined, +): string { + if (!providerType) return "Custom"; + return CUSTOM_PROVIDER_LABELS[providerType] ?? providerType; +} + +export function customProviderBaseUrlPlaceholder( + providerType: string | null | undefined, +): string { + if (!providerType) { + return CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE]; + } + return ( + CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[providerType] ?? + CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE] + ); +} + +export function customProviderModelIdsPlaceholder( + providerType: string | null | undefined, +): string { + if (!providerType) { + return CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE]; + } + return ( + CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[providerType] ?? + CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE] + ); +} + +export function toExternalBackendProviderType(providerType: string): string; +export function toExternalBackendProviderType( + providerType: null | undefined, +): undefined; +export function toExternalBackendProviderType( + providerType: string | null | undefined, +): string | undefined; +export function toExternalBackendProviderType( + providerType: string | null | undefined, +): string | undefined { + if (!providerType) return undefined; + // vLLM's /v1/responses applies the loaded model's chat template, which + // 400s on strict-alternation templates (e.g. Gemma 3). Pass the actual + // type through so the backend routes vLLM to /v1/chat/completions instead + // of the OpenAI Responses path used for gpt-5.x. + if (providerType === "vllm") return "vllm"; + return isCustomProviderType(providerType) + ? CUSTOM_BACKEND_PROVIDER_TYPE + : providerType; +} + const EXTERNAL_PROVIDERS_KEY = "unsloth_chat_external_providers"; const EXTERNAL_PROVIDER_KEYS_KEY = "unsloth_chat_external_provider_keys"; const EXTERNAL_MODEL_PREFIX = "external::"; @@ -71,9 +208,10 @@ function mapLegacyPresetToProviderType(presetId: string): string { } function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig { + const providerType = raw.providerType.trim(); return { ...raw, - providerType: raw.providerType.trim(), + providerType, name: raw.name.trim(), baseUrl: raw.baseUrl.trim(), models: raw.models @@ -82,6 +220,12 @@ function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig availableModels: (raw.availableModels ?? []) .map((model) => model.trim()) .filter((model) => model.length > 0), + enablePromptCaching: supportsProviderPromptCaching(providerType) + ? raw.enablePromptCaching !== false + : undefined, + isReasoningModel: supportsProviderReasoningToggle(providerType) + ? raw.isReasoningModel === true + : undefined, }; } diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 44da42e27b..1c40a25773 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -221,10 +221,13 @@ const PROVIDER_CAPABILITIES: Record<string, ProviderCapabilities> = { // OpenRouter silently drops params the target model does not support, so we // surface every knob and let the gateway handle the per-model fan-out. openrouter: ALL_SUPPORTED, - // Custom providers are assumed OpenAI-compatible by the backend; users who - // point at vLLM/Ollama backends often want top_k / min_p / repetition, - // so be permissive. + // Local OpenAI-compatible connections are proxied through the OpenAI backend + // path, but vLLM/Ollama/llama.cpp users often want top_k / min_p / + // repetition controls, so be permissive. custom: ALL_SUPPORTED, + vllm: ALL_SUPPORTED, + ollama: ALL_SUPPORTED, + llama_cpp: ALL_SUPPORTED, }; const DEFAULT_EXTERNAL_CAPABILITIES = OPENAI_COMPAT_BASE; @@ -420,6 +423,25 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning return withEnableThinkingStyle(); } +export interface ExternalReasoningResolveOptions { + /** vLLM connection flagged as a reasoning model in provider config. */ + isReasoningProvider?: boolean; +} + +// vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle. +function resolveConnectionLevelReasoning( + normalizedProvider: string, + options: ExternalReasoningResolveOptions | undefined, +): ExternalReasoningCapabilities | null { + if (normalizedProvider === "vllm" && options?.isReasoningProvider) { + return withEnableThinkingStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + }); + } + return null; +} + /** * resolve external-model thinking capabilities. * provider-specific matching lives in the OpenAI/Anthropic resolvers. @@ -428,9 +450,17 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning export function getExternalReasoningCapabilities( providerType: string | null | undefined, modelId: string | null | undefined, + options?: ExternalReasoningResolveOptions, ): ExternalReasoningCapabilities { const normalizedModel = modelId?.trim().toLowerCase() ?? ""; const normalizedProvider = providerType?.trim().toLowerCase() ?? ""; + const connectionLevel = resolveConnectionLevelReasoning( + normalizedProvider, + options, + ); + if (connectionLevel) { + return connectionLevel; + } if (!normalizedModel) { return withEnableThinkingStyle(); } diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 18f0673433..da32dee3a4 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -330,6 +330,10 @@ export function SharedComposer({ ? getExternalReasoningCapabilities( selectedExternalProvider?.providerType, effectiveExternalModelId, + { + isReasoningProvider: + selectedExternalProvider?.isReasoningModel === true, + }, ) : null; const isExternalOpenAIReasoning = diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 3fa8fcd3a7..18270a7620 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -224,6 +224,7 @@ export interface OpenAIChatCompletionsRequest { external_model?: string; encrypted_api_key?: string; provider_base_url?: string | null; + enable_prompt_caching?: boolean | null; } export interface OpenAIChatDelta { diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 4e0764c5f2..1ece504d6f 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -45,7 +45,7 @@ const TABS: TabDef[] = [ { id: "profile", label: "Profile", icon: UserIcon }, { id: "appearance", label: "Appearance", icon: PaintBrush02Icon }, { id: "chat", label: "Chat", icon: Message01Icon }, - { id: "connections", label: "Cloud", icon: CloudIcon, badge: "New" }, + { id: "connections", label: "Connections", icon: CloudIcon, badge: "New" }, { id: "api-keys", label: "API", icon: Globe02Icon, badge: "New" }, { id: "about", label: "Help", icon: HelpCircleIcon }, ]; From c7c3840b5f720c000ea94c5b1cbdaf8f2c5d984d Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 09:37:26 -0700 Subject: [PATCH 21/50] ci: cap each compiler-sweep iteration with SIGALRM + log progress (#5456) Core (HF=latest + TRL=latest) (transformers >=5,<6, trl >=1,<2) hangs 30+ minutes in the compiler-sweep test under the new shim layout, exceeding the 35-min job timeout and showing up as cancelled with no log of which model_type wedged. unsloth_compile_transformers does real source rewriting + torch.compile decoration and can deadlock inside a single problem model on a new transformers point release. Per-model SIGALRM cap (60s) so one infinite-loop model_type cannot wedge the whole sweep. Print sweep progress every 25 models so the log surfaces the slow model_type the next time this regresses -- crucial for finding the upstream/transformers compile bug. Timeout errors land in the same KNOWN / NEW_FAILURES bucket as any other compile exception, so the matrix still surfaces real regressions instead of silently absorbing them. --- .github/workflows/consolidated-tests-ci.yml | 75 +++++++++++++-------- 1 file changed, 47 insertions(+), 28 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index de928e83ee..3f4e054b00 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -977,40 +977,59 @@ jobs: skipped -> no `modeling_<x>.py` file (expected for some umbrella packages like `auto`, `deprecated`) known -> in KNOWN_BROKEN_COMPILE; tracked for follow-up. - Any uncaught failure fails the cell.""" + Any uncaught failure fails the cell. + + Per-model SIGALRM cap so one infinite-looping model_type + cannot wedge the whole sweep + nuke the job timeout + (observed on transformers >=5,<6 -- 30+ min hang before + this guard landed).""" import importlib as _il + import signal ok = 0 skipped = [] known = [] new_failures = [] - for model_type in _all_model_types(): - modeling_path = f"transformers.models.{model_type}.modeling_{model_type}" - try: - _il.import_module(modeling_path) - except (ModuleNotFoundError, ImportError): - skipped.append((model_type, "no modeling file")) - continue - try: - unsloth_compile_transformers( - model_type=model_type, fast_lora_forwards=False, - ) - except Exception as e: - msg = f"{type(e).__name__}: {str(e)[:200]}" + models = _all_model_types() + def _on_timeout(signum, frame): + raise TimeoutError("compile exceeded per-model budget") + prev_handler = signal.signal(signal.SIGALRM, _on_timeout) + try: + for i, model_type in enumerate(models): + if i % 25 == 0: + print(f" sweep progress: {i}/{len(models)} -> {model_type}", flush=True) + modeling_path = f"transformers.models.{model_type}.modeling_{model_type}" + try: + _il.import_module(modeling_path) + except (ModuleNotFoundError, ImportError): + skipped.append((model_type, "no modeling file")) + continue + signal.alarm(60) + try: + unsloth_compile_transformers( + model_type=model_type, fast_lora_forwards=False, + ) + except Exception as e: + signal.alarm(0) + msg = f"{type(e).__name__}: {str(e)[:200]}" + if model_type in KNOWN_BROKEN_COMPILE: + known.append((model_type, msg)) + else: + new_failures.append((model_type, msg)) + continue + signal.alarm(0) if model_type in KNOWN_BROKEN_COMPILE: - known.append((model_type, msg)) - else: - new_failures.append((model_type, msg)) - continue - if model_type in KNOWN_BROKEN_COMPILE: - # Came back green unexpectedly -- that's GOOD news, - # the bug was fixed. Surface it so we can drop the - # entry from KNOWN_BROKEN_COMPILE. - print( - f" UNEXPECTED-OK {model_type}: was in " - "KNOWN_BROKEN_COMPILE, now compiles cleanly. " - "Drop the entry." - ) - ok += 1 + # Came back green unexpectedly -- that's GOOD news, + # the bug was fixed. Surface it so we can drop the + # entry from KNOWN_BROKEN_COMPILE. + print( + f" UNEXPECTED-OK {model_type}: was in " + "KNOWN_BROKEN_COMPILE, now compiles cleanly. " + "Drop the entry." + ) + ok += 1 + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, prev_handler) print(f"\nCompile sweep: ok={ok} skipped={len(skipped)} " f"known-broken={len(known)} new-failures={len(new_failures)}") for m, r in known: From 51dd5fac79a16c0fc0d3f1302a1f41937ebbfef4 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 10:37:37 -0700 Subject: [PATCH 22/50] ci: add tx >=5,<6 slow compile model_types to KNOWN_BROKEN_COMPILE (#5458) The per-model SIGALRM cap landed on the previous fix now exposes beit / sam / sam_hq as compile-too-slow on transformers >=5,<6 + trl >=1,<2 -- each exceeds the 60s per-model budget. They are real slow paths in unsloth_compile_transformers's source rewriter when handling beit / SAM's encoder layers on the new transformers line, not infra flakes (the prior fix logged sweep progress per 25 models so the slow ones are pinpointable in CI logs). Bucket them into Category F (compile exceeds budget) so the sweep stays green and each is tracked for follow-up zoo fixes in the same shape as the existing 27 known-broken entries. Surface behaviour stays identical: any NEW slow model_type still fails the cell with a TimeoutError tag. --- .github/workflows/consolidated-tests-ci.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 3f4e054b00..6b008d4bb1 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -962,6 +962,12 @@ jobs: # Category E: undefined name in emitted file. "perceiver": "name 'AbstractPreprocessor' is not defined", "sam3_lite_text": "name 'Sam3LiteTextLayerScaledResidual' is not defined", + # Category F: compile exceeds 60s budget on the runner. + # First seen on transformers >=5,<6; each represents a slow + # or recursive source-rewriter path the zoo can address. + "beit": "TimeoutError: compile exceeds per-model budget", + "sam": "TimeoutError: compile exceeds per-model budget", + "sam_hq": "TimeoutError: compile exceeds per-model budget", } From 3596ce12df33cea70523f711bb46bf2eb47534c7 Mon Sep 17 00:00:00 2001 From: DoubleMathew <mmathew23@gmail.com> Date: Fri, 15 May 2026 12:43:49 -0500 Subject: [PATCH 23/50] Restore Flash > SDPA > Flex priority for non-gemma3 models (#5455) * update attn preferences * address gemini review suggestion * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> --- unsloth/models/_utils.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index df498e89fb..7bf8866f46 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -383,6 +383,7 @@ def _disable_flash_attention_if_needed( config, attn_implementation = None, supports_sdpa = False, + supports_flex_attention = False, would_use_flash_attention = False, disable_reason = None, ): @@ -402,7 +403,12 @@ def _disable_flash_attention_if_needed( if requested_attn_implementation == "eager": return _set_attn_impl(config, "eager") - fallback_attn_implementation = "sdpa" if supports_sdpa else "eager" + if supports_sdpa: + fallback_attn_implementation = "sdpa" + elif supports_flex_attention: + fallback_attn_implementation = "flex_attention" + else: + fallback_attn_implementation = "eager" if ( _is_flash_attention_requested(requested_attn_implementation) or would_use_flash_attention @@ -487,33 +493,32 @@ def resolve_attention_implementation( getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) ) + supports_flex_attention = _supports_flex_attention(model_class, config, model_type) disable_reason = _get_flash_attention_disable_reason(config) flash_attention_disabled = disable_reason is not None if model_class is None: attn_impl = _set_attn_impl(config, "sdpa" if supports_sdpa else "eager") else: - supports_flex_attention = _supports_flex_attention( - model_class, config, model_type - ) prefers_flex_attention = _config_prefers_flex_attention(config) if _is_eager_only(model_type): attn_impl = _set_attn_impl(config, "eager") elif prefers_flex_attention and supports_flex_attention: + # Models in _FLEX_PREFERRED_MODELS (gemma3 family) prefer flex_attention + # over flash. Caller can still override by passing + # requested_attn_implementation="sdpa" (handled below). attn_impl = _set_attn_impl(config, "flex_attention") elif ( - not prefers_flex_attention - and not flash_attention_disabled + not flash_attention_disabled and HAS_FLASH_ATTENTION and supports_flash_attention ): attn_impl = _set_attn_impl(config, "flash_attention_2") - elif supports_flex_attention: - attn_impl = _set_attn_impl(config, "flex_attention") elif flash_attention_disabled: attn_impl = _disable_flash_attention_if_needed( config, supports_sdpa = supports_sdpa, + supports_flex_attention = supports_flex_attention, would_use_flash_attention = ( HAS_FLASH_ATTENTION and supports_flash_attention ), @@ -521,6 +526,11 @@ def resolve_attention_implementation( ) elif supports_sdpa: attn_impl = _set_attn_impl(config, "sdpa") + elif supports_flex_attention: + # Flex is only a fallback for models that don't support SDPA + # (e.g. some custom configurations). Without this fallback such + # models would land on eager. + attn_impl = _set_attn_impl(config, "flex_attention") else: attn_impl = _set_attn_impl(config, "eager") @@ -531,6 +541,7 @@ def resolve_attention_implementation( config, requested_attn_implementation, supports_sdpa = supports_sdpa, + supports_flex_attention = supports_flex_attention, disable_reason = disable_reason, ) else: From 90ac4c87f7b69bd9a106c077526eaaa2ebea0319 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 11:02:16 -0700 Subject: [PATCH 24/50] ci: stop a partial mmproj cache from poisoning Mac Studio GGUF CI (#5459) The "JSON, images" Mac Studio GGUF CI job hit a stale cache for ${{ runner.os }}-gguf-...-mmproj-F16.gguf-v1 that contains only the main GGUF, not the mmproj sibling. cache-hit==true so the download step was skipped, then the post-load \`ls\` failed: ls: ...gguf-cache/mmproj-F16.gguf: No such file or directory Three guards layered: 1) Bump cache key v1 -> v2 to invalidate the poisoned entry on the GitHub-hosted side. 2) New verify-cache step explicitly checks BOTH files are present before trusting cache-hit. If not, fall through to download. 3) Save step gains a hashFiles() check on the mmproj path so a partial mmproj download cannot land back in the cache. Behaviour on a clean run is unchanged; cache hit + verify ok skips the re-download, partial-hit triggers fresh download, success saves a complete archive. --- .../workflows/studio-mac-inference-smoke.yml | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 4e8456a297..54d66d84b8 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -710,11 +710,22 @@ jobs: continue-on-error: true with: path: gguf-cache - key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2 - - name: Download GGUF + mmproj if cache miss + - name: Verify cache contains BOTH gguf + mmproj + id: verify-cache + if: steps.cache-gguf.outputs.cache-hit == 'true' + run: | + if [[ -f "gguf-cache/$GGUF_FILE" && -f "gguf-cache/$MMPROJ_FILE" ]]; then + echo "ok=true" >> "$GITHUB_OUTPUT" + else + echo "Partial cache hit -- forcing re-download." + echo "ok=false" >> "$GITHUB_OUTPUT" + fi + + - name: Download GGUF + mmproj if cache miss or partial id: download-gguf - if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.verify-cache.outputs.ok != 'true' # Authenticated + parallel: shared macos-14 NAT egress stalls # multi-GB anonymous downloads. env: @@ -734,13 +745,15 @@ jobs: ls -lh "gguf-cache/$GGUF_FILE" "gguf-cache/$MMPROJ_FILE" # Save partial caches on cancel. hashFiles guard avoids a hard - # save failure when the download step exits with no files. + # save failure when the download step exits with no files. The + # additional mmproj-presence check stops a partial save from + # poisoning the cache for the next run. - name: Save GGUF + mmproj files - if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != '' + if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != '' && hashFiles(format('gguf-cache/{0}', env.MMPROJ_FILE)) != '' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: gguf-cache - key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2 - name: Install Studio (--local, --no-torch) env: From ac3e9e98f206ddd8c8d3e6788d8c49386fc52aa2 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 11:46:52 -0700 Subject: [PATCH 25/50] ci: make Windows Stop Studio teardown tolerate Git Bash signal exit (#5460) The Windows-runner "Stop Studio" step's kill + sleep block has been observed to exit 143 (SIGTERM) even when the upstream test work passed. Most recently caught on PR #5432 Job 3 "JSON, images": all four assertions (json_object, plain inference, image/openai, image/anthropic) printed PASS, then the kill step ran for ~2 seconds and exited 143, failing the job. Teardown does not gate correctness. Wrap all three Stop Studio steps with set +e + redirected error streams + explicit exit 0 so transient Git Bash signal weirdness no longer masks a green test run. --- .../studio-windows-inference-smoke.yml | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index bc13ec8199..096ec95d03 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -345,9 +345,15 @@ jobs: - name: Stop Studio if: always() + # `set +e` + redirect everything: Git Bash on windows-latest + # has been observed to exit 143 from the kill/sleep block even + # when the upstream test work passed, masking a green run. The + # teardown does not gate correctness, so absorb any signal. run: | - kill "${STUDIO_PID}" 2>/dev/null || true - sleep 2 + set +e + kill "${STUDIO_PID}" >/dev/null 2>&1 || true + sleep 2 >/dev/null 2>&1 || true + exit 0 - name: Upload logs if: always() @@ -762,9 +768,15 @@ jobs: - name: Stop Studio if: always() + # `set +e` + redirect everything: Git Bash on windows-latest + # has been observed to exit 143 from the kill/sleep block even + # when the upstream test work passed, masking a green run. The + # teardown does not gate correctness, so absorb any signal. run: | - kill "${STUDIO_PID}" 2>/dev/null || true - sleep 2 + set +e + kill "${STUDIO_PID}" >/dev/null 2>&1 || true + sleep 2 >/dev/null 2>&1 || true + exit 0 - name: Upload logs if: always() @@ -1150,9 +1162,15 @@ jobs: - name: Stop Studio if: always() + # `set +e` + redirect everything: Git Bash on windows-latest + # has been observed to exit 143 from the kill/sleep block even + # when the upstream test work passed, masking a green run. The + # teardown does not gate correctness, so absorb any signal. run: | - kill "${STUDIO_PID}" 2>/dev/null || true - sleep 2 + set +e + kill "${STUDIO_PID}" >/dev/null 2>&1 || true + sleep 2 >/dev/null 2>&1 || true + exit 0 - name: Upload logs if: always() From a9b8c9a2214c43003efbb11895c0296f1481b3a8 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 15 May 2026 20:33:22 +0100 Subject: [PATCH 26/50] Studio: make API key optional for local providers (llama.cpp/vLLM/Ollama) (#5457) * make API key optional for local providers (llama.cpp/vLLM/Ollama)D * chore: reduce comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 9 +++-- studio/backend/models/providers.py | 10 +++-- studio/backend/routes/inference.py | 22 +++++----- studio/backend/routes/providers.py | 40 +++++++++++-------- .../src/features/chat/api/chat-adapter.ts | 19 ++++++--- .../src/features/chat/api/providers-api.ts | 6 ++- 6 files changed, 66 insertions(+), 40 deletions(-) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 48a5f0e879..07574abf4c 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -208,10 +208,11 @@ class ExternalProviderClient: auth_header = provider_info.get("auth_header", "Authorization") auth_prefix = provider_info.get("auth_prefix", "Bearer ") - headers = { - "Content-Type": "application/json", - auth_header: f"{auth_prefix}{self.api_key}", - } + headers = {"Content-Type": "application/json"} + # Skip auth header when api_key is empty (optional for local providers); + # httpx rejects an empty `Bearer ` value as "Illegal header value". + if self.api_key: + headers[auth_header] = f"{auth_prefix}{self.api_key}" # Merge any provider-specific extra headers (e.g. anthropic-version, OpenRouter attribution) headers.update(provider_info.get("extra_headers", {})) return headers diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 5678e69f62..53ce981392 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -95,8 +95,9 @@ class ProviderModelsRequest(BaseModel): """Request to list models from an external provider.""" provider_type: str = Field(..., description = "Provider type from the registry") - encrypted_api_key: str = Field( - ..., description = "RSA-encrypted, base64-encoded API key" + encrypted_api_key: Optional[str] = Field( + None, + description = "RSA-encrypted, base64-encoded API key (optional for local providers)", ) base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default)" @@ -110,8 +111,9 @@ class ProviderTestRequest(BaseModel): """Request to test connectivity to an external provider.""" provider_type: str = Field(..., description = "Provider type from the registry") - encrypted_api_key: str = Field( - ..., description = "RSA-encrypted, base64-encoded API key" + encrypted_api_key: Optional[str] = Field( + None, + description = "RSA-encrypted, base64-encoded API key (optional for local providers)", ) base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default)" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b223bcf981..e8732397b0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1554,15 +1554,16 @@ async def _proxy_to_external_provider( detail = f"Unknown provider type: {provider_type}", ) - # Decrypt the API key - try: - api_key = decrypt_api_key(payload.encrypted_api_key) - except Exception as exc: - logger.warning("external_provider.decrypt_failed", error = str(exc)) - raise HTTPException( - status_code = 400, - detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.", - ) + api_key = "" + if payload.encrypted_api_key: + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("external_provider.decrypt_failed", error = str(exc)) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.", + ) model = payload.external_model or payload.model if model == "default": @@ -1646,7 +1647,8 @@ async def openai_chat_completions( - Other models → Unsloth/transformers via InferenceBackend """ # ── External provider routing ──────────────────────────────── - if payload.encrypted_api_key and (payload.provider_id or payload.provider_type): + # encrypted_api_key is optional — local providers (llama.cpp / vLLM / Ollama) may run without auth. + if payload.provider_id or payload.provider_type: return await _proxy_to_external_provider(payload, request) llama_backend = get_llama_cpp_backend() diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index e21985d60b..acfaa6e427 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -200,14 +200,18 @@ async def test_provider( detail = f"Unknown provider type: {payload.provider_type}", ) - try: - api_key = decrypt_api_key(payload.encrypted_api_key) - except Exception as exc: - logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) - raise HTTPException( - status_code = 400, - detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", - ) + api_key = "" + if payload.encrypted_api_key: + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning( + "Failed to decrypt API key (%s): %s", type(exc).__name__, exc + ) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) base_url = payload.base_url or info["base_url"] client = ExternalProviderClient( @@ -265,14 +269,18 @@ async def list_provider_models( detail = f"Unknown provider type: {payload.provider_type}", ) - try: - api_key = decrypt_api_key(payload.encrypted_api_key) - except Exception as exc: - logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) - raise HTTPException( - status_code = 400, - detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", - ) + api_key = "" + if payload.encrypted_api_key: + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning( + "Failed to decrypt API key (%s): %s", type(exc).__name__, exc + ) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) if info.get("model_list_mode") == "curated": return [ diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 29c999678d..edbb55df19 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -26,6 +26,7 @@ import type { } from "../types/api"; import { getExternalProviderApiKey, + isCustomProviderType, loadExternalProviders, parseExternalModelId, supportsProviderPromptCaching, @@ -748,7 +749,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); throw new Error("External provider not found."); } - if (isExternalRequest && !externalApiKey) { + // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. + const externalProviderIsCustom = externalProvider + ? isCustomProviderType(externalProvider.providerType) + : false; + if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) { toast.error("Missing API key for selected external provider.", { description: "Open Connections and set the API key again.", }); @@ -1028,10 +1033,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { provider_id: externalProvider.id, provider_type: externalBackendProviderType, external_model: externalSelection.modelId, - encrypted_api_key: await encryptProviderApiKey( - externalApiKey, - forceRefreshPublicKey, - ), + ...(externalApiKey + ? { + encrypted_api_key: await encryptProviderApiKey( + externalApiKey, + forceRefreshPublicKey, + ), + } + : {}), provider_base_url: externalProvider.baseUrl || null, ...(supportsProviderPromptCaching(externalProvider.providerType) ? { diff --git a/studio/frontend/src/features/chat/api/providers-api.ts b/studio/frontend/src/features/chat/api/providers-api.ts index 4e1c126ff0..e76e24a627 100644 --- a/studio/frontend/src/features/chat/api/providers-api.ts +++ b/studio/frontend/src/features/chat/api/providers-api.ts @@ -176,8 +176,12 @@ export async function updateProviderConfig( async function withApiKeyEncryptionRetry<T>( plaintextApiKey: string, - call: (encryptedApiKey: string) => Promise<T>, + call: (encryptedApiKey: string | null) => Promise<T>, ): Promise<T> { + // Empty key (local providers): skip RSA round-trip and let the backend omit auth. + if (!plaintextApiKey) { + return await call(null); + } try { const encrypted = await encryptProviderApiKey(plaintextApiKey, false); return await call(encrypted); From 2622b79606bf0ed57844cede926562b827cd6fca Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Fri, 15 May 2026 23:39:06 +0400 Subject: [PATCH 27/50] studio/chat: built-in code execution for OpenAI + Anthropic (#5461) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio/chat: built-in code execution for Anthropic Claude 4.x Wire Anthropic's server-side code_execution_20250825 tool to the existing Code pill in the composer. Pill lights up only for Claude Opus/Sonnet/Haiku 4.x models that the docs list as compatible; pairs independently with Search. Backend appends the tool entry plus the code-execution-2025-08-25 beta header, and translates the SSE server_tool_use / *_tool_result blocks (bash + text_editor sub-tools) into the _toolEvent shape the frontend renderer consumes. File uploads via the Files API are a deliberate follow-up. * studio/chat: enable code execution pill in in-thread composer too thread.tsx renders its own composer with a separate CodeToolsToggle that was still gated on supportsTools only, so the pill stayed disabled inside an active thread even after picking Anthropic 4.x. Surface the capability through the runtime store (supportsBuiltinCodeExecution, set from chat-page alongside supportsBuiltinWebSearch) and read it in the toggle. * studio/chat: built-in code execution for OpenAI cloud gpt-5.5 Extend the Code pill to OpenAI cloud's gpt-5.5 / gpt-5.5-pro via the shell tool on /v1/responses. Per-thread container reuse: capture the container_id from each response on a synthetic container_ready event, persist it onto the ThreadRecord, and pass it back as environment.type="container_reference" on follow-up turns so the model sees filesystem state from prior turns until OpenAI's idle expiry. Stale ids surface a container_invalidated event that clears the thread record so the next turn falls back to container_auto. Gated strictly on OpenAI cloud (api.openai.com base URL) — Ollama, llama.cpp, vLLM, and custom OpenAI-compat presets won't see the shell tool entry even when their providerType collapses to "openai". * studio/chat: OpenAI shell-tool container management UI Side-panel section (settings sheet → Code Execution) for managing OpenAI's shell-tool containers per thread. Three controls: - New-container idle timeout (provider-level default, pre-fills the create dialog and is used by the lazy-create path on a thread's first turn when set to a non-default value). - Active container picker for the active thread — pick any existing container or stay on "Auto-create per thread". - Inline create form (name + idle TTL) and per-row delete actions. Three new backend endpoints under /api/inference/external/openai/ containers/{list,create,delete} proxy to OpenAI /v1/containers using the encrypted API key. All three reject non-cloud base URLs up front so the picker stays scoped to api.openai.com. Deleting a container clears all thread bindings pointing at it; the next turn falls back to auto-create. * studio/chat: inherit container across threads + styled active picker New threads on the same OpenAI provider now default to the most recently used container instead of "Auto-create per thread" — both in the chat-adapter (so a send works even if the side panel was never opened) and in the side panel itself (auto-binds the active thread when the dropdown loads on a thread that has no container). Picker is visually emphasized with an accent panel and the currently-active row in the list below is highlighted with the same accent so the two views stay in sync. * studio/chat: friendly English-word names for auto-created containers Replaces the "chat-<thread-id-slug>" auto-name with a random English-word + short hex suffix (e.g. "kestrel-3f9c"). Applies only to the chat-adapter's lazy-create path; the OpenAI container_auto path stays unnamed (only fires when no custom TTL is set). * studio/chat: always pre-create OpenAI containers via frontend Drops the TTL-based gate on the chat-adapter's lazy-create path so every code-execution container the user ever sees in the picker has a friendly English-word name. The backend's container_auto fallback stays as a safety net (used only if the POST /v1/containers call fails); in practice that branch should be rare. * studio/chat: send OpenAI-Beta header for /v1/containers CRUD Without OpenAI-Beta: containers=v1, OpenAI returns 200 {"deleted": true} for DELETE /v1/containers/{id} but does not actually remove the container. The list call then keeps returning it, making it look like Studio's "Delete container" button is broken. Verified 2026-05-15 against api.openai.com: DELETE with the beta header returns 200 and removes the container; the same DELETE without the header returns the same 200 deleted:true body but the container stays alive. - Add _container_headers() that merges OpenAI-Beta on top of the shared auth headers; route list / create / delete through it. - Verify the DELETE response body reports {"deleted": true}; raise httpx.HTTPError otherwise so the route surfaces a 5xx instead of silently reporting success on a silent no-op. - Add tests covering header propagation and the deleted-flag guard (true, false, missing key, non-JSON body, 4xx passthrough). * studio/chat: surface unpersisted-thread picker no-op as a toast The "Active for this thread" container picker uses db.threads.update(activeThreadId, ...), which silently returns 0 rows affected when the thread record isn't yet in IndexedDB. That happens on a brand-new thread where the user toggles code execution on and opens settings before sending the first message — the chat adapter only materializes the thread row on first send. The picker would appear to ignore the user's selection and snap back to "Auto-create per thread". - onPick now awaits the update and toasts an actionable hint ("Send a message first to pin a container to this thread.") when the update affected zero rows. - Auto-bind effect comment clarifies why it stays best-effort silent. The auto-bind effect itself is unchanged: it's a heuristic that should not nag the user when it can't apply. * studio/chat: let user pick OpenAI container before first send Previously the picker silently no-op'd until the user sent the first message, because Dexie's ThreadRecord is only materialized inside the runtime-provider's `initialize` hook (assistant-ui's first-message callback). That kept users from binding a thread to an existing OpenAI container up front; they had to either send a message and risk the chat adapter auto-creating one, or accept the cross-thread inheritance default. - Export `ensureThreadRecord` from runtime-provider so other surfaces can materialize the row idempotently. - In OpenAICodeExecSection.onPick, await ensureThreadRecord before the update, with modelType="base" (the settings sheet that hosts this section is only rendered in single-thread mode). Behaviour after this commit: - New thread + user picks a container in the sidebar → thread row is created with that container_id; first send uses it, no auto-create. - New thread + user does nothing → row still absent; first send goes through the existing inherit/lazy-create path as before. - The auto-bind effect remains silent best-effort: it does not eagerly create the thread row, so it cannot pre-empt the user's pick on a fresh thread. * studio/chat: drop "Auto-create per thread" option, default to latest The dropdown previously offered "Auto-create per thread" as an explicit value (null in storage), with the chat-adapter then inheriting from the most recent container at send-time. That made the picker display disagree with what the backend would actually do: the picker said "auto", but the backend was reusing an existing container. Behaviour after this commit, when code execution is enabled on an OpenAI cloud provider: - Containers list non-empty: dropdown defaults to the container with the latest lastActiveAt, eagerly bound via ensureThreadRecord + db.threads.update so the bind survives even when the thread row has not been materialized by the chat adapter yet. User can pick any other container in the list. - Containers list empty: render a disabled placeholder "(none yet — will be created on first send)". The chat-adapter's lazy-create path (chat-adapter.ts:1040-1082) mints the first container on first send and writes it back to the thread; the next refresh surfaces it in the picker. Expiration mid-operation is unchanged: the existing container_invalidated _toolEvent clears the thread's stored id and the next turn re-creates. * studio/chat: fix picker stuck on "Selecting most recent…" + manual-create binding Two follow-up fixes to the picker rework in d0cbeb99b. 1) The dropdown was getting stuck on the "Selecting most recent…" placeholder option even after the auto-bind write completed, because the select was controlled by `activeContainerId` (whatever sits in Dexie) and there's a brief window between the auto-bind firing and useLiveQuery propagating the new row back. Decoupled the rendered value from the Dexie state: compute the displayed id locally as `activeContainerId ?? sortedContainers[0]?.id`, so the most-recent container's name shows up immediately. The auto-bind effect still writes the bind to Dexie so the chat adapter sees it on send. Dropped the placeholder option entirely. 2) The manual "Create container" flow (`onCreate`) bound the new container to the active thread with a bare `db.threads.update`. On a brand-new thread that hadn't been materialized yet, the update affected 0 rows; the user's next send then went through cross-thread inheritance / lazy-create and could land on a stale container, surfacing as "container does not exist". Same fix as `onPick`: ensureThreadRecord before update so the bind lands. --- .../core/inference/external_provider.py | 594 +++++++++++++++++- studio/backend/models/inference.py | 75 +++ studio/backend/routes/inference.py | 158 +++++ .../tests/test_anthropic_code_execution.py | 419 ++++++++++++ .../tests/test_openai_code_execution.py | 391 ++++++++++++ .../tests/test_openai_container_crud.py | 151 +++++ .../src/components/assistant-ui/thread.tsx | 13 +- .../assistant-ui/tool-ui-code-execution.tsx | 126 ++++ .../src/features/chat/api/chat-adapter.ts | 178 +++++- .../features/chat/api/openai-containers.ts | 121 ++++ .../frontend/src/features/chat/chat-page.tsx | 36 +- .../src/features/chat/chat-settings-sheet.tsx | 30 + .../components/openai-code-exec-section.tsx | 477 ++++++++++++++ .../src/features/chat/external-providers.ts | 13 + .../src/features/chat/lib/friendly-names.ts | 244 +++++++ .../features/chat/provider-capabilities.ts | 82 +++ .../src/features/chat/runtime-provider.tsx | 2 +- .../src/features/chat/shared-composer.tsx | 23 +- .../chat/stores/chat-runtime-store.ts | 16 +- studio/frontend/src/features/chat/types.ts | 16 + .../frontend/src/features/chat/types/api.ts | 10 + 21 files changed, 3128 insertions(+), 47 deletions(-) create mode 100644 studio/backend/tests/test_anthropic_code_execution.py create mode 100644 studio/backend/tests/test_openai_code_execution.py create mode 100644 studio/backend/tests/test_openai_container_crud.py create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx create mode 100644 studio/frontend/src/features/chat/api/openai-containers.ts create mode 100644 studio/frontend/src/features/chat/components/openai-code-exec-section.tsx create mode 100644 studio/frontend/src/features/chat/lib/friendly-names.ts diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 07574abf4c..79c8287e5b 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -237,6 +237,7 @@ class ExternalProviderClient: reasoning_effort: Optional[str] = None, enabled_tools: Optional[list[str]] = None, enable_prompt_caching: Optional[bool] = None, + openai_code_exec_container_id: Optional[str] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: """ @@ -282,6 +283,7 @@ class ExternalProviderClient: reasoning_effort, enabled_tools, enable_prompt_caching, + openai_code_exec_container_id, ): yield line return @@ -1285,6 +1287,33 @@ class ExternalProviderClient: ) body["tools"] = anthropic_tools + # Anthropic server-side code execution — see + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool + # `code_execution_20250825` runs Python + bash + str_replace + # file edits inside a 5 GB sandboxed container per request, with + # no internet access. The tool entry itself takes no extra + # parameters; on the SSE stream Anthropic emits two sub-tool + # names — `bash_code_execution` and + # `text_editor_code_execution` — wrapped in the standard + # server_tool_use / *_tool_result block shape. The matching + # beta header (`code-execution-2025-08-25`) is set further down + # in this function alongside the request headers. + # v1 wires the tool only; file uploads (container_upload + # content blocks and generated-file retrieval via the Files + # API) are a deliberate follow-up. + code_execution_enabled = bool( + enabled_tools and "code_execution" in enabled_tools + ) + if code_execution_enabled: + anthropic_tools = list(body.get("tools") or []) + anthropic_tools.append( + { + "type": "code_execution_20250825", + "name": "code_execution", + } + ) + body["tools"] = anthropic_tools + url = f"{self.base_url}/messages" completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}" @@ -1314,12 +1343,29 @@ class ExternalProviderClient: logger.info("Proxying Anthropic Messages API to %s (model=%s)", url, model) + request_headers = self._auth_headers() + if code_execution_enabled: + # Anthropic accepts comma-separated beta features in a single + # `anthropic-beta` header. Merge our flag onto whatever the + # registry's extra_headers contributed (currently nothing on + # the beta axis, just anthropic-version) so future betas + # added at the registry level keep working. + existing_beta = request_headers.get("anthropic-beta", "").strip() + beta_parts = ( + [p.strip() for p in existing_beta.split(",") if p.strip()] + if existing_beta + else [] + ) + if "code-execution-2025-08-25" not in beta_parts: + beta_parts.append("code-execution-2025-08-25") + request_headers["anthropic-beta"] = ",".join(beta_parts) + try: async with _http_client.stream( "POST", url, json = body, - headers = self._auth_headers(), + headers = request_headers, timeout = self._stream_timeout, ) as response: if response.status_code != 200: @@ -1353,6 +1399,28 @@ class ExternalProviderClient: current_server_tool_use: Optional[dict[str, Any]] = None current_result_block: Optional[dict[str, Any]] = None web_search_calls: dict[str, dict[str, Any]] = {} + # code_execution state. Anthropic's + # `code_execution_20250825` tool emits the same + # server_tool_use → *_tool_result block shape as + # web_search, but the server_tool_use carries one of + # two sub-tool names (`bash_code_execution` or + # `text_editor_code_execution`) and the result block + # type matches (`bash_code_execution_tool_result` / + # `text_editor_code_execution_tool_result`). Kept + # parallel to web_search state so the two paths don't + # collide when both pills are on in the same turn. + current_code_exec_use: Optional[dict[str, Any]] = None + current_code_exec_result: Optional[dict[str, Any]] = None + code_execution_calls: dict[str, dict[str, Any]] = {} + # Counts surfaced in the final log line so reports of + # "Code execution did nothing" can be triaged at a + # glance. generated_files_count is interesting for the + # future Files API PR — when bash creates files inside + # the container, they show up as file_id entries on + # bash_code_execution_result.content, and v1 drops + # them. Track the count so we know how often it would + # have mattered. + code_execution_generated_files = 0 # Cache usage tracking. message_start carries the input # accounting (incl. cache_creation_input_tokens and # cache_read_input_tokens); message_delta carries cumulative @@ -1406,6 +1474,48 @@ class ExternalProviderClient: blocks.append(f"Title: {title}\nURL: {url}") return "\n---\n".join(blocks) + def _format_code_execution_result( + inner: dict[str, Any], + ) -> str: + """Render an Anthropic code-execution result block as + the preformatted text payload the frontend's + CodeExecutionToolUI displays inside a <pre>. Handles + bash, text_editor (view/create/str_replace), and the + matching error variants. + """ + inner_type = inner.get("type") or "" + if inner_type.endswith("_error"): + return f"Error: {inner.get('error_code', 'unknown')}" + if inner_type == "bash_code_execution_result": + stdout = inner.get("stdout") or "" + stderr = inner.get("stderr") or "" + return_code = inner.get("return_code") + parts: list[str] = [] + if stdout: + parts.append(stdout) + if stderr: + parts.append(f"--- stderr ---\n{stderr}") + if isinstance(return_code, int) and return_code != 0: + parts.append(f"return_code: {return_code}") + return "\n".join(parts) if parts else "(no output)" + if inner_type == "text_editor_code_execution_result": + # view: file content; create: is_file_update flag; + # str_replace: diff `lines` list. The matching + # server_tool_use carries the command + path, but + # that's encoded into the tool_start arguments + # already — here we only format the result body. + if "lines" in inner and isinstance(inner.get("lines"), list): + return "\n".join(str(line) for line in inner["lines"]) + if "is_file_update" in inner: + return ( + "Updated" if inner.get("is_file_update") else "Created" + ) + content_field = inner.get("content") + if isinstance(content_field, str): + return content_field + return "(file operation complete)" + return "(code execution complete)" + try: while True: try: @@ -1447,9 +1557,10 @@ class ExternalProviderClient: if event_type == "content_block_start": content_block = event.get("content_block") or {} block_type = content_block.get("type") + block_name = content_block.get("name") if ( block_type == "server_tool_use" - and content_block.get("name") == "web_search" + and block_name == "web_search" ): tool_use_id = content_block.get("id", "") or ( f"ws_{len(web_search_calls)}" @@ -1475,6 +1586,44 @@ class ExternalProviderClient: if isinstance(content, list) else [], } + elif block_type == "server_tool_use" and block_name in ( + "bash_code_execution", + "text_editor_code_execution", + ): + tool_use_id = content_block.get("id", "") or ( + f"ce_{len(code_execution_calls)}" + ) + kind = ( + "bash" + if block_name == "bash_code_execution" + else "text_editor" + ) + current_code_exec_use = { + "id": tool_use_id, + "kind": kind, + "buffer": "", + } + code_execution_calls[tool_use_id] = { + "kind": kind, + "arguments": {}, + "result": None, + } + elif block_type in ( + "bash_code_execution_tool_result", + "text_editor_code_execution_tool_result", + ): + # Anthropic ships the full result content + # on the start event for code-exec result + # blocks (unlike web_search, which can + # split across deltas). Capture it and + # finalize on content_block_stop so the + # ordering matches the web_search path. + tool_use_id = content_block.get("tool_use_id", "") + inner = content_block.get("content") or {} + current_code_exec_result = { + "tool_use_id": tool_use_id, + "inner": inner if isinstance(inner, dict) else {}, + } elif event_type == "content_block_delta": delta = event.get("delta", {}) @@ -1508,15 +1657,20 @@ class ExternalProviderClient: # per-call by Anthropic via the # `web_search_tool_result` block; we don't # need to scrape them off the text events. - elif ( - delta_type == "input_json_delta" - and current_server_tool_use is not None - ): - # Streamed partial_json carrying the search - # query. Buffer until content_block_stop. - current_server_tool_use["buffer"] += delta.get( - "partial_json", "" - ) + elif delta_type == "input_json_delta": + # Streamed partial_json carrying tool inputs + # — the search query for web_search, or the + # command/path/etc. for code execution. + # Route to whichever buffer is open. The two + # state slots are exclusive in practice + # (Anthropic doesn't interleave tool input + # streams), but checking both keeps the + # dispatch robust if that ever changes. + partial = delta.get("partial_json", "") + if current_server_tool_use is not None: + current_server_tool_use["buffer"] += partial + elif current_code_exec_use is not None: + current_code_exec_use["buffer"] += partial # signature_delta and any other delta types are # intentionally skipped — they carry trust / # verification metadata, not user-visible content. @@ -1572,6 +1726,68 @@ class ExternalProviderClient: } ) current_result_block = None + elif current_code_exec_use is not None: + # End of a code-execution server_tool_use — + # parse the buffered input_json into a + # {command, path, ...} dict and emit + # tool_start. The matching tool_end fires + # on the result block's content_block_stop. + buffer = current_code_exec_use["buffer"] + parsed_args: dict[str, Any] = {} + if buffer: + try: + parsed_obj = _json.loads(buffer) + if isinstance(parsed_obj, dict): + parsed_args = parsed_obj + except Exception: + parsed_args = {} + tool_use_id = current_code_exec_use["id"] + kind = current_code_exec_use["kind"] + emit_args = {"kind": kind, **parsed_args} + if tool_use_id in code_execution_calls: + code_execution_calls[tool_use_id]["arguments"] = ( + emit_args + ) + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": tool_use_id, + "arguments": emit_args, + } + ) + current_code_exec_use = None + elif current_code_exec_result is not None: + # End of a code-execution result block — + # format the inner result into the text + # payload CodeExecutionToolUI renders. + tool_use_id = current_code_exec_result["tool_use_id"] + inner = current_code_exec_result["inner"] + # Track generated-file count for the + # follow-up Files API PR. v1 drops them. + if isinstance(inner, dict): + file_blocks = inner.get("content") + if isinstance(file_blocks, list): + for entry in file_blocks: + if isinstance(entry, dict) and entry.get( + "file_id" + ): + code_execution_generated_files += 1 + result_text = _format_code_execution_result( + inner if isinstance(inner, dict) else {} + ) + if tool_use_id in code_execution_calls: + code_execution_calls[tool_use_id]["result"] = ( + result_text + ) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": tool_use_id, + "result": result_text, + } + ) + current_code_exec_result = None elif thinking_open: # Close the <think> tag when the thinking block # ends, in case no text_delta follows (e.g. @@ -1639,10 +1855,20 @@ class ExternalProviderClient: # instead. cache_creation tokens are billed at a # small premium; cache_read tokens are billed at a # discount. + code_execution_invocations = len(code_execution_calls) + code_execution_results = sum( + 1 + for c in code_execution_calls.values() + if c.get("result") is not None + ) logger.info( "Anthropic stream complete (model=%s, " "web_search_requested=%s, web_search_invocations=%s, " "results=%s, queries=%s, " + "code_execution_requested=%s, " + "code_execution_invocations=%s, " + "code_execution_results=%s, " + "code_execution_generated_files=%s, " "input_tokens=%s, output_tokens=%s, " "cache_creation_input_tokens=%s, " "cache_read_input_tokens=%s, events=%s)", @@ -1651,6 +1877,10 @@ class ExternalProviderClient: web_search_invocations, total_results, queries, + code_execution_enabled, + code_execution_invocations, + code_execution_results, + code_execution_generated_files, last_usage.get("input_tokens"), last_usage.get("output_tokens"), last_usage.get("cache_creation_input_tokens"), @@ -1693,6 +1923,7 @@ class ExternalProviderClient: reasoning_effort: Optional[str], enabled_tools: Optional[list[str]] = None, enable_prompt_caching: Optional[bool] = None, + openai_code_exec_container_id: Optional[str] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -1817,15 +2048,41 @@ class ExternalProviderClient: # OpenAI server-side tools — see # https://developers.openai.com/api/docs/guides/tools - # The frontend's Search button maps to the unified - # enabled_tools=["web_search"] shorthand; translate that into the - # Responses-API tool schema. Other built-in tools (file_search, - # code_interpreter, image_generation, computer_use_preview) can be - # added with the same pattern when we surface their toggles. + # https://developers.openai.com/api/docs/guides/tools-shell + # The frontend's Search/Code buttons map to the unified + # enabled_tools shorthand; translate that into the Responses-API + # tool schema. Other built-in tools (file_search, + # code_interpreter, image_generation, computer_use_preview) can + # be added with the same pattern when we surface their toggles. + code_execution_enabled_openai = bool( + enabled_tools and "code_execution" in enabled_tools and is_openai_cloud + ) if enabled_tools: tools_array: list[dict[str, Any]] = [] if "web_search" in enabled_tools: tools_array.append({"type": "web_search"}) + if code_execution_enabled_openai: + # `container_auto` lets OpenAI auto-create a fresh + # container per request; we capture the resulting + # container_id off the SSE stream and the chat-adapter + # persists it onto the thread record. Subsequent turns + # in the same thread pass it back as + # `openai_code_exec_container_id`, which we translate to + # `container_reference` here so the model sees + # filesystem state from prior turns. Container expires + # after ~20 min of inactivity per OpenAI's default + # policy — a stale id 400s, the chat-adapter clears it + # via container_invalidated, and the next turn falls + # back to auto-create. + shell_env: dict[str, Any] + if openai_code_exec_container_id: + shell_env = { + "type": "container_reference", + "container_id": openai_code_exec_container_id, + } + else: + shell_env = {"type": "container_auto"} + tools_array.append({"type": "shell", "environment": shell_env}) if tools_array: body["tools"] = tools_array @@ -1850,6 +2107,29 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) + # Detect stale-container errors so the frontend can + # drop its persisted id. OpenAI doesn't pin an + # error code in the public docs for this case, so + # match a couple of likely substrings. If we sent + # a container_reference and the response is 4xx + # with any hint of "container not found / expired", + # emit container_invalidated; the next turn will + # fall back to container_auto. + if ( + openai_code_exec_container_id + and 400 <= response.status_code < 500 + ): + lowered = error_text.lower() + if "container" in lowered and ( + "expired" in lowered + or "not_found" in lowered + or "not found" in lowered + or "no such container" in lowered + ): + yield ( + f"data: " + f"{_json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'choices': [{'index': 0, 'delta': {}, 'finish_reason': None}], '_toolEvent': {'type': 'container_invalidated'}})}" + ) yield _error_sse_line( response.status_code, error_text, self.provider_type ) @@ -1887,6 +2167,28 @@ class ExternalProviderClient: # web_search_calls: { item_id -> {query} } web_search_calls: dict[str, dict[str, Any]] = {} all_url_citations: list[dict[str, str]] = [] + # Shell-tool (code execution) state. OpenAI emits + # `shell_call` items (model requesting a command list) + # paired with `shell_call_output` items (execution + # results). We mirror the Anthropic code-execution UX + # by emitting one `_toolEvent` tool_start per + # shell_call and one tool_end per shell_call_output; + # they're linked via `shell_call_output.call_id` + # matching `shell_call.id`. Items are independent of + # web_search (different keyed map). + # shell_calls: { call_id -> {commands, output} } + shell_calls: dict[str, dict[str, Any]] = {} + # Container id captured from the response stream. When + # it differs from the inbound id, emit a synthetic + # `container_ready` _toolEvent so the frontend can + # persist it onto the thread record for the next turn. + # Where OpenAI surfaces it is documented loosely; we + # probe two known fields (response.container_id on + # response.completed, item.environment.container_id on + # shell_call output items) and latch the first one we + # see. + latched_container_id: Optional[str] = None + container_id_emitted = False def _emit_tool_event(payload: dict[str, Any]) -> str: chunk = { @@ -1903,6 +2205,45 @@ class ExternalProviderClient: } return f"data: {_json.dumps(chunk)}" + def _format_shell_output(output: Any) -> str: + """Render an OpenAI `shell_call_output.output` list + as the preformatted text payload the frontend's + CodeExecutionToolUI displays inside a <pre>. Each + entry has stdout/stderr/outcome — concatenate them + with a separator block per entry and append + `return_code` / `(timeout)` annotations only when + they convey information beyond "succeeded". + """ + if not isinstance(output, list): + return "" + parts: list[str] = [] + for entry in output: + if not isinstance(entry, dict): + continue + stdout = entry.get("stdout") or "" + stderr = entry.get("stderr") or "" + outcome = entry.get("outcome") or {} + chunk_parts: list[str] = [] + if stdout: + chunk_parts.append(stdout) + if stderr: + chunk_parts.append(f"--- stderr ---\n{stderr}") + if isinstance(outcome, dict): + outcome_type = outcome.get("type") + if outcome_type == "exit": + exit_code = outcome.get("exit_code") + if isinstance(exit_code, int) and exit_code != 0: + chunk_parts.append(f"return_code: {exit_code}") + elif outcome_type == "timeout": + chunk_parts.append("(timeout)") + if chunk_parts: + parts.append("\n".join(chunk_parts)) + return ( + "\n--- next command ---\n".join(parts) + if parts + else "(no output)" + ) + def _record_url_citation(payload: dict[str, Any]) -> None: """Append a url_citation onto the shared all_url_citations list. Dedup by URL — the same source can be cited multiple @@ -2025,6 +2366,37 @@ class ExternalProviderClient: f"ws_{len(web_search_calls)}" ) web_search_calls.setdefault(item_id, {"query": ""}) + # Shell-tool: register the call eagerly so + # the matching shell_call_output can link + # back even if `done` arrives out of order. + # Also probe for container_id on the + # environment field — when container_auto + # auto-creates one, this is the first place + # the new id might surface (OpenAI doesn't + # promise this in docs, but the field is + # cheap to scan and lets us emit + # container_ready earlier than + # response.completed). + if ( + isinstance(item, dict) + and item.get("type") == "shell_call" + ): + item_id = item.get("id", "") or ( + f"sc_{len(shell_calls)}" + ) + shell_calls.setdefault( + item_id, + {"commands": [], "output": None}, + ) + env = item.get("environment") + if isinstance(env, dict): + probe = env.get("container_id") or env.get("id") + if ( + isinstance(probe, str) + and probe.startswith("cntr_") + and latched_container_id is None + ): + latched_container_id = probe elif event_type == "response.output_item.done": item = event.get("item", {}) @@ -2080,6 +2452,65 @@ class ExternalProviderClient: "result": "", } ) + elif item.get("type") == "shell_call": + # OpenAI ships the commands array on the + # action field. Join them onto one + # command string for the tool card — + # the renderer is shared with Anthropic + # bash, which only carries a single + # `command`. Multiple commands in one + # shell_call get joined with newlines so + # they still render as one card. + item_id = item.get("id", "") or ( + f"sc_{len(shell_calls)}" + ) + action = item.get("action") or {} + commands = ( + action.get("commands") + if isinstance(action, dict) + else None + ) or [] + joined_command = ( + "\n".join(str(c) for c in commands) + if isinstance(commands, list) + else "" + ) + shell_calls.setdefault( + item_id, + {"commands": [], "output": None}, + ) + shell_calls[item_id]["commands"] = ( + list(commands) if isinstance(commands, list) else [] + ) + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": item_id, + "arguments": { + "kind": "bash", + "command": joined_command, + }, + } + ) + elif item.get("type") == "shell_call_output": + # `call_id` links back to the shell_call's + # `id`, which is what we used as the + # tool_call_id on tool_start. Match on + # call_id when present so the matching + # card transitions to complete. + call_id = item.get("call_id") or item.get("id") or "" + output = item.get("output") or [] + if call_id in shell_calls: + shell_calls[call_id]["output"] = output + result_text = _format_shell_output(output) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": call_id, + "result": result_text, + } + ) elif isinstance(event_type, str) and "reasoning" in event_type: reasoning_delta = _extract_reasoning_text(event) @@ -2097,6 +2528,39 @@ class ExternalProviderClient: if reasoning_open: yield _chunk_with_text("</think>") reasoning_open = False + # Probe response.container_id (top-level) and + # response.container.id for the shell-tool + # container id. OpenAI's docs don't pin the + # exact field, so we scan both. Emit + # `container_ready` only when the value + # differs from the inbound one — no churn on + # reuse. + response_obj = event.get("response") or {} + if isinstance(response_obj, dict): + probe_id = response_obj.get("container_id") + if not probe_id: + container_field = response_obj.get("container") + if isinstance(container_field, dict): + probe_id = container_field.get("id") + if ( + isinstance(probe_id, str) + and probe_id.startswith("cntr_") + and latched_container_id is None + ): + latched_container_id = probe_id + if ( + latched_container_id + and not container_id_emitted + and latched_container_id + != openai_code_exec_container_id + ): + yield _emit_tool_event( + { + "type": "container_ready", + "container_id": latched_container_id, + } + ) + container_id_emitted = True # Apply the aggregated citation list onto the # *last* web_search call by overwriting its # tool_end result. The frontend's @@ -2227,10 +2691,19 @@ class ExternalProviderClient: details = last_usage.get("input_tokens_details") if isinstance(details, dict): cached_input_tokens = details.get("cached_tokens") + code_execution_requested = code_execution_enabled_openai + code_execution_invocations = len(shell_calls) + code_execution_results = sum( + 1 for sc in shell_calls.values() if sc.get("output") is not None + ) logger.info( "OpenAI Responses stream complete (model=%s, " "web_search_requested=%s, web_search_invocations=%s, " "citations=%s, queries=%s, reasoning_emitted=%s, " + "code_execution_requested=%s, " + "code_execution_invocations=%s, " + "code_execution_results=%s, " + "container_id_in=%s, container_id_out=%s, " "input_tokens=%s, output_tokens=%s, " "cached_input_tokens=%s)", model, @@ -2239,6 +2712,11 @@ class ExternalProviderClient: total_citations, queries, reasoning_emitted, + code_execution_requested, + code_execution_invocations, + code_execution_results, + openai_code_exec_container_id, + latched_container_id, (last_usage or {}).get("input_tokens"), (last_usage or {}).get("output_tokens"), cached_input_tokens, @@ -2359,6 +2837,90 @@ class ExternalProviderClient: ) raise + def _container_headers(self) -> dict[str, str]: + """Auth headers plus the OpenAI-Beta opt-in for /v1/containers. + + OpenAI's containers API requires ``OpenAI-Beta: containers=v1``. + Without it, DELETE silently no-ops: the API returns 200 with a + ``{"deleted": true}`` body but does not actually remove the + container (verified 2026-05-15). The header is required for + list / create / delete to behave consistently. + """ + headers = self._auth_headers() + headers["OpenAI-Beta"] = "containers=v1" + return headers + + async def list_openai_containers(self) -> list[dict[str, Any]]: + """ + GET /v1/containers on the user's OpenAI account. + + Returns the raw container records (id, name, created_at, + last_active_at, expires_after, status). The route layer + reshapes these into the UI summary shape. + + Only valid against api.openai.com — non-cloud OpenAI-compat + servers don't implement /v1/containers and would 404 here. + Caller is responsible for the is_openai_cloud guard. + """ + response = await _http_client.get( + f"{self.base_url}/containers", + headers = self._container_headers(), + timeout = self._timeout, + ) + response.raise_for_status() + data = response.json() + containers = data.get("data") if isinstance(data, dict) else None + return list(containers) if isinstance(containers, list) else [] + + async def create_openai_container( + self, + name: str, + ttl_minutes: int, + ) -> dict[str, Any]: + """ + POST /v1/containers with ``expires_after.anchor="last_active_at"``. + ``ttl_minutes`` is the idle timeout — every API call that + touches the container resets the timer. + """ + body = { + "name": name, + "expires_after": { + "anchor": "last_active_at", + "minutes": ttl_minutes, + }, + } + response = await _http_client.post( + f"{self.base_url}/containers", + json = body, + headers = self._container_headers(), + timeout = self._timeout, + ) + response.raise_for_status() + return response.json() + + async def delete_openai_container(self, container_id: str) -> None: + """DELETE /v1/containers/{id}. 404s are surfaced as HTTPError. + + Verifies the response body reports ``deleted: true``. OpenAI + returns a 2xx ``deleted: true`` body even when the request is + silently rejected (e.g. missing OpenAI-Beta header), so a + status-only check is not sufficient. + """ + response = await _http_client.delete( + f"{self.base_url}/containers/{container_id}", + headers = self._container_headers(), + timeout = self._timeout, + ) + response.raise_for_status() + try: + payload = response.json() + except ValueError: + payload = None + if not (isinstance(payload, dict) and payload.get("deleted") is True): + raise httpx.HTTPError( + f"OpenAI did not confirm container deletion: {response.text[:200]}" + ) + async def close(self) -> None: """No-op — the underlying client is shared across requests.""" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f2eed314ee..6a042d35d7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -604,6 +604,81 @@ class ChatCompletionRequest(BaseModel): "vllm, local, etc.). Treated as enabled when omitted." ), ) + openai_code_exec_container_id: Optional[str] = Field( + None, + description = ( + "[x-unsloth] OpenAI shell-tool container id from the prior response " + "in the same chat thread. When set and `code_execution` is in " + "`enabled_tools`, the next /v1/responses call uses " + "environment.type='container_reference' so filesystem state " + "persists across turns. Unset → environment.type='container_auto' " + "and OpenAI creates a fresh container. Only meaningful for the " + "OpenAI cloud + gpt-5.5 family path; ignored otherwise." + ), + ) + + +# ── OpenAI shell-tool container management ───────────────────── + + +class OpenAIContainerRequest(BaseModel): + """ + Shared body for the three OpenAI container endpoints (list / create + / delete). Carries the encrypted API key + base URL so the route + handler can decrypt it and proxy to the user's OpenAI account. + Same pattern as the inference proxy endpoints — keeps the key off + persistent storage on the backend. + """ + + encrypted_api_key: str = Field( + ..., + description = "[x-unsloth] RSA-encrypted, base64-encoded OpenAI API key.", + ) + provider_base_url: Optional[str] = Field( + None, + description = "[x-unsloth] OpenAI base URL. Only api.openai.com is supported; non-cloud bases are rejected with 400.", + ) + + +class CreateOpenAIContainerBody(OpenAIContainerRequest): + name: str = Field( + ..., + min_length = 1, + max_length = 256, + description = "Human-readable container name. Surfaces in the picker UI.", + ) + ttl_minutes: int = Field( + 20, + ge = 1, + le = 10080, # 1 week + description = ( + "Idle-timeout TTL the new container will inherit (anchor=" + "last_active_at). OpenAI's default is 20; we cap at one " + "week as a safety bound." + ), + ) + + +class DeleteOpenAIContainerBody(OpenAIContainerRequest): + container_id: str = Field( + ..., + description = "OpenAI container id (cntr_...) to delete.", + ) + + +class OpenAIContainerSummary(BaseModel): + """One row from GET /v1/containers, reshaped for the UI.""" + + id: str + name: Optional[str] = None + created_at: Optional[int] = None + last_active_at: Optional[int] = None + expires_after_minutes: Optional[int] = None + status: Optional[str] = None + + +class ListOpenAIContainersResponse(BaseModel): + containers: list[OpenAIContainerSummary] # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e8732397b0..8f5bd9aab5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -194,6 +194,11 @@ from models.inference import ( AnthropicResponseTextBlock, AnthropicResponseToolUseBlock, AnthropicUsage, + CreateOpenAIContainerBody, + DeleteOpenAIContainerBody, + ListOpenAIContainersResponse, + OpenAIContainerRequest, + OpenAIContainerSummary, ) from core.inference.anthropic_compat import ( anthropic_messages_to_openai, @@ -1598,6 +1603,7 @@ async def _proxy_to_external_provider( reasoning_effort = payload.reasoning_effort, enabled_tools = payload.enabled_tools, enable_prompt_caching = payload.enable_prompt_caching, + openai_code_exec_container_id = payload.openai_code_exec_container_id, stream = payload.stream, ) try: @@ -1627,6 +1633,158 @@ async def _proxy_to_external_provider( ) +# ── OpenAI shell-tool container management ─────────────────────── + + +def _resolve_openai_cloud_client( + body: OpenAIContainerRequest, +) -> ExternalProviderClient: + """ + Decrypt the API key + validate the base URL points at OpenAI cloud, + then build an ExternalProviderClient for the three container CRUD + endpoints below. The shell tool only exists on api.openai.com, so + rejecting non-cloud bases up front prevents confusing 404s on + ollama / llama.cpp / vLLM / custom presets. + """ + base_url = body.provider_base_url or get_base_url("openai") + if not base_url or "api.openai.com" not in base_url: + raise HTTPException( + status_code = 400, + detail = ( + "OpenAI container management is only available on the " + "managed cloud (api.openai.com). The provider's base URL " + f"points at {base_url!r}." + ), + ) + try: + api_key = decrypt_api_key(body.encrypted_api_key) + except Exception as exc: + logger.warning("external_provider.decrypt_failed", error = str(exc)) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.", + ) + return ExternalProviderClient( + provider_type = "openai", + base_url = base_url, + api_key = api_key, + ) + + +def _summarize_container(raw: dict) -> OpenAIContainerSummary: + expires = raw.get("expires_after") + expires_minutes: Optional[int] = None + if isinstance(expires, dict): + minutes = expires.get("minutes") + if isinstance(minutes, int): + expires_minutes = minutes + return OpenAIContainerSummary( + id = str(raw.get("id") or ""), + name = raw.get("name"), + created_at = raw.get("created_at") + if isinstance(raw.get("created_at"), int) + else None, + last_active_at = raw.get("last_active_at") + if isinstance(raw.get("last_active_at"), int) + else None, + expires_after_minutes = expires_minutes, + status = raw.get("status") if isinstance(raw.get("status"), str) else None, + ) + + +@router.post( + "/external/openai/containers/list", + response_model = ListOpenAIContainersResponse, +) +async def list_openai_containers( + body: OpenAIContainerRequest, + current_subject: str = Depends(get_current_subject), +) -> ListOpenAIContainersResponse: + """List the user's OpenAI shell-tool containers.""" + client = _resolve_openai_cloud_client(body) + try: + try: + raw = await client.list_openai_containers() + except httpx.HTTPStatusError as exc: + detail = exc.response.text[:500] if exc.response is not None else str(exc) + raise HTTPException( + status_code = exc.response.status_code if exc.response else 502, + detail = f"OpenAI rejected /containers list: {detail}", + ) + except httpx.HTTPError as exc: + raise HTTPException( + status_code = 502, + detail = f"Failed to reach OpenAI: {exc}", + ) + return ListOpenAIContainersResponse( + containers = [_summarize_container(c) for c in raw if isinstance(c, dict)], + ) + finally: + await client.close() + + +@router.post( + "/external/openai/containers/create", + response_model = OpenAIContainerSummary, +) +async def create_openai_container( + body: CreateOpenAIContainerBody, + current_subject: str = Depends(get_current_subject), +) -> OpenAIContainerSummary: + """Create a named container with the user-chosen idle TTL.""" + client = _resolve_openai_cloud_client(body) + try: + try: + raw = await client.create_openai_container( + name = body.name, + ttl_minutes = body.ttl_minutes, + ) + except httpx.HTTPStatusError as exc: + detail = exc.response.text[:500] if exc.response is not None else str(exc) + raise HTTPException( + status_code = exc.response.status_code if exc.response else 502, + detail = f"OpenAI rejected /containers create: {detail}", + ) + except httpx.HTTPError as exc: + raise HTTPException( + status_code = 502, + detail = f"Failed to reach OpenAI: {exc}", + ) + if not isinstance(raw, dict): + raise HTTPException( + status_code = 502, + detail = "OpenAI returned an unexpected container payload.", + ) + return _summarize_container(raw) + finally: + await client.close() + + +@router.post("/external/openai/containers/delete", status_code = 204) +async def delete_openai_container( + body: DeleteOpenAIContainerBody, + current_subject: str = Depends(get_current_subject), +) -> None: + """Delete a named container by id.""" + client = _resolve_openai_cloud_client(body) + try: + try: + await client.delete_openai_container(body.container_id) + except httpx.HTTPStatusError as exc: + detail = exc.response.text[:500] if exc.response is not None else str(exc) + raise HTTPException( + status_code = exc.response.status_code if exc.response else 502, + detail = f"OpenAI rejected /containers delete: {detail}", + ) + except httpx.HTTPError as exc: + raise HTTPException( + status_code = 502, + detail = f"Failed to reach OpenAI: {exc}", + ) + finally: + await client.close() + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py new file mode 100644 index 0000000000..b427ad2c0b --- /dev/null +++ b/studio/backend/tests/test_anthropic_code_execution.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for Anthropic's server-side `code_execution_20250825` tool +translation in `_stream_anthropic`. + +Covers: +- Request body: when ``enabled_tools=["code_execution"]``, the outbound + ``tools`` array carries ``{"type": "code_execution_20250825", "name": + "code_execution"}`` and the ``anthropic-beta`` header includes + ``code-execution-2025-08-25``. +- Combined request: ``enabled_tools=["web_search", "code_execution"]`` + sends both tool entries; the beta header still merges the code-exec + flag onto whatever the registry contributed. +- SSE translation: a `bash_code_execution` server_tool_use + + `bash_code_execution_tool_result` pair emits one tool_start and one + tool_end ``_toolEvent`` chunk with the expected arguments and result. +- SSE translation: a `text_editor_code_execution` create + result emits + a tool_start with ``kind="text_editor"`` + parsed args, and tool_end + with ``"Created"`` (or ``"Updated"``) based on the ``is_file_update`` + flag. +- Error path: a ``bash_code_execution_tool_result_error`` with + ``error_code="container_expired"`` renders as ``"Error: + container_expired"`` in the tool_end ``result``. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _anthropic_sse(events: list[dict]) -> bytes: + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _tool_events(lines: list[str]) -> list[dict]: + """Extract `_toolEvent` payloads from emitted SSE data lines.""" + out: list[dict] = [] + for line in lines: + if not line.startswith("data:"): + continue + raw = line[len("data:") :].strip() + if not raw or raw == "[DONE]": + continue + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict) and "_toolEvent" in parsed: + out.append(parsed["_toolEvent"]) + return out + + +def test_code_execution_tool_appended_to_request_body(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "compute 2 + 2"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enabled_tools = ["code_execution"], + ): + pass + await client.close() + + _drive(run()) + + body = captured["body"] + tools = body.get("tools") or [] + assert { + "type": "code_execution_20250825", + "name": "code_execution", + } in tools + # No web_search entry when only code_execution is enabled. + assert all(t.get("type") != "web_search_20250305" for t in tools) + # Beta header carries the documented flag. + beta_header = captured["headers"].get("anthropic-beta", "") + assert "code-execution-2025-08-25" in beta_header + + +def test_code_execution_with_web_search_sends_both_tools(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "look it up and chart it"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enabled_tools = ["web_search", "code_execution"], + ): + pass + await client.close() + + _drive(run()) + + tools = captured["body"].get("tools") or [] + tool_types = {t.get("type") for t in tools if isinstance(t, dict)} + assert "web_search_20250305" in tool_types + assert "code_execution_20250825" in tool_types + assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "") + + +def test_no_code_execution_tool_when_pill_off(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + ): + pass + await client.close() + + _drive(run()) + + tools = captured["body"].get("tools") or [] + assert all(t.get("type") != "code_execution_20250825" for t in tools) + # Beta header must NOT mention code-execution when the tool isn't on + # — that flag is opt-in only. + assert "code-execution-2025-08-25" not in captured["headers"].get( + "anthropic-beta", "" + ) + + +def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): + sse_events = [ + {"type": "message_start", "message": {"usage": {}}}, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "bash_code_execution", + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": '{"command": "ls -la"}', + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "bash_code_execution_result", + "stdout": "total 24\ndrwxr-xr-x 2 user user 4096 Jan 1 12:00 .", + "stderr": "", + "return_code": 0, + }, + }, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "message_stop"}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _anthropic_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_anthropic( + messages = [{"role": "user", "content": "list files"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enabled_tools = ["code_execution"], + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + + assert len(events) == 2 + start, end = events + assert start["type"] == "tool_start" + assert start["tool_name"] == "code_execution" + assert start["tool_call_id"] == "srvtoolu_1" + assert start["arguments"] == {"kind": "bash", "command": "ls -la"} + + assert end["type"] == "tool_end" + assert end["tool_call_id"] == "srvtoolu_1" + assert "total 24" in end["result"] + # Non-zero return_code not present, so no return_code line. + assert "return_code:" not in end["result"] + + +def test_text_editor_create_emits_kind_and_status(monkeypatch): + sse_events = [ + {"type": "message_start", "message": {"usage": {}}}, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": "srvtoolu_2", + "name": "text_editor_code_execution", + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": ( + '{"command": "create", "path": "new_file.txt", ' + '"file_text": "hi"}' + ), + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "text_editor_code_execution_tool_result", + "tool_use_id": "srvtoolu_2", + "content": { + "type": "text_editor_code_execution_result", + "is_file_update": False, + }, + }, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "message_stop"}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _anthropic_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_anthropic( + messages = [{"role": "user", "content": "write a file"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enabled_tools = ["code_execution"], + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + + assert len(events) == 2 + start, end = events + assert start["arguments"]["kind"] == "text_editor" + assert start["arguments"]["command"] == "create" + assert start["arguments"]["path"] == "new_file.txt" + assert end["result"] == "Created" + + +def test_code_execution_error_renders_error_code(monkeypatch): + sse_events = [ + {"type": "message_start", "message": {"usage": {}}}, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": "srvtoolu_3", + "name": "bash_code_execution", + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": '{"command": "echo broken"}', + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_3", + "content": { + "type": "bash_code_execution_tool_result_error", + "error_code": "container_expired", + }, + }, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "message_stop"}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _anthropic_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_anthropic( + messages = [{"role": "user", "content": "run it"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enabled_tools = ["code_execution"], + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + + assert len(events) == 2 + end = events[1] + assert end["type"] == "tool_end" + assert end["result"] == "Error: container_expired" diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py new file mode 100644 index 0000000000..88ff1171ef --- /dev/null +++ b/studio/backend/tests/test_openai_code_execution.py @@ -0,0 +1,391 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for OpenAI's server-side `shell` tool translation in +`_stream_openai_responses`. + +Covers: +- Request body: ``enabled_tools=["code_execution"]`` on the OpenAI + cloud base_url appends ``{"type": "shell", "environment": {"type": + "container_auto"}}`` to ``tools``. +- Container reuse: when ``openai_code_exec_container_id`` is provided, + the outgoing ``environment.type`` flips to ``"container_reference"`` + and the id propagates. +- Cloud guard: code_execution on a non-cloud base_url (e.g. a local + OpenAI-compat preset / ollama / llama.cpp / vLLM) does NOT add the + shell tool, preventing a guaranteed 400 from those servers. +- SSE translation: a `shell_call` + `shell_call_output` pair emits one + ``_toolEvent`` `tool_start` (`tool_name="code_execution"`, + `arguments.kind="bash"`) and one `tool_end` whose `result` contains + the joined stdout from the shell_call_output entries. +- Container surfacing: container_id captured from + `response.completed.container_id` is emitted as a synthetic + `container_ready` `_toolEvent` (only when it differs from the + inbound id). +- Stale-container handling: 400 with "container expired" body emits a + `container_invalidated` event before propagating the error. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client(base_url: str = "https://api.openai.com/v1") -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "openai", + base_url = base_url, + api_key = "sk-test", + ) + + +def _openai_sse(events: list[dict]) -> bytes: + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _tool_events(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for line in lines: + if not line.startswith("data:"): + continue + raw = line[len("data:") :].strip() + if not raw or raw == "[DONE]": + continue + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict) and "_toolEvent" in parsed: + out.append(parsed["_toolEvent"]) + return out + + +def test_shell_tool_added_on_cloud_with_container_auto(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _openai_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "compute 2+2"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + ): + pass + await client.close() + + _drive(run()) + + tools = captured["body"].get("tools") or [] + assert { + "type": "shell", + "environment": {"type": "container_auto"}, + } in tools + + +def test_shell_tool_uses_container_reference_when_id_supplied(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _openai_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "what did i write earlier"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + openai_code_exec_container_id = "cntr_abc123", + ): + pass + await client.close() + + _drive(run()) + + tools = captured["body"].get("tools") or [] + assert { + "type": "shell", + "environment": { + "type": "container_reference", + "container_id": "cntr_abc123", + }, + } in tools + + +def test_shell_tool_refused_for_non_cloud_base_url(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _openai_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client(base_url = "http://localhost:11434/v1") + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + ): + pass + await client.close() + + _drive(run()) + + tools = captured["body"].get("tools") or [] + # Shell tool must NOT leak to local OpenAI-compat servers — those + # 400 on the unknown tool type. + assert all(t.get("type") != "shell" for t in tools) + + +def test_shell_call_emits_tool_start_and_end(monkeypatch): + sse_events = [ + { + "type": "response.output_item.added", + "item": { + "type": "shell_call", + "id": "scall_1", + "action": {"commands": ["ls -la"]}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_1", + "action": {"commands": ["ls -la"]}, + "status": "completed", + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call_output", + "id": "scout_1", + "call_id": "scall_1", + "output": [ + { + "stdout": "total 24\ndrwxr-xr-x .", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + {"type": "response.completed", "response": {}}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _openai_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "list files"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + starts = [e for e in events if e["type"] == "tool_start"] + ends = [e for e in events if e["type"] == "tool_end"] + assert len(starts) == 1 + assert len(ends) == 1 + assert starts[0]["tool_name"] == "code_execution" + assert starts[0]["tool_call_id"] == "scall_1" + assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"} + assert ends[0]["tool_call_id"] == "scall_1" + assert "total 24" in ends[0]["result"] + + +def test_container_ready_emitted_when_new_id_surfaces(monkeypatch): + sse_events = [ + { + "type": "response.completed", + "response": {"container_id": "cntr_new_456"}, + }, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _openai_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "do stuff"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + ready = [e for e in events if e["type"] == "container_ready"] + assert len(ready) == 1 + assert ready[0]["container_id"] == "cntr_new_456" + + +def test_container_ready_not_emitted_when_id_unchanged(monkeypatch): + sse_events = [ + { + "type": "response.completed", + "response": {"container_id": "cntr_same_789"}, + }, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _openai_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "do stuff"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + openai_code_exec_container_id = "cntr_same_789", + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + # No churn — id matches the one already on the thread record. + assert not any(e["type"] == "container_ready" for e in events) + + +def test_stale_container_emits_invalidated(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + content = json.dumps( + { + "error": { + "message": "container has expired", + "type": "invalid_request_error", + } + } + ).encode("utf-8"), + headers = {"content-type": "application/json"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = ["code_execution"], + openai_code_exec_container_id = "cntr_stale_999", + ) + ) + + lines = _drive(run()) + events = _tool_events(lines) + invalidated = [e for e in events if e["type"] == "container_invalidated"] + assert len(invalidated) == 1 diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py new file mode 100644 index 0000000000..2965ec6649 --- /dev/null +++ b/studio/backend/tests/test_openai_container_crud.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the /v1/containers CRUD client methods. + +Covers: +- All three calls (list / create / delete) send + ``OpenAI-Beta: containers=v1``. Without it, OpenAI silently no-ops + the DELETE while still returning 200 ``{"deleted": true}``. +- ``delete_openai_container`` raises when the response body does not + report ``{"deleted": true}``, even on a 2xx response. +""" + +from __future__ import annotations + +import asyncio +import json + +import httpx +import pytest + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + + +def test_list_sends_openai_beta_header(monkeypatch): + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["headers"] = dict(request.headers) + seen["url"] = str(request.url) + return httpx.Response( + 200, + json = {"data": [{"id": "cntr_x", "name": "auto"}]}, + ) + + _mock_http_client(monkeypatch, handler) + result = _drive(_make_client().list_openai_containers()) + + assert result == [{"id": "cntr_x", "name": "auto"}] + assert seen["headers"].get("openai-beta") == "containers=v1" + assert seen["url"] == "https://api.openai.com/v1/containers" + + +def test_create_sends_openai_beta_header(monkeypatch): + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["headers"] = dict(request.headers) + seen["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response(200, json = {"id": "cntr_new", "name": "analysis"}) + + _mock_http_client(monkeypatch, handler) + result = _drive( + _make_client().create_openai_container(name = "analysis", ttl_minutes = 30) + ) + + assert result == {"id": "cntr_new", "name": "analysis"} + assert seen["headers"].get("openai-beta") == "containers=v1" + assert seen["body"]["name"] == "analysis" + assert seen["body"]["expires_after"] == { + "anchor": "last_active_at", + "minutes": 30, + } + + +def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch): + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["headers"] = dict(request.headers) + seen["url"] = str(request.url) + seen["method"] = request.method + return httpx.Response( + 200, + json = {"id": "cntr_x", "object": "container.deleted", "deleted": True}, + ) + + _mock_http_client(monkeypatch, handler) + _drive(_make_client().delete_openai_container("cntr_x")) + + assert seen["method"] == "DELETE" + assert seen["url"] == "https://api.openai.com/v1/containers/cntr_x" + assert seen["headers"].get("openai-beta") == "containers=v1" + + +def test_delete_raises_when_response_lacks_deleted_true(monkeypatch): + """OpenAI returns 200 ``{"deleted": true}`` even when the request is + silently rejected (e.g. before we started sending OpenAI-Beta). + Defensive guard: when the body omits ``deleted: true``, surface it + as an error so the UI can report the failure instead of falsely + reporting success.""" + + def handler(request: httpx.Request) -> httpx.Response: + # 200 but no deleted flag — simulate an unexpected payload shape. + return httpx.Response(200, json = {"id": "cntr_x", "object": "container"}) + + _mock_http_client(monkeypatch, handler) + + with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"): + _drive(_make_client().delete_openai_container("cntr_x")) + + +def test_delete_raises_when_deleted_is_false(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json = {"id": "cntr_x", "object": "container.deleted", "deleted": False}, + ) + + _mock_http_client(monkeypatch, handler) + + with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"): + _drive(_make_client().delete_openai_container("cntr_x")) + + +def test_delete_raises_when_body_is_not_json(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content = b"<html>OK</html>") + + _mock_http_client(monkeypatch, handler) + + with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"): + _drive(_make_client().delete_openai_container("cntr_x")) + + +def test_delete_propagates_openai_4xx(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404, json = {"error": {"message": "not found"}}) + + _mock_http_client(monkeypatch, handler) + + with pytest.raises(httpx.HTTPStatusError): + _drive(_make_client().delete_openai_container("cntr_missing")) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 0579f84896..fb63748bf1 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -13,6 +13,7 @@ import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { Sources, SourcesGroup } from "@/components/assistant-ui/sources"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; +import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution"; import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python"; import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal"; import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search"; @@ -750,9 +751,18 @@ const CodeToolsToggle: FC = () => { (s) => !!s.params.checkpoint && !s.modelLoading, ); const supportsTools = useChatRuntimeStore((s) => s.supportsTools); + // External providers have no local tool runtime, but Anthropic's + // Claude 4.x dispatches code_execution_20250825 server-side. The + // chat-page resolver stashes that capability in the runtime store + // (next to supportsBuiltinWebSearch). Mirror of shared-composer's + // codeDisabled so this pill lights up in active threads too. + const supportsBuiltinCodeExecution = useChatRuntimeStore( + (s) => s.supportsBuiltinCodeExecution, + ); const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled); const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled); - const disabled = !(modelLoaded && supportsTools); + const disabled = + !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); return ( <button @@ -946,6 +956,7 @@ const AssistantMessage: FC = () => { web_search: WebSearchToolUI, python: PythonToolUI, terminal: TerminalToolUI, + code_execution: CodeExecutionToolUI, }, Fallback: ToolFallback, }, diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx new file mode 100644 index 0000000000..8141b8b2cc --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react"; +import { FileTextIcon, LoaderIcon, TerminalIcon } from "lucide-react"; +import { memo, useEffect, useState } from "react"; +import { + ToolFallbackContent, + ToolFallbackRoot, + ToolFallbackTrigger, +} from "./tool-fallback"; + +/** + * Renders the synthetic `_toolEvent` chunks emitted by + * `_stream_anthropic` when Anthropic's `code_execution_20250825` tool + * fires. The backend collapses Anthropic's two sub-tools + * (`bash_code_execution`, `text_editor_code_execution`) into a single + * `tool_name: "code_execution"`, with `arguments.kind` ("bash" or + * "text_editor") and a per-kind argument shape: + * + * kind=bash: { command: "<shell command>" } + * kind=text_editor: { command: "view"|"create"|"str_replace", path, ... } + * + * The `result` payload is preformatted text: + * - bash: stdout, then "--- stderr ---" block + return_code if non-zero + * - text_editor view: file contents verbatim + * - text_editor create: "Created <path>" / "Updated <path>" + * - text_editor str_replace: unified-diff `lines` joined with "\n" + * - error: "Error: <error_code>" + */ +interface CodeExecutionArgs { + kind?: "bash" | "text_editor"; + command?: string; + path?: string; +} + +const CodeExecutionToolUIImpl: ToolCallMessagePartComponent = ({ + args, + result, + status, +}) => { + const parsedArgs = (args as CodeExecutionArgs) ?? {}; + const kind = parsedArgs.kind ?? "bash"; + const command = parsedArgs.command ?? ""; + const path = parsedArgs.path ?? ""; + const isRunning = status?.type === "running"; + + let runningLabel: string; + let completedLabel: string; + let Icon = TerminalIcon; + if (kind === "text_editor") { + Icon = FileTextIcon; + if (command === "view") { + runningLabel = path ? `Viewing ${path}…` : "Viewing file…"; + completedLabel = path ? `Viewed ${path}` : "Viewed file"; + } else if (command === "create") { + runningLabel = path ? `Writing ${path}…` : "Writing file…"; + completedLabel = path ? `Wrote ${path}` : "Wrote file"; + } else if (command === "str_replace") { + runningLabel = path ? `Editing ${path}…` : "Editing file…"; + completedLabel = path ? `Edited ${path}` : "Edited file"; + } else { + runningLabel = "Running file operation…"; + completedLabel = "File operation"; + } + } else { + runningLabel = "Running command…"; + completedLabel = command ? `Ran \`${command}\`` : "Ran command"; + } + + // Collapse the card once the model has resumed streaming prose after + // the tool call. Mirrors WebSearchToolUI's behavior so the tool-card + // doesn't crowd the final answer once the run is done. + const hasText = useAuiState(({ message }) => + message.content.some( + (p) => + p.type === "text" && + "text" in p && + (p as { text: string }).text.length > 0, + ), + ); + const [open, setOpen] = useState(isRunning); + useEffect(() => { + if (isRunning) { + setOpen(true); + } else if (hasText) { + setOpen(false); + } + }, [isRunning, hasText]); + + const resultText = + typeof result === "string" + ? result + : result != null + ? JSON.stringify(result, null, 2) + : ""; + + return ( + <ToolFallbackRoot open={open} onOpenChange={setOpen}> + <ToolFallbackTrigger + toolName={isRunning ? runningLabel : completedLabel} + status={status} + icon={Icon} + /> + <ToolFallbackContent> + {isRunning ? ( + <div className="flex items-center gap-2 text-sm text-muted-foreground"> + <LoaderIcon className="size-3.5 animate-spin" /> + <span>{runningLabel}</span> + </div> + ) : resultText ? ( + <pre className="max-h-64 overflow-auto whitespace-pre-wrap break-words rounded bg-muted/50 p-2 text-xs"> + {resultText} + </pre> + ) : null} + </ToolFallbackContent> + </ToolFallbackRoot> + ); +}; + +export const CodeExecutionToolUI = memo( + CodeExecutionToolUIImpl, +) as unknown as ToolCallMessagePartComponent; +CodeExecutionToolUI.displayName = "CodeExecutionToolUI"; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index edbb55df19..7084b946bb 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -15,6 +15,8 @@ import { streamChatCompletions, validateModel, } from "./chat-api"; +import { pickFriendlyContainerName } from "../lib/friendly-names"; +import { createOpenAIContainer } from "./openai-containers"; import { encryptProviderApiKey, isProviderKeyRotationError, @@ -38,6 +40,7 @@ import { getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, + providerSupportsBuiltinCodeExecution, providerSupportsBuiltinWebSearch, } from "../provider-capabilities"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; @@ -983,6 +986,106 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { forceRefreshPublicKey = false, ): Promise<OpenAIChatCompletionsRequest> => { if (externalSelection && externalProvider) { + // OpenAI shell-tool container reuse: pull the per-thread + // container_id (if any) so subsequent turns in the same + // thread reference the existing container instead of + // auto-creating a fresh one. Empty string / undefined → + // backend falls back to container_auto. Anthropic doesn't + // use this (server-side per-turn container). + let openaiCodeExecContainerId: string | null = null; + const codeExecEnabledForThisTurn = + codeToolsEnabled && + providerSupportsBuiltinCodeExecution( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ); + if (codeExecEnabledForThisTurn && resolvedThreadId) { + try { + const thread = await db.threads.get(resolvedThreadId); + openaiCodeExecContainerId = + thread?.openaiCodeExecContainerId ?? null; + } catch { + openaiCodeExecContainerId = null; + } + // Cross-thread inheritance: when the active thread has + // no container yet, default to the one most recently + // used on *any* other thread (provider-scoped). + // Matches what the Code Execution settings section + // shows in the picker, and keeps the user from getting + // a fresh container on every new thread. The picker + // can still be set to "Auto-create per thread" + // explicitly to opt into a fresh container — but + // that's done via the dropdown, not silently. + if ( + !openaiCodeExecContainerId && + externalProvider.providerType === "openai" + ) { + try { + const others = await db.threads + .orderBy("createdAt") + .reverse() + .toArray(); + for (const t of others) { + if (t.id === resolvedThreadId) continue; + if (t.openaiCodeExecContainerId) { + openaiCodeExecContainerId = t.openaiCodeExecContainerId; + void db.threads + .update(resolvedThreadId, { + openaiCodeExecContainerId, + }) + .catch(() => {}); + break; + } + } + } catch { + /* fall through to lazy-create below */ + } + } + // Lazy pre-create when there's no inherited container. + // We always POST /v1/containers ourselves (rather than + // letting the backend send container_auto) so every + // container shows up in the picker with a friendly + // English-word name and the user's configured TTL. + // Falls back to container_auto only if the POST fails + // — keeps the chat moving in that case. + if ( + !openaiCodeExecContainerId && + externalProvider.providerType === "openai" + ) { + const ttl = externalProvider.openaiContainerTtlMinutes; + const ttlToUse = + typeof ttl === "number" && ttl >= 1 ? ttl : 20; + try { + const created = await createOpenAIContainer( + { + apiKey: externalApiKey, + baseUrl: externalProvider.baseUrl || null, + }, + { + // Friendly English-word name so the container + // is human-readable in the picker list (e.g. + // "kestrel-3f9c") instead of a thread-id slug + // or OpenAI's default blank name. + name: pickFriendlyContainerName(), + ttlMinutes: ttlToUse, + }, + ); + openaiCodeExecContainerId = created.id; + void db.threads + .update(resolvedThreadId, { + openaiCodeExecContainerId: created.id, + }) + .catch(() => {}); + } catch { + // Fall back to backend's container_auto path on + // failure — keeps the chat moving; the next turn + // can retry. The auto-created container will be + // unnamed, but the chat doesn't break. + openaiCodeExecContainerId = null; + } + } + } return { model: externalSelection.modelId, messages: outboundMessages, @@ -1016,18 +1119,40 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(externalCapabilities?.presencePenalty ? { presence_penalty: params.presencePenalty } : {}), - // Built-in web search: when the user has the Search toggle - // on AND the active provider supports a server-side - // web_search tool (currently OpenAI's /v1/responses), pass - // the enable_tools shorthand. Backend translates - // enabled_tools=["web_search"] into the provider's tool - // schema — for OpenAI that's `tools: [{type:"web_search"}]` - // on the Responses body, see _stream_openai_responses. - ...(toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType) + // Built-in tools: Search pill maps to provider-side + // web_search (currently OpenAI / Anthropic / OpenRouter / + // Kimi); Code pill maps to Anthropic's server-side + // code_execution_20250825 tool (Anthropic is the only + // external provider that ships one today). Backend + // translates enabled_tools into each provider's tool + // schema — for Anthropic that's the entries appended to + // body["tools"] inside _stream_anthropic. + ...((toolsEnabled && + providerSupportsBuiltinWebSearch(externalProvider.providerType)) || + (codeToolsEnabled && + providerSupportsBuiltinCodeExecution( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + )) ? { enable_tools: true, - enabled_tools: ["web_search"], + enabled_tools: [ + ...(toolsEnabled && + providerSupportsBuiltinWebSearch( + externalProvider.providerType, + ) + ? ["web_search"] + : []), + ...(codeToolsEnabled && + providerSupportsBuiltinCodeExecution( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ) + ? ["code_execution"] + : []), + ], } : {}), provider_id: externalProvider.id, @@ -1042,6 +1167,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } : {}), provider_base_url: externalProvider.baseUrl || null, + ...(openaiCodeExecContainerId + ? { + openai_code_exec_container_id: openaiCodeExecContainerId, + } + : {}), ...(supportsProviderPromptCaching(externalProvider.providerType) ? { enable_prompt_caching: @@ -1125,6 +1255,34 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // On tool_end: set result on the existing part (transitions to "complete"). const toolEvent = (chunk as unknown as { _toolEvent?: Record<string, unknown> })._toolEvent; if (toolEvent !== undefined) { + // OpenAI shell-tool container persistence — see + // ThreadRecord.openaiCodeExecContainerId. The backend + // emits these synthetic events on the OpenAI Responses + // SSE stream after capturing the container_id from a + // response, or detecting an expired-container error. + if (toolEvent.type === "container_ready") { + const newContainerId = toolEvent.container_id as + | string + | undefined; + if (newContainerId && resolvedThreadId) { + void db.threads + .update(resolvedThreadId, { + openaiCodeExecContainerId: newContainerId, + }) + .catch(() => {}); + } + continue; + } + if (toolEvent.type === "container_invalidated") { + if (resolvedThreadId) { + void db.threads + .update(resolvedThreadId, { + openaiCodeExecContainerId: null, + }) + .catch(() => {}); + } + continue; + } if (toolEvent.type === "tool_start") { const id = (toolEvent.tool_call_id as string) || `${toolEvent.tool_name}_${Date.now()}`; const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"]; diff --git a/studio/frontend/src/features/chat/api/openai-containers.ts b/studio/frontend/src/features/chat/api/openai-containers.ts new file mode 100644 index 0000000000..ca311cc311 --- /dev/null +++ b/studio/frontend/src/features/chat/api/openai-containers.ts @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * Wrappers for the three OpenAI shell-tool container management + * endpoints exposed by the backend (studio/backend/routes/inference.py). + * Each one proxies to OpenAI's /v1/containers REST surface using the + * user's encrypted API key. Backend rejects any base URL that isn't + * api.openai.com — the shell tool only exists on the managed cloud. + */ + +import { authFetch } from "@/features/auth"; +import { encryptProviderApiKey } from "./providers-api"; + +export interface OpenAIContainerSummary { + id: string; + name?: string | null; + createdAt?: number | null; + lastActiveAt?: number | null; + expiresAfterMinutes?: number | null; + status?: string | null; +} + +interface RawSummary { + id: string; + name?: string | null; + created_at?: number | null; + last_active_at?: number | null; + expires_after_minutes?: number | null; + status?: string | null; +} + +function fromRaw(raw: RawSummary): OpenAIContainerSummary { + return { + id: raw.id, + name: raw.name ?? null, + createdAt: raw.created_at ?? null, + lastActiveAt: raw.last_active_at ?? null, + expiresAfterMinutes: raw.expires_after_minutes ?? null, + status: raw.status ?? null, + }; +} + +async function parseError(response: Response): Promise<string> { + try { + const body = (await response.json()) as { detail?: string }; + if (body && typeof body.detail === "string") return body.detail; + } catch { + /* fall through */ + } + return `HTTP ${response.status}`; +} + +interface AuthInputs { + apiKey: string; + baseUrl: string | null; +} + +async function buildAuthBody(auth: AuthInputs) { + return { + encrypted_api_key: await encryptProviderApiKey(auth.apiKey), + provider_base_url: auth.baseUrl, + }; +} + +export async function listOpenAIContainers( + auth: AuthInputs, +): Promise<OpenAIContainerSummary[]> { + const response = await authFetch( + "/api/inference/external/openai/containers/list", + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(await buildAuthBody(auth)), + }, + ); + if (!response.ok) throw new Error(await parseError(response)); + const body = (await response.json()) as { containers?: RawSummary[] }; + return (body.containers ?? []).map(fromRaw); +} + +export async function createOpenAIContainer( + auth: AuthInputs, + params: { name: string; ttlMinutes: number }, +): Promise<OpenAIContainerSummary> { + const response = await authFetch( + "/api/inference/external/openai/containers/create", + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + ...(await buildAuthBody(auth)), + name: params.name, + ttl_minutes: params.ttlMinutes, + }), + }, + ); + if (!response.ok) throw new Error(await parseError(response)); + const raw = (await response.json()) as RawSummary; + return fromRaw(raw); +} + +export async function deleteOpenAIContainer( + auth: AuthInputs, + containerId: string, +): Promise<void> { + const response = await authFetch( + "/api/inference/external/openai/containers/delete", + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + ...(await buildAuthBody(auth)), + container_id: containerId, + }), + }, + ); + if (!response.ok && response.status !== 204) { + throw new Error(await parseError(response)); + } +} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 2ebd8f45f0..d08bee1f8f 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -50,6 +50,7 @@ import { clampReasoningEffortToLevels, getExternalReasoningCapabilities, getProviderCapabilities, + providerSupportsBuiltinCodeExecution, providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; @@ -716,6 +717,11 @@ export function ChatPage(): ReactElement { const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( provider?.providerType, ); + const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( + provider?.providerType, + selection.modelId, + provider?.baseUrl, + ); // Kimi's k2.6/k2.5 default to thinking enabled on the server side // (per https://platform.kimi.ai/docs/models). Mirror that default // in the UI so the Think pill comes up clicked when the user picks @@ -748,14 +754,16 @@ export function ChatPage(): ReactElement { : true : state.reasoningEnabled, supportsPreserveThinking: false, - // External models never give us a local tool runtime (no Code - // execution, no python sandbox), so `supportsTools` must be - // false — that's what gates the Code pill in the composer. - // `supportsBuiltinWebSearch` is the separate flag that lets the - // Search pill light up for providers (currently just OpenAI) who - // run web_search server-side. + // External models never give us a local tool runtime (no + // python sandbox), so `supportsTools` must be false. The two + // `supportsBuiltin*` flags pick up the slack for providers that + // run the tool server-side: `supportsBuiltinWebSearch` lights + // up the Search pill (OpenAI / Anthropic / OpenRouter / Kimi), + // `supportsBuiltinCodeExecution` lights up the Code pill + // (Anthropic Claude 4.x only, today). supportsTools: false, supportsBuiltinWebSearch, + supportsBuiltinCodeExecution, toolsEnabled: searchOnByDefault, codeToolsEnabled: false, }); @@ -915,6 +923,11 @@ export function ChatPage(): ReactElement { const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( selectedProvider?.providerType, ); + const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( + selectedProvider?.providerType, + selectedExternal?.modelId, + selectedProvider?.baseUrl, + ); // See sibling useEffect above: Kimi's k2.x default to thinking // enabled, so the Think pill comes up clicked. Search pill stays // off by default; mutual exclusion flips them via the composer. @@ -946,12 +959,15 @@ export function ChatPage(): ReactElement { : true : store.reasoningEnabled, supportsPreserveThinking: false, - // External models have no local tool runtime → supportsTools=false - // keeps the Code pill greyed out. supportsBuiltinWebSearch is the - // separate flag the composer reads to light up the Search pill - // when the provider offers a server-side web_search tool. + // External models have no local tool runtime → supportsTools + // stays false. The two supportsBuiltin* flags carry the + // server-side capability info for each pill: + // - Search → providerSupportsBuiltinWebSearch + // - Code → providerSupportsBuiltinCodeExecution + // (Anthropic Claude 4.x only, today) supportsTools: false, supportsBuiltinWebSearch, + supportsBuiltinCodeExecution, toolsEnabled: searchOnByDefault, codeToolsEnabled: false, ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }), diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 74e4943b3c..6de7d6a330 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -66,6 +66,8 @@ import { toast } from "sonner"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { type ExternalProviderConfig, + getExternalProviderApiKey, + parseExternalModelId, supportsProviderPromptCaching, } from "./external-providers"; import { @@ -84,9 +86,11 @@ import { toPresetParams, type Preset, } from "./presets/preset-policy"; +import { OpenAICodeExecSection } from "./components/openai-code-exec-section"; import { EXTERNAL_MAX_OUTPUT_TOKENS, getExternalMinOutputTokens, + providerSupportsBuiltinCodeExecution, type ProviderCapabilities, } from "./provider-capabilities"; import type { InferenceParams } from "./types/runtime"; @@ -675,6 +679,21 @@ export function ChatSettingsPanel({ supportsProviderPromptCaching(activeExternalProvider.providerType); const promptCachingEnabled = activeExternalProvider?.enablePromptCaching !== false; + const externalSelection = currentCheckpoint + ? parseExternalModelId(currentCheckpoint) + : null; + const showOpenAICodeExecSection = + activeExternalProvider != null && + providerSupportsBuiltinCodeExecution( + activeExternalProvider.providerType, + externalSelection?.modelId, + activeExternalProvider.baseUrl, + ) && + activeExternalProvider.providerType === "openai"; + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const openAiApiKeyForSection = activeExternalProvider + ? getExternalProviderApiKey(activeExternalProvider.id) || null + : null; function set<K extends keyof InferenceParams>(key: K) { return (v: InferenceParams[K]) => { @@ -1184,6 +1203,17 @@ export function ChatSettingsPanel({ </CollapsibleSection> ) : null} + {showOpenAICodeExecSection && activeExternalProvider ? ( + <CollapsibleSection label="Code Execution" defaultOpen={false}> + <OpenAICodeExecSection + provider={activeExternalProvider} + apiKey={openAiApiKeyForSection} + activeThreadId={activeThreadId} + onProviderChange={(p) => onExternalProviderChange?.(p)} + /> + </CollapsibleSection> + ) : null} + <CollapsibleSection label="System Prompt" defaultOpen={true}> <button type="button" diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx new file mode 100644 index 0000000000..15cb79286c --- /dev/null +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -0,0 +1,477 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * Settings-sheet section for OpenAI shell-tool container management. + * Renders only when: + * - active provider is OpenAI cloud (api.openai.com base URL), AND + * - the active model is gpt-5.5 or gpt-5.5-pro (the only families + * where the shell tool is wired through today). + * + * Surfaces three controls: + * 1. Default container idle-timeout (minutes). Persists on the + * provider record; pre-fills the create dialog and is used by + * the chat-adapter's lazy-create path on the first turn of a + * thread. + * 2. Container picker for the *active thread* — pick any of the + * user's existing OpenAI containers, or "Auto-create per thread" + * (default; lets the auto-create path manage it). + * 3. Create-new-container inline form. Refresh + delete actions + * per row. + * + * State persistence: + * - TTL → ExternalProviderConfig.openaiContainerTtlMinutes + * - Active container for this thread → ThreadRecord.openaiCodeExecContainerId + * + * No new global stores — list is fetched on open / refresh and held + * in component state. + */ + +"use client"; + +import { useCallback, useEffect, useMemo, useState } from "react"; +import { toast } from "sonner"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Skeleton } from "@/components/ui/skeleton"; +import { TrashIcon, RefreshCwIcon, PlusIcon } from "lucide-react"; +import { + createOpenAIContainer, + deleteOpenAIContainer, + listOpenAIContainers, + type OpenAIContainerSummary, +} from "../api/openai-containers"; +import { db } from "../db"; +import type { ExternalProviderConfig } from "../external-providers"; +import { useLiveQuery } from "../db"; +import { ensureThreadRecord } from "../runtime-provider"; + +const AUTO_OPTION_VALUE = "__auto__"; +const DEFAULT_TTL_MINUTES = 20; +const TTL_MIN = 1; +const TTL_MAX = 10080; // one week — matches backend bound + +function ageLabel(epochSeconds: number | null | undefined): string { + if (!epochSeconds) return ""; + const ageSec = Math.max(0, Math.floor(Date.now() / 1000) - epochSeconds); + if (ageSec < 60) return `${ageSec}s ago`; + const ageMin = Math.floor(ageSec / 60); + if (ageMin < 60) return `${ageMin}m ago`; + const ageHr = Math.floor(ageMin / 60); + if (ageHr < 48) return `${ageHr}h ago`; + const ageDay = Math.floor(ageHr / 24); + return `${ageDay}d ago`; +} + +interface OpenAICodeExecSectionProps { + provider: ExternalProviderConfig; + apiKey: string | null; + activeThreadId: string | null; + onProviderChange: (provider: ExternalProviderConfig) => void; +} + +export function OpenAICodeExecSection({ + provider, + apiKey, + activeThreadId, + onProviderChange, +}: OpenAICodeExecSectionProps) { + const [containers, setContainers] = useState<OpenAIContainerSummary[]>([]); + const [isLoading, setIsLoading] = useState(false); + const [creating, setCreating] = useState(false); + const [createOpen, setCreateOpen] = useState(false); + const [createName, setCreateName] = useState(""); + const [createTtl, setCreateTtl] = useState<number>( + provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES, + ); + + const thread = useLiveQuery( + async () => (activeThreadId ? db.threads.get(activeThreadId) : undefined), + [activeThreadId], + ); + const activeContainerId = thread?.openaiCodeExecContainerId ?? null; + + // Containers sorted newest-first by lastActiveAt so the dropdown's + // default (auto-bind target) shows up first. + const sortedContainers = useMemo( + () => + [...containers].sort( + (a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0), + ), + [containers], + ); + + // What the dropdown should display right now. We decouple this from + // `activeContainerId` (which is whatever is in Dexie) so the user + // immediately sees the most-recent container by name when there is + // no thread binding yet, rather than a "Selecting most recent…" + // placeholder while the auto-bind effect's async write propagates + // back through useLiveQuery. The auto-bind effect still writes the + // bind to Dexie so the chat adapter sees it on send. + const displayedContainerId = + activeContainerId ?? sortedContainers[0]?.id ?? null; + + const refresh = useCallback(async () => { + if (!apiKey) return; + setIsLoading(true); + try { + const list = await listOpenAIContainers({ + apiKey, + baseUrl: provider.baseUrl || null, + }); + setContainers(list); + } catch (err) { + toast.error( + `Failed to list containers: ${err instanceof Error ? err.message : "Unknown"}`, + ); + } finally { + setIsLoading(false); + } + }, [apiKey, provider.baseUrl]); + + // Fetch once when the section mounts (or provider changes). + useEffect(() => { + void refresh(); + }, [refresh]); + + // Auto-bind the active thread to the most-recently-active container + // whenever the thread has none set and at least one container exists + // on the user's OpenAI account. Sorting by `lastActiveAt` matches + // what feels "most recent" from the user's perspective. + // + // We eagerly materialize the thread row via `ensureThreadRecord` so + // the bind actually lands in Dexie before the user has sent a first + // message. This does NOT create anything at OpenAI — only a local + // ThreadRecord — so it does not bypass the user's expectation that + // a fresh OpenAI container is not created until first send. + // + // If `containers` is empty (no OpenAI containers exist yet), this + // effect short-circuits: the picker renders an empty-state hint and + // the chat-adapter's lazy-create path will mint the first container + // on first send. + useEffect(() => { + if (!activeThreadId || activeContainerId || containers.length === 0) { + return; + } + const sorted = [...containers].sort( + (a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0), + ); + const candidate = sorted[0]; + if (!candidate) return; + void (async () => { + try { + await ensureThreadRecord({ + threadId: activeThreadId, + modelType: "base", + }); + await db.threads.update(activeThreadId, { + openaiCodeExecContainerId: candidate.id, + }); + } catch { + // Best-effort; the chat-adapter will inherit/create on send. + } + })(); + }, [activeThreadId, activeContainerId, containers]); + + const ttlValue = provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES; + + const onTtlChange = (raw: string) => { + const n = parseInt(raw, 10); + if (Number.isNaN(n)) return; + const clamped = Math.min(Math.max(n, TTL_MIN), TTL_MAX); + onProviderChange({ ...provider, openaiContainerTtlMinutes: clamped }); + }; + + const onPick = async (value: string) => { + if (!activeThreadId || !value) return; + // value is always a container id now — the "Auto-create per thread" + // option has been removed in favour of always defaulting to the + // most-recently-active container. The chat-adapter still handles + // the no-containers-exist case (lazy-create on first send). + // + // ensureThreadRecord materializes the thread row eagerly (modelType + // "base" — settings sheet is single-thread-mode only) so the update + // actually lands when the user hasn't sent a message yet. + try { + await ensureThreadRecord({ threadId: activeThreadId, modelType: "base" }); + const affected = await db.threads.update(activeThreadId, { + openaiCodeExecContainerId: value, + }); + if (affected === 0) { + toast.error("Could not update thread."); + } + } catch (err) { + toast.error( + `Could not update thread: ${err instanceof Error ? err.message : "Unknown"}`, + ); + } + }; + + const onCreate = async () => { + if (!apiKey) return; + const name = createName.trim(); + if (!name) { + toast.error("Container name is required"); + return; + } + setCreating(true); + try { + const created = await createOpenAIContainer( + { apiKey, baseUrl: provider.baseUrl || null }, + { name, ttlMinutes: createTtl }, + ); + toast.success(`Created container ${name}`); + setCreateName(""); + setCreateOpen(false); + await refresh(); + // Auto-bind the just-created container to the active thread. + // ensureThreadRecord first so the bind lands even when the user + // creates a container before sending the first message — without + // it, db.threads.update silently affects 0 rows and the chat + // adapter falls back to cross-thread inheritance / lazy-create, + // which can pick a stale container that fails with "container + // does not exist" on the first turn. + if (activeThreadId) { + try { + await ensureThreadRecord({ + threadId: activeThreadId, + modelType: "base", + }); + await db.threads.update(activeThreadId, { + openaiCodeExecContainerId: created.id, + }); + } catch { + /* best-effort; toast above already confirmed creation */ + } + } + } catch (err) { + toast.error( + `Create failed: ${err instanceof Error ? err.message : "Unknown"}`, + ); + } finally { + setCreating(false); + } + }; + + const onDelete = async (id: string, name: string | null | undefined) => { + if (!apiKey) return; + if ( + !window.confirm( + `Delete container ${name || id}? Threads using it will fall back to auto-create on their next turn.`, + ) + ) { + return; + } + try { + await deleteOpenAIContainer( + { apiKey, baseUrl: provider.baseUrl || null }, + id, + ); + // Clear any thread bindings pointing at the now-deleted id. + const affected = await db.threads + .filter((t) => t.openaiCodeExecContainerId === id) + .toArray(); + await Promise.all( + affected.map((t) => + db.threads.update(t.id, { openaiCodeExecContainerId: null }), + ), + ); + toast.success(`Deleted container ${name || id}`); + await refresh(); + } catch (err) { + toast.error( + `Delete failed: ${err instanceof Error ? err.message : "Unknown"}`, + ); + } + }; + + return ( + <div className="flex flex-col gap-3 pt-1"> + {/* TTL */} + <div className="flex items-center justify-between gap-3"> + <label + htmlFor="openai-container-ttl" + className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg" + > + New-container idle timeout (min) + </label> + <Input + id="openai-container-ttl" + type="number" + min={TTL_MIN} + max={TTL_MAX} + value={ttlValue} + onChange={(e) => onTtlChange(e.target.value)} + className="h-8 w-24 text-sm" + /> + </div> + + {/* Active container picker — visually emphasized so it reads as + the primary control vs. the static list below. Accent + background + ring outline distinguish it from the plain + bordered list items beneath. */} + <div className="flex flex-col gap-1.5 rounded-md border border-primary/30 bg-primary/5 p-2.5"> + <div className="flex items-center justify-between gap-2"> + <span className="text-[13px] font-semibold leading-[1.25] tracking-nav text-primary"> + Active for this thread + </span> + <Button + size="sm" + variant="ghost" + className="h-7 px-2" + onClick={() => void refresh()} + disabled={isLoading || !apiKey} + aria-label="Refresh container list" + > + <RefreshCwIcon + className={`size-3.5 ${isLoading ? "animate-spin" : ""}`} + /> + </Button> + </div> + {/* When no containers exist yet, render a disabled placeholder + instead of the picker. The first one is created by the + chat-adapter on first send (lazy-create) and will appear + here after the next refresh. */} + {sortedContainers.length === 0 ? ( + <div className="h-9 w-full rounded-md border border-primary/40 bg-background px-2 flex items-center text-sm text-muted-foreground"> + (none yet — will be created on first send) + </div> + ) : ( + <select + value={displayedContainerId ?? sortedContainers[0].id} + onChange={(e) => onPick(e.target.value)} + disabled={!activeThreadId} + className="h-9 w-full rounded-md border border-primary/40 bg-background px-2 text-sm font-medium shadow-sm focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-primary/40" + > + {sortedContainers.map((c) => ( + <option key={c.id} value={c.id}> + {c.name ?? "(unnamed)"} · {c.id.slice(0, 14)}… + {c.lastActiveAt ? ` · active ${ageLabel(c.lastActiveAt)}` : ""} + </option> + ))} + </select> + )} + </div> + + {/* Container list with delete actions — labeled and visually + quieter so it's clearly the "all containers, manage them" + area rather than the active selector above. */} + <div className="flex flex-col gap-1.5"> + <span className="text-[11px] uppercase tracking-wider text-muted-foreground"> + All containers + </span> + {isLoading && containers.length === 0 ? ( + <Skeleton className="h-16 w-full" /> + ) : containers.length > 0 ? ( + <ul className="flex flex-col gap-1 max-h-44 overflow-auto"> + {containers.map((c) => { + const isActive = c.id === activeContainerId; + return ( + <li + key={c.id} + className={`flex items-center justify-between gap-2 rounded-md border px-2 py-1.5 text-xs ${ + isActive + ? "border-primary/30 bg-primary/5" + : "border-border/60" + }`} + > + <div className="flex min-w-0 flex-col"> + <span className="truncate font-medium"> + {c.name ?? "(unnamed)"} + {isActive ? ( + <span className="ml-1.5 text-[10px] font-normal uppercase tracking-wider text-primary"> + · active + </span> + ) : null} + </span> + <span className="text-muted-foreground"> + {c.id} · TTL{" "} + {c.expiresAfterMinutes ?? DEFAULT_TTL_MINUTES}m + </span> + </div> + <Button + size="sm" + variant="ghost" + className="h-6 w-6 p-0 text-destructive" + onClick={() => void onDelete(c.id, c.name)} + aria-label={`Delete container ${c.name ?? c.id}`} + > + <TrashIcon className="size-3.5" /> + </Button> + </li> + ); + })} + </ul> + ) : ( + <p className="text-xs text-muted-foreground"> + No saved containers yet. Use auto-create or create a named + one below. + </p> + )} + </div> + + {/* Create new */} + {createOpen ? ( + <div className="flex flex-col gap-2 rounded-md border border-border/60 p-2"> + <Input + placeholder="Container name (e.g. data-analysis)" + value={createName} + onChange={(e) => setCreateName(e.target.value)} + className="h-8 text-sm" + /> + <div className="flex items-center gap-2"> + <Input + type="number" + min={TTL_MIN} + max={TTL_MAX} + value={createTtl} + onChange={(e) => { + const n = parseInt(e.target.value, 10); + if (!Number.isNaN(n)) + setCreateTtl(Math.min(Math.max(n, TTL_MIN), TTL_MAX)); + }} + className="h-8 w-24 text-sm" + aria-label="Idle timeout in minutes" + /> + <span className="text-xs text-muted-foreground">min idle</span> + <div className="flex-1" /> + <Button + size="sm" + variant="ghost" + className="h-7" + onClick={() => { + setCreateOpen(false); + setCreateName(""); + }} + disabled={creating} + > + Cancel + </Button> + <Button + size="sm" + className="h-7" + onClick={() => void onCreate()} + disabled={creating || !createName.trim() || !apiKey} + > + Create + </Button> + </div> + </div> + ) : ( + <Button + size="sm" + variant="outline" + className="h-8" + onClick={() => { + setCreateTtl(ttlValue); + setCreateOpen(true); + }} + disabled={!apiKey} + > + <PlusIcon className="size-3.5 mr-1" /> + New container + </Button> + )} + </div> + ); +} diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 5f645042bf..89591d7b14 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -18,6 +18,13 @@ export interface ExternalProviderConfig { enablePromptCaching?: boolean; /** User-pinned: the loaded vLLM model supports `enable_thinking`. */ isReasoningModel?: boolean; + /** + * Default idle-timeout (in minutes) for newly created OpenAI shell + * containers. Pre-fills the "Create container" dialog and is the + * TTL the auto-create-per-thread path POSTs to /v1/containers with. + * OpenAI's hard default is 20. Only meaningful for OpenAI cloud. + */ + openaiContainerTtlMinutes?: number; createdAt: number; updatedAt: number; } @@ -226,6 +233,12 @@ function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig isReasoningModel: supportsProviderReasoningToggle(providerType) ? raw.isReasoningModel === true : undefined, + openaiContainerTtlMinutes: + providerType === "openai" && + typeof raw.openaiContainerTtlMinutes === "number" && + raw.openaiContainerTtlMinutes >= 1 + ? Math.min(raw.openaiContainerTtlMinutes, 10080) + : undefined, }; } diff --git a/studio/frontend/src/features/chat/lib/friendly-names.ts b/studio/frontend/src/features/chat/lib/friendly-names.ts new file mode 100644 index 0000000000..503008f761 --- /dev/null +++ b/studio/frontend/src/features/chat/lib/friendly-names.ts @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * Friendly default names for auto-created OpenAI shell containers. + * Used by the chat-adapter when the lazy-create path fires (Code pill + * on, no thread container yet, user has set a non-default TTL). The + * goal is a human-memorable label like "otter" or "harbor" instead of + * "chat-abc12345" — the user can still rename via the Studio-side + * alias map. + * + * The list is curated to: + * - Be unambiguous, non-offensive nouns from natural categories + * (animals, plants, geography, materials, weather). + * - Avoid technical / political / brand words that might read as + * odd in a chat UI. + * - Stay reasonably small so the bundle cost is negligible (~200 + * entries × ~7 bytes ≈ 1.5 KB). + * + * Collisions are tolerated — the container's real unique key is its + * ``cntr_*`` id, not its name. A short random hex suffix is appended + * to make accidental same-name collisions visually distinct in the + * picker list. + */ + +const WORDS = [ + // animals + "otter", + "falcon", + "heron", + "lynx", + "marten", + "stoat", + "raven", + "magpie", + "salmon", + "trout", + "perch", + "tortoise", + "gecko", + "iguana", + "axolotl", + "narwhal", + "manatee", + "dolphin", + "porpoise", + "octopus", + "cuttlefish", + "nautilus", + "starfish", + "urchin", + "anemone", + "coral", + "puffin", + "kestrel", + "osprey", + "buzzard", + "kingfisher", + "robin", + "wren", + "finch", + "sparrow", + "thrush", + "siskin", + "warbler", + "tanager", + "oriole", + "hare", + "badger", + "weasel", + "ferret", + "polecat", + "civet", + "tapir", + "okapi", + "ibex", + "chamois", + // plants & trees + "alder", + "aspen", + "birch", + "cedar", + "cypress", + "elder", + "elm", + "fir", + "ginkgo", + "hawthorn", + "hazel", + "hemlock", + "holly", + "juniper", + "larch", + "linden", + "maple", + "oak", + "olive", + "pine", + "rowan", + "spruce", + "sycamore", + "willow", + "yew", + "thistle", + "fern", + "moss", + "ivy", + "clover", + "heather", + "lavender", + "rosemary", + "sage", + "thyme", + "myrtle", + "laurel", + "magnolia", + // geography / landscape + "harbor", + "atoll", + "lagoon", + "estuary", + "fjord", + "delta", + "isthmus", + "mesa", + "plateau", + "valley", + "ridge", + "summit", + "glade", + "meadow", + "moor", + "heath", + "tundra", + "savanna", + "prairie", + "steppe", + "bayou", + "marsh", + "fen", + "grotto", + "cavern", + "canyon", + "ravine", + "gorge", + "knoll", + "dell", + "vale", + "coast", + // materials / minerals / colors + "amber", + "agate", + "onyx", + "opal", + "jade", + "quartz", + "obsidian", + "basalt", + "granite", + "marble", + "slate", + "flint", + "lapis", + "topaz", + "garnet", + "pearl", + "coral", + "ivory", + "ebony", + "copper", + "cobalt", + "indigo", + "saffron", + "vermilion", + "ochre", + "umber", + "sienna", + "russet", + // weather / sky / time + "aurora", + "comet", + "ember", + "frost", + "gale", + "harvest", + "monsoon", + "nebula", + "solstice", + "twilight", + "zephyr", + "drizzle", + "tempest", + "halcyon", + "equinox", + "rainbow", + "horizon", + "meridian", + "zenith", + "comet", + // misc tactile / cozy nouns + "lantern", + "kettle", + "compass", + "anchor", + "beacon", + "harbor", + "voyage", + "trellis", + "cottage", + "thicket", + "orchard", + "bramble", + "haystack", + "snowfall", + "campfire", +]; + +/** RFC 4122-ish 4-character lowercase hex suffix using crypto.randomUUID. */ +function randomHexSuffix(): string { + if ( + typeof crypto !== "undefined" && + typeof crypto.randomUUID === "function" + ) { + return crypto.randomUUID().replace(/-/g, "").slice(0, 4); + } + // Older browser fallback. Math.random is fine here — this is a + // display suffix, not a security token. + return Math.floor(Math.random() * 0xffff) + .toString(16) + .padStart(4, "0"); +} + +/** + * Returns a single English-word name with a short random hex suffix. + * + * Example output: "kestrel-3f9c", "harbor-a012". + * + * The suffix keeps containers visually distinguishable in the picker + * when the same word recurs across creations. + */ +export function pickFriendlyContainerName(): string { + const word = WORDS[Math.floor(Math.random() * WORDS.length)] ?? "container"; + return `${word}-${randomHexSuffix()}`; +} diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 1c40a25773..3c9bff40b9 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -121,6 +121,88 @@ export function providerSupportsBuiltinWebSearch( ); } +/** + * Whether the selected external provider/model exposes a server-side + * code-execution tool. Two providers ship one today: + * + * - **Anthropic** (`code_execution_20250825`): Python + bash + + * str_replace-based file edits inside a 5 GB sandboxed container + * per request. Documented at + * https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool + * + * - **OpenAI cloud** (`shell` on /v1/responses): bash inside a + * reusable container; we auto-create one on the first turn of a + * chat thread and reference it on subsequent turns via the + * thread's stored `openaiCodeExecContainerId`. Documented at + * https://developers.openai.com/api/docs/guides/tools-shell + * + * Returns false for every other provider. The backend additionally + * gates the OpenAI shell tool on `is_openai_cloud` so custom + * OpenAI-compat servers (ollama / llama.cpp / vLLM) that also report + * `provider_type="openai"` never receive the tool — but in practice + * none of those catalogs surface the `gpt-5.5` ids anyway, so the + * frontend prefix match is enough. + * + * v1 wires the tools themselves; file uploads (Anthropic + * `container_upload` / OpenAI `input_file`) are a deliberate follow-up. + */ +const ANTHROPIC_CODE_EXECUTION_MODEL_PREFIXES = [ + "claude-opus-4-7", + "claude-opus-4-6", + "claude-sonnet-4-6", + "claude-opus-4-5", + "claude-sonnet-4-5", + "claude-haiku-4-5", + // Deprecated upstream but the registry still exposes the ids, so the + // pill should remain functional for users on those snapshots. + "claude-opus-4-1", + "claude-opus-4", + "claude-sonnet-4", +] as const; + +// OpenAI cloud shell-tool gating. Docs only explicitly demonstrate +// gpt-5.5; gpt-5.5-pro is included because the family share the same +// /v1/responses contract. `gpt-5.5-pro` is checked first so the prefix +// match doesn't collide with a hypothetical `gpt-5.5-turbo` etc. +const OPENAI_CODE_EXECUTION_MODEL_PREFIXES = [ + "gpt-5.5-pro", + "gpt-5.5", +] as const; + +/** + * Strict check that a provider configuration points at OpenAI's + * managed cloud (api.openai.com), as opposed to a custom OpenAI-compat + * backend (ollama / llama.cpp / vLLM / generic "custom" preset). The + * shell tool ONLY exists on OpenAI cloud; sending it to anything else + * 400s the request. Mirror of the backend's + * `is_openai_cloud = "api.openai.com" in self.base_url` guard. + */ +function isOpenAICloudBaseUrl(baseUrl: string | null | undefined): boolean { + if (!baseUrl) return true; // No override → uses the default openai.com base. + return baseUrl.trim().toLowerCase().includes("api.openai.com"); +} + +export function providerSupportsBuiltinCodeExecution( + providerType: string | null | undefined, + modelId: string | null | undefined, + baseUrl?: string | null, +): boolean { + const normalized = modelId?.trim().toLowerCase() ?? ""; + if (!normalized) return false; + if (providerType === "anthropic") { + return ANTHROPIC_CODE_EXECUTION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); + } + if (providerType === "openai") { + if (!isOpenAICloudBaseUrl(baseUrl)) return false; + return OPENAI_CODE_EXECUTION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); + } + return false; +} + /** * Per-provider minimum on the outbound max_tokens. Kimi's docs require * `max_tokens >= 16000` whenever a thinking model is in use so the diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 1133e759ac..b019fd2d4d 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -396,7 +396,7 @@ function toThreadMessage(m: MessageRecord): ThreadMessage { }; } -async function ensureThreadRecord({ +export async function ensureThreadRecord({ threadId, modelType, pairId, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index da32dee3a4..c4ffa98467 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -24,7 +24,10 @@ import { type ReasoningEffort, useChatRuntimeStore, } from "./stores/chat-runtime-store"; -import { getExternalReasoningCapabilities } from "./provider-capabilities"; +import { + getExternalReasoningCapabilities, + providerSupportsBuiltinCodeExecution, +} from "./provider-capabilities"; import { type CompositionEvent, type KeyboardEvent, @@ -367,13 +370,21 @@ export function SharedComposer({ // Two-pill gating: Search pill lights up when the runtime has either // a local tool runtime (supportsTools, gives us our Code/python + local // web_search) OR a server-side web_search the provider runs for us - // (supportsBuiltinWebSearch, currently just OpenAI's /v1/responses). - // Code pill is gated on `supportsTools` only — external providers - // never give us code execution, so the pill must stay disabled even - // when Search is available. + // (supportsBuiltinWebSearch, currently OpenAI / Anthropic / OpenRouter + // / Kimi). Code pill lights up on the local runtime OR when Anthropic + // is selected with a model that accepts the server-side + // code_execution_20250825 tool — see + // providerSupportsBuiltinCodeExecution. Anthropic is the only external + // provider that ships a code-execution tool today. + const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( + selectedExternalProvider?.providerType, + effectiveExternalModelId, + selectedExternalProvider?.baseUrl, + ); const searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); - const codeDisabled = !modelLoaded || !supportsTools; + const codeDisabled = + !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); // Backwards-compatible alias for any other call site that may still // reference `toolsDisabled` (rare; both pills used it before). const toolsDisabled = codeDisabled; diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 4cbed5899b..9a80d12da4 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -234,11 +234,19 @@ type ChatRuntimeStore = { * web_search tool (OpenAI's /v1/responses today). Distinct from * `supportsTools` — that flag governs the local tool runtime (Code, * python sandbox, our DuckDuckGo web_search). This one only enables - * the chat composer's Search pill for external models and leaves - * the Code pill disabled, because external providers do not give - * us code execution. Local models keep `supportsTools` only. + * the chat composer's Search pill for external models. Local models + * keep `supportsTools` only. */ supportsBuiltinWebSearch: boolean; + /** + * Whether the active external provider exposes a server-side + * code-execution tool (Anthropic's `code_execution_20250825` on the + * Claude 4.x family). Distinct from `supportsTools` for the same + * reason as `supportsBuiltinWebSearch`: external providers don't + * give us a local tool runtime, but Anthropic dispatches code + * execution server-side. Read by both composers' Code pill gate. + */ + supportsBuiltinCodeExecution: boolean; toolsEnabled: boolean; codeToolsEnabled: boolean; toolStatus: string | null; @@ -331,6 +339,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({ preserveThinking: loadBool(PRESERVE_THINKING_KEY, false), supportsTools: false, supportsBuiltinWebSearch: false, + supportsBuiltinCodeExecution: false, toolsEnabled: false, codeToolsEnabled: false, toolStatus: null, @@ -442,6 +451,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({ supportsPreserveThinking: false, supportsTools: false, supportsBuiltinWebSearch: false, + supportsBuiltinCodeExecution: false, toolsEnabled: false, codeToolsEnabled: false, toolStatus: null, diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 1f370b6ac1..eb88da635d 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -15,6 +15,22 @@ export interface ThreadRecord { pairId?: string; archived: boolean; createdAt: number; + /** + * OpenAI shell tool container id captured from a prior response on + * this thread. When set, the next turn reuses it via + * `environment.type="container_reference"` so the model can read + * files it wrote earlier in the conversation. When null/undefined, + * the next turn auto-creates a fresh container. + * + * OpenAI containers expire after ~20 min of inactivity by default; + * if a stale id is sent, the backend surfaces an + * `_toolEvent.type="container_invalidated"` and the chat-adapter + * clears this field so the following turn falls back to auto-create. + * + * Anthropic's code-execution path doesn't need this — each turn + * gets a fresh container server-side. + */ + openaiCodeExecContainerId?: string | null; } export interface MessageRecord { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 18270a7620..ec4a5ea355 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -225,6 +225,16 @@ export interface OpenAIChatCompletionsRequest { encrypted_api_key?: string; provider_base_url?: string | null; enable_prompt_caching?: boolean | null; + /** + * OpenAI shell-tool container id captured from the prior response in + * this chat thread. When set and the Code pill is on, the backend + * routes the next /v1/responses call with + * `environment.type="container_reference"` so filesystem state + * persists across turns. Unset → backend uses + * `environment.type="container_auto"` and OpenAI creates a fresh + * container. Only meaningful for OpenAI cloud + gpt-5.5 family. + */ + openai_code_exec_container_id?: string | null; } export interface OpenAIChatDelta { From 85cf0a41ea6981a4e8526c01ede49a7d4b034b26 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 13:14:34 -0700 Subject: [PATCH 28/50] ci: switch Windows Stop Studio to a cmd no-op marker (#5462) The prior set +e + redirect + exit 0 fix in #5460 did not stop the Stop Studio step from exiting 143 (SIGTERM) on Git Bash; bash on windows-latest exits with that signal before any inline guard runs, regardless of redirection. The teardown does not gate correctness -- the runner reclaims the Studio child process at job end -- so swap the shell from Git Bash to cmd and just emit a marker line. After this, Job 3 (JSON, images) and the two other Windows GGUF CI jobs cannot fail at the teardown step. --- .../studio-windows-inference-smoke.yml | 48 ++++++++----------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 096ec95d03..188cdf5a19 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -345,15 +345,13 @@ jobs: - name: Stop Studio if: always() - # `set +e` + redirect everything: Git Bash on windows-latest - # has been observed to exit 143 from the kill/sleep block even - # when the upstream test work passed, masking a green run. The - # teardown does not gate correctness, so absorb any signal. - run: | - set +e - kill "${STUDIO_PID}" >/dev/null 2>&1 || true - sleep 2 >/dev/null 2>&1 || true - exit 0 + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Upload logs if: always() @@ -768,15 +766,13 @@ jobs: - name: Stop Studio if: always() - # `set +e` + redirect everything: Git Bash on windows-latest - # has been observed to exit 143 from the kill/sleep block even - # when the upstream test work passed, masking a green run. The - # teardown does not gate correctness, so absorb any signal. - run: | - set +e - kill "${STUDIO_PID}" >/dev/null 2>&1 || true - sleep 2 >/dev/null 2>&1 || true - exit 0 + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Upload logs if: always() @@ -1162,15 +1158,13 @@ jobs: - name: Stop Studio if: always() - # `set +e` + redirect everything: Git Bash on windows-latest - # has been observed to exit 143 from the kill/sleep block even - # when the upstream test work passed, masking a green run. The - # teardown does not gate correctness, so absorb any signal. - run: | - set +e - kill "${STUDIO_PID}" >/dev/null 2>&1 || true - sleep 2 >/dev/null 2>&1 || true - exit 0 + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Upload logs if: always() From 4b23af48b116785f95505a07f192e5f67618ca8e Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 14:18:04 -0700 Subject: [PATCH 29/50] tests: raise pwsh/bash subprocess timeout from 10s to 60s (#5463) CI surfaced a flaky failure on Linux 'Repo tests (CPU)': TestPwshPrForcePromotion.test_baked_in_pr_force_promotes -> subprocess.TimeoutExpired after 10s on /usr/bin/pwsh startup. The scripts under test run in well under a second; the 10s budget only covered pwsh / bash launch time, which spikes on heavily- loaded GitHub-hosted runners. Raise the default helper timeout to 60s for both run_bash and run_pwsh. Real bugs in the script logic will still surface as wrong output or non-zero exit; this just absorbs runner-side launch jitter. --- .../install/test_llama_pr_force_and_source.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 114680f458..2d7c038861 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -35,9 +35,11 @@ requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not availa # Helpers # --------------------------------------------------------------------------- def run_bash( - script: str, *, timeout: int = 10, env: dict | None = None + script: str, *, timeout: int = 60, env: dict | None = None ) -> subprocess.CompletedProcess: - """Run a bash script fragment and return the CompletedProcess.""" + """Run a bash script fragment and return the CompletedProcess. + 60s default tolerates slow shell startup on heavily-loaded CI + runners; the scripts themselves run in well under a second.""" run_env = os.environ.copy() if env: run_env.update(env) @@ -51,9 +53,12 @@ def run_bash( def run_pwsh( - script: str, *, timeout: int = 10, env: dict | None = None + script: str, *, timeout: int = 60, env: dict | None = None ) -> subprocess.CompletedProcess: - """Run a PowerShell script fragment and return the CompletedProcess.""" + """Run a PowerShell script fragment and return the CompletedProcess. + 60s default tolerates slow pwsh startup on heavily-loaded CI + runners; the scripts themselves run in well under a second. + A 10s budget previously surfaced as a flaky TimeoutExpired.""" run_env = os.environ.copy() run_env["NO_COLOR"] = "1" if env: From 4f59c8e539db39ce40da1609aaa8a01d960be297 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 14:44:52 -0700 Subject: [PATCH 30/50] studio/install: repair upstream llama.cpp prebuilt mangled symlinks (#5465) The macos-arm64 prebuilt tarball for llama.cpp b9165 and b9169 ships symlinks whose linkname is missing both the directory separator AND the leading character of the target basename: llama-b9165/libggml-rpc.0.dylib -> llama-b9165ibggml-rpc.0.11.1.dylib extract_tar_safely correctly classified those as unresolved and made install.sh fall back to source-build, which Mac CI then fails as a hard error (Studio must use the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon). Add _try_repair_missing_slash inside safe_link_target: when a linkname starts with the member's top-level dir but no following slash, search the archive for an entry under that dir whose name ends with the mangled suffix. Accept only when the suffix uniquely identifies a real archive entry, so legitimate archives are untouched. Verified against /tmp/llama-b9165.tar.gz: all 18 link entries repair to real files in the archive. --- studio/install_llama_prebuilt.py | 48 ++++++++++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 7c933bd612..89322c83ee 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3433,10 +3433,52 @@ def extract_archive(archive_path: Path, destination: Path) -> None: ) from exc return target + def _try_repair_missing_slash( + member_name: str, link_name: str, archive_names: set[str] + ) -> str | None: + """Some upstream llama.cpp Mac releases (e.g. b9165, b9169) ship + symlinks whose linkname is missing the directory separator AND + the leading character of the file basename between the + top-level dir and the rest of the path: + + llama-b9165/libggml-rpc.0.dylib -> llama-b9165ibggml-rpc.0.11.1.dylib + + That cannot be resolved as written. Detect the pattern + (linkname starts with the top-level dir name but no following + slash) and search archive entries under that dir for a real + file whose basename ends with the mangled suffix. Only accept + when the suffix uniquely identifies a real archive entry.""" + if "/" not in member_name or "/" in link_name: + return None + top, _, _ = member_name.partition("/") + if not link_name.startswith(top) or len(link_name) <= len(top): + return None + bad_suffix = link_name[len(top) :] + if not bad_suffix or bad_suffix.startswith("/"): + return None + prefix = f"{top}/" + candidates = [ + name + for name in archive_names + if name.startswith(prefix) + and "/" not in name[len(prefix) :] + and name[len(prefix) :].endswith(bad_suffix) + ] + if len(candidates) != 1: + return None + return candidates[0] + def safe_link_target( - base: Path, member_name: str, link_name: str, target: Path + base: Path, + member_name: str, + link_name: str, + target: Path, + archive_names: set[str], ) -> tuple[str, Path]: normalized = link_name.replace("\\", "/") + repaired = _try_repair_missing_slash(member_name, normalized, archive_names) + if repaired is not None: + normalized = repaired link_path = Path(normalized) if link_path.is_absolute(): raise PrebuiltFallback( @@ -3473,8 +3515,10 @@ def extract_archive(archive_path: Path, destination: Path) -> None: def extract_tar_safely(source: Path, base: Path) -> None: pending_links: list[tuple[tarfile.TarInfo, Path]] = [] + archive_names: set[str] = set() with tarfile.open(source, "r:gz") as archive: for member in archive.getmembers(): + archive_names.add(member.name) target = safe_extract_path(base, member.name) if member.isdir(): target.mkdir(parents = True, exist_ok = True) @@ -3501,7 +3545,7 @@ def extract_archive(archive_path: Path, destination: Path) -> None: progressed = False for member, target in unresolved: normalized_link, resolved_target = safe_link_target( - base, member.name, member.linkname, target + base, member.name, member.linkname, target, archive_names ) if not resolved_target.exists() and not resolved_target.is_symlink(): next_round.append((member, target)) From a70bf02bb88957d6b3617f9220ecf9fb5ae68528 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Sat, 16 May 2026 01:53:13 +0400 Subject: [PATCH 31/50] studio/chat: OpenAI container picker delete reliability (#5466) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio/chat: fix OpenAI container delete UX (expired filter, TTL cap, idempotent 404, refresh-on-error) - Filter status="expired" from /containers/list so the picker only shows usable containers. OpenAI keeps expired entries in the list indefinitely, which made delete look broken. - Cap ttl_minutes at 20 (backend Field + frontend TTL_MAX + persistence clamp). OpenAI's actual hard limit is 20; the prior 10080 cap caused integer_above_max_value rejections on create. - Treat 404 on delete as idempotent success in the frontend client so already-gone containers don't surface a scary error toast. - Run refresh() in finally for onCreate/onDelete so the picker stays in sync with OpenAI even when the call errors. - Add route-level test for the expired filter. * studio/chat: add diagnostic logging for OpenAI /containers DELETE Trace what arrives at /external/openai/containers/delete (subject, container_id, base_url) and what we send to OpenAI (URL, presence of Authorization, value of OpenAI-Beta) plus the full response status + body (capped at 300 chars). Helps confirm whether the beta header is on the wire and whether OpenAI's response actually reports deleted=true, when users report the delete "not taking". No secrets are logged — Authorization is reported as a boolean. * studio/chat: log raw /containers list response from OpenAI Sibling to the delete diagnostics. After a confirmed delete (deleted=true on the wire), we want to see whether the very next list call returns the just-deleted id — that distinguishes "OpenAI eventually-consistent list" from "frontend stale state". Logs each entry's id + status only; no names, no timestamps. * studio/chat: fingerprint decrypted API key for container CRUD Logs kind (sk-proj-/sk-/other), length, and last-4 chars only — never the full secret. Lets us compare what the backend actually uses against the key the user expects, since the same DELETE request shape can produce different results across keys (project-scoped containers: list is permissive but delete requires the owning project's key). * studio/chat: use fresh httpx client for /v1/containers DELETE Same key, same headers, same URL via the shared _http_client returned deleted=true but the container persisted in subsequent list calls. A fresh httpx.AsyncClient with the identical request shape (verified with a standalone reproducer) deleted the same container cleanly. Suspect connection-pool state from earlier chat-completion streams interferes at the edge — switching to a per-call client side-steps it entirely. Scoped to delete only; list/create keep using the shared pool until we can confirm the same fix is needed there. * studio/chat: log OpenAI response headers on container DELETE Adds cf-ray / x-request-id / openai-organization / openai-project / openai-processing-ms to the delete-response diagnostic line. Lets us cross-reference a failing delete against OpenAI support (or against a working standalone reproducer) using the unique request-id and edge node. * studio/chat: client-side tombstone for just-deleted OpenAI containers OpenAI's /v1/containers DELETE returns {"deleted": true} but the list endpoint can keep returning the same container for several minutes (replica lag or in-use silent no-op — undocumented per developers.openai.com/api/docs/guides/tools-shell). Our backend sends the correct DELETE with OpenAI-Beta: containers=v1 and a standalone reproducer shows the same behavior, so the right fix is UI-side rather than waiting on OpenAI. After a successful delete, the id goes into a per-component tombstone map with a 5-minute expiry. visibleContainers (now the single chokepoint feeding sortedContainers, auto-bind, and the all-containers list) filters those ids out. A 30s sweep clears expired tombstones so the picker recovers automatically if OpenAI eventually catches up (or the container's TTL elapses). * studio/chat: tombstones live for the page lifetime; drop API key fingerprint log - Tombstones change from Map<id, expiry> to Set<id>: once tombstoned, the id stays hidden from the picker until page reload. OpenAI's list can keep returning a deleted id for an undocumented and variable amount of time; automatically un-tombstoning after a fixed window surfaces it again and creates more confusion than it solves. The container's own TTL eventually expires the entry on OpenAI's side, and the expired-status filter at the backend list route hides it anyway. - Remove the periodic sweep effect (dead code without expiries). - Remove the api-key fingerprint log added during debugging — it served its purpose (confirmed parity) and isn't needed long-term. --- .../core/inference/external_provider.py | 43 ++++++++++++-- studio/backend/models/inference.py | 6 +- studio/backend/routes/inference.py | 30 +++++++++- .../tests/test_openai_container_crud.py | 38 ++++++++++++ .../features/chat/api/openai-containers.ts | 5 +- .../components/openai-code-exec-section.tsx | 58 +++++++++++++++---- .../src/features/chat/external-providers.ts | 2 +- 7 files changed, 159 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 79c8287e5b..de5f5c5500 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -2870,7 +2870,17 @@ class ExternalProviderClient: response.raise_for_status() data = response.json() containers = data.get("data") if isinstance(data, dict) else None - return list(containers) if isinstance(containers, list) else [] + result = list(containers) if isinstance(containers, list) else [] + logger.info( + "openai_container_list.response count=%s items=%s", + len(result), + [ + {"id": c.get("id"), "status": c.get("status")} + for c in result + if isinstance(c, dict) + ], + ) + return result async def create_openai_container( self, @@ -2901,15 +2911,38 @@ class ExternalProviderClient: async def delete_openai_container(self, container_id: str) -> None: """DELETE /v1/containers/{id}. 404s are surfaced as HTTPError. + Uses a fresh httpx client (not the shared ``_http_client``) so + connection-pool state from earlier chat requests cannot + interfere — observed in the wild that DELETEs over the shared + pool returned ``deleted: true`` while the container persisted + in subsequent /containers list calls, even though the same + DELETE issued from a fresh client genuinely removed it. + Verifies the response body reports ``deleted: true``. OpenAI returns a 2xx ``deleted: true`` body even when the request is silently rejected (e.g. missing OpenAI-Beta header), so a status-only check is not sufficient. """ - response = await _http_client.delete( - f"{self.base_url}/containers/{container_id}", - headers = self._container_headers(), - timeout = self._timeout, + url = f"{self.base_url}/containers/{container_id}" + headers = self._container_headers() + logger.info( + "openai_container_delete.outbound url=%s has_auth=%s openai_beta=%s", + url, + "Authorization" in headers, + headers.get("OpenAI-Beta"), + ) + async with httpx.AsyncClient(timeout = self._timeout) as fresh_client: + response = await fresh_client.delete(url, headers = headers) + logger.info( + "openai_container_delete.response status=%s cf_ray=%s " + "request_id=%s organization=%s project=%s processing_ms=%s body=%s", + response.status_code, + response.headers.get("cf-ray"), + response.headers.get("x-request-id"), + response.headers.get("openai-organization"), + response.headers.get("openai-project"), + response.headers.get("openai-processing-ms"), + response.text[:300], ) response.raise_for_status() try: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 6a042d35d7..3aa89cc934 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -650,11 +650,11 @@ class CreateOpenAIContainerBody(OpenAIContainerRequest): ttl_minutes: int = Field( 20, ge = 1, - le = 10080, # 1 week + le = 20, description = ( "Idle-timeout TTL the new container will inherit (anchor=" - "last_active_at). OpenAI's default is 20; we cap at one " - "week as a safety bound." + "last_active_at). OpenAI hard-caps this at 20 minutes and " + "rejects larger values with integer_above_max_value." ), ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8f5bd9aab5..76bbb59c94 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1716,8 +1716,15 @@ async def list_openai_containers( status_code = 502, detail = f"Failed to reach OpenAI: {exc}", ) + # OpenAI keeps expired containers in /v1/containers indefinitely + # with status="expired" — they're effectively dead but still + # listed. Hide them so the picker only shows usable containers. return ListOpenAIContainersResponse( - containers = [_summarize_container(c) for c in raw if isinstance(c, dict)], + containers = [ + _summarize_container(c) + for c in raw + if isinstance(c, dict) and c.get("status") != "expired" + ], ) finally: await client.close() @@ -1766,17 +1773,38 @@ async def delete_openai_container( current_subject: str = Depends(get_current_subject), ) -> None: """Delete a named container by id.""" + logger.info( + "openai_container_delete.request subject=%s container_id=%s base_url=%s", + current_subject, + body.container_id, + body.provider_base_url, + ) client = _resolve_openai_cloud_client(body) try: try: await client.delete_openai_container(body.container_id) + logger.info( + "openai_container_delete.success container_id=%s", + body.container_id, + ) except httpx.HTTPStatusError as exc: detail = exc.response.text[:500] if exc.response is not None else str(exc) + logger.warning( + "openai_container_delete.openai_rejected container_id=%s status=%s body=%s", + body.container_id, + exc.response.status_code if exc.response else None, + detail, + ) raise HTTPException( status_code = exc.response.status_code if exc.response else 502, detail = f"OpenAI rejected /containers delete: {detail}", ) except httpx.HTTPError as exc: + logger.warning( + "openai_container_delete.transport_error container_id=%s error=%s", + body.container_id, + exc, + ) raise HTTPException( status_code = 502, detail = f"Failed to reach OpenAI: {exc}", diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py index 2965ec6649..fbc0393677 100644 --- a/studio/backend/tests/test_openai_container_crud.py +++ b/studio/backend/tests/test_openai_container_crud.py @@ -149,3 +149,41 @@ def test_delete_propagates_openai_4xx(monkeypatch): with pytest.raises(httpx.HTTPStatusError): _drive(_make_client().delete_openai_container("cntr_missing")) + + +def test_list_route_filters_expired_containers(monkeypatch): + """OpenAI keeps containers in /v1/containers indefinitely with + status="expired" after their idle TTL passes — they can't be + used but still show up. The list route must drop them so the + picker only surfaces usable containers.""" + from routes import inference as inf_mod + from models.inference import OpenAIContainerRequest + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json = { + "data": [ + {"id": "cntr_active", "name": "live", "status": "running"}, + {"id": "cntr_dead", "name": "old", "status": "expired"}, + {"id": "cntr_unknown", "name": "no-status"}, + ], + }, + ) + + _mock_http_client(monkeypatch, handler) + + def fake_resolve(_body): + return _make_client() + + monkeypatch.setattr(inf_mod, "_resolve_openai_cloud_client", fake_resolve) + + body = OpenAIContainerRequest( + encrypted_api_key = "enc", + provider_base_url = "https://api.openai.com/v1", + ) + response = _drive(inf_mod.list_openai_containers(body, current_subject = "u")) + ids = [c.id for c in response.containers] + assert "cntr_active" in ids + assert "cntr_unknown" in ids # missing status is treated as usable + assert "cntr_dead" not in ids diff --git a/studio/frontend/src/features/chat/api/openai-containers.ts b/studio/frontend/src/features/chat/api/openai-containers.ts index ca311cc311..29d292f7cf 100644 --- a/studio/frontend/src/features/chat/api/openai-containers.ts +++ b/studio/frontend/src/features/chat/api/openai-containers.ts @@ -115,7 +115,10 @@ export async function deleteOpenAIContainer( }), }, ); - if (!response.ok && response.status !== 204) { + // 404 = container already gone (deleted elsewhere, or expired-then-purged). + // Treat as idempotent success so a stale list entry doesn't surface as a + // confusing error — the caller will refresh and the entry will disappear. + if (!response.ok && response.status !== 204 && response.status !== 404) { throw new Error(await parseError(response)); } } diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx index 15cb79286c..2d3d225069 100644 --- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -49,7 +49,7 @@ import { ensureThreadRecord } from "../runtime-provider"; const AUTO_OPTION_VALUE = "__auto__"; const DEFAULT_TTL_MINUTES = 20; const TTL_MIN = 1; -const TTL_MAX = 10080; // one week — matches backend bound +const TTL_MAX = 20; // OpenAI hard cap on expires_after.minutes function ageLabel(epochSeconds: number | null | undefined): string { if (!epochSeconds) return ""; @@ -84,6 +84,13 @@ export function OpenAICodeExecSection({ const [createTtl, setCreateTtl] = useState<number>( provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES, ); + // Ids that have been deleted in this session. Once tombstoned, an id + // stays hidden from the picker for the lifetime of the page — OpenAI's + // /containers list can keep returning a freshly-deleted id for an + // undocumented and variable amount of time, and an automatic re-show + // creates more confusion than it solves. Refreshing the page resets + // the tombstone naturally. + const [tombstones, setTombstones] = useState<Set<string>>(() => new Set()); const thread = useLiveQuery( async () => (activeThreadId ? db.threads.get(activeThreadId) : undefined), @@ -91,14 +98,22 @@ export function OpenAICodeExecSection({ ); const activeContainerId = thread?.openaiCodeExecContainerId ?? null; + // Hide just-deleted containers even if OpenAI's list still returns them. + // This is the single chokepoint — every downstream view (sorted picker, + // auto-bind candidate, all-containers list) derives from visibleContainers. + const visibleContainers = useMemo(() => { + if (tombstones.size === 0) return containers; + return containers.filter((c) => !tombstones.has(c.id)); + }, [containers, tombstones]); + // Containers sorted newest-first by lastActiveAt so the dropdown's // default (auto-bind target) shows up first. const sortedContainers = useMemo( () => - [...containers].sort( + [...visibleContainers].sort( (a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0), ), - [containers], + [visibleContainers], ); // What the dropdown should display right now. We decouple this from @@ -150,10 +165,14 @@ export function OpenAICodeExecSection({ // the chat-adapter's lazy-create path will mint the first container // on first send. useEffect(() => { - if (!activeThreadId || activeContainerId || containers.length === 0) { + if ( + !activeThreadId || + activeContainerId || + visibleContainers.length === 0 + ) { return; } - const sorted = [...containers].sort( + const sorted = [...visibleContainers].sort( (a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0), ); const candidate = sorted[0]; @@ -171,7 +190,7 @@ export function OpenAICodeExecSection({ // Best-effort; the chat-adapter will inherit/create on send. } })(); - }, [activeThreadId, activeContainerId, containers]); + }, [activeThreadId, activeContainerId, visibleContainers]); const ttlValue = provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES; @@ -223,7 +242,6 @@ export function OpenAICodeExecSection({ toast.success(`Created container ${name}`); setCreateName(""); setCreateOpen(false); - await refresh(); // Auto-bind the just-created container to the active thread. // ensureThreadRecord first so the bind lands even when the user // creates a container before sending the first message — without @@ -250,6 +268,10 @@ export function OpenAICodeExecSection({ ); } finally { setCreating(false); + // Refresh even on failure: the request may have partially succeeded + // server-side (created container, lost response), and a re-fetch + // keeps the picker in sync with OpenAI's actual state. + await refresh(); } }; @@ -267,6 +289,14 @@ export function OpenAICodeExecSection({ { apiKey, baseUrl: provider.baseUrl || null }, id, ); + // Tombstone the id so the picker hides it immediately even if + // OpenAI's list keeps returning it for a while. + setTombstones((prev) => { + if (prev.has(id)) return prev; + const next = new Set(prev); + next.add(id); + return next; + }); // Clear any thread bindings pointing at the now-deleted id. const affected = await db.threads .filter((t) => t.openaiCodeExecContainerId === id) @@ -277,11 +307,15 @@ export function OpenAICodeExecSection({ ), ); toast.success(`Deleted container ${name || id}`); - await refresh(); } catch (err) { toast.error( `Delete failed: ${err instanceof Error ? err.message : "Unknown"}`, ); + } finally { + // Always refresh so a stale list entry (e.g. container deleted + // elsewhere, or already expired) is purged from the UI even when + // the delete call itself errored. + await refresh(); } }; @@ -293,7 +327,7 @@ export function OpenAICodeExecSection({ htmlFor="openai-container-ttl" className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg" > - New-container idle timeout (min) + New-container idle timeout (min, max 20) </label> <Input id="openai-container-ttl" @@ -360,11 +394,11 @@ export function OpenAICodeExecSection({ <span className="text-[11px] uppercase tracking-wider text-muted-foreground"> All containers </span> - {isLoading && containers.length === 0 ? ( + {isLoading && visibleContainers.length === 0 ? ( <Skeleton className="h-16 w-full" /> - ) : containers.length > 0 ? ( + ) : visibleContainers.length > 0 ? ( <ul className="flex flex-col gap-1 max-h-44 overflow-auto"> - {containers.map((c) => { + {visibleContainers.map((c) => { const isActive = c.id === activeContainerId; return ( <li diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 89591d7b14..d455ff138e 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -237,7 +237,7 @@ function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig providerType === "openai" && typeof raw.openaiContainerTtlMinutes === "number" && raw.openaiContainerTtlMinutes >= 1 - ? Math.min(raw.openaiContainerTtlMinutes, 10080) + ? Math.min(raw.openaiContainerTtlMinutes, 20) : undefined, }; } From 2de99a23d85077e7583e87ac7a79afa8383a4636 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 15:09:50 -0700 Subject: [PATCH 32/50] studio/install: strip top-level dir from repaired symlink target (#5467) The repair in 5465 returned the full archive entry name (e.g. "llama-b9165 libggml-rpc.0.11.1.dylib") but safe_link_target joins the return value with target.parent (which already lives under base llama-b9165). That doubled the prefix to base llama-b9165 llama-b9165 libggml-rpc.0.11.1.dylib, the resolved path never existed, and extract_tar_safely still raised 'tar archive contained unresolved link entries'. Strip the top-level dir before returning so the linkname is relative to target.parent, mirroring how unmangled symlinks are stored in the tar (basename-only relative to the symlink). Verified end-to-end against the upstream b9165 tarball: extraction succeeds and every symlink resolves to an existing file. --- studio/install_llama_prebuilt.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 89322c83ee..ac1d2aded8 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3447,7 +3447,11 @@ def extract_archive(archive_path: Path, destination: Path) -> None: (linkname starts with the top-level dir name but no following slash) and search archive entries under that dir for a real file whose basename ends with the mangled suffix. Only accept - when the suffix uniquely identifies a real archive entry.""" + when the suffix uniquely identifies a real archive entry. + Returns the corrected linkname expressed relative to the + member's parent directory -- callers join it with + `target.parent`, so a full `top/file` path would double the + prefix into `top/top/file`.""" if "/" not in member_name or "/" in link_name: return None top, _, _ = member_name.partition("/") @@ -3466,7 +3470,10 @@ def extract_archive(archive_path: Path, destination: Path) -> None: ] if len(candidates) != 1: return None - return candidates[0] + # Strip the top-level dir so the caller's `target.parent / Path(...)` + # composition resolves inside the staging dir, not into a duplicate + # `top/top/...` path. + return candidates[0][len(prefix) :] def safe_link_target( base: Path, From ba0cae1aff53ba3c0348887918c8e2136822454a Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 15 May 2026 23:17:03 +0100 Subject: [PATCH 33/50] Stop: drop Ollama API key, clean up code execution UI (#5464) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chat: drop Ollama API key, clean up code execution UI * studio/chat: fix undefined candidateId + keyboard a11y on container list - Auto-bind effect referenced `candidateId`, which is not declared in this scope (only `candidate` is) — would fail the TS/Next build. Use `candidate.id` to match the variable that's actually defined. - Container list items get `role="button"` when `canActivate` is true but had no keyboard activation. Add `onKeyDown` for Enter/Space and `tabIndex={0}` so the row is focusable and activatable from the keyboard, matching the existing onClick behavior. * studio/chat: restore declarations dropped by the main merge The 75646444d auto-merge with main (#5466) silently dropped the declarations a4f19171c added in regions #5466 also rewrote, while leaving the usages further down in the file. No textual conflict markers, but the result referenced undeclared names: - REFRESH_POLL_MS constant (drives the 30s list refresh interval). - pendingDelete / setPendingDelete / deleting / setDeleting state (drives the in-sheet AlertDialog delete confirm — replaces the window.confirm() that landed via #5466). - Per-row locals inside the container list .map callback: running, isActive (recomputed with running), ttlMinutes, canActivate, statusLabel (drive click-to-activate, expired/active badges, and the muted styling for expired containers). Also wire setDeleting(false) + setPendingDelete(null) into the confirmDelete finally so the AlertDialog closes after the delete call resolves; previously the busy state never cleared. The all-containers list now iterates sortedContainers (matches the picker above and the "newest-active first" UX) instead of the unsorted visibleContainers. --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai> --- .../features/chat/chat-providers-dialog.tsx | 82 +++--- .../components/openai-code-exec-section.tsx | 242 ++++++++++++++---- 2 files changed, 241 insertions(+), 83 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 0127369eb4..ea99ada0e0 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -188,6 +188,11 @@ export function ChatProvidersSettings({ const [isReasoningModel, setIsReasoningModel] = useState(false); const reduceMotion = useReducedMotion(); const isCustomProvider = isCustomProviderType(providerType); + // Ollama runs locally and does not require an API key. Hide the input + // entirely rather than just marking it optional so users aren't prompted + // for a credential the provider never uses. + const isOllamaProvider = providerType === "ollama"; + const showApiKeyField = !isOllamaProvider; const showReasoningToggle = supportsProviderReasoningToggle(providerType); const registryByType = useMemo( @@ -734,7 +739,10 @@ export function ChatProvidersSettings({ async function testProvider(provider: ExternalProviderConfig) { const savedKey = getExternalProviderApiKey(provider.id).trim(); - if (!savedKey) { + // Ollama runs locally and never requires a key — fall through to the + // real connection check instead of prompting for credentials the form + // no longer exposes. + if (!savedKey && provider.providerType !== "ollama") { if (isCustomProviderType(provider.providerType)) { await editProvider(provider); toast.info(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE); @@ -884,42 +892,44 @@ export function ChatProvidersSettings({ </Select> </div> - <div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1"> - <div className="flex min-w-0 flex-col gap-0.5"> - <Label - htmlFor="provider-api-key" - className="text-sm font-medium" - > - API key {isCustomProvider ? "(optional)" : ""} - </Label> - <p className="text-xs leading-snug text-muted-foreground"> - Stored locally. - </p> + {showApiKeyField ? ( + <div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1"> + <div className="flex min-w-0 flex-col gap-0.5"> + <Label + htmlFor="provider-api-key" + className="text-sm font-medium" + > + API key {isCustomProvider ? "(optional)" : ""} + </Label> + <p className="text-xs leading-snug text-muted-foreground"> + Stored locally. + </p> + </div> + <div className="relative min-w-0"> + <Input + id="provider-api-key" + type={showApiKey ? "text" : "password"} + value={apiKey} + onChange={(event) => setApiKey(event.target.value)} + placeholder="Enter API key" + className="h-9 pr-9 text-sm" + /> + <button + type="button" + onClick={() => setShowApiKey((visible) => !visible)} + className="absolute top-1/2 right-1.5 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" + aria-label={showApiKey ? "Hide API key" : "Show API key"} + aria-pressed={showApiKey} + > + {showApiKey ? ( + <Eye className="size-3.5" /> + ) : ( + <EyeOff className="size-3.5" /> + )} + </button> + </div> </div> - <div className="relative min-w-0"> - <Input - id="provider-api-key" - type={showApiKey ? "text" : "password"} - value={apiKey} - onChange={(event) => setApiKey(event.target.value)} - placeholder="Enter API key" - className="h-9 pr-9 text-sm" - /> - <button - type="button" - onClick={() => setShowApiKey((visible) => !visible)} - className="absolute top-1/2 right-1.5 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label={showApiKey ? "Hide API key" : "Show API key"} - aria-pressed={showApiKey} - > - {showApiKey ? ( - <Eye className="size-3.5" /> - ) : ( - <EyeOff className="size-3.5" /> - )} - </button> - </div> - </div> + ) : null} {isCustomProvider ? ( <div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1"> diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx index 2d3d225069..d76b4a282f 100644 --- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -31,6 +31,16 @@ import { useCallback, useEffect, useMemo, useState } from "react"; import { toast } from "sonner"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Skeleton } from "@/components/ui/skeleton"; @@ -50,6 +60,11 @@ const AUTO_OPTION_VALUE = "__auto__"; const DEFAULT_TTL_MINUTES = 20; const TTL_MIN = 1; const TTL_MAX = 20; // OpenAI hard cap on expires_after.minutes +// Cadence for re-fetching the container list while the section is +// mounted. OpenAI's container TTL flips at minute granularity, so 30s +// is fast enough that an expired container loses its ACTIVE pill within +// half a minute without hammering /v1/containers. +const REFRESH_POLL_MS = 30_000; function ageLabel(epochSeconds: number | null | undefined): string { if (!epochSeconds) return ""; @@ -63,6 +78,21 @@ function ageLabel(epochSeconds: number | null | undefined): string { return `${ageDay}d ago`; } +function shortContainerId(id: string): string { + // Mid-truncate keeps the "cntr_" prefix readable and still surfaces the + // tail digits users sometimes copy off OpenAI's dashboard. + if (id.length <= 18) return id; + return `${id.slice(0, 12)}…${id.slice(-4)}`; +} + +function isContainerRunning(c: OpenAIContainerSummary): boolean { + // OpenAI's containers API reports `status: "running"` while idle TTL is + // valid and `status: "expired"` once the idle window has passed. Treat + // a missing status as running so we don't false-positive on any older + // payloads that didn't include the field. + return c.status == null || c.status === "running"; +} + interface OpenAICodeExecSectionProps { provider: ExternalProviderConfig; apiKey: string | null; @@ -91,6 +121,12 @@ export function OpenAICodeExecSection({ // creates more confusion than it solves. Refreshing the page resets // the tombstone naturally. const [tombstones, setTombstones] = useState<Set<string>>(() => new Set()); + // Target row for the destructive confirmation dialog. Held in state + // (rather than blocking with window.confirm) so the dialog sits inside + // the settings sheet instead of a native browser alert. + const [pendingDelete, setPendingDelete] = + useState<OpenAIContainerSummary | null>(null); + const [deleting, setDeleting] = useState(false); const thread = useLiveQuery( async () => (activeThreadId ? db.threads.get(activeThreadId) : undefined), @@ -116,15 +152,27 @@ export function OpenAICodeExecSection({ [visibleContainers], ); - // What the dropdown should display right now. We decouple this from - // `activeContainerId` (which is whatever is in Dexie) so the user - // immediately sees the most-recent container by name when there is - // no thread binding yet, rather than a "Selecting most recent…" - // placeholder while the auto-bind effect's async write propagates - // back through useLiveQuery. The auto-bind effect still writes the - // bind to Dexie so the chat adapter sees it on send. + // First running container by lastActiveAt — the auto-bind target and + // also what we surface visually before Dexie catches up. + const firstRunningContainer = useMemo( + () => sortedContainers.find(isContainerRunning) ?? null, + [sortedContainers], + ); + + // What the picker should treat as "active" right now. We decouple + // this from `activeContainerId` (Dexie state) so the user immediately + // sees the most-recent running container while the auto-bind effect's + // async write propagates. If the Dexie-bound container has since + // expired, fall back to the first running candidate — the stale-bind + // sweeper below will clear Dexie shortly after. + const boundContainer = useMemo( + () => sortedContainers.find((c) => c.id === activeContainerId) ?? null, + [sortedContainers, activeContainerId], + ); const displayedContainerId = - activeContainerId ?? sortedContainers[0]?.id ?? null; + (boundContainer && isContainerRunning(boundContainer) + ? boundContainer.id + : firstRunningContainer?.id) ?? null; const refresh = useCallback(async () => { if (!apiKey) return; @@ -144,9 +192,28 @@ export function OpenAICodeExecSection({ } }, [apiKey, provider.baseUrl]); - // Fetch once when the section mounts (or provider changes). + // Fetch once when the section mounts (or provider changes), then + // poll on a low cadence so an expired container's ACTIVE pill clears + // without the user clicking the refresh button. Also re-fetch when + // the tab regains visibility — covers the common case of leaving the + // sheet open across a long idle period. useEffect(() => { void refresh(); + const interval = window.setInterval(() => { + if (document.visibilityState === "visible") { + void refresh(); + } + }, REFRESH_POLL_MS); + const onVisibility = () => { + if (document.visibilityState === "visible") { + void refresh(); + } + }; + document.addEventListener("visibilitychange", onVisibility); + return () => { + window.clearInterval(interval); + document.removeEventListener("visibilitychange", onVisibility); + }; }, [refresh]); // Auto-bind the active thread to the most-recently-active container @@ -275,15 +342,10 @@ export function OpenAICodeExecSection({ } }; - const onDelete = async (id: string, name: string | null | undefined) => { - if (!apiKey) return; - if ( - !window.confirm( - `Delete container ${name || id}? Threads using it will fall back to auto-create on their next turn.`, - ) - ) { - return; - } + const confirmDelete = async () => { + if (!apiKey || !pendingDelete) return; + const { id, name } = pendingDelete; + setDeleting(true); try { await deleteOpenAIContainer( { apiKey, baseUrl: provider.baseUrl || null }, @@ -312,6 +374,8 @@ export function OpenAICodeExecSection({ `Delete failed: ${err instanceof Error ? err.message : "Unknown"}`, ); } finally { + setDeleting(false); + setPendingDelete(null); // Always refresh so a stale list entry (e.g. container deleted // elsewhere, or already expired) is purged from the UI even when // the delete call itself errored. @@ -319,6 +383,8 @@ export function OpenAICodeExecSection({ } }; + const displayActiveId = displayedContainerId; + return ( <div className="flex flex-col gap-3 pt-1"> {/* TTL */} @@ -336,23 +402,23 @@ export function OpenAICodeExecSection({ max={TTL_MAX} value={ttlValue} onChange={(e) => onTtlChange(e.target.value)} - className="h-8 w-24 text-sm" + className="h-8 w-14 px-2 text-center text-sm tabular-nums" /> </div> - {/* Active container picker — visually emphasized so it reads as - the primary control vs. the static list below. Accent - background + ring outline distinguish it from the plain - bordered list items beneath. */} - <div className="flex flex-col gap-1.5 rounded-md border border-primary/30 bg-primary/5 p-2.5"> + {/* Single container list. The previously-separate "Active for + this thread" picker collapses into this list: clicking a row + binds it to the active thread, and the ACTIVE pill marks + which one. Avoids duplicating state across two controls. */} + <div className="flex flex-col gap-1.5"> <div className="flex items-center justify-between gap-2"> - <span className="text-[13px] font-semibold leading-[1.25] tracking-nav text-primary"> - Active for this thread + <span className="text-[11px] uppercase tracking-wider text-muted-foreground"> + Containers </span> <Button size="sm" variant="ghost" - className="h-7 px-2" + className="-mr-1 h-6 w-6 p-0 text-muted-foreground" onClick={() => void refresh()} disabled={isLoading || !apiKey} aria-label="Refresh container list" @@ -396,38 +462,84 @@ export function OpenAICodeExecSection({ </span> {isLoading && visibleContainers.length === 0 ? ( <Skeleton className="h-16 w-full" /> - ) : visibleContainers.length > 0 ? ( - <ul className="flex flex-col gap-1 max-h-44 overflow-auto"> - {visibleContainers.map((c) => { - const isActive = c.id === activeContainerId; + ) : sortedContainers.length > 0 ? ( + <ul className="flex max-h-52 flex-col gap-1 overflow-auto"> + {sortedContainers.map((c) => { + const running = isContainerRunning(c); + const isActive = running && c.id === displayActiveId; + const ttlMinutes = c.expiresAfterMinutes ?? DEFAULT_TTL_MINUTES; + const canActivate = + activeThreadId != null && !isActive && running; + const statusLabel = !running ? (c.status ?? "expired") : null; return ( <li key={c.id} - className={`flex items-center justify-between gap-2 rounded-md border px-2 py-1.5 text-xs ${ + className={`flex items-center gap-2 rounded-md border px-2 py-1.5 text-xs transition-colors ${ isActive ? "border-primary/30 bg-primary/5" - : "border-border/60" + : "border-border/60 hover:bg-muted/40" + } ${canActivate ? "cursor-pointer" : ""} ${ + running ? "" : "opacity-60" }`} + onClick={() => { + if (canActivate) void onPick(c.id); + }} + onKeyDown={(e) => { + if (!canActivate) return; + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + void onPick(c.id); + } + }} + tabIndex={canActivate ? 0 : undefined} + role={canActivate ? "button" : undefined} + aria-pressed={isActive} + title={ + canActivate + ? "Use this container for the active thread" + : !running + ? `Container is ${statusLabel}` + : undefined + } > - <div className="flex min-w-0 flex-col"> - <span className="truncate font-medium"> - {c.name ?? "(unnamed)"} + {/* min-w-0 + truncate keeps long OpenAI container ids + from spilling under the trash button on narrow + settings sheets. */} + <div className="flex min-w-0 flex-1 flex-col gap-0.5"> + <div className="flex min-w-0 items-center gap-1.5"> + <span className="min-w-0 truncate font-medium"> + {c.name ?? "(unnamed)"} + </span> {isActive ? ( - <span className="ml-1.5 text-[10px] font-normal uppercase tracking-wider text-primary"> - · active + <span className="shrink-0 rounded-sm bg-primary/15 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-primary"> + Active + </span> + ) : statusLabel ? ( + <span className="shrink-0 rounded-sm bg-muted px-1 py-px text-[9px] font-medium uppercase tracking-wider text-muted-foreground"> + {statusLabel} </span> ) : null} - </span> - <span className="text-muted-foreground"> - {c.id} · TTL{" "} - {c.expiresAfterMinutes ?? DEFAULT_TTL_MINUTES}m - </span> + </div> + <div + className="flex min-w-0 items-center gap-1.5 text-muted-foreground" + title={c.id} + > + <span className="min-w-0 truncate font-mono text-[11px]"> + {shortContainerId(c.id)} + </span> + <span className="shrink-0 text-[10px] uppercase tracking-wider"> + · {ttlMinutes}m + </span> + </div> </div> <Button size="sm" variant="ghost" - className="h-6 w-6 p-0 text-destructive" - onClick={() => void onDelete(c.id, c.name)} + className="h-6 w-6 shrink-0 p-0 text-muted-foreground hover:text-destructive" + onClick={(e) => { + e.stopPropagation(); + setPendingDelete(c); + }} aria-label={`Delete container ${c.name ?? c.id}`} > <TrashIcon className="size-3.5" /> @@ -438,8 +550,7 @@ export function OpenAICodeExecSection({ </ul> ) : ( <p className="text-xs text-muted-foreground"> - No saved containers yet. Use auto-create or create a named - one below. + None yet — one will be created on first send. </p> )} </div> @@ -506,6 +617,43 @@ export function OpenAICodeExecSection({ New container </Button> )} + + <AlertDialog + open={pendingDelete !== null} + onOpenChange={(nextOpen) => { + if (!nextOpen && deleting) return; + if (!nextOpen) setPendingDelete(null); + }} + > + <AlertDialogContent size="sm"> + <AlertDialogHeader> + <AlertDialogTitle> + Delete{" "} + <span className="font-mono"> + {pendingDelete?.name ?? "container"} + </span> + ? + </AlertDialogTitle> + <AlertDialogDescription> + Threads using this container will fall back to auto-create on + their next turn. This cannot be undone. + </AlertDialogDescription> + </AlertDialogHeader> + <AlertDialogFooter> + <AlertDialogCancel disabled={deleting}>Cancel</AlertDialogCancel> + <AlertDialogAction + variant="destructive" + disabled={deleting} + onClick={(e) => { + e.preventDefault(); + void confirmDelete(); + }} + > + {deleting ? "Deleting…" : "Delete"} + </AlertDialogAction> + </AlertDialogFooter> + </AlertDialogContent> + </AlertDialog> </div> ); } From e775f941a42ec6daf06ed95ad94d5adc7bb37247 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 15:53:54 -0700 Subject: [PATCH 34/50] tests/openai: patch httpx.AsyncClient ctor so delete tests hit mock (#5469) delete_openai_container intentionally creates a fresh httpx.AsyncClient per call (see external_provider docstring: shared pool produced false 'deleted: true' responses while the container survived). The existing _mock_http_client only swapped the shared module-level _http_client, so the four delete tests bypassed the mock entirely and hit the real OpenAI API, returning 401 Unauthorized on Python 3.10 / 3.12 / 3.13. Extend the helper to also monkey-patch httpx.AsyncClient itself to a factory that injects the test's MockTransport into any freshly constructed client. List/create paths still use the shared client and pass unchanged. Verified locally: pytest tests/test_openai_container_crud.py -> 8 passed. --- studio/backend/tests/test_openai_container_crud.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py index fbc0393677..161a6fab83 100644 --- a/studio/backend/tests/test_openai_container_crud.py +++ b/studio/backend/tests/test_openai_container_crud.py @@ -28,8 +28,20 @@ def _drive(coro): def _mock_http_client(monkeypatch, handler): + """Wire `handler` for both the shared `_http_client` AND any + per-call `httpx.AsyncClient(...)` instances. delete_openai_container + intentionally creates a fresh AsyncClient (see comment in + external_provider.delete_openai_container) so the test must + also intercept that constructor.""" transport = httpx.MockTransport(handler) monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + real_async_client = httpx.AsyncClient + + def _patched_async_client(*args, **kwargs): + kwargs["transport"] = transport + return real_async_client(*args, **kwargs) + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", _patched_async_client) def _make_client() -> ExternalProviderClient: From 77e09297353c36dac34af6f1f5688611d80902d0 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 19:41:09 -0700 Subject: [PATCH 35/50] revert: stop touching DEVICE_TYPE == "cuda" branches for CPU CI (#5473) #5429 (cb15a7a5) tightened three production-path branches to DEVICE_TYPE == "cuda" and torch.cuda.is_available() and added a new else: SUPPORTS_BFLOAT16 = False arm to let `import unsloth.trainer` survive on a CPU-only CI host. We already ship the package on Intel XPU / AMD HIP / NVIDIA CUDA and don't want any extra branching in those hot paths. Move the entire CPU-CI handling to one place -- the top of unsloth/device_type.get_device_type() -- so the UNSLOTH_ALLOW_CPU=1 sentinel short-circuits detection and returns "cuda" before any torch probe runs. Every downstream DEVICE_TYPE == "cuda" branch then behaves identically to a real CUDA host, with no additional checks. The two existing duplicate UNSLOTH_ALLOW_CPU returns later in the function are dropped (the new top-of-function check covers both). Revert the three call-site changes: - unsloth/_gpu_init.py:212 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/_gpu_init.py:247 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/models/_utils.py:1207 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/_gpu_init.py: drop the new `else: SUPPORTS_BFLOAT16 = False` branch (dead under the top-of-function short-circuit). Keep the two env-var gates that are needed for zoo's drift detectors to inspect pristine TRL source (no behavioural change on production hosts that never set UNSLOTH_ALLOW_CPU): - unsloth/_gpu_init.py: `if env != "1": _patch_trl_trainer()` - unsloth/models/rl.py:PatchFastRL: `if env == "1": return` Verified: - CUDA_VISIBLE_DEVICES=5 python -c "import unsloth.trainer" produces UnslothSFTTrainer.__init__ (TRL still patched on real hosts). - UNSLOTH_ALLOW_CPU=1 + aggressive cuda spoof import succeeds and trl.SFTTrainer.__init__.__qualname__ stays SFTTrainer.__init__. - pytest tests/_zoo_compiler_cache_shim.py -> 5 passed, 1 skipped. --- unsloth/_gpu_init.py | 10 ++-------- unsloth/device_type.py | 13 +++++++------ unsloth/models/_utils.py | 2 +- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index a30111b529..df446195fb 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -209,7 +209,7 @@ del fix_peft_transformers_weight_conversion_import del patch_peft_weight_converter_compatibility # Torch 2.4 has including_emulation -if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): +if DEVICE_TYPE == "cuda": major_version, minor_version = torch.cuda.get_device_capability() SUPPORTS_BFLOAT16 = major_version >= 8 @@ -233,18 +233,12 @@ elif DEVICE_TYPE == "xpu": # torch.xpu.is_bf16_supported() does not have including_emulation # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported() -else: - # CPU-only CI under UNSLOTH_ALLOW_CPU=1. We can't probe device - # capability, so assume no bf16 -- training won't run on this host - # anyway, this branch only exists to let `import unsloth.trainer` - # succeed for source-inspection tests. - SUPPORTS_BFLOAT16 = False # For Gradio HF Spaces? # if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ: import triton -if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): +if DEVICE_TYPE == "cuda": libcuda_dirs = lambda: None if Version(triton.__version__) >= Version("3.0.0"): try: diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 6a82e42e8c..f7a330594b 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -52,6 +52,13 @@ def is_hip(): @functools.cache def get_device_type(): + # Test-only CPU fallback. Short-circuits the detection chain so the + # rest of the function -- and every DEVICE_TYPE == "cuda" branch in + # the codebase -- behaves identically to a real CUDA host. The env + # var is read exactly once per process because get_device_type is + # @functools.cache'd, so production hosts pay no runtime cost. + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return "cuda" if _IS_MLX: return "mlx" if hasattr(torch, "cuda") and torch.cuda.is_available(): @@ -63,10 +70,6 @@ def get_device_type(): # Check torch.accelerator if hasattr(torch, "accelerator"): if not torch.accelerator.is_available(): - # Test-only CPU fallback. The env var is read exactly once per - # process because get_device_type is @functools.cache'd. - if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": - return "cuda" raise NotImplementedError( "Unsloth cannot find any torch accelerator? You need a GPU." ) @@ -77,8 +80,6 @@ def get_device_type(): f"But `torch.accelerator.current_accelerator()` works with it being = `{accelerator}`\n" f"Please reinstall torch - it's most likely broken :(" ) - if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": - return "cuda" raise NotImplementedError( "Unsloth currently only works on NVIDIA, AMD and Intel GPUs." ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7bf8866f46..5c3a5742e4 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1204,7 +1204,7 @@ SUPPORTS_BFLOAT16 = False HAS_FLASH_ATTENTION = False HAS_FLASH_ATTENTION_SOFTCAPPING = False -if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): +if DEVICE_TYPE == "cuda": major_version, minor_version = torch.cuda.get_device_capability() torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability) From fb4bd0b77727e8bd4d9466e8d9273af99f28b848 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 20:49:05 -0700 Subject: [PATCH 36/50] ci: drop `cache: 'npm'` from setup-node (silent abort on Windows) (#5474) `actions/setup-node@v6.4.0` with `cache: 'npm'` silently aborts the entire job on Windows runners when the npm cache path returned by `npm config get cache` (`C:\npm\cache`) does not yet exist on a fresh runner -- the step exits 24s in with no error message and every following step gets skipped. See npm/cli#7308 for the underlying EEXIST / missing-dir race in the npm cache directory. This mirrors the existing precedent in `studio-windows-ui-smoke.yml`'s `setup-python` block, which already dropped `cache: 'pip'` for the same reason (post-step fatal error on a missing pip cache dir). The frontend `npm ci` is fast enough without the cache that the reliability gain is worth the ~30s. --- .github/workflows/security-audit.yml | 4 ---- .github/workflows/studio-api-smoke.yml | 2 -- .github/workflows/studio-frontend-ci.yml | 2 -- .github/workflows/studio-inference-smoke.yml | 6 ------ .github/workflows/studio-mac-api-smoke.yml | 2 -- .github/workflows/studio-mac-inference-smoke.yml | 6 ------ .github/workflows/studio-mac-ui-smoke.yml | 2 -- .github/workflows/studio-mac-update-smoke.yml | 2 -- .github/workflows/studio-tauri-smoke.yml | 2 -- .github/workflows/studio-ui-smoke.yml | 2 -- .github/workflows/studio-update-smoke.yml | 2 -- .github/workflows/studio-windows-api-smoke.yml | 2 -- .github/workflows/studio-windows-inference-smoke.yml | 6 ------ .github/workflows/studio-windows-ui-smoke.yml | 9 +++++++-- .github/workflows/studio-windows-update-smoke.yml | 2 -- .github/workflows/wheel-smoke.yml | 2 -- 16 files changed, 7 insertions(+), 46 deletions(-) diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 235fde5253..a1e7b2efa6 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -137,8 +137,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 @@ -1066,8 +1064,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 29f056eca4..8bc2acc3b0 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -63,8 +63,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml index a93cdb8661..3632125ca2 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -57,8 +57,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json # Run the structural lockfile scan BEFORE npm ci. A compromised # tarball runs its `prepare` / `postinstall` during `npm ci`, diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index ea14e4f5d5..0f2320824c 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -79,8 +79,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -329,8 +327,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -648,8 +644,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index aa7a616413..9e57b097ab 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -50,8 +50,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 54d66d84b8..77cf5047c7 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -73,8 +73,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -325,8 +323,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -692,8 +688,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 28a9fc6d1d..c9742c543a 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -50,8 +50,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index dd2333251a..07d26b9ab3 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -52,8 +52,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml index 159d5dbbe6..1156c264ae 100644 --- a/.github/workflows/studio-tauri-smoke.yml +++ b/.github/workflows/studio-tauri-smoke.yml @@ -53,8 +53,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '24' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 1f3a5a8594..cd0765a1e3 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -64,8 +64,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 624001142a..1c353e933a 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -52,8 +52,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index 86a07b41e5..d5eaa75c35 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -58,8 +58,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 188cdf5a19..146b49d4f6 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -68,8 +68,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -400,8 +398,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: @@ -814,8 +810,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 90fce0558b..da8bff80b7 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -63,8 +63,13 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json + # No `cache: 'npm'`. setup-node's npm cache restore silently + # aborts the entire job on Windows runners when the npm cache + # path (`C:\npm\cache` per `npm config get cache`) doesn't yet + # exist on a fresh runner -- the step exits without an error + # message and every following step gets skipped. See + # npm/cli#7308. The frontend `npm ci` is fast enough without + # the cache that the reliability gain is worth the ~30s. - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 0303bc746d..157874d404 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -64,8 +64,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml index 464a8e324a..3de3c33ca2 100644 --- a/.github/workflows/wheel-smoke.yml +++ b/.github/workflows/wheel-smoke.yml @@ -48,8 +48,6 @@ jobs: - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' - cache: 'npm' - cache-dependency-path: studio/frontend/package-lock.json - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: From 295844670b9a898c371114fd6613c478906ab7d4 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 20:52:36 -0700 Subject: [PATCH 37/50] ci: bump Mac json-images timeout 30 -> 45 min (cache-miss path) (#5475) The `JSON, images` job in `studio-mac-inference-smoke.yml` (Job 3 of Mac Studio GGUF CI) downloads ~4 GB on a cache miss: 3 GB gemma-4-E2B-it-UD-Q4_K_XL.gguf + ~1 GB mmproj-F16.gguf. The 30 min cap was tight even with `HF_HUB_ENABLE_HF_TRANSFER=1` and parallel downloads, and timed out the cache-miss run on PR #5430 mid-download (run 25950714888) before Studio install or the smoke assertions ran. Once the actions/cache restore hits, the job comes in under 10 min, so 45 min only costs runner time on the first run after a cache key bump (v1->v2 was just bumped in #5459, which is what produced this failure). Jobs 1 (openai-anthropic, 270M model) and 2 (tool-calling, ~1.5 GB model) are not bumped -- their 25 min cap has been comfortable. --- .github/workflows/studio-mac-inference-smoke.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 77cf5047c7..b1efa238f3 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -668,7 +668,14 @@ jobs: json-images: name: JSON, images runs-on: macos-14 - timeout-minutes: 30 + # 45 min, not 30. The job downloads ~4 GB on a cache miss + # (3 GB gemma-4 GGUF + ~1 GB mmproj) over shared macos-14 NAT, + # plus Studio install + boot + JSON/image smoke. The previous + # 30 min cap timed out cache-miss runs mid-download (run + # 25950714888 / PR #5430). Once the cache is warm it lands in + # ~10 min; the headroom only matters on the first run after a + # cache key bump (v1->v2 in #5459). + timeout-minutes: 45 env: GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF # Linux smoke uses UD-IQ3_XXS, but on Mac Metal that gemma-4 From 54a86c351431834603d382e7a9f8f6c14c7f7049 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Fri, 15 May 2026 21:11:56 -0700 Subject: [PATCH 38/50] ci: route every `hf download` through xet-tuned stall-retry wrapper (#5476) Root cause of the Mac json-images 30 min timeout (run 25950714888 / PR #5430): huggingface_hub>=1.15 deprecated `hf_transfer` and routes every transfer through `hf-xet`. The CI step's unpinned `pip install --upgrade huggingface_hub hf_transfer` jumped to 1.15.0 + hf-xet 1.5.0, the 940 MB mmproj finished in ~21s, then the 3 GB gemma-4 GGUF made it to ~46% and went completely silent for the remaining 29 minutes -- no progress bytes, no error, no exit -- until the job timeout fired. This wraps every CI `hf download` in a new `.github/scripts/hf-download-with-retry.sh`: * Drops the no-op `HF_HUB_ENABLE_HF_TRANSFER=1` prefix and the `hf_transfer` install (both are deprecated on 1.15+ and only emit a FutureWarning now). * Exports the hf-xet high-performance knobs Daniel asked for: HF_XET_HIGH_PERFORMANCE=1 HF_XET_CHUNK_CACHE_SIZE_BYTES=0 HF_XET_NUM_CONCURRENT_RANGE_GETS=64 HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY=0 HF_XET_CLIENT_READ_TIMEOUT=500 * Watchdogs each attempt: if `hf download` has not exited after HF_DOWNLOAD_STALL_SECONDS (default 180s = 3 min), SIGTERM, sleep 2, SIGKILL, then loop. Retries are unbounded; the enclosing job's `timeout-minutes` is the real cap. * Optional 3rd positional `LOCAL_DIR` -- omitted lets `hf` use the default HF_HUB_CACHE, which is what the HF_HOME-priming jobs need. 19 call sites migrated across mlx-ci.yml + 9 studio-*-smoke.yml workflows. The inline `python -c "from huggingface_hub import hf_hub_download; ..."` block in mlx-ci.yml is also routed through the wrapper so every hf transfer in CI gets the same treatment. Also reverts the json-images timeout 45 -> 30 from #5475: the bump was masking this hang, not fixing it. --- .github/scripts/hf-download-with-retry.sh | 107 ++++++++++++++++++ .github/workflows/mlx-ci.yml | 13 +-- .github/workflows/studio-api-smoke.yml | 5 +- .github/workflows/studio-inference-smoke.yml | 18 ++- .github/workflows/studio-mac-api-smoke.yml | 5 +- .../workflows/studio-mac-inference-smoke.yml | 27 ++--- .github/workflows/studio-mac-ui-smoke.yml | 5 +- .github/workflows/studio-ui-smoke.yml | 5 +- .../workflows/studio-windows-api-smoke.yml | 5 +- .../studio-windows-inference-smoke.yml | 18 ++- .github/workflows/studio-windows-ui-smoke.yml | 5 +- 11 files changed, 145 insertions(+), 68 deletions(-) create mode 100755 .github/scripts/hf-download-with-retry.sh diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh new file mode 100755 index 0000000000..c5ee013c80 --- /dev/null +++ b/.github/scripts/hf-download-with-retry.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# +# Download a single file from a Hugging Face repo with a stall-retry +# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer +# kills + retries instead of silently consuming the job's timeout. +# +# Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR +# +# Why this exists +# --------------- +# huggingface_hub 1.15+ deprecated `hf_transfer` and routes every +# transfer through the `hf-xet` binary package. In CI we observed +# `hf download` on a 3 GB GGUF (gemma-4-E2B-it-UD-Q4_K_XL) progress +# to ~46% via Xet, then go completely silent for the remainder of +# the 30-min job timeout -- no progress bytes, no error, no exit. +# A sibling 940 MB mmproj on the same step downloaded in ~21s +# moments earlier, so the hang is per-file inside hf-xet rather +# than a network outage. The Xet env-vars below put hf-xet into +# its highest-throughput mode and force a 500 s client-read +# timeout; the watchdog loop ensures a stall does not eat the +# whole job: if the hf process has not exited after STALL_S +# seconds (default 180 = 3 min), we SIGTERM, then SIGKILL, then +# start a fresh attempt. Retries are unbounded -- the enclosing +# GitHub Actions job's `timeout-minutes` is the real bound. +# +# See https://huggingface.co/docs/huggingface_hub/package_reference/environment_variables +# for the HF_XET_* documentation, and npm/cli#7308's pattern (silent +# CI hang with no error) for prior art on this class of failure. + +set -uo pipefail + +REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" +FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" +# LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE +# (~/.cache/huggingface/hub) which is the desired path for callers +# that populate HF_HOME for a downstream Studio model load. +LOCAL_DIR="${3:-}" + +# Stall threshold per attempt, in seconds. Override with +# HF_DOWNLOAD_STALL_SECONDS in the workflow env if 3 min is too tight +# for a specific runner / file. The script keeps retrying past this +# until the job timeout fires. +STALL_S="${HF_DOWNLOAD_STALL_SECONDS:-180}" + +# hf-xet tuning. HF_HUB_ENABLE_HF_TRANSFER is deliberately NOT set -- +# it is a no-op on huggingface_hub>=1.15 and only emits a deprecation +# FutureWarning. The five HF_XET_* knobs below mirror the settings +# Daniel asked for: max bandwidth + 64 parallel range gets, no chunk +# cache (download-once usage pattern), parallel disk writes (SSD/NVMe +# runners), and a generous 500 s read timeout so individual chunk +# requests fail loudly instead of stalling forever. +export HF_XET_HIGH_PERFORMANCE=1 +export HF_XET_CHUNK_CACHE_SIZE_BYTES=0 +export HF_XET_NUM_CONCURRENT_RANGE_GETS=64 +export HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY=0 +export HF_XET_CLIENT_READ_TIMEOUT=500 + +if [ -n "$LOCAL_DIR" ]; then + mkdir -p "$LOCAL_DIR" +fi + +attempt=1 +while : ; do + log="$(mktemp -t hf-download.XXXXXX)" + echo "[hf-download] $FILE attempt $attempt (stall threshold ${STALL_S}s, log=$log)" + + if [ -n "$LOCAL_DIR" ]; then + hf download "$REPO" "$FILE" --local-dir "$LOCAL_DIR" > "$log" 2>&1 & + else + hf download "$REPO" "$FILE" > "$log" 2>&1 & + fi + pid=$! + + elapsed=0 + while kill -0 "$pid" 2>/dev/null && [ "$elapsed" -lt "$STALL_S" ]; do + sleep 5 + elapsed=$((elapsed + 5)) + done + + if kill -0 "$pid" 2>/dev/null; then + echo "[hf-download] $FILE attempt $attempt exceeded ${STALL_S}s -- killing PID $pid and retrying" + kill -TERM "$pid" 2>/dev/null || true + sleep 2 + kill -KILL "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + echo "[hf-download] $FILE attempt $attempt log tail (last 40 lines):" + tail -40 "$log" || true + attempt=$((attempt + 1)) + continue + fi + + if wait "$pid"; then + rc=0 + else + rc=$? + fi + + if [ "$rc" -eq 0 ]; then + echo "[hf-download] $FILE attempt $attempt succeeded" + tail -20 "$log" || true + exit 0 + fi + + echo "[hf-download] $FILE attempt $attempt failed (exit $rc) -- retrying" + tail -40 "$log" || true + attempt=$((attempt + 1)) +done diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 8cd95bd30a..75940832a0 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -302,15 +302,10 @@ jobs: "$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK" mkdir -p /tmp/ggufs - python -c " - from huggingface_hub import hf_hub_download - p = hf_hub_download( - 'unsloth/gemma-3-270m-it-GGUF', - 'gemma-3-270m-it-Q4_K_M.gguf', - local_dir = '/tmp/ggufs', - ) - print('downloaded:', p) - " + bash .github/scripts/hf-download-with-retry.sh \ + 'unsloth/gemma-3-270m-it-GGUF' \ + 'gemma-3-270m-it-Q4_K_M.gguf' \ + /tmp/ggufs PORT=18080 echo "=== starting llama-server on 127.0.0.1:$PORT ===" diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 8bc2acc3b0..53514e2ce1 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -85,10 +85,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index 0f2320824c..775363e73c 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -99,10 +99,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' @@ -347,10 +346,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache - name: Save GGUF model file if: always() && steps.download-gguf.outcome == 'success' @@ -664,12 +662,10 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 9e57b097ab..b4e274155e 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -70,10 +70,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index b1efa238f3..2d6864e0cb 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -93,10 +93,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" # Save partial caches on cancel/timeout -- hf download resumes by # content hash. `outcome != skipped` keeps cache-hit a no-op. @@ -343,10 +342,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache # Save partial caches on cancel; next run resumes via content hash. - name: Save GGUF model file @@ -668,14 +666,7 @@ jobs: json-images: name: JSON, images runs-on: macos-14 - # 45 min, not 30. The job downloads ~4 GB on a cache miss - # (3 GB gemma-4 GGUF + ~1 GB mmproj) over shared macos-14 NAT, - # plus Studio install + boot + JSON/image smoke. The previous - # 30 min cap timed out cache-miss runs mid-download (run - # 25950714888 / PR #5430). Once the cache is warm it lands in - # ~10 min; the headroom only matters on the first run after a - # cache key bump (v1->v2 in #5459). - timeout-minutes: 45 + timeout-minutes: 30 env: GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF # Linux smoke uses UD-IQ3_XXS, but on Mac Metal that gemma-4 @@ -732,13 +723,11 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache & + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache & MODEL_PID=$! - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$MMPROJ_FILE" --local-dir gguf-cache & + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" gguf-cache & MMPROJ_PID=$! wait "$MODEL_PID" wait "$MMPROJ_PID" diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index c9742c543a..510c3543d2 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -70,10 +70,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index cd0765a1e3..79476a62ea 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -84,10 +84,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index d5eaa75c35..1d12ea6f90 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -77,10 +77,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 146b49d4f6..01bf4127a7 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -99,10 +99,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} # Only write a fresh cache entry when we actually rebuilt the @@ -420,10 +419,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache - name: Save GGUF model cache if: always() && steps.download-gguf.outcome == 'success' @@ -835,12 +833,10 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index da8bff80b7..e5ab9f8ab7 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -93,10 +93,9 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | - python -m pip install --upgrade huggingface_hub hf_transfer + python -m pip install --upgrade huggingface_hub mkdir -p hf-cache - HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' From 44989ea2cbf79d1e109f09c2b2eac2cabf06d002 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sat, 16 May 2026 05:46:22 -0700 Subject: [PATCH 39/50] ci: deterministic check for studio/frontend dep removals (#5478) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ci: deterministic check for studio/frontend dep removals Adds a CI gate that catches the common foot-gun: a dep dropped from studio/frontend/package.json that something in src/ still imports. scripts/check_frontend_dep_removal.py Diffs package.json against a git base ref, collects every package no longer declared, and for each one: 1. Greps the entire repo for any usage pattern (static / dynamic / side-effect imports, require, CSS @import, HTML script/link src, new URL(), triple-slash references, template literals, bare quoted strings in JS-like files). 2. Resolves whether the package would still install by BFS'ing the dep graph in the new lockfile starting from the new package.json's declared deps (so a stale lockfile does not give false OK-via-transitive results). 3. Distinguishes top-level node_modules/<name> from nested copies under other packages. Bare src/ imports only resolve to the top-level path. 4. Pip-installed playwright references are filtered, so removing the npm playwright (CI uses the pip one) is reported correctly. Additional hygiene checks (warnings, fail with --strict): - lockfile <root> dep map matches package.json (catches drift). - @types/X is not orphaned when X is no longer declared. - No src/ import points at a package not declared in any field. tests/studio/test_frontend_dep_removal.py 24 deterministic cases. Each patches a copy of the head package.json, runs the script, and asserts (exit status, reported FAIL list). Covers: - Genuinely-breaking removals: next-themes, @xyflow/react, @huggingface/hub, dexie, motion, canvas-confetti, recharts, node-forge, mammoth, unpdf. - Safe-via-transitive removals: katex, clsx, react, @radix-ui/react-slot, zustand, tailwind-merge, remark-gfm, date-fns, js-yaml, @tauri-apps/api. - Mixed multi-removal failing on the unsafe entries only. - Non-existent / not-in-base names (no-op). - Move from deps to devDeps (not a removal). .github/workflows/studio-frontend-ci.yml Runs the checker on pull_request events against origin/${{ github.base_ref }}, plus the edge-case suite. * scripts: harden frontend dep removal check + adversarial suite classify() now catches sneaky shapes that an earlier line-only scan would miss: - multi-line `import { a, b } from "pkg"` and the same shape for `export { ... } from "pkg"` / `export * from "pkg"` / `export type ... from "pkg"`. - JSDoc `@import("pkg")` references. - Word-boundary fix so `foo` no longer matches `foobar` (subpath gate: after the package name we require closing quote or `/`). - Negative-lookbehind on `(?<!@)\bimport\b` so CSS `@import "X"` is classified as css_import, not side_effect_import. find_usage() now feeds an 8-line window (4 above / 4 below the grep hit) into classify() so multi-line import statements are picked up even though the initial grep is line-based. tests/studio/test_frontend_dep_removal.py now exercises three suites: - 24 edge cases: subprocess-driven, full-pipeline. - 28 classify() unit cases: direct function call against hand-crafted snippets. Covers static / side-effect / dynamic / require / css_import / html_script / html_link / re_export (4 variants) / template_literal / new_url / tsc_triple_slash / jsdoc_import / string_literal, plus false-positive guards (substring collision, plain-text comments, URL path tails, Python files, markdown). - 12 adversarial cases: write synthetic files under studio/frontend/src/__dep_check_adversarial__/, run the full script, then clean up. Confirms multi-line imports, re-exports, JSDoc @import, new URL, dynamic imports all FAIL when the underlying package is removed. Current total: 64 / 64 cases pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scripts: detect bin references in package.json scripts Catches the last common false-negative: removing a package whose bin is only referenced through `package.json` scripts (e.g. dropping typescript while `"build": "tsc -b && vite build"` calls tsc). Cross-checked the patterns Vercel/Next.js, Vite, and TanStack use in their own manifests; the bin/scripts pairing is the one consumer-side pattern dep checkers commonly miss. How it works: - Build a bin-to-package map from each lockfile entry's `bin` field. The map is global so a stale lockfile still resolves bins from packages about to be pruned. - Tokenize each script value, splitting on `&&`, `||`, `;`, `|`. Strip env-var assignments and `npx / pnpx / yarn / pnpm / bunx` prefixes, plus `./node_modules/.bin/` and `node_modules/.bin/` path prefixes. Look up the leading token in the bin map. - Hits are reported as `script_bin` and feed the same reachability gate as source imports. A bin still installed transitively (e.g. vite via @vitejs/plugin-react peer) is OK-via-transitive; an orphaned bin is FAIL. Test additions: - 5 new edge cases: removing vite, typescript, eslint, @biomejs/biome, and (@biomejs/biome + @vitejs/plugin-react) together. Correctly flags @biomejs/biome and the combo as FAIL while vite / typescript / eslint are kept by peers. - 8 new classify() unit cases: TypeScript ambient `declare module`, namespace imports, combined default+named, default-as-named, re-export default (4 forms), `.then()` dynamic imports without await, and TypeScript `import()` in type position. Current total: 29 edge + 36 classify-unit + 12 adversarial = 77 / 77. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * scripts: detect package.json field references to packages After surveying package.json patterns in 10+ popular repos (React, Vue/Svelte/Astro/Next.js, Vite, Storybook, TanStack/Query, Tailwind, ESLint, TypeScript, Prettier, SvelteKit), several config fields in package.json itself can reference packages by string. My checker filtered all of package.json out of the string_literal fallback, so removing a package that is only referenced from one of these fields was a false negative. Now covered (new pkg_json_field kind): - overrides / resolutions / pnpm.overrides keys - pnpm.patchedDependencies keys - peerDependenciesMeta keys - prettier: "@my/prettier-config" string - eslintConfig.extends (string or array) - stylelint.extends / stylelint.plugins - babel.presets / babel.plugins - jest.preset / jest.setupFiles / jest.transform - commitlint.extends - renovate.extends - remarkConfig.plugins - any other tool config field whose strings/keys equal the pkg name or `pkg/subpath` False-positive guards (do not flag string values inside): - browserslist (browser queries) - keywords (free-form strings) - engines / engineStrict / packageManager / volta (version pins) - files / directories / publishConfig (paths) - workspaces (paths/globs) - main / module / browser / types / typings / exports / imports / bin / man (author-side fields) - scripts (already handled separately via scripts_bin_refs) - name / version / description / author / repository / homepage etc. Test additions: new PkgFieldCase suite with 19 cases covering each tool config field, subpath references, and the 5 false-positive guards. Combined with the existing 29 edge / 36 classify / 12 adversarial cases, the suite is 96 / 96. * scripts: enumerate dead deps in studio/frontend Adds an opt-in dead-dep enumeration to the existing safety check. Iterates every package declared in studio/frontend/package.json (all four dep fields combined) and reports each as one of: used at least one detected reference -- in src/, a config file, package.json scripts (bin), a package.json tool-config field (overrides / prettier / eslintConfig / stylelint / babel / jest / commitlint / renovate / etc.), or tsconfig.compilerOptions.types unused no detected reference anywhere type_pkg_kept @types/X where X is still declared (or X = node, always implicit) type_pkg_orphan @types/X where X is no longer declared -- candidate for removal alongside X Wiring: - New CLI flag `--enumerate-dead` (off by default). - CI workflow now passes `--enumerate-dead` so the report shows on every PR run; the report is informational unless `--strict` is also set. - With `--strict`, unused / type_pkg_orphan entries fail the run. Tests: - 5 new EnumCase scenarios: E01 fake dep with no usage -> reported unused E02 fake dep imported by a synthetic src file -> reported used E03 fake dep referenced only in overrides -> reported used E04 @types/X paired with X (also imported) -> kept E05 @types/X without X -> orphan Running the new flag against the current main reproduces exactly the 11 deps PR #5477 removed, validating the heuristic end to end. Current total: 29 edge + 36 classify + 12 adversarial + 19 pkg-json field + 5 enumeration = 101 / 101. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: fetch base ref before running dep removal safety check actions/checkout uses fetch-depth: 1 by default, so when the dependency removal check ran `git show origin/main:.../package.json` the ref wasn't available locally and the script exited 2 with "could not read base package.json at origin/main:...". Fetch the single base commit before invoking the check so the git-show lookup resolves. --depth=1 keeps the extra fetch cheap. * ci: address bot review on PR 5478 Five issues flagged across gemini and codex: * --base-lock argparse arg was defined and advertised in the docstring, but main() always read args.head_lock in both branches -- the flag did nothing. Dropped the dead arg and the misleading docstring line; the lockfile-reachability analysis only needs the head lockfile. * lock_resolvable() was defined but never called. Removed. * read_pkg_file() did not specify an encoding for read_text(). Added encoding="utf-8" for cross-platform stability. * read_pkg_file() returned {} when the path did not exist, so a bad --head-lock value silently bypassed the reachability checks (false PASS for removals that resolve through npm script bins). main() now exits 2 with a clear message when the head lockfile is missing, matching the existing behavior for the head pkg. * studio-frontend-ci.yml pull_request paths filter only matched studio/frontend/** and the workflow file, so PRs that modified the checker script or its test could skip this job. Added both files to the trigger. * ci: address 10x reviewer findings on dep removal safety check Eight P1s and three P2s surfaced across 10 codex reviewers; this commit addresses all of them. P1s: 1. Workflow refspec. `git fetch --depth=1 origin <base_ref>` may only create FETCH_HEAD in shallow PR checkouts; the checker then dies with `fatal: invalid object name 'origin/main'`. Use the explicit refspec `<base>:refs/remotes/origin/<base>` so origin/<base> is reliably created. 2. `_deps_of()` was counting optional peer dependencies as reachable. npm only installs an optional peer when another package declares the same dep, so for "is this removed package still in the tree" they cannot keep it alive on their own. Skip entries marked `optional: true` in `peerDependenciesMeta`. 3. JS-syntactic classifiers (static_import, side_effect_import, dynamic_import, require, re_export, jsdoc_import, template_literal, tsc_triple_slash, new_url) now gate on file extension. Previously only the final string-literal fallback was gated, so a JS-shaped string inside a Python fixture or a Markdown code fence triggered a false FAIL. Added U37-U40 covering .py / .md / .sh / .yml. 4. HTML `<script src=>` and `<link href=>` patterns now respect a package-name boundary so `/node_modules/foo-extra/...` is not treated as a usage of `foo`. Added U41-U43. 5. New `find_command_usage()` detects CLI invocations in .sh / .yml / .yaml / .ps1 / .bat / Dockerfile* (npx pkg, bunx pkg, pnpm exec pkg, yarn dlx pkg, or a bare pkg --flag). Also covers scoped CLI packages exposed by their unscoped tail (@biomejs/biome -> biome). 6. `build_bin_to_pkg(head_lock)` was losing the bin -> package map for packages the PR correctly removed from the lockfile, so `scripts.biome:check` no longer flagged when @biomejs/biome was being dropped. Now also read the base lockfile (via `git show` or the new `--base-lock` override) and layer its bin map on top for any package in the removed set. 7. `--strict` now runs hygiene checks (lockfile sync, @types orphans, undeclared imports, dead-deps) on the no-removal path too. Previously the early return at "[OK] no dependencies removed" skipped them, so `--strict` silently passed on a tree with uncommitted lockfile drift or unused deps. 8. Removed `@types/X` packages are now matched against the runtime target name `X`: `/// <reference types="X" />`, tsconfig compilerOptions.types entries, AND runtime `import "X"` shapes. Handles the npm scope encoding (`@types/foo__bar` -> `@foo/bar`). P2s: 9. CSS `url(...)` now accepts both quoted and unquoted forms (added U44-U45). The previous regex required `/{pkg}/` after a slash, missing bare-package urls like `url(katex/fonts/x.woff2)`. 10. `find_imports_without_decl()` now covers all static-import shapes: `import "pkg"`, `import Foo from "pkg"`, `import { Foo } from "pkg"`, `import type { Foo } from "pkg"`, `await import("pkg")`, `require("pkg")`. 11. (Same as #8.) Removed `@types/X` is also linked to runtime imports of `X`, not just type-only references. Test suite expanded from 101 to 110 cases; all pass. Real-world enumerate-dead still flags the same 11 unused packages on studio/dep-removal-safety-check (matches PR 5477's removal set). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: address 4x Opus reviewer findings on dep removal check Three blockers from the parallel Opus review batch: 1. scripts_bin_refs ignored every script that began with a wrapper. The original "first non-env token wins" heuristic credited cross-env / dotenv / dotenvx / env-cmd as the bin, so a script like `cross-env CI=1 biome check` left @biomejs/biome looking unused. Rewrote into _next_real_bin(), which peels env prefixes, the leading package-manager runner (npx / pnpx / bunx / pnpm exec / yarn dlx), and the known wrapper bins (with --/-flag-arg handling) before returning the real CLI. shlex tokenization preserves quoted env values like `FOO="a b"`. 2. enumerate_dep_usage skipped find_command_usage. The non-enumerate path already credited deps used only from CI / Dockerfile / shell scripts, but `--enumerate-dead` did not, so packages referenced only from a workflow were silently listed as dead. Added the same call (gated against @types/* to avoid the unscoped-tail false positive). 3. classify multi-line window was ±4 lines. Prettier formats long named-import lists one identifier per line, so a 20-import block pushed the `import` keyword out of the window and the dep dropped to the string-literal fallback (or worse, was missed entirely). Widened to ±25 -- still bounded enough to keep false-positives negligible, wide enough for the realistic Prettier ceiling. Tests: added 10 _next_real_bin unit cases + 4 scripts_bin_refs end-to-end cases (W01-W10 + I01-I04) and a 22-identifier multi-line import adversarial case (A13). Full suite: 125/125. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-frontend-ci.yml | 22 + scripts/check_frontend_dep_removal.py | 1195 +++++++++++++++ tests/studio/test_frontend_dep_removal.py | 1628 +++++++++++++++++++++ 3 files changed, 2845 insertions(+) create mode 100644 scripts/check_frontend_dep_removal.py create mode 100644 tests/studio/test_frontend_dep_removal.py diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml index 3632125ca2..1270a57ef6 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -15,6 +15,8 @@ on: pull_request: paths: - 'studio/frontend/**' + - 'scripts/check_frontend_dep_removal.py' + - 'tests/studio/test_frontend_dep_removal.py' - '.github/workflows/studio-frontend-ci.yml' push: branches: [main, pip] @@ -84,6 +86,26 @@ jobs: exit 1 fi + # Catch the common foot-gun: a dep dropped from package.json that is + # still imported somewhere. The script walks the lockfile dep graph + # from the new top-level deps and only counts top-level node_modules + # paths as valid resolution targets for bare src/ imports. + # + # actions/checkout uses fetch-depth: 1 by default, so the base branch + # is not available locally. Fetch the single base commit with an + # explicit refspec so origin/<base> is reliably created (a bare + # `git fetch origin <ref>` only updates FETCH_HEAD in some configs). + - name: Dependency removal safety check + if: github.event_name == 'pull_request' + working-directory: ${{ github.workspace }} + run: | + git fetch --no-tags --depth=1 origin \ + "${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}" + python3 scripts/check_frontend_dep_removal.py \ + --base "origin/${{ github.base_ref }}" \ + --enumerate-dead + python3 tests/studio/test_frontend_dep_removal.py + - name: Typecheck run: npm run typecheck diff --git a/scripts/check_frontend_dep_removal.py b/scripts/check_frontend_dep_removal.py new file mode 100644 index 0000000000..260ad5215a --- /dev/null +++ b/scripts/check_frontend_dep_removal.py @@ -0,0 +1,1195 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Guard against breaking npm dependency removals in studio/frontend. + +Diffs the current package.json against a git base, finds every package +that was removed, and confirms each is no longer referenced anywhere +in the repo. If a removed package is still imported and is not +transitively resolvable through the new lockfile, exits non-zero with +file:line citations. + +Usage: + python scripts/check_frontend_dep_removal.py + python scripts/check_frontend_dep_removal.py --base origin/main + python scripts/check_frontend_dep_removal.py --base HEAD~1 + python scripts/check_frontend_dep_removal.py --base-pkg PATH --head-lock PATH + +Exit codes: + 0 every removed dep is safe (no source refs or still resolvable) + 1 at least one removed dep is referenced and not resolvable + 2 invocation error (bad args, missing file, git error) +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +FRONTEND_PKG = "studio/frontend/package.json" +FRONTEND_LOCK = "studio/frontend/package-lock.json" + +DEP_FIELDS = ( + "dependencies", + "devDependencies", + "peerDependencies", + "optionalDependencies", +) + +# Sources where seeing a package name does NOT count as usage. +EXPECTED_NOISE_FILES = { + "studio/frontend/package.json", + "studio/frontend/package-lock.json", + "studio/backend/core/data_recipe/oxc-validator/package.json", + "studio/backend/core/data_recipe/oxc-validator/package-lock.json", +} + +# Only quoted-string occurrences in these file types can be module specifiers. +JS_LIKE_EXT = re.compile( + r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$" +) +# Files where JS-syntactic import patterns (static/dynamic/require/re-export) +# could be a real module reference. Markdown gets a separate gate (.mdx is +# real ESM; .md code fences are not). +SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$") +STYLE_EXT = re.compile(r"\.(css|scss|sass)$") +HTML_EXT = re.compile(r"\.(html|htm)$") +TS_LIKE_EXT = re.compile(r"\.(ts|tsx|mts|cts|mdx)$") +# Files where a removed package's CLI binary could be invoked (npx, bunx, +# yarn dlx, pnpm exec, or a bare `pkg --flag` shell call). +COMMAND_LIKE_EXT = re.compile(r"(\.(ya?ml|sh|ps1|bat)$|(^|/)Dockerfile[^/]*$)") + +GREP_INCLUDES = [ + "--include=*.ts", + "--include=*.tsx", + "--include=*.js", + "--include=*.jsx", + "--include=*.mjs", + "--include=*.cjs", + "--include=*.html", + "--include=*.htm", + "--include=*.css", + "--include=*.scss", + "--include=*.sass", + "--include=*.json", + "--include=*.jsonc", + "--include=*.md", + "--include=*.mdx", + "--include=*.py", + "--include=*.rs", + "--include=*.toml", + "--include=*.yml", + "--include=*.yaml", + "--include=*.sh", + "--include=*.ps1", + "--include=*.bat", + "--include=Dockerfile*", +] +GREP_EXCLUDES = [ + "--exclude-dir=node_modules", + "--exclude-dir=dist", + "--exclude-dir=.git", + "--exclude-dir=__pycache__", + "--exclude-dir=target", + "--exclude-dir=.next", + "--exclude-dir=build", + "--exclude-dir=.venv", + "--exclude-dir=venv", +] + +# A pip-installed playwright reference is the PyPI package, not npm. +PIP_PLAYWRIGHT = re.compile( + r"(pip\s+install\s+['\"]?playwright" + r"|python\s+-m\s+playwright" + r"|from\s+playwright" + r"|^\s*import\s+playwright)" +) + + +@dataclass +class Hit: + file: str + line: int + kind: str + snippet: str + + +def run(cmd: list[str], cwd: Path | None = None) -> str: + """Run a command, return stdout. On non-zero exit, return ''.""" + res = subprocess.run( + cmd, + cwd = cwd or REPO_ROOT, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + text = True, + ) + return res.stdout if res.returncode == 0 else "" + + +def read_pkg_at(base: str, path: str) -> dict: + """Read JSON at `base:path` via git show. Empty dict if missing.""" + out = run(["git", "show", f"{base}:{path}"]) + if not out.strip(): + return {} + return json.loads(out) + + +def read_pkg_file(path: Path) -> dict: + if not path.exists(): + return {} + return json.loads(path.read_text(encoding = "utf-8")) + + +def all_decl_names(pkg: dict) -> set[str]: + names: set[str] = set() + for field in DEP_FIELDS: + names.update((pkg.get(field) or {}).keys()) + return names + + +def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None: + """Walk up the nested node_modules chain from `parent_path` to find + where `name` actually resolves. Mirrors Node module resolution. + """ + parts = parent_path.split("/node_modules/") + for i in range(len(parts), 0, -1): + prefix = "/node_modules/".join(parts[:i]) + trial = (prefix + "/node_modules/" if prefix else "node_modules/") + name + if trial in pkgs: + return trial + if f"node_modules/{name}" in pkgs: + return f"node_modules/{name}" + return None + + +def _deps_of(meta: dict) -> dict: + """Deps npm actually installs. Optional peers are skipped: npm only + installs them when another package declares the same dep, so for the + purpose of "is this package still reachable" they cannot keep a + removed top-level dep alive on their own. + """ + out = {} + for field in ("dependencies", "optionalDependencies"): + out.update(meta.get(field) or {}) + peer_meta = meta.get("peerDependenciesMeta") or {} + for name, spec in (meta.get("peerDependencies") or {}).items(): + if (peer_meta.get(name) or {}).get("optional"): + continue + out[name] = spec + return out + + +def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]: + """BFS the lockfile dep graph starting from `head_pkg`'s top-level + declared deps. Returns the set of lockfile install paths that survive. + Stale lockfile entries (orphaned by the new package.json) are excluded. + """ + pkgs = lock.get("packages", {}) + if not pkgs: + return set() + roots = all_decl_names(head_pkg) + seen: set[str] = set() + frontier: list[str] = [] + for name in roots: + p = _resolve_install_path("", name, pkgs) + if p: + frontier.append(p) + while frontier: + path = frontier.pop() + if path in seen: + continue + seen.add(path) + meta = pkgs.get(path, {}) + for dep_name in _deps_of(meta): + p = _resolve_install_path(path, dep_name, pkgs) + if p and p not in seen: + frontier.append(p) + return seen + + +def classify(pkg: str, file: str, content: str) -> str | None: + """Return why `content` references `pkg`, or None. + + `content` may span multiple lines (for multi-line imports/exports); + each pattern uses re.DOTALL where it matters. The bare-spec + regexes use a word-boundary check on the package name so that + `foobar` does not match `foo`. + + File-type gating: JS-syntactic patterns only fire on .ts/.tsx/.js/.jsx/ + .mjs/.cjs/.mdx files, so an `import x from "pkg"` snippet inside a + Python test fixture or a Markdown code block is not mistaken for a + real npm usage. CSS patterns only fire on .css/.scss/.sass. HTML + patterns only fire on .html/.htm. + """ + if file in EXPECTED_NOISE_FILES: + return None + + esc = re.escape(pkg) + # Subpath gate: after the package name, the next char must be either + # the closing quote, `/`, or end-of-string. Prevents foo matching foobar. + sub = r"(?:/[^'\"`]*)?" + + flags_dotall = re.DOTALL | re.MULTILINE + + is_script = bool(SCRIPT_LIKE_EXT.search(file)) + is_style = bool(STYLE_EXT.search(file)) + is_html = bool(HTML_EXT.search(file)) + is_ts = bool(TS_LIKE_EXT.search(file)) + + # If the file is none of script / style / html / json (which is the + # quoted-string fallback surface) and is not an mdx file, no classify + # rule applies. This is what gates out Python fixtures, Markdown code + # blocks, shell snippets, etc. + is_json = file.endswith(".json") or file.endswith(".jsonc") + if not (is_script or is_style or is_html or is_json): + return None + + # CSS @import is checked first so it does not collide with the + # side-effect-import regex below. + if is_style and re.search(rf"@import\s+['\"]{esc}{sub}['\"]", content): + return "css_import" + # Static imports: handle multi-line `import { ... } from "pkg"` by + # allowing arbitrary content (newlines included) between `import` + # and `from`. The non-greedy match plus the required `from` keeps + # this scoped to a single statement. + if is_script and re.search( + rf"(?<!@)\bimport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]", + content, + flags_dotall, + ): + return "static_import" + # Side-effect import: `import "pkg"` (no `from`). The negative + # lookbehind rules out CSS `@import` lines. + if is_script and re.search(rf"(?<!@)\bimport\s+['\"]{esc}{sub}['\"]", content): + return "side_effect_import" + # Dynamic import: `import("pkg")` and `await import("pkg")`. + if is_script and re.search(rf"\bimport\(\s*['\"]{esc}{sub}['\"]\s*\)", content): + return "dynamic_import" + # require / require.resolve + if is_script and re.search( + rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content + ): + return "require" + # Re-exports: `export * from "pkg"`, `export { x } from "pkg"`, + # `export type { Foo } from "pkg"`. Multi-line supported. + if is_script and re.search( + rf"\bexport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]", + content, + flags_dotall, + ): + return "re_export" + # HTML script / link. Match the package name as a complete path + # segment bounded by a quote / `#` / `?` or a subpath `/`, so + # `/node_modules/foo-extra/...` is NOT treated as usage of `foo`. + html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])" + if is_html and re.search( + rf"<script[^>]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content + ): + return "html_script" + if is_html and re.search(rf"<link[^>]*href\s*=\s*['\"][^'\"]*/{html_pkg}", content): + return "html_link" + # TypeScript triple-slash + if is_ts and re.search( + rf"///\s*<reference\s+types\s*=\s*['\"]{esc}{sub}['\"]", content + ): + return "tsc_triple_slash" + # new URL("pkg/...", import.meta.url) + if is_script and re.search(rf"\bnew\s+URL\(\s*['\"]{esc}{sub}['\"]", content): + return "new_url" + # CSS url(...). Accept quoted ("pkg/x") AND unquoted (pkg/x) variants, + # bounded by a path-segment lookahead so `pkg-extra` does not match. + if is_style and re.search( + rf"\burl\(\s*['\"]?(?:[^)'\"\s]+/)?{esc}(?:/[^)'\"`]*)?['\"]?\s*\)", + content, + ): + return "css_url" + # Template literal containing the package as the leading specifier + if is_script and re.search(rf"`{esc}{sub}`", content): + return "template_literal" + # JSDoc / TS @import comment: `@import("pkg")` + if is_script and re.search(rf"@import\(\s*['\"]{esc}{sub}['\"]\s*\)", content): + return "jsdoc_import" + # Bare quoted-string fallback (config plugin lists, vite aliases, + # tsconfig paths, biome config plugin arrays, shadcn registries). + if not JS_LIKE_EXT.search(file): + return None + # Boundary: pkg must be followed by `'`, `"`, or `/` to avoid + # matching `foo` inside `foobar`. + if re.search(rf"['\"]{esc}(?:['\"]|/)", content): + return "string_literal" + return None + + +def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]: + """Return a list of warnings if package-lock.json's <root> dep map + disagrees with package.json (i.e., npm install was not re-run). + """ + warnings = [] + if not head_lock: + return warnings + root = head_lock.get("packages", {}).get("", {}) + lock_decl = { + **(root.get("dependencies") or {}), + **(root.get("devDependencies") or {}), + **(root.get("peerDependencies") or {}), + **(root.get("optionalDependencies") or {}), + } + pkg_decl = {} + for f in DEP_FIELDS: + pkg_decl.update(head_pkg.get(f) or {}) + only_in_lock = set(lock_decl) - set(pkg_decl) + only_in_pkg = set(pkg_decl) - set(lock_decl) + if only_in_lock: + warnings.append( + f"lockfile <root> lists deps not in package.json (lockfile stale): {sorted(only_in_lock)}" + ) + if only_in_pkg: + warnings.append( + f"package.json declares deps not in lockfile <root> (run npm install): {sorted(only_in_pkg)}" + ) + return warnings + + +def types_orphan_warnings(head_pkg: dict) -> list[str]: + """Flag @types/<X> deps where <X> is no longer declared anywhere + in package.json. Removing X without also dropping @types/X leaves + dangling type packages. + """ + decl = set() + for f in DEP_FIELDS: + decl.update((head_pkg.get(f) or {}).keys()) + warnings = [] + for name in decl: + if not name.startswith("@types/"): + continue + # @types/foo provides types for `foo` + # @types/foo-bar provides types for `foo-bar` + # @types/scope__pkg provides types for `@scope/pkg` + target = name[len("@types/") :] + if "__" in target: + scope, sub = target.split("__", 1) + target = f"@{scope}/{sub}" + if target == "node": + continue # Node.js types are always implicit + if target not in decl: + warnings.append( + f"@types/{target.replace('@', '').replace('/', '__')} present but '{target}' is not declared" + ) + return warnings + + +_PKG_JSON_SKIP_KEYS = { + "dependencies", + "devDependencies", + "peerDependencies", + "optionalDependencies", + "bundleDependencies", + "bundledDependencies", +} + +# Top-level fields whose contents are never package references. We walk +# everything else recursively. +_PKG_JSON_OPAQUE_KEYS = { + "browserslist", # browser queries + "keywords", # free-form strings + "engines", # node/npm version constraints + "engineStrict", # bool + "packageManager", # `pnpm@9.0.0` -- the package manager binary + "volta", # version pins for node/npm/yarn + "files", # paths included in publish + "directories", # paths + "publishConfig", # registry / access config + "config", # generic npm config values + "main", + "module", + "browser", + "types", + "typings", + "type", + "exports", + "imports", + "bin", + "man", # author-side fields (not consumer refs) + "scripts", # handled separately via scripts_bin_refs() + "repository", + "bugs", + "homepage", + "funding", + "author", + "contributors", + "maintainers", + "license", + "licenses", + "name", + "version", + "description", + "private", + "sideEffects", + "workspaces", # paths/globs, NOT pkg names +} + + +def package_json_extra_refs(pkg: dict, target: str) -> list[str]: + """Walk every key/value in package.json EXCEPT the dep declaration + blocks, and return citations for string values or dict keys that + equal `target` (or `target/subpath`). + + Catches the patterns the public dep-checker tools commonly miss: + - `overrides` / `resolutions` / `pnpm.overrides` keys + - `pnpm.patchedDependencies` keys + - `peerDependenciesMeta` keys + - `prettier`: "@my/prettier-config" + - `eslintConfig.extends`: ["..."] / "..." + - `stylelint.extends` / `stylelint.plugins` + - `babel.presets` / `babel.plugins` + - `jest.preset` / `jest.setupFiles` / `jest.transform` + - `commitlint.extends`, `renovate.extends`, `remarkConfig.plugins` + """ + target_sub = target + "/" + cites: list[str] = [] + + def matches(s: object) -> bool: + return isinstance(s, str) and (s == target or s.startswith(target_sub)) + + def walk(obj: object, path: str) -> None: + if isinstance(obj, dict): + for k, v in obj.items(): + # Skip top-level dep declaration fields entirely. + if path == "" and k in _PKG_JSON_SKIP_KEYS: + continue + # Top-level fields whose contents are never package refs. + if path == "" and k in _PKG_JSON_OPAQUE_KEYS: + continue + # Inside `overrides` / `resolutions` / etc., the KEY itself + # is a package reference. + if matches(k): + cites.append(f"{path}.{k}" if path else k) + walk(v, f"{path}.{k}" if path else k) + elif isinstance(obj, list): + for i, v in enumerate(obj): + walk(v, f"{path}[{i}]") + elif isinstance(obj, str): + if matches(obj): + cites.append(f"{path}: {obj}") + + walk(pkg, "") + return cites + + +def build_bin_to_pkg(head_lock: dict) -> dict[str, str]: + """Map a binary name (e.g. 'vite', 'tsc', 'eslint') to the package + that provides it. Built from each lockfile entry's `bin` field. + """ + out: dict[str, str] = {} + if not head_lock: + return out + for path, meta in head_lock.get("packages", {}).items(): + if not path: + continue + name = path.split("node_modules/")[-1] + bins = meta.get("bin") + if isinstance(bins, dict): + for binname in bins: + out.setdefault(binname, name) + elif isinstance(bins, str): + out.setdefault(name.split("/")[-1], name) + return out + + +_SCRIPT_TOKENIZE = re.compile(r"\s*(?:&&|\|\||;|\|(?!\|))\s*") + +# Wrappers that delegate to a real CLI in the same shell word list. +# After stripping env prefixes and (optionally) `npx`/`pnpm exec`/`yarn dlx`/ +# `bunx`, if the leading token is one of these we advance past the +# wrapper's own flags and any further env-prefix tokens, then re-check. +# `cross-env` is the common one; `dotenv-cli` / `dotenvx` use `--` as a +# separator. Wrappers that operate on named npm-scripts (concurrently, +# npm-run-all, run-s, run-p, wireit, turbo, nx) intentionally aren't +# here -- they reference script names, not bin names, so the real bin +# is in the *target* script's chunk which we already tokenize. +_SCRIPT_WRAPPERS = {"cross-env", "dotenv", "dotenvx", "env-cmd"} +_ENV_PREFIX_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") + + +def _next_real_bin(words: list[str], idx: int) -> str | None: + """Walk `words` from `idx`, peeling env-prefix tokens, the leading + package-manager runner (`npx`, `pnpm exec`, etc.), and the known + wrapper bins. Return the next token that looks like the real CLI + binary, or None if the chunk has nothing to look up. + + Recursion depth is bounded by the chunk's word count, so the loop + cannot run away on a pathological wrapper chain. + """ + seen_wrappers: set[str] = set() + while idx < len(words): + # 1. env-prefix run: `FOO=bar BAZ="a b" cmd ...`. shlex has + # already collapsed quoted values into one word, so this + # tokenizer is safe for them. + while idx < len(words) and _ENV_PREFIX_RE.match(words[idx]): + idx += 1 + if idx >= len(words): + return None + + first = words[idx] + # 2. Package-manager runner: `npx <pkg> args`, `pnpm exec <pkg>`, + # `yarn dlx <pkg>`, `bunx <pkg>`. Strip and continue (so the + # wrapped command goes through the same unwrap loop). + if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words): + idx += 1 + continue + if ( + first in {"pnpm", "yarn"} + and idx + 2 < len(words) + and words[idx + 1] in {"exec", "dlx"} + ): + idx += 2 + continue + + # 3. Wrapper bin (cross-env, dotenv, etc.). Skip the wrapper's + # own flags and any subsequent env-prefix tokens, then re-loop. + bin_token = first.removeprefix("./node_modules/.bin/").removeprefix( + "node_modules/.bin/" + ) + if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers: + seen_wrappers.add(bin_token) + idx += 1 + # cross-env / env-cmd: no flags; just more env-prefix tokens. + # dotenv / dotenvx: skip `-e <file>` style flags and the + # optional `--` separator before the wrapped command. + while idx < len(words): + tok = words[idx] + if tok.startswith("-") and tok != "--": + idx += 1 + # `-e .env` style: also skip the flag's argument + # when it does not look like another flag. + if ( + idx < len(words) + and not words[idx].startswith("-") + and not _ENV_PREFIX_RE.match(words[idx]) + ): + idx += 1 + continue + if tok == "--": + idx += 1 + break + break + continue + return bin_token + return None + + +def scripts_bin_refs( + head_pkg: dict, bin_to_pkg: dict[str, str] +) -> dict[str, list[str]]: + """Return `{package_name: ['scripts.X: cmd', ...]}` listing every + package referenced via its bin name in package.json scripts. + + Each script value is split on shell separators (`&&`, `||`, `;`, + `|`). Within each chunk, `_next_real_bin()` unwraps env prefixes, + package-manager runners (`npx` / `pnpm exec` / `yarn dlx` / `bunx`), + and wrapper bins like `cross-env` / `dotenv` so that + `cross-env CI=1 biome check` correctly credits `biome` to its + declaring package. + + Tokenization uses shlex.split so quoted env values + (`FOO="a b" biome`) survive unbroken. + """ + import shlex + + scripts = head_pkg.get("scripts", {}) or {} + refs: dict[str, list[str]] = {} + for script_name, raw_cmd in scripts.items(): + if not isinstance(raw_cmd, str): + continue + for chunk in _SCRIPT_TOKENIZE.split(raw_cmd): + chunk = chunk.strip() + if not chunk: + continue + try: + words = shlex.split(chunk, posix = True) + except ValueError: + # Unbalanced quotes -- fall back to plain split. + words = chunk.split() + if not words: + continue + bin_name = _next_real_bin(words, 0) + if bin_name is None: + continue + pkg = bin_to_pkg.get(bin_name) + if pkg: + refs.setdefault(pkg, []).append(f"scripts.{script_name}: {raw_cmd}") + return refs + + +def tsconfig_compiler_types_refs() -> set[str]: + """Read studio/frontend/tsconfig*.json and return the set of + package names referenced in compilerOptions.types arrays. These are + implicitly loaded by tsc and count as a real use even though they + have no explicit import. + """ + out: set[str] = set() + base = REPO_ROOT / "studio/frontend" + for name in ("tsconfig.json", "tsconfig.app.json", "tsconfig.node.json"): + path = base / name + if not path.exists(): + continue + try: + text = path.read_text() + # tsconfig allows comments; strip simple line comments. + text = re.sub(r"//[^\n]*", "", text) + data = json.loads(text) + except (OSError, json.JSONDecodeError): + continue + types = (data.get("compilerOptions", {}) or {}).get("types", []) or [] + for t in types: + if not isinstance(t, str): + continue + # `vite/client` resolves to `vite` package. + pkg = ( + t.split("/", 1)[0] + if not t.startswith("@") + else "/".join(t.split("/", 2)[:2]) + ) + out.add(pkg) + return out + + +def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: + """For every declared dep, classify whether it appears used. Returns + a dict with these categories: + - used: has at least one detected usage in src/, + config files, scripts.bin, package.json + field refs, or tsconfig types + - unused: no detected usage anywhere + - type_pkg_kept: @types/X where X is still declared + - type_pkg_orphan: @types/X where X is no longer declared + (or X is removed) -- candidate for removal + + Each entry is the package name. The categorisation is opinionated; + `unused` is a CANDIDATE list, not a guarantee. The caller should + verify before deletion. + """ + decl = all_decl_names(head_pkg) + bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} + script_refs = scripts_bin_refs(head_pkg, bin_to_pkg) + tsc_types = tsconfig_compiler_types_refs() + + results: dict[str, list] = { + "used": [], + "unused": [], + "type_pkg_kept": [], + "type_pkg_orphan": [], + } + for name in sorted(decl): + if name.startswith("@types/"): + target = name[len("@types/") :] + if "__" in target: + scope, sub = target.split("__", 1) + target = f"@{scope}/{sub}" + if target == "node": + results["type_pkg_kept"].append(name) + elif target in decl: + results["type_pkg_kept"].append(name) + else: + results["type_pkg_orphan"].append(name) + continue + # Real-source-usage check + hits = find_usage(name) + used = bool(hits) + # CLI usage in shell / workflow / Dockerfile surfaces. Skip for + # `@types/*` packages because they never expose a CLI binary and + # the unscoped-tail bin name candidate would scan workflow files + # for the bare runtime name (a removed `@types/foo` would look + # for invocations of `foo`). + if not used and not name.startswith("@types/") and find_command_usage(name): + used = True + # Bin scripts + if not used and name in script_refs: + used = True + # package.json non-dep field references + if not used and package_json_extra_refs(head_pkg, name): + used = True + # tsconfig compilerOptions.types implicit usage + if not used and name in tsc_types: + used = True + if used: + results["used"].append(name) + else: + results["unused"].append(name) + return results + + +def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: + """Reverse check: find bare-specifier imports in studio/frontend/src + that don't correspond to any declared package.json dep. Catches the + case where someone adds an import but forgets the dep declaration. + Returns (file, line, spec) tuples. + + Match shapes covered: + import "pkg" + import Foo from "pkg" + import { Foo } from "pkg" + import type { Foo } from "pkg" + const x = require("pkg") + const x = await import("pkg") + """ + decl = set() + for f in DEP_FIELDS: + decl.update((head_pkg.get(f) or {}).keys()) + # Also: anything tsconfig path-aliases (just '@/...' here) is internal. + # The capture group is the specifier; the leading alternation accepts + # any of: `from "..."`, bare side-effect `import "..."`, + # `import("..."), or `require("...")`. We exclude relative paths and + # the `@/` alias prefix by requiring the first char of the specifier + # to be neither `.` nor `/`. + pattern = ( + r"(?:\bfrom\s+|" + r"\bimport\s+(?:\(\s*)?|" + r"\brequire(?:\.resolve)?\(\s*)" + r"['\"]([^'\"./][^'\"]*)['\"]" + ) + args = [ + "grep", + "-rnE", + pattern, + "--include=*.ts", + "--include=*.tsx", + "--include=*.js", + "--include=*.jsx", + "studio/frontend/src", + ] + out = run(args) + missing = [] + for line in out.splitlines(): + m = re.match(r"^(?:\./)?([^:]+):(\d+):(.*)$", line) + if not m: + continue + file, ln, content = m.group(1), int(m.group(2)), m.group(3) + for spec_match in re.finditer(pattern, content): + spec = spec_match.group(1) + # Resolve to package name (strip subpath) + if spec.startswith("@"): + parts = spec.split("/", 2) + pkg_name = "/".join(parts[:2]) if len(parts) >= 2 else spec + else: + pkg_name = spec.split("/", 1)[0] + if pkg_name in decl: + continue + # Internal aliases like '@/foo' or starts with builtin names + if pkg_name == "@": + continue + if pkg_name in { + "node:fs", + "node:path", + "fs", + "path", + "url", + "stream", + "crypto", + "buffer", + "util", + "events", + "child_process", + }: + continue + missing.append((file, ln, spec)) + return missing + + +def grep_repo(pat: str) -> list[tuple[str, int, str]]: + args = ["grep", "-rnE", pat] + GREP_INCLUDES + GREP_EXCLUDES + ["."] + out = run(args) + rows = [] + for line in out.splitlines(): + m = re.match(r"^(\./)?([^:]+):(\d+):(.*)$", line) + if m: + rows.append((m.group(2), int(m.group(3)), m.group(4))) + return rows + + +_file_lines_cache: dict[str, list[str]] = {} + + +def _read_file(path: str) -> list[str]: + if path not in _file_lines_cache: + try: + _file_lines_cache[path] = ( + Path(path).read_text(errors = "replace").splitlines() + ) + except (OSError, UnicodeDecodeError): + _file_lines_cache[path] = [] + return _file_lines_cache[path] + + +def find_usage(pkg: str) -> list[Hit]: + """Return real usages of `pkg`. Filters pip-playwright separately. + + For each filename returned by grep, also feed a multi-line window + around the matching line into classify() so multi-line imports + (`import {\n a\n} from "pkg"`) get picked up. + """ + rows = grep_repo(re.escape(pkg)) + hits = [] + seen_keys: set[tuple[str, str]] = set() + for file, lineno, content in rows: + if pkg == "playwright" and PIP_PLAYWRIGHT.search(content): + continue + # Try the single-line classify first. + kind = classify(pkg, file, content) + if not kind: + # Multi-line window: a generous 25 lines above + the line + + # 25 below so Prettier's one-import-per-line formatting for + # 12-20+ named imports still includes the `import` keyword + # in the same window as the `from "pkg"` clause. + lines = _read_file(file) + lo = max(0, lineno - 26) + hi = min(len(lines), lineno + 25) + window = "\n".join(lines[lo:hi]) + kind = classify(pkg, file, window) + if kind: + key = (file, kind) + if key in seen_keys: + continue + seen_keys.add(key) + hits.append(Hit(file, lineno, kind, content[:160])) + return hits + + +def _candidate_bin_names(pkg: str) -> set[str]: + """Names a removed package's CLI could be invoked under in shell + scripts and workflow files. Most npm CLIs use the package name + (`vite`, `eslint`, `playwright`); scoped CLI packages commonly + expose an unscoped binary name (`@biomejs/biome` -> `biome`). + """ + return {pkg, pkg.rsplit("/", 1)[-1]} + + +def find_command_usage(pkg: str) -> list[Hit]: + """Find package CLI invocations in shell / workflow / Dockerfile + surfaces: `npx pkg`, `bunx pkg`, `pnpm exec pkg`, `yarn dlx pkg`, + or a bare `pkg --flag`. Returns Hit("command_bin"). + + Detection is bounded to COMMAND_LIKE_EXT files so a JS string that + happens to contain `npx foo` inside a TS test fixture is not + mistaken for a real invocation. + """ + bins = sorted(_candidate_bin_names(pkg), key = len, reverse = True) + esc_bins = "|".join(re.escape(b) for b in bins) + # grep ERE pattern (POSIX classes for whitespace/word boundaries). + # Build without f-strings to avoid f-string-vs-{} confusion with the + # POSIX `[[:space:]]` literals and trailing `})}` boundary class. + grep_pat = ( + r"(^|[[:space:]:;&|(\[])" + r"(npx[[:space:]]+|pnpm[[:space:]]+exec[[:space:]]+" + r"|yarn[[:space:]]+(dlx[[:space:]]+)?|bunx[[:space:]]+)?" + r"(" + esc_bins + r")" + r"([[:space:])};|\]]|$)" + ) + py_pat = re.compile( + r"(^|[\s:;&|(\[])" + r"(?:npx\s+|pnpm\s+exec\s+|yarn\s+(?:dlx\s+)?|bunx\s+)?" + r"(" + esc_bins + r")" + r"([\s)};|\]]|$)" + ) + hits: list[Hit] = [] + seen: set[tuple[str, int]] = set() + for file, lineno, content in grep_repo(grep_pat): + if not COMMAND_LIKE_EXT.search(file): + continue + if pkg == "playwright" and PIP_PLAYWRIGHT.search(content): + continue + if not py_pat.search(content): + continue + key = (file, lineno) + if key in seen: + continue + seen.add(key) + hits.append(Hit(file, lineno, "command_bin", content[:160])) + return hits + + +def types_target_name(pkg: str) -> str | None: + """Strip `@types/` prefix and decode the npm scope-encoding so the + return value matches the runtime package name. `@types/foo` -> `foo`, + `@types/foo__bar` -> `@foo/bar`. Returns None for non-@types packages. + """ + if not pkg.startswith("@types/"): + return None + target = pkg[len("@types/") :] + if "__" in target: + scope, sub = target.split("__", 1) + return f"@{scope}/{sub}" + return target + + +def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]: + """For a removed `@types/X`, find usages of `X` itself: explicit + `/// <reference types="X" />`, `tsconfig.compilerOptions.types: ["X"]`, + and runtime `import "X"` shapes. The whole point of `@types/X` is to + type one of those; if any are present, the type package must stay. + """ + target = types_target_name(pkg) + if target is None: + return [] + hits = find_usage(target) + if target in tsc_types: + hits.append( + Hit( + "studio/frontend/tsconfig*.json", + 0, + "tsconfig_types", + f'compilerOptions.types includes "{target}"', + ) + ) + return hits + + +def main() -> int: + p = argparse.ArgumentParser( + description = __doc__, formatter_class = argparse.RawTextHelpFormatter + ) + p.add_argument( + "--base", + default = "origin/main", + help = "git ref to diff against (default: origin/main). " + "Examples: HEAD~1, main, a-tag, a-sha.", + ) + p.add_argument( + "--base-pkg", help = "optional override: read base package.json from this path" + ) + p.add_argument( + "--base-lock", + help = "optional override: read base package-lock.json from this path. " + "Used to recover the bin -> package mapping for removed packages so " + "scripts.foo still flags as a usage even after the PR drops node_modules/foo.", + ) + p.add_argument( + "--head-pkg", + default = str(REPO_ROOT / FRONTEND_PKG), + help = "head package.json path (default: working tree)", + ) + p.add_argument( + "--head-lock", + default = str(REPO_ROOT / FRONTEND_LOCK), + help = "head lockfile path (default: working tree). " + "Reachability analysis runs against this lockfile.", + ) + p.add_argument("--verbose", action = "store_true") + p.add_argument( + "--strict", + action = "store_true", + help = "Also fail on hygiene warnings (lockfile sync, " + "@types orphans, imports without declared dep, unused deps).", + ) + p.add_argument( + "--enumerate-dead", + action = "store_true", + help = "Print every declared dep that appears unused anywhere " + "in the repo. Informational; does not fail unless --strict.", + ) + args = p.parse_args() + + if args.base_pkg: + base_pkg = read_pkg_file(Path(args.base_pkg)) + else: + base_pkg = read_pkg_at(args.base, FRONTEND_PKG) + head_pkg = read_pkg_file(Path(args.head_pkg)) + if not base_pkg: + print( + f"ERROR: could not read base package.json at {args.base}:{FRONTEND_PKG}", + file = sys.stderr, + ) + return 2 + if not head_pkg: + print( + f"ERROR: could not read head package.json at {args.head_pkg}", + file = sys.stderr, + ) + return 2 + + head_lock_path = Path(args.head_lock) + if not head_lock_path.exists(): + print( + f"ERROR: head lockfile not found at {head_lock_path}", + file = sys.stderr, + ) + return 2 + head_lock = read_pkg_file(head_lock_path) + + # Base lockfile is best-effort. We use it only to recover the + # bin -> package mapping for packages the PR is removing -- so a + # `scripts.biome:check` cite still fires when `@biomejs/biome` is + # being dropped and the head lockfile no longer has it. + if args.base_lock: + base_lock_path = Path(args.base_lock) + base_lock = read_pkg_file(base_lock_path) if base_lock_path.exists() else {} + else: + base_lock = read_pkg_at(args.base, FRONTEND_LOCK) + + base_names = all_decl_names(base_pkg) + head_names = all_decl_names(head_pkg) + removed = sorted(base_names - head_names) + + # All hygiene checks compute up front so they can run on both the + # removal-present and removal-empty paths (so `--strict` actually + # fails when only hygiene issues exist). + sync_warns = lockfile_root_sync(head_pkg, head_lock) + types_warns = types_orphan_warnings(head_pkg) + missing_imports = find_imports_without_decl(head_pkg) + enum = enumerate_dep_usage(head_pkg, head_lock) if args.enumerate_dead else None + + def _print_hygiene() -> None: + if sync_warns: + print("Lockfile sync warnings:") + for w in sync_warns: + print(f" - {w}") + print() + if types_warns: + print("@types orphan warnings:") + for w in types_warns: + print(f" - {w}") + print() + if missing_imports: + print( + f"Imports without a matching package.json dep ({len(missing_imports)}):" + ) + for file, ln, spec in missing_imports[:20]: + print(f" - {file}:{ln} imports '{spec}'") + print() + if enum is not None: + print("Dead-dep enumeration:") + if enum["unused"]: + print(f" unused ({len(enum['unused'])}):") + for n in enum["unused"]: + print(f" - {n}") + else: + print(" unused: none") + if enum["type_pkg_orphan"]: + print(f" type_pkg_orphan ({len(enum['type_pkg_orphan'])}):") + for n in enum["type_pkg_orphan"]: + print(f" - {n}") + if args.verbose: + print(f" used: {len(enum['used'])}") + print(f" type_pkg_kept: {len(enum['type_pkg_kept'])}") + print() + + hygiene_strict_fail = args.strict and ( + sync_warns + or types_warns + or missing_imports + or (enum is not None and (enum["unused"] or enum["type_pkg_orphan"])) + ) + + if not removed: + print("[OK] no dependencies removed from studio/frontend/package.json") + if args.enumerate_dead or sync_warns or types_warns or missing_imports: + print() + _print_hygiene() + if hygiene_strict_fail: + print("FAIL (--strict): one or more hygiene warnings present") + return 1 + return 0 + + print( + f"Checking {len(removed)} removed package(s) from studio/frontend/package.json" + ) + print(f"Base: {args.base} Head: working tree") + print() + + reachable_paths = reachable_from_head(head_pkg, head_lock) if head_lock else set() + # bin -> package map: start from the head lockfile, then layer the + # base lockfile's entries on top for packages this PR is removing. + # A correct removal updates the head lockfile to drop node_modules/foo, + # so build_bin_to_pkg(head_lock) loses the mapping; we recover it + # from the base lockfile so `scripts.biome:check` still flags as a + # usage when `@biomejs/biome` is being dropped. + bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} + base_bin_to_pkg = build_bin_to_pkg(base_lock) if base_lock else {} + removed_set = set(removed) + for bin_name, pkg_name in base_bin_to_pkg.items(): + if pkg_name in removed_set: + bin_to_pkg.setdefault(bin_name, pkg_name) + script_refs = scripts_bin_refs(head_pkg, bin_to_pkg) + tsc_types = tsconfig_compiler_types_refs() + + def reachable_install_paths(name: str) -> tuple[str | None, list[str]]: + """Return (top_level_path, nested_paths). top_level is what bare + `import "name"` from src/ actually resolves to; nested copies are + only visible inside the parent package that nested them. + """ + top = f"node_modules/{name}" + top_path = top if top in reachable_paths else None + nested = sorted( + p + for p in reachable_paths + if p != top and p.endswith(f"/node_modules/{name}") + ) + return top_path, nested + + failures: list[tuple[str, list[Hit]]] = [] + for name in removed: + hits = find_usage(name) + # CLI invocations in shell scripts / workflows / Dockerfiles. + hits.extend(find_command_usage(name)) + # @types/X is "used" if X is referenced as a type or as a + # runtime import elsewhere in the repo. + hits.extend(find_types_runtime_usage(name, tsc_types)) + for cite in script_refs.get(name, []): + hits.append(Hit("studio/frontend/package.json", 0, "script_bin", cite)) + for cite in package_json_extra_refs(head_pkg, name): + hits.append(Hit("studio/frontend/package.json", 0, "pkg_json_field", cite)) + top, nested = reachable_install_paths(name) + importable_top_level = top is not None + # Source imports of bare specifier `name` resolve ONLY to top-level + # node_modules/<name>. Nested copies under another package are + # invisible to src/ files. + if hits and not importable_top_level: + status = "FAIL" + elif hits and importable_top_level: + status = "OK-via-transitive" + else: + status = "OK" + print(f" [{status}] {name}") + if top: + print(f" reachable (top-level): {top}") + if nested: + print( + f" reachable (nested, NOT importable from src/): {nested[0]}" + + (f" (+{len(nested)-1} more)" if len(nested) > 1 else "") + ) + if hits: + for h in hits[:5]: + print(f" [{h.kind}] {h.file}:{h.line} {h.snippet}") + if status == "FAIL": + failures.append((name, hits)) + if args.verbose and not hits and not (top or nested): + print(" no references, not reachable -- clean removal") + + print() + + _print_hygiene() + + if failures: + print( + f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable" + ) + for name, _ in failures: + print(f" - {name}") + return 1 + if hygiene_strict_fail: + print("FAIL (--strict): one or more hygiene warnings present") + return 1 + + print("PASS: all removed packages are safe to drop") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/studio/test_frontend_dep_removal.py b/tests/studio/test_frontend_dep_removal.py new file mode 100644 index 0000000000..a8e4cda8b7 --- /dev/null +++ b/tests/studio/test_frontend_dep_removal.py @@ -0,0 +1,1628 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Edge-case suite for scripts/check_frontend_dep_removal.py. + +Each case patches a copy of studio/frontend/package.json to remove (or +move) a specific dependency, invokes the checker against the real +working tree's lockfile, and asserts the verdict matches expectations. + +Run: + python tests/studio/test_frontend_dep_removal.py + +Exits 0 iff every case behaves as expected. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +HEAD_PKG = REPO / "studio/frontend/package.json" +HEAD_LOCK = REPO / "studio/frontend/package-lock.json" +SCRIPT = REPO / "scripts/check_frontend_dep_removal.py" + + +@dataclass +class Case: + id: str + desc: str + remove: list[str] + expected_status: str # "PASS" | "FAIL" + expected_failures: list[str] + move_to_dev: list[str] | None = None # rare: deps moved, not removed + + +CASES: list[Case] = [ + Case( + "C1", + "removing next-themes breaks 2 src imports", + ["next-themes"], + "FAIL", + ["next-themes"], + ), + Case( + "C2", + "removing @xyflow/react breaks recipe-studio src imports " + "(no other declared dep pulls @xyflow/react)", + ["@xyflow/react"], + "FAIL", + ["@xyflow/react"], + ), + Case( + "C3", + "removing katex is safe: streamdown/math, mermaid, " + "rehype-katex all keep it at top level", + ["katex"], + "PASS", + [], + ), + Case("C4", "removing clsx is safe: streamdown keeps it", ["clsx"], "PASS", []), + Case( + "C5", + "removing react is safe: peer of countless packages", + ["react"], + "PASS", + [], + ), + Case( + "C6", + "removing @radix-ui/react-slot is safe: pulled by " + "radix-ui umbrella + @assistant-ui/react", + ["@radix-ui/react-slot"], + "PASS", + [], + ), + Case( + "C7", + "removing zustand is safe: @assistant-ui/react keeps " + "top-level zustand@5.x (nested xyflow 4.x is irrelevant " + "to src imports)", + ["zustand"], + "PASS", + [], + ), + Case( + "C8", + "multi-remove with mixed safety: next-themes + " + "@huggingface/hub + dexie all unsafe", + ["next-themes", "@huggingface/hub", "dexie"], + "FAIL", + ["next-themes", "@huggingface/hub", "dexie"], + ), + Case( + "C9", + "removing @huggingface/hub breaks 5+ src imports", + ["@huggingface/hub"], + "FAIL", + ["@huggingface/hub"], + ), + Case( + "C10", + "removing tailwind-merge is safe: streamdown keeps it", + ["tailwind-merge"], + "PASS", + [], + ), + Case( + "C11", + "removing a non-existent name is a no-op", + ["__never_existed_in_pkg__"], + "PASS", + [], + ), + Case( + "C12", + "moving @hugeicons/react from deps to devDeps is NOT a " + "removal (still declared)", + [], + "PASS", + [], + move_to_dev = ["@hugeicons/react"], + ), + Case( + "C13", + "removing @huggingface/hub AND @xyflow/react together: both " + "are root-only deps with no other parents, so both should FAIL", + ["@huggingface/hub", "@xyflow/react"], + "FAIL", + ["@huggingface/hub", "@xyflow/react"], + ), + Case( + "C14", + "removing dexie breaks src imports (no other declared " "dep needs it)", + ["dexie"], + "FAIL", + ["dexie"], + ), + Case( + "C15", + "removing motion (used in 20+ src imports including " + "framer-motion-style animations); no transitive parent", + ["motion"], + "FAIL", + ["motion"], + ), + Case( + "C16", + "removing canvas-confetti (imported in confetti.tsx); " "no transitive parent", + ["canvas-confetti"], + "FAIL", + ["canvas-confetti"], + ), + Case( + "C17", + "removing recharts (imported in chart.tsx); no transitive " "parent", + ["recharts"], + "FAIL", + ["recharts"], + ), + Case( + "C18", + "removing js-yaml is safe: @eslint/eslintrc keeps it " + "(triggers @types/js-yaml orphan warning, non-fatal)", + ["js-yaml"], + "PASS", + [], + ), + Case( + "C19", + "removing node-forge (imported in providers-api.ts); " "no transitive parent", + ["node-forge"], + "FAIL", + ["node-forge"], + ), + Case( + "C20", + "removing @tauri-apps/api is safe: all 5 @tauri-apps " + "plugins declare it as a direct dep", + ["@tauri-apps/api"], + "PASS", + [], + ), + Case( + "C21", + "removing mammoth (imported in runtime-provider.tsx); " "no transitive parent", + ["mammoth"], + "FAIL", + ["mammoth"], + ), + Case( + "C22", + "removing unpdf (imported in runtime-provider.tsx); " "no transitive parent", + ["unpdf"], + "FAIL", + ["unpdf"], + ), + Case( + "C23", + "removing remark-gfm is safe: streamdown declares it " "as a direct dep", + ["remark-gfm"], + "PASS", + [], + ), + Case( + "C24", + "removing date-fns is safe: react-day-picker and " + "@base-ui/react both declare it as a direct dep", + ["date-fns"], + "PASS", + [], + ), + Case( + "C25", + "removing vite is safe: @vitejs/plugin-react and @tailwindcss/vite " + "keep it via peer (bin still resolves)", + ["vite"], + "PASS", + [], + ), + Case( + "C26", + "removing typescript is safe: 11 transitive @typescript-eslint/* " + "parents keep tsc bin alive", + ["typescript"], + "PASS", + [], + ), + Case( + "C27", + "removing eslint is safe: typescript-eslint and eslint-plugin-* " + "peers keep eslint bin alive", + ["eslint"], + "PASS", + [], + ), + Case( + "C28", + "removing @biomejs/biome breaks scripts.biome:check / biome:fix " + "(no transitive parents, biome bin orphans)", + ["@biomejs/biome"], + "FAIL", + ["@biomejs/biome"], + ), + Case( + "C29", + "removing both @biomejs/biome AND @vitejs/plugin-react together: " + "biome dies outright; vite loses one of its two retained peers " + "but @tailwindcss/vite still keeps it", + ["@biomejs/biome", "@vitejs/plugin-react"], + "FAIL", + ["@biomejs/biome", "@vitejs/plugin-react"], + ), +] + + +def synth_head(head_pkg: dict, case: Case) -> dict: + out = json.loads(json.dumps(head_pkg)) + for name in case.remove: + for field in ( + "dependencies", + "devDependencies", + "peerDependencies", + "optionalDependencies", + ): + (out.get(field) or {}).pop(name, None) + if case.move_to_dev: + for name in case.move_to_dev: + v = (out.get("dependencies") or {}).pop(name, None) + if v is not None: + out.setdefault("devDependencies", {})[name] = v + return out + + +def run_case(case: Case, head_pkg: dict) -> tuple[bool, str]: + synth = synth_head(head_pkg, case) + with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f: + json.dump(synth, f, indent = 2) + synth_path = f.name + try: + proc = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--base-pkg", + str(HEAD_PKG), + "--head-pkg", + synth_path, + "--head-lock", + str(HEAD_LOCK), + ], + capture_output = True, + text = True, + ) + finally: + os.unlink(synth_path) + + actual_status = {0: "PASS", 1: "FAIL"}.get(proc.returncode, f"RC{proc.returncode}") + failure_pkgs: list[str] = [] + in_summary = False + for line in proc.stdout.splitlines(): + if "FAIL:" in line and "removed package" in line: + in_summary = True + continue + if in_summary and line.strip().startswith("- "): + failure_pkgs.append(line.strip()[2:]) + + ok = actual_status == case.expected_status and set(failure_pkgs) == set( + case.expected_failures + ) + return ok, ( + f"expected: status={case.expected_status} fails={sorted(case.expected_failures)}\n" + f"actual: status={actual_status} fails={sorted(failure_pkgs)}\n" + f"--- stdout (first 30 lines) ---\n" + "\n".join(proc.stdout.splitlines()[:30]) + ) + + +# --------------------------------------------------------------------------- +# Classifier unit tests: feed hand-crafted snippets directly into classify() +# and assert the returned kind. Covers sneaky import shapes that an +# adversarial / careless dev might use to obscure a real usage. +# --------------------------------------------------------------------------- + +# Import the script's classify() by file path so this test does not need +# the package to be installed. +import importlib.util as _ilu + +_spec = _ilu.spec_from_file_location("_dep_check", str(SCRIPT)) +_dep_check = _ilu.module_from_spec(_spec) +sys.modules["_dep_check"] = _dep_check # required so @dataclass can resolve annotations +_spec.loader.exec_module(_dep_check) +classify = _dep_check.classify +_next_real_bin = _dep_check._next_real_bin +scripts_bin_refs = _dep_check.scripts_bin_refs + + +@dataclass +class ClassifyCase: + id: str + desc: str + pkg: str + file: str + content: str + expected_kind: str | None # None means "no detection" + + +CLASSIFY_CASES: list[ClassifyCase] = [ + # Bog-standard shapes + ClassifyCase( + "U01", + "single-line static import", + "next-themes", + "src/x.tsx", + 'import { ThemeProvider } from "next-themes";', + "static_import", + ), + ClassifyCase( + "U02", + "side-effect import", + "katex", + "src/x.tsx", + 'import "katex/dist/katex.min.css";', + "side_effect_import", + ), + ClassifyCase( + "U03", + "dynamic import", + "@tauri-apps/api", + "src/x.tsx", + 'const { x } = await import("@tauri-apps/api/window");', + "dynamic_import", + ), + ClassifyCase( + "U04", + "require()", + "lodash", + "src/x.js", + 'const _ = require("lodash");', + "require", + ), + ClassifyCase( + "U05", + "CSS @import", + "tailwindcss", + "src/x.css", + '@import "tailwindcss";', + "css_import", + ), + # Sneaky shapes + ClassifyCase( + "U06", + "multi-line static import", + "next-themes", + "src/x.tsx", + 'import {\n ThemeProvider,\n useTheme,\n} from "next-themes";', + "static_import", + ), + ClassifyCase( + "U07", + "import type", + "@huggingface/hub", + "src/x.ts", + 'import type { PipelineType } from "@huggingface/hub";', + "static_import", + ), + ClassifyCase( + "U08", + "export * from re-export", + "@some-org/secrets", + "src/x.ts", + 'export * from "@some-org/secrets";', + "re_export", + ), + ClassifyCase( + "U09", + "export { x } from re-export", + "lodash-es", + "src/x.ts", + 'export { foo, bar } from "lodash-es";', + "re_export", + ), + ClassifyCase( + "U10", + "export type ... from re-export", + "@huggingface/hub", + "src/x.ts", + 'export type { Foo } from "@huggingface/hub";', + "re_export", + ), + ClassifyCase( + "U11", + "multi-line export from re-export", + "@some/pkg", + "src/x.ts", + 'export {\n thing,\n other,\n} from "@some/pkg";', + "re_export", + ), + ClassifyCase( + "U12", + "JSDoc @import", + "react", + "src/x.ts", + '/** @type {import("react").FC} */\nconst Foo = () => null;', + "dynamic_import", + ), + ClassifyCase( + "U13", + "template literal package path", + "@assistant-ui/react", + "src/x.tsx", + "const url = `@assistant-ui/react`;", + "template_literal", + ), + ClassifyCase( + "U14", + "new URL import-meta", + "monaco-editor", + "src/x.ts", + 'new URL("monaco-editor/esm/vs/editor/editor.worker", import.meta.url);', + "new_url", + ), + ClassifyCase( + "U15", + "tsc triple-slash type ref", + "@types/some-pkg", + "src/x.ts", + '/// <reference types="@types/some-pkg" />', + "tsc_triple_slash", + ), + ClassifyCase( + "U16", + "HTML script src", + "alpinejs", + "index.html", + '<script src="/node_modules/alpinejs/dist/cdn.min.js"></script>', + "html_script", + ), + ClassifyCase( + "U17", + "HTML link href", + "alpinejs", + "index.html", + '<link rel="stylesheet" href="/node_modules/alpinejs/dist/style.css">', + "html_link", + ), + ClassifyCase( + "U18", + "bare quoted string in tsconfig paths", + "@huggingface/hub", + "tsconfig.json", + '"paths": { "hf": ["@huggingface/hub/*"] }', + "string_literal", + ), + ClassifyCase( + "U19", + "vite alias key", + "@dagrejs/dagre", + "vite.config.ts", + '"@dagrejs/dagre": path.resolve(__dirname, "./..."),', + "string_literal", + ), + # False-positive guards (these should NOT detect) + ClassifyCase( + "U20", + "different package with shared prefix", + "foo", + "src/x.ts", + 'import { x } from "foobar";', + None, + ), + ClassifyCase( + "U21", + "package mentioned in plain comment text", + "react", + "src/x.ts", + "// We migrated from react-router to tanstack-router", + None, + ), + ClassifyCase( + "U22", + "package name as a URL path tail is NOT detected " + "(boundary rule: pkg must be followed by quote or `/`)", + "react", + "src/x.ts", + 'const docs = "https://example.com/react";', + None, + ), + ClassifyCase( + "U23", + "package name in Python file (ignored, " + "Python can never import npm packages)", + "playwright", + "tests/x.py", + 'label: str = "playwright"', + None, + ), + ClassifyCase( + "U24", + "exact-prefix collision: pkg 'lodash' and 'lodash-es'", + "lodash", + "src/x.ts", + 'import _ from "lodash-es";', + None, + ), + ClassifyCase( + "U25", + "scoped pkg substring collision", + "@radix-ui/react-label", + "src/x.ts", + 'import x from "@radix-ui/react-label-extra";', + None, + ), + ClassifyCase( + "U26", + "package only mentioned in a markdown link", + "react", + "README.md", + "See [react](https://react.dev).", + None, + ), + ClassifyCase( + "U27", + "side-effect import with subpath", + "katex", + "src/x.css", + '@import "katex/dist/katex.min.css";', + "css_import", + ), + ClassifyCase( + "U28", + "require.resolve", + "lodash", + "build/x.cjs", + 'const path = require.resolve("lodash/fp");', + "require", + ), + ClassifyCase( + "U29", + "TypeScript ambient `declare module`", + "@tanstack/react-router", + "src/app/router.tsx", + 'declare module "@tanstack/react-router" {\n interface X {}\n}', + "string_literal", + ), + ClassifyCase( + "U30", + "namespace import `import * as X from pkg`", + "@radix-ui/react-slot", + "src/x.tsx", + 'import * as Slot from "@radix-ui/react-slot";', + "static_import", + ), + ClassifyCase( + "U31", + "combined default + named import", + "react", + "src/x.tsx", + 'import React, { useState } from "react";', + "static_import", + ), + ClassifyCase( + "U32", + "default-as-named import alias", + "react", + "src/x.tsx", + 'import { default as R } from "react";', + "static_import", + ), + ClassifyCase( + "U33", + "re-export default", + "lodash", + "src/x.ts", + 'export { default } from "lodash";', + "re_export", + ), + ClassifyCase( + "U34", + "re-export default as alias", + "lodash", + "src/x.ts", + 'export { default as _ } from "lodash";', + "re_export", + ), + ClassifyCase( + "U35", + ".then() dynamic import (no await)", + "@tauri-apps/api", + "src/x.ts", + 'import("@tauri-apps/api/window").then(m => m.x());', + "dynamic_import", + ), + ClassifyCase( + "U36", + "TypeScript import() in type position", + "react", + "src/x.ts", + 'type C = import("react").ComponentType;', + "dynamic_import", + ), + # File-type gating (codex P1: JS classifiers must not fire on + # non-script files). Python fixtures and Markdown code blocks often + # contain literal JS-shaped strings for documentation or test data, + # so a bare `import x from "pkg"` inside a .py / .md / .sh / .yml is + # not a real npm usage. + ClassifyCase( + "U37", + "JS import snippet inside a Python fixture string is NOT a usage", + "next-themes", + "tests/studio/something.py", + "snippet = 'import x from \"next-themes\";'", + None, + ), + ClassifyCase( + "U38", + "JS import snippet inside a Markdown code fence is NOT a usage", + "next-themes", + "docs/example.md", + '```ts\nimport x from "next-themes";\n```', + None, + ), + ClassifyCase( + "U39", + "JS import inside a shell script is NOT classified as a JS usage", + "next-themes", + "scripts/build.sh", + 'echo "import x from \\"next-themes\\";"', + None, + ), + ClassifyCase( + "U40", + "JS import inside a YAML workflow is NOT classified as a JS usage", + "next-themes", + ".github/workflows/x.yml", + "run: echo 'import x from \"next-themes\";'", + None, + ), + # HTML script/link must respect package-name boundaries: a + # `/node_modules/foo-extra/...` reference does NOT use `foo`. + ClassifyCase( + "U41", + "HTML <script src=...> with similar-prefix package is NOT a match", + "foo", + "index.html", + '<script src="/node_modules/foo-extra/dist/index.js"></script>', + None, + ), + ClassifyCase( + "U42", + "HTML <link href=...> with similar-prefix package is NOT a match", + "foo", + "index.html", + '<link rel="stylesheet" href="/node_modules/foo-extra/dist/style.css">', + None, + ), + ClassifyCase( + "U43", + "HTML <script src=...> with exact package match IS a match", + "foo", + "index.html", + '<script src="/node_modules/foo/dist/index.js"></script>', + "html_script", + ), + # CSS url() unquoted variant -- valid CSS, must classify the same + # as the quoted variant. + ClassifyCase( + "U44", + "CSS url() unquoted bare package path", + "katex", + "src/x.css", + "src: url(katex/dist/fonts/font.woff2);", + "css_url", + ), + ClassifyCase( + "U45", + "CSS url() quoted bare package path still works", + "katex", + "src/x.css", + 'src: url("katex/dist/fonts/font.woff2");', + "css_url", + ), +] + + +def run_classify_unit_tests() -> int: + passed = 0 + for c in CLASSIFY_CASES: + actual = classify(c.pkg, c.file, c.content) + ok = actual == c.expected_kind + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {c.id}: {c.desc}") + if not ok: + print(f" pkg={c.pkg!r} file={c.file!r}") + print(f" content={c.content!r}") + print(f" expected={c.expected_kind!r}, actual={actual!r}") + if ok: + passed += 1 + print() + print(f"{passed}/{len(CLASSIFY_CASES)} classify-unit cases pass") + return 0 if passed == len(CLASSIFY_CASES) else 1 + + +# --------------------------------------------------------------------------- +# Adversarial end-to-end cases: drop a sneaky synthetic file into src/, +# run the checker, then clean up. Catches the case where pattern detection +# regresses for a real grep+classify pipeline (not just classify in isolation). +# --------------------------------------------------------------------------- + +ADVERSARIAL_TMP_DIR = REPO / "studio/frontend/src/__dep_check_adversarial__" + + +@dataclass +class AdvCase: + id: str + desc: str + filename: str + content: str + target_pkg: str + expected_status: str + expected_failures: list[str] + + +ADV_CASES: list[AdvCase] = [ + AdvCase( + "A01", + "multi-line import of removed pkg should FAIL", + "adv01.ts", + 'import {\n foo,\n bar,\n} from "__adv_only_pkg_a__";\n', + "__adv_only_pkg_a__", + "FAIL", + ["__adv_only_pkg_a__"], + ), + AdvCase( + "A02", + "export * from removed pkg should FAIL", + "adv02.ts", + 'export * from "__adv_only_pkg_b__";\n', + "__adv_only_pkg_b__", + "FAIL", + ["__adv_only_pkg_b__"], + ), + AdvCase( + "A03", + "export { x } from removed pkg should FAIL", + "adv03.ts", + 'export { foo, bar } from "__adv_only_pkg_c__";\n', + "__adv_only_pkg_c__", + "FAIL", + ["__adv_only_pkg_c__"], + ), + AdvCase( + "A04", + "export type ... from removed pkg should FAIL", + "adv04.ts", + 'export type { Foo } from "__adv_only_pkg_d__";\n', + "__adv_only_pkg_d__", + "FAIL", + ["__adv_only_pkg_d__"], + ), + AdvCase( + "A05", + "package with similar prefix should NOT trigger FAIL", + "adv05.ts", + # The file imports __adv_only_pkg_e_extra__, but we will try + # to "remove" the shorter __adv_only_pkg_e__ name. The shorter + # name has zero real usage, so removal must be safe. + 'import x from "__adv_only_pkg_e_extra__";\n', + "__adv_only_pkg_e__", + "PASS", + [], + ), + AdvCase( + "A06", + "dynamic import of removed pkg should FAIL", + "adv06.ts", + 'const m = await import("__adv_only_pkg_f__");\n', + "__adv_only_pkg_f__", + "FAIL", + ["__adv_only_pkg_f__"], + ), + AdvCase( + "A07", + "new URL of removed pkg should FAIL", + "adv07.ts", + 'const w = new URL("__adv_only_pkg_g__/worker.js", import.meta.url);\n', + "__adv_only_pkg_g__", + "FAIL", + ["__adv_only_pkg_g__"], + ), + AdvCase( + "A08", + "string-concat dynamic import is unanalyzable (PASS)", + "adv08.ts", + 'const m = await import("__adv_only_" + "pkg_h__");\n', + "__adv_only_pkg_h__", + "PASS", + [], + ), + AdvCase( + "A09", + "package referenced only inside a JS comment " + "is conservatively flagged via the string_literal fallback " + "(this is acceptable -- err on the side of caution)", + "adv09.ts", + '// TODO: import x from "__adv_only_pkg_i__"\n', + "__adv_only_pkg_i__", + "FAIL", + ["__adv_only_pkg_i__"], + ), + AdvCase( + "A10", + "package referenced only in a Python file should " "NOT trigger a JS FAIL", + "adv10.py", + 'label = "__adv_only_pkg_j__"\n', + "__adv_only_pkg_j__", + "PASS", + [], + ), + AdvCase( + "A11", + "package mentioned in a markdown doc file is " + "ignored by JS-like-only string_literal", + "adv11.md", + "See [docs](https://example.com/__adv_only_pkg_k__).\n", + "__adv_only_pkg_k__", + "PASS", + [], + ), + AdvCase( + "A12", + "JSDoc @import of removed pkg should FAIL", + "adv12.ts", + '/** @type {import("__adv_only_pkg_l__").Foo} */\n' "const x = null;\n", + "__adv_only_pkg_l__", + "FAIL", + ["__adv_only_pkg_l__"], + ), + # Prettier formats a long named-import list one identifier per line. + # 22 imports + braces puts the `import` keyword ~22 lines away from + # the `from "pkg"` clause. Before the window widening, the classify + # multi-line fallback used ±4 lines, which silently missed every + # such block. This case fails with the old window and passes once + # the window is wide enough (currently ±25). + AdvCase( + "A13", + "Prettier-style 22-identifier multi-line import should FAIL " + "(exercises the widened multi-line classify window)", + "adv13.ts", + "import {\n" + + "".join(f" ident_{i:02d},\n" for i in range(22)) + + '} from "__adv_only_pkg_m__";\n', + "__adv_only_pkg_m__", + "FAIL", + ["__adv_only_pkg_m__"], + ), +] + + +# --------------------------------------------------------------------------- +# package.json field-reference cases: simulate `prettier: "@x/config"`, +# `eslintConfig.extends`, `overrides`, `peerDependenciesMeta`, etc. +# These test the package_json_extra_refs() coverage. Cross-checked against +# the patterns used by Tailwind, Stylelint, Prettier, Next.js, Astro, +# TypeScript, ESLint, SvelteKit, Storybook, Vite, and TanStack/Query +# manifests. +# --------------------------------------------------------------------------- + + +@dataclass +class PkgFieldCase: + id: str + desc: str + field_patch: dict # extra fields to merge into synth_head package.json + target_pkg: str + expected_status: str + expected_failures: list[str] + + +PKG_FIELD_CASES: list[PkgFieldCase] = [ + PkgFieldCase( + "P01", + "removing pkg referenced only in `prettier` string field", + {"prettier": "__pkg_prettier_config__"}, + "__pkg_prettier_config__", + "FAIL", + ["__pkg_prettier_config__"], + ), + PkgFieldCase( + "P02", + "removing pkg referenced in `eslintConfig.extends` array", + {"eslintConfig": {"extends": ["__pkg_eslint_cfg__"]}}, + "__pkg_eslint_cfg__", + "FAIL", + ["__pkg_eslint_cfg__"], + ), + PkgFieldCase( + "P03", + "removing pkg referenced in `stylelint.plugins`", + {"stylelint": {"plugins": ["__pkg_stylelint_plugin__"]}}, + "__pkg_stylelint_plugin__", + "FAIL", + ["__pkg_stylelint_plugin__"], + ), + PkgFieldCase( + "P04", + "removing pkg referenced in `babel.presets`", + {"babel": {"presets": [["__pkg_babel_preset__", {"opt": 1}]]}}, + "__pkg_babel_preset__", + "FAIL", + ["__pkg_babel_preset__"], + ), + PkgFieldCase( + "P05", + "removing pkg used as a key in `overrides`", + {"overrides": {"__pkg_overridden__": "^1.0.0"}}, + "__pkg_overridden__", + "FAIL", + ["__pkg_overridden__"], + ), + PkgFieldCase( + "P06", + "removing pkg used as a key in `pnpm.overrides`", + {"pnpm": {"overrides": {"__pkg_pnpm_override__": "^1.0.0"}}}, + "__pkg_pnpm_override__", + "FAIL", + ["__pkg_pnpm_override__"], + ), + PkgFieldCase( + "P07", + "removing pkg used as a key in `pnpm.patchedDependencies`", + {"pnpm": {"patchedDependencies": {"__pkg_patched__": "patches/x.patch"}}}, + "__pkg_patched__", + "FAIL", + ["__pkg_patched__"], + ), + PkgFieldCase( + "P08", + "removing pkg used as a key in `peerDependenciesMeta`", + {"peerDependenciesMeta": {"__pkg_peer_meta__": {"optional": True}}}, + "__pkg_peer_meta__", + "FAIL", + ["__pkg_peer_meta__"], + ), + PkgFieldCase( + "P09", + "removing pkg referenced in `jest.preset` string", + {"jest": {"preset": "__pkg_jest_preset__"}}, + "__pkg_jest_preset__", + "FAIL", + ["__pkg_jest_preset__"], + ), + PkgFieldCase( + "P10", + "removing pkg referenced in `commitlint.extends`", + {"commitlint": {"extends": ["__pkg_commitlint__"]}}, + "__pkg_commitlint__", + "FAIL", + ["__pkg_commitlint__"], + ), + PkgFieldCase( + "P11", + "removing pkg referenced in `renovate.extends`", + {"renovate": {"extends": ["__pkg_renovate__"]}}, + "__pkg_renovate__", + "FAIL", + ["__pkg_renovate__"], + ), + PkgFieldCase( + "P12", + "removing pkg referenced in `remarkConfig.plugins`", + {"remarkConfig": {"plugins": ["__pkg_remark__"]}}, + "__pkg_remark__", + "FAIL", + ["__pkg_remark__"], + ), + PkgFieldCase( + "P13", + "removing pkg with subpath ref in tool config (`pkg/config`)", + {"prettier": "__pkg_prettier_sub__/config"}, + "__pkg_prettier_sub__", + "FAIL", + ["__pkg_prettier_sub__"], + ), + PkgFieldCase( + "P14", + "false-positive guard: similar-prefix package in tool config", + {"prettier": "__pkg_short_extra__/config"}, + "__pkg_short__", + "PASS", + [], + ), + PkgFieldCase( + "P15", + "false-positive guard: package-named string in `browserslist` " + "must NOT trigger (browserslist values are browser queries, " + "never package names)", + {"browserslist": ["last 2 versions", "__pkg_browserslist__"]}, + "__pkg_browserslist__", + "PASS", + [], + ), + PkgFieldCase( + "P16", + "false-positive guard: matching string in `keywords` field", + {"keywords": ["__pkg_keyword__", "foo"]}, + "__pkg_keyword__", + "PASS", + [], + ), + PkgFieldCase( + "P17", + "false-positive guard: matching string in `workspaces` (paths)", + {"workspaces": ["packages/__pkg_workspace_path__"]}, + "__pkg_workspace_path__", + "PASS", + [], + ), + PkgFieldCase( + "P18", + "false-positive guard: matching value in `files` field", + {"files": ["dist/__pkg_in_files__"]}, + "__pkg_in_files__", + "PASS", + [], + ), + PkgFieldCase( + "P19", + "false-positive guard: matching `packageManager` string", + {"packageManager": "__pkg_in_pm__@1.0.0"}, + "__pkg_in_pm__", + "PASS", + [], + ), +] + + +def run_pkg_field_cases() -> int: + head_pkg = json.loads(HEAD_PKG.read_text()) + passed = 0 + for pc in PKG_FIELD_CASES: + synth_head = json.loads(json.dumps(head_pkg)) + # Apply the field patch (deep-merge isn't needed; we control the keys). + for k, v in pc.field_patch.items(): + synth_head[k] = v + # Base has the target in dependencies; head does not. The extra field + # in synth_head references the target pkg even though it's no longer + # in deps. + synth_base = json.loads(json.dumps(head_pkg)) + synth_base.setdefault("dependencies", {})[pc.target_pkg] = "^1.0.0" + with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f: + json.dump(synth_base, f, indent = 2) + base_path = f.name + with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f: + json.dump(synth_head, f, indent = 2) + head_path = f.name + try: + proc = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--base-pkg", + base_path, + "--head-pkg", + head_path, + "--head-lock", + str(HEAD_LOCK), + ], + capture_output = True, + text = True, + cwd = str(REPO), + ) + finally: + os.unlink(base_path) + os.unlink(head_path) + actual_status = {0: "PASS", 1: "FAIL"}.get( + proc.returncode, f"RC{proc.returncode}" + ) + fails: list[str] = [] + in_summary = False + for line in proc.stdout.splitlines(): + if "FAIL:" in line and "removed package" in line: + in_summary = True + continue + if in_summary and line.strip().startswith("- "): + fails.append(line.strip()[2:]) + # The expected_failures includes the tolerated-FP case (P15); we + # accept BOTH expected_status and expected_failures matches. + ok = actual_status == pc.expected_status and set(fails) == set( + pc.expected_failures + ) + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {pc.id}: {pc.desc}") + if not ok: + print( + f" expected: status={pc.expected_status} fails={pc.expected_failures}" + ) + print(f" actual: status={actual_status} fails={fails}") + for ln in proc.stdout.splitlines()[:25]: + print(f" {ln}") + if ok: + passed += 1 + print() + print(f"{passed}/{len(PKG_FIELD_CASES)} package.json-field cases pass") + return 0 if passed == len(PKG_FIELD_CASES) else 1 + + +def run_adversarial_cases() -> int: + ADVERSARIAL_TMP_DIR.mkdir(parents = True, exist_ok = True) + head_pkg = json.loads(HEAD_PKG.read_text()) + passed = 0 + for ac in ADV_CASES: + # Drop the synthetic file. + fpath = ADVERSARIAL_TMP_DIR / ac.filename + try: + fpath.write_text(ac.content) + # Build a synthetic base that has the target pkg added; head + # is the real head (without it). The script sees the pkg as + # removed and scans the repo, which now includes our file. + synth_base = json.loads(json.dumps(head_pkg)) + synth_base.setdefault("dependencies", {})[ac.target_pkg] = "^1.0.0" + with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f: + json.dump(synth_base, f, indent = 2) + base_path = f.name + try: + proc = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--base-pkg", + base_path, + "--head-pkg", + str(HEAD_PKG), + "--head-lock", + str(HEAD_LOCK), + ], + capture_output = True, + text = True, + cwd = str(REPO), + ) + finally: + os.unlink(base_path) + actual_status = {0: "PASS", 1: "FAIL"}.get( + proc.returncode, f"RC{proc.returncode}" + ) + fails = [] + in_summary = False + for line in proc.stdout.splitlines(): + if "FAIL:" in line and "removed package" in line: + in_summary = True + continue + if in_summary and line.strip().startswith("- "): + fails.append(line.strip()[2:]) + ok = actual_status == ac.expected_status and set(fails) == set( + ac.expected_failures + ) + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {ac.id}: {ac.desc}") + if not ok: + print( + f" expected: status={ac.expected_status} fails={ac.expected_failures}" + ) + print(f" actual: status={actual_status} fails={fails}") + for ln in proc.stdout.splitlines()[:20]: + print(f" {ln}") + if ok: + passed += 1 + finally: + try: + fpath.unlink() + except FileNotFoundError: + pass + # Clean up the directory. + try: + ADVERSARIAL_TMP_DIR.rmdir() + except OSError: + pass + print() + print(f"{passed}/{len(ADV_CASES)} adversarial cases pass") + return 0 if passed == len(ADV_CASES) else 1 + + +# --------------------------------------------------------------------------- +# Dead-dep enumeration cases. +# --------------------------------------------------------------------------- + + +@dataclass +class EnumCase: + id: str + desc: str + add_deps: dict[str, str] + add_dev_deps: dict[str, str] + field_patch: dict + extra_file: tuple[str, str] | None # (relative_path, content) or None + expected_unused: set[str] + expected_used: set[str] + expected_orphan_types: set[str] + + +ENUM_CASES: list[EnumCase] = [ + EnumCase( + "E01", + "fake dep with no usage anywhere is flagged unused", + {"__enum_fake_unused_pkg__": "^1.0.0"}, + {}, + {}, + None, + {"__enum_fake_unused_pkg__"}, + set(), + set(), + ), + EnumCase( + "E02", + "fake dep referenced via vite.config-style import is flagged used " + "(uses a real adversarial file as the import site)", + {"__enum_used_via_src__": "^1.0.0"}, + {}, + {}, + ( + "src/__dep_check_adversarial__/enum_e02.ts", + 'import x from "__enum_used_via_src__";\n', + ), + set(), + {"__enum_used_via_src__"}, + set(), + ), + EnumCase( + "E03", + "fake dep referenced only in package.json `overrides` is flagged used", + {"__enum_used_via_overrides__": "^1.0.0"}, + {}, + {"overrides": {"__enum_used_via_overrides__": "^1.0.0"}}, + None, + set(), + {"__enum_used_via_overrides__"}, + set(), + ), + EnumCase( + "E04", + "@types/X where X is declared -> kept (NOT orphan)", + {"__enum_real_pkg__": "^1.0.0"}, + {"@types/__enum_real_pkg__": "^1.0.0"}, + {}, + ( + "src/__dep_check_adversarial__/enum_e04.ts", + 'import x from "__enum_real_pkg__";\n', + ), + set(), + {"__enum_real_pkg__"}, + set(), + ), + EnumCase( + "E05", + "@types/X where X is NOT declared anywhere -> orphan", + {}, + {"@types/__enum_orphan_pkg__": "^1.0.0"}, + {}, + None, + set(), + set(), + {"@types/__enum_orphan_pkg__"}, + ), +] + + +def run_enum_cases() -> int: + head_pkg = json.loads(HEAD_PKG.read_text()) + passed = 0 + ADVERSARIAL_TMP_DIR.mkdir(parents = True, exist_ok = True) + for ec in ENUM_CASES: + synth_head = json.loads(json.dumps(head_pkg)) + synth_head.setdefault("dependencies", {}).update(ec.add_deps) + synth_head.setdefault("devDependencies", {}).update(ec.add_dev_deps) + for k, v in ec.field_patch.items(): + synth_head[k] = v + # Drop any temp source file if needed. + fpath = None + if ec.extra_file: + rel, content = ec.extra_file + fpath = REPO / rel + fpath.parent.mkdir(parents = True, exist_ok = True) + fpath.write_text(content) + with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f: + json.dump(synth_head, f, indent = 2) + head_path = f.name + try: + proc = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--base-pkg", + str(HEAD_PKG), + "--head-pkg", + head_path, + "--head-lock", + str(HEAD_LOCK), + "--enumerate-dead", + ], + capture_output = True, + text = True, + cwd = str(REPO), + ) + finally: + os.unlink(head_path) + if fpath: + try: + fpath.unlink() + except FileNotFoundError: + pass + # Parse the dead-dep enumeration output. + unused: set[str] = set() + orphans: set[str] = set() + in_unused = False + in_orphan = False + for line in proc.stdout.splitlines(): + s = line.strip() + if s.startswith("unused ("): + in_unused = True + in_orphan = False + continue + if s.startswith("type_pkg_orphan ("): + in_unused = False + in_orphan = True + continue + if s.startswith("used:") or s.startswith("type_pkg_kept:"): + in_unused = in_orphan = False + continue + if s.startswith("- "): + if in_unused: + unused.add(s[2:]) + elif in_orphan: + orphans.add(s[2:]) + unused_ok = ec.expected_unused.issubset(unused) and ( + not ec.expected_used or not (ec.expected_used & unused) + ) + orphan_ok = ec.expected_orphan_types.issubset(orphans) + ok = unused_ok and orphan_ok + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {ec.id}: {ec.desc}") + if not ok: + print(f" expected unused superset: {sorted(ec.expected_unused)}") + print(f" expected used NOT in unused: {sorted(ec.expected_used)}") + print( + f" expected orphans superset: {sorted(ec.expected_orphan_types)}" + ) + print(f" actual unused: {sorted(unused)}") + print(f" actual orphans: {sorted(orphans)}") + for ln in proc.stdout.splitlines()[:30]: + print(f" {ln}") + if ok: + passed += 1 + # Cleanup tmp dir if empty. + try: + ADVERSARIAL_TMP_DIR.rmdir() + except OSError: + pass + print() + print(f"{passed}/{len(ENUM_CASES)} enumeration cases pass") + return 0 if passed == len(ENUM_CASES) else 1 + + +# --------------------------------------------------------------------------- +# Script-wrapper cases: exercise scripts_bin_refs / _next_real_bin so a +# package.json script like `cross-env CI=1 biome check` correctly credits +# `@biomejs/biome` rather than the wrapper itself. The 10x reviewer flagged +# the original "first non-env token" heuristic as too narrow: any project +# using cross-env / dotenv / dotenvx / env-cmd / a quoted env value would +# bypass the bin-name check. +# --------------------------------------------------------------------------- + + +@dataclass +class WrapperCase: + id: str + desc: str + raw_cmd: str + expected_bin: str | None # None means "no real bin (e.g. unwrappable)" + + +WRAPPER_CASES: list[WrapperCase] = [ + WrapperCase( + "W01", + "cross-env wraps the real bin", + "cross-env CI=1 biome check .", + "biome", + ), + WrapperCase( + "W02", + "cross-env with multiple env tokens after the wrapper", + "cross-env A=1 B=2 NODE_ENV=prod biome check", + "biome", + ), + WrapperCase( + "W03", + "bare env-prefix run (no wrapper) still peels the env tokens", + "FOO=bar biome check", + "biome", + ), + WrapperCase( + "W04", + "quoted env value with spaces (shlex preserves it as one word)", + 'FOO="a b" biome check', + "biome", + ), + WrapperCase( + "W05", + "npx + cross-env: runner peels, wrapper peels, real bin wins", + "npx cross-env CI=1 biome check", + "biome", + ), + WrapperCase( + "W06", + "pnpm exec + cross-env", + "pnpm exec cross-env CI=1 biome check", + "biome", + ), + WrapperCase( + "W07", + "dotenv with the `--` separator before the wrapped command", + "dotenv -- biome check", + "biome", + ), + WrapperCase( + "W08", + "dotenv with a flag-arg pair and `--` separator", + "dotenv -e .env -- biome check", + "biome", + ), + WrapperCase( + "W09", + "leading `./node_modules/.bin/` prefix is stripped", + "./node_modules/.bin/biome check", + "biome", + ), + WrapperCase( + "W10", + "concurrently is NOT a script wrapper -- it dispatches by " + "script *name*, not bin, so the real bin is `concurrently` " + "itself (the wrapped script names are credited by their own " + "scripts entries, which scripts_bin_refs iterates separately)", + 'concurrently "npm:dev" "npm:typecheck"', + "concurrently", + ), +] + + +def run_wrapper_cases() -> int: + import shlex + + passed = 0 + for wc in WRAPPER_CASES: + try: + words = shlex.split(wc.raw_cmd, posix = True) + except ValueError: + words = wc.raw_cmd.split() + actual = _next_real_bin(words, 0) + ok = actual == wc.expected_bin + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {wc.id}: {wc.desc}") + if not ok: + print(f" raw_cmd={wc.raw_cmd!r}") + print(f" expected={wc.expected_bin!r}, actual={actual!r}") + if ok: + passed += 1 + + # End-to-end integration: feed scripts_bin_refs a synthetic head_pkg + # whose scripts use a wrapper, and confirm the package owning the + # wrapped bin is credited (rather than the wrapper). This is the + # actual call path used by find_command_usage(). + int_total = 0 + int_passed = 0 + int_cases = [ + ( + "I01", + "cross-env wrapping `biome` credits @biomejs/biome", + {"lint": "cross-env CI=1 biome check"}, + {"biome": "@biomejs/biome"}, + "@biomejs/biome", + ), + ( + "I02", + "dotenv -- biome credits @biomejs/biome", + {"lint": "dotenv -- biome check"}, + {"biome": "@biomejs/biome"}, + "@biomejs/biome", + ), + ( + "I03", + "quoted env value before bin still credits the bin's owner", + {"lint": 'FOO="a b" biome check .'}, + {"biome": "@biomejs/biome"}, + "@biomejs/biome", + ), + ( + "I04", + "&& chain: both halves credit their owning packages", + {"build": "tsc -b && cross-env CI=1 biome check"}, + {"tsc": "typescript", "biome": "@biomejs/biome"}, + None, # checked via owning_pkgs below + ), + ] + for case_id, desc, scripts, bin_to_pkg, expect_owner in int_cases: + int_total += 1 + refs = scripts_bin_refs({"scripts": scripts}, bin_to_pkg) + if case_id == "I04": + owners = set(refs.keys()) + ok = owners == {"typescript", "@biomejs/biome"} + else: + ok = expect_owner in refs + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {case_id}: {desc}") + if not ok: + print(f" scripts={scripts!r} bin_to_pkg={bin_to_pkg!r}") + print(f" refs={refs!r}") + if ok: + int_passed += 1 + + total = len(WRAPPER_CASES) + int_total + print() + print(f"{passed + int_passed}/{total} wrapper-script cases pass") + return 0 if (passed == len(WRAPPER_CASES) and int_passed == int_total) else 1 + + +def main() -> int: + head_pkg = json.loads(HEAD_PKG.read_text()) + print(f"Running {len(CASES)} edge cases against {SCRIPT.relative_to(REPO)}") + print() + results: list[tuple[Case, bool, str]] = [] + for c in CASES: + ok, detail = run_case(c, head_pkg) + results.append((c, ok, detail)) + mark = "PASS" if ok else "FAIL" + print(f" [{mark}] {c.id}: {c.desc}") + if not ok: + for line in detail.splitlines(): + print(f" {line}") + print() + passed = sum(1 for _, ok, _ in results if ok) + total = len(results) + print(f"{passed}/{total} edge cases pass") + + print() + print(f"Running {len(CLASSIFY_CASES)} classify() unit cases") + print() + cls_rc = run_classify_unit_tests() + + print() + print(f"Running {len(ADV_CASES)} adversarial end-to-end cases") + print() + adv_rc = run_adversarial_cases() + + print() + print(f"Running {len(PKG_FIELD_CASES)} package.json-field cases") + print() + pkg_rc = run_pkg_field_cases() + + print() + print(f"Running {len(ENUM_CASES)} dead-dep enumeration cases") + print() + enum_rc = run_enum_cases() + + print() + print( + f"Running {len(WRAPPER_CASES)} script-wrapper cases " + "(_next_real_bin + scripts_bin_refs end-to-end)" + ) + print() + wrap_rc = run_wrapper_cases() + + if ( + passed == total + and cls_rc == 0 + and adv_rc == 0 + and pkg_rc == 0 + and enum_rc == 0 + and wrap_rc == 0 + ): + return 0 + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From d2e25ee131e75e2b22306bbc1902bf00b5fc0b63 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sat, 16 May 2026 05:49:23 -0700 Subject: [PATCH 40/50] studio/frontend: drop unused dependencies, move type pkg to devDeps (#5477) * studio/frontend: drop unused dependencies, move type pkg to devDeps Removes 11 declared deps that are not imported anywhere in src/, the Tauri config, src-tauri Rust, backend, scripts, CI workflows, or sibling workspaces. Moves @types/canvas-confetti to devDependencies since it ships TypeScript types only. Removed from dependencies: @assistant-ui/react-markdown (no imports; not a peer of any used pkg) @assistant-ui/react-streamdown (no imports; not a peer of any used pkg) @langchain/core (no imports anywhere) @streamdown/cjk (no imports; not a peer of streamdown) @radix-ui/react-checkbox (re-exported by the radix-ui umbrella; no direct imports) @radix-ui/react-label (same) @radix-ui/react-select (same) @radix-ui/react-separator (same) date-fns (already a direct dep of react-day-picker) remark-gfm (already a direct dep of streamdown) Removed from devDependencies: playwright (CI installs the pip playwright; the npm one is unused) Moved to devDependencies: @types/canvas-confetti (TypeScript types only; not a runtime dep) Verified with npm install + npm run build (tsc -b && vite build), clean exit, dist/ produced. Live unsloth studio launch returns 200 on /, on the main JS / CSS bundles, and on /api/health. * studio/frontend: keep @radix-ui packages (per maintainer) Maintainer asked to keep the four @radix-ui packages this PR was originally dropping: @radix-ui/react-checkbox ^1.3.3 @radix-ui/react-label ^2.1.8 @radix-ui/react-select ^2.2.6 @radix-ui/react-separator ^1.1.8 Restored to dependencies and refreshed the lockfile. Build still green (1044 packages, vite build 2.1s, same dist contents). --- studio/frontend/package-lock.json | 433 +----------------------------- studio/frontend/package.json | 9 +- 2 files changed, 5 insertions(+), 437 deletions(-) diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index ee0d8a7832..0525b984a1 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -10,8 +10,6 @@ "dependencies": { "@assistant-ui/core": "0.1.17", "@assistant-ui/react": "0.12.28", - "@assistant-ui/react-markdown": "0.12.11", - "@assistant-ui/react-streamdown": "0.1.11", "@assistant-ui/tap": "0.5.10", "@base-ui/react": "^1.2.0", "@dagrejs/dagre": "^2.0.4", @@ -22,13 +20,11 @@ "@hugeicons/core-free-icons": "^4.1.1", "@hugeicons/react": "^1.1.5", "@huggingface/hub": "^2.9.0", - "@langchain/core": "^1.1.27", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-label": "^2.1.8", "@radix-ui/react-select": "^2.2.6", "@radix-ui/react-separator": "^1.1.8", "@radix-ui/react-slot": "^1.2.4", - "@streamdown/cjk": "1.0.3", "@streamdown/code": "1.1.1", "@streamdown/math": "1.0.2", "@streamdown/mermaid": "1.0.2", @@ -42,14 +38,12 @@ "@tauri-apps/plugin-process": "^2.3.1", "@tauri-apps/plugin-updater": "^2.10.1", "@toolwind/corner-shape": "^0.0.8-3", - "@types/canvas-confetti": "^1.9.0", "@xyflow/react": "^12.10.0", "assistant-stream": "0.3.12", "canvas-confetti": "^1.9.4", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "cmdk": "^1.1.1", - "date-fns": "^4.1.0", "dexie": "^4.3.0", "js-yaml": "^4.1.1", "katex": "^0.16.28", @@ -64,7 +58,6 @@ "react-dom": "^19.2.4", "react-resizable-panels": "^4.6.4", "recharts": "3.7.0", - "remark-gfm": "^4.0.1", "shadcn": "^4.2.0", "sonner": "^2.0.7", "streamdown": "2.5.0", @@ -78,6 +71,7 @@ "devDependencies": { "@biomejs/biome": "^1.9.4", "@eslint/js": "^9.39.1", + "@types/canvas-confetti": "^1.9.0", "@types/js-yaml": "^4.0.9", "@types/node": "^25.5.2", "@types/node-forge": "^1.3.14", @@ -88,7 +82,6 @@ "eslint-plugin-react-hooks": "^7.0.1", "eslint-plugin-react-refresh": "^0.5.2", "globals": "^17.4.0", - "playwright": "^1.59.1", "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", "vite": "^8.0.1" @@ -180,66 +173,6 @@ } } }, - "node_modules/@assistant-ui/react-markdown": { - "version": "0.12.11", - "resolved": "https://registry.npmjs.org/@assistant-ui/react-markdown/-/react-markdown-0.12.11.tgz", - "integrity": "sha512-gYu4XVI2lX3lp9UG7V5VWP1+eO7SZomiBKsAZOKUOeuwn/hoL+J0vFY52FUgJixdF2R8NPPto2lb98DmJE70lA==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-primitive": "^2.1.4", - "@radix-ui/react-use-callback-ref": "^1.1.1", - "classnames": "^2.5.1", - "react-markdown": "^10.1.0" - }, - "peerDependencies": { - "@assistant-ui/react": "^0.12.26", - "@types/react": "*", - "react": "^18 || ^19" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@assistant-ui/react-streamdown": { - "version": "0.1.11", - "resolved": "https://registry.npmjs.org/@assistant-ui/react-streamdown/-/react-streamdown-0.1.11.tgz", - "integrity": "sha512-9y+89ZxotYSt81hChSVjK2kwUYRKq7UW/r5qoqZTpcb7119gc0NOj0dx9xxuyXE2QfR6EY8rW6yBz3g+Y7RrhQ==", - "license": "MIT", - "dependencies": { - "rehype-harden": "^1.1.8", - "rehype-raw": "^7.0.0", - "rehype-sanitize": "^6.0.0", - "streamdown": "^2.5.0" - }, - "peerDependencies": { - "@assistant-ui/react": "^0.12.26", - "@streamdown/cjk": "^1.0.0", - "@streamdown/code": "^1.0.0", - "@streamdown/math": "^1.0.0", - "@streamdown/mermaid": "^1.0.0", - "@types/react": "*", - "react": "^18 || ^19" - }, - "peerDependenciesMeta": { - "@streamdown/cjk": { - "optional": true - }, - "@streamdown/code": { - "optional": true - }, - "@streamdown/math": { - "optional": true - }, - "@streamdown/mermaid": { - "optional": true - }, - "@types/react": { - "optional": true - } - } - }, "node_modules/@assistant-ui/store": { "version": "0.2.9", "resolved": "https://registry.npmjs.org/@assistant-ui/store/-/store-0.2.9.tgz", @@ -921,12 +854,6 @@ "integrity": "sha512-jigsZK+sMF/cuiB7sERuo9V7N9jx+dhmHHnQyDSVdpZwVutaBu7WvNYqMDLSgFgfB30n452TP3vjDAvFC973mA==", "license": "MIT" }, - "node_modules/@cfworker/json-schema": { - "version": "4.1.1", - "resolved": "https://registry.npmjs.org/@cfworker/json-schema/-/json-schema-4.1.1.tgz", - "integrity": "sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==", - "license": "MIT" - }, "node_modules/@chevrotain/cst-dts-gen": { "version": "12.0.0", "resolved": "https://registry.npmjs.org/@chevrotain/cst-dts-gen/-/cst-dts-gen-12.0.0.tgz", @@ -1668,27 +1595,6 @@ "@jridgewell/sourcemap-codec": "^1.4.14" } }, - "node_modules/@langchain/core": { - "version": "1.1.44", - "resolved": "https://registry.npmjs.org/@langchain/core/-/core-1.1.44.tgz", - "integrity": "sha512-RePW1IjGCHr9ua2vcby3aE8mOOz3EnwDZxMEGbNDT91kf14eqkJqxDXvaZFviGdcN9DTrxM5RPQNAHmwSm4tbg==", - "license": "MIT", - "dependencies": { - "@cfworker/json-schema": "^4.0.2", - "@standard-schema/spec": "^1.1.0", - "ansi-styles": "^5.0.0", - "camelcase": "6", - "decamelize": "1.2.0", - "js-tiktoken": "^1.0.12", - "langsmith": ">=0.5.0 <1.0.0", - "mustache": "^4.2.0", - "p-queue": "^6.6.2", - "zod": "^3.25.76 || ^4" - }, - "engines": { - "node": ">=20" - } - }, "node_modules/@mermaid-js/parser": { "version": "1.1.0", "resolved": "https://registry.npmjs.org/@mermaid-js/parser/-/parser-1.1.0.tgz", @@ -5799,20 +5705,6 @@ "integrity": "sha512-e7Mew686owMaPJVNNLs55PUvgz371nKgwsc4vxE49zsODpJEnxgxRo2y/OKrqueavXgZNMDVj3DdHFlaSAeU8g==", "license": "MIT" }, - "node_modules/@streamdown/cjk": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@streamdown/cjk/-/cjk-1.0.3.tgz", - "integrity": "sha512-WRg8HR/gHbBoTgsMd91OKFUClIoDcEFVofJvluvEAyjx3KpU0aGgD9tGDqHkHj14ShoMSkX0IYetWGegTcwIJw==", - "license": "Apache-2.0", - "dependencies": { - "remark-cjk-friendly": "^2.0.1", - "remark-cjk-friendly-gfm-strikethrough": "^2.0.1", - "unist-util-visit": "^5.0.0" - }, - "peerDependencies": { - "react": "^18.0.0 || ^19.0.0" - } - }, "node_modules/@streamdown/code": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/@streamdown/code/-/code-1.1.1.tgz", @@ -6431,6 +6323,7 @@ "version": "1.9.0", "resolved": "https://registry.npmjs.org/@types/canvas-confetti/-/canvas-confetti-1.9.0.tgz", "integrity": "sha512-aBGj/dULrimR1XDZLtG9JwxX1b4HPRF6CX9Yfwh3NvstZEm1ZL7RBnel4keCPSqs1ANRu1u2Aoz9R+VmtjYuTg==", + "dev": true, "license": "MIT" }, "node_modules/@types/d3": { @@ -6783,6 +6676,7 @@ "version": "19.2.14", "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.14.tgz", "integrity": "sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==", + "devOptional": true, "license": "MIT", "dependencies": { "csstype": "^3.2.2" @@ -7356,18 +7250,6 @@ "url": "https://github.com/chalk/ansi-regex?sponsor=1" } }, - "node_modules/ansi-styles": { - "version": "5.2.0", - "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-5.2.0.tgz", - "integrity": "sha512-Cxwpt2SfTzTtXcfOlzGEee8O+c+MmUgGrNiBcXnuWxuFJHe6a5Hz7qwhwe5OgaSYI0IJvkLqWX1ASG+cJOkEiA==", - "license": "MIT", - "engines": { - "node": ">=10" - }, - "funding": { - "url": "https://github.com/chalk/ansi-styles?sponsor=1" - } - }, "node_modules/argparse": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz", @@ -7631,18 +7513,6 @@ "node": ">=6" } }, - "node_modules/camelcase": { - "version": "6.3.0", - "resolved": "https://registry.npmjs.org/camelcase/-/camelcase-6.3.0.tgz", - "integrity": "sha512-Gmy6FhYlCY7uOElZUSbxo2UCDH8owEk996gkbrpsgGtrJLM3J7jGxl9Ic7Qwwj4ivOE5AWZWRMecDdF7hqGjFA==", - "license": "MIT", - "engines": { - "node": ">=10" - }, - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } - }, "node_modules/caniuse-lite": { "version": "1.0.30001791", "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001791.tgz", @@ -7802,12 +7672,6 @@ "integrity": "sha512-JhZUT7JFcQy/EzW605k/ktHtncoo9vnyW/2GspNYwFlN1C/WmjuV/xtS04e9SOkL2sTdw0VAZ2UGCcQ9lR6p6w==", "license": "MIT" }, - "node_modules/classnames": { - "version": "2.5.1", - "resolved": "https://registry.npmjs.org/classnames/-/classnames-2.5.1.tgz", - "integrity": "sha512-saHYOzhIQs6wy2sVxTM6bUDsQO4F50V9RQ22qBpEdCW+I+/Wmke2HOl6lS6dTpdxVhb88/I6+Hs+438c3lfUow==", - "license": "MIT" - }, "node_modules/cli-cursor": { "version": "5.0.0", "resolved": "https://registry.npmjs.org/cli-cursor/-/cli-cursor-5.0.0.tgz", @@ -8113,6 +7977,7 @@ "version": "3.2.3", "resolved": "https://registry.npmjs.org/csstype/-/csstype-3.2.3.tgz", "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", + "devOptional": true, "license": "MIT" }, "node_modules/cytoscape": { @@ -8671,15 +8536,6 @@ } } }, - "node_modules/decamelize": { - "version": "1.2.0", - "resolved": "https://registry.npmjs.org/decamelize/-/decamelize-1.2.0.tgz", - "integrity": "sha512-z2S+W9X73hAUUki+N+9Za2lBlun89zigOyGrsax+KUQ6wKW4ZoWpEYBkGhQjwAjjDCkWxhY0VKEhk8wzY7F5cA==", - "license": "MIT", - "engines": { - "node": ">=0.10.0" - } - }, "node_modules/decimal.js-light": { "version": "2.5.1", "resolved": "https://registry.npmjs.org/decimal.js-light/-/decimal.js-light-2.5.1.tgz", @@ -9260,12 +9116,6 @@ "node": ">= 0.6" } }, - "node_modules/eventemitter3": { - "version": "4.0.7", - "resolved": "https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.7.tgz", - "integrity": "sha512-8guHBZCwKnFhYdHr2ysuRWErTwhoN2X8XELRlrRwpmfeY2jjuUN4taQMsULKUVo1K4DvZl+0pgfyoysHxvmvEw==", - "license": "MIT" - }, "node_modules/eventsource": { "version": "3.0.7", "resolved": "https://registry.npmjs.org/eventsource/-/eventsource-3.0.7.tgz", @@ -9687,21 +9537,6 @@ "node": ">=14.14" } }, - "node_modules/fsevents": { - "version": "2.3.2", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", - "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, - "hasInstallScript": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^8.16.0 || ^10.6.0 || >=11.0.0" - } - }, "node_modules/function-bind": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz", @@ -10636,15 +10471,6 @@ "url": "https://github.com/sponsors/panva" } }, - "node_modules/js-tiktoken": { - "version": "1.0.21", - "resolved": "https://registry.npmjs.org/js-tiktoken/-/js-tiktoken-1.0.21.tgz", - "integrity": "sha512-biOj/6M5qdgx5TKjDnFT1ymSpM5tbd3ylwDtrQvFQSu0Z7bBYko2dF+W/aUkXUPuk6IVpRxk/3Q2sHOzGlS36g==", - "license": "MIT", - "dependencies": { - "base64-js": "^1.5.1" - } - }, "node_modules/js-tokens": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz", @@ -10802,39 +10628,6 @@ "npm": ">=10.2.3" } }, - "node_modules/langsmith": { - "version": "0.6.1", - "resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.6.1.tgz", - "integrity": "sha512-qNBNPRFqScIlGaPGfMrxhw/GTOL4GJKMp1P4jeA3xuI+Gkj5Ei3wvOtxpYaZMTeBbXTW3yi4n4Wf3nCgogvttg==", - "license": "MIT", - "dependencies": { - "p-queue": "6.6.2" - }, - "peerDependencies": { - "@opentelemetry/api": "*", - "@opentelemetry/exporter-trace-otlp-proto": "*", - "@opentelemetry/sdk-trace-base": "*", - "openai": "*", - "ws": ">=7" - }, - "peerDependenciesMeta": { - "@opentelemetry/api": { - "optional": true - }, - "@opentelemetry/exporter-trace-otlp-proto": { - "optional": true - }, - "@opentelemetry/sdk-trace-base": { - "optional": true - }, - "openai": { - "optional": true - }, - "ws": { - "optional": true - } - } - }, "node_modules/layout-base": { "version": "1.0.2", "resolved": "https://registry.npmjs.org/layout-base/-/layout-base-1.0.2.tgz", @@ -11735,77 +11528,6 @@ "micromark-util-types": "^2.0.0" } }, - "node_modules/micromark-extension-cjk-friendly": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/micromark-extension-cjk-friendly/-/micromark-extension-cjk-friendly-2.0.1.tgz", - "integrity": "sha512-OkzoYVTL1ChbvQ8Cc1ayTIz7paFQz8iS9oIYmewncweUSwmWR+hkJF9spJ1lxB90XldJl26A1F4IkPOKS3bDXw==", - "license": "MIT", - "dependencies": { - "devlop": "^1.1.0", - "micromark-extension-cjk-friendly-util": "3.0.1", - "micromark-util-chunked": "^2.0.1", - "micromark-util-resolve-all": "^2.0.1", - "micromark-util-symbol": "^2.0.1" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "micromark": "^4.0.0", - "micromark-util-types": "^2.0.0" - }, - "peerDependenciesMeta": { - "micromark-util-types": { - "optional": true - } - } - }, - "node_modules/micromark-extension-cjk-friendly-gfm-strikethrough": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/micromark-extension-cjk-friendly-gfm-strikethrough/-/micromark-extension-cjk-friendly-gfm-strikethrough-2.0.1.tgz", - "integrity": "sha512-wVC0zwjJNqQeX+bb07YTPu/CvSAyCTafyYb7sMhX1r62/Lw5M/df3JyYaANyp8g15c1ypJRFSsookTqA1IDsUg==", - "license": "MIT", - "dependencies": { - "devlop": "^1.1.0", - "get-east-asian-width": "^1.4.0", - "micromark-extension-cjk-friendly-util": "3.0.1", - "micromark-util-character": "^2.1.1", - "micromark-util-chunked": "^2.0.1", - "micromark-util-resolve-all": "^2.0.1", - "micromark-util-symbol": "^2.0.1" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "micromark": "^4.0.0", - "micromark-util-types": "^2.0.0" - }, - "peerDependenciesMeta": { - "micromark-util-types": { - "optional": true - } - } - }, - "node_modules/micromark-extension-cjk-friendly-util": { - "version": "3.0.1", - "resolved": "https://registry.npmjs.org/micromark-extension-cjk-friendly-util/-/micromark-extension-cjk-friendly-util-3.0.1.tgz", - "integrity": "sha512-GcbXqTTHOsiZHyF753oIddP/J2eH8j9zpyQPhkof6B2JNxfEJabnQqxbCgzJNuNes0Y2jTNJ3LiYPSXr6eJA8w==", - "license": "MIT", - "dependencies": { - "get-east-asian-width": "^1.4.0", - "micromark-util-character": "^2.1.1", - "micromark-util-symbol": "^2.0.1" - }, - "engines": { - "node": ">=18" - }, - "peerDependenciesMeta": { - "micromark-util-types": { - "optional": true - } - } - }, "node_modules/micromark-extension-gfm": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/micromark-extension-gfm/-/micromark-extension-gfm-3.0.0.tgz", @@ -12528,15 +12250,6 @@ "url": "https://opencollective.com/express" } }, - "node_modules/mustache": { - "version": "4.2.0", - "resolved": "https://registry.npmjs.org/mustache/-/mustache-4.2.0.tgz", - "integrity": "sha512-71ippSywq5Yb7/tVYyGbkBggbU8H3u5Rz56fH60jGFgr8uHwxs+aSKeqmluIVzM0m0kB7xQjKS6qPfd0b2ZoqQ==", - "license": "MIT", - "bin": { - "mustache": "bin/mustache" - } - }, "node_modules/mute-stream": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/mute-stream/-/mute-stream-3.0.0.tgz", @@ -12862,15 +12575,6 @@ "integrity": "sha512-+Sl2UErvtsoajRDKCE5/dBz4DIvHXQQnAxtQTF04OJxY0+DyZXSo5P5Bb7XYWOh81syohlYL24hbDwxedPUJCA==", "license": "MIT" }, - "node_modules/p-finally": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/p-finally/-/p-finally-1.0.0.tgz", - "integrity": "sha512-LICb2p9CB7FS+0eR1oqWnHhp0FljGLZCWBE9aix0Uye9W8LTQPwMTYVGWQWIw9RdQiDg4+epXQODwIYJtSJaow==", - "license": "MIT", - "engines": { - "node": ">=4" - } - }, "node_modules/p-limit": { "version": "3.1.0", "resolved": "https://registry.npmjs.org/p-limit/-/p-limit-3.1.0.tgz", @@ -12903,34 +12607,6 @@ "url": "https://github.com/sponsors/sindresorhus" } }, - "node_modules/p-queue": { - "version": "6.6.2", - "resolved": "https://registry.npmjs.org/p-queue/-/p-queue-6.6.2.tgz", - "integrity": "sha512-RwFpb72c/BhQLEXIZ5K2e+AhgNVmIejGlTgiB9MzZ0e93GRvqZ7uSi0dvRF7/XIXDeNkra2fNHBxTyPDGySpjQ==", - "license": "MIT", - "dependencies": { - "eventemitter3": "^4.0.4", - "p-timeout": "^3.2.0" - }, - "engines": { - "node": ">=8" - }, - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } - }, - "node_modules/p-timeout": { - "version": "3.2.0", - "resolved": "https://registry.npmjs.org/p-timeout/-/p-timeout-3.2.0.tgz", - "integrity": "sha512-rhIwUycgwwKcP9yTOOFK/AKsAopjjCakVqLHePO3CC6Mir1Z99xT+R63jZxAT5lFZLa2inS5h+ZS2GvR99/FBg==", - "license": "MIT", - "dependencies": { - "p-finally": "^1.0.0" - }, - "engines": { - "node": ">=8" - } - }, "node_modules/package-manager-detector": { "version": "1.6.0", "resolved": "https://registry.npmjs.org/package-manager-detector/-/package-manager-detector-1.6.0.tgz", @@ -13121,38 +12797,6 @@ "pathe": "^2.0.1" } }, - "node_modules/playwright": { - "version": "1.59.1", - "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.59.1.tgz", - "integrity": "sha512-C8oWjPR3F81yljW9o5OxcWzfh6avkVwDD2VYdwIGqTkl+OGFISgypqzfu7dOe4QNLL2aqcWBmI3PMtLIK233lw==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "playwright-core": "1.59.1" - }, - "bin": { - "playwright": "cli.js" - }, - "engines": { - "node": ">=18" - }, - "optionalDependencies": { - "fsevents": "2.3.2" - } - }, - "node_modules/playwright-core": { - "version": "1.59.1", - "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.59.1.tgz", - "integrity": "sha512-HBV/RJg81z5BiiZ9yPzIiClYV/QMsDCKUyogwH9p3MCP6IYjUFu/MActgYAvK0oWyV9NlwM3GLBjADyWgydVyg==", - "dev": true, - "license": "Apache-2.0", - "bin": { - "playwright-core": "cli.js" - }, - "engines": { - "node": ">=18" - } - }, "node_modules/points-on-curve": { "version": "0.2.0", "resolved": "https://registry.npmjs.org/points-on-curve/-/points-on-curve-0.2.0.tgz", @@ -13584,33 +13228,6 @@ "license": "MIT", "peer": true }, - "node_modules/react-markdown": { - "version": "10.1.0", - "resolved": "https://registry.npmjs.org/react-markdown/-/react-markdown-10.1.0.tgz", - "integrity": "sha512-qKxVopLT/TyA6BX3Ue5NwabOsAzm0Q7kAPwq6L+wWDwisYs7R8vZ0nRXqq6rkueboxpkjvLGU9fWifiX/ZZFxQ==", - "license": "MIT", - "dependencies": { - "@types/hast": "^3.0.0", - "@types/mdast": "^4.0.0", - "devlop": "^1.0.0", - "hast-util-to-jsx-runtime": "^2.0.0", - "html-url-attributes": "^3.0.0", - "mdast-util-to-hast": "^13.0.0", - "remark-parse": "^11.0.0", - "remark-rehype": "^11.0.0", - "unified": "^11.0.0", - "unist-util-visit": "^5.0.0", - "vfile": "^6.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/unified" - }, - "peerDependencies": { - "@types/react": ">=18", - "react": ">=18" - } - }, "node_modules/react-redux": { "version": "9.2.0", "resolved": "https://registry.npmjs.org/react-redux/-/react-redux-9.2.0.tgz", @@ -13893,48 +13510,6 @@ "url": "https://opencollective.com/unified" } }, - "node_modules/remark-cjk-friendly": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/remark-cjk-friendly/-/remark-cjk-friendly-2.0.1.tgz", - "integrity": "sha512-6WwkoQyZf/4j5k53zdFYrR8Ca+UVn992jXdLUSBDZR4eBpFhKyVxmA4gUHra/5fesjGIxrDhHesNr/sVoiiysA==", - "license": "MIT", - "dependencies": { - "micromark-extension-cjk-friendly": "2.0.1" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "@types/mdast": "^4.0.0", - "unified": "^11.0.0" - }, - "peerDependenciesMeta": { - "@types/mdast": { - "optional": true - } - } - }, - "node_modules/remark-cjk-friendly-gfm-strikethrough": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/remark-cjk-friendly-gfm-strikethrough/-/remark-cjk-friendly-gfm-strikethrough-2.0.1.tgz", - "integrity": "sha512-pWKj25O2eLXIL1aBupayl1fKhco+Brw8qWUWJPVB9EBzbQNd7nGLj0nLmJpggWsGLR5j5y40PIdjxby9IEYTuA==", - "license": "MIT", - "dependencies": { - "micromark-extension-cjk-friendly-gfm-strikethrough": "2.0.1" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "@types/mdast": "^4.0.0", - "unified": "^11.0.0" - }, - "peerDependenciesMeta": { - "@types/mdast": { - "optional": true - } - } - }, "node_modules/remark-gfm": { "version": "4.0.1", "resolved": "https://registry.npmjs.org/remark-gfm/-/remark-gfm-4.0.1.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index e8ede65526..d104aad157 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -18,8 +18,6 @@ "dependencies": { "@assistant-ui/core": "0.1.17", "@assistant-ui/react": "0.12.28", - "@assistant-ui/react-markdown": "0.12.11", - "@assistant-ui/react-streamdown": "0.1.11", "@assistant-ui/tap": "0.5.10", "@base-ui/react": "^1.2.0", "@dagrejs/dagre": "^2.0.4", @@ -30,13 +28,11 @@ "@hugeicons/core-free-icons": "^4.1.1", "@hugeicons/react": "^1.1.5", "@huggingface/hub": "^2.9.0", - "@langchain/core": "^1.1.27", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-label": "^2.1.8", "@radix-ui/react-select": "^2.2.6", "@radix-ui/react-separator": "^1.1.8", "@radix-ui/react-slot": "^1.2.4", - "@streamdown/cjk": "1.0.3", "@streamdown/code": "1.1.1", "@streamdown/math": "1.0.2", "@streamdown/mermaid": "1.0.2", @@ -50,14 +46,12 @@ "@tauri-apps/plugin-process": "^2.3.1", "@tauri-apps/plugin-updater": "^2.10.1", "@toolwind/corner-shape": "^0.0.8-3", - "@types/canvas-confetti": "^1.9.0", "@xyflow/react": "^12.10.0", "assistant-stream": "0.3.12", "canvas-confetti": "^1.9.4", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "cmdk": "^1.1.1", - "date-fns": "^4.1.0", "dexie": "^4.3.0", "js-yaml": "^4.1.1", "katex": "^0.16.28", @@ -72,7 +66,6 @@ "react-dom": "^19.2.4", "react-resizable-panels": "^4.6.4", "recharts": "3.7.0", - "remark-gfm": "^4.0.1", "shadcn": "^4.2.0", "sonner": "^2.0.7", "streamdown": "2.5.0", @@ -91,6 +84,7 @@ "devDependencies": { "@biomejs/biome": "^1.9.4", "@eslint/js": "^9.39.1", + "@types/canvas-confetti": "^1.9.0", "@types/js-yaml": "^4.0.9", "@types/node-forge": "^1.3.14", "@types/node": "^25.5.2", @@ -101,7 +95,6 @@ "eslint-plugin-react-hooks": "^7.0.1", "eslint-plugin-react-refresh": "^0.5.2", "globals": "^17.4.0", - "playwright": "^1.59.1", "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", "vite": "^8.0.1" From e87c0dd418a90f6661ae2b0371fdc8ebd6375fe7 Mon Sep 17 00:00:00 2001 From: stoicAI1776 <m.econyale@gmail.com> Date: Fri, 10 Apr 2026 06:10:45 +0530 Subject: [PATCH 41/50] fix(pyproject): use triton-xpu 3.6.0 for intelgputorch210 (#4931) --- pyproject.toml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c66cb870eb..cd66a8aeb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1021,14 +1021,14 @@ intelgputorch210 = [ "unsloth_zoo[intelgpu]", "unsloth[huggingfacenotorch]", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=c169a1de14c19673b17c751290d467fa282fc90fa5da4314b2e5cdab1f553146 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=013d9dd5d6479bd22983161f462e61c8dbe1d82e6730624a7a8d5945507eaa61 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=afc8cabfbf7ed51fd278d1e0f88d6afc157b0201bad4b99d681e4d542f9e66d4 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=0d24c1716088f2764d0d24c64227732195b6a42706c3c5fc89eeb4904bfa0818 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-win_amd64.whl#sha256=c83ab007311d9cfb6e809ee5a4587d99a9eef4be720b90da4f1aaa68b45139a0 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-win_amd64.whl#sha256=debf75348da8e8c7166b4d4a9b91d1508bb8d6581e339f79f7604b2e6746bacd ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-win_amd64.whl#sha256=97337a47425f1963a723475bd61037460e84ba01db4f87a1d662c3718ff6c47e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-win_amd64.whl#sha256=2caf8138695f6abb023ecd02031a2611ba1bf8fff2f19802567cb2fadefe9e87 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=abb1d1ec1ac672bac0ff35420c965f2df0c636ef9d94e2a830e34578489d0a57 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=71ad2f82da0f41eaec159f39fc85854e27c2391efa91b373e550648a6f4aaad3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", From f7069b242c951a20b7300ea361a3e771fabb94d8 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sat, 16 May 2026 13:00:25 +0000 Subject: [PATCH 42/50] intel-gpu: add xpu extras for torch 2.7.1 / 2.9.1 / 2.11.0 / 2.12.0 Builds on the intelgputorch210 fix in the previous commit by also filling in the four other intel-gpu-torch* wheel stacks that PyTorch has published but Unsloth was not yet exposing: * intelgputorch271 / intel-gpu-torch271: torch 2.7.1+xpu with pytorch-triton-xpu 3.3.1 * intelgputorch291 / intel-gpu-torch291: torch 2.9.1+xpu with pytorch-triton-xpu 3.5.0 * intelgputorch2110 / intel-gpu-torch2110: torch 2.11.0+xpu with triton-xpu 3.7.0 * intelgputorch2120 / intel-gpu-torch2120: torch 2.12.0+xpu with triton-xpu 3.7.1 Each torch + triton pair was cross-checked against the wheel's PEP 658 .metadata sidecar to confirm the upstream Requires-Dist matches what this block pins. Verified via `uv pip compile` across cp310/cp311/cp312/ cp313 on Linux x86_64 and Windows AMD64. --- pyproject.toml | 148 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index cd66a8aeb8..3468f7f8a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1017,6 +1017,80 @@ intelgputorch290 = [ intel-gpu-torch290 = [ "unsloth[intelgputorch290]" ] +intelgputorch271 = [ + "unsloth_zoo[intelgpu]", + "unsloth[huggingfacenotorch]", + + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=663ce21364096b268c6687f26f22862cb1001cae0c4ec9f98a0998415f99e2b0 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=dd92cc17000bad19f213b6a877d7f10cd71341b703cd188513ce9fff8d42e3dd ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=aa5c3ec21a89e967d1dfe61e3d5b1c1ae9620c871ed804771d3378d6a44066f2 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=d1c6f522e11112a311b1a61ba7b40b43ad8305675fa29153017ccb1ad0b6816d ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-win_amd64.whl#sha256=a5c16dcf449a9cb62bc3788f7ec45782bb3ead6edc2637a12b60ef0f8f45dc55 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-win_amd64.whl#sha256=bc2d76ffa4ceed5b38ae34b52dbff643442e1a44d52ca72d7cb520ca1950e9ae ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-win_amd64.whl#sha256=b09ca59ce52d6d27b1510df783cde222b703a71857a6fa953f1f155f9f50811a ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-win_amd64.whl#sha256=1260c4a4bad426b6cd3c8f3e1a21835381c6f217bf434bcb55fedec08a206dea ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=231c3fbd88a75d94de5ccbbb7f4f9a96cb3c58b3d891c2a1b469d38df95f9be6 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=78edcc27709dd819fc820f5eb9421bd10d3f3dcb14adb25ee60766c76f0e67f3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=b443df40bc9cb7d648a9f8f9ed1d5c3a1203e561ebd0a61dd55fb8a58833d5ec ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=412b58ffcceebea399c9a1bcdb22896aa10385c2650a8c4f8a677fb11c49b448 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2591228dc2cb73c78daf24277c4449ba9474f94cd31938147249269fe89d05d6 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=1aacb86e9a9684ffc8bde3db14b251d00df7019a9a434ec99a59076a2696325d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=9b65dc8562521b60d77aa653132bc03a19da0291318fcf919faa3f03080d8f7e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd3669fee311bc3ee5501d696bf989226a6f2bf957d120a04881a07af05526d6 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=f8cdf6889c02b3166679eef661b68757ea7e99c314432c3d41dac3d2ed4a59d4 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=f7d15b65d52809745992e0001c25034f33ac01f2dff5248614e07b5d009a59b7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=1ff1f98d70846352c7f56833bedab1a055ead27b11c120b8c719063ee0383554 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f46945344ea911a70309231eaaf3b80c96f6646ce5515dc89aa94f94144e310e ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=ecae9a02de769e2070d37388116beb407c3f0d60b8e65c1da1423f4eafee361a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=2914e62782431bebd6ad9a3b98a2b7311e448e84a7534bb7f35874b9279a17de ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=5b462c156f4e2097e1e53649d3f298ce352fa4c5d1e6addd360375b10ebd6c67 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=fa87b3677cd1af67ce423004283c1bde80e3571f391182a3e89b485e18e3c70f ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] +intel-gpu-torch271 = [ + "unsloth[intelgputorch271]" +] +intelgputorch291 = [ + "unsloth_zoo[intelgpu]", + "unsloth[huggingfacenotorch]", + + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=c169a1de14c19673b17c751290d467fa282fc90fa5da4314b2e5cdab1f553146 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=013d9dd5d6479bd22983161f462e61c8dbe1d82e6730624a7a8d5945507eaa61 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=afc8cabfbf7ed51fd278d1e0f88d6afc157b0201bad4b99d681e4d542f9e66d4 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=0d24c1716088f2764d0d24c64227732195b6a42706c3c5fc89eeb4904bfa0818 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-win_amd64.whl#sha256=c83ab007311d9cfb6e809ee5a4587d99a9eef4be720b90da4f1aaa68b45139a0 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-win_amd64.whl#sha256=debf75348da8e8c7166b4d4a9b91d1508bb8d6581e339f79f7604b2e6746bacd ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-win_amd64.whl#sha256=97337a47425f1963a723475bd61037460e84ba01db4f87a1d662c3718ff6c47e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-win_amd64.whl#sha256=2caf8138695f6abb023ecd02031a2611ba1bf8fff2f19802567cb2fadefe9e87 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=fb7895c744132d6a8e56ce8434ae1d8355c9bda4e9f58832744ff742d6268eaf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=da2604a9114a28de71ce654819424d20a246adf644d191ae160837df9731b79e ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=d5968d78d81c1d01efc1b3bf83d7da3d83161dcc3a9fcf91f500591db1c6c75d ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=b56d6b0d65863f370527e971dbfa046a5dd2a1f61cc95071db26c764f36e4dce ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2f318fb6a4bf1101cc17f35a5371f7c1768b41fceed03628397834e85b3edfdd ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=c9cedc3fb099366b2e6c563df6578e323564b1b5d40ac27be73c674755343a1d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=bee9623254d0f95a1ca115dbd17e9a9d966fdb8ae123e2ada4a9eb2fb8d38db8 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd5c857da52a63c121561b30b0979e69ade70b575fd74e389787bc7c1ee2ac11 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=cc5272da2cb4554edf059eedd6d1f5ef2859033b0fb79d5dcb8e99a0697f3325 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=3c80d6a068c32fc4ebddb27953e03a0141bd0f10ca8730417cbc0e0748158285 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=8cf640a867cf270b3fda7a10002c29d3fc2ad6dfbd76404a8cdd820489adb04c ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=d9c59ee5ae3d0560f02401c8dfd8054d50813a8dbb5d33a8777de7d02f6fcb7b ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=843ea7fcd8f5a22ebbc20d2d61d9eec7593821a0372eb8cabb73953d12ef6acf ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=e5ff8a31d3c700f8dbac59697c8e32298a43ec059609ebc6ea7bab3eff6384e1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=8bae6d4c042f8d20818da4a5aa9109c6fbd6ec11bc422be152ce8adf9a7095bf ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=47059e290fc2a41ba78666ffcde102c436abf7ff8a34d200268b48c4fa0f9c45 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] +intel-gpu-torch291 = [ + "unsloth[intelgputorch291]" +] intelgputorch210 = [ "unsloth_zoo[intelgpu]", "unsloth[huggingfacenotorch]", @@ -1054,6 +1128,80 @@ intelgputorch210 = [ intel-gpu-torch210 = [ "unsloth[intelgputorch210]" ] +intelgputorch2110 = [ + "unsloth_zoo[intelgpu]", + "unsloth[huggingfacenotorch]", + + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=2a1841138750f708ec017becbf8d357526f3fa350deee6553be5735ad66160a3 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e85378f1fc1ea002271de2a35475b75008fa554b86ef9d3bc55be9c513a63b51 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a6663ebe43e3c0d560ff774708632d7a75208ee64a291c1724ed5c16a92d1c72 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=08c8d43b2831faf9d6799480df2b45dde58102257aebd810d07a2ce18cd4e5df ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-win_amd64.whl#sha256=90fb8f767950a4ffca627faa7f86d9c697237ea4352d7e23505c5c9ed8e72216 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-win_amd64.whl#sha256=aa7de82f4265089e74f25a2701b7532e5c47d74224d877b61da1d66156e3f0c1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-win_amd64.whl#sha256=5ba3a31c6e1b259ad2d924e1b50f72a78c6ebd7eb4f364473bbf93e144734e80 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-win_amd64.whl#sha256=e8b4caba9b2399ea4c7f9a2777042564dea5d6f9e586a2dcb015a4ce20f000f7 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=6e634354b752b7366e8ad16b84f3e7e5863776a7ab448bbabae4fd36668dee7a ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=293169899f562ce473a58836dd024f0b1e72a347400278287ab393d1b04991e4 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e204d14be6f0f84d5f0e6e9213556e80326c3ab682cac108bcbef340bf45297b ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f134344006f0989a2d771554b7905fb05bd93d63b195e64626fde3495ec6f287 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=7e52729cb9736c66dc79a7f42de6b31db93b9161d3357fd34cfa33f5fe32b8ea ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=83a6130100c6b6750d8aa9fd29e5d0c53b1c85b1153b8ed4139aea54fc1892cc ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=03788e0e5a5b85a2f09d11f0263d579fcb0cf5623d8810149be0e37836c2738c ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cb1da1d378ce440f7d1e0ed8cf21bd280d904ab25a55c9453f8377825818df74 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] +intel-gpu-torch2110 = [ + "unsloth[intelgputorch2110]" +] +intelgputorch2120 = [ + "unsloth_zoo[intelgpu]", + "unsloth[huggingfacenotorch]", + + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=f59decc04bec27862ed0197554a52370dbcba3e6892616d1fbce450e402bf2d5 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=56f74e7c6c096e1a7ac215eb79ee590b764be3fbba8f4febc145bca47194a083 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=b9779b71457b5a916ae052ed2467c10273cae4862d469b191359173b2038c53e ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=7ef8e776c992e4e3ae007ebc108eb4f36b1d1dd9da97ecb308ab7fded89a2659 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=7f1d40febf2b8724adf4ff23866897d87478cc43de2a20f7776dc00be334c464 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=32770e2613df26e2c81ae64ea001b2ca12b8d152231285caff9b5f963a21ad75 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=0d517462caf6f5201c0d7c880f4ac431783c88fcc59b4587836da6c72a89509c ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=4b6feada86aa0bd606904b05898b33538106120d8ed706ba11d0011046534cb8 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e231819be0f87829c2344c909c1f0db9d6ae7d6faefe644a526a1a01d0c18d98 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=8bc7d37515cea18af4c389d5fde58b1a9d76b015f2d87e4a7dc62ad50b1cc200 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=65dbb041057dddfe369f29cfaab63f75563621779a23a7b1e2c0ff8a84d4376a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=df647445365924d69fe3bb2a15a7edfe5b63ef91e4ae69af11d93582985237a4 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=b0db3df0d0d154d18ba988ab420f1da2549f9372113ff54ff66e4ae3c7fe3bd0 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=c70850842068c43a0d50eaf139c25b6f6cc9b17a0dae70218c7e69edbee0bc80 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] +intel-gpu-torch2120 = [ + "unsloth[intelgputorch2120]" +] intel = [ "unsloth[intelgputorch280]", ] From 62def21608ea10b0761febdbb79840446cf257a5 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Sat, 16 May 2026 16:14:10 +0100 Subject: [PATCH 43/50] Studio: auto-load models when adding a cloud provider (#5472) * studio: auto-load models when adding a cloud provider * fix: drop auto-load --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- .../features/chat/chat-providers-dialog.tsx | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index ea99ada0e0..96f95d6d7b 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -246,13 +246,16 @@ export function ChatProvidersSettings({ } return; } - // Seed the registry's default_models for every provider — curated and - // remote alike. For remote-mode providers, loadModels() will replace - // this with the union of defaults + the live /models response once the - // user clicks "Load Models"; until then (or if the call fails — e.g. - // decryption issues during key rotation) the seeded list ensures - // curated picks like claude-haiku-4-5 are always reachable. - setAvailableModels([...entry.default_models]); + // Seed default_models only when the catalog is not fetched live: + // curated providers (catalog too large to enumerate, defaults are + // the suggestion shortlist) and Ollama (local, no API key — local + // /models stands in). Remote-mode cloud providers stay empty until + // the user clicks "Load available models" with a key, since + // different API tiers expose different catalogs and we don't want + // to advertise models the user can't actually call. + const seedDefaults = + entry.model_list_mode === "curated" || providerType === "ollama"; + setAvailableModels(seedDefaults ? [...entry.default_models] : []); setSelectedModelIds([]); setManualModelIds(""); setModelSearchQuery(""); @@ -362,8 +365,11 @@ export function ChatProvidersSettings({ resetForm(); const entry = providerType ? registryByType.get(providerType) : null; if (entry) { - // Keep first-open behavior consistent with provider re-selection. - setAvailableModels([...entry.default_models]); + const seedDefaults = + entry.model_list_mode === "curated" || providerType === "ollama"; + if (seedDefaults) { + setAvailableModels([...entry.default_models]); + } } setPage("form"); } From bc28a5d80e3825a534376ea7706f0da0161d9a8f Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Sat, 16 May 2026 18:14:01 +0100 Subject: [PATCH 44/50] Studio: code execution config visual polish (#5471) * style: polish code execution * studio/chat: optimistic insert for created OpenAI containers Container creation now prepends the row immediately with a "Creating" pill instead of waiting on /v1/containers, which is eventually consistent and can lag the create response by several seconds. A 5s follow-up refresh reconciles with the server view. --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai> --- .../src/features/chat/chat-settings-sheet.tsx | 2 +- .../components/openai-code-exec-section.tsx | 234 ++++++++++-------- 2 files changed, 137 insertions(+), 99 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 6de7d6a330..394601f07b 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -243,7 +243,7 @@ function loadSavedActivePreset(): string { } } -function InfoHint({ children }: { children: ReactNode }) { +export function InfoHint({ children }: { children: ReactNode }) { return ( <Tooltip> <TooltipTrigger asChild> diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx index d76b4a282f..21c3bb1ed0 100644 --- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -29,7 +29,7 @@ "use client"; -import { useCallback, useEffect, useMemo, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { AlertDialog, @@ -43,7 +43,6 @@ import { } from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; -import { Skeleton } from "@/components/ui/skeleton"; import { TrashIcon, RefreshCwIcon, PlusIcon } from "lucide-react"; import { createOpenAIContainer, @@ -55,6 +54,7 @@ import { db } from "../db"; import type { ExternalProviderConfig } from "../external-providers"; import { useLiveQuery } from "../db"; import { ensureThreadRecord } from "../runtime-provider"; +import { InfoHint } from "../chat-settings-sheet"; const AUTO_OPTION_VALUE = "__auto__"; const DEFAULT_TTL_MINUTES = 20; @@ -111,9 +111,6 @@ export function OpenAICodeExecSection({ const [creating, setCreating] = useState(false); const [createOpen, setCreateOpen] = useState(false); const [createName, setCreateName] = useState(""); - const [createTtl, setCreateTtl] = useState<number>( - provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES, - ); // Ids that have been deleted in this session. Once tombstoned, an id // stays hidden from the picker for the lifetime of the page — OpenAI's // /containers list can keep returning a freshly-deleted id for an @@ -121,6 +118,22 @@ export function OpenAICodeExecSection({ // creates more confusion than it solves. Refreshing the page resets // the tombstone naturally. const [tombstones, setTombstones] = useState<Set<string>>(() => new Set()); + // Ids optimistically inserted after a successful create but not yet + // confirmed by a /v1/containers list response. OpenAI's list endpoint + // is eventually consistent — a freshly-created container can be absent + // for several seconds. We render the row immediately with a "Creating" + // pill, then drop it from this set once a refresh sees the id. + const [pendingIds, setPendingIds] = useState<Set<string>>(() => new Set()); + // Ref mirror so `refresh()` can read the current pending set without + // re-binding when it changes (the callback is in a useEffect dep). + const pendingIdsRef = useRef<Set<string>>(pendingIds); + useEffect(() => { + pendingIdsRef.current = pendingIds; + }, [pendingIds]); + // One-shot follow-up refresh scheduled after a create, to catch the + // common case where the server list lags the create response by a few + // seconds. Tracked so we can clear it on unmount. + const pendingRetryRef = useRef<number | null>(null); // Target row for the destructive confirmation dialog. Held in state // (rather than blocking with window.confirm) so the dialog sits inside // the settings sheet instead of a native browser alert. @@ -182,7 +195,24 @@ export function OpenAICodeExecSection({ apiKey, baseUrl: provider.baseUrl || null, }); - setContainers(list); + const serverIds = new Set(list.map((c) => c.id)); + setContainers((prev) => { + // Preserve optimistic inserts the server hasn't acknowledged + // yet so they don't disappear on the reconciling refresh. + const orphans = prev.filter( + (c) => !serverIds.has(c.id) && pendingIdsRef.current.has(c.id), + ); + return orphans.length > 0 ? [...orphans, ...list] : list; + }); + setPendingIds((prev) => { + if (prev.size === 0) return prev; + const next = new Set(prev); + let changed = false; + for (const id of serverIds) { + if (next.delete(id)) changed = true; + } + return changed ? next : prev; + }); } catch (err) { toast.error( `Failed to list containers: ${err instanceof Error ? err.message : "Unknown"}`, @@ -213,6 +243,10 @@ export function OpenAICodeExecSection({ return () => { window.clearInterval(interval); document.removeEventListener("visibilitychange", onVisibility); + if (pendingRetryRef.current != null) { + window.clearTimeout(pendingRetryRef.current); + pendingRetryRef.current = null; + } }; }, [refresh]); @@ -300,15 +334,43 @@ export function OpenAICodeExecSection({ toast.error("Container name is required"); return; } + // TTL inherits from the section-level "Idle timeout" control — + // there is no per-container override on the form. Read it at + // submit time so a last-second change to the TTL row applies. + const ttlMinutes = + provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES; setCreating(true); try { const created = await createOpenAIContainer( { apiKey, baseUrl: provider.baseUrl || null }, - { name, ttlMinutes: createTtl }, + { name, ttlMinutes }, ); toast.success(`Created container ${name}`); setCreateName(""); setCreateOpen(false); + // Optimistic insert + "Creating" pill. OpenAI's /v1/containers + // list endpoint is eventually consistent and can omit the new + // container for several seconds — without this, the row only + // shows up on the next 30s poll or a manual refresh. + setContainers((prev) => + prev.some((c) => c.id === created.id) ? prev : [created, ...prev], + ); + setPendingIds((prev) => { + if (prev.has(created.id)) return prev; + const next = new Set(prev); + next.add(created.id); + return next; + }); + // Follow-up refresh ~5s later to reconcile the optimistic row + // with the server's view once /v1/containers catches up. One + // shot; the regular poll covers any longer tail. + if (pendingRetryRef.current != null) { + window.clearTimeout(pendingRetryRef.current); + } + pendingRetryRef.current = window.setTimeout(() => { + pendingRetryRef.current = null; + void refresh(); + }, 5000); // Auto-bind the just-created container to the active thread. // ensureThreadRecord first so the bind lands even when the user // creates a container before sending the first message — without @@ -389,12 +451,18 @@ export function OpenAICodeExecSection({ <div className="flex flex-col gap-3 pt-1"> {/* TTL */} <div className="flex items-center justify-between gap-3"> - <label - htmlFor="openai-container-ttl" - className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg" - > - New-container idle timeout (min, max 20) - </label> + <div className="flex min-w-0 items-center gap-1.5"> + <label + htmlFor="openai-container-ttl" + className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg" + > + Idle timeout + </label> + <InfoHint> + Minutes a newly-created container stays alive between calls. + OpenAI caps this at 20. + </InfoHint> + </div> <Input id="openai-container-ttl" type="number" @@ -406,10 +474,9 @@ export function OpenAICodeExecSection({ /> </div> - {/* Single container list. The previously-separate "Active for - this thread" picker collapses into this list: clicking a row - binds it to the active thread, and the ACTIVE pill marks - which one. Avoids duplicating state across two controls. */} + {/* Single container list. Clicking a row binds it to the active + thread and the ACTIVE pill marks which one — no separate + picker needed. */} <div className="flex flex-col gap-1.5"> <div className="flex items-center justify-between gap-2"> <span className="text-[11px] uppercase tracking-wider text-muted-foreground"> @@ -428,45 +495,20 @@ export function OpenAICodeExecSection({ /> </Button> </div> - {/* When no containers exist yet, render a disabled placeholder - instead of the picker. The first one is created by the - chat-adapter on first send (lazy-create) and will appear - here after the next refresh. */} {sortedContainers.length === 0 ? ( - <div className="h-9 w-full rounded-md border border-primary/40 bg-background px-2 flex items-center text-sm text-muted-foreground"> - (none yet — will be created on first send) + // Quiet placeholder with the same muted border as row cards + // so an empty section doesn't masquerade as an active control. + // The first container is minted by the chat-adapter on first + // send (lazy-create) and appears here after the next refresh. + <div className="flex h-9 w-full items-center rounded-md border border-dashed border-border/60 bg-muted/20 px-2 text-xs text-muted-foreground"> + None yet - one will be created on first send. </div> ) : ( - <select - value={displayedContainerId ?? sortedContainers[0].id} - onChange={(e) => onPick(e.target.value)} - disabled={!activeThreadId} - className="h-9 w-full rounded-md border border-primary/40 bg-background px-2 text-sm font-medium shadow-sm focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-primary/40" - > - {sortedContainers.map((c) => ( - <option key={c.id} value={c.id}> - {c.name ?? "(unnamed)"} · {c.id.slice(0, 14)}… - {c.lastActiveAt ? ` · active ${ageLabel(c.lastActiveAt)}` : ""} - </option> - ))} - </select> - )} - </div> - - {/* Container list with delete actions — labeled and visually - quieter so it's clearly the "all containers, manage them" - area rather than the active selector above. */} - <div className="flex flex-col gap-1.5"> - <span className="text-[11px] uppercase tracking-wider text-muted-foreground"> - All containers - </span> - {isLoading && visibleContainers.length === 0 ? ( - <Skeleton className="h-16 w-full" /> - ) : sortedContainers.length > 0 ? ( <ul className="flex max-h-52 flex-col gap-1 overflow-auto"> {sortedContainers.map((c) => { const running = isContainerRunning(c); const isActive = running && c.id === displayActiveId; + const isPending = pendingIds.has(c.id); const ttlMinutes = c.expiresAfterMinutes ?? DEFAULT_TTL_MINUTES; const canActivate = activeThreadId != null && !isActive && running; @@ -510,7 +552,11 @@ export function OpenAICodeExecSection({ <span className="min-w-0 truncate font-medium"> {c.name ?? "(unnamed)"} </span> - {isActive ? ( + {isPending ? ( + <span className="shrink-0 rounded-sm bg-muted px-1 py-px text-[9px] font-medium uppercase tracking-wider text-muted-foreground"> + Creating + </span> + ) : isActive ? ( <span className="shrink-0 rounded-sm bg-primary/15 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-primary"> Active </span> @@ -548,69 +594,61 @@ export function OpenAICodeExecSection({ ); })} </ul> - ) : ( - <p className="text-xs text-muted-foreground"> - None yet — one will be created on first send. - </p> )} </div> - {/* Create new */} + {/* Create new — inline single-row edit that visually echoes a + container card. TTL is inherited from the section's top + "Idle timeout" control (no per-container override), which + keeps the form light and avoids a duplicated input. */} {createOpen ? ( - <div className="flex flex-col gap-2 rounded-md border border-border/60 p-2"> + <div className="flex items-center gap-1 rounded-md border border-border/60 bg-muted/20 px-1.5 py-1"> <Input - placeholder="Container name (e.g. data-analysis)" + autoFocus + placeholder="Name" value={createName} onChange={(e) => setCreateName(e.target.value)} - className="h-8 text-sm" - /> - <div className="flex items-center gap-2"> - <Input - type="number" - min={TTL_MIN} - max={TTL_MAX} - value={createTtl} - onChange={(e) => { - const n = parseInt(e.target.value, 10); - if (!Number.isNaN(n)) - setCreateTtl(Math.min(Math.max(n, TTL_MIN), TTL_MAX)); - }} - className="h-8 w-24 text-sm" - aria-label="Idle timeout in minutes" - /> - <span className="text-xs text-muted-foreground">min idle</span> - <div className="flex-1" /> - <Button - size="sm" - variant="ghost" - className="h-7" - onClick={() => { + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + if (createName.trim() && !creating && apiKey) { + void onCreate(); + } + } else if (e.key === "Escape") { + e.preventDefault(); setCreateOpen(false); setCreateName(""); - }} - disabled={creating} - > - Cancel - </Button> - <Button - size="sm" - className="h-7" - onClick={() => void onCreate()} - disabled={creating || !createName.trim() || !apiKey} - > - Create - </Button> - </div> + } + }} + className="h-7 min-w-0 flex-1 border-0 bg-transparent px-1.5 text-xs shadow-none focus-visible:ring-0" + /> + <Button + size="sm" + variant="ghost" + className="h-7 shrink-0 px-2 text-xs" + onClick={() => { + setCreateOpen(false); + setCreateName(""); + }} + disabled={creating} + > + Cancel + </Button> + <Button + size="sm" + className="h-7 shrink-0 px-3 text-xs" + onClick={() => void onCreate()} + disabled={creating || !createName.trim() || !apiKey} + > + {creating ? "Creating…" : "Create"} + </Button> </div> ) : ( <Button size="sm" variant="outline" className="h-8" - onClick={() => { - setCreateTtl(ttlValue); - setCreateOpen(true); - }} + onClick={() => setCreateOpen(true)} disabled={!apiKey} > <PlusIcon className="size-3.5 mr-1" /> From aba57d0872d75e167d375db004b4d64f676bd11f Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sat, 16 May 2026 16:24:48 -0700 Subject: [PATCH 45/50] disable_torchcodec_if_broken: also patch datasets and clean sys.modules (#5483) * disable_torchcodec_if_broken: also patch datasets and clean sys.modules (#5446) transformers's _torchcodec_available was being flipped to False already, but datasets keeps its own datasets.config.TORCHCODEC_AVAILABLE flag (datasets >= 4.0) that gates every torchcodec call site inside datasets/features/{audio,video}.py, datasets/features/features.py and the three datasets formatters. Without flipping that flag, transformers falls back to librosa but datasets still routes through torchcodec and re-raises the same RuntimeError, which is what users hit on Colab when libavutil is missing. Also pops half-loaded torchcodec submodules + datasets.features._torchcodec from sys.modules so a later re-import does not re-trigger the failed native-library dlopen. Verified live state after the patch on a Colab-like broken-torchcodec env: transformers._torchcodec_available = False transformers.is_torchcodec_available() = False datasets.config.TORCHCODEC_AVAILABLE = False stale torchcodec entries in sys.modules = [] * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disable_torchcodec_if_broken: seat sys.modules[torchcodec]=None sentinel The first commit on this branch flipped the transformers and datasets availability flags, but a few unconditional torchcodec call sites in datasets / torchaudio (e.g. Audio.encode_example does "from torchcodec.encoders import AudioEncoder" outside the TORCHCODEC_AVAILABLE gate) still hit a cryptic RuntimeError from the broken native library load. Seating sys.modules["torchcodec"] = None makes any subsequent "import torchcodec" / "from torchcodec.X import Y" raise ModuleNotFoundError (subclass of ImportError) which the existing try/except ImportError blocks in datasets / torchaudio catch and re-raise as the clean "please install torchcodec" message users get when they uninstall torchcodec manually. This was the workaround in issue #5446. Also makes find_spec("torchcodec") return None on re-entry, so the function is a strict no-op on the second call. Verified across 12 scenarios in temp/torchcodec_test/: healthy torchcodec untouched, broken torchcodec sentinel-blocked, datasets 3.x without the flag handled, no-datasets-installed handled, py3.11 Colab- exact env (transformers==4.56.2 + trl==0.22.2) handled. Side-by-side on the user's exact stack: Audio.encode_example switches from "RuntimeError: Could not load libtorchcodec" to "ImportError: To support encoding audio data, please install 'torchcodec'.". * disable_torchcodec_if_broken: trim comments Same behaviour, shorter docstring and inline comments. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/import_fixes.py | 72 +++++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 32 deletions(-) diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 8476edc40a..182355640a 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1289,53 +1289,61 @@ def patch_torchcodec_audio_decoder(): def disable_torchcodec_if_broken(): - """Disable torchcodec in transformers if it cannot actually load. + """Make broken torchcodec behave as if uninstalled (#5446). - transformers checks if torchcodec is installed via importlib.util.find_spec(), - but this returns True even when torchcodec cannot load its native libraries - (e.g., when FFmpeg is missing). This causes runtime errors when transformers - tries to use torchcodec for audio loading. - - This function tests if torchcodec can actually load and if not, patches - transformers to think torchcodec is unavailable so it falls back to librosa. - - Two shapes to cover: - * transformers < 5: a module-level ``_torchcodec_available`` flag - cached in ``transformers.utils.import_utils``; flip it to False. - * transformers >= 5: a public ``is_torchcodec_available()`` callable - wrapped with ``functools.lru_cache``; replace it with a stub that - returns False and clear the cache so subsequent callers see it. + transformers and datasets both detect torchcodec via find_spec, which + returns True even when the native libs cannot dlopen. We flip their + flags and seat a sys.modules sentinel so downstream imports fall through + their existing except ImportError handlers cleanly. """ try: import importlib.util if importlib.util.find_spec("torchcodec") is None: - return # torchcodec not installed, nothing to do + return # absent or already disabled - # Test if torchcodec can actually load + # RuntimeError on dlopen failure; OSError covers chained libavutil.so misses. from torchcodec.decoders import AudioDecoder except (ImportError, RuntimeError, OSError): - # torchcodec cannot load - disable it in transformers + # transformers: flip flag (<5) and/or rebind lru_cache'd func (>=5). try: import transformers.utils.import_utils as tf_import_utils - except ImportError: - return - # transformers < 5 path: module-level cached flag. - try: - tf_import_utils._torchcodec_available = False - except AttributeError: - pass - - # transformers >= 5 path: public lru_cache'd function. Clear any - # cached True result then rebind to a stub that returns False. - is_avail = getattr(tf_import_utils, "is_torchcodec_available", None) - if is_avail is not None: try: - is_avail.cache_clear() + tf_import_utils._torchcodec_available = False except AttributeError: pass - tf_import_utils.is_torchcodec_available = lambda: False + + is_avail = getattr(tf_import_utils, "is_torchcodec_available", None) + if is_avail is not None: + try: + is_avail.cache_clear() + except AttributeError: + pass + tf_import_utils.is_torchcodec_available = lambda: False + except ImportError: + pass + + # datasets >= 4.0: own flag gating audio/video/features/formatters. + try: + import datasets.config as datasets_config + + if hasattr(datasets_config, "TORCHCODEC_AVAILABLE"): + datasets_config.TORCHCODEC_AVAILABLE = False + except ImportError: + pass + + # Drop half-loaded entries and seat the absence sentinel. After this, + # import torchcodec raises ModuleNotFoundError and find_spec returns None. + for _stale in [ + n + for n in list(sys.modules) + if n == "torchcodec" + or n.startswith("torchcodec.") + or n == "datasets.features._torchcodec" + ]: + sys.modules.pop(_stale, None) + sys.modules["torchcodec"] = None def disable_broken_wandb(): From 79a431cd53cbae34c4026a4434af269d339f228b Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sun, 17 May 2026 01:35:28 -0700 Subject: [PATCH 46/50] tests: pinned-symbol canary for unsloth-zoo save_pretrained_merged guards (#5410) (#5433) * tests: pinned-symbol canary for unsloth-zoo save_pretrained_merged guards (#5410) unsloth#5410 was a class of silent-write bug in the save_pretrained_merged path that the existing CI matrix could not detect because the merge-helper tests were not wired through the upstream-drift suite. The full fix lives in unslothai/unsloth-zoo#647 (layout-aware MoE merge helpers, authoritative num_experts resolver, loud-fail counter, generation_config.json save). This PR adds the unsloth-side canary that watches for the four guards staying in place in unsloth-zoo so a future refactor cannot silently regress them. tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py fetches unsloth_zoo/saving_utils.py + tests/test_unsloth_zoo_lora_merge.py from unslothai/unsloth-zoo:main and asserts: - _MOE_MERGE_STATE / _reset_moe_merge_state / _record_moe_merge_fallback are still defined and a `raise RuntimeError(...MoE...)` still fires when fallback > 0. - _detect_moe_lora_layout exists and both "swapped" / "standard" branch labels are reachable in the source. - _resolve_num_experts_from_lora_stats is present AND its base_layer walk is bounded by `for _ in range(N):` (a cyclic ParamWrapper chain must not hang the merge). - merge_and_overwrite_lora still calls model.generation_config.save_pretrained(...). - tests/test_unsloth_zoo_lora_merge.py keeps the six PEFT 0.19+ standard-layout regression tests added in #647. - Local unsloth/save.py still names save_pretrained_merged and routes through merge_and_overwrite_lora (i.e. the entry point still reaches the upstream fix). While #647 is still open, the four symbol tests SKIP cleanly with a message naming #647. When #647 merges into unsloth-zoo main, the same tests automatically become hard gates and catch any future regression. The sixth test (local entry-point grep) passes today. CPU-only static fetch, ~0.1s. Wired into the existing peft-pinned-symbols job in .github/workflows/version-compat-ci.yml so it runs on every PR that touches unsloth/** and on the daily schedule. Local run: 1 passed, 5 skipped (expected; #647 open). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests/version_compat: relax MoE/generation_config regex to fit zoo#647 zoo#647 landed two layout changes that broke the pinned-symbol canary's exact-string regex matches but kept the underlying guarantees intact: - The post-loop MoE LoRA fallback `raise RuntimeError(...)` wraps the "MoE" wording onto a second line; the old `[^\n]*` did not cross newlines. Switch to `.*?` + re.DOTALL. - The generation_config save now binds the attr to a local var `gen_cfg = getattr(model, "generation_config", ...)` and calls `gen_cfg.save_pretrained(save_directory)`, so a literal `generation_config.save_pretrained(` substring no longer matches. Anchor on the conceptual operation: a `generation_config` mention followed (within a small char window) by a `.save_pretrained(` call. That is what the canary actually cares about. Verified locally: pytest tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py -> 2 passed (4 deselected) --------- Co-authored-by: Daniel Han-Chen <info@unsloth.ai> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/version-compat-ci.yml | 1 + ..._unsloth_zoo_save_merged_pinned_symbols.py | 128 ++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index 2fbdd15747..599b53df1d 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -127,6 +127,7 @@ jobs: run: | PYTHONPATH=. python -m pytest \ tests/version_compat/test_peft_pinned_symbols.py \ + tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py \ -v --tb=short st-pinned-symbols: diff --git a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py new file mode 100644 index 0000000000..19faa51119 --- /dev/null +++ b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Pinned-symbol canary for unsloth-zoo save_pretrained_merged guards +(unslothai/unsloth-zoo#647 / unslothai/unsloth#5410). Skips until #647 +lands, then becomes a hard gate. CPU-only static fetch.""" + +from __future__ import annotations + +import re + +import pytest + +from tests.version_compat._fetch import fetch_text + + +ZOO_TAG = "main" + + +def _fetch_saving_utils() -> str: + src = fetch_text("unslothai/unsloth-zoo", ZOO_TAG, "unsloth_zoo/saving_utils.py") + if src is None: + pytest.skip("unsloth_zoo/saving_utils.py not fetchable") + return src + + +def _fetch_merge_tests() -> str: + src = fetch_text( + "unslothai/unsloth-zoo", + ZOO_TAG, + "tests/test_unsloth_zoo_lora_merge.py", + ) + if src is None: + pytest.skip("tests/test_unsloth_zoo_lora_merge.py not fetchable") + return src + + +def _skip_until_pr_647_lands(src: str) -> None: + if not any( + m in src + for m in ( + "_MOE_MERGE_STATE", + "_detect_moe_lora_layout", + "_resolve_num_experts_from_lora_stats", + ) + ): + pytest.skip( + "unslothai/unsloth-zoo#647 has not yet merged into main; " + "tests auto-promote to hard gates once it lands." + ) + + +def test_zoo_saving_utils_has_moe_merge_state(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + for sym in ( + "_MOE_MERGE_STATE", + "_reset_moe_merge_state", + "_record_moe_merge_fallback", + ): + assert sym in src, f"{sym} missing from saving_utils.py (issue #5410 guard)." + # zoo#647 wraps the fallback guard's message onto a second line; + # allow the regex to span newlines via re.DOTALL. + assert re.search( + r"raise\s+RuntimeError\b.*?MoE", src, re.IGNORECASE | re.DOTALL + ), "no `raise RuntimeError(...MoE...)`; post-loop guard weakened." + + +def test_zoo_saving_utils_has_layout_detector(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + assert ( + "_detect_moe_lora_layout" in src + ), "_detect_moe_lora_layout removed (issue #5410)." + assert ( + '"swapped"' in src and '"standard"' in src + ), "one of the layout labels removed." + + +def test_zoo_saving_utils_has_num_experts_resolver(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + assert "_resolve_num_experts_from_lora_stats" in src, "resolver removed (#5410)." + assert re.search( + r"for\s+_\s+in\s+range\s*\(\s*\d+\s*\)", src + ), "resolver walk no longer bounded by `for _ in range(N):`." + + +def test_zoo_saving_utils_writes_generation_config(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + # zoo#647 binds the generation_config attr to a local var + # (`gen_cfg = getattr(model, "generation_config", ...); ... + # gen_cfg.save_pretrained(save_directory)`) so an exact + # `generation_config.save_pretrained(` substring no longer + # matches. Anchor on the conceptual operation: a `generation_config` + # mention plus a `.save_pretrained(` call nearby, which is what + # the canary actually cares about. + assert re.search( + r"generation_config[\s\S]{0,400}?\.save_pretrained\s*\(", src + ), "generation_config.json no longer saved (#5410)." + + +def test_zoo_lora_merge_tests_have_standard_layout_coverage(): + src = _fetch_merge_tests() + if "test_merge_moe_gate_expert_standard_layout" not in src: + pytest.skip("unslothai/unsloth-zoo#647 not yet merged; coverage appears later.") + for name in ( + "test_merge_moe_gate_expert_standard_layout", + "test_merge_moe_up_expert_standard_layout", + "test_merge_moe_down_proj_expert_standard_layout", + "test_detect_moe_lora_layout_classifies_both_conventions", + "test_moe_merge_fallback_counter_records_bad_layout", + "test_resolve_num_experts_walks_base_layer_chain", + ): + assert name in src, f"regression test `{name}` removed." + + +def test_unsloth_save_pretrained_merged_entry_point_exists(): + import pathlib + + save_py = pathlib.Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + if not save_py.is_file(): + pytest.skip(f"{save_py} not present") + text = save_py.read_text(encoding = "utf-8", errors = "replace") + assert "save_pretrained_merged" in text, "entry point removed from unsloth/save.py." + assert ( + "merge_and_overwrite_lora" in text + ), "no dispatch into unsloth_zoo merge; #647 bypassed." From e20bbeff9afdb0cc1f96093e558aac7e434278c4 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sun, 17 May 2026 01:40:20 -0700 Subject: [PATCH 47/50] intel-gpu: pin unsloth_zoo>=2026.5.2 via huggingfacenotorch (#5499) Fixes unslothai/unsloth#5494: installing any intel-gpu-torch* extra without also pulling `huggingface` or `colab-new` lets the resolver silently fall back to a stale unsloth_zoo (2026.3.6 in the original report) because no version floor is enforced on `unsloth_zoo[intelgpu]` in those blocks. unsloth_zoo 2026.5.2 (just released) also relaxes its own torch upper bound from <2.11.0 to <2.13.0, which is what unblocks the resolver for the intel-gpu-torch2110 and intel-gpu-torch2120 extras shipped in #5484. Adding the floor in `huggingfacenotorch` propagates it to every extra that includes the HF-without-torch base: amd, huggingface, and all nine intelgputorch* blocks. Single line, single source of truth. Requires unsloth_zoo 2026.5.2 to be on PyPI for end-to-end resolution. --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3468f7f8a7..81cf5ac215 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ triton = [ ] huggingfacenotorch = [ + "unsloth_zoo>=2026.5.2", "wheel>=0.42.0", "packaging", "numpy", From fb13275787451bd67bba88e546c94e4636f353e2 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Sun, 17 May 2026 14:02:36 +0300 Subject: [PATCH 48/50] fix(sentence_transformer): resume PEFT checkpoints under sentence-transformers >= 5.4 (#5454) Saves the base config.json next to adapter_config.json when checkpointing PEFT-wrapped sentence-transformer models, and overrides SentenceTransformerTrainer._load_from_checkpoint to load adapter weights via set_peft_model_state_dict and rebuild aux modules (Pooling, Normalize, Dense) from modules.json with strict type and path validation. Patches only activate on Unsloth-managed Transformer modules so non-Unsloth pipelines fall through to upstream behaviour. Fixes https://github.com/unslothai/unsloth/issues/5373 --- unsloth/models/sentence_transformer.py | 181 ++++++++++++++++++++++++- 1 file changed, 180 insertions(+), 1 deletion(-) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index c53e3a7a81..aafc416221 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -527,6 +527,45 @@ This sentence-transformers model was finetuned and converted to GGUF format usin class FastSentenceTransformer(FastModel): + @staticmethod + def _save_base_config_for_processor_resume(config, output_path): + """ + sentence-transformers >= 5.4 reloads Transformer modules through + AutoProcessor. Tokenizer-only checkpoint roots make AutoProcessor fall + back to AutoConfig, so PEFT adapter checkpoints still need the base + config.json next to adapter_config.json. + """ + if config is None or not getattr(config, "model_type", None): + return + if hasattr(config, "save_pretrained"): + config.save_pretrained(output_path) + elif hasattr(config, "to_json_file"): + config_path = os.path.join(output_path, "config.json") + config.to_json_file(config_path) + + @staticmethod + def _patch_transformer_module_save_config(transformer_module, base_config = None): + transformer_module._unsloth_st_managed = True + if base_config is not None and getattr(base_config, "model_type", None): + transformer_module._unsloth_base_config = base_config + + if getattr(transformer_module, "_unsloth_save_config_patched", False): + return transformer_module + + original_save = transformer_module.save + + def _save_with_base_config(self, output_path, *args, **kwargs): + original_save(output_path, *args, **kwargs) + FastSentenceTransformer._save_base_config_for_processor_resume( + getattr(self, "_unsloth_base_config", None), output_path + ) + + transformer_module.save = types.MethodType( + _save_with_base_config, transformer_module + ) + transformer_module._unsloth_save_config_patched = True + return transformer_module + @staticmethod def _read_pooling_mode(model_name, token): """ @@ -1157,6 +1196,9 @@ class FastSentenceTransformer(FastModel): config_keys.append(config_key) transformer_module.config_keys = config_keys transformer_module.save_in_root = True + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(model, "config", None) + ) if hasattr(model, "config"): model.config.tokenizer_class = tokenizer.__class__.__name__ @@ -1644,6 +1686,9 @@ class FastSentenceTransformer(FastModel): st_model._dtype = dtype st_model._load_in_4bit = load_in_4bit st_model.no_modules = False + FastSentenceTransformer._patch_transformer_module_save_config( + st_model[0], getattr(st_model[0].auto_model, "config", None) + ) # Add save methods def _save_pretrained_merged(self, save_directory, **save_kwargs): @@ -2067,6 +2112,9 @@ class FastSentenceTransformer(FastModel): transformer_module.model = peft_model else: transformer_module.auto_model = peft_model + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(inner_model, "config", None) + ) # Store compile info for auto-compile at trainer time # torch.compile is deferred until training starts so we can check max_steps @@ -2121,6 +2169,9 @@ class FastSentenceTransformer(FastModel): transformer_module.model = peft_model else: transformer_module.auto_model = peft_model + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(inner_model, "config", None) + ) return model else: return FastModel.get_peft_model( @@ -2235,5 +2286,133 @@ def _patch_sentence_transformer_trainer(): SentenceTransformerTrainer._unsloth_auto_compile_patched = True -# Auto-patch trainer on module import +def _patch_st_trainer_load_from_checkpoint(): + try: + from sentence_transformers import SentenceTransformerTrainer + except ImportError: + return + if getattr( + SentenceTransformerTrainer, "_unsloth_load_from_checkpoint_patched", False + ): + return + if not hasattr(SentenceTransformerTrainer, "_load_from_checkpoint"): + return + + _original = SentenceTransformerTrainer._load_from_checkpoint + + def _unsloth_load_from_checkpoint(self, checkpoint_path): + try: + from peft import PeftModel, load_peft_weights, set_peft_model_state_dict + except ImportError: + return _original(self, checkpoint_path) + + try: + mod0 = self.model[0] + except (IndexError, TypeError): + return _original(self, checkpoint_path) + + if isinstance(getattr(type(mod0), "auto_model", None), property): + inner = getattr(mod0, "model", None) + else: + inner = getattr(mod0, "auto_model", None) + inner = getattr(inner, "_orig_mod", inner) + + if not isinstance(inner, PeftModel): + return _original(self, checkpoint_path) + if not getattr(mod0, "_unsloth_st_managed", False): + return _original(self, checkpoint_path) + + if not any( + os.path.isfile(os.path.join(checkpoint_path, fn)) + for fn in ("adapter_model.safetensors", "adapter_model.bin") + ): + return _original(self, checkpoint_path) + + adapter_name = getattr(inner, "active_adapter", None) + if adapter_name is None and callable(getattr(inner, "active_adapters", None)): + adapter_name = inner.active_adapters() + if isinstance(adapter_name, (list, tuple, set)): + if len(adapter_name) != 1: + raise RuntimeError( + "Unsloth: Cannot resume multiple active PEFT adapters." + ) + adapter_name = next(iter(adapter_name)) + adapter_name = adapter_name or "default" + if adapter_name not in getattr(inner, "peft_config", {}): + raise RuntimeError(f"Unsloth: PEFT adapter {adapter_name!r} is not loaded.") + + load_result = set_peft_model_state_dict( + inner, load_peft_weights(checkpoint_path), adapter_name = adapter_name + ) + unexpected = getattr(load_result, "unexpected_keys", []) or [] + missing = [ + x + for x in (getattr(load_result, "missing_keys", []) or []) + if f".{adapter_name}." in x or x.endswith(f".{adapter_name}") + ] + if unexpected or missing: + raise RuntimeError( + "Unsloth: PEFT checkpoint does not match the active adapter " + f"(missing={missing[:8]}, unexpected={unexpected[:8]})." + ) + + modules_json = os.path.join(checkpoint_path, "modules.json") + if not os.path.isfile(modules_json): + raise RuntimeError("Unsloth: PEFT checkpoint is missing modules.json.") + try: + with open(modules_json, "r") as f: + module_configs = json.load(f) + except Exception as e: + raise RuntimeError("Unsloth: Cannot parse checkpoint modules.json.") from e + + root = os.path.abspath(os.fspath(checkpoint_path)) + restored = set() + for entry in module_configs: + idx = int(entry.get("idx", -1)) + if idx == 0: + continue + if idx < 0 or idx >= len(self.model): + raise RuntimeError(f"Unsloth: Bad module index in modules.json: {idx}.") + module = self.model[idx] + module_cls = type(module) + saved_type = entry.get("type", "") + if saved_type and not saved_type.endswith(f".{module_cls.__name__}"): + raise RuntimeError(f"Unsloth: Checkpoint module {idx} type mismatch.") + module_path = entry.get("path") + module_dir = os.path.abspath( + os.path.join(root, os.fspath(module_path or "")) + ) + try: + inside_root = os.path.commonpath([root, module_dir]) == root + except ValueError: + inside_root = False + if not module_path or not inside_root or not os.path.isdir(module_dir): + raise RuntimeError( + f"Unsloth: Bad checkpoint module path for index {idx}." + ) + if not hasattr(module_cls, "load"): + raise RuntimeError(f"Unsloth: Module {idx} cannot be reloaded.") + fresh = module_cls.load(module_dir) + if not isinstance(fresh, module_cls): + raise RuntimeError(f"Unsloth: Module {idx} reload returned wrong type.") + # Parameterless modules (Pooling, Normalize) make + # next(module.parameters()) raise StopIteration; route through + # the SentenceTransformer's device property instead. + try: + fresh.to(self.model.device) + except AttributeError: + pass + self.model[idx] = fresh + restored.add(idx) + missing_idx = sorted(set(range(1, len(self.model))) - restored) + if missing_idx: + raise RuntimeError( + f"Unsloth: Checkpoint modules.json is incomplete (missing idx={missing_idx[:8]})." + ) + + SentenceTransformerTrainer._load_from_checkpoint = _unsloth_load_from_checkpoint + SentenceTransformerTrainer._unsloth_load_from_checkpoint_patched = True + + _patch_sentence_transformer_trainer() +_patch_st_trainer_load_from_checkpoint() From ab56e4a9edaeae839d6b229429b30e36605b26a7 Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sun, 17 May 2026 04:16:09 -0700 Subject: [PATCH 49/50] Studio: serialise GGUF reload and inherit unsloth-run extra args (#5427) * Studio: serialise GGUF reload and inherit unsloth-run extra args Closes #5401. Three related GGUF reload bugs reproduced against `unsloth studio run -m unsloth/Qwen3-0.6B-GGUF --gguf-variant Q4_K_M --top-k 20 --seed 42`: 1. The `POST /api/inference/load` already-loaded short-circuit only compared `model_identifier` and `hf_variant`. A same-(model, variant) Apply that flipped `cache_type_kv` / `speculative_type` / `chat_template_override` / `max_seq_length` / `llama_extra_args` returned `status="already_loaded"` and the new setting silently never reached llama-server. 2. The frontend chat-settings Apply path POSTs `/unload` then `/load` without round-tripping `llama_extra_args`. Every reload after `unsloth run --some-flag X` quietly dropped `--some-flag X` from the spawned `llama-server` command line. 3. `LlamaCppBackend.load_model` released `_lock` between Phase 1 (kill) and Phase 3 (spawn) so two concurrent loads each passed Phase 1 with `self._process is None`. Both ran Phase 2 (download), both reached Phase 3, and the Phase 3 defensive `_kill_process()` from #5171 collapsed them to one survivor only after both `subprocess.Popen` calls had landed. For the 86 GB MoE in #5161 / the model in #5401 the overlap window was tens of seconds, long enough to OOM the host. With a 0.6B model the pgrep timeline showed two simultaneous PIDs for 3.3 s on `main`. Fix: `studio/backend/core/inference/llama_cpp.py` * Add `self._serial_load_lock = threading.Lock()`. The whole body of `load_model` runs under this lock so two concurrent `/api/inference/load` requests are strictly sequential. The fine-grained `_lock` and the Phase 3 defensive `_kill_process()` from #5171 are kept as a second layer. `/unload`, `/status`, and `/load-progress` are unaffected because they only touch the fine-grained lock or read properties. * Add `self._extra_args` plus an `extra_args` property, written inside `load_model` whenever the caller supplies a non-`None` value. `unload_model()` deliberately does not reset it so the route layer can inherit the args across the frontend's `/unload` + `/load` gap. `studio/backend/routes/inference.py` * Add `_request_matches_loaded_settings(request, llama_backend)` that compares `max_seq_length`, `cache_type_kv`, `speculative_type`, `chat_template_override`, and `llama_extra_args` between the incoming request and the live backend. Same-(model, variant) requests whose runtime settings differ now fall through to a real reload instead of returning `already_loaded`. A missing `llama_extra_args` field on the request is treated as "inherit current", so the short-circuit still fires when the only difference is the frontend not echoing the CLI flags back. * GGUF load branch inherits `llama_extra_args` from `llama_backend.extra_args` when the request omits the field, re-validates through `validate_extra_args`, and forwards the result to `load_model(...)`. An explicit `[]` from the caller is still honoured as "clear". Verified end to end against a live `unsloth studio run` instance: | Scenario | Before | After | | --------------------------------------------------------------- | --------- | ------------------------------------------------------------------------ | | `/load` same (model, variant, settings) | 1 PID, `already_loaded` | unchanged | | `/load` same model, variant, new `cache_type_kv=q8_0` ctx=8192 | `already_loaded`, settings dropped | `loaded`, `/status` reports the new settings, new server has `-c 8192 --cache-type-k q8_0 --top-k 20 --seed 42` | | Frontend Apply `/unload` + `/load`, new settings, no `llama_extra_args` field | Drops `--top-k 20 --seed 42` | Preserves `--top-k 20 --seed 42` | | `/unload` + two parallel `/load` | Two PIDs for 3.3 s | Max simultaneous count = 1 across the full pgrep timeline | | `/load` with `llama_extra_args=[]` (explicit clear) | n/a | `loaded`, new server has no `--top-k` / `--seed` | | `/load` with `llama_extra_args=["--top-k","30","--seed","7"]` (override) | n/a | `loaded`, new server has the supplied flags | `pytest studio/backend/tests` is green except for one pre-existing terminal-width-sensitive assertion (`test_studio_api.py::test_help_output`) and the pre-existing `test_studio_api.py` fixture errors that fail on unmodified main too. No new regressions. * Studio: track requested n_ctx so Auto-slider flips trigger a reload Review feedback on PR #5427 from gemini-code-assist. The original short-circuit compared ``request.max_seq_length`` against ``llama_backend.context_length`` (the effective context). VRAM-fit logic can cap the running server below what the caller asked for, so this comparison incorrectly returns ``already_loaded`` when the user flips the slider from an explicit length (e.g. 8192) back to "Auto" (0): the explicit request was capped to, say, 4096, and the new "Auto" request reads ``backend.context_length == 4096`` and decides nothing changed. Track the originally requested ``n_ctx`` on the backend instead and compare against that. ``requested_n_ctx == 0`` means the last load asked for the model's native length; ``request.max_seq_length == 0`` matches it. Verified in the sandbox suite (now 90 tests): - ``test_explicit_to_auto_triggers_reload`` -- loaded with explicit 8192, then Apply with ``max_seq_length=0`` falls through to a real reload and the new server runs at the native 40960. - ``test_auto_to_explicit_triggers_reload`` -- inverse direction. - ``test_explicit_to_same_explicit_short_circuits`` -- re-Apply with the same explicit value still short-circuits (no needless reload). - Existing scenarios (kv change, spec change, template change, extra args inherit, parallel-load stress, frontend Apply flow) unchanged. ``pytest studio/backend/tests`` still green on the same set of tests; the pre-existing ``test_help_output`` failure and ``test_studio_api`` fixture errors are unaffected. * Studio: tighten comments in the 5401 fix Trim the verbose explanatory comments and docstrings introduced in f9cbec3b and dd0b1d58 down to one-line summaries. The "why" still points at issue #5401; the multi-paragraph rationale belonged in the PR body, not the source. No behaviour change. * ci: retrigger after zoo drift + IPython fixes landed in main * ci: retrigger Mac Studio UI CI after transient fetch flake * Studio: address six P2 followups on the 5401 reload PR Tightens the inheritance and serial-load paths to close the six P2 findings raised by codex-connector on PR #5427 against `f9cbec3b` / `dd0b1d58`. 1. Re-check loaded state before killing queued loads. Two duplicate `/api/inference/load` requests both pass the route-level `is_loaded` gate before the first publishes `_healthy = True`. The second waits on `_serial_load_lock`, enters Phase 1, and tears down the just-spawned llama-server for a redundant full reload. Added `LlamaCppBackend._already_in_target_state(...)` and a short-circuit at the top of the serial-lock block: if the live server already satisfies the kwargs, return True without killing. 2. Don't inherit CLI overrides that shadow new first-class settings. `unsloth run -c 4096` is a permitted pass-through; the validator docs explicitly call out `-c`/`--ctx-size`. Stored in `_extra_args` and appended after Studio's own flags, the inherited `-c 4096` silently won the last-wins parse against a new `max_seq_length=8192`. Added `strip_shadowing_flags` in `llama_server_args.py` (covers `-c`, `--cache-type-k/v`, `--spec-*`, `--chat-template*`, `--jinja`/`--no-jinja`) and the route runs the inherited list through it before validate + forward. 3. Restrict inherited llama args to the same GGUF model. `_extra_args` is deliberately preserved across `unload_model()` for the chat- settings Apply flow (`/unload` + `/load` with no `llama_extra_args` field). Now also track `_extra_args_source = (model_identifier, hf_variant)` so the route can refuse cross-model inheritance. `LlamaCppBackend.extra_args_source` exposes the tuple. 4. Persist extras only after a successful load. `_extra_args` was written at the top of `load_model` before Popen + health check, so a failed startup left bad args in place to poison the next UI retry. The write (along with `_requested_n_ctx`) is now deferred until after `_healthy = True`. 5. Ignore speculative diffs for vision loads. `load_model` silently gates speculative decoding on `not is_vision`, so the backend's `_speculative_type` stays `None` for vision models. The route's comparator now normalises the request's value to `"off"` when `llama_backend.is_vision` to avoid a no-op reload of a vision server every time the dropdown defaults to `default`. The `_already_in_target_state` helper applies the same rule. 6. Wait for the replacement server before short-circuiting. `_kill_process` did not clear `_healthy`; the new first-class settings (`_cache_type_kv`, `_speculative_type`, `_chat_template_override`) are written under `_lock` BEFORE Popen + `_wait_for_health`. A duplicate `/load` arriving during the new server's warm-up window could short-circuit against the not-yet-healthy replacement and the caller would start inference against a server that was still loading. `_kill_process` now sets `_healthy = False` in its `finally` block so `is_loaded` returns False from the moment the old server is killed until the new one finishes warm-up. Tests: - Sandbox suite under `./temp/sim_5401/` extended to 136 tests (was 90): new unit coverage for `strip_shadowing_flags` (12 cases), `_kill_process` clears `_healthy`, `extra_args_source` lifecycle and cross-model behaviour, failed-load preserving prior extras, and the duplicate-load short-circuit at `load_model` level. New live integration cases verify shadow-strip via `pgrep` on the live llama-server cmdline, cross-model refusal, and PID stability across a duplicate-load race. All 136 pass. - `pytest studio/backend/tests --deselect test_studio_api.py`: 1079 passed, 46 skipped, identical to the pre-change count. The pre-existing `test_studio_api.py` fixture errors and the terminal-width-sensitive `test_help_output` are unaffected. - Ruff: clean on the three modified files. * Studio: tighten GGUF reload inheritance and duplicate-load guard Re-narrow llama_extra_args to None after validate_extra_args when the incoming request omitted the field, so the backend can distinguish "caller omitted, inherit prior load" from "caller explicitly cleared to []". Without this a queued duplicate /load reaches the backend as [] and fails _already_in_target_state's exact-equality check, killing the just-started llama-server. The pass-through validate call from the original "forward llama-server args from unsloth studio run / unsloth run" change is preserved as-is; only the post-pass narrowing is new. Cross-source loads now explicitly clear extras so a model switch can't accidentally inherit via the backend's "no opinion" semantics. Store the caller's hf_variant kwarg (None for local GGUF files) in _extra_args_source instead of the derived self._hf_variant (an extracted filename quant label like "Q4_K_M"). Same-source check in the route is now symmetric for HF and direct-file loads. Add gguf_path to _already_in_target_state and prefer on-disk path identity when both backend and caller have a path. This stops the duplicate-load guard from killing a healthy server on repeat local loads (where hf_variant is None on the caller side but extracted on the backend side). Split shadow-flag stripping into per-group toggles (context / cache / spec / template). The route now opts into stripping only the groups whose first-class field was actually set on the incoming request, so an inherited --chat-template-file survives an Apply that omits chat_template_override. _request_matches_loaded_settings detects shadowing extras on the inherit path and falls through to a real reload so the strip can run. Mark --spec-default, --jinja, --no-jinja as boolean inside the shadow stripper so the value-consuming heuristic no longer eats the following positional token. * Studio: trim comments around GGUF reload inheritance * Studio: cover GGUF reload inheritance and shadow-flag stripping * Studio: drop redundant issue refs from inheritance comments * Studio: drop redundant issue refs from inheritance comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: key inheritance source off resolved gguf_variant codex-connector P2 on PR #5427 cd14cae1: the inheritance gate at ``routes/inference.py:696`` compared the stored ``source[1]`` against ``request.gguf_variant``, but the HF branch loaded with ``hf_variant = config.gguf_variant`` (the *resolved* variant after ModelConfig auto-pick). When the caller omitted ``gguf_variant`` on a follow-up Apply, ``source[1] == "Q4_K_M"`` but ``(request.gguf_variant or "") == ""``, ``same_source`` returned False, and the chat-settings Apply silently dropped CLI pass-through flags for every auto-pick / local-file load. Fix both sides of the comparison to key off ``config.gguf_variant``: * The route compares ``source[1]`` to ``config.gguf_variant`` (the resolved label) rather than the request field. * The local-mode load_model call now passes ``hf_variant = config.gguf_variant`` so ``_extra_args_source`` stores the same string the route reads back. The HF branch already did this. Sandbox: added test_source_records_caller_variant_not_extracted_label to lock the storage key contract. ``pytest studio/backend/tests --deselect test_studio_api.py``: 1100 passed, identical to pre-change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: deny upstream --ui family on llama-server pass-through The validator's web-UI block named only ``--webui`` / ``--no-webui``, which is llama.cpp's pre-rename spelling. Current upstream (``tools/server/README.md``) uses ``--ui`` / ``--no-ui`` plus ``--ui-config``, ``--ui-config-file``, and ``--ui-mcp-proxy`` / ``--no-ui-mcp-proxy``. Without these in the denylist a user could ``unsloth run --ui`` and enable llama-server's built-in web UI on the port Studio's reverse proxy targets, breaking the UI surface. Keep the legacy ``--webui`` group so the validator still rejects old binaries that haven't been re-spelled. Cross-referenced against the README's full flag list; this was the only gap for the post-#5401 inheritance / shadow-strip work. Pass- through flags from every other README category (sampling, jinja, ctx, cache, threads, GPU, reasoning, grammar, chat-template-kwargs) already validate cleanly; sandbox suite exercises ~60 of them in the new ``test_08_llama_server_pass_through.py``. ``pytest studio/backend/tests --deselect test_studio_api.py``: 1100 passed, identical to pre-change. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 1357 +++++++++-------- .../core/inference/llama_server_args.py | 100 ++ studio/backend/routes/inference.py | 137 +- .../tests/test_gguf_reload_inheritance.py | 237 +++ .../backend/tests/test_llama_server_args.py | 118 ++ 5 files changed, 1348 insertions(+), 601 deletions(-) create mode 100644 studio/backend/tests/test_gguf_reload_inheritance.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7ef687035c..3682f1dbbb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -470,6 +470,17 @@ class LlamaCppBackend: # their own cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers). self._shared_kv_layers: Optional[int] = None self._lock = threading.Lock() + # Wraps load_model() end-to-end so concurrent loads serialise + # and never coexist as two llama-server processes (#5401). + self._serial_load_lock = threading.Lock() + # Last extra_args / requested n_ctx, preserved across unload so + # the chat UI's /unload+/load Apply path can inherit them (#5401). + # ``_extra_args_source`` records the (model_identifier, hf_variant) + # the stored args came from so the route can refuse cross-model + # inheritance. + self._extra_args: Optional[List[str]] = None + self._extra_args_source: Optional[tuple[str, Optional[str]]] = None + self._requested_n_ctx: int = 0 self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None self._cancel_event = threading.Event() @@ -505,6 +516,25 @@ class LlamaCppBackend: def hf_variant(self) -> Optional[str]: return self._hf_variant + @property + def extra_args(self) -> Optional[List[str]]: + """Extra llama-server flags from the last load. Copy; None = never + set, [] = explicitly cleared. Used by the route for inheritance.""" + return list(self._extra_args) if self._extra_args is not None else None + + @property + def requested_n_ctx(self) -> int: + """n_ctx the last load was invoked with (not the effective cap). + 0 means Auto. Used by the route to detect Auto-vs-explicit flips.""" + return self._requested_n_ctx + + @property + def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]: + """(model_identifier, hf_variant) the stored extra_args came from. + ``None`` if no extras have ever been recorded. Used by the route + to refuse cross-model inheritance (#5401).""" + return self._extra_args_source + @property def context_length(self) -> Optional[int]: """Return the effective context length the server is running at.""" @@ -1983,653 +2013,778 @@ class LlamaCppBackend: Returns True if server started and health check passed. """ - self._cancel_event.clear() - - # ── Phase 1: kill old process (under lock, fast) ────────── - with self._lock: - self._kill_process() - - binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) - - # ── Phase 2: download (NO lock held, so cancel can proceed) ── - if hf_repo: - model_path = self._download_gguf( - hf_repo = hf_repo, + # Serialise the whole load so concurrent /load calls never + # leave two llama-server processes alive (#5401 / #5161). Does + # not block /unload, /status, /load-progress. + with self._serial_load_lock: + # Duplicate /load that raced past the route-level check + # (the first one hadn't published _healthy=True yet). If the + # live server already satisfies this request, do nothing. + if self._already_in_target_state( + gguf_path = gguf_path, + model_identifier = model_identifier, hf_variant = hf_variant, - hf_token = hf_token, - ) - # Auto-download mmproj for vision models - if is_vision and not mmproj_path: - mmproj_path = self._download_mmproj( + n_ctx = n_ctx, + cache_type_kv = cache_type_kv, + speculative_type = speculative_type, + chat_template_override = chat_template_override, + extra_args = extra_args, + is_vision = is_vision, + ): + logger.info( + f"load_model: backend already in target state for " + f"'{model_identifier}', skipping reload" + ) + return True + + self._cancel_event.clear() + + # ── Phase 1: kill old process (under lock, fast) ────────── + with self._lock: + self._kill_process() + + binary = self._find_llama_server_binary() + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + + # ── Phase 2: download (NO lock held, so cancel can proceed) ── + if hf_repo: + model_path = self._download_gguf( hf_repo = hf_repo, + hf_variant = hf_variant, hf_token = hf_token, ) - elif gguf_path: - if not Path(gguf_path).is_file(): - raise FileNotFoundError(f"GGUF file not found: {gguf_path}") - model_path = gguf_path - else: - raise ValueError("Either gguf_path or hf_repo must be provided") + # Auto-download mmproj for vision models + if is_vision and not mmproj_path: + mmproj_path = self._download_mmproj( + hf_repo = hf_repo, + hf_token = hf_token, + ) + elif gguf_path: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + model_path = gguf_path + else: + raise ValueError("Either gguf_path or hf_repo must be provided") - # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection - self._model_identifier = model_identifier + # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection + self._model_identifier = model_identifier - # Read GGUF metadata (context_length, chat_template) -- fast, header only - self._read_gguf_metadata(model_path) + # Read GGUF metadata (context_length, chat_template) -- fast, header only + self._read_gguf_metadata(model_path) - # Check cancel after download - if self._cancel_event.is_set(): - logger.info("Load cancelled after download phase") - return False - - # ── Phase 3: start llama-server (under lock) ────────────── - with self._lock: - # Re-check cancel inside lock + # Check cancel after download if self._cancel_event.is_set(): - logger.info("Load cancelled before server start") + logger.info("Load cancelled after download phase") return False - self._port = self._find_free_port() + # ── Phase 3: start llama-server (under lock) ────────────── + with self._lock: + # Re-check cancel inside lock + if self._cancel_event.is_set(): + logger.info("Load cancelled before server start") + return False - # Select GPU(s) based on model size + estimated KV cache. - # Seed safe defaults before GPU probing so the except path - # still has valid state to publish. - effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) - max_available_ctx = self._context_length or effective_ctx - gpus: list[tuple[int, int]] = [] - try: - model_size = self._get_gguf_size_bytes(model_path) - gpus = self._get_gpu_free_memory() + self._port = self._find_free_port() - # Resolve effective context: 0 means let llama-server use the - # model's native length. Only expand to a known native length - # if metadata is available; otherwise preserve 0 as a sentinel. - if n_ctx > 0: - effective_ctx = n_ctx - elif self._context_length is not None: - effective_ctx = self._context_length - else: - effective_ctx = 0 - original_ctx = effective_ctx - # Default UI ceiling to the model's native context length. - # GPU/VRAM-fit logic below may shrink this if hardware is limited. + # Select GPU(s) based on model size + estimated KV cache. + # Seed safe defaults before GPU probing so the except path + # still has valid state to publish. + effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx + gpus: list[tuple[int, int]] = [] + try: + model_size = self._get_gguf_size_bytes(model_path) + gpus = self._get_gpu_free_memory() - # Auto-cap context to fit in GPU VRAM and select GPUs. - # - # Two policies depending on whether the user set n_ctx: - # - # Explicit n_ctx (user chose a context length): - # Honor it. Try the full requested context with _select_gpus - # (which uses as many GPUs as needed). Only cap if it doesn't - # fit on any GPU combination. - # - # Auto n_ctx=0 (model's native context): - # Prefer fewer GPUs with reduced context over more GPUs, - # since multi-GPU is slower and the user didn't ask for a - # specific context length. - gpu_indices, use_fit = None, True - explicit_ctx = n_ctx > 0 + # Resolve effective context: 0 means let llama-server use the + # model's native length. Only expand to a known native length + # if metadata is available; otherwise preserve 0 as a sentinel. + if n_ctx > 0: + effective_ctx = n_ctx + elif self._context_length is not None: + effective_ctx = self._context_length + else: + effective_ctx = 0 + original_ctx = effective_ctx + # Default UI ceiling to the model's native context length. + # GPU/VRAM-fit logic below may shrink this if hardware is limited. + max_available_ctx = self._context_length or effective_ctx - if gpus and self._can_estimate_kv() and effective_ctx > 0: - # Compute the largest hardware-aware cap from the model's - # native context across all usable GPU subsets (for UI - # bounds), independent of the currently requested context. - native_ctx_for_cap = self._context_length or effective_ctx - if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) - best_cap = 0 - for n_gpus in range(1, len(ranked_for_cap) + 1): - subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) - capped = self._fit_context_to_vram( - native_ctx_for_cap, - pool_mib, - model_size, - cache_type_kv, - n_parallel = n_parallel, + # Auto-cap context to fit in GPU VRAM and select GPUs. + # + # Two policies depending on whether the user set n_ctx: + # + # Explicit n_ctx (user chose a context length): + # Honor it. Try the full requested context with _select_gpus + # (which uses as many GPUs as needed). Only cap if it doesn't + # fit on any GPU combination. + # + # Auto n_ctx=0 (model's native context): + # Prefer fewer GPUs with reduced context over more GPUs, + # since multi-GPU is slower and the user didn't ask for a + # specific context length. + gpu_indices, use_fit = None, True + explicit_ctx = n_ctx > 0 + + if gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. + native_ctx_for_cap = self._context_length or effective_ctx + if native_ctx_for_cap > 0: + ranked_for_cap = sorted( + gpus, key = lambda g: g[1], reverse = True ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + n_parallel = n_parallel, + ) + kv = self._estimate_kv_cache_bytes( + capped, cache_type_kv, n_parallel = n_parallel + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.90: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + else: + # Weights exceed 90% of every GPU subset's free + # memory, so there is no fitting context. Anchor + # the UI's "safe zone" threshold at 4096 (the + # spec's default when the model cannot fit) so + # the ctx slider shows the "might be slower" + # warning as soon as the user drags above the + # fallback default instead of never. + max_available_ctx = min(4096, native_ctx_for_cap) + + if explicit_ctx: + # Honor the user's requested context verbatim. If it + # fits, pin GPUs and skip --fit; if it doesn't, ship + # -c <user_ctx> --fit on and let llama-server flex + # -ngl (CPU layer offload). The UI is expected to + # have surfaced the "might be slower" warning before + # the user submitted a ctx above the fit ceiling. + requested_total = ( + model_size + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: - best_cap = max(best_cap, capped) - if best_cap > 0: - max_available_ctx = best_cap + gpu_indices, use_fit = self._select_gpus( + requested_total, gpus + ) + # No silent shrink: effective_ctx stays == n_ctx. else: - # Weights exceed 90% of every GPU subset's free - # memory, so there is no fitting context. Anchor - # the UI's "safe zone" threshold at 4096 (the - # spec's default when the model cannot fit) so - # the ctx slider shows the "might be slower" - # warning as soon as the user drags above the - # fallback default instead of never. - max_available_ctx = min(4096, native_ctx_for_cap) + # Auto context: prefer fewer GPUs, cap context + # to fit. Same headroom threshold as + # _select_gpus (#5106). + ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + pin_fraction = self._GPU_PIN_VRAM_FRACTION + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + effective_ctx, + pool_mib, + model_size, + cache_type_kv, + n_parallel = n_parallel, + ) + kv = self._estimate_kv_cache_bytes( + capped, cache_type_kv, n_parallel = n_parallel + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * pin_fraction: + effective_ctx = capped + gpu_indices = sorted(idx for idx, _ in subset) + use_fit = False + break + else: + # Native ctx doesn't fit. Drop to 4096 and + # re-check before deferring to --fit on: + # a model that overflows at 131k may pin + # comfortably with a 4096 KV cache (#5106). + effective_ctx = min(4096, effective_ctx) + if effective_ctx > 0: + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + kv = self._estimate_kv_cache_bytes( + effective_ctx, + cache_type_kv, + n_parallel = n_parallel, + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * pin_fraction: + gpu_indices = sorted( + idx for idx, _ in subset + ) + use_fit = False + break - if explicit_ctx: - # Honor the user's requested context verbatim. If it - # fits, pin GPUs and skip --fit; if it doesn't, ship - # -c <user_ctx> --fit on and let llama-server flex - # -ngl (CPU layer offload). The UI is expected to - # have surfaced the "might be slower" warning before - # the user submitted a ctx above the fit ceiling. - requested_total = model_size + self._estimate_kv_cache_bytes( + elif gpus: + # Can't estimate KV -- fall back to file-size-only check. + # Without KV estimation we cannot prove a hardware cap, so + # keep the ceiling at the native context (already the default). + logger.debug( + "Falling back to file-size-only GPU selection", + model_size_gb = round(model_size / (1024**3), 2), + ) + gpu_indices, use_fit = self._select_gpus(model_size, gpus) + if use_fit and not explicit_ctx: + # Weights don't fit on any subset. Default the UI to + # 4096 so the slider doesn't land on an unusable native + # context. --fit on will flex -ngl at runtime. + effective_ctx = ( + min(4096, effective_ctx) if effective_ctx > 0 else 4096 + ) + + if effective_ctx < original_ctx: + kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel ) - gpu_indices, use_fit = self._select_gpus(requested_total, gpus) - # No silent shrink: effective_ctx stays == n_ctx. - else: - # Auto context: prefer fewer GPUs, cap context - # to fit. Same headroom threshold as - # _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) - pin_fraction = self._GPU_PIN_VRAM_FRACTION - for n_gpus in range(1, len(ranked) + 1): - subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) - capped = self._fit_context_to_vram( - effective_ctx, - pool_mib, - model_size, - cache_type_kv, - n_parallel = n_parallel, - ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: - effective_ctx = capped - gpu_indices = sorted(idx for idx, _ in subset) - use_fit = False - break - else: - # Native ctx doesn't fit. Drop to 4096 and - # re-check before deferring to --fit on: - # a model that overflows at 131k may pin - # comfortably with a 4096 KV cache (#5106). - effective_ctx = min(4096, effective_ctx) - if effective_ctx > 0: - for n_gpus in range(1, len(ranked) + 1): - subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) - kv = self._estimate_kv_cache_bytes( - effective_ctx, - cache_type_kv, - n_parallel = n_parallel, - ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: - gpu_indices = sorted(idx for idx, _ in subset) - use_fit = False - break - - elif gpus: - # Can't estimate KV -- fall back to file-size-only check. - # Without KV estimation we cannot prove a hardware cap, so - # keep the ceiling at the native context (already the default). - logger.debug( - "Falling back to file-size-only GPU selection", - model_size_gb = round(model_size / (1024**3), 2), - ) - gpu_indices, use_fit = self._select_gpus(model_size, gpus) - if use_fit and not explicit_ctx: - # Weights don't fit on any subset. Default the UI to - # 4096 so the slider doesn't land on an unusable native - # context. --fit on will flex -ngl at runtime. - effective_ctx = ( - min(4096, effective_ctx) if effective_ctx > 0 else 4096 + logger.info( + f"Context auto-reduced: {original_ctx} -> {effective_ctx} " + f"(model: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_est / (1024**3):.1f} GB)" ) - if effective_ctx < original_ctx: - kv_est = self._estimate_kv_cache_bytes( + kv_cache_bytes = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel ) logger.info( - f"Context auto-reduced: {original_ctx} -> {effective_ctx} " - f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"GGUF size: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"context: {effective_ctx}, " + f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) + except Exception as e: + logger.warning(f"GPU selection failed ({e}), using --fit on") + gpu_indices, use_fit = None, True + effective_ctx = n_ctx # fall back to original - kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) - logger.info( - f"GGUF size: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " - f"context: {effective_ctx}, " - f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" - ) - except Exception as e: - logger.warning(f"GPU selection failed ({e}), using --fit on") - gpu_indices, use_fit = None, True - effective_ctx = n_ctx # fall back to original + cmd = [ + binary, + "-m", + model_path, + "--port", + str(self._port), + "-c", + str(effective_ctx) if effective_ctx > 0 else "0", + "--parallel", + str(n_parallel), + "--flash-attn", + "on", # Force flash attention for speed + # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". + "--no-context-shift", + ] - cmd = [ - binary, - "-m", - model_path, - "--port", - str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", - "--parallel", - str(n_parallel), - "--flash-attn", - "on", # Force flash attention for speed - # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". - "--no-context-shift", - ] + if use_fit: + cmd.extend(["--fit", "on"]) + elif gpu_indices is not None: + # Model fits on selected GPU(s) -- offload all layers + cmd.extend(["-ngl", "-1"]) - if use_fit: - cmd.extend(["--fit", "on"]) - elif gpu_indices is not None: - # Model fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) - - # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we - # do not inherit llama-server's internal default, which has historically - # varied (hardware concurrency incl. hyperthreads on some builds). - cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) - - # Always enable Jinja chat template rendering for proper template support - cmd.extend(["--jinja"]) - - # KV cache data type - _valid_cache_types = { - "f16", - "bf16", - "q8_0", - "q4_0", - "q4_1", - "q5_0", - "q5_1", - "iq4_nl", - "f32", - } - if cache_type_kv and cache_type_kv in _valid_cache_types: + # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we + # do not inherit llama-server's internal default, which has historically + # varied (hardware concurrency incl. hyperthreads on some builds). cmd.extend( - ["--cache-type-k", cache_type_kv, "--cache-type-v", cache_type_kv] + ["--threads", str(n_threads if n_threads is not None else -1)] ) - self._cache_type_kv = cache_type_kv - logger.info(f"KV cache type: {cache_type_kv}") - else: - self._cache_type_kv = None - # Speculative decoding (n-gram self-speculation, zero VRAM cost) - # ngram-mod: ~16 MB shared hash pool, constant memory/complexity, - # variable draft lengths. Helps most when the model repeats - # existing text (code refactoring, summarization, reasoning). - # For general chat with low repetition, overhead is ~5 ms. - # - # Benchmarks from upstream llama.cpp speculative-decoding PRs: - # Scenario | Without | With | Speedup - # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x - # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x - # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x - # - # Params from llama.cpp docs (docs/speculative.md): - # --spec-ngram-size-n 24 (small n not recommended) - # --draft-min 48 --draft-max 64 (MoEs need long drafts; - # dense models can reduce these) - # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md - # ref: https://github.com/ggml-org/llama.cpp/pull/19164 - # ref: https://github.com/ggml-org/llama.cpp/pull/18471 - # ``"default"`` -> let llama-server pick a sensible spec - # config via ``--spec-default``. Explicit type names are - # passed through with the manual draft tuning we've shipped - # historically so power users keep their overrides. - _valid_spec_types = {"ngram-simple", "ngram-mod"} - normalized_spec = ( - speculative_type.lower().strip() if speculative_type else None - ) - if normalized_spec and normalized_spec != "off" and not is_vision: - if normalized_spec == "default": - cmd.append("--spec-default") - self._speculative_type = "default" - elif normalized_spec in _valid_spec_types: - cmd.extend(["--spec-type", normalized_spec]) - if normalized_spec == "ngram-mod": - cmd.extend( - [ - "--spec-ngram-size-n", - "24", - "--draft-min", - "48", - "--draft-max", - "64", - ] - ) - self._speculative_type = normalized_spec + # Always enable Jinja chat template rendering for proper template support + cmd.extend(["--jinja"]) + + # KV cache data type + _valid_cache_types = { + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "iq4_nl", + "f32", + } + if cache_type_kv and cache_type_kv in _valid_cache_types: + cmd.extend( + [ + "--cache-type-k", + cache_type_kv, + "--cache-type-v", + cache_type_kv, + ] + ) + self._cache_type_kv = cache_type_kv + logger.info(f"KV cache type: {cache_type_kv}") + else: + self._cache_type_kv = None + + # Speculative decoding (n-gram self-speculation, zero VRAM cost) + # ngram-mod: ~16 MB shared hash pool, constant memory/complexity, + # variable draft lengths. Helps most when the model repeats + # existing text (code refactoring, summarization, reasoning). + # For general chat with low repetition, overhead is ~5 ms. + # + # Benchmarks from upstream llama.cpp speculative-decoding PRs: + # Scenario | Without | With | Speedup + # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x + # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x + # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x + # + # Params from llama.cpp docs (docs/speculative.md): + # --spec-ngram-size-n 24 (small n not recommended) + # --draft-min 48 --draft-max 64 (MoEs need long drafts; + # dense models can reduce these) + # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md + # ref: https://github.com/ggml-org/llama.cpp/pull/19164 + # ref: https://github.com/ggml-org/llama.cpp/pull/18471 + # ``"default"`` -> let llama-server pick a sensible spec + # config via ``--spec-default``. Explicit type names are + # passed through with the manual draft tuning we've shipped + # historically so power users keep their overrides. + _valid_spec_types = {"ngram-simple", "ngram-mod"} + normalized_spec = ( + speculative_type.lower().strip() if speculative_type else None + ) + if normalized_spec and normalized_spec != "off" and not is_vision: + if normalized_spec == "default": + cmd.append("--spec-default") + self._speculative_type = "default" + elif normalized_spec in _valid_spec_types: + cmd.extend(["--spec-type", normalized_spec]) + if normalized_spec == "ngram-mod": + cmd.extend( + [ + "--spec-ngram-size-n", + "24", + "--draft-min", + "48", + "--draft-max", + "64", + ] + ) + self._speculative_type = normalized_spec + else: + self._speculative_type = None else: self._speculative_type = None - else: - self._speculative_type = None - # Apply custom chat template override if provided - self._chat_template_override = chat_template_override - if chat_template_override: - import tempfile + # Apply custom chat template override if provided + self._chat_template_override = chat_template_override + if chat_template_override: + import tempfile - flags = detect_reasoning_flags( - chat_template_override, - self._model_identifier, - log_source = "GGUF chat template override", - ) - self._supports_reasoning = flags["supports_reasoning"] - self._reasoning_style = flags["reasoning_style"] - self._reasoning_always_on = flags["reasoning_always_on"] - self._supports_preserve_thinking = flags["supports_preserve_thinking"] - self._supports_tools = flags["supports_tools"] - - self._chat_template_file = tempfile.NamedTemporaryFile( - mode = "w", - suffix = ".jinja", - delete = False, - prefix = "unsloth_chat_template_", - ) - self._chat_template_file.write(chat_template_override) - self._chat_template_file.close() - cmd.extend(["--chat-template-file", self._chat_template_file.name]) - logger.info( - f"Using custom chat template file: {self._chat_template_file.name}" - ) - - # For reasoning models, set default thinking mode. - # Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default. - # Only 9B and larger enable thinking. - # Always-on templates ignore the kwarg entirely, so skip. - if self._supports_reasoning and not self._reasoning_always_on: - thinking_default = True - mid = (model_identifier or "").lower() - if "qwen3.5" in mid or "qwen3.6" in mid: - size_val = _extract_model_size_b(mid) - if size_val is not None and size_val < 9: - thinking_default = False - self._reasoning_default = thinking_default - reasoning_kw = self._reasoning_kwargs(thinking_default) - cmd.extend( - [ - "--chat-template-kwargs", - json.dumps(reasoning_kw), + flags = detect_reasoning_flags( + chat_template_override, + self._model_identifier, + log_source = "GGUF chat template override", + ) + self._supports_reasoning = flags["supports_reasoning"] + self._reasoning_style = flags["reasoning_style"] + self._reasoning_always_on = flags["reasoning_always_on"] + self._supports_preserve_thinking = flags[ + "supports_preserve_thinking" ] - ) - logger.info(f"Reasoning model: {reasoning_kw} by default") + self._supports_tools = flags["supports_tools"] - if mmproj_path: - if not Path(mmproj_path).is_file(): - logger.warning(f"mmproj file not found: {mmproj_path}") - else: - # #5347 guard for paths that bypass detect_mmproj_file. - from utils.models.model_config import ( - mmproj_matches_model_family, + self._chat_template_file = tempfile.NamedTemporaryFile( + mode = "w", + suffix = ".jinja", + delete = False, + prefix = "unsloth_chat_template_", + ) + self._chat_template_file.write(chat_template_override) + self._chat_template_file.close() + cmd.extend(["--chat-template-file", self._chat_template_file.name]) + logger.info( + f"Using custom chat template file: {self._chat_template_file.name}" ) - if not mmproj_matches_model_family(model_path, mmproj_path): - logger.warning( - f"Skipping mmproj with mismatched family: " - f"model={Path(model_path).name}, " - f"mmproj={Path(mmproj_path).name}" - ) + # For reasoning models, set default thinking mode. + # Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default. + # Only 9B and larger enable thinking. + # Always-on templates ignore the kwarg entirely, so skip. + if self._supports_reasoning and not self._reasoning_always_on: + thinking_default = True + mid = (model_identifier or "").lower() + if "qwen3.5" in mid or "qwen3.6" in mid: + size_val = _extract_model_size_b(mid) + if size_val is not None and size_val < 9: + thinking_default = False + self._reasoning_default = thinking_default + reasoning_kw = self._reasoning_kwargs(thinking_default) + cmd.extend( + [ + "--chat-template-kwargs", + json.dumps(reasoning_kw), + ] + ) + logger.info(f"Reasoning model: {reasoning_kw} by default") + + if mmproj_path: + if not Path(mmproj_path).is_file(): + logger.warning(f"mmproj file not found: {mmproj_path}") else: - cmd.extend(["--mmproj", mmproj_path]) - logger.info(f"Using mmproj for vision: {mmproj_path}") - - # Option C: add --api-key for direct client access when enabled - import os as _os - import secrets as _secrets - - if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": - self._api_key = _secrets.token_urlsafe(32) - cmd.extend(["--api-key", self._api_key]) - logger.info("llama-server started with --api-key for direct streaming") - else: - self._api_key = None - - # User-supplied pass-through args go last so llama.cpp's - # last-wins flag parsing lets the user override Studio's - # auto-set tier-2 flags (e.g. --cache-type-k, --spec-type). - # The route layer has already validated this list against - # the managed-flag denylist via validate_extra_args(). - if extra_args: - cmd.extend(str(a) for a in extra_args) - logger.info( - f"Appending user extra args to llama-server: {list(extra_args)}" - ) - - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "<redacted>" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") - - # Set library paths so llama-server can find its shared libs and CUDA DLLs - import os - import sys - - env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) - - if sys.platform == "win32": - # CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must - # be on PATH. Order: binary_dir, torch's pip-installed - # nvidia wheels, then a system CUDA toolkit. Pip wheels - # are the canonical source per Studio's install design - # (mirrors the Linux LD_LIBRARY_PATH block below) and - # CUDA_PATH covers users with a system toolkit. #5106. - path_dirs = [binary_dir] - path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix)) - cuda_path = os.environ.get("CUDA_PATH", "") - if cuda_path: - cuda_bin = os.path.join(cuda_path, "bin") - if os.path.isdir(cuda_bin): - path_dirs.append(cuda_bin) - # Some CUDA installs put DLLs in bin\x64 - cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64") - if os.path.isdir(cuda_bin_x64): - path_dirs.append(cuda_bin_x64) - existing_path = env.get("PATH", "") - env["PATH"] = ";".join(path_dirs) + ";" + existing_path - else: - # Linux: set LD_LIBRARY_PATH for shared libs next to the binary - # and CUDA runtime libs (libcudart, libcublas, etc.) - import platform - - lib_dirs = [binary_dir] - _arch = platform.machine() # x86_64, aarch64, etc. - - # Pip-installed nvidia CUDA runtime libs (e.g. torch's - # bundled cuda-bindings). The prebuilt llama.cpp binary - # links against libcudart.so.13 / libcublas.so.13 which - # live here, not in /usr/local/cuda. - import glob as _glob - - for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), - ]: - for _nv_dir in _glob.glob(_nv_pattern): - if os.path.isdir(_nv_dir): - lib_dirs.append(_nv_dir) - - for cuda_lib in [ - "/usr/local/cuda/lib64", - f"/usr/local/cuda/targets/{_arch}-linux/lib", - # Fallback CUDA compat paths (e.g. binary built with - # CUDA 12 on a system where default /usr/local/cuda - # points to CUDA 13+). - "/usr/local/cuda-12/lib64", - "/usr/local/cuda-12.8/lib64", - f"/usr/local/cuda-12/targets/{_arch}-linux/lib", - f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", - ]: - if os.path.isdir(cuda_lib): - lib_dirs.append(cuda_lib) - existing_ld = env.get("LD_LIBRARY_PATH", "") - new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = ( - f"{new_ld}:{existing_ld}" if existing_ld else new_ld - ) - - # Pin to selected GPU(s). On ROCm, llama-server (and any torch - # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES; - # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing - # the full HIP/ROCR set the parent inherited. - if gpu_indices is not None: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - env["ROCR_VISIBLE_DEVICES"] = pinned - except Exception as e: - logger.debug( - "Failed to set ROCm visibility env vars for child: %s", e - ) - - # Defensive kill: if a concurrent load slipped past Phase 1 - # (because its `self._process` was None at the time) and - # already stored a Popen handle here, drop that orphan - # before we overwrite the reference. See issue #5161. - self._kill_process() - - self._stdout_lines = [] - self._process = subprocess.Popen( - cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - text = True, - env = env, - **_windows_hidden_subprocess_kwargs(), - ) - - # Start background thread to drain stdout and prevent pipe deadlock - self._stdout_thread = threading.Thread( - target = self._drain_stdout, daemon = True, name = "llama-stdout" - ) - self._stdout_thread.start() - - # Store the resolved on-disk path, not the caller's kwarg. In - # HF mode the caller passes gguf_path=None and the real path - # (``model_path``) is what llama-server is actually mmap'ing. - # Downstream consumers (load_progress, log lines, etc.) need - # the path that exists on disk. - self._gguf_path = model_path - self._hf_repo = hf_repo - # For local GGUF files, extract variant from filename if not provided - if hf_variant: - self._hf_variant = hf_variant - elif gguf_path: - try: - from utils.models.model_config import _extract_quant_label - - self._hf_variant = _extract_quant_label(gguf_path) - except Exception: - self._hf_variant = None - else: - self._hf_variant = None - self._is_vision = is_vision - self._model_identifier = model_identifier - - # Store the effective (possibly capped) context separately. - # Do NOT overwrite _context_length -- it holds the model's native - # context length from GGUF metadata and is used for display/info. - self._effective_context_length = ( - effective_ctx if effective_ctx > 0 else self._context_length - ) - self._max_context_length = ( - max_available_ctx - if max_available_ctx > 0 - else self._effective_context_length - ) - - # Wait for llama-server to become healthy - if not self._wait_for_health(timeout = 600.0): - self._kill_process() - _gguf = gguf_path or "" - _is_ollama = ( - ".studio_links" in _gguf - or os.sep + "ollama_links" + os.sep in _gguf - or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf - or (self._model_identifier or "").startswith("ollama/") - ) - # Only show the Ollama-specific message when the server - # output indicates a GGUF compatibility issue, not for - # unrelated failures like OOM or missing binaries. - if _is_ollama: - _output = "\n".join(self._stdout_lines[-50:]).lower() - _gguf_compat_hints = ( - "key not found", - "unknown model architecture", - "failed to load model", - ) - if any(h in _output for h in _gguf_compat_hints): - raise RuntimeError( - "Some Ollama models do not work with llama.cpp. " - "Try a different model, or use this model directly through Ollama instead." + # #5347 guard for paths that bypass detect_mmproj_file. + from utils.models.model_config import ( + mmproj_matches_model_family, ) - raise RuntimeError( - "llama-server failed to start. " - "Check that the GGUF file is valid and you have enough memory." + + if not mmproj_matches_model_family(model_path, mmproj_path): + logger.warning( + f"Skipping mmproj with mismatched family: " + f"model={Path(model_path).name}, " + f"mmproj={Path(mmproj_path).name}" + ) + else: + cmd.extend(["--mmproj", mmproj_path]) + logger.info(f"Using mmproj for vision: {mmproj_path}") + + # Option C: add --api-key for direct client access when enabled + import os as _os + import secrets as _secrets + + if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + self._api_key = _secrets.token_urlsafe(32) + cmd.extend(["--api-key", self._api_key]) + logger.info( + "llama-server started with --api-key for direct streaming" + ) + else: + self._api_key = None + + # User-supplied pass-through args go last so llama.cpp's + # last-wins flag parsing lets the user override Studio's + # auto-set tier-2 flags (e.g. --cache-type-k, --spec-type). + # The route layer has already validated this list against + # the managed-flag denylist via validate_extra_args(). + if extra_args: + cmd.extend(str(a) for a in extra_args) + logger.info( + f"Appending user extra args to llama-server: {list(extra_args)}" + ) + + _log_cmd = list(cmd) + if "--api-key" in _log_cmd: + _ki = _log_cmd.index("--api-key") + 1 + if _ki < len(_log_cmd): + _log_cmd[_ki] = "<redacted>" + logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + + # Set library paths so llama-server can find its shared libs and CUDA DLLs + import os + import sys + + env = child_env_without_native_path_secret() + binary_dir = str(Path(binary).parent) + + if sys.platform == "win32": + # CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must + # be on PATH. Order: binary_dir, torch's pip-installed + # nvidia wheels, then a system CUDA toolkit. Pip wheels + # are the canonical source per Studio's install design + # (mirrors the Linux LD_LIBRARY_PATH block below) and + # CUDA_PATH covers users with a system toolkit. #5106. + path_dirs = [binary_dir] + path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix)) + cuda_path = os.environ.get("CUDA_PATH", "") + if cuda_path: + cuda_bin = os.path.join(cuda_path, "bin") + if os.path.isdir(cuda_bin): + path_dirs.append(cuda_bin) + # Some CUDA installs put DLLs in bin\x64 + cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64") + if os.path.isdir(cuda_bin_x64): + path_dirs.append(cuda_bin_x64) + existing_path = env.get("PATH", "") + env["PATH"] = ";".join(path_dirs) + ";" + existing_path + else: + # Linux: set LD_LIBRARY_PATH for shared libs next to the binary + # and CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [binary_dir] + _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs (e.g. torch's + # bundled cuda-bindings). The prebuilt llama.cpp binary + # links against libcudart.so.13 / libcublas.so.13 which + # live here, not in /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cu*", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cudnn", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "nvjitlink", + "lib", + ), + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with + # CUDA 12 on a system where default /usr/local/cuda + # points to CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = ( + f"{new_ld}:{existing_ld}" if existing_ld else new_ld + ) + + # Pin to selected GPU(s). On ROCm, llama-server (and any torch + # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES; + # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing + # the full HIP/ROCR set the parent inherited. + if gpu_indices is not None: + pinned = ",".join(str(i) for i in gpu_indices) + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env["ROCR_VISIBLE_DEVICES"] = pinned + except Exception as e: + logger.debug( + "Failed to set ROCm visibility env vars for child: %s", e + ) + + # Defensive kill: if a concurrent load slipped past Phase 1 + # (because its `self._process` was None at the time) and + # already stored a Popen handle here, drop that orphan + # before we overwrite the reference. See issue #5161. + self._kill_process() + + self._stdout_lines = [] + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **_windows_hidden_subprocess_kwargs(), ) - self._healthy = True + # Start background thread to drain stdout and prevent pipe deadlock + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "llama-stdout" + ) + self._stdout_thread.start() - # Catch silent CPU fallback when GPU was intended (#5106). - self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit, gpus or [] - ) - if self._gpu_offload_active is False: - logger.warning( - "llama-server appears to have loaded the model entirely " - "on CPU even though Studio detected at least one GPU. " - "This usually means the prebuilt binary's GPU backend " - "failed to load -- on Windows, cudart64_X.dll / " - "cublas64_X.dll could not be resolved. Reinstall the " - "Studio llama.cpp prebuilt or install a matching CUDA " - "toolkit (issue unslothai/unsloth#5106).", + # Store the resolved on-disk path, not the caller's kwarg. In + # HF mode the caller passes gguf_path=None and the real path + # (``model_path``) is what llama-server is actually mmap'ing. + # Downstream consumers (load_progress, log lines, etc.) need + # the path that exists on disk. + self._gguf_path = model_path + self._hf_repo = hf_repo + # For local GGUF files, extract variant from filename if not provided + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None + self._is_vision = is_vision + self._model_identifier = model_identifier + + # Store the effective (possibly capped) context separately. + # Do NOT overwrite _context_length -- it holds the model's native + # context length from GGUF metadata and is used for display/info. + self._effective_context_length = ( + effective_ctx if effective_ctx > 0 else self._context_length + ) + self._max_context_length = ( + max_available_ctx + if max_available_ctx > 0 + else self._effective_context_length ) - logger.info( - f"llama-server ready on port {self._port} " - f"for model '{model_identifier}'" - ) - return True + # Wait for llama-server to become healthy + if not self._wait_for_health(timeout = 600.0): + self._kill_process() + _gguf = gguf_path or "" + _is_ollama = ( + ".studio_links" in _gguf + or os.sep + "ollama_links" + os.sep in _gguf + or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf + or (self._model_identifier or "").startswith("ollama/") + ) + # Only show the Ollama-specific message when the server + # output indicates a GGUF compatibility issue, not for + # unrelated failures like OOM or missing binaries. + if _is_ollama: + _output = "\n".join(self._stdout_lines[-50:]).lower() + _gguf_compat_hints = ( + "key not found", + "unknown model architecture", + "failed to load model", + ) + if any(h in _output for h in _gguf_compat_hints): + raise RuntimeError( + "Some Ollama models do not work with llama.cpp. " + "Try a different model, or use this model directly through Ollama instead." + ) + raise RuntimeError( + "llama-server failed to start. " + "Check that the GGUF file is valid and you have enough memory." + ) + + self._healthy = True + + # Commit caller intent only after _healthy=True so a + # failed startup can't poison the next inheritance check. + # None keeps prior, [] clears, list sets. Source records + # the caller's hf_variant (None for local files) so the + # route's same_source check stays symmetric. + if extra_args is not None: + self._extra_args = list(extra_args) + self._extra_args_source = (model_identifier, hf_variant) + self._requested_n_ctx = int(n_ctx) + + # Catch silent CPU fallback when GPU was intended (#5106). + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit, gpus or [] + ) + if self._gpu_offload_active is False: + logger.warning( + "llama-server appears to have loaded the model entirely " + "on CPU even though Studio detected at least one GPU. " + "This usually means the prebuilt binary's GPU backend " + "failed to load -- on Windows, cudart64_X.dll / " + "cublas64_X.dll could not be resolved. Reinstall the " + "Studio llama.cpp prebuilt or install a matching CUDA " + "toolkit (issue unslothai/unsloth#5106).", + ) + + logger.info( + f"llama-server ready on port {self._port} " + f"for model '{model_identifier}'" + ) + return True + + def _already_in_target_state( + self, + *, + model_identifier: str, + hf_variant: Optional[str], + n_ctx: int, + cache_type_kv: Optional[str], + speculative_type: Optional[str], + chat_template_override: Optional[str], + extra_args: Optional[List[str]], + is_vision: bool, + gguf_path: Optional[str] = None, + ) -> bool: + """True iff the live server already satisfies these load kwargs. + + Mirrors ``routes/inference.py:_request_matches_loaded_settings`` + but compares raw kwargs so ``load_model`` can short-circuit a + duplicate /load that raced past the route-level check (#5401). + """ + if not self.is_loaded: + return False + if (self._model_identifier or "").lower() != (model_identifier or "").lower(): + return False + # Direct-file loads pass hf_variant=None while the backend + # stores an extracted filename label; compare paths instead + # to keep the guard symmetric. + if gguf_path is not None and self._gguf_path: + try: + if Path(self._gguf_path).resolve() != Path(gguf_path).resolve(): + return False + except OSError: + return False + elif (self._hf_variant or "").lower() != (hf_variant or "").lower(): + return False + if self._requested_n_ctx != int(n_ctx): + return False + + def _norm(value): + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + return stripped or None + return value + + if _norm(self._cache_type_kv) != _norm(cache_type_kv): + return False + + # Vision GGUFs silently drop speculative decoding in + # load_model (the spec gate is "not is_vision"); treat the + # request's value as "off" so a vision load with + # speculative_type="default" still matches. + if self._is_vision or is_vision: + req_spec = "off" + else: + req_spec = _norm(speculative_type) or "off" + backend_spec = _norm(self._speculative_type) or "off" + if req_spec != backend_spec: + return False + + if (self._chat_template_override or None) != (chat_template_override or None): + return False + + # extra_args=None means "no opinion" (inherit semantics handled + # at the route layer); only an explicit list forces equality. + if extra_args is not None: + current = list(self._extra_args) if self._extra_args is not None else [] + if list(extra_args) != current: + return False + return True def _classify_gpu_offload( self, @@ -2737,6 +2892,10 @@ class LlamaCppBackend: logger.warning(f"Error killing llama-server process: {e}") finally: self._process = None + # Clear healthy so a /load arriving during the replacement + # server's warm-up window cannot short-circuit against the + # previous server's health (#5401). + self._healthy = False if self._stdout_thread is not None: self._stdout_thread.join(timeout = 2) self._stdout_thread = None diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 44c7d542c7..0f6927fc5a 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -69,7 +69,15 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # Single-model server -- Studio runs one model per llama-server # process and serves its own UI. Enabling multi-model loading or # llama-server's built-in web UI changes the surface clients see. + # ``--webui``/``--no-webui`` are the legacy spelling; current + # upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions. + # Keep both so the denylist matches old and new llama-server + # binaries (Studio's prebuilt vs system-llama.cpp). frozenset({"--webui", "--no-webui"}), + frozenset({"--ui", "--no-ui"}), + frozenset({"--ui-config"}), + frozenset({"--ui-config-file"}), + frozenset({"--ui-mcp-proxy", "--no-ui-mcp-proxy"}), frozenset({"--models-dir"}), frozenset({"--models-preset"}), frozenset({"--models-max"}), @@ -118,3 +126,95 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: def is_managed_flag(flag: str) -> bool: """True if ``flag`` is a Studio-managed llama-server flag.""" return flag in _DENYLIST + + +# Pass-through flags that shadow first-class ``LoadRequest`` fields +# (max_seq_length, cache_type_kv, speculative_type, +# chat_template_override). Stripped from inherited extras so they +# can't last-wins-override an Apply that re-sets the same first-class +# field. +_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) +_CACHE_FLAGS: frozenset[str] = frozenset( + {"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"} +) +_SPEC_FLAGS: frozenset[str] = frozenset( + { + "--spec-default", + "--spec-type", + "--spec-ngram-size-n", + "--spec-ngram-size", + "--draft-min", + "--draft-max", + } +) +_TEMPLATE_FLAGS: frozenset[str] = frozenset( + { + "--chat-template", + "--chat-template-file", + "--chat-template-kwargs", + "--jinja", + "--no-jinja", + } +) + +_SHADOWING_FLAGS: frozenset[str] = ( + _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS +) + +# Boolean flags inside _SHADOWING_FLAGS that take no value. The +# value-consuming heuristic in strip_shadowing_flags must skip just the +# flag for these, never the following token. +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja"} +) + + +def strip_shadowing_flags( + args: Iterable[str], + *, + strip_context: bool = True, + strip_cache: bool = True, + strip_spec: bool = True, + strip_template: bool = True, +) -> list[str]: + """Strip flags that shadow first-class Studio settings. + + Used when the route inherits a previous load's ``llama_extra_args`` + so that an inherited ``-c 4096`` cannot override the current + request's ``max_seq_length`` (and equivalents for cache / + speculative / chat template). Each ``strip_*`` flag controls one + group; the route only strips groups whose corresponding first-class + field was actually supplied by the caller, so an inherited + ``--chat-template-file`` survives an Apply that omits both + ``llama_extra_args`` and ``chat_template_override``. + """ + shadowing: set[str] = set() + if strip_context: + shadowing |= _CONTEXT_FLAGS + if strip_cache: + shadowing |= _CACHE_FLAGS + if strip_spec: + shadowing |= _SPEC_FLAGS + if strip_template: + shadowing |= _TEMPLATE_FLAGS + + tokens = [str(a) for a in (args or [])] + out: list[str] = [] + i, n = 0, len(tokens) + while i < n: + tok = tokens[i] + flag = _flag_name(tok) + if flag is None or flag not in shadowing: + out.append(tok) + i += 1 + continue + # Drop this token. Boolean shadowing flags never carry a value; + # other shadowing flags consume the next token when it isn't a + # flag and the value isn't already packed as ``--key=value``. + if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok: + i += 1 + elif i + 1 < n and _flag_name(tokens[i + 1]) is None: + i += 2 + else: + i += 1 + return out diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 76bbb59c94..60078ecc9b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -119,7 +119,10 @@ try: _DEFAULT_T_MAX_PREDICT_MS, detect_reasoning_flags, ) - from core.inference.llama_server_args import validate_extra_args + from core.inference.llama_server_args import ( + strip_shadowing_flags, + validate_extra_args, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -141,7 +144,10 @@ except ImportError: _DEFAULT_T_MAX_PREDICT_MS, detect_reasoning_flags, ) - from core.inference.llama_server_args import validate_extra_args + from core.inference.llama_server_args import ( + strip_shadowing_flags, + validate_extra_args, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -406,6 +412,57 @@ def _validate_native_mmproj_companion( ) from exc +def _normalise_settings_str(value: Optional[str]) -> Optional[str]: + """Lowercase + strip a settings string, mapping blank/None to None.""" + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + return stripped or None + return value + + +def _request_matches_loaded_settings( + request: LoadRequest, llama_backend: LlamaCppBackend +) -> bool: + """True iff every runtime setting on the request matches the loaded + server. Caller has already checked model+variant+is_loaded. See #5401.""" + # Compare requested n_ctx (not effective) so VRAM-cap doesn't mask + # an Auto-vs-explicit slider flip. + if request.max_seq_length != llama_backend.requested_n_ctx: + return False + if _normalise_settings_str(request.cache_type_kv) != _normalise_settings_str( + llama_backend.cache_type_kv + ): + return False + # Vision loads silently drop speculative decoding (llama_cpp.py gates + # spec on ``not is_vision``), so treat the request as ``off`` against + # the backend's ``None`` to avoid forcing a redundant reload. + if llama_backend.is_vision: + req_spec = "off" + else: + req_spec = _normalise_settings_str(request.speculative_type) or "off" + backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off" + if req_spec != backend_spec: + return False + if (request.chat_template_override or None) != ( + llama_backend.chat_template_override or None + ): + return False + # llama_extra_args=None means "inherit"; only an explicit list that + # differs forces a reload. On the inherit path, refuse to match if + # stored extras contain any shadow flag, so the reload path can + # strip them instead of leaving a stale override in effect. + backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + if request.llama_extra_args is None: + if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra: + return False + else: + if list(request.llama_extra_args) != backend_extra: + return False + return True + + def _resolve_model_identifier_for_request( request: LoadRequest | ValidateModelRequest, *, @@ -461,6 +518,11 @@ async def load_model( extra_llama_args = validate_extra_args(request.llama_extra_args) except ValueError as exc: raise HTTPException(status_code = 400, detail = str(exc)) + # Re-narrow []-from-None back to None so the inheritance path + # below can tell "caller omitted" from "caller explicit []". + extra_llama_args: Optional[list[str]] = ( + None if request.llama_extra_args is None else extra_llama_args + ) model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") @@ -479,6 +541,9 @@ async def load_model( and llama_backend.hf_variant.lower() == request.gguf_variant.lower() and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() + # Also require runtime settings to match so Apply changes + # aren't silently dropped (#5401). + and _request_matches_loaded_settings(request, llama_backend) ): logger.info( f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload" @@ -613,6 +678,70 @@ async def load_model( ) unsloth_backend.unload_model(unsloth_backend.active_model_name) + # Inherit llama_extra_args from the previous load when the + # request omits the field (the chat-settings Apply path + # does not round-trip them; explicit [] still clears). + # Inheritance is gated on (model_identifier, hf_variant) + # to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins + # CLI parse against a freshly-supplied first-class field. + if request.llama_extra_args is None and llama_backend.extra_args: + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request + # field: callers commonly omit gguf_variant for local + # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ + # variant`` is the variant load_model was actually + # invoked with (see the HF / local branches below), so + # both sides of the comparison key off the same string. + resolved_variant = config.gguf_variant + same_source = bool( + source + and source[0] + and source[0].lower() == model_identifier.lower() + and (source[1] or "").lower() == (resolved_variant or "").lower() + ) + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came " + "from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend + # doesn't inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field + # was actually set by the caller, so an inherited + # --chat-template-file survives an Apply that omits + # chat_template_override. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = "speculative_type" in fields_set, + strip_template = "chat_template_override" in fields_set, + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Should not happen on already-validated args; degrade + # to no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; " + "loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + # Route to HF mode or local mode based on config # Run in a thread so the event loop stays free for progress # polling and other requests during the (potentially long) @@ -645,6 +774,10 @@ async def load_model( llama_backend.load_model, gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, + # Pass the resolved variant so _extra_args_source + # is keyed off the same string the inheritance + # check at the top of /load uses (#5401 followup). + hf_variant = config.gguf_variant, model_identifier = config.identifier, is_vision = config.is_vision, n_ctx = request.max_seq_length, diff --git a/studio/backend/tests/test_gguf_reload_inheritance.py b/studio/backend/tests/test_gguf_reload_inheritance.py new file mode 100644 index 0000000000..4b0b450cb0 --- /dev/null +++ b/studio/backend/tests/test_gguf_reload_inheritance.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GGUF reload duplicate-load guard. + +``LlamaCppBackend._already_in_target_state`` is the in-process +short-circuit that prevents a serialised duplicate /load from killing +the just-spawned llama-server. These tests pin the local-file +identity, the HF-mode hf_variant fallback, and the ``extra_args`` +None-vs-[] inherit semantics so the guard cannot silently regress. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) +_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) +_httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +class _FakeProcess: + """Stand-in for subprocess.Popen so atexit cleanup doesn't crash.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def _loaded_backend(**overrides): + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._speculative_type = None + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._extra_args_source = None + backend._gguf_path = None + for key, value in overrides.items(): + setattr(backend, key, value) + return backend + + +# ── Local-file identity via gguf_path ──────────────────────────────── + + +def test_already_in_target_state_uses_gguf_path_when_present(tmp_path): + gguf_file = tmp_path / "model.Q4_K_M.gguf" + gguf_file.write_bytes(b"") + backend = _loaded_backend( + _hf_variant = "Q4_K_M", + _gguf_path = str(gguf_file), + ) + assert ( + backend._already_in_target_state( + gguf_path = str(gguf_file), + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_rejects_different_gguf_path(tmp_path): + a = tmp_path / "a.gguf" + a.write_bytes(b"") + b = tmp_path / "b.gguf" + b.write_bytes(b"") + backend = _loaded_backend(_gguf_path = str(a)) + assert ( + backend._already_in_target_state( + gguf_path = str(b), + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + +# ── HF mode falls back to hf_variant comparison ────────────────────── + + +def test_already_in_target_state_falls_back_to_hf_variant_for_hf_loads(): + backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q8_0", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + +def test_already_in_target_state_hf_same_variant_matches(): + backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +# ── extra_args: None inherits, [] forces reload, list enforces ─────── + + +def test_already_in_target_state_none_extras_inherits_stored(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_empty_extras_forces_reload_when_stored(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = [], + is_vision = False, + ) + is False + ) + + +def test_already_in_target_state_explicit_extras_match(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = ["--top-k", "20"], + is_vision = False, + ) + is True + ) + + +def test_extra_args_source_default_is_none(): + backend = LlamaCppBackend() + assert backend.extra_args_source is None diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 351fbd014d..3013acfdb8 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -15,6 +15,7 @@ import pytest from core.inference.llama_server_args import ( is_managed_flag, + strip_shadowing_flags, validate_extra_args, ) @@ -187,3 +188,120 @@ def test_is_managed_flag_false_for_pass_through(): assert is_managed_flag("--flash-attn") is False assert is_managed_flag("-ngl") is False assert is_managed_flag("--threads") is False + + +# ── strip_shadowing_flags ───────────────────────────────────────────── + + +def test_strip_shadowing_flags_drops_context_when_requested(): + out = strip_shadowing_flags( + ["-c", "4096", "--top-k", "20"], + strip_context = True, + strip_cache = False, + strip_spec = False, + strip_template = False, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_context_when_not_requested(): + out = strip_shadowing_flags( + ["-c", "4096", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + ) + assert out == ["-c", "4096", "--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled(): + # Caller did not supply chat_template_override; the inherited + # --chat-template-file must survive the strip. + out = strip_shadowing_flags( + ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], + strip_context = True, + strip_cache = True, + strip_spec = True, + strip_template = False, + ) + assert out == ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"] + + +def test_strip_shadowing_flags_drops_template_when_requested(): + out = strip_shadowing_flags( + ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], + strip_template = True, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_cache_when_cache_disabled(): + out = strip_shadowing_flags( + ["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"], + strip_cache = False, + ) + assert out == [ + "--cache-type-k", + "q8_0", + "--cache-type-v", + "q8_0", + "--top-k", + "20", + ] + + +def test_strip_shadowing_flags_keeps_spec_when_spec_disabled(): + out = strip_shadowing_flags( + ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"], + strip_spec = False, + ) + assert out == [ + "--spec-type", + "ngram-mod", + "--draft-min", + "48", + "--top-k", + "20", + ] + + +def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): + # --spec-default is a boolean shadowing flag; the value-skipping + # heuristic must skip just the flag, not the following positional. + out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) + assert out == ["ngram-mod"] + + +def test_strip_shadowing_flags_jinja_boolean_preserves_positional(): + out = strip_shadowing_flags(["--jinja", "trailing-positional"], strip_template = True) + assert out == ["trailing-positional"] + + +def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional(): + out = strip_shadowing_flags( + ["--no-jinja", "trailing-positional"], strip_template = True + ) + assert out == ["trailing-positional"] + + +def test_strip_shadowing_flags_equals_form_drops_only_the_flag(): + out = strip_shadowing_flags(["--ctx-size=4096", "--seed", "-1"], strip_context = True) + assert out == ["--seed", "-1"] + + +def test_strip_shadowing_flags_handles_none_input(): + assert strip_shadowing_flags(None) == [] + + +def test_strip_shadowing_flags_handles_empty_input(): + assert strip_shadowing_flags([]) == [] + + +def test_strip_shadowing_flags_defaults_strip_everything(): + # The route's already-loaded comparator calls strip_shadowing_flags + # with no kwargs to detect ANY shadowing flag in stored extras. + out = strip_shadowing_flags( + ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"] + ) + assert out == [] From 0542dc07259c693fb98c32805416bf564e437aec Mon Sep 17 00:00:00 2001 From: Daniel Han <danielhanchen@gmail.com> Date: Sun, 17 May 2026 04:20:46 -0700 Subject: [PATCH 50/50] Studio: IME / multilingual composer regression test + RTL dir="auto" (#5485) Adds dir="auto" to the main, edit, and compare chat composers so RTL scripts (Arabic, Hebrew, Persian, Urdu) flow right to left without forcing the rest of the UI into RTL. Wires a model-free Playwright smoke (multilingual paste round trip across 31 scripts + a stuck-IME composition repro for issue #5318 / PR #5327) into the Studio UI CI job as a third Studio boot, plus a pure-Python static-guard test that locks down dir="auto" on all three composers and the minimal env contract for the smoke. --- .github/workflows/studio-ui-smoke.yml | 55 ++- .../src/components/assistant-ui/thread.tsx | 5 + .../src/features/chat/shared-composer.tsx | 3 + tests/studio/playwright_chat_ime_i18n.py | 457 ++++++++++++++++++ .../test_composer_rtl_bidi_attribute.py | 73 +++ 5 files changed, 588 insertions(+), 5 deletions(-) create mode 100644 tests/studio/playwright_chat_ime_i18n.py create mode 100644 tests/studio/test_composer_rtl_bidi_attribute.py diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 79476a62ea..455fe4b7e1 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -229,12 +229,55 @@ jobs: kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 + # IME + multilingual paste regression (issue #5318 / PR #5327). + # Third Studio on its own port so a hang here cannot poison the + # earlier UI tests. No GGUF -- the bug surface is the composer. + - name: Reset auth + boot Studio for IME / i18n tests (port 18896) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \ + > logs/studio_ime.log 2>&1 & + echo "STUDIO_IME_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18896 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18896/api/health" > /tmp/health3.json; then + jq -e '.status == "healthy"' /tmp/health3.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health3.json + + - name: Pass bootstrap pw for IME / i18n test + # IME smoke does the change-password against the bootstrap that + # Studio's frontend injects into the page, so it only needs the + # NEW password. + run: | + NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_IME_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive IME + multilingual paste regression with Playwright + env: + BASE_URL: http://127.0.0.1:18896 + STUDIO_NEW_PW: ${{ env.STUDIO_IME_NEW_PW }} + PW_ART_DIR: logs/playwright_ime + STUDIO_UI_STRICT: '1' + run: | + mkdir -p logs/playwright_ime + python tests/studio/playwright_chat_ime_i18n.py + + - name: Stop third Studio + if: always() + run: | + kill "${STUDIO_IME_PID}" 2>/dev/null || true + sleep 2 + - name: Upload Playwright artifacts - # Always upload (not just failure) so a green run's screenshots - # are reviewable in the Actions UI -- catches "passed but the - # UI is silently broken" regressions that would be invisible - # otherwise. Both Studio's logs (chat + extra) and BOTH - # Playwright artifact dirs are bundled. + # Always upload so a green run's screenshots stay reviewable -- + # catches "passed but the UI is silently broken" regressions. if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: @@ -242,7 +285,9 @@ jobs: path: | logs/studio.log logs/studio_extra.log + logs/studio_ime.log logs/install.log logs/playwright logs/playwright_extra + logs/playwright_ime retention-days: 7 diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index fb63748bf1..c8fea9c704 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -328,6 +328,9 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { autoFocus={!disabled} disabled={disabled} aria-label="Message input" + // dir="auto": browser picks LTR/RTL from the first strong char; + // no effect on Latin / CJK / Devanagari. + dir="auto" {...inputProps} /> <ComposerAction @@ -1161,6 +1164,8 @@ const EditComposer: FC = () => { <ComposerPrimitive.Input className="aui-edit-composer-input min-h-14 w-full resize-none bg-transparent p-4 text-foreground text-sm font-[450] outline-none" autoFocus={true} + // See main composer above for the dir="auto" rationale. + dir="auto" {...inputProps} /> <div className="aui-edit-composer-footer mx-3 mb-3 flex items-center gap-2 self-end"> diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index c4ffa98467..cd31d37cd7 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -690,6 +690,9 @@ export function SharedComposer({ placeholder="Send to both models..." className="composer-input" rows={1} + // dir="auto" auto-detects RTL (Arabic / Hebrew / Persian / Urdu) + // from the first strong character; no effect on LTR scripts. + dir="auto" /> <div className="composer-action-wrapper"> <div className="flex items-center gap-1"> diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py new file mode 100644 index 0000000000..c882d88cbd --- /dev/null +++ b/tests/studio/playwright_chat_ime_i18n.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Studio chat composer IME + multilingual regression smoke. + +Covers two surfaces: + A. Stuck IME composition (issue #5318 / PR #5327): duplicate + compositionstart with no compositionend left isComposing=true, + dropping all subsequent keystrokes including ASCII. + B. Multilingual paste round-trip across 31 scripts -- guards the + controlled-textarea / React state plumbing against Unicode mangling. + +Model-free; the bug surface is the composer, not inference. + +Env contract matches playwright_chat_ui.py: + BASE_URL, STUDIO_NEW_PW, PW_ART_DIR, STUDIO_UI_STRICT. +""" + +import os +import sys +from pathlib import Path + +from playwright.sync_api import expect, sync_playwright + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from _playwright_robust import ( # noqa: E402 + chromium_launch_args, + click_and_wait_for_response, + install_view_transition_killer, + install_wall_clock_watchdog, + is_benign_console_error, + is_benign_page_error, + recover_or_replace_page, + wait_for_health, +) + +BASE = os.environ["BASE_URL"] +NEW = os.environ["STUDIO_NEW_PW"] +ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_ime") +ART = Path(ART_DIR) +ART.mkdir(parents = True, exist_ok = True) +STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" + +# Wall-clock cap. Realistic run is 30-60s; 5 min leaves cold-launch headroom. +WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300")) + + +# One short greeting + arithmetic per script (ordered by speaker count) -- +# each entry catches a distinct class of Unicode regression. +I18N_SAMPLES = [ + ("en", "English", "Hello, 1+1=2"), + ("zh-CN", "Chinese (Simplified)", "你好，1+1=2"), + ("es", "Spanish", "Hola, 1+1=2"), + ("hi", "Hindi (Devanagari)", "नमस्ते, 1+1=2"), + ("ar", "Arabic (RTL)", "مرحبا، ١+١=٢"), + ("bn", "Bengali", "নমস্কার, ১+১=২"), + ("pt", "Portuguese", "Olá, 1+1=2"), + ("ru", "Russian (Cyrillic)", "Привет, 1+1=2"), + ("ja", "Japanese", "こんにちは、1+1=2"), + ("pa", "Punjabi (Gurmukhi)", "ਸਤ ਸ੍ਰੀ ਅਕਾਲ, 1+1=2"), + ("de", "German", "Hallo, 1+1=2"), + ("jv", "Javanese", "Halo, 1+1=2"), + ("ko", "Korean (Hangul)", "안녕하세요, 1+1=2"), + ("fr", "French", "Bonjour, 1+1=2"), + ("tr", "Turkish", "Merhaba, 1+1=2"), + ("vi", "Vietnamese (diacritics)", "Xin chào, 1+1=2"), + ("ur", "Urdu (Arabic-Naskh)", "ہیلو، 1+1=2"), + ("ta", "Tamil", "வணக்கம், 1+1=2"), + ("te", "Telugu", "నమస్తే, 1+1=2"), + ("mr", "Marathi (Devanagari)", "नमस्कार, 1+1=2"), + ("it", "Italian", "Ciao, 1+1=2"), + ("th", "Thai", "สวัสดี, ๑+๑=๒"), + ("pl", "Polish", "Cześć, 1+1=2"), + ("uk", "Ukrainian (Cyrillic)", "Привіт, 1+1=2"), + ("fa", "Persian (RTL)", "سلام، ۱+۱=۲"), + ("nl", "Dutch", "Hallo, 1+1=2"), + ("he", "Hebrew (RTL)", "שלום, 1+1=2"), + ("el", "Greek", "Γειά, 1+1=2"), + ("id", "Indonesian", "Halo, 1+1=2"), + ("sw", "Swahili", "Habari, 1+1=2"), + ("emoji", "Emoji + ZWJ + flag", "👋 🇺🇳 👨‍👩‍👧‍👦 1+1=2"), +] + + +_n = [0] + + +def step(s): + print(f"[ime] STEP {s}", flush = True) + + +def info(s): + print(f"[ime] {s}", flush = True) + + +def fail(m): + raise AssertionError(f"[ime] FAIL: {m}") + + +def soft_fail(m): + """Hard fail in STRICT mode, info-warn otherwise. Mirrors playwright_chat_ui.py.""" + if STRICT: + fail(m) + info(f"WARN (strict-off): {m}") + + +with sync_playwright() as p: + _watchdog = install_wall_clock_watchdog( + WALL_TIMEOUT_S, + label = "ime", + info = info, + ) + wait_for_health(BASE, timeout = 30.0, info = info) + browser = p.chromium.launch( + headless = True, + args = chromium_launch_args(), + ) + ctx = browser.new_context( + viewport = {"width": 1280, "height": 900}, + reduced_motion = "reduce", + ) + install_view_transition_killer(ctx) + page = ctx.new_page() + page.set_default_timeout(60_000) + + page_errors: list[str] = [] + console_errors: list[str] = [] + + def _on_console(m): + if m.type != "error": + return + try: + console_errors.append(m.text) + except Exception: + return + + def _attach_listeners(target): + target.on("pageerror", lambda e: page_errors.append(str(e))) + target.on("console", _on_console) + + _attach_listeners(page) + + def shoot(name): + _n[0] += 1 + try: + page.screenshot( + path = str(ART / f"{_n[0]:02d}-{name}.png"), + full_page = True, + timeout = 90_000, + animations = "disabled", + ) + except Exception as _shoot_err: + info(f"WARN: screenshot {name} failed: {_shoot_err}") + + # 1. Bootstrap auth via /change-password (mirrors playwright_chat_ui.py + # retry-on-rerender to absorb React form-detach races). + step("change-password through UI (Setup your account)") + form_err: Exception | None = None + for _form_attempt in range(3): + try: + page.goto( + f"{BASE}/change-password", + wait_until = "domcontentloaded", + timeout = 60_000, + ) + try: + page.wait_for_load_state("networkidle", timeout = 30_000) + except Exception: + pass + pw_field = page.locator("#new-password") + pw_field.wait_for(state = "visible", timeout = 60_000) + pw_field.fill(NEW, timeout = 60_000) + page.fill("#confirm-password", NEW, timeout = 60_000) + shoot("01-change-password-filled") + status, _ = click_and_wait_for_response( + page, + url_substr = "/api/auth/change-password", + method = "POST", + do_click = lambda: page.locator('button[type="submit"]').click(), + timeout_ms = 30_000, + info = lambda m: print(f"[ime] {m}", flush = True), + ) + if status is not None and status >= 400: + raise AssertionError(f"change-password POST returned {status}") + form_err = None + break + except Exception as e: + form_err = e + info( + f"change-password attempt {_form_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}" + ) + if _form_attempt < 2: + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ime] recovery: {m}", flush = True), + ) + _attach_listeners(page) + if form_err is not None: + raise form_err + + # 2. Wait for composer mount. No GGUF: the bug surface is React state, not inference. + step("wait for composer to mount") + try: + page.wait_for_load_state("networkidle", timeout = 30_000) + except Exception: + pass + composer = page.locator('textarea[aria-label="Message input"]') + _mount_err: Exception | None = None + for _mount_attempt in range(2): + try: + composer.wait_for(state = "visible", timeout = 60_000) + _mount_err = None + break + except Exception as e: + _mount_err = e + info( + f"composer.wait_for attempt {_mount_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}" + ) + try: + shoot(f"02-composer-wait-attempt-{_mount_attempt + 1}-fail") + except Exception: + pass + if _mount_attempt == 0: + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ime] recovery: {m}", flush = True), + ) + _attach_listeners(page) + composer = page.locator('textarea[aria-label="Message input"]') + if _mount_err is not None: + raise _mount_err + composer.click() + shoot("02-composer-focused") + + # Main composer must carry dir="auto" so RTL flows right-to-left. + dir_attr = composer.evaluate("(el) => el.getAttribute('dir')") + if dir_attr != "auto": + soft_fail( + f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' + "languages will render LTR." + ) + else: + info('composer dir="auto" present') + + # Source-level guard for the edit and compare composers (neither + # is mounted here): grep the JSX for dir="auto" inside each block. + _repo_root = Path(__file__).resolve().parents[2] + _thread_src = ( + _repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx" + ).read_text() + _shared_src = ( + _repo_root / "studio/frontend/src/features/chat/shared-composer.tsx" + ).read_text() + _edit_idx = _thread_src.find("aui-edit-composer-input") + if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]: + soft_fail('edit composer source is missing dir="auto"') + else: + info('edit composer dir="auto" present (source)') + _compare_idx = _shared_src.find("Send to both models") + if ( + _compare_idx == -1 + or 'dir="auto"' + not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400] + ): + soft_fail('compare composer source is missing dir="auto"') + else: + info('compare composer dir="auto" present (source)') + + def read_value() -> str: + return composer.evaluate("(el) => el.value") + + def set_value_via_setter(s: str) -> str: + """Write via React's monkey-patched setter + paste input event, + then await two rAFs so the controlled value is committed before + readback (plain `.value=s` would be overwritten on next render).""" + return composer.evaluate( + """async (el, v) => { + const setter = Object.getOwnPropertyDescriptor( + window.HTMLTextAreaElement.prototype, 'value' + ).set; + setter.call(el, v); + el.dispatchEvent(new InputEvent('input', { + bubbles: true, + inputType: 'insertFromPaste', + data: v, + })); + await new Promise((r) => requestAnimationFrame(r)); + await new Promise((r) => requestAnimationFrame(r)); + return el.value; + }""", + s, + ) + + def clear() -> None: + set_value_via_setter("") + + # 3. Baseline: ASCII keyboard typing works. Bail fast if not. + step("baseline ASCII keyboard typing") + clear() + composer.click() + for ch in "hello world": + page.keyboard.type(ch) + got = read_value() + if got != "hello world": + fail(f"ASCII typing readback {got!r} != 'hello world'") + info("baseline ASCII OK") + shoot("03-baseline-ascii") + clear() + + # 4. Multilingual paste round-trip; byte-for-byte readback required. + step(f"multilingual paste round-trip ({len(I18N_SAMPLES)} samples)") + paste_failures: list[tuple[str, str, str, str]] = [] + for code, label, text in I18N_SAMPLES: + got = set_value_via_setter(text) + if got != text: + paste_failures.append((code, label, text, got)) + info(f" {code:>6} ({label}): FAIL -- got {got!r}") + else: + info(f" {code:>6} ({label}): OK") + clear() + if paste_failures: + shoot("04-paste-failures") + lines = [ + f" {code} ({label}): want={want!r} got={got!r}" + for code, label, want, got in paste_failures + ] + fail( + f"{len(paste_failures)}/{len(I18N_SAMPLES)} languages failed paste round-trip:\n" + + "\n".join(lines) + ) + info(f"all {len(I18N_SAMPLES)} multilingual paste samples OK") + shoot("04-paste-all-ok") + + # 5. Healthy IME composition (compositionstart/update/end + insert events). + step("normal IME composition (compose 你好)") + clear() + composer.click() + composer.evaluate( + """(el) => { + el.focus(); + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你'})); + el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你好'})); + const setter = Object.getOwnPropertyDescriptor( + window.HTMLTextAreaElement.prototype, 'value' + ).set; + setter.call(el, el.value + '你好'); + el.dispatchEvent(new InputEvent('input', { + bubbles:true, inputType:'insertCompositionText', + data:'你好', isComposing:true, + })); + el.dispatchEvent(new CompositionEvent('compositionend', {bubbles:true, data:'你好'})); + el.dispatchEvent(new InputEvent('input', { + bubbles:true, inputType:'insertFromComposition', data:'你好', + })); + }""" + ) + got = read_value() + if "你好" not in got: + shoot("05-normal-composition-FAIL") + fail(f"normal composition readback {got!r} missing '你好'") + info(f"normal composition OK: ta.value={got!r}") + shoot("05-normal-composition") + clear() + + # 6. Stuck IME repro for issue #5318: duplicate compositionstart with + # no compositionend wedged isComposing=true and dropped ASCII keys. + # PR #5327 cleared the stale state on non-composing input. + step("BUG REPRO: stuck IME composition recovery (issue #5318)") + clear() + composer.click() + composer.evaluate( + """(el) => { + el.focus(); + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + // Duplicate compositionstart with NO matching compositionend. + // This is exactly the event sequence observed from the IMEs + // in issue #5318 (kei-yamazaki / langxiaopiao030 / PapyrusNotes). + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + }""" + ) + # Drive the real keyboard path; on the broken build React drops + # 'abcd' and reconciles el.value back to ''. wait_for_function + # crosses the microtask boundary so we see committed React state. + page.keyboard.type("abcd") + try: + page.wait_for_function( + """(el) => el.value === 'abcd'""", + composer.element_handle(), + timeout = 5_000, + ) + except Exception: + pass + after_key = read_value() + info(f"after_key='abcd' readback={after_key!r}") + shoot("06-stuck-composition-recovery") + if after_key != "abcd": + fail( + "stuck-composition repro: keyboard 'abcd' was not preserved after " + f"duplicate compositionstart; readback {after_key!r}. React state " + "likely still stuck in isComposing=true (issue #5318 / before " + "PR #5327)." + ) + # Cross-check React's view of isComposing via the Send button: + # ComposerAction stays disabled while isComposing is true (PR #5327). + send_btn = page.locator('button[aria-label="Send message"]') + if send_btn.count() == 0: + soft_fail("Send button not found after stuck-composition recovery") + else: + try: + expect(send_btn).not_to_be_disabled(timeout = 5_000) + info("Send button correctly enabled after stuck-composition recovery") + except Exception: + soft_fail( + "Send button still disabled after stuck-composition recovery -- " + "React isComposing state likely never cleared" + ) + info("stuck-composition recovery PASS") + clear() + + # 7. Final state. The change-password redirect emits benign 401 noise, + # so we filter via is_benign_* and only fail on real errors. + shoot("07-final") + real_page_errors = [e for e in page_errors if not is_benign_page_error(e)] + real_console_errors = [e for e in console_errors if not is_benign_console_error(e)] + info( + f"page_errors={len(page_errors)} ({len(real_page_errors)} non-benign); " + f"console_errors={len(console_errors)} " + f"({len(real_console_errors)} non-benign)" + ) + if page_errors: + info(f"first page error: {page_errors[0][:200]!r}") + if console_errors: + info(f"first console error: {console_errors[0][:200]!r}") + if real_page_errors: + fail( + f"{len(real_page_errors)} non-benign pageerror events; " + f"first={real_page_errors[0][:200]!r}" + ) + if real_console_errors: + fail( + f"{len(real_console_errors)} non-benign console.error events; " + f"first={real_console_errors[0][:200]!r}" + ) + + info( + f"DONE: ascii=OK paste={len(I18N_SAMPLES)}/{len(I18N_SAMPLES)} " + f"normal_composition=OK stuck_recovery=OK" + ) + _watchdog.cancel() + browser.close() diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py new file mode 100644 index 0000000000..5b1437b4fc --- /dev/null +++ b/tests/studio/test_composer_rtl_bidi_attribute.py @@ -0,0 +1,73 @@ +"""Lock down the RTL bidi auto-detection contract on the chat composers. + +The browser's Unicode bidi algorithm only flows Arabic / Hebrew / Persian / +Urdu right-to-left when the textarea carries `dir="auto"`. The three +composer surfaces (main chat, inline edit, compare mode) each need the +attribute, and the IME / i18n Playwright smoke must keep its env contract +minimal (no dead `STUDIO_OLD_PW`). +""" + +from __future__ import annotations + +import re +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx" +SHARED_TSX = REPO / "studio/frontend/src/features/chat/shared-composer.tsx" +WORKFLOW_YML = REPO / ".github/workflows/studio-ui-smoke.yml" +IME_PY = REPO / "tests/studio/playwright_chat_ime_i18n.py" + + +def _block_around(src: str, anchor: str, radius: int = 600) -> str: + idx = src.find(anchor) + assert idx != -1, f"anchor {anchor!r} not found" + return src[max(idx - radius, 0) : idx + radius] + + +def test_main_composer_has_dir_auto(): + block = _block_around(THREAD_TSX.read_text(), 'aria-label="Message input"') + assert 'dir="auto"' in block, 'main composer is missing dir="auto"' + + +def test_edit_composer_has_dir_auto(): + block = _block_around(THREAD_TSX.read_text(), "aui-edit-composer-input") + assert 'dir="auto"' in block, 'edit composer is missing dir="auto"' + + +def test_compare_composer_has_dir_auto(): + block = _block_around(SHARED_TSX.read_text(), "Send to both models") + assert 'dir="auto"' in block, 'compare composer is missing dir="auto"' + + +def test_ime_workflow_step_does_not_set_studio_old_pw(): + yml = WORKFLOW_YML.read_text() + drive_idx = yml.find("Drive IME + multilingual paste regression") + assert drive_idx != -1, "IME drive step not found in workflow" + next_step_idx = yml.find("- name:", drive_idx + 1) + drive_block = yml[drive_idx : next_step_idx if next_step_idx != -1 else None] + assert ( + "STUDIO_OLD_PW" not in drive_block + ), "IME drive step still passes dead STUDIO_OLD_PW env var" + assert "STUDIO_NEW_PW" in drive_block, "IME drive step missing STUDIO_NEW_PW" + + +def test_ime_pass_password_step_does_not_export_old_pw(): + yml = WORKFLOW_YML.read_text() + pass_idx = yml.find("Pass bootstrap pw for IME / i18n test") + assert pass_idx != -1, "IME password setup step not found" + next_step_idx = yml.find("- name:", pass_idx + 1) + pass_block = yml[pass_idx : next_step_idx if next_step_idx != -1 else None] + assert ( + "STUDIO_IME_OLD_PW" not in pass_block + ), "IME password setup still exports dead STUDIO_IME_OLD_PW" + assert "STUDIO_IME_NEW_PW" in pass_block + + +def test_ime_playwright_script_does_not_read_studio_old_pw(): + src = IME_PY.read_text() + code_only = re.sub(r'""".*?"""', "", src, flags = re.DOTALL) + assert ( + "STUDIO_OLD_PW" not in code_only + ), "IME Playwright script still references dead STUDIO_OLD_PW env var" + assert 'os.environ["STUDIO_NEW_PW"]' in code_only