unslothai--unsloth/studio/backend/startup_banner.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

"""Terminal banner for Studio startup.

Stdlib only -- safe to import without the rest of the backend.
"""

from __future__ import annotations

import os
import sys


def _safe_print(text: str) -> None:
    """Print text without crashing on terminals that cannot encode Unicode."""
    try:
        print(text)
    except UnicodeEncodeError:
        encoding = getattr(sys.stdout, "encoding", None) or "ascii"
        try:
            print(text.encode(encoding, errors = "replace").decode(encoding))
        except LookupError:
            print(text.encode("ascii", errors = "replace").decode("ascii"))


def stdout_supports_color() -> bool:
    """True if we should emit ANSI colors."""
    if os.environ.get("NO_COLOR", "").strip():
        return False
    if os.environ.get("FORCE_COLOR", "").strip():
        return True
    try:
        return sys.stdout.isatty()
    except (AttributeError, OSError, ValueError):
        return False


def print_port_in_use_notice(original_port: int, new_port: int) -> None:
    """Message when the requested port is taken and another is chosen."""
    msg = f"Port {original_port} is in use, using port {new_port} instead."
    if stdout_supports_color():
        _safe_print(f"\033[38;5;245m{msg}\033[0m")
    else:
        _safe_print(msg)


def print_studio_stop_hint() -> None:
    """Print the trailing stop hint + closing divider, separate from the
    banner so callers can interleave content (e.g. a reachability check)."""
    use_color = stdout_supports_color()
    dim = "\033[38;5;245m"
    stop_hint_style = "\033[38;5;215;1m"
    reset = "\033[0m"

    def style(text: str, code: str) -> str:
        return f"{code}{text}{reset}" if use_color else text

    _safe_print(
        "\n".join(
            [
                "",
                style(
                    "  To stop Unsloth Studio: press Ctrl+C "
                    "(Control+C, not Command+C, on macOS).",
                    stop_hint_style,
                ),
                style("─" * 52, dim),
                "",
            ]
        )
    )


def print_studio_access_banner(
    *,
    port: int,
    bind_host: str,
    display_host: str,
    include_stop_hint: bool = True,
) -> None:
    """Pretty-print URLs once the server is listening. Set
    ``include_stop_hint=False`` to omit the trailing stop block; pair with
    :func:`print_studio_stop_hint` after inserting your own content."""
    use_color = stdout_supports_color()
    dim = "\033[38;5;245m"
    title = "\033[38;5;150m"
    local_url_style = "\033[38;5;108;1m"
    secondary = "\033[38;5;109m"
    stop_hint_style = "\033[38;5;215;1m"
    reset = "\033[0m"

    def style(text: str, code: str) -> str:
        return f"{code}{text}{reset}" if use_color else text

    ipv6_bind = bind_host in ("::", "::1")
    if ipv6_bind:
        loopback_url = f"http://[::1]:{port}"
        alt_local = f"http://localhost:{port}"
    else:
        loopback_url = f"http://127.0.0.1:{port}"
        alt_local = f"http://localhost:{port}"
    if ":" in display_host:
        external_url = f"http://[{display_host}]:{port}"
    else:
        external_url = f"http://{display_host}:{port}"

    listen_all = bind_host in ("0.0.0.0", "::")
    # The exact aliases the canned loopback_url below is valid for; any other bind
    # (e.g. a specific LAN IP) must show its real address, not http://127.0.0.1.
    loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1")

    # Use the loopback URL only when reachable on loopback; otherwise show
    # the actual bound address.
    primary_url = loopback_url if listen_all or loopback_bind else external_url
    api_base = primary_url

    lines: list[str] = [
        "",
        style("🦥 Unsloth Studio is running", title),
        style("─" * 52, dim),
        style("  On this machine -- open this in your browser:", dim),
        style(f"    {primary_url}", local_url_style),
    ]

    if (listen_all or loopback_bind) and primary_url != alt_local:
        lines.append(style(f"    (same as {alt_local})", dim))

    if listen_all and display_host not in (
        "127.0.0.1",
        "localhost",
        "::1",
        "0.0.0.0",
        "::",
    ):
        lines.extend(
            [
                "",
                style("  From another device on your network / to share:", dim),
                style(f"    {external_url}", secondary),
            ]
        )
    elif not listen_all and not loopback_bind and external_url != primary_url:
        lines.extend(
            [
                "",
                style("  Bound address:", dim),
                style(f"    {external_url}", secondary),
            ]
        )

    lines.extend(
        [
            "",
            style("  API & health:", dim),
            style(f"    {api_base}/api", secondary),
            style(f"    {api_base}/api/health", secondary),
            style("─" * 52, dim),
        ]
    )

    if loopback_bind and not listen_all:
        lines.extend(
            [
                "",
                style(
                    "  Reachable on this machine only (bound to 127.0.0.1).",
                    secondary,
                ),
                style(
                    f"  To expose it, stop and relaunch with:  unsloth studio -H 0.0.0.0 -p {port}",
                    secondary,
                ),
                style(
                    "  Only on trusted networks -- anyone who reaches this machine can use Studio.",
                    secondary,
                ),
            ]
        )

    if include_stop_hint:
        lines.extend(
            [
                "",
                style(
                    "  To stop Unsloth Studio: press Ctrl+C "
                    "(Control+C, not Command+C, on macOS).",
                    stop_hint_style,
                ),
                style("─" * 52, dim),
                "",
            ]
        )

    _safe_print("\n".join(lines))