unsloth/unsloth_cli/_inference.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

"""Model loading and streaming shared by `inference` and `chat`."""

import asyncio
import json
import os
import re
import sys
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from pathlib import Path
from typing import List, Literal, Optional

import typer

# Canonical speculative-decoding modes, mirroring the backend's
# _CANONICAL_SPEC_MODES. Named once so the CLI's option annotations, the HTTP
# payload builders and the in-process loader cannot drift apart when a mode is
# added; typer reads it at runtime to validate --speculative-type.
SpeculativeType = Literal[
    "auto", "mtp", "dspark", "dflash", "ngram", "mtp+ngram", "off", "ngram-simple"
]

_THINK_OPEN = "<think>"
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)
_STREAMED_ERROR_PREFIX = "Error: "

# Cloudflare (in front of remote Unsloth 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"
_MPI_ENV_PAIRS = (
    ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
    ("PMI_RANK", "PMI_SIZE"),
    ("PMIX_RANK", "PMIX_SIZE"),
    ("MPI_RANK", "MPI_WORLD_SIZE"),
    ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
)

# Built lazily; urllib stays function-local to match this module.
_no_redirect_opener = None


def urlopen_no_redirect(request, timeout):
    """urlopen that errors on any redirect: following a 3xx would send a bearer
    token (or accept an identity proof) to a base we never vetted, letting a port
    squatter relay a real Unsloth's response."""
    global _no_redirect_opener
    if _no_redirect_opener is None:
        import urllib.error
        import urllib.request

        class _NoRedirect(urllib.request.HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):
                raise urllib.error.HTTPError(
                    req.full_url, code, f"refusing redirect to {newurl}", headers, fp
                )

        _no_redirect_opener = urllib.request.build_opener(_NoRedirect)
    return _no_redirect_opener.open(request, timeout = timeout)


# /api/inference/load and /unload pad their body so a proxy cannot time a slow load
# out, committing the 200 before the work finishes. A failure found after that travels
# only in-band under this key (studio/backend/routes/inference.py), so a client that
# treats any 200 as success reports a failed load as a successful one.
_DEFERRED_ERROR_KEY = "_deferred_error"


def raise_for_deferred_error(url: str, body):
    """Raise the late failure a padded 200 body carries; else return ``body``.

    ``urllib.error.HTTPError`` specifically: it is the class every CLI caller already
    handles for a plain HTTP failure, so existing ``except`` blocks, messages and exit
    codes keep working, and ``.read()`` yields the same ``{"detail": ...}`` shape.
    """
    if not isinstance(body, dict):
        return body
    deferred = body.get(_DEFERRED_ERROR_KEY)
    if not isinstance(deferred, dict):
        return body

    import email.message
    import io
    import urllib.error

    status = deferred.get("status_code")
    if not isinstance(status, int) or isinstance(status, bool):
        status = 500
    detail = deferred.get("detail")
    if not isinstance(detail, str) or not detail:
        detail = "unknown error" if detail is None else json.dumps(detail)
    headers = email.message.Message()
    headers["Content-Type"] = "application/json"
    raise urllib.error.HTTPError(
        url, status, detail, headers, io.BytesIO(json.dumps({"detail": detail}).encode())
    )


def require_completed_padded_body(url: str, body):
    """Return ``body``, or raise if it is not the payload a padded route promised.

    A proxy that gives up mid-pad leaves a 200 with an empty or truncated body, so
    accepting it reports an unfinished load or unload as completed. Only the two padded
    routes commit their status that early, so only they require a payload; ``{}`` is
    rejected too, since that is what a blank body decodes to here. Mirrored by
    ``assertCompletedPaddedBody`` in studio/frontend/src/features/chat/api/padded-response.ts.
    """
    if isinstance(body, dict) and body:
        return body
    raise RuntimeError(
        f"{url} did not report completion: the connection closed before the "
        "server's reply arrived. Check the model's status before retrying."
    )


def read_json_checking_deferred_error(url: str, response):
    """Drain ``response``, then raise any deferred error its body carries.

    Draining matters on its own: stopping at the headers of a padded /load leaves the
    load running, so the caller resumes too early. An incomplete JSON payload is a
    truncated padded reply, not a success (see ``require_completed_padded_body``).
    """
    try:
        raw = response.read()
    finally:
        response.close()
    try:
        body = json.loads(raw.decode(errors = "replace") or "{}")
    except ValueError:
        body = None
    return require_completed_padded_body(url, raise_for_deferred_error(url, body))


