Studio: allow --secure with --api-only (headless secure API server) and add --api-only to unsloth studio run

studio/backend/main.py

@@ -855,24 +855,16 @@ async def _recipes_redirect(rest: str = ""):
     return _RedirectResponse(url = target, status_code = 308)
 
 
-_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
-_cors_origins = ["*"]
-if _api_only:
-    _cors_origins = [
-        "tauri://localhost",  # Linux/macOS Tauri webview
-        "http://tauri.localhost",  # Windows Tauri webview
-        "http://localhost",  # dev fallback
-        "http://localhost:5173",  # Tauri dev/Vite
-        "http://127.0.0.1:5173",  # Tauri dev/Vite fallback
-    ]
-    _cors_origin_regex = None
-else:
-    _cors_origin_regex = None
+from utils.host_policy import cors_origins_for_mode  # noqa: E402
+
+_cors_origins = cors_origins_for_mode(
+    api_only = os.environ.get("UNSLOTH_API_ONLY") == "1",
+    secure = os.environ.get("UNSLOTH_SECURE") == "1",
+)
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins = _cors_origins,
-    allow_origin_regex = _cors_origin_regex,
     allow_credentials = True,
     allow_methods = ["*"],
     allow_headers = ["*"],

studio/backend/run.py

@@ -893,9 +893,14 @@ def _setup_server_disk_logging():
 def _cloudflare_tunnel_should_start(
     *, cloudflare: bool, host: str, secure: bool, api_only: bool, is_colab: bool
 ) -> bool:
-    """Whether to start the Cloudflare tunnel. --secure tunnels a loopback bind too;
-    non-secure keeps the 0.0.0.0-only rule. Colab/api-only never tunnel."""
-    return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab
+    """Whether to start the Cloudflare tunnel. --secure exposes only the tunnel
+    (loopback bind), so it tunnels even api-only (headless secure API serving);
+    otherwise tunnel only a 0.0.0.0 bind, never api-only (Tauri) or Colab."""
+    if is_colab or not cloudflare:
+        return False
+    if secure:
+        return True
+    return host == "0.0.0.0" and not api_only
 
 
 def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
