# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Tool definitions and executors for LLM tool calling: web search (DuckDuckGo), Python code execution, and terminal commands.""" import ast import http.client import os import signal os.environ["UNSLOTH_IS_PRESENT"] = "1" import asyncio import random import re import shlex import ssl import subprocess import sys import tempfile import threading import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, call_tool_sync, is_stdio, list_tools_async, parse_server_headers, probe_timeout, stdio_mcp_enabled, ) from storage import mcp_servers_db from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes # Splits the UI source-map from the result; loops strip it (like __IMAGES__). RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:" # Import these at module level so the preexec_fn closure triggers no imports in # the forked child (which can deadlock multi-threaded servers). _libc = None if sys.platform == "linux": try: import ctypes import ctypes.util _libc_name = ctypes.util.find_library("c") if _libc_name: _libc = ctypes.CDLL(_libc_name, use_errno = True) except (OSError, AttributeError): pass _resource = None if sys.platform != "win32": try: import resource as _resource except ImportError: pass # Raster-image allowlist for sandbox file serving. # No .svg (XSS via embedded scripts), no .html, no .pdf. _IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output _BLOCKED_COMMANDS_COMMON = frozenset( { "rm", "dd", "chmod", "chown", "mkfs", "mount", "umount", "fdisk", "sudo", "su", "doas", "pkexec", "shutdown", "reboot", "halt", "poweroff", "kill", "killall", "pkill", "passwd", "curl", "wget", "nc", "ncat", "netcat", "socat", "ssh", "scp", "sftp", "rsync", "eval", "source", } ) _BLOCKED_COMMANDS_WIN = frozenset( { "rmdir", "takeown", "icacls", "runas", "powershell", "pwsh", } ) _BLOCKED_COMMANDS = ( _BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN if sys.platform == "win32" else _BLOCKED_COMMANDS_COMMON ) _SHELL_SEPARATORS = frozenset({";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"}) # Bash keywords starting a new command position (then $cmd, do $cmd, etc.). _SHELL_KEYWORDS_AS_SEP = frozenset({"then", "do", "else", "elif"}) # Wrappers whose next non-flag argument is the command Bash will exec. _COMMAND_PREFIXES = frozenset( { "env", "command", "builtin", "exec", "time", "nohup", "nice", "setsid", "stdbuf", "timeout", "ionice", "chroot", "sudo", "doas", "su", "xargs", } ) _ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") _FIND_EXEC_FLAGS = frozenset({"-exec", "-execdir", "-ok", "-okdir"}) def _find_blocked_commands(command: str) -> set[str]: """Detect blocked commands at shell command position only. A token is at command position if it is the first token, or follows a shell separator / brace-group opener / new-command keyword (`then`, `do`, etc.), or a command-prefix wrapper like `env` / `time` / `xargs` (next token is the real command). Tokens in argument position (`grep -r curl .`, `echo source the data`, `ls /usr/bin/curl`) pass through. Also scans `find ... -exec CMD` and recurses into bash -c / cmd /c. """ blocked: set[str] = set() # punctuation_chars splits separators into their own tokens, so command # position is detected even in `echo done; rm -rf x` (no whitespace) or # quote-split names (`r''m` collapses to `rm` after `;`). try: if sys.platform == "win32": tokens = shlex.split(command, posix = False) else: lexer = shlex.shlex(command, posix = True, punctuation_chars = ";&|()`") lexer.whitespace_split = True tokens = list(lexer) except ValueError: tokens = command.split() def _token_basename(tok: str) -> str: # Strip glued-on meta-chars (`rm;`) so the basename still matches `rm`. tok = tok.strip(";&|()`{}") base = os.path.basename(tok).lower() stem, ext = os.path.splitext(base) if ext in {".exe", ".com", ".bat", ".cmd"}: base = stem return base expect_command = True # start of string is a command position prefix_pending = False # last cmd-position token was a wrapper (env/time/xargs/...) for token in tokens: if token in _SHELL_SEPARATORS or token in _SHELL_KEYWORDS_AS_SEP: expect_command = True prefix_pending = False continue if token.startswith("-"): # Flags belong to the active command, but keep expect_command while a # wrapper prefix awaits its command (`stdbuf -oL cmd`, `xargs -- cmd`). if not prefix_pending: expect_command = False continue if not expect_command: continue # FOO=bar assignment prefix; next non-assignment token is the command. if _ASSIGNMENT_RE.match(token): continue # Numeric wrapper arg: `timeout 1 cmd` / `nice -n 5 cmd`. if prefix_pending and token.lstrip("-").isdigit(): continue base = _token_basename(token) if base in _BLOCKED_COMMANDS: blocked.add(base) # Wrappers (env/time/xargs/sudo) consume one command; the next non-flag, # non-numeric token is the real command. sudo is also in _BLOCKED_COMMANDS. if base in _COMMAND_PREFIXES: prefix_pending = True continue expect_command = False prefix_pending = False # `find ... -exec CMD ... ;` and `-execdir CMD ... ;` invoke CMD directly. for i, tok in enumerate(tokens): if tok in _FIND_EXEC_FLAGS and i + 1 < len(tokens): base = _token_basename(tokens[i + 1]) if base in _BLOCKED_COMMANDS: blocked.add(base) # Regex catches blocked words at command boundaries shlex misses: inside # $(rm -rf), <(rm), backtick chains, or "foo;rm". Anchored to command-position # delimiters, so it doesn't match in argument position. lowered = command.lower() if _BLOCKED_COMMANDS: words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS)) pattern = ( rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)" rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?" rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b" ) blocked.update(re.findall(pattern, lowered)) # Nested shell invocations (bash -c '...', bash -lc '...', cmd /c '...'): # on a -c/-/c flag, look back for a shell name (skipping flags) and # recursively scan the nested command string. _SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"} _SHELLS_WIN = {"cmd", "cmd.exe"} for i, token in enumerate(tokens): tok_lower = token.lower() # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc) is_unix_c = tok_lower == "-c" or ( tok_lower.startswith("-") and tok_lower.endswith("c") and not tok_lower.startswith("--") ) is_win_c = tok_lower == "/c" if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens): continue # Look back past flags for the shell binary. Windows flags and absolute # paths both start with /, so only skip short /X flags (not /bin/bash). for j in range(i - 1, -1, -1): prev = tokens[j] if prev.startswith("-"): continue # skip Unix flags like --login, -l if is_win_c and prev.startswith("/") and len(prev) <= 3: continue # skip Windows flags like /s, /q (not /bin/bash) prev_base = os.path.basename(prev).lower() if is_unix_c and prev_base in _SHELLS: blocked |= _find_blocked_commands(tokens[i + 1]) elif is_win_c and prev_base in _SHELLS_WIN: blocked |= _find_blocked_commands(tokens[i + 1]) break # stop at first non-flag token return blocked def _build_safe_env(workdir: str) -> dict[str, str]: """Build a minimal, credential-free environment for sandboxed subprocesses. Whitelist-built from scratch (parent env NOT inherited): only PATH/HOME/ TMPDIR/LANG/TERM/PYTHONIOENCODING (+VIRTUAL_ENV or Windows SystemRoot) reach the child; all credential vars (HF_TOKEN, AWS_*, etc.) are absent. HOME points at the sandbox workdir so SDKs can't read the operator's cached creds. """ # Start from the running interpreter's dir so 'python'/'pip' resolve to the # same environment the Studio server runs in. exe_dir = os.path.dirname(sys.executable) path_entries = [exe_dir] if exe_dir else [] # If a virtualenv is active, include its bin/Scripts directory. venv = os.environ.get("VIRTUAL_ENV") if venv: venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin") if venv_bin not in path_entries: path_entries.append(venv_bin) if sys.platform == "win32": sysroot = os.environ.get("SystemRoot", r"C:\Windows") path_entries.extend([os.path.join(sysroot, "System32"), sysroot]) else: path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) # Deduplicate, preserving order. deduped = list(dict.fromkeys(p for p in path_entries if p)) env = { "PATH": os.pathsep.join(deduped), "HOME": workdir, "TMPDIR": workdir, "LANG": os.environ.get("LANG", "C.UTF-8"), "TERM": "dumb", "PYTHONIOENCODING": "utf-8", } if venv: env["VIRTUAL_ENV"] = venv # Windows needs SystemRoot for Python/subprocess to work. if sys.platform == "win32": env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") return env def _sandbox_preexec(): """Best-effort sandbox setup for sandboxed subprocesses (modules are resolved at import time so the forked child runs no imports).""" try: os.setsid() except OSError: pass try: os.umask(0o077) except OSError: pass if _libc is not None: try: _libc.prctl(38, 1, 0, 0, 0) # PR_SET_NO_NEW_PRIVS except (OSError, AttributeError): pass try: _libc.prctl(1, 9, 0, 0, 0) # PR_SET_PDEATHSIG = SIGKILL except (OSError, AttributeError): pass # CLONE_NEWNET not applied: with userns enabled it blocks all egress, # including allowlisted hosts. Network policy is enforced by the AST # host check and the bash blocklist. if _resource is not None: # RLIMIT_NPROC is per-real-UID, so the cap is well above normal usage. try: nproc = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NPROC", "10000")) _resource.setrlimit(_resource.RLIMIT_NPROC, (nproc, nproc)) except (ValueError, OSError, AttributeError): pass try: _resource.setrlimit(_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)) except (ValueError, OSError): pass try: as_bytes = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8")) * 1024 * 1024 * 1024 _resource.setrlimit(_resource.RLIMIT_AS, (as_bytes, as_bytes)) except (ValueError, OSError, AttributeError): pass try: cpu_s = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_CPU_S", "600")) _resource.setrlimit(_resource.RLIMIT_CPU, (cpu_s, cpu_s)) except (ValueError, OSError, AttributeError): pass try: # High enough for multi-shard safetensors mmaps; tunable via env. # Clamp to the inherited hard limit so setrlimit doesn't ValueError # when the parent's hard cap is below the request. nofile = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NOFILE", "16384")) _soft_cur, hard_cur = _resource.getrlimit(_resource.RLIMIT_NOFILE) target = nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur) _resource.setrlimit(_resource.RLIMIT_NOFILE, (target, target)) except (ValueError, OSError, AttributeError): pass def _get_shell_cmd(command: str) -> list[str]: """Return the platform-appropriate shell invocation for a command string.""" if sys.platform == "win32": return ["cmd", "/c", command] return ["bash", "-c", command] # Per-session working directories so each chat thread gets its own sandbox. # Falls back to ~/studio_sandbox/_default for callers without a session_id. _workdirs: dict[str, str] = {} # Non-matching session_ids collapse to ``_invalid`` to block cross-session escapes. _SESSION_ID_RE = re.compile(r"\A[A-Za-z0-9_\-]{1,64}\Z") _PROJECT_SESSION_PREFIX = "project-" def _get_project_workdir(session_id: str) -> str | None: if not session_id.startswith(_PROJECT_SESSION_PREFIX): return None project_id = session_id[len(_PROJECT_SESSION_PREFIX) :] if not project_id or not _SESSION_ID_RE.match(project_id): return None try: from storage.studio_db import ensure_chat_project_workspace project = ensure_chat_project_workspace(project_id) except Exception: logger.warning("Failed to resolve project sandbox for %s", session_id, exc_info = True) return None if not project: return None root_path = project.get("rootPath") sandbox_path = project.get("sandboxPath") if not root_path or not sandbox_path: return None root_real = os.path.realpath(root_path) sandbox_real = os.path.realpath(sandbox_path) if sandbox_real != root_real and not sandbox_real.startswith(root_real + os.sep): return None return sandbox_real def _get_workdir(session_id: str | None = None) -> str: """Return a per-session sandbox dir at mode 0o700.""" global _workdirs key = session_id or "_default" if key not in _workdirs or not os.path.isdir(_workdirs[key]): home = os.path.expanduser("~") sandbox_root = os.path.join(home, "studio_sandbox") project_workdir = ( _get_project_workdir(session_id) if session_id and _SESSION_ID_RE.match(session_id) else None ) if project_workdir: workdir = project_workdir elif session_id and _SESSION_ID_RE.match(session_id): workdir = os.path.join(sandbox_root, session_id) if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root) + os.sep): workdir = os.path.join(sandbox_root, "_invalid") elif session_id: workdir = os.path.join(sandbox_root, "_invalid") else: workdir = os.path.join(sandbox_root, "_default") os.makedirs(workdir, exist_ok = True) try: os.chmod(sandbox_root, 0o700) except OSError: pass try: os.chmod(workdir, 0o700) except OSError: pass _workdirs[key] = workdir return _workdirs[key] def get_sandbox_workdir(session_id: str | None = None) -> str: return _get_workdir(session_id) WEB_SEARCH_TOOL = { "type": "function", "function": { "name": "web_search", "description": ( "Search the web and fetch page content. Returns snippets for all results. " "Use the url parameter to fetch full page text from a specific URL." ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The search query", }, "url": { "type": "string", "description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.", }, }, "required": [], }, }, } PYTHON_TOOL = { "type": "function", "function": { "name": "python", "description": "Execute Python code in a sandbox and return stdout/stderr.", "parameters": { "type": "object", "properties": { "code": { "type": "string", "description": "The Python code to run", } }, "required": ["code"], }, }, } TERMINAL_TOOL = { "type": "function", "function": { "name": "terminal", "description": "Execute a terminal command and return stdout/stderr.", "parameters": { "type": "object", "properties": { "command": { "type": "string", "description": "The command to run", } }, "required": ["command"], }, }, } RENDER_HTML_TOOL = { "type": "function", "function": { "name": "render_html", "description": ( "Render a self-contained HTML/CSS/JavaScript artifact for the user. " "Call this at most once per assistant response unless the user " "explicitly asks for changes in that response. Future user requests " "for new artifacts may call render_html once. Put the entire document " "in code, including any CSS in