def ensure_studio_backend_path() -> None:
    backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend")
    if backend_dir not in sys.path:
        sys.path.insert(0, backend_dir)


def configure_quiet_logging() -> None:
    import logging

    # The CLI never configures structlog, so without this every backend INFO
    # line prints. LOG_LEVEL is exported so the worker subprocess inherits it.
    level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper()
    level = getattr(logging, level_name, logging.WARNING)
    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")

    # Quieting logs must not fail a command before the import that really needs
    # structlog gets to report itself.
    try:
        import structlog
    except ModuleNotFoundError:
        return
    structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level))


def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]:
    if value is None:
        return None
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        return None
    return parsed if parsed >= 0 else None


def _first_mpi_env_pair() -> tuple[Optional[int], Optional[int]]:
    for rank_name, size_name in _MPI_ENV_PAIRS:
        rank = _parse_nonnegative_int(os.environ.get(rank_name))
        world_size = _parse_nonnegative_int(os.environ.get(size_name))
        if rank is not None and world_size is not None and world_size > 1 and rank < world_size:
            return rank, world_size
    return None, None


def _json_rank_count_from_env(name: str) -> Optional[int]:
    value = os.environ.get(name)
    if not value:
        return None
    try:
        if value.lstrip().startswith(("[", "{")):
            data = json.loads(value)
        else:
            with open(value, "r", encoding = "utf-8") as f:
                data = json.load(f)
    except (json.JSONDecodeError, OSError, UnicodeDecodeError):
        return None
    if isinstance(data, list):
        return len(data)
    if isinstance(data, dict) and isinstance(data.get("hosts"), list):
        return len(data["hosts"])
    return None


def mlx_distributed_info() -> tuple[bool, int, Optional[int]]:
    """Return launch-context metadata without initializing MLX distributed."""
    rank = _parse_nonnegative_int(os.environ.get("MLX_RANK"))
    world_size = _parse_nonnegative_int(os.environ.get("MLX_WORLD_SIZE"))
    if rank is not None:
        if (
            world_size is not None
            and world_size > 1
            and rank < world_size
            and os.environ.get("NCCL_HOST_IP")
            and os.environ.get("NCCL_PORT")
        ):
            return True, rank, world_size
        inferred_size = _json_rank_count_from_env("MLX_HOSTFILE")
        if inferred_size is not None and inferred_size > 1 and rank < inferred_size:
            return True, rank, inferred_size
        inferred_size = _json_rank_count_from_env("MLX_IBV_DEVICES")
        if (
            inferred_size is not None
            and inferred_size > 1
            and rank < inferred_size
            and os.environ.get("MLX_JACCL_COORDINATOR")
        ):
            return True, rank, inferred_size
        return False, 0, None

    mpi_rank, mpi_world_size = _first_mpi_env_pair()
    return mpi_rank is not None, mpi_rank or 0, mpi_world_size


def mlx_distributed_uses_mpi() -> bool:
    """Whether the current distributed context was launched through MPI."""
    return (
        _parse_nonnegative_int(os.environ.get("MLX_RANK")) is None
        and _first_mpi_env_pair()[0] is not None
    )


@contextmanager
def quiet_if_nonzero_mlx_rank():
    """Silence parent and child-process stdout/stderr on nonzero ranks."""
    if mlx_distributed_info()[1] == 0:
        yield
        return

    sys.stdout.flush()
    sys.stderr.flush()
    saved_stdout_fd = os.dup(1)
    saved_stderr_fd = os.dup(2)
    with open(os.devnull, "w", encoding = "utf-8") as devnull:
        try:
            os.dup2(devnull.fileno(), 1)
            os.dup2(devnull.fileno(), 2)
            with redirect_stdout(devnull), redirect_stderr(devnull):
                yield
        finally:
            sys.stdout.flush()
            sys.stderr.flush()
            os.dup2(saved_stdout_fd, 1)
            os.dup2(saved_stderr_fd, 2)
            os.close(saved_stdout_fd)
            os.close(saved_stderr_fd)


