unsloth/unsloth_cli/tests/test_connect.py

# 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}


# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and
# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form.
def _assert_env_set(output: str, name: str, value: str) -> None:
    needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}"
    assert needle in output, f"{needle!r} not found in:\n{output}"


def _assert_env_unset(output: str, name: str) -> None:
    needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}"
    assert needle in output, f"{needle!r} not found in:\n{output}"


@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)
    # Identity handshake has its own tests; trust the loopback server here.
    monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True)
    # _studio_token / api-keys are faked so the mint flow stays offline.
    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_env_unset(result.output, "ANTHROPIC_API_KEY")
    _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN")
    _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE)
    _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface")
    _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"])
    _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1")
    _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1")
    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"]


@pytest.mark.skipif(
    os.name == "nt",
    reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); "
    "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.",
)
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(":")


@pytest.mark.skipif(
    os.name == "nt",
    reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); "
    "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.",
)
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_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface")
    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"])
    # First run mints; second reuses the minted key cached for this server.
    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["servers"][BASE]["minted"] == ["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
    # Reused, not re-minted (a mint would return the feedface stand-in).
    _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef")
    cached = json.loads((tmp_path / "agent_api_key.json").read_text())
    # An explicit key is remembered as "saved" so it replays without the handshake.
    assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"]


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(
            {"servers": {BASE: {"minted": ["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_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface")
    # The working key moves to the front so the next run tries it first.
    cached = json.loads(cache.read_text())
    assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"]


def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path):
    # Legacy unscoped caches have no server binding (could leak across servers),
    # so they're ignored: a fresh key is minted and stored scoped to this server.
    (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_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface")
    cached = json.loads((tmp_path / "agent_api_key.json").read_text())
    assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"]
    assert "key" not in cached  # legacy field collapsed away


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_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B")


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_env_set(result.output, "ANTHROPIC_MODEL", canonical)


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_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch):
    # A server known only by URL + health check is unverified: keyless connect
    # must refuse and make no request at all.
    monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888")
    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
    assert fake_studio == []  # no HTTP request of any kind (no mint, no /v1/models)


def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch):
    # User named both server and key, so it's their choice; only auto-send is blocked.
    monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888")
    result = CliRunner().invoke(
        connect.connect_app,
        ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"],
    )
    assert result.exit_code == 0, result.output


def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch):
    # A key saved for a remote (non-loopback) Studio is replayed on keyless runs;
    # auto-minting stays blocked for non-loopback.
    remote = "http://studio.example:8888"
    monkeypatch.setattr(connect, "find_studio_server", lambda: remote)
    (tmp_path / "agent_api_key.json").write_text(
        json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})
    )
    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
    assert result.exit_code == 0, result.output
    _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef")
    assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio)  # never minted


def test_connect_studio_server_errors_on_explicit_remote(monkeypatch):
    # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an
    # error, not a silent local model load (which they did not ask for).
    import typer

    import unsloth_cli._inference as inference

    monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888")
    monkeypatch.setattr(
        inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888"
    )
    with pytest.raises(typer.Exit):
        inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False)


def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch):
    # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback
    # server can't be verified, fall back to a local load rather than erroring.
    import unsloth_cli._inference as inference

    monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False)
    monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888")
    monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False)
    assert (
        inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False)
        is None
    )


def test_connect_unverified_loopback_without_cached_key_refuses_to_mint(
    fake_studio, tmp_path, monkeypatch
):
    # With no saved key, the next step would auto-mint; an unverified loopback
    # server (port squatter) must be refused, with nothing sent.
    monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False)
    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
    assert result.exit_code == 1
    assert "--api-key" in result.output
    assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio)  # never minted


def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch):
    # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match)
    # replays on keyless runs without the handshake, scoped to its own base.
    cache = tmp_path / "agent_api_key.json"
    cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}))
    monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False)
    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
    assert result.exit_code == 0, result.output
    _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef")
    assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio)  # reused, not minted


def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch):
    # A "minted" key is NOT replayed to an unverified loopback server: minting and
    # minted-key replay both sit behind the handshake, so a squatter can't grab it.
    cache = tmp_path / "agent_api_key.json"
    cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}}))
    monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False)
    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
    assert result.exit_code == 1
    assert "--api-key" in result.output
    assert not any(c[1].endswith("/v1/models") for c in fake_studio)  # minted key never sent


def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch):
    # An explicit key is the user's deliberate choice, so it does not require
    # the automatic identity handshake.
    monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False)
    result = CliRunner().invoke(
        connect.connect_app,
        ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"],
    )
    assert result.exit_code == 0, result.output
    _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef")


