unsloth/unsloth_cli/_tool_policy.py
oobabooga 4176448fb8
Studio: enable stdio MCP servers on a loopback bind (#6295)
* Studio: enable stdio MCP servers on a loopback bind

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: address codex review on stdio MCP loopback gate

* Studio: fix banner URL and preserve stdio MCP env opt-in on network binds

* Studio: scope loopback to exact aliases and honor force-disable on run_server reuse

* Studio: cover force-disable across a public re-bind and fix a stale test comment

* Studio: keep stdio MCP off on Colab loopback launches

* Studio: set tool policy before server startup

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-15 03:02:32 +01:00

68 lines
2.1 KiB
Python

# 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
# 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.
_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.
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