def visible_text(text: str, show_thinking: bool) -> str:
    if show_thinking:
        return text
    text = _THINK_BLOCK.sub("", text)
    # Hold back an unclosed trailing <think> so reasoning never leaks mid-stream.
    open_idx = text.find(_THINK_OPEN)
    if open_idx != -1:
        text = text[:open_idx]
    max_prefix = min(len(text), len(_THINK_OPEN) - 1)
    for size in range(max_prefix, 0, -1):
        if _THINK_OPEN.startswith(text[-size:]):
            return text[:-size]
    return text


def stream_to_stdout(stream, show_thinking: bool) -> str:
    # Backends yield the full text-so-far on each step (llama.cpp ends with a
    # metadata dict, skipped); print the growing tail, return the raw text.
    raw = ""
    shown = ""
    for chunk in stream:
        if not isinstance(chunk, str):
            continue
        raw = chunk
        rendered = visible_text(chunk, show_thinking)
        delta = rendered[len(shown) :]
        if delta:
            sys.stdout.write(delta)
            sys.stdout.flush()
        shown = rendered
    sys.stdout.write("\n")
    sys.stdout.flush()
    return raw


def stream_markdown(stream, show_thinking: bool, *, console) -> str:
    from rich.live import Live
    from rich.markdown import Markdown
    from rich.text import Text

    raw = ""
    with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live:
        for chunk in stream:
            if not isinstance(chunk, str):
                continue
            raw = chunk
            visible = visible_text(chunk, show_thinking)
            live.update(Markdown(visible) if visible.strip() else Text(""))
    return raw


def collect_stream(stream, show_thinking: bool) -> str:
    raw = ""
    for chunk in stream:
        if isinstance(chunk, str):
            raw = chunk
    return visible_text(raw, show_thinking)


def raise_on_streamed_error(stream):
    # Match real backend errors by type (GenStreamError), not the "Error:" text
    # prefix, so a completion whose text opens with "Error:" is not misread as a
    # backend failure.
    try:
        ensure_studio_backend_path()
        from core.inference.orchestrator import GenStreamError
    except Exception:
        GenStreamError = None
    for chunk in stream:
        if GenStreamError is not None and isinstance(chunk, GenStreamError):
            raise RuntimeError(str(chunk)[len(_STREAMED_ERROR_PREFIX) :].strip() or "Unknown error")
        yield chunk


def render_columns(
    left_label: str,
    left_text: str,
    right_label: str,
    right_text: str,
    *,
    console = None,
) -> None:
    from rich import box
    from rich.console import Console
    from rich.table import Table

    table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False)
    table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold")
    table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold")
    table.add_row(left_text or "", right_text or "")
    (console or Console()).print(table)


class ChatBackend:
    """Uniform stream()/close() over the llama-server and Unsloth backends."""

    def __init__(self, kind: str, backend) -> None:
        self._kind = kind  # "gguf" | "unsloth"
        self._backend = backend

    def stream(
        self,
        messages: list,
        *,
        system_prompt: str,
        temperature: float,
        top_p: float,
        top_k: int,
        max_new_tokens: int,
        repetition_penalty: float,
        enable_thinking: bool,
        use_adapter: Optional[bool] = None,
    ):
        if self._kind == "gguf":
            # llama-server takes the system prompt as the first message.
            msgs = list(messages)
            if system_prompt:
                msgs = [{"role": "system", "content": system_prompt}, *msgs]
            return self._backend.generate_chat_completion(
                messages = msgs,
                temperature = temperature,
                top_p = top_p,
                top_k = top_k,
                max_tokens = max_new_tokens,
                repetition_penalty = repetition_penalty,
                enable_thinking = enable_thinking,
            )
        gen_kwargs = dict(
            messages = messages,
            system_prompt = system_prompt,
            temperature = temperature,
            top_p = top_p,
            top_k = top_k,
            max_new_tokens = max_new_tokens,
            repetition_penalty = repetition_penalty,
            enable_thinking = enable_thinking,
        )
        if use_adapter is not None:
            return self._backend.generate_with_adapter_control(
                use_adapter = use_adapter, **gen_kwargs
            )
        return self._backend.generate_chat_response(**gen_kwargs)

    def close(self) -> None:
        # Shut the worker down directly: the graceful unload_model waits for
        # an ack that compare mode can swallow, hanging exit for minutes.
        try:
            if self._kind == "gguf":
                self._backend.unload_model()
            else:
                self._backend._shutdown_subprocess(timeout = 2.0)
        except Exception:
            pass

    def share_distributed_object(
        self,
        obj,
        *,
        timeout = 300.0,
    ):
        if self._kind != "unsloth" or not hasattr(self._backend, "share_distributed_object"):
            raise RuntimeError(
                "Distributed MLX chat requires the Unsloth MLX backend; "
                f"backend '{self._kind}' cannot broadcast chat turns."
            )
        return self._backend.share_distributed_object(obj, timeout = timeout)


