unsloth/tests/python/test_unsloth_run_tool_policy_resolver.py
Daniel Han ce193c243d
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 <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 05:52:40 -07:00

167 lines
4.6 KiB
Python

# Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
"""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
from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy
def _never_prompt(_msg: str) -> bool:
raise AssertionError("resolve_tool_policy must not prompt")
class TestLocalhostHost:
@pytest.mark.parametrize("flag", [None, True, False])
def test_no_prompt(self, flag):
# localhost never prompts regardless of flag
result = resolve_tool_policy(
host = "127.0.0.1",
flag = flag,
yes = False,
silent = False,
prompt = _never_prompt,
)
assert result is (True if flag in (None, True) else False)
def test_default_is_on(self):
assert (
resolve_tool_policy(
host = "127.0.0.1",
flag = None,
yes = False,
silent = False,
prompt = _never_prompt,
)
is True
)
def test_explicit_off(self):
assert (
resolve_tool_policy(
host = "127.0.0.1",
flag = False,
yes = False,
silent = False,
prompt = _never_prompt,
)
is False
)
class TestZeroHost:
def test_default_is_on(self):
# Network bind defaults ON now (operator owns network security).
assert (
resolve_tool_policy(
host = "0.0.0.0",
flag = None,
yes = False,
silent = False,
prompt = _never_prompt,
)
is True
)
def test_explicit_off_no_prompt(self):
assert (
resolve_tool_policy(
host = "0.0.0.0",
flag = False,
yes = False,
silent = False,
prompt = _never_prompt,
)
is False
)
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
)
class TestIsExternalHost:
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"])
def test_loopback_aliases_are_local(self, host):
assert is_external_host(host) is False
@pytest.mark.parametrize(
"host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.5", "10.0.0.1", "example.com"]
)
def test_non_loopback_is_external(self, host):
assert is_external_host(host) is True
class TestSpecificNetworkIP:
"""Binding to a specific LAN IP follows the same default-on rules as 0.0.0.0."""
def test_default_is_on(self):
assert (
resolve_tool_policy(
host = "192.168.1.5",
flag = None,
yes = False,
silent = False,
prompt = _never_prompt,
)
is True
)
def test_explicit_on_no_prompt(self):
assert (
resolve_tool_policy(
host = "192.168.1.5",
flag = True,
yes = False,
silent = False,
prompt = _never_prompt,
)
is True
)
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 (
resolve_tool_policy(
host = "localhost",
flag = True,
yes = False,
silent = False,
prompt = _never_prompt,
)
is True
)