unsloth/unsloth_cli/_tool_policy.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

35 lines
1.2 KiB
Python

# Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
"""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.
"""
from typing import Callable, Optional
import typer
# 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.
_LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})
def is_external_host(host: str) -> bool:
"""True when `host` is reachable from beyond loopback."""
return host.lower() not in _LOOPBACK_HOSTS
def resolve_tool_policy(
host: str,
flag: Optional[bool],
yes: bool,
silent: bool,
prompt: Callable[[str], bool] = typer.confirm,
) -> bool:
"""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