def resolve_model_config(model: str, *, hf_token: Optional[str]):
    ensure_studio_backend_path()
    from utils.models import ModelConfig

    model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token)
    if not model_config:
        typer.echo("Could not resolve model config", err = True)
        raise typer.Exit(code = 1)
    return model_config


def _validate_llama_extra_args_or_exit(llama_extra_args: Optional[List[str]]) -> list[str]:
    from core.inference.llama_server_args import validate_extra_args
    try:
        return validate_extra_args(llama_extra_args)
    except ValueError as exc:
        typer.echo(f"Error: {exc}", err = True)
        raise typer.Exit(code = 1)


def _load_gguf_backend(
    model_config,
    *,
    hf_token,
    max_seq_length,
    tensor_parallel: bool = False,
    speculative_type: Optional[SpeculativeType] = None,
    spec_draft_n_max: Optional[int] = None,
    llama_extra_args: Optional[List[str]] = None,
):
    ensure_studio_backend_path()
    from core.inference.llama_cpp import GgufLoadIntent, LlamaCppBackend
    from core.inference.tensor_fallback import load_with_tensor_fallback

    llama_backend = LlamaCppBackend()
    extra_args = _validate_llama_extra_args_or_exit(llama_extra_args)
    intent_fields = dict(
        hf_variant = model_config.gguf_variant,
        model_identifier = model_config.identifier,
        is_vision = model_config.is_vision,
        n_ctx = max_seq_length,
    )
    if model_config.gguf_hf_repo:
        intent_fields.update(hf_repo = model_config.gguf_hf_repo, hf_token = hf_token)
    else:
        intent_fields.update(
            gguf_path = model_config.gguf_file,
            mmproj_path = model_config.gguf_mmproj_file,
            mtp_draft_path = model_config.gguf_mtp_file,
            dspark_draft_path = model_config.gguf_dspark_file,
            dflash_draft_path = model_config.gguf_dflash_file,
        )
    if speculative_type is not None:
        intent_fields["speculative_type"] = speculative_type
    if spec_draft_n_max is not None:
        intent_fields["spec_draft_n_max"] = spec_draft_n_max

    async def _attempt_gguf_load(
        requested_tensor_parallel: bool, attempt_extra_args: Optional[List[str]]
    ) -> bool:
        return llama_backend.load_model(
            GgufLoadIntent(
                **intent_fields,
                tensor_parallel = requested_tensor_parallel,
                extra_args = attempt_extra_args,
            )
        )

    loaded = asyncio.run(
        load_with_tensor_fallback(
            _attempt_gguf_load,
            requested_tensor = tensor_parallel,
            extra_args = extra_args,
            label = model_config.identifier,
        )
    )
    if not loaded:
        typer.echo("Model load failed", err = True)
        raise typer.Exit(code = 1)
    return ChatBackend("gguf", llama_backend)