def _serve_identity(proof_for):
    """Start a localhost HTTP server answering /api/auth/identity with
    proof_for(nonce_bytes). Returns (base_url, shutdown)."""
    import base64
    import threading
    from http.server import BaseHTTPRequestHandler, HTTPServer
    from urllib.parse import parse_qs, urlparse

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            parsed = urlparse(self.path)
            if parsed.path != "/api/auth/identity":
                self.send_response(404)
                self.end_headers()
                return
            nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0])
            host, port = self.server.server_address[0], self.server.server_address[1]
            body = json.dumps({"proof": proof_for(nonce, host, port)}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(body)

        def log_message(self, *a):
            pass

    server = HTTPServer(("127.0.0.1", 0), Handler)
    threading.Thread(target = server.serve_forever, daemon = True).start()
    base = f"http://127.0.0.1:{server.server_address[1]}"
    return base, server.shutdown


def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch):
    # Real crypto end to end: verify_studio_identity reads the install secret from
    # an isolated DB; a "good" server proves the same secret, a spoofing one can't.
    import unsloth_cli._inference as inference

    inference.ensure_studio_backend_path()
    try:
        from studio.backend.auth import storage
    except Exception as exc:  # backend not importable here (e.g. missing deps)
        pytest.skip(f"studio backend not importable: {exc}")

    monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
    monkeypatch.setattr(storage, "_identity_secret_cache", None)

    good = lambda nonce, host, port: storage.compute_identity_proof(
        nonce, host, port
    )  # real secret
    bad = lambda nonce, host, port: "00" * 32  # spoofer without the secret
    base_ok, stop_ok = _serve_identity(good)
    base_bad, stop_bad = _serve_identity(bad)
    try:
        assert inference.verify_studio_identity(base_ok) is True
        assert inference.verify_studio_identity(base_bad) is False
    finally:
        stop_ok()
        stop_bad()


def _serve_redirect(target):
    """Start a localhost server that 302-redirects every GET to target+path."""
    import threading
    from http.server import BaseHTTPRequestHandler, HTTPServer

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            self.send_response(302)
            self.send_header("Location", target + self.path)
            self.end_headers()

        def log_message(self, *a):
            pass

    server = HTTPServer(("127.0.0.1", 0), Handler)
    threading.Thread(target = server.serve_forever, daemon = True).start()
    base = f"http://127.0.0.1:{server.server_address[1]}"
    return base, server.shutdown


def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch):
    # A squatter could 302 /api/auth/identity to the real Studio and relay its
    # proof; redirects must be refused so the squatter's base isn't accepted.
    import unsloth_cli._inference as inference

    inference.ensure_studio_backend_path()
    try:
        from studio.backend.auth import storage
    except Exception as exc:
        pytest.skip(f"studio backend not importable: {exc}")

    monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
    monkeypatch.setattr(storage, "_identity_secret_cache", None)

    real_base, stop_real = _serve_identity(
        lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port)
    )
    squatter_base, stop_squatter = _serve_redirect(real_base)
    try:
        assert inference.verify_studio_identity(real_base) is True  # direct: ok
        assert inference.verify_studio_identity(squatter_base) is False  # relayed: refused
    finally:
        stop_real()
        stop_squatter()


def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch):
    # A squatter that proxies the nonce to the real Studio on another port gets a
    # proof bound to *that* port; the client expects one bound to the port it
    # connected to, so the relayed proof is rejected.
    import unsloth_cli._inference as inference

    inference.ensure_studio_backend_path()
    try:
        from studio.backend.auth import storage
    except Exception as exc:
        pytest.skip(f"studio backend not importable: {exc}")

    monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
    monkeypatch.setattr(storage, "_identity_secret_cache", None)

    real_base, stop_real = _serve_identity(
        lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port)
    )
    real_port = int(real_base.rsplit(":", 1)[1])
    # The squatter answers on its own port but returns the proof for the real port.
    squatter_base, stop_squatter = _serve_identity(
        lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port)
    )
    try:
        assert inference.verify_studio_identity(real_base) is True
        assert inference.verify_studio_identity(squatter_base) is False
    finally:
        stop_real()
        stop_squatter()


@pytest.mark.parametrize(
    "url, loopback",
    [
        ("http://127.0.0.1:8888", True),
        ("http://localhost:8888", True),
        ("http://[::1]:8888", True),
        ("http://127.0.0.5:9001", True),  # SSH tunnels can land anywhere in 127/8
        ("http://0.0.0.0:8888", False),
        ("http://10.0.0.5:8888", False),
        ("http://studio.evil.example:8888", False),
        ("https://studio.example.com", False),
    ],
)
def test_is_loopback_url(url, loopback):
    assert connect.is_loopback_url(url) is loopback


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_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef")
    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_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface")
    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)