diff --git a/studio/backend/run.py b/studio/backend/run.py index 469365f588..15bb075527 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -395,7 +395,38 @@ def _verify_global_reachability(display_host: str, port: int) -> None: pass -def _emit_secure_startup_output(port: int) -> None: +def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str: + """One-line tool-policy summary for the plain-server startup banner, so a + network-reachable launch is never silent about code execution.""" + if enable_tools is False: + return "Server-side tools are DISABLED (--disable-tools)." + state = ( + "ENABLED (--enable-tools)" + if enable_tools + else "ENABLED by default (per-request setting honored)" + ) + if secure: + return ( + f"Server-side tools are {state}, reachable via the authenticated " + "Cloudflare HTTPS tunnel. Anyone with the API key can run code on " + "this machine. Do not share the API key. Pass --disable-tools to turn off." + ) + from utils.host_policy import is_external_host + + if host in ("0.0.0.0", "::") or is_external_host(host): + return ( + f"Server-side tools are {state} and this port is network-reachable. " + "Anyone who can reach it with the API key can run code on this " + "machine. Do not share the API key. Pass --disable-tools to turn off." + ) + return f"Server-side tools are {state} for loopback. Pass --disable-tools to turn off." + + +def _emit_tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> None: + print(_tool_policy_notice(host, secure, enable_tools), flush = True) + + +def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None) -> None: """Secure-mode banner: only the Cloudflare link (loopback has no public raw URL).""" print("") print("🦥 Unsloth Studio is running (secure)") @@ -403,6 +434,7 @@ def _emit_secure_startup_output(port: int) -> None: _print_cloudflare_line() print(f" On this machine only: http://127.0.0.1:{port}/") print("─" * 52) + _emit_tool_policy_notice("127.0.0.1", True, enable_tools) print_studio_stop_hint() @@ -411,35 +443,28 @@ def _emit_startup_output( port: int, display_host: str, secure: bool = False, + enable_tools: "Optional[bool]" = None, ) -> None: - """Print the access banner plus any post-startup warnings. - - Extracted from ``_run`` so the banner/warning wiring is testable. The - ``localhost``-to-::1 mismatch warning and the wildcard reachability - check are mutually exclusive (the mismatch helper returns None for any - non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the - trailing stop hint is emitted exactly once. - """ + """Print the access banner, post-startup warnings, the tool-policy notice, + then a single stop hint. Extracted from ``_run`` so the wiring is testable.""" if secure: - _emit_secure_startup_output(port) + _emit_secure_startup_output(port, enable_tools) return wildcard_bind = host in ("0.0.0.0", "::") localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port) - # For wildcard binds, run the reachability check between the URL - # section and the stop hint so the stop hint stays last. print_studio_access_banner( port = port, bind_host = host, display_host = display_host, - include_stop_hint = not wildcard_bind and not localhost_mismatch_url, + include_stop_hint = False, ) if localhost_mismatch_url: _print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port) - print_studio_stop_hint() elif wildcard_bind: _verify_global_reachability(display_host, port) _print_cloudflare_line() - print_studio_stop_hint() + _emit_tool_policy_notice(host, False, enable_tools) + print_studio_stop_hint() def _print_cloudflare_line() -> None: @@ -845,15 +870,15 @@ def _cloudflare_tunnel_should_start( return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab -def _apply_default_tool_policy(host: str, secure: bool) -> None: - """Force server-side tools off on network-reachable launches (0.0.0.0 or --secure) - so a public endpoint can't run code via a client's `enable_tools`. `unsloth studio - run` installs its own resolved policy and bypasses this.""" - if not (secure or host == "0.0.0.0"): +def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None: + """Honor an explicit --enable-tools/--disable-tools; None leaves the policy + unset (tools default on, per-request enable_tools honored). Host is never + inspected here.""" + if enable_tools is None: return from state.tool_policy import set_tool_policy - set_tool_policy(False) + set_tool_policy(enable_tools) def run_server( @@ -865,6 +890,7 @@ def run_server( llama_parallel_slots: int = 1, cloudflare: bool = True, secure: bool = False, + enable_tools: "Optional[bool]" = None, ): """ Start the FastAPI server. @@ -876,6 +902,8 @@ def run_server( silent: Suppress startup messages api_only: API server only, no frontend (for Tauri desktop app) 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) Note: Signal handlers are NOT registered here so embedders (e.g. Colab) keep @@ -898,8 +926,8 @@ def run_server( if secure: host = "127.0.0.1" - # `unsloth studio run` overrides this afterward with its resolved policy. - _apply_default_tool_policy(host, secure) + # `unsloth studio run` installs its own resolved policy and passes None here. + _apply_cli_tool_policy(enable_tools) # Windows cp1252 can't encode emoji; reconfigure stdout to UTF-8. if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"): @@ -1141,7 +1169,7 @@ def run_server( sys.exit(1) if not silent: - _emit_startup_output(host, port, display_host, secure = secure) + _emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools) return app @@ -1193,6 +1221,23 @@ if __name__ == "__main__": "if the tunnel can't start. Without it, --not-secure also serves the raw " "0.0.0.0 port, which is reachable from anywhere on the network", ) + # Tri-state tool policy: no flag -> None (tools on, per-request honored); + # --enable-tools/--disable-tools force on/off. + parser.add_argument( + "--enable-tools", + dest = "enable_tools", + action = "store_true", + default = None, + help = "Force server-side tools (web search, code execution) on for " + "every request. Default: on for every bind, per-request setting honored.", + ) + parser.add_argument( + "--disable-tools", + dest = "enable_tools", + action = "store_false", + default = None, + help = "Force server-side tools off for every request.", + ) # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct # backend launches; `unsloth studio run` always passes its own value (4). _PARALLEL_MIN = 1 @@ -1225,6 +1270,7 @@ if __name__ == "__main__": llama_parallel_slots = args.parallel, cloudflare = args.cloudflare, secure = args.secure, + enable_tools = args.enable_tools, ) if args.frontend is not None: kwargs["frontend_path"] = Path(args.frontend) diff --git a/studio/backend/tests/test_secure_tools_execute.py b/studio/backend/tests/test_secure_tools_execute.py new file mode 100644 index 0000000000..d8c76091e4 --- /dev/null +++ b/studio/backend/tests/test_secure_tools_execute.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Server-side tools actually EXECUTE when policy leaves them on (the `--secure` +contract). A fake llama-server stream emits native tool calls and the real +``execute_tool`` runs them: python counts 1..100, terminal returns a UTC +datetime, and web_search is exercised through real ``_web_search`` with only the +``ddgs`` network boundary mocked. No model, GPU, or live network. The policy +tie-in proves the post-fix secure path (policy ``None`` + per-request +``enable_tools``) is what keeps these executions reachable. +""" + +from __future__ import annotations + +import contextlib +import copy +import json +import re +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.llama_cpp import LlamaCppBackend +from state.tool_policy import get_tool_policy, reset_tool_policy, set_tool_policy + + +# ── Fake llama-server stream (mirrors test_llama_cpp_tool_loop.py) ── + + +def _sse(delta: dict) -> str: + return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n" + + +def _done() -> str: + return "data: [DONE]\n" + + +def _tool_call_stream(tool_name: str, arguments: dict, call_id: str) -> list[str]: + return [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": {"name": tool_name, "arguments": json.dumps(arguments)}, + } + ] + } + ), + _done(), + ] + + +def _final_stream(text: str = "Done.") -> list[str]: + return [_sse({"content": text}), _done()] + + +def _make_backend(monkeypatch, streams: list[list[str]]): + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48851 + backend._api_key = None + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + + @contextlib.contextmanager + def fake_stream_with_retry( + _client, + _url, + payload, + _cancel_event, + headers = None, + first_token_deadline = None, + ): + yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() + + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): + yield from response.chunks + + monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) + monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable) + return backend + + +def _tool_schema(name: str) -> dict: + return { + "type": "function", + "function": { + "name": name, + "description": f"{name} tool", + "parameters": {"type": "object", "properties": {}}, + }, + } + + +def _run_one_tool(monkeypatch, tool_name: str, arguments: dict) -> str: + """Drive the agentic loop with one tool call and return its real result.""" + backend = _make_backend( + monkeypatch, + [_tool_call_stream(tool_name, arguments, f"call_{tool_name}"), _final_stream()], + ) + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": f"use the {tool_name} tool"}], + tools = [_tool_schema(tool_name)], + max_tool_iterations = 1, + ) + ) + tool_ends = [ + e for e in events if e.get("type") == "tool_end" and e.get("tool_name") == tool_name + ] + assert tool_ends, f"loop never executed {tool_name}; events={[e.get('type') for e in events]}" + return tool_ends[0]["result"] + + +@pytest.fixture(autouse = True) +def _reset_policy(): + reset_tool_policy() + yield + reset_tool_policy() + + +# ── Real tool execution under the loop ── + + +def test_python_tool_counts_to_100(monkeypatch): + # "Use the python tool to count from 1 to 100." + expected = " ".join(str(i) for i in range(1, 101)) + result = _run_one_tool( + monkeypatch, "python", {"code": "print(' '.join(str(i) for i in range(1, 101)))"} + ) + assert expected in result, result # real subprocess produced the full sequence + + +def test_bash_tool_returns_current_datetime(monkeypatch): + # "Use the bash tool to provide today's datetime." Bound the parsed UTC time + # to the call window rather than a hard-coded date (survives midnight/TZ). + before = datetime.now(timezone.utc) - timedelta(seconds = 5) + result = _run_one_tool(monkeypatch, "terminal", {"command": "date -u +%Y-%m-%dT%H:%M:%SZ"}) + after = datetime.now(timezone.utc) + timedelta(seconds = 5) + + match = re.search(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z", result) + assert match, f"no UTC datetime in terminal result: {result!r}" + parsed = datetime.strptime(match.group(), "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo = timezone.utc) + assert before <= parsed <= after, f"{parsed} not in [{before}, {after}]" + + +def test_web_search_tool_runs_with_mocked_fetch(monkeypatch): + # "Web search for the weather for San Francisco's weather." Mock only the + # ddgs network boundary; real _web_search formats the canned hit. + class _FakeDDGS: + def __init__(self, *a, **k): + pass + + def text( + self, + query, + max_results = 5, + ): + return [ + { + "title": "San Francisco Weather", + "href": "https://example.test/sf", + "body": "San Francisco: sunny, 68F.", + } + ] + + monkeypatch.setattr("ddgs.DDGS", _FakeDDGS) + result = _run_one_tool(monkeypatch, "web_search", {"query": "weather in San Francisco"}) + assert "San Francisco: sunny, 68F." in result, result + assert "https://example.test/sf" in result + + +# ── Policy tie-in: the post-fix `--secure` path keeps tools reachable ── + + +class _Payload: + def __init__(self, enable_tools): + self.enable_tools = enable_tools + + +def test_effective_enable_tools_honors_secure_policy(): + from routes.inference import _effective_enable_tools + + # Post-fix --secure leaves policy None, so the request's flag is honored. + set_tool_policy(None) + assert _effective_enable_tools(_Payload(True)) is True + assert _effective_enable_tools(_Payload(False)) is False + assert get_tool_policy() is None + + # --enable-tools forces on; the old --secure (forced off) suppresses tools. + set_tool_policy(True) + assert _effective_enable_tools(_Payload(False)) is True + set_tool_policy(False) + assert _effective_enable_tools(_Payload(True)) is False diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index 6d9b6be892..469a1619f5 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -60,24 +60,79 @@ def test_run_server_accepts_secure_kwarg(): assert inspect.signature(run.run_server).parameters["secure"].default is False -def test_plain_network_launch_forces_tools_off(): - # Network-reachable launches (0.0.0.0 or --secure) must force server-side tools off. +def test_run_server_accepts_enable_tools_kwarg(): + import inspect + + import run + + params = inspect.signature(run.run_server).parameters + assert "enable_tools" in params + assert params["enable_tools"].default is None # default: leave policy unset + + +def test_tool_policy_not_auto_disabled_by_bind(): + # Tools default on for every bind; the backend only changes the policy from + # an explicit --enable-tools/--disable-tools, never from host/secure. import run from state.tool_policy import get_tool_policy, reset_tool_policy - reset_tool_policy() - run._apply_default_tool_policy("127.0.0.1", False) - assert get_tool_policy() is None # loopback: untouched (per-request honored) - - run._apply_default_tool_policy("0.0.0.0", False) - assert get_tool_policy() is False # network bind: forced off + for host in ("127.0.0.1", "localhost", "0.0.0.0"): + reset_tool_policy() + run._apply_cli_tool_policy(None) # no flag, on any bind + assert get_tool_policy() is None, host # untouched: per-request honored reset_tool_policy() - run._apply_default_tool_policy("127.0.0.1", True) # --secure (public tunnel) + run._apply_cli_tool_policy(True) # --enable-tools: forced on + assert get_tool_policy() is True + + reset_tool_policy() + run._apply_cli_tool_policy(False) # --disable-tools: forced off assert get_tool_policy() is False reset_tool_policy() +def test_tool_policy_notice_wording(): + # The plain-server startup banner states the resolved policy for every bind. + import run + + loopback = run._tool_policy_notice("127.0.0.1", False, None) + assert "ENABLED by default" in loopback and "loopback" in loopback + + network = run._tool_policy_notice("0.0.0.0", False, None) + assert "ENABLED by default" in network and "network-reachable" in network + + secure = run._tool_policy_notice("127.0.0.1", True, None) + assert "Cloudflare HTTPS tunnel" in secure + + assert run._tool_policy_notice("0.0.0.0", False, False) == ( + "Server-side tools are DISABLED (--disable-tools)." + ) + assert "ENABLED (--enable-tools)" in run._tool_policy_notice("0.0.0.0", False, True) + + +def test_startup_output_emits_tool_notice_on_network_bind(capsys, monkeypatch): + # Plain `unsloth studio -H 0.0.0.0` must not be silent about tools now. + import run + + monkeypatch.setattr(run, "_verify_global_reachability", lambda *a, **k: None) + monkeypatch.setattr(run, "_print_cloudflare_line", lambda: None) + monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None) + + run._emit_startup_output("0.0.0.0", 8000, "0.0.0.0", secure = False, enable_tools = None) + out = capsys.readouterr().out + assert "Server-side tools" in out + assert "network-reachable" in out + + +def test_startup_output_emits_disabled_notice(capsys, monkeypatch): + import run + + monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None) + run._emit_startup_output("127.0.0.1", 8000, "127.0.0.1", secure = False, enable_tools = False) + out = capsys.readouterr().out + assert "Server-side tools are DISABLED" in out + + def test_run_server_rejects_secure_without_cloudflare(): # Direct backend callers (not just the CLI) must reject the contradictory combo. import run diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py index 00d4996368..2f13a17b69 100644 --- a/tests/python/test_unsloth_run_tool_policy_resolver.py +++ b/tests/python/test_unsloth_run_tool_policy_resolver.py @@ -1,23 +1,16 @@ # Copyright 2025-present the Unsloth AI Inc. team. All rights reserved. -"""Truth-table tests for `resolve_tool_policy` behind `unsloth run --enable-tools/--disable-tools`.""" +"""Truth-table tests for `resolve_tool_policy`: tools default on for every bind +(loopback, --secure tunnel, raw network), explicit on/off wins, and the resolver +never prompts (yes/silent/prompt kept for compatibility).""" import pytest -import typer from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy def _never_prompt(_msg: str) -> bool: - raise AssertionError("prompt should not have been called") - - -def _prompt_yes(_msg: str) -> bool: - return True - - -def _prompt_no(_msg: str) -> bool: - return False + raise AssertionError("resolve_tool_policy must not prompt") class TestLocalhostHost: @@ -59,7 +52,8 @@ class TestLocalhostHost: class TestZeroHost: - def test_default_is_off(self): + def test_default_is_on(self): + # Network bind defaults ON now (operator owns network security). assert ( resolve_tool_policy( host = "0.0.0.0", @@ -68,7 +62,7 @@ class TestZeroHost: silent = False, prompt = _never_prompt, ) - is False + is True ) def test_explicit_off_no_prompt(self): @@ -83,53 +77,31 @@ class TestZeroHost: is False ) - def test_explicit_on_silent_skips_prompt(self): + def test_explicit_on_no_prompt(self): assert ( resolve_tool_policy( host = "0.0.0.0", flag = True, yes = False, + silent = False, + prompt = _never_prompt, + ) + is True + ) + + def test_yes_and_silent_accepted_but_do_not_change_result(self): + # Retained for backward compatibility; they no longer gate the result. + assert ( + resolve_tool_policy( + host = "0.0.0.0", + flag = None, + yes = True, silent = True, prompt = _never_prompt, ) is True ) - def test_explicit_on_yes_skips_prompt(self): - assert ( - resolve_tool_policy( - host = "0.0.0.0", - flag = True, - yes = True, - silent = False, - prompt = _never_prompt, - ) - is True - ) - - def test_explicit_on_prompt_yes(self): - assert ( - resolve_tool_policy( - host = "0.0.0.0", - flag = True, - yes = False, - silent = False, - prompt = _prompt_yes, - ) - is True - ) - - def test_explicit_on_prompt_no_aborts(self): - with pytest.raises(typer.Exit) as exc_info: - resolve_tool_policy( - host = "0.0.0.0", - flag = True, - yes = False, - silent = False, - prompt = _prompt_no, - ) - assert exc_info.value.exit_code == 1 - class TestIsExternalHost: @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"]) @@ -144,9 +116,9 @@ class TestIsExternalHost: class TestSpecificNetworkIP: - """Binding to a specific LAN IP must follow the same rules as 0.0.0.0.""" + """Binding to a specific LAN IP follows the same default-on rules as 0.0.0.0.""" - def test_default_is_off(self): + def test_default_is_on(self): assert ( resolve_tool_policy( host = "192.168.1.5", @@ -155,27 +127,32 @@ class TestSpecificNetworkIP: silent = False, prompt = _never_prompt, ) - is False + is True ) - def test_explicit_on_prompts(self): - seen = [] - - def _prompt(msg: str) -> bool: - seen.append(msg) - return True - + def test_explicit_on_no_prompt(self): assert ( resolve_tool_policy( host = "192.168.1.5", flag = True, yes = False, silent = False, - prompt = _prompt, + prompt = _never_prompt, ) is True ) - assert any("192.168.1.5" in m for m in seen) + + def test_explicit_off(self): + assert ( + resolve_tool_policy( + host = "192.168.1.5", + flag = False, + yes = False, + silent = False, + prompt = _never_prompt, + ) + is False + ) def test_localhost_alias_does_not_prompt(self): assert ( diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index bb82811614..64715a6327 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -422,9 +422,11 @@ class TestStudioLocalhostIpv6Warning: run_module._emit_startup_output("127.0.0.1", 8888, "127.0.0.1") - assert calls["banner"][0]["include_stop_hint"] is True + # The stop hint no longer rides inside the banner: it is printed once at the + # end so it stays the final line after the tool-policy notice. + assert calls["banner"][0]["include_stop_hint"] is False assert calls["warning"] == [] - assert calls["stop_hint"] == 0 + assert calls["stop_hint"] == 1 assert calls["reachability"] == [] @pytest.mark.parametrize("host", ["0.0.0.0", "::"]) diff --git a/unsloth_cli/_tool_policy.py b/unsloth_cli/_tool_policy.py index f74be6d461..5d854c0837 100644 --- a/unsloth_cli/_tool_policy.py +++ b/unsloth_cli/_tool_policy.py @@ -1,6 +1,6 @@ # Copyright 2025-present the Unsloth AI Inc. team. All rights reserved. -"""Pure resolver for `unsloth run --enable-tools/--disable-tools`. +"""Pure resolver for `unsloth studio [run] --enable-tools/--disable-tools`. Kept as a standalone module so the truth table can be unit-tested without spinning up Typer or the studio venv. @@ -10,9 +10,6 @@ from typing import Callable, Optional import typer -# Orange so the security warning stands out in a crowded terminal. -_PROMPT_FG = (217, 119, 87) - # Loopback aliases; any other bind address is treated as network-reachable. # Mirrored in studio/backend/utils/host_policy.py (kept separate because the # backend is self-contained); keep the two in sync. @@ -24,19 +21,6 @@ def is_external_host(host: str) -> bool: return host.lower() not in _LOOPBACK_HOSTS -def _build_prompt_text(host: str) -> str: - return typer.style( - ( - f"Tools include arbitrary code execution (Python, terminal). " - f"You're binding to {host}, which is reachable from your network. " - f"If your API key leaks, anyone with it can run code on this machine. " - f"Do not share the API key. Continue?" - ), - fg = _PROMPT_FG, - bold = True, - ) - - def resolve_tool_policy( host: str, flag: Optional[bool], @@ -44,25 +28,8 @@ def resolve_tool_policy( silent: bool, prompt: Callable[[str], bool] = typer.confirm, ) -> bool: - """Return the resolved server-side tool policy. - - Args: - host: The bind address. - flag: Tri-state from `--enable-tools/--disable-tools` (None if neither passed). - yes: True if `--yes/-y` was passed. - silent: True if `--silent/-q` was passed. - prompt: Confirmation callable (injected for testability). - - Raises: - typer.Exit: when the operator declines the confirmation. - """ - is_external = is_external_host(host) - default = not is_external # loopback defaults on, network defaults off - - resolved = default if flag is None else flag - - if is_external and resolved is True and not yes and not silent: - if not prompt(_build_prompt_text(host)): - raise typer.Exit(1) - - return resolved + """Resolve the server-side tool policy. Tools default on for every bind; + an explicit --enable-tools/--disable-tools (`flag`) forces on/off. `host`, + `yes`, `silent`, `prompt` are kept for signature compatibility and no longer + affect the result (network binds no longer prompt).""" + return True if flag is None else flag diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index fa7a103797..58a03bd12b 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -682,6 +682,12 @@ def studio_default( help = "Log every API request, including the high-frequency polling that is " "deduplicated by default.", ), + enable_tools: Optional[bool] = typer.Option( + None, + "--enable-tools/--disable-tools", + help = "Force server-side tools (web search, code execution) on or off for " + "every request. Default: on for every bind, with the per-chat UI toggle honored.", + ), ): """Launch the Unsloth Studio server.""" # Runs before every subcommand (run/setup/update/...). @@ -729,6 +735,17 @@ def studio_default( err = True, ) raise typer.Exit(2) + # Same for --enable-tools/--disable-tools: it would not reach the subcommand. + if enable_tools is not None: + _tool_flag = "--enable-tools" if enable_tools else "--disable-tools" + typer.echo( + f"Error: {_tool_flag} 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} {_tool_flag} ...`", + err = True, + ) + raise typer.Exit(2) return # --secure requires the tunnel; force a loopback bind. @@ -782,6 +799,11 @@ def studio_default( # Forward the explicit polarity (matches run.py's BooleanOptionalAction). args.append("--cloudflare" if cloudflare else "--no-cloudflare") args.append("--secure" if secure else "--not-secure") + # Forward an explicit tool policy; None -> run.py leaves it unset (tools on). + if enable_tools is True: + args.append("--enable-tools") + elif enable_tools is False: + args.append("--disable-tools") # On Windows os.execvp keeps the parent alive, so Ctrl+C # would orphan the child; use Popen+wait instead. if sys.platform == "win32": @@ -825,6 +847,7 @@ def studio_default( llama_parallel_slots = parallel, cloudflare = cloudflare, secure = secure, + enable_tools = enable_tools, ) if frontend is not None: run_kwargs["frontend_path"] = frontend @@ -981,15 +1004,15 @@ def run( None, "--enable-tools/--disable-tools", help = ( - "Force server-side tools on/off for all requests. " - "Default: on for 127.0.0.1, off for 0.0.0.0." + "Force server-side tools (web search, code execution) on or off for " + "every request. Default: on for every bind." ), ), yes: bool = typer.Option( False, "--yes", "-y", - help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.", + help = "Accepted for backward compatibility; the tool policy no longer prompts.", ), parallel: int = typer.Option( _PARALLEL_DEFAULT_RUN, @@ -1103,16 +1126,13 @@ def run( raise typer.Exit(2) host = "127.0.0.1" - # Gate tools on the *public* exposure: --secure is public via the tunnel, so - # tools default off even though the bind is loopback. - tool_policy_host = "0.0.0.0" if secure else host - - # Resolve tool policy here so the re-exec'd child inherits a - # concrete decision and never re-prompts. + # Tool policy no longer depends on the bind: tools default on everywhere + # (--secure is a loopback tunnel; the operator owns a raw bind). Resolve here + # so the re-exec'd child inherits a concrete decision. from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy enable_tools = resolve_tool_policy( - host = tool_policy_host, + host = host, flag = enable_tools, yes = yes, silent = silent, @@ -1161,8 +1181,8 @@ def run( args.append("--enable-tools") else: args.append("--disable-tools") - # Forward --yes if the parent already cleared the network-bind prompt. - if yes or (enable_tools and is_external_host(tool_policy_host)): + # Forward --yes only if the user passed it; resolution no longer prompts. + if yes: args.append("--yes") # Typer claims --parallel outside ctx.args; without this the # child reverts to its default and silently drops the value. @@ -1265,27 +1285,25 @@ def run( # Orange so the tool-policy notice stands out; printed under # --silent / --yes too so the policy is never invisible. _tool_notice_fg = (217, 119, 87) - _is_external = is_external_host(tool_policy_host) - _exposure = "the public Cloudflare tunnel" if secure else host - if _is_external and enable_tools: + _is_external = is_external_host(host) + if not enable_tools: + _tool_notice = "Server-side tools are DISABLED (--disable-tools)." + elif secure: _tool_notice = ( - f"Server-side tools are ENABLED on {_exposure} (network-reachable). " - f"Anyone with the API key can run code on this machine. " - f"Do not share the API key." + "Server-side tools are ENABLED, reachable via the authenticated " + "Cloudflare HTTPS tunnel. Anyone with the API key can run code on " + "this machine. Do not share the API key. Pass --disable-tools to turn off." ) elif _is_external: _tool_notice = ( - f"Server-side tools are disabled by default on {_exposure} " - f"(network-reachable). Pass --enable-tools to turn on " - f"(you will be warned about API-key risk)." - ) - elif enable_tools: - _tool_notice = ( - "Server-side tools are enabled by default for loopback. " - "Pass --disable-tools to turn off." + "Server-side tools are ENABLED and this port is network-reachable. " + "Anyone who can reach it with the API key can run code on this " + "machine. Do not share the API key. Pass --disable-tools to turn off." ) else: - _tool_notice = "Server-side tools are disabled." + _tool_notice = ( + "Server-side tools are ENABLED for loopback. Pass --disable-tools to turn off." + ) if not silent: typer.echo("") diff --git a/unsloth_cli/tests/test_studio_secure_flag.py b/unsloth_cli/tests/test_studio_secure_flag.py index 2c6f1ebf9e..9ed9325011 100644 --- a/unsloth_cli/tests/test_studio_secure_flag.py +++ b/unsloth_cli/tests/test_studio_secure_flag.py @@ -237,11 +237,12 @@ def test_studio_default_rejects_secure_with_subcommand(): assert "--secure" in combined, combined -# ── secure resolves tools against the PUBLIC exposure, not the loopback bind ── +# ── secure resolves tools against the loopback bind (tools stay ON) ── -def test_run_secure_resolves_tools_against_public_host(monkeypatch): - # --secure is public via the tunnel, so tools resolve against 0.0.0.0 (OFF), not loopback (ON). +def test_run_secure_resolves_tools_against_loopback(monkeypatch): + # --secure is a loopback bind behind an authenticated tunnel, so tools resolve + # against 127.0.0.1 (ON): the child gets --enable-tools, not --disable-tools. studio_mod = _studio() monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") fake_venv = Path("/fake/studio/venv/unsloth_studio") @@ -261,7 +262,7 @@ def test_run_secure_resolves_tools_against_public_host(monkeypatch): def rec(host, flag, yes, silent): calls.append(host) - return (not _tp_mod.is_external_host(host)) if flag is None else bool(flag) + return True if flag is None else bool(flag) # default ON everywhere monkeypatch.setattr(_tp_mod, "resolve_tool_policy", rec) @@ -281,14 +282,74 @@ def test_run_secure_resolves_tools_against_public_host(monkeypatch): )(studio_mod.run) CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0", "--secure"], catch_exceptions = True) - assert calls and calls[0] == "0.0.0.0", calls + # Resolved against the forced-loopback bind, not the public 0.0.0.0 exposure. + assert calls and calls[0] == "127.0.0.1", calls + assert len(captured) == 1, captured + assert "--enable-tools" in captured[0] and "--disable-tools" not in captured[0], captured[0] + + +def test_run_secure_enable_tools_no_auto_yes(monkeypatch): + # No prompt now, so a secure --enable-tools forwards --enable-tools but not + # --yes (only an explicit --yes is forwarded). + captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure", "--enable-tools"]) + assert len(captured) == 1, captured + argv = captured[0] + assert "--enable-tools" in argv, argv + assert "--yes" not in argv, argv + + +# ── plain `unsloth studio` exposes + forwards --enable-tools/--disable-tools ── + + +def test_studio_default_exposes_enable_tools_option_default_none(): + import inspect + + opt = inspect.signature(_studio().studio_default).parameters["enable_tools"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "--enable-tools/--disable-tools" in decls + assert opt.default is None # tri-state: omitted -> leave policy unset (tools on) + + +def test_studio_default_forwards_disable_tools(monkeypatch): + captured = _invoke_studio_default(monkeypatch, ["--disable-tools"]) assert len(captured) == 1, captured assert "--disable-tools" in captured[0] and "--enable-tools" not in captured[0], captured[0] -def test_run_secure_enable_tools_forwards_yes(monkeypatch): - # Enabling tools on a secure endpoint forwards --yes so the child doesn't re-prompt. - captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure", "--enable-tools"]) +def test_studio_default_forwards_enable_tools(monkeypatch): + captured = _invoke_studio_default(monkeypatch, ["--enable-tools"]) assert len(captured) == 1, captured - argv = captured[0] - assert "--enable-tools" in argv and "--yes" in argv, argv + assert "--enable-tools" in captured[0] and "--disable-tools" not in captured[0], captured[0] + + +def test_studio_default_no_tool_flag_omits_both(monkeypatch): + # No flag -> neither flag forwarded; run.py leaves the policy unset (tools on). + captured = _invoke_studio_default(monkeypatch, []) + assert len(captured) == 1, captured + assert "--enable-tools" not in captured[0] and "--disable-tools" not in captured[0], captured[0] + + +def test_studio_default_rejects_enable_tools_with_subcommand(): + import typer as _typer + + studio_mod = _studio() + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + result = CliRunner().invoke(app, ["studio", "--enable-tools", "run", "--model", "X"]) + assert result.exit_code == 2, result.output + combined = (result.output or "") + (getattr(result, "stderr", "") or "") + assert "--enable-tools" in combined, combined + + +def test_run_tool_help_reflects_default_on_everywhere(): + # Help must match the new policy (tools on everywhere, no prompt). + import inspect + + params = inspect.signature(_studio().run).parameters + tools_help = params["enable_tools"].default.help or "" + assert "on for every bind" in tools_help, tools_help + assert "0.0.0.0" not in tools_help, tools_help + + yes_help = params["yes"].default.help or "" + assert "Skip the 0.0.0.0" not in yes_help, yes_help + assert "no longer prompts" in yes_help, yes_help