def load_chat_backend(
    model: str,
    *,
    hf_token: Optional[str],
    max_seq_length: int,
    load_in_4bit: bool,
    tensor_parallel: bool = False,
    speculative_type: Optional[SpeculativeType] = None,
    spec_draft_n_max: Optional[int] = None,
    llama_extra_args: Optional[List[str]] = None,
    model_config = None,
    fresh_backend: bool = False,
):
    """Load `model` in-process: GGUF via llama-server, else the orchestrator.

    fresh_backend uses a private orchestrator so a second model (compare's
    base column) can run alongside the main one.
    """
    from unsloth_cli._studio_deps import studio_backend_imports

    with studio_backend_imports("unsloth inference", studio_only = True), quiet_if_nonzero_mlx_rank():
        is_mlx_distributed, rank, _world_size = mlx_distributed_info()
        if model_config is None:
            model_config = resolve_model_config(model, hf_token = hf_token)

        if is_mlx_distributed and model_config.is_gguf:
            if rank == 0:
                typer.echo(
                    "Distributed MLX inference does not support GGUF/llama.cpp models. "
                    "Use a non-GGUF MLX model under mlx.launch, or run GGUF without "
                    "mlx.launch.",
                    err = True,
                )
            raise typer.Exit(code = 1)

        if rank == 0:
            typer.echo(f"Loading {model}", err = True)

        if model_config.is_gguf:
            return _load_gguf_backend(
                model_config,
                hf_token = hf_token,
                max_seq_length = max_seq_length,
                tensor_parallel = tensor_parallel,
                speculative_type = speculative_type,
                spec_draft_n_max = spec_draft_n_max,
                llama_extra_args = llama_extra_args,
            )

        if fresh_backend:
            ensure_studio_backend_path()
            from core.inference import InferenceOrchestrator
            backend = InferenceOrchestrator()
        else:
            ensure_studio_backend_path()
            from core.inference import get_inference_backend
            backend = get_inference_backend()
        try:
            loaded = backend.load_model(
                config = model_config,
                max_seq_length = max_seq_length,
                load_in_4bit = load_in_4bit,
                hf_token = hf_token,
                tensor_parallel = tensor_parallel,
                mlx_distributed = is_mlx_distributed,
            )
        except Exception as exc:
            if not is_mlx_distributed:
                raise
            if rank == 0:
                typer.echo(str(exc) or "Model load failed", err = True)
            raise typer.Exit(code = 1)
        if not loaded:
            typer.echo("Model load failed", err = True)
            raise typer.Exit(code = 1)
    return ChatBackend("unsloth", backend)


def _loopback_candidate_bases(base: str) -> list:
    """For a bare ``localhost`` base, the concrete IP bases to try, IPv4
    127.0.0.1 first (where ``unsloth studio`` binds by default). Pinning to one
    address up front means discovery, the identity check, and the credential we
    then send all target the same endpoint instead of racing IPv4/IPv6
    resolution -- which would otherwise let the health probe land on one address
    and the identity check on another. A literal IP or remote name is unchanged.
    """
    from urllib.parse import urlparse

    parsed = urlparse(base)
    if (parsed.hostname or "").lower() != "localhost":
        return [base]
    import socket

    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    try:
        ips = {
            ai[4][0] for ai in socket.getaddrinfo(parsed.hostname, port, type = socket.SOCK_STREAM)
        }
    except Exception:
        return [base]
    ordered = sorted(ips, key = lambda ip: (ip != "127.0.0.1", ip))
    bases = [
        f"{parsed.scheme}://" + (f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}")
        for ip in ordered
    ]
    return bases or [base]


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("/")
    # Try the concrete loopback addresses in order and return the first that
    # answers, so the rest of the flow talks to that exact address.
    for candidate in _loopback_candidate_bases(base):
        request = urllib.request.Request(
            f"{candidate}/api/health", headers = {"User-Agent": _USER_AGENT}
        )
        try:
            with urllib.request.urlopen(request, timeout = timeout):
                return candidate
        except Exception:
            continue
    return None


def is_loopback_url(base: str) -> bool:
    """True only when *base* resolves to loopback. find_studio_server() trusts a
    base after only a health probe, so credentials are auto-sent only to loopback
    (a local Unsloth or an SSH tunnel on 127.0.0.1), the targets the auto flows mean."""
    from urllib.parse import urlparse

    host = (urlparse(base).hostname or "").lower()
    if host in ("localhost", "127.0.0.1", "::1"):
        return True
    try:
        import ipaddress
        return ipaddress.ip_address(host).is_loopback
    except ValueError:
        return False