@@ -919,6 +924,7 @@ def run_server(
     cloudflare: bool = True,
     secure: bool = False,
     enable_tools: "Optional[bool]" = None,
+    emit_tauri_port: bool = True,
 ):
     """
     Start the FastAPI server.
@@ -932,6 +938,9 @@ def run_server(
         llama_parallel_slots: parallel slots for llama-server
         enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
             the default (tools on, per-request enable_tools honored)
+        emit_tauri_port: print the machine-readable TAURI_PORT line the desktop
+            app parses from stdout; the headless `run --api-only` path turns it
+            off so it does not pollute the documented URL/API-key banner
 
     Note:
         Signal handlers are NOT registered here so embedders (e.g. Colab) keep
@@ -974,9 +983,13 @@ def run_server(
     if _session_log is not None and not silent:
         print(f"Session log: {_session_log}")
 
-    # Set env var BEFORE importing main so CORS middleware picks it up.
+    # Set env vars BEFORE importing main so CORS middleware picks them up.
+    # secure api-only is a remote server behind Cloudflare, so it keeps the
+    # any-origin CORS profile; plain api-only stays locked to the Tauri app.
     if api_only:
         os.environ["UNSLOTH_API_ONLY"] = "1"
+    if secure:
+        os.environ["UNSLOTH_SECURE"] = "1"
 
     import nest_asyncio
 
@@ -1158,7 +1171,8 @@ def _run():
     atexit.register(terminate_all)
 
     # Output port for Tauri (api-only), only after sockets bind and startup done.
-    if api_only:
+    # The headless `run --api-only` path opts out so it does not leak this line.
+    if api_only and emit_tauri_port:
         print(f"TAURI_PORT={port}", flush = True)
 
     # Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often

studio/backend/tests/test_secure_tunnel_gate.py

@@ -31,11 +31,14 @@
         # --no-cloudflare always wins.
         (False, "0.0.0.0", False, False, False, False),
         (False, "127.0.0.1", True, False, False, False),
-        # api-only and Colab never tunnel.
+        # Non-secure api-only never tunnels (Tauri).
         (True, "0.0.0.0", False, True, False, False),
-        (True, "127.0.0.1", True, True, False, False),
+        # --secure tunnels even api-only (headless secure API server).
+        (True, "127.0.0.1", True, True, False, True),
+        # Colab never tunnels, even --secure.
         (True, "0.0.0.0", False, False, True, False),
         (True, "127.0.0.1", True, False, True, False),
+        (True, "127.0.0.1", True, True, True, False),
     ],
 )
 def test_cloudflare_gate(cloudflare, host, secure, api_only, is_colab, expected):
@@ -162,3 +165,46 @@ def test_failclosed_message_present_in_source():
         "A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link"
         in src
     )
+
+
+@pytest.mark.parametrize(
+    "api_only,secure,expected",
+    [
+        (False, False, ["*"]),  # plain server: any origin
+        (False, True, ["*"]),  # secure UI server: any origin
+        (True, True, ["*"]),  # secure api-only: remote browsers need any origin
+        (True, False, "tauri"),  # local api-only: locked to the Tauri app
+    ],
+)
+def test_cors_origins_for_mode(api_only, secure, expected):
+    from utils.host_policy import cors_origins_for_mode
+    origins = cors_origins_for_mode(api_only = api_only, secure = secure)
+    if expected == "tauri":
+        assert origins != ["*"] and any(o.startswith("tauri://") for o in origins)
+    else:
+        assert origins == expected
+
+
+def test_run_server_exports_secure_env_for_cors():
+    # run_server must export UNSLOTH_SECURE before importing main so the CORS
+    # profile can tell remote secure serving from local Tauri use.
+    src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
+    assert 'os.environ["UNSLOTH_SECURE"] = "1"' in src
+
+
+def test_run_server_emit_tauri_port_defaults_on():
+    # Default on keeps the desktop app's stdout contract; the headless
+    # `run --api-only` path opts out explicitly.
+    import inspect
+
+    import run
+
+    params = inspect.signature(run.run_server).parameters
+    assert "emit_tauri_port" in params
+    assert params["emit_tauri_port"].default is True
+
+
+def test_tauri_port_print_is_gated_in_source():
+    # The TAURI_PORT line must depend on emit_tauri_port, not api_only alone.
+    src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
+    assert "if api_only and emit_tauri_port:" in src

studio/backend/utils/host_policy.py

@@ -34,6 +34,26 @@ def is_external_host(host: str) -> bool:
     return host.lower() not in _LOOPBACK_HOSTS
 
 
+# Tauri desktop webview origins. api-only serving (the desktop app calling a
+# local backend) locks CORS to these.
+_TAURI_CORS_ORIGINS = (
+    "tauri://localhost",  # Linux/macOS Tauri webview
+    "http://tauri.localhost",  # Windows Tauri webview
+    "http://localhost",  # dev fallback
+    "http://localhost:5173",  # Tauri dev/Vite
+    "http://127.0.0.1:5173",  # Tauri dev/Vite fallback
+)
+
+
+def cors_origins_for_mode(*, api_only: bool, secure: bool) -> list[str]:
+    """Allowed CORS origins. Default is any-origin (["*"]); api-only locks down
+    to the Tauri desktop app, except in secure mode where the API is published
+    over Cloudflare and must stay reachable from remote browser origins."""
+    if api_only and not secure:
+        return list(_TAURI_CORS_ORIGINS)
+    return ["*"]
+
+
 def apply_stdio_mcp_loopback_default(host: str, *, is_colab: bool = False) -> None:
     """Default stdio MCP servers on when bound to loopback.
 

unsloth_cli/commands/studio.py

@@ -776,6 +776,16 @@ def studio_default(
                 err = True,
             )
             raise typer.Exit(2)
+        # Same for --api-only: dropping it here would silently serve the UI.
+        if api_only:
+            typer.echo(
+                f"Error: --api-only on `unsloth studio` applies to the "
+                f"plain-server path only. For `unsloth studio "
+                f"{ctx.invoked_subcommand}`, put it after the subcommand: "
+                f"`unsloth studio {ctx.invoked_subcommand} --api-only ...`",
+                err = True,
+            )
+            raise typer.Exit(2)
         return
 
     # --secure requires the tunnel; force a loopback bind.
