diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0690cc6b79..a6b00360af 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -215,6 +215,18 @@ router = APIRouter() studio_router = APIRouter() +def _effective_enable_tools(payload) -> Optional[bool]: + """Resolve `payload.enable_tools` against the process-level tool policy. + + Returns the policy value when set (CLI hard-override from `unsloth run`), + otherwise the per-request value. + """ + from state.tool_policy import get_tool_policy + + policy = get_tool_policy() + return policy if policy is not None else payload.enable_tools + + # Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts # so is_disconnected() never fires. POST /inference/cancel looks up # in-flight cancel_events here by cancel_id (per-run) or session_id / @@ -1654,7 +1666,7 @@ async def openai_chat_completions( ) if ( using_gguf - and not payload.enable_tools + and not _effective_enable_tools(payload) and (_tools_passthrough or _has_response_format) ): # Preserve the vision guard that would otherwise run in the @@ -1745,8 +1757,13 @@ async def openai_chat_completions( created = int(time.time()) # ── Tool-calling path (agentic loop) ────────────────── + # `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools` + # hard-override the per-request value. Without a CLI override, falls + # back to `payload.enable_tools` (existing behavior). use_tools = ( - payload.enable_tools and llama_backend.supports_tools and not image_b64 + _effective_enable_tools(payload) + and llama_backend.supports_tools + and not image_b64 ) if use_tools: @@ -3470,7 +3487,9 @@ async def anthropic_messages( # Server-side agentic loop doesn't support multimodal input — matches # the `not image_b64` gate in /v1/chat/completions. server_tools = ( - payload.enable_tools and llama_backend.supports_tools and not _has_image + _effective_enable_tools(payload) + and llama_backend.supports_tools + and not _has_image ) client_tools = ( not server_tools diff --git a/studio/backend/state/tool_policy.py b/studio/backend/state/tool_policy.py new file mode 100644 index 0000000000..9343a39806 --- /dev/null +++ b/studio/backend/state/tool_policy.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Process-level server-side tool policy. + +Set by `unsloth run` at startup; consulted by the inference route gates. + + None -> no CLI override (default). Per-request `enable_tools` is honored. + True -> CLI forced tools on for every request. + False -> CLI forced tools off for every request. +""" + +from typing import Optional + +_tool_policy: Optional[bool] = None + + +def get_tool_policy() -> Optional[bool]: + return _tool_policy + + +def set_tool_policy(value: Optional[bool]) -> None: + if value is not None and not isinstance(value, bool): + raise TypeError( + f"tool_policy must be Optional[bool], got {type(value).__name__}" + ) + global _tool_policy + _tool_policy = value + + +def reset_tool_policy() -> None: + global _tool_policy + _tool_policy = None diff --git a/studio/backend/tests/test_tool_policy_gates.py b/studio/backend/tests/test_tool_policy_gates.py new file mode 100644 index 0000000000..01f6bbbc3f --- /dev/null +++ b/studio/backend/tests/test_tool_policy_gates.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for `_effective_enable_tools` -- the helper that folds the +process-level `tool_policy` over a request's `enable_tools` field. + +Truth table (policy x payload.enable_tools -> effective): + policy=None + payload=None -> None + policy=None + payload=True -> True + policy=None + payload=False -> False + policy=True + payload=* -> True + policy=False + payload=* -> False +""" + +import os +import sys +from types import SimpleNamespace + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import pytest + +from routes.inference import _effective_enable_tools +from state.tool_policy import reset_tool_policy, set_tool_policy + + +@pytest.fixture(autouse = True) +def _reset(): + reset_tool_policy() + yield + reset_tool_policy() + + +def _payload(value): + return SimpleNamespace(enable_tools = value) + + +class TestEffectiveEnableTools: + @pytest.mark.parametrize( + "payload_value,expected", + [(None, None), (True, True), (False, False)], + ) + def test_no_policy_falls_through_to_payload(self, payload_value, expected): + assert _effective_enable_tools(_payload(payload_value)) == expected + + @pytest.mark.parametrize("payload_value", [None, True, False]) + def test_policy_true_overrides_any_payload(self, payload_value): + set_tool_policy(True) + assert _effective_enable_tools(_payload(payload_value)) is True + + @pytest.mark.parametrize("payload_value", [None, True, False]) + def test_policy_false_overrides_any_payload(self, payload_value): + set_tool_policy(False) + assert _effective_enable_tools(_payload(payload_value)) is False diff --git a/studio/backend/tests/test_tool_policy_state.py b/studio/backend/tests/test_tool_policy_state.py new file mode 100644 index 0000000000..5f6b228281 --- /dev/null +++ b/studio/backend/tests/test_tool_policy_state.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for the process-level server-side tool policy used by `unsloth run`. + +The policy has three states: + None -> no CLI override (default; honor per-request enable_tools) + True -> CLI forced tools on + False -> CLI forced tools off +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import pytest + +from state.tool_policy import ( + get_tool_policy, + reset_tool_policy, + set_tool_policy, +) + + +@pytest.fixture(autouse = True) +def _reset(): + reset_tool_policy() + yield + reset_tool_policy() + + +class TestToolPolicy: + def test_default_is_none(self): + assert get_tool_policy() is None + + def test_set_true_then_get(self): + set_tool_policy(True) + assert get_tool_policy() is True + + def test_set_false_then_get(self): + set_tool_policy(False) + assert get_tool_policy() is False + + def test_set_none_clears(self): + set_tool_policy(True) + set_tool_policy(None) + assert get_tool_policy() is None + + def test_reset_clears(self): + set_tool_policy(False) + reset_tool_policy() + assert get_tool_policy() is None + + def test_rejects_non_optional_bool(self): + with pytest.raises(TypeError): + set_tool_policy("true") # type: ignore[arg-type] diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py new file mode 100644 index 0000000000..6aff02494b --- /dev/null +++ b/tests/python/test_unsloth_run_tool_policy_resolver.py @@ -0,0 +1,201 @@ +# Copyright 2025-present the Unsloth AI Inc. team. All rights reserved. + +""" +Truth-table tests for `resolve_tool_policy` -- the pure resolver behind +`unsloth run --enable-tools/--disable-tools`. + +Covers: + - 127.0.0.1 default-on, explicit on, explicit off + - 0.0.0.0 default-off, explicit off + - 0.0.0.0 + explicit on: confirm prompt unless --silent or --yes, + abort on negative answer. +""" + +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 + + +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_off(self): + assert ( + resolve_tool_policy( + host = "0.0.0.0", + flag = None, + yes = False, + silent = False, + prompt = _never_prompt, + ) + is False + ) + + 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_silent_skips_prompt(self): + assert ( + resolve_tool_policy( + host = "0.0.0.0", + flag = True, + yes = False, + 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"] + ) + def test_loopback_aliases_are_local(self, host): + assert is_external_host(host) is False + + @pytest.mark.parametrize( + "host", ["0.0.0.0", "::", "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 must follow the same rules as 0.0.0.0.""" + + def test_default_is_off(self): + assert ( + resolve_tool_policy( + host = "192.168.1.5", + flag = None, + yes = False, + silent = False, + prompt = _never_prompt, + ) + is False + ) + + def test_explicit_on_prompts(self): + seen = [] + + def _prompt(msg: str) -> bool: + seen.append(msg) + return True + + assert ( + resolve_tool_policy( + host = "192.168.1.5", + flag = True, + yes = False, + silent = False, + prompt = _prompt, + ) + is True + ) + assert any("192.168.1.5" in m for m in seen) + + 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 + ) diff --git a/unsloth_cli/_tool_policy.py b/unsloth_cli/_tool_policy.py new file mode 100644 index 0000000000..5dd6c70d9f --- /dev/null +++ b/unsloth_cli/_tool_policy.py @@ -0,0 +1,71 @@ +# Copyright 2025-present the Unsloth AI Inc. team. All rights reserved. + +"""Pure resolver for `unsloth 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 + +# Claude Code brand orange — used so the security warning stands out +# in a crowded terminal. +_PROMPT_FG = (217, 119, 87) + +# Loopback aliases. Any other bind address (0.0.0.0, ::, a specific +# LAN IP, a hostname, ...) is treated as network-reachable. +_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 _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], + yes: bool, + silent: bool, + prompt: Callable[[str], bool] = typer.confirm, +) -> bool: + """Return the resolved server-side tool policy. + + Args: + host: The bind address (e.g. "127.0.0.1", "0.0.0.0", or a + specific network IP). + flag: Tri-state from `--enable-tools/--disable-tools` + (None when neither was passed). + yes: True if the operator passed `--yes/-y`. + silent: True if the operator passed `--silent/-q`. + prompt: Callable used for the network-bind + on confirmation + (injected for testability; defaults to ``typer.confirm``). + + 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 diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index e28f0f11f1..140940209f 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -582,6 +582,20 @@ def run( host: str = typer.Option("127.0.0.1", "--host", "-H"), frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"), silent: bool = typer.Option(False, "--silent", "-q"), + enable_tools: Optional[bool] = typer.Option( + 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." + ), + ), + yes: bool = typer.Option( + False, + "--yes", + "-y", + help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.", + ), ): """Start Studio, load a model, and print an API key -- one-liner server. @@ -614,6 +628,17 @@ def run( model = parsed_repo gguf_variant = gguf_variant or embedded_variant + # ── Resolve the server-side tool policy. The y/N prompt (if any) + # runs in the outer process so the re-exec'd child never re-prompts. + from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy + + enable_tools = resolve_tool_policy( + host = host, + flag = enable_tools, + yes = yes, + silent = silent, + ) + # ── 1. Venv re-exec (same pattern as studio_default) ────────────── studio_venv_dir = STUDIO_HOME / "unsloth_studio" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) @@ -653,6 +678,18 @@ def run( args.extend(["--frontend", str(frontend)]) if silent: args.append("--silent") + # Forward the resolved tool policy (always concrete True/False + # at this point — the resolver above ran before the re-exec). + if enable_tools: + args.append("--enable-tools") + else: + args.append("--disable-tools") + # Forward --yes whenever the parent already cleared the prompt + # (either operator passed --yes, or the parent's resolver + # accepted the network-bind confirmation). Otherwise the child + # re-runs the resolver and prompts a second time. + if yes or (enable_tools and is_external_host(host)): + args.append("--yes") # Forward unknown args (llama-server pass-through) to the # re-exec'd command so the studio venv sees them in ctx.args # and the re-execed run() can include them in the load payload. @@ -678,6 +715,17 @@ def run( app = run_server(**run_kwargs) actual_port = getattr(app.state, "server_port", port) or port + # ── Apply the resolved tool policy as a process-level override. + # Must use the same import path the route handlers use -- + # `studio/backend/run.py` adds `studio/backend/` to sys.path so the + # routes import this module as top-level `state.tool_policy`. If we + # imported via `studio.backend.state.tool_policy` instead, Python + # would cache two different module objects with two different + # `_tool_policy` globals, and the gates would never see our value. + from state.tool_policy import set_tool_policy + + set_tool_policy(enable_tools) + # ── 3. Wait for server health ───────────────────────────────────── if not silent: typer.echo("Starting Unsloth Studio...") @@ -713,6 +761,32 @@ def run( base_url = f"http://{display_host}:{actual_port}" sdk_base_url = f"{base_url}/v1" + # Claude orange (Claude Code's brand color) for tool-policy notices + # so they stand out from the surrounding banner. Always printed -- + # even under --silent / --yes -- so the operator never misses the + # current tool-execution status. + _tool_notice_fg = (217, 119, 87) + _is_external = is_external_host(host) + if _is_external and enable_tools: + _tool_notice = ( + f"Server-side tools are ENABLED on {host} (network-reachable). " + f"Anyone with the API key can run code on this machine. " + f"Do not share the API key." + ) + elif _is_external: + _tool_notice = ( + f"Server-side tools are disabled by default on {host} " + 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." + ) + else: + _tool_notice = "Server-side tools are disabled." + if not silent: typer.echo("") typer.echo("=" * 56) @@ -723,6 +797,7 @@ def run( typer.echo(" OpenAI / Anthropic SDK base URL:") typer.echo(f" {sdk_base_url}") typer.echo("=" * 56) + typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) typer.echo("") typer.echo("OpenAI Chat Completions:") typer.echo(f" curl {sdk_base_url}/chat/completions \\") @@ -746,6 +821,13 @@ def run( typer.echo(' -H "Content-Type: application/json" \\') typer.echo(""" -d '{"input": "Hello", "stream": true}'""") typer.echo("") + else: + # Silent mode still prints the essentials (URL, API key) plus + # the orange tool-status notice so the operator never loses + # visibility into the security-relevant policy. + typer.echo(f"URL: {base_url}") + typer.echo(f"API Key: {api_key}") + typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) # ── 7. Wait for Ctrl+C ──────────────────────────────────────────── from studio.backend.run import _shutdown_event, _graceful_shutdown, _server