def verify_studio_identity(base: str, timeout: float = 3.0) -> bool:
    """Confirm `base` is really this machine's Unsloth before sending a secret.

    Send a random nonce to /api/auth/identity and check the returned HMAC against
    the one computed from the local same-user secret; an endpoint without that
    secret (port squatter, remote/fake) can't match. Fails closed on any error."""
    import base64
    import hmac as _hmac
    import json
    import secrets as _secrets
    import socket
    import urllib.request
    from urllib.parse import urlparse

    try:
        import studio.backend.core  # noqa: F401  puts studio/backend on sys.path
        from studio.backend.auth import storage
    except Exception:
        return False

    parsed = urlparse(base)
    host = parsed.hostname or ""
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    # Resolve to one concrete address and talk to *that* address, then bind the
    # proof to (address, port). A name like localhost can resolve to a squatter on
    # ::1 while the real Unsloth is on 127.0.0.1; connecting to the resolved IP and
    # binding to it means a proof relayed from a different address/port won't match.
    try:
        ip = socket.getaddrinfo(host, port, type = socket.SOCK_STREAM)[0][4][0]
    except Exception:
        return False
    netloc = f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}"
    nonce = _secrets.token_bytes(32)
    query = base64.urlsafe_b64encode(nonce).decode()
    request = urllib.request.Request(
        f"{parsed.scheme}://{netloc}/api/auth/identity?nonce={query}",
        headers = {"User-Agent": _USER_AGENT, "Host": parsed.netloc},
    )
    try:
        # No redirects: a 302 could relay a real Unsloth's proof (see urlopen_no_redirect).
        # Cap the read: the server is still unverified, so don't trust its length.
        with urlopen_no_redirect(request, timeout = timeout) as response:
            proof = json.loads(response.read(65536).decode() or "{}").get("proof")
    except Exception:
        return False
    if not isinstance(proof, str):
        return False
    try:
        expected = storage.compute_identity_proof(nonce, ip, port)
    except Exception:
        return False
    return _hmac.compare_digest(proof, expected)


def _studio_token() -> Optional[str]:
    """Self-issue a JWT: the CLI runs as the same OS user as the server, so it
    signs with the same stored secret the server validates against."""
    try:
        import studio.backend.core  # noqa: F401  puts studio/backend on sys.path

        from studio.backend.auth import storage
        from studio.backend.auth.authentication import create_access_token

        row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone()
        return create_access_token(row[0], desktop = True) if row else None
    except Exception:
        return None


