vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Keep server-side tools enabled under --secure (#6403)

Browse source 
* 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 <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-18 05:52:40 -07:00 • committed by GitHub
parent 7fecce4e49
commit ce193c243d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 509 additions and 172 deletions
Download patch file Download diff file Expand all files Collapse all files

94
studio/backend/run.py
View file

@ -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)

211
studio/backend/tests/test_secure_tools_execute.py Normal file
View file

@ -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

73
studio/backend/tests/test_secure_tunnel_gate.py
View file

@ -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

99
tests/python/test_unsloth_run_tool_policy_resolver.py
View file

@ -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 (

6
tests/studio/install/test_selection_logic.py
View file

@ -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", "::"])

45
unsloth_cli/_tool_policy.py
View file

@ -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

72
unsloth_cli/commands/studio.py
View file

@ -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("")

81
unsloth_cli/tests/test_studio_secure_flag.py
View file

@ -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