diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 08d7e17e8..440b6276c 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -11,6 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.chat import chat +from unsloth_cli.commands.connect import connect_app from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -77,6 +78,11 @@ app.command()(chat) app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") +app.add_typer( + connect_app, + name = "connect", + help = "Connect a coding agent (Claude Code, Codex) to Studio.", +) # Top-level `unsloth run` aliases `unsloth studio run`; same context # so unknown flags still pass through to llama-server. diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 5e734ee2a..1faa63292 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -14,6 +14,10 @@ import typer _THINK_OPEN = "" _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) +# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default +# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request. +_USER_AGENT = "unsloth-cli" + def ensure_studio_backend_path() -> None: backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend") @@ -254,11 +258,13 @@ def load_chat_backend( return ChatBackend("unsloth", backend) -def find_studio_server(timeout: float = 0.4) -> Optional[str]: +def find_studio_server(timeout: float = 3.0) -> Optional[str]: import urllib.request + base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) try: - with urllib.request.urlopen(f"{base}/api/health", timeout = timeout): + with urllib.request.urlopen(request, timeout = timeout): return base except Exception: return None @@ -306,6 +312,7 @@ class HttpChatBackend: headers = { "Authorization": f"Bearer {self._token}", "Content-Type": "application/json", + "User-Agent": _USER_AGENT, }, method = method, ) diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py new file mode 100644 index 000000000..05f75d263 --- /dev/null +++ b/unsloth_cli/commands/connect.py @@ -0,0 +1,721 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth connect` — launch a coding agent against a running Studio server.""" + +import json +import os +import re +import shlex +import shutil +import signal +import subprocess +import urllib.error +import urllib.request +from pathlib import Path +from typing import NoReturn, Optional + +import typer + +from unsloth_cli._inference import ( + _USER_AGENT, + _studio_token, + ensure_studio_backend_path, + find_studio_server, +) + +connect_app = typer.Typer( + help = "Connect a coding agent to a running Studio server.", + no_args_is_help = True, + context_settings = {"help_option_names": ["-h", "--help"]}, +) + +_CODEX_PROFILE = "unsloth_api" +_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" +_HERMES_ENV_KEY = "UNSLOTH_API_KEY" +_HERMES_PROVIDER = "unsloth" +_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" +_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} +_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") + +# Shared by every agent command; only the config/env/command differ. +_MODEL_OPTION = typer.Option( + None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." +) +_KEY_OPTION = typer.Option( + None, + "--api-key", + envvar = "UNSLOTH_API_KEY", + help = "Studio API key; minted automatically when omitted. Keys are remembered, so passing one once is enough.", +) +_LAUNCH_OPTION = typer.Option( + True, + "--launch/--no-launch", + help = "--no-launch prints the env and command instead (remote shells, WSL).", +) + + +def _fail(message: str) -> NoReturn: + typer.echo(message, err = True) + raise typer.Exit(code = 1) + + +def _http_error_detail(exc: urllib.error.HTTPError) -> str: + try: + body = json.loads(exc.read().decode()) + return body.get("detail") or body["error"]["message"] + except Exception: + return str(exc) + + +def _http_json( + method: str, + url: str, + token: str, + payload = None, + timeout = 30, + error = None, +): + """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" + request = urllib.request.Request( + url, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": _USER_AGENT, + }, + method = method, + ) + try: + with urllib.request.urlopen(request, timeout = timeout) as response: + return json.loads(response.read().decode() or "{}") + except urllib.error.HTTPError as exc: + if error is None: + raise + _fail(f"{error}: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + if error is None: + raise + _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") + + +def _require_studio() -> str: + base = find_studio_server() + if base is None: + expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") + _fail( + f"No running Studio server found at {expected}. Start one with " + "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." + ) + return base + + +def _key_cache_path() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agent_api_key.json" + + +def _cached_keys(cache: Path) -> list: + try: + data = json.loads(cache.read_text()) + except Exception: + return [] + if not isinstance(data, dict): + return [] + keys = [k for k in data.get("keys", []) if isinstance(k, str)] + legacy = data.get("key") # pre-multi-key cache format + if isinstance(legacy, str) and legacy not in keys: + keys.append(legacy) + return keys + + +def _write_private_json(path: Path, data: dict) -> None: + # O_CREAT with 0o600 so a file holding an API key is never world-readable, + # even briefly (existing files keep whatever perms the user set). + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as handle: + handle.write(json.dumps(data, indent = 2) + "\n") + + +def _read_json_object(path: Path) -> Optional[dict]: + # {} when missing, None when it can't be parsed as an object (so the caller + # leaves a user-managed file untouched rather than clobbering it). + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + return None + return data if isinstance(data, dict) else None + + +def _subdict(parent: dict, key: str) -> dict: + child = parent.get(key) + if not isinstance(child, dict): + child = parent[key] = {} + return child + + +def _remember_key(cache: Path, key: str) -> None: + existing = _cached_keys(cache) + keys = ([key] + [k for k in existing if k != key])[:8] + if keys == existing: + return + try: + _write_private_json(cache, {"keys": keys}) + except OSError: + pass # worst case the next launch mints another key + + +def _agent_api_key(base: str, explicit: Optional[str]) -> str: + cache = _key_cache_path() + if explicit: + _remember_key(cache, explicit) + return explicit + + # Keys are per-server, so when switching between Studios (local one day, + # an SSH-tunnelled remote the next) the right key is whichever validates. + for key in _cached_keys(cache): + try: + _http_json("GET", f"{base}/v1/models", key) + except Exception: + continue + _remember_key(cache, key) + return key + + token = _studio_token() + auth_help = ( + "Couldn't authenticate with the Studio server automatically (it may be " + "remote, or running as a different OS user). Create an API key in " + "Studio → Settings → API and pass it once with --api-key; it is " + "remembered for next time." + ) + if token is None: + _fail(auth_help) + try: + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth connect)"}, + )["key"] + except urllib.error.HTTPError as exc: + # A self-issued token only validates against a local, same-OS-user + # server; a remote Studio signs with a different secret and rejects it + # (401/403). Point at --api-key instead of the raw "expired token". + if exc.code in (401, 403): + _fail(auth_help) + _fail(f"Couldn't create an API key: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + _fail(f"Couldn't create an API key: {getattr(exc, 'reason', None) or exc}") + _remember_key(cache, key) + return key + + +def _loaded_models(base: str, key: str) -> list: + return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) + + +def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict: + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] == requested), None) + if requested and match is None: + typer.echo(f"Loading {requested} on the Studio server (this can take a while)…") + loaded = _http_json( + "POST", + f"{base}/api/inference/load", + key, + {"model_path": requested}, + timeout = 3600, + error = "Model load failed", + ) + # Studio registers the model under a canonical id (resolved identifier, + # casing) that /v1/models echoes but which may differ from the path we + # passed; match on the id the load reports so we don't silently fall + # through to models[0] and connect to a different loaded model. + wanted = {requested} + if isinstance(loaded, dict): + wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] in wanted), None) + if match is not None: + return match + if requested: + # We asked Studio to load it and it didn't surface in /v1/models; don't + # silently hand back an unrelated loaded model. + _fail( + f"Studio didn't report '{requested}' as loaded. Double-check the model " + "id, or load it from the model dropdown in the UI." + ) + if not models: + _fail( + "No model is loaded in Studio. Load one from the model dropdown in " + "the UI, or pass --model to load it from here." + ) + return models[0] + + +def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: + # Codex always streams, and Studio only streams /v1/responses from llama-server. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return # older server without the endpoint; don't block the launch + raise + if status.get("is_gguf"): + return + hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" + _fail( + f"Codex needs a GGUF model served by llama-server, but {model_id} is on " + f"the transformers backend. Try: unsloth connect codex --model {hint}" + ) + + +def claude_settings_path() -> Path: + return Path.home() / ".claude" / "settings.json" + + +def ensure_claude_attribution_header() -> None: + # The header invalidates the llama.cpp KV cache (~90% slower) and Claude + # Code only honors this setting from settings.json, not the env var. + path = claude_settings_path() + settings = {} + if path.exists(): + try: + settings = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + settings = None + if not isinstance(settings, dict): + typer.echo( + f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " + 'to "0" in its "env" section yourself, or local inference will be much slower.', + err = True, + ) + return + env = settings.get("env") + if not isinstance(env, dict): + env = settings["env"] = {} + if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0": + return + env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0" + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8") + except OSError: + typer.echo( + f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " + 'to "0" in its "env" section yourself, or local inference will be much slower.', + err = True, + ) + return + typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).") + + +_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" + + +def _claude_cache_flags() -> list: + # The flag moves per-machine context (cwd, env info, git status) out of + # the system prompt, where it changes every session and defeats llama.cpp + # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p` + # passed through ctx.args); interactive sessions accept and ignore it. + # Claude Code < 2.1.98 aborts on the unknown flag, so check the version + # first; no local binary means a --no-launch printout for another machine. + executable = shutil.which("claude") + if executable is None: + return [_DYNAMIC_SECTIONS_FLAG] + try: + result = subprocess.run( + [executable, "--version"], capture_output = True, text = True, timeout = 10 + ) + version = tuple(int(part) for part in result.stdout.split()[0].split(".")) + except Exception: + return [] + return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else [] + + +def codex_home() -> Path: + return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") + + +def _merge_codex_config(existing: str, base: str) -> str: + chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table + if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): + if chunks[0] and not chunks[0].endswith("\n"): + chunks[0] += "\n" + chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' + # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. + stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") + text = "".join(c for c in chunks if not c.startswith(stale)) + if not text.endswith("\n"): + text += "\n" + if not text.endswith("\n\n"): + text += "\n" + return text + ( + f"{_PROVIDER_HEADER}\n" + 'name = "Unsloth Studio"\n' + f"base_url = {json.dumps(base + '/v1')}\n" + f'env_key = "{_CODEX_ENV_KEY}"\n' + 'wire_api = "responses"\n' + "requires_openai_auth = false\n" + ) + + +def write_codex_config(base: str, model: dict) -> None: + home = codex_home() + home.mkdir(parents = True, exist_ok = True) + + config = home / "config.toml" + existing = config.read_text(encoding = "utf-8") if config.exists() else "" + merged = _merge_codex_config(existing, base) + if merged != existing: + config.write_text(merged, encoding = "utf-8") + typer.echo(f"Updated {config}") + + # oss_provider here too: codex --oss picks the provider from it, and the + # profile layer must beat a user-set value (e.g. "ollama") in config.toml. + profile_text = ( + f'oss_provider = "{_CODEX_PROFILE}"\n' + f'model_provider = "{_CODEX_PROFILE}"\n' + f"model = {json.dumps(model['id'])}\n" + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + profile_text += f"model_context_window = {int(window)}\n" + profile = home / f"{_CODEX_PROFILE}.config.toml" + if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: + profile.write_text(profile_text, encoding = "utf-8") + typer.echo(f"Updated {profile}") + + +def _wsl_windows_executable(command: list) -> Optional[str]: + if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): + return None + executable = shutil.which(command[0]) + if executable and executable.startswith("/mnt/"): + return executable + return None + + +def _merge_wslenv(current: str, names: tuple) -> str: + entries = [entry for entry in current.split(":") if entry] + existing = {entry.split("/", 1)[0] for entry in entries} + entries.extend(name for name in names if name not in existing) + return ":".join(entries) + + +def _print_env( + env: dict, + command: list, + unset_env: tuple = (), + wsl_env_bridge: tuple = (), +) -> None: + if os.name == "nt": + for name in unset_env: + typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") + for name, value in env.items(): + # PowerShell: ` is the escape char, and $ triggers expansion inside "". + escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") + typer.echo(f'$env:{name} = "{escaped}"') + typer.echo(subprocess.list2cmdline(command)) + return + for name in unset_env: + typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") + for name, value in env.items(): + typer.echo(f"export {name}={shlex.quote(value)}") + if wsl_env_bridge: + typer.echo( + f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + typer.echo(shlex.join(command)) + + +def _launch( + command: list, + env: dict, + install_hint: str, + unset_env: tuple = (), +) -> NoReturn: + executable = shutil.which(command[0]) + if executable is None: + _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") + wsl_env_bridge = ( + tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () + ) + child_env = dict(os.environ) + if wsl_env_bridge: + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) + for name in unset_env: + child_env[name] = "" + else: + for name in unset_env: + child_env.pop(name, None) + child_env.update(env) + # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + code = subprocess.run([executable, *command[1:]], env = child_env).returncode + finally: + signal.signal(signal.SIGINT, previous) + # Negative returncode means killed by signal N; shells expect 128+N. + raise typer.Exit(code = code if code >= 0 else 128 - code) + + +def _connect(api_key: Optional[str], model: Optional[str]) -> tuple: + base = _require_studio() + key = _agent_api_key(base, api_key) + return base, key, _resolve_model(base, key, model) + + +def _run( + base: str, + entry: dict, + env: dict, + command: list, + *, + launch: bool, + install_hint: str, + unset_env: tuple = (), +) -> None: + typer.echo(f"Studio {base} · model {entry['id']}") + wsl_env_bridge = ( + tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () + ) + if not launch: + _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + return + _launch(command, env, install_hint = install_hint, unset_env = unset_env) + + +def openclaw_config_path() -> Path: + return Path.home() / ".openclaw" / "openclaw.json" + + +def write_openclaw_config(base: str, key: str, model: dict) -> None: + path = openclaw_config_path() + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). + provider_model = {"id": model["id"], "name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + provider_model["contextWindow"] = int(window) + models = _subdict(config, "models") + models.setdefault("mode", "merge") + _subdict(models, "providers")["unsloth"] = { + "baseUrl": f"{base}/v1", + "apiKey": key, + "api": "openai-completions", + "models": [provider_model], + } + # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). + defaults = _subdict(_subdict(config, "agents"), "defaults") + _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # Unauthenticated loopback gateway: without auth.mode=none the client won't open + # the websocket. The daemon must still be started separately (`openclaw gateway`). + gateway = _subdict(config, "gateway") + gateway.setdefault("mode", "local") + _subdict(gateway, "auth").setdefault("mode", "none") + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def opencode_config_path() -> Path: + config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config" + return Path(config_home) / "opencode" / "opencode.json" + + +def write_opencode_config(base: str, key: str, model: dict) -> None: + path = opencode_config_path() + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + config.setdefault("$schema", "https://opencode.ai/config.json") + _subdict(config, "provider")["unsloth"] = { + "npm": "@ai-sdk/openai-compatible", + "name": "Unsloth Studio", + "options": {"baseURL": f"{base}/v1", "apiKey": key}, + "models": {model["id"]: {"name": model["id"]}}, + } + # OpenCode selects a model by "/". + config["model"] = f"unsloth/{model['id']}" + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def hermes_config_path() -> Path: + return Path.home() / ".hermes" / "config.yaml" + + +def write_hermes_config(base: str, model: dict) -> None: + import yaml + + path = hermes_config_path() + config: dict = {} + if path.exists(): + try: + loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) + except (yaml.YAMLError, OSError): + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + if isinstance(loaded, dict): + config = loaded + elif loaded is not None: + # Non-empty, non-mapping YAML is a user-managed file; leave it. + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + # Hermes only reads the key for a *named* custom provider (a bare + # `provider: custom` ignores it), so register it under providers.*. + _subdict(config, "model").update( + provider = f"custom:{_HERMES_PROVIDER}", + default = model["id"], + api_mode = "openai", + ) + _subdict(config, "providers")[_HERMES_PROVIDER] = { + "base_url": f"{base}/v1", + "api_mode": "openai", + "key_env": _HERMES_ENV_KEY, + } + text = yaml.safe_dump(config, sort_keys = False) + if not path.exists() or path.read_text(encoding = "utf-8") != text: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(text, encoding = "utf-8") + typer.echo(f"Updated {path}") + + +@connect_app.command("claude", context_settings = _PASSTHROUGH) +def claude( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, +): + """Point Claude Code at the running Studio server and start it.""" + base, key, entry = _connect(api_key, model) + model_id = entry["id"] + ensure_claude_attribution_header() + + env = { + "ANTHROPIC_BASE_URL": base, + "ANTHROPIC_AUTH_TOKEN": key, + "ANTHROPIC_MODEL": model_id, + # Update checks, beta features, and other background requests either + # stall against a local server or evict the conversation from + # llama-server's KV-cache slots, so turn off everything nonessential. + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + } + command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args] + install_hint = ( + "irm https://claude.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://claude.ai/install.sh | bash" + ) + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + unset_env = _CLAUDE_ENV_UNSET, + ) + + +@connect_app.command("codex", context_settings = _PASSTHROUGH) +def codex( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, +): + """Point OpenAI Codex at the running Studio server and start it.""" + base, key, entry = _connect(api_key, model) + _require_gguf_for_codex(base, key, entry["id"]) + write_codex_config(base, entry) + + env = {_CODEX_ENV_KEY: key} + command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args] + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") + + +@connect_app.command("openclaw", context_settings = _PASSTHROUGH) +def openclaw( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, +): + """Point OpenClaw at the running Studio server and start it.""" + base, key, entry = _connect(api_key, model) + write_openclaw_config(base, key, entry) # key lives in the config, not the env + + command = ["openclaw", *ctx.args] + install_hint = ( + "iwr -useb https://openclaw.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://openclaw.ai/install.sh | bash" + ) + _run(base, entry, {}, command, launch = launch, install_hint = install_hint) + + +@connect_app.command("opencode", context_settings = _PASSTHROUGH) +def opencode( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, +): + """Point OpenCode at the running Studio server and start it.""" + base, key, entry = _connect(api_key, model) + write_opencode_config(base, key, entry) # key lives in the config, not the env + + command = ["opencode", *ctx.args] + _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai") + + +@connect_app.command("hermes", context_settings = _PASSTHROUGH) +def hermes( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, +): + """Point Hermes (Nous Research) at the running Studio server and start it.""" + base, key, entry = _connect(api_key, model) + write_hermes_config(base, entry) + + env = {_HERMES_ENV_KEY: key} + command = ["hermes", *ctx.args] + install_hint = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash" + ) + _run(base, entry, env, command, launch = launch, install_hint = install_hint) diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py new file mode 100644 index 000000000..06806c8fd --- /dev/null +++ b/unsloth_cli/tests/test_connect.py @@ -0,0 +1,679 @@ +# 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 `unsloth connect` — config merging and launch env, no network.""" + +from __future__ import annotations + +import json +import os +import sys +import urllib.error +from pathlib import Path +from types import SimpleNamespace + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import pytest +from typer.testing import CliRunner + +import unsloth_cli.commands.connect as connect + +BASE = "http://127.0.0.1:8888" +MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} + + +@pytest.fixture() +def claude_settings(tmp_path, monkeypatch): + path = tmp_path / "claude" / "settings.json" + monkeypatch.setattr(connect, "claude_settings_path", lambda: path) + return path + + +def test_claude_settings_created_when_missing(claude_settings): + connect.ensure_claude_attribution_header() + settings = json.loads(claude_settings.read_text()) + assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +def test_claude_settings_merge_preserves_existing(claude_settings): + claude_settings.parent.mkdir(parents = True) + claude_settings.write_text( + json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}}) + ) + connect.ensure_claude_attribution_header() + settings = json.loads(claude_settings.read_text()) + assert settings["effortLevel"] == "high" + assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0" + assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +def test_claude_settings_already_set_untouched(claude_settings): + claude_settings.parent.mkdir(parents = True) + original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}}) + claude_settings.write_text(original) + connect.ensure_claude_attribution_header() + assert claude_settings.read_text() == original + + +def test_claude_settings_bad_json_left_alone(claude_settings, capsys): + claude_settings.parent.mkdir(parents = True) + claude_settings.write_text("{not json") + connect.ensure_claude_attribution_header() + assert claude_settings.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def _fake_claude(monkeypatch, version_output: str) -> None: + monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr( + connect.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(stdout = version_output), + ) + + +def test_cache_flags_passed_to_supported_claude(monkeypatch): + _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") + assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"] + + +def test_cache_flags_skipped_on_old_claude(monkeypatch): + _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") + assert connect._claude_cache_flags() == [] + + +def test_cache_flags_skipped_on_unparseable_version(monkeypatch): + _fake_claude(monkeypatch, "weird build string\n") + assert connect._claude_cache_flags() == [] + + +def _parse_toml(text: str) -> dict: + tomllib = pytest.importorskip("tomllib") + return tomllib.loads(text) + + +def test_merge_codex_config_fresh(): + merged = connect._merge_codex_config("", BASE) + parsed = _parse_toml(merged) + assert parsed["oss_provider"] == "unsloth_api" + provider = parsed["model_providers"]["unsloth_api"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["wire_api"] == "responses" + assert provider["requires_openai_auth"] is False + + +def test_merge_codex_config_replaces_stale_block(): + existing = ( + 'model = "gpt-5"\n' + "\n" + "[model_providers.unsloth_api]\n" + 'base_url = "http://old-host:9999/v1"\n' + 'wire_api = "chat"\n' + "\n" + "[model_providers.unsloth_api.http_headers]\n" + 'x-old = "1"\n' + "\n" + "[model_providers.ollama]\n" + 'base_url = "http://localhost:11434/v1"\n' + ) + merged = connect._merge_codex_config(existing, BASE) + parsed = _parse_toml(merged) + assert parsed["model"] == "gpt-5" + assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" + assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" + assert "http_headers" not in parsed["model_providers"]["unsloth_api"] + assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" + assert connect._merge_codex_config(merged, BASE) == merged + + +def test_merge_codex_config_keeps_user_oss_provider(): + merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE) + assert _parse_toml(merged)["oss_provider"] == "ollama" + + +def test_write_codex_config_profile(tmp_path, monkeypatch): + monkeypatch.setenv("CODEX_HOME", str(tmp_path)) + connect.write_codex_config(BASE, MODEL) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert profile["oss_provider"] == "unsloth_api" + assert profile["model_provider"] == "unsloth_api" + assert profile["model"] == MODEL["id"] + assert profile["model_context_window"] == 131072 + config = _parse_toml((tmp_path / "config.toml").read_text()) + assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" + + +@pytest.fixture() +def fake_studio(tmp_path, monkeypatch, claude_settings): + calls = [] + state = {"models": [MODEL]} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return {"object": "list", "data": state["models"]} + if url.endswith("/api/inference/status"): + return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} + if url.endswith("/api/auth/api-keys"): + return {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/inference/load"): + state["models"] = [{"id": payload["model_path"], "context_length": 4096}] + return {} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) + monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") + monkeypatch.setattr(connect, "_http_json", http_json) + monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") + # No `claude` on PATH, so _claude_cache_flags never probes the real binary. + monkeypatch.setattr(connect.shutil, "which", lambda _: None) + monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) + monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) + return calls + + +def test_connect_claude_no_launch(fake_studio, claude_settings): + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "unset ANTHROPIC_API_KEY" in result.output + assert "unset CLAUDE_CODE_OAUTH_TOKEN" in result.output + assert f"export ANTHROPIC_BASE_URL={BASE}" in result.output + assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output + assert f"export ANTHROPIC_MODEL={MODEL['id']}" in result.output + assert "export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1" in result.output + assert "export CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1" in result.output + assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output + settings = json.loads(claude_settings.read_text()) + assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(connect.subprocess, "run", run) + result = CliRunner().invoke(connect.connect_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] + assert "ANTHROPIC_API_KEY" not in captured["env"] + assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + + +def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr( + connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(connect.subprocess, "run", run) + result = CliRunner().invoke(connect.connect_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == [ + "/mnt/c/Users/samle/AppData/Roaming/npm/claude", + "--model", + MODEL["id"], + ] + assert captured["env"]["ANTHROPIC_API_KEY"] == "" + assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + for name in ( + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_MODEL", + "ANTHROPIC_API_KEY", + "CLAUDE_CODE_OAUTH_TOKEN", + ): + assert name in captured["env"]["WSLENV"].split(":") + + +def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_API_KEY=" in result.output + assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output + assert "export WSLENV=" in result.output + assert "ANTHROPIC_AUTH_TOKEN" in result.output + assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output + + +def test_connect_codex_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "export UNSLOTH_STUDIO_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output + assert "codex --oss --profile unsloth_api" in result.output + assert (tmp_path / "codex" / "config.toml").exists() + assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() + + +def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): + CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert len(mints) == 1 + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["keys"] == ["sk-unsloth-feedfacefeedface"] + + +def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): + CliRunner().invoke( + connect.connect_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"keys": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]})) + inner = connect._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-stale": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(connect, "_http_json", http_json) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + # The working key moves to the front so the next run tries it first. + cached = json.loads(cache.read_text()) + assert cached["keys"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + + +def test_connect_reads_legacy_single_key_cache(fake_studio, tmp_path): + (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-oldformat" in result.output + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_model_flag_loads_on_server(fake_studio): + result = CliRunner().invoke( + connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) + ] + assert "export ANTHROPIC_MODEL=unsloth/Qwen3.5-35B-A3B" in result.output + + +def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): + # Studio registers a loaded model under a canonical id (resolved identifier + # / casing) that can differ from the path we passed. The agent must connect + # to that model, not silently fall through to the first loaded one. + requested = "Unsloth/Qwen3.5-35B-A3B" + canonical = "unsloth/Qwen3.5-35B-A3B" + inner = connect._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {"model": canonical, "display_name": canonical} + if url.endswith("/v1/models"): + # Decoy sorts first, so models[0] is the wrong pick on the old code. + return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(connect, "_http_json", http_json) + result = CliRunner().invoke( + connect.connect_app, ["claude", "--no-launch", "--model", requested] + ) + assert result.exit_code == 0, result.output + assert f"export ANTHROPIC_MODEL={canonical}" in result.output + + +def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): + monkeypatch.setattr( + connect, + "_http_json", + lambda method, url, token, payload = None, timeout = 30, error = None: ( + {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/auth/api-keys") + else {"object": "list", "data": []} + ), + ) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No model is loaded" in result.output + + +def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): + # Studio never surfaces the requested model; fail loudly rather than + # silently connecting to whatever else happens to be loaded. + inner = connect._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {} + if url.endswith("/v1/models"): + return {"object": "list", "data": [MODEL]} # decoy; request never appears + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(connect, "_http_json", http_json) + result = CliRunner().invoke( + connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] + ) + assert result.exit_code == 1 + assert "unsloth/Missing-7B" in result.output + + +def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): + inner = connect._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(connect, "_http_json", http_json) + result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) + assert result.exit_code == 1 + assert "GGUF" in result.output + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + + +def test_connect_remote_token_rejected_points_at_api_key(fake_studio, monkeypatch): + # A self-issued token is invalid against a remote Studio (different secret); + # the auto-mint 401 should become actionable --api-key guidance. + inner = connect._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/auth/api-keys"): + raise urllib.error.HTTPError(url, 401, "Invalid or expired token", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(connect, "_http_json", http_json) + result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) + assert result.exit_code == 1 + assert "Settings → API" in result.output + assert "--api-key" in result.output + + +def test_connect_no_studio_errors(fake_studio, monkeypatch): + monkeypatch.setattr(connect, "find_studio_server", lambda: None) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + + +def test_connect_explicit_api_key_skips_mint(fake_studio): + result = CliRunner().invoke( + connect.connect_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) + + +# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── + + +def test_write_openclaw_config_fresh(tmp_path, monkeypatch): + path = tmp_path / "openclaw.json" + monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) + connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) + config = json.loads(path.read_text()) + provider = config["models"]["providers"]["unsloth"] + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + assert provider["api"] == "openai-completions" + assert provider["models"] == [ + {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} + ] + # The default model must be pinned or OpenClaw has nothing active. + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["gateway"]["mode"] == "local" + assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway + if os.name != "nt": # the file holds an API key + assert path.stat().st_mode & 0o777 == 0o600 + + +def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch): + path = tmp_path / "openclaw.json" + monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) + path.write_text( + json.dumps( + { + "theme": "dark", + "agents": {"defaults": {"temperature": 0.5}}, + "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, + } + ) + ) + connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) + config = json.loads(path.read_text()) + assert config["theme"] == "dark" + assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["models"]["mode"] == "replace" # user's mode is left as-is + assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" + assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) + assert path.read_text() == before + + +def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys): + path = tmp_path / "openclaw.json" + monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) + path.write_text("{not json") + connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) + assert path.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch): + path = tmp_path / "openclaw.json" + monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) + result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "openclaw" in result.output + assert "export" not in result.output # key lives in the config, not the env + config = json.loads(path.read_text()) + assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + # OpenAI /v1/chat/completions works on either backend — no GGUF gate. + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── + + +def test_write_opencode_config_fresh(tmp_path, monkeypatch): + path = tmp_path / "opencode.json" + monkeypatch.setattr(connect, "opencode_config_path", lambda: path) + connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) + config = json.loads(path.read_text()) + provider = config["provider"]["unsloth"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} + assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}} + assert config["model"] == f"unsloth/{MODEL['id']}" + + +def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch): + path = tmp_path / "opencode.json" + monkeypatch.setattr(connect, "opencode_config_path", lambda: path) + path.write_text( + json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + ) + connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) + config = json.loads(path.read_text()) + assert config["theme"] == "tokyonight" + assert config["provider"]["anthropic"]["name"] == "Anthropic" + assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + before = path.read_text() + connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) + assert path.read_text() == before + + +def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch): + path = tmp_path / "opencode.json" + monkeypatch.setattr(connect, "opencode_config_path", lambda: path) + result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "opencode" in result.output + config = json.loads(path.read_text()) + assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"unsloth/{MODEL['id']}" + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── + + +@pytest.fixture() +def hermes_config(tmp_path, monkeypatch): + path = tmp_path / "config.yaml" + monkeypatch.setattr(connect, "hermes_config_path", lambda: path) + return path + + +def test_write_hermes_config_fresh(hermes_config): + yaml = pytest.importorskip("yaml") + connect.write_hermes_config(BASE, MODEL) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes only honors the key for a *named* custom provider, so the endpoint + # is registered under providers.* and model.provider points at it. + assert config["model"]["provider"] == "custom:unsloth" + assert config["model"]["default"] == MODEL["id"] + assert config["model"]["api_mode"] == "openai" + provider = config["providers"]["unsloth"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["api_mode"] == "openai" + assert provider["key_env"] == "UNSLOTH_API_KEY" + # The key is resolved from the launch env, never written to disk. + assert "sk-unsloth" not in hermes_config.read_text() + + +def test_write_hermes_config_preserves_and_idempotent(hermes_config): + yaml = pytest.importorskip("yaml") + hermes_config.write_text( + yaml.safe_dump( + { + "terminal": {"backend": "local"}, + "model": {"temperature": 0.7}, + "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, + } + ) + ) + connect.write_hermes_config(BASE, MODEL) + config = yaml.safe_load(hermes_config.read_text()) + assert config["terminal"] == {"backend": "local"} # unrelated sections kept + assert config["model"]["temperature"] == 0.7 # unrelated model keys kept + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + before = hermes_config.read_text() + connect.write_hermes_config(BASE, MODEL) + assert hermes_config.read_text() == before + + +def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): + pytest.importorskip("yaml") + original = "- just\n- a\n- list\n" # valid YAML, but not a mapping + hermes_config.write_text(original) + connect.write_hermes_config(BASE, MODEL) + assert hermes_config.read_text() == original # user-managed file left untouched + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_hermes_no_launch(fake_studio, hermes_config): + yaml = pytest.importorskip("yaml") + result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "export UNSLOTH_API_KEY=sk-unsloth-feedfacefeedface" in result.output + assert "hermes" in result.output + config = yaml.safe_load(hermes_config.read_text()) + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + assert config["model"]["default"] == MODEL["id"] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)