class HttpChatBackend:
    """Chat against a running Unsloth server over its OpenAI-compatible API.

    close() leaves the model loaded on purpose — the next session (or the
    UI) starts instantly.
    """

    def __init__(self, base_url: str, token: str) -> None:
        self._base = base_url
        self._token = token

    def _request(
        self,
        method: str,
        path: str,
        payload = None,
        timeout = None,
    ):
        import json
        import urllib.request

        request = urllib.request.Request(
            self._base + path,
            data = None if payload is None else json.dumps(payload).encode(),
            headers = {
                "Authorization": f"Bearer {self._token}",
                "Content-Type": "application/json",
                "User-Agent": _USER_AGENT,
            },
            method = method,
        )
        # No redirects: this carries a bearer token (see urlopen_no_redirect).
        return urlopen_no_redirect(request, timeout = timeout)

    def ensure_loaded(
        self,
        model: str,
        *,
        hf_token,
        max_seq_length,
        load_in_4bit,
        tensor_parallel: bool = False,
        speculative_type: Optional[SpeculativeType] = None,
        spec_draft_n_max: Optional[int] = None,
        llama_extra_args: Optional[List[str]] = None,
    ) -> None:
        typer.echo(f"Loading {model} on the Unsloth server", err = True)
        payload = {
            "model_path": model,
            "hf_token": hf_token,
            "max_seq_length": max_seq_length,
            "load_in_4bit": load_in_4bit,
            "tensor_parallel": tensor_parallel,
        }
        if llama_extra_args:
            payload["llama_extra_args"] = llama_extra_args
        if speculative_type is not None:
            payload["speculative_type"] = speculative_type
        if spec_draft_n_max is not None:
            payload["spec_draft_n_max"] = spec_draft_n_max
        try:
            # Read the body, don't close at the headers: a slow load commits its 200
            # early and pads until done, so closing here would generate mid-load and
            # discard the only report of a late failure.
            read_json_checking_deferred_error(
                self._base + "/api/inference/load",
                self._request("POST", "/api/inference/load", payload),
            )
        except Exception as exc:
            typer.echo(f"Model load failed: {exc}", err = True)
            raise typer.Exit(code = 1)

    def stream(
        self,
        messages: list,
        *,
        system_prompt: str,
        temperature: float,
        top_p: float,
        top_k: int,
        max_new_tokens: int,
        repetition_penalty: float,
        enable_thinking: bool,
        use_adapter: Optional[bool] = None,
    ):
        import json

        msgs = list(messages)
        if system_prompt:
            msgs = [{"role": "system", "content": system_prompt}, *msgs]
        resp = self._request(
            "POST",
            "/v1/chat/completions",
            {
                "model": "default",
                "messages": msgs,
                "stream": True,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "max_tokens": max_new_tokens,
                "repetition_penalty": repetition_penalty,
                "enable_thinking": enable_thinking,
            },
        )

        def cumulative():
            # Accumulate SSE deltas into the full-text-so-far convention the
            # stream helpers expect.
            text = ""
            with resp:
                for raw_line in resp:
                    line = raw_line.decode("utf-8", "replace").strip()
                    if not line.startswith("data:"):
                        continue
                    data = line[len("data:") :].strip()
                    if data == "[DONE]":
                        break
                    try:
                        parsed = json.loads(data)
                    except ValueError:
                        continue
                    if "error" in parsed:
                        raise RuntimeError(
                            f"Server error: {parsed['error'].get('message', 'Unknown server error')}"
                        )
                    try:
                        delta = parsed["choices"][0]["delta"].get("content")
                    except (KeyError, IndexError):
                        continue
                    if not delta:
                        continue
                    text += delta
                    # An emoji can arrive split across two deltas as lone
                    # surrogate halves: hold back a trailing half, merge pairs.
                    visible = text
                    if "\ud800" <= visible[-1] <= "\udbff":
                        visible = visible[:-1]
                    yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace")

        return cumulative()

    def close(self) -> None:
        pass


def connect_studio_server(
    model: str,
    *,
    hf_token,
    max_seq_length,
    load_in_4bit,
    tensor_parallel: bool = False,
    speculative_type: Optional[SpeculativeType] = None,
    spec_draft_n_max: Optional[int] = None,
    llama_extra_args: Optional[List[str]] = None,
):
    """Backend on a running Unsloth server, or None (caller loads locally)."""
    base_url = find_studio_server()
    if not base_url:
        return None

    # Explicit server (UNSLOTH_STUDIO_URL) we can't safely attach to -> fail loudly;
    # opportunistic local discovery just falls back to a local load.
    explicit = bool(os.environ.get("UNSLOTH_STUDIO_URL"))

    def _refuse(reason: str):
        if not explicit:
            return None
        typer.echo(
            f"Can't attach to the Unsloth server at {base_url}: {reason} Run Unsloth "
            "on this machine, or unset UNSLOTH_STUDIO_URL to load the model locally.",
            err = True,
        )
        raise typer.Exit(code = 1)

    # Only hand the self-issued JWT (signed with the local secret) to loopback: a
    # remote URL is unverified and a real remote Unsloth would reject it anyway.
    if not is_loopback_url(base_url):
        return _refuse(
            "it isn't a local Unsloth, so a self-issued token can't "
            "authenticate to it and must not be sent to it."
        )
    # Confirm the loopback responder is really our Unsloth (not a port squatter).
    if not verify_studio_identity(base_url):
        return _refuse(
            "its identity couldn't be verified (it may be running as a "
            "different OS user, or another process took the port)."
        )
    token = _studio_token()
    if not token:
        return _refuse("couldn't self-issue an Unsloth token (is Unsloth set up here?).")
    backend = HttpChatBackend(base_url, token)
    backend.ensure_loaded(
        model,
        hf_token = hf_token,
        max_seq_length = max_seq_length,
        load_in_4bit = load_in_4bit,
        tensor_parallel = tensor_parallel,
        speculative_type = speculative_type,
        spec_draft_n_max = spec_draft_n_max,
        llama_extra_args = llama_extra_args,
    )
    return backend