diff --git a/studio/backend/main.py b/studio/backend/main.py index 017e78fb3..a56bd46c4 100644 --- a/studio/backend/main.py +++ b/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 = ["*"], diff --git a/studio/backend/run.py b/studio/backend/run.py index 709efc209..bb57db9bc 100644 --- a/studio/backend/run.py +++ b/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_server( 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 diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index d254121c1..bf9836e8f 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -31,11 +31,14 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402 # --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 diff --git a/studio/backend/utils/host_policy.py b/studio/backend/utils/host_policy.py index bd9ebd68b..f506eadc0 100644 --- a/studio/backend/utils/host_policy.py +++ b/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. diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index cde3d2cdb..5975e7e35 100644 --- a/unsloth_cli/commands/studio.py +++ b/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 diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py index 8caf31a43..558b268a4 100644 --- a/unsloth_cli/tests/test_studio_run_parallel_flag.py +++ b/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 test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value): 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"