diff --git a/unsloth_cli/claude_subagent_mcp.py b/unsloth_cli/claude_subagent_mcp.py index e044d78705..b66d3c7ab9 100644 --- a/unsloth_cli/claude_subagent_mcp.py +++ b/unsloth_cli/claude_subagent_mcp.py @@ -29,6 +29,10 @@ from unsloth_cli.commands.start import ( _MAX_RESULT_CHARACTERS = 100_000 _CANCEL_POLL_SECONDS = 0.1 _CANCEL_GRACE_SECONDS = 2.0 +# A local server that accepts the connection and then never answers leaves the +# child, and the parent waiting on it, blocked forever. Generous enough not to cut +# a long legitimate run short; 0 restores the unbounded wait. +_DEFAULT_TIMEOUT_SECONDS = 1800.0 def _required_env(name: str) -> str: @@ -38,6 +42,18 @@ def _required_env(name: str) -> str: return value +def _timeout_seconds() -> float: + """Wall-clock cap on one child run; 0 or unparsable means wait forever.""" + raw = os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT") + if raw is None or not raw.strip(): + return _DEFAULT_TIMEOUT_SECONDS + try: + parsed = float(raw.strip()) + except ValueError: + return _DEFAULT_TIMEOUT_SECONDS + return parsed if parsed > 0 else 0.0 + + def _bounded(text: str) -> str: if len(text) <= _MAX_RESULT_CHARACTERS: return text @@ -153,6 +169,17 @@ def run_local_agent( "--output-format", "json", "--no-session-persistence", + # Strip human-blocking tools so the child runs unattended. Only the read-only + # child's writers bite today, since a --print child is never offered the plan + # or prompt tools; those are listed anyway so a version that starts offering + # them cannot stall the subagent. Bash is denied read-only side because plan + # mode gates it through the same local model, which is not a write barrier. + "--disallowedTools", + ( + "AskUserQuestion,EnterPlanMode,Edit,Write,NotebookEdit,Bash" + if read_only + else "AskUserQuestion,EnterPlanMode,ExitPlanMode" + ), "--append-system-prompt", _SUBAGENT_PLAN_INSTRUCTIONS if read_only else _SUBAGENT_INSTRUCTIONS, f"Task: {task}", @@ -185,6 +212,8 @@ def run_local_agent( [executable, *command[1:]], **popen_kwargs, ) + deadline = _timeout_seconds() + started_at = time.monotonic() try: while True: try: @@ -194,6 +223,13 @@ def run_local_agent( if cancel_event.is_set(): _stop_child(process) raise RuntimeError("The local Claude agent was cancelled.") + waited = time.monotonic() - started_at + if deadline and waited > deadline: + _stop_child(process) + raise RuntimeError( + f"The local Claude agent produced nothing after {waited:.0f}s. " + "The local server is likely wedged; check that a model is loaded." + ) except BaseException: if process.poll() is None: _stop_child(process) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 079c7850e5..85e303215b 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -4,6 +4,7 @@ """`unsloth start` — launch a coding agent against a running Unsloth server.""" import atexit +import base64 import contextlib import json import os @@ -2052,6 +2053,32 @@ def _opencode_subagent_inline_config(path: Path, permission: dict) -> dict: return inline +def _b64_path(path: Path) -> str: + """Path as base64, so it can cross a shell without being expanded.""" + return base64.b64encode(str(path).encode("utf-8")).decode("ascii") + + +_CLAUDE_PLAN_GATE_SCRIPT = '''\ +"""Deny the editing agent while the parent session is in plan mode.""" +import json, sys + +try: + mode = (json.load(sys.stdin) or {}).get("permission_mode") +except Exception: + sys.exit(0) # fail open: a hook error must never block the parent session +if mode == "plan": + print(json.dumps({"hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": ( + "Plan mode is active. Call the read-only Unsloth plan agent " + "(unsloth_plan_agent) instead of unsloth_agent." + ), + }})) +sys.exit(0) +''' + + def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path: """Write a session plugin that exposes the local Claude child through MCP.""" plugin = path / "unsloth-local-agent" @@ -2094,6 +2121,52 @@ def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path: } }, ) + # Claude already refuses the editing tool in plan mode, since it advertises + # readOnlyHint false. This PreToolUse hook replaces that dead end with a reason + # naming the read-only tool to call instead. Skipped under the WSL bridge, where + # the gate is a Linux path but the hook would run beside the Windows claude. + gate = plugin / "hooks" / "plan_gate.py" + if command == "wsl.exe": + # A persisted plugin dir may still hold a gate from an earlier non-WSL run. + for stale in (gate, plugin / "hooks" / "hooks.json"): + stale.unlink(missing_ok = True) + else: + _write_private_text(gate, _CLAUDE_PLAN_GATE_SCRIPT) + _write_private_json( + plugin / "hooks" / "hooks.json", + { + "hooks": { + "PreToolUse": [ + { + "matcher": _CLAUDE_SUBAGENT_TOOL, + "hooks": [ + { + "type": "command", + # Run through runpy rather than handing the path to + # the interpreter: a missing gate is then an + # ordinary traceback (exit 1, fails open) instead + # of exit 2, which Claude treats as a blocking + # error and would deny the tool in every mode. + # The path is base64'd because this string goes + # through a shell: a temp root holding $(..) or a + # backtick expands under sh, %VAR% under cmd, and + # the gate then silently fails open. base64's + # alphabet has no metacharacter in either. + "command": ( + f'"{sys.executable}" -c ' + f'"import base64,runpy; runpy.run_path(' + f"base64.b64decode('{_b64_path(gate)}').decode())\"" + ), + # A hook with no timeout stalls the parent for as + # long as it hangs; measured unbounded past 400s. + "timeout": 10, + } + ], + } + ] + } + }, + ) skill = plugin / "skills" / "local-agent" / "SKILL.md" skill.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) skill.write_text( diff --git a/unsloth_cli/tests/test_claude_plan_gate.py b/unsloth_cli/tests/test_claude_plan_gate.py new file mode 100644 index 0000000000..85336c8dfe --- /dev/null +++ b/unsloth_cli/tests/test_claude_plan_gate.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic plan-mode routing for the local Claude subagent. + +SKILL.md asks the parent model to pick the read-only tool in plan mode, which a +small local model can forget. The generated plugin also ships a PreToolUse hook +that reads permission_mode directly, so the editing agent is denied by rule. +""" + +from __future__ import annotations + +import json +import subprocess +import sys + +import pytest + +from unsloth_cli.commands import start + + +def _plugin(tmp_path): + return start.write_claude_subagent_plugin(tmp_path, {"UNSLOTH_CLAUDE_SUBAGENT_MODEL": "m"}) + + +def _run_gate(script, payload): + return subprocess.run( + [sys.executable, str(script)], + input = payload, + capture_output = True, + text = True, + timeout = 30, + ) + + +def test_plugin_registers_a_pretooluse_hook_on_the_editing_tool(tmp_path): + plugin = _plugin(tmp_path) + + hooks = json.loads((plugin / "hooks" / "hooks.json").read_text())["hooks"]["PreToolUse"] + + [entry] = hooks + # Only the destructive tool is gated; the read-only agent stays reachable. + assert entry["matcher"] == start._CLAUDE_SUBAGENT_TOOL + assert start._CLAUDE_SUBAGENT_PLAN_TOOL not in json.dumps(hooks) + [hook] = entry["hooks"] + assert hook["type"] == "command" + assert sys.executable in hook["command"] + # The interpreter is quoted: unquoted, any space in the path splits the command. + assert f'"{sys.executable}"' in hook["command"] + # The gate path rides as base64, never as a literal the shell can expand. + encoded = start._b64_path(plugin / "hooks" / "plan_gate.py") + assert encoded in hook["command"] + assert str(plugin / "hooks" / "plan_gate.py") not in hook["command"] + # A hook with no timeout stalls the parent for as long as it hangs. + assert 0 < hook["timeout"] <= 30 + + +def test_gate_script_is_written_and_compiles(tmp_path): + plugin = _plugin(tmp_path) + gate = plugin / "hooks" / "plan_gate.py" + + compile(gate.read_text(), str(gate), "exec") # syntax-valid as shipped + + +def test_gate_denies_the_editing_tool_in_plan_mode(tmp_path): + gate = _plugin(tmp_path) / "hooks" / "plan_gate.py" + + result = _run_gate(gate, json.dumps({"permission_mode": "plan"})) + + assert result.returncode == 0 + output = json.loads(result.stdout)["hookSpecificOutput"] + assert output["hookEventName"] == "PreToolUse" + assert output["permissionDecision"] == "deny" + # The reason is shown to the model, so it must name the tool to call instead. + assert "unsloth_plan_agent" in output["permissionDecisionReason"] + + +@pytest.mark.parametrize("mode", ["default", "acceptEdits", "bypassPermissions", "dontAsk", "auto"]) +def test_gate_allows_every_non_plan_mode(tmp_path, mode): + gate = _plugin(tmp_path) / "hooks" / "plan_gate.py" + + result = _run_gate(gate, json.dumps({"permission_mode": mode})) + + assert result.returncode == 0 + assert result.stdout.strip() == "" # no decision -> normal permission flow + + +@pytest.mark.parametrize("payload", ["", "not json", "[]", "null", "{}"]) +def test_gate_fails_open_on_unusable_input(tmp_path, payload): + # A hook crash would block the parent session, so anything unparsable allows. + gate = _plugin(tmp_path) / "hooks" / "plan_gate.py" + + result = _run_gate(gate, payload) + + assert result.returncode == 0 + assert result.stdout.strip() == "" + + +def test_plugin_still_writes_the_mcp_server_and_skill(tmp_path): + # The hook is additive; the existing wiring must be untouched. + plugin = _plugin(tmp_path) + + assert (plugin / ".mcp.json").exists() + assert (plugin / "skills" / "local-agent" / "SKILL.md").exists() + assert (plugin / ".claude-plugin" / "plugin.json").exists() + + +def test_wsl_run_clears_a_gate_left_by_an_earlier_windows_run(tmp_path, monkeypatch): + # The plugin dir survives across runs when persisted, so a gate written by a + # Windows run would otherwise be shipped into the distro with an interpreter + # path it cannot execute. + plugin = _plugin(tmp_path) + gate = plugin / "hooks" / "plan_gate.py" + hooks = plugin / "hooks" / "hooks.json" + assert gate.exists() and hooks.exists() + + monkeypatch.setattr(start, "_wsl_windows_executable", lambda _argv: True) + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + _plugin(tmp_path) + + assert not gate.exists() + assert not hooks.exists() + + +def test_hook_command_survives_a_missing_gate_and_a_path_with_spaces(tmp_path): + # Handing the path straight to the interpreter makes a missing gate exit 2, + # which Claude treats as a blocking error: the editing tool would then be + # denied in every mode, not just plan. Going through runpy makes it exit 1. + plugin = _plugin(tmp_path / "dir with space") + hook = json.loads((plugin / "hooks" / "hooks.json").read_text()) + command = hook["hooks"]["PreToolUse"][0]["hooks"][0]["command"] + + # Works normally through the real shell path Claude uses. + denied = subprocess.run( + command, + input = json.dumps({"permission_mode": "plan"}), + shell = True, + capture_output = True, + text = True, + timeout = 30, + ) + assert denied.returncode == 0 + assert json.loads(denied.stdout)["hookSpecificOutput"]["permissionDecision"] == "deny" + + (plugin / "hooks" / "plan_gate.py").unlink() + gone = subprocess.run( + command, + input = json.dumps({"permission_mode": "default"}), + shell = True, + capture_output = True, + text = True, + timeout = 30, + ) + assert gone.returncode != 2, "exit 2 blocks the tool in every mode" + assert gone.stdout.strip() == "" + + +@pytest.mark.parametrize("hostile", ["sub$(echo X)", "tick`echo X`", "var$HOME", "pct%TEMP%pct"]) +def test_gate_survives_shell_metacharacters_in_its_path(tmp_path, hostile): + # The hook command is run by a shell. A temp root holding these expands under + # sh (or cmd, for %VAR%) before Python sees the path, so the gate is not found + # and exits 1, which fails open and silently drops the routing message. + plugin = _plugin(tmp_path / hostile) + command = json.loads((plugin / "hooks" / "hooks.json").read_text())["hooks"]["PreToolUse"][0][ + "hooks" + ][0]["command"] + + denied = subprocess.run( + command, + input = json.dumps({"permission_mode": "plan"}), + shell = True, + capture_output = True, + text = True, + timeout = 60, + ) + + assert denied.returncode == 0, denied.stderr + decision = json.loads(denied.stdout)["hookSpecificOutput"]["permissionDecision"] + assert decision == "deny" diff --git a/unsloth_cli/tests/test_claude_subagent_mcp.py b/unsloth_cli/tests/test_claude_subagent_mcp.py index 568dc76ff5..2155ca050d 100644 --- a/unsloth_cli/tests/test_claude_subagent_mcp.py +++ b/unsloth_cli/tests/test_claude_subagent_mcp.py @@ -15,6 +15,16 @@ import pytest import unsloth_cli.claude_subagent_mcp as bridge +def _stub_env(monkeypatch, tmp_path): + """Minimum env + claude lookup for driving run_local_agent under a fake Popen.""" + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M") + monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path)) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(bridge, "_claude_flags", lambda model: ["--settings", "{}"]) + + def test_protocol_lists_and_calls_local_agent(): initialized = bridge._response( {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, @@ -229,6 +239,8 @@ def test_local_child_uses_unsloth_without_overwriting_parent_auth( assert command[:3] == ["/usr/local/bin/claude", "--model", "unsloth/model-GGUF:Q4_K_M"] assert command[command.index("--permission-mode") + 1] == permission assert "--no-session-persistence" in command + disallowed = command[command.index("--disallowedTools") + 1] + assert disallowed == "AskUserQuestion,EnterPlanMode,ExitPlanMode" assert captured["cwd"] == str(tmp_path) assert captured["stdin"] is bridge.subprocess.DEVNULL assert captured["stdout"] is bridge.subprocess.PIPE @@ -275,6 +287,10 @@ def test_read_only_local_child_uses_plan_mode(monkeypatch, tmp_path): assert bridge.run_local_agent("plan this", read_only = True) == "PLAN_OK" command = captured["command"] assert command[command.index("--permission-mode") + 1] == "plan" + disallowed = command[command.index("--disallowedTools") + 1] + assert disallowed == "AskUserQuestion,EnterPlanMode,Edit,Write,NotebookEdit,Bash" + # Bash matters: plan mode routes it through a classifier served by this same + # local model, so without the deny a "read-only" child can still write files. prompt = command[command.index("--append-system-prompt") + 1] assert "read-only local coding subagent" in prompt @@ -403,3 +419,37 @@ def test_stop_child_kills_survivors_after_leader_exit(monkeypatch, tmp_path): def test_result_parser_accepts_diagnostics_before_json(): output = "connector warning\n" + json.dumps({"is_error": False, "result": "OK"}) assert bridge._result_text(output) == "OK" + + +def test_child_is_stopped_when_it_produces_nothing_before_the_deadline(monkeypatch, tmp_path): + # A local server that accepts and never answers used to block the child, and + # the parent waiting on it, indefinitely. Measured past 400s before this. + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", "0.3") + _stub_env(monkeypatch, tmp_path) + stopped = [] + + class _Hanging: + returncode = None + + def communicate(self, timeout = None): + raise subprocess.TimeoutExpired("claude", timeout) + + def poll(self): + return None + + monkeypatch.setattr(bridge, "_stop_child", lambda proc: stopped.append(proc)) + monkeypatch.setattr(bridge.subprocess, "Popen", lambda *a, **k: _Hanging()) + + with pytest.raises(RuntimeError, match = "produced nothing"): + bridge.run_local_agent("hello") + assert stopped, "a timed-out child must be killed, not left running" + + +def test_timeout_can_be_disabled(monkeypatch): + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", "0") + assert bridge._timeout_seconds() == 0.0 + for bad in ("", " ", "abc"): + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", bad) + assert bridge._timeout_seconds() == bridge._DEFAULT_TIMEOUT_SECONDS + monkeypatch.delenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT") + assert bridge._timeout_seconds() == bridge._DEFAULT_TIMEOUT_SECONDS