From ce193c243da51f0a07d3e866779407f2e349e55b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 05:52:40 -0700 Subject: [PATCH] Keep server-side tools enabled under --secure (#6403) * Keep server-side tools enabled under --secure and on every bind --secure binds loopback and exposes Studio only through an authenticated Cloudflare HTTPS tunnel, but it was grouped with a raw 0.0.0.0 bind and force-disabled all server-side tools (web search, Python, terminal). The process tool policy overrode the client's enable_tools request, so the model was never told the tools existed and answered in plain text. The plain 'unsloth studio' command had no way to re-enable and printed nothing. Tools now default on for every bind. The bind host and --secure no longer change the tool policy; only an explicit --enable-tools/--disable-tools forces it on or off. Both 'unsloth studio' and 'unsloth studio run' accept the flags and the startup banner states the resolved policy. - run.py: replace _apply_default_tool_policy(host, secure) with _apply_cli_tool_policy(enable_tools); add an enable_tools kwarg to run_server and --enable-tools/--disable-tools to the argparse. - _tool_policy.py: resolve_tool_policy defaults to on for every host and no longer prompts on a network bind. - studio.py: drop the secure-as-public tool gating, add the flags to the plain command, and reword the startup banner. - Update and extend the secure-flag and tool-policy tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tool-policy notice to plain server banner and refresh run --help Follow-up to PR review: - run.py: the plain 'unsloth studio' / --secure / direct run.py path went through _emit_startup_output without any tool-policy line, so a network-reachable launch was silent about code execution now that tools default on. Thread enable_tools through _emit_startup_output / _emit_secure_startup_output and print a one-line policy notice, followed by a single stop hint. - studio.py: the 'unsloth studio run' --enable-tools/--disable-tools and --yes help still described the removed loopback-on/network-off default and the confirmation prompt; reword to match the new policy. - Add tests for the banner notice and the refreshed help text. * Update CI tool-policy resolver tests for default-on behavior tests/python/test_unsloth_run_tool_policy_resolver.py still asserted the removed network-bind policy (0.0.0.0 and LAN IP default off, explicit enable prompts and aborts on a declined prompt), so it failed the Python CI jobs. Rewrite the truth table: every bind defaults on, explicit on/off always wins, and the resolver never prompts (yes/silent/prompt kept for compatibility). * Trim comments for the tool-policy change Shorten the verbose docstrings and block comments added for --secure tool handling; keep the security-relevant intent. Verified comment-only via an AST diff (code unchanged). * Add deterministic test that server-side tools execute under --secure Drive the GGUF agentic tool loop with a fake llama-server stream and let the real execute_tool run: python counts 1..100, terminal returns a UTC datetime, and web_search runs through real _web_search with only the ddgs network boundary mocked. A policy assertion pins that the post-fix --secure path (policy None + per-request enable_tools) is what keeps these executions reachable. No model, GPU, or live network; runs in the existing backend CI. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Align _emit_startup_output banner test with the moved stop hint The tool-policy notice now prints between the access banner and the stop hint, so the stop hint is emitted once at the end instead of inline in the banner (include_stop_hint is False and print_studio_stop_hint runs once). Update the plain-localhost case to match; the mismatch and wildcard cases already asserted this wiring. --------- Co-authored-by: Michael Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/run.py | 94 ++++++-- .../tests/test_secure_tools_execute.py | 211 ++++++++++++++++++ .../backend/tests/test_secure_tunnel_gate.py | 73 +++++- .../test_unsloth_run_tool_policy_resolver.py | 99 ++++---- tests/studio/install/test_selection_logic.py | 6 +- unsloth_cli/_tool_policy.py | 45 +--- unsloth_cli/commands/studio.py | 72 +++--- unsloth_cli/tests/test_studio_secure_flag.py | 81 ++++++- 8 files changed, 509 insertions(+), 172 deletions(-) create mode 100644 studio/backend/tests/test_secure_tools_execute.py diff --git a/studio/backend/run.py b/studio/backend/run.py index 469365f58..15bb07552 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 000000000..d8c76091e --- /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 6d9b6be89..469a1619f 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 00d499636..2f13a17b6 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 bb8281161..64715a632 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 f74be6d46..5d854c083 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 fa7a10379..58a03bd12 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 2c6f1ebf9..9ed932501 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