@@ -1031,6 +1041,12 @@ def run(
     host: str = typer.Option("127.0.0.1", "--host", "-H"),
     # `-f` removed (clustered `-fa`/`-fit*`); studio_default keeps it.
     frontend: Optional[Path] = typer.Option(None, "--frontend"),
+    api_only: bool = typer.Option(
+        False,
+        "--api-only",
+        help = "Serve only the API (no web UI), for a headless model server. "
+        "Pairs with --secure to expose the API over the Cloudflare link alone.",
+    ),
     silent: bool = typer.Option(False, "--silent", "-q"),
     enable_tools: Optional[bool] = typer.Option(
         None,
@@ -1214,6 +1230,8 @@ def run(
         args.append("--load-in-4bit" if load_in_4bit else "--no-load-in-4bit")
         if frontend:
             args.extend(["--frontend", str(frontend)])
+        if api_only:
+            args.append("--api-only")
         if silent:
             args.append("--silent")
         # Forward the resolved tool policy so the child doesn't re-resolve.
@@ -1262,9 +1280,13 @@ def run(
         host = host,
         port = port,
         silent = True,
+        api_only = api_only,
         llama_parallel_slots = parallel,
         cloudflare = cloudflare,
         secure = secure,
+        # Headless serving prints its own URL/API-key banner; the Tauri-only
+        # TAURI_PORT line would corrupt that machine-parseable output.
+        emit_tauri_port = False,
     )
     if frontend is not None:
         run_kwargs["frontend_path"] = frontend

unsloth_cli/tests/test_studio_run_parallel_flag.py

@@ -361,6 +361,29 @@ def test_studio_default_rejects_parallel_when_subcommand_invoked():
     ), f"error message must show the corrected invocation; got: {combined}"
 
 
+def test_studio_default_rejects_api_only_when_subcommand_invoked():
+    """`unsloth studio --api-only run ...` would silently serve the UI (the
+    parent's --api-only never reaches run). The callback rejects with exit 2
+    and points at the subcommand flag."""
+    studio_mod = _load_run_command()
+    import typer as _typer
+
+    app = _typer.Typer()
+    app.add_typer(studio_mod.studio_app, name = "studio")
+
+    runner = CliRunner()
+    result = runner.invoke(app, ["studio", "--api-only", "run", "--model", "X"])
+    assert result.exit_code == 2, (
+        f"expected exit 2 when --api-only is on studio group with a "
+        f"subcommand invoked; got {result.exit_code}; output={result.output!r}"
+    )
+    combined = (result.output or "") + (getattr(result, "stderr", "") or "")
+    assert "--api-only" in combined, combined
+    assert (
+        "run --api-only" in combined
+    ), f"error message must show the corrected invocation; got: {combined}"
+
+
 def test_studio_default_default_parallel_with_subcommand_does_not_error():
     """Omitting --parallel on the group must still let subcommands
     run; the group's default 1 is benign."""
@@ -445,3 +468,81 @@ def fake_run_server(**kwargs):
     assert (
         captured.get("llama_parallel_slots") == value
     ), f"run_server got llama_parallel_slots={captured.get('llama_parallel_slots')!r}, expected {value}"
+
+
+# --api-only: serve API only (no UI). Both re-exec and in-venv paths must carry it.
+
+
+def test_api_only_option_is_registered():
+    studio_mod = _load_run_command()
+    import inspect
+
+    opt = inspect.signature(studio_mod.run).parameters["api_only"].default
+    assert "--api-only" in set(getattr(opt, "param_decls", []) or [])
+    assert getattr(opt, "default", None) is False  # opt-in; plain run keeps the UI
+
+
+@pytest.mark.parametrize(
+    "extra,present",
+    [
+        (["--api-only"], True),
+        (["--secure", "--api-only"], True),  # secure headless path
+        ([], False),
+    ],
+)
+def test_reexec_forwards_api_only(monkeypatch, extra, present):
+    """`--api-only` (and only when typed) must reach the re-exec'd child."""
+    result, captured = _invoke_run(monkeypatch, _BASE + extra)
+    assert len(captured) == 1, result.output
+    argv = captured[0]["argv"]
+    assert ("--api-only" in argv) is present, argv
+
+
+@pytest.mark.parametrize("extra,expected", [(["--api-only"], True), ([], False)])
+def test_in_venv_path_passes_api_only_to_run_server(monkeypatch, extra, expected):
+    """In-venv path must forward --api-only to run_server(api_only=...)."""
+    studio_mod = _load_run_command()
+
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(sys, "prefix", str(fake_venv))
+    monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
+
+    from unsloth_cli import _tool_policy as _tp_mod
+
+    monkeypatch.setattr(
+        _tp_mod,
+        "resolve_tool_policy",
+        lambda host, flag, yes, silent: False if flag is None else bool(flag),
+    )
+
+    captured: dict = {}
+
+    def fake_run_server(**kwargs):
+        captured.update(kwargs)
+        raise _RunServerCaptured(kwargs)
+
+    fake_backend_run = sys.modules.setdefault(
+        "studio.backend.run", _types_module("studio.backend.run")
+    )
+    fake_backend_run.run_server = fake_run_server
+    fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
+    monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run)
+
+    import typer as _typer
+
+    app = _typer.Typer()
+    app.command(
+        context_settings = {
+            "allow_extra_args": True,
+            "ignore_unknown_options": True,
+        },
+    )(studio_mod.run)
+    CliRunner().invoke(app, _BASE + extra, catch_exceptions = True)
+
+    assert (
+        captured.get("api_only") is expected
+    ), f"run_server got api_only={captured.get('api_only')!r}, expected {expected}"
+    # Headless serving must suppress the Tauri-only TAURI_PORT line.
+    assert (
+        captured.get("emit_tauri_port") is False
+    ), f"run_server got emit_tauri_port={captured.get('emit_tauri_port')!r}, expected False"