unsloth/unsloth_cli/tests/test_claude_subagent_mcp.py
Daniel Han 7a9749eb4f
unsloth start: keep the local subagent unattended and out of plan mode (#7437)
* unsloth start: keep the local subagent unattended and out of plan mode

The local subagent child could stall waiting on a permission prompt, and a
parent session in plan mode could still reach the editing agent.

- Drop human-blocking tools from the child so it runs unattended. The
  read-only child also drops the file writers.
- Emit a PreToolUse hook that reads permission_mode itself and denies the
  editing agent under plan mode, so routing holds when the model ignores
  SKILL.md. Fails open, and is skipped under the WSL bridge where a Windows
  interpreter path is not runnable in the distro.

* Make the read-only subagent actually read-only, and drop stale WSL gates

From the first review of this branch, which drove the real code against a fake
HOME holding a pre-existing Claude install and diffed the tree before and after.
No config, agent, MCP server or CLAUDE.md of the user's was touched in either
arm, and the session dir is removed on exit, Ctrl-C and exception.

Three real findings came out of it:

- The read-only child could still write. Plan mode routes Bash through a safety
  classifier served by the same local model, so a small model saying yes is what
  authorised the write; a child spawned with read_only created a file. Denying
  Bash there makes the label true, at the cost of shell exploration while
  planning. Read, Grep and Glob still cover the search it needs.
- A persisted plugin dir kept a plan_gate.py from an earlier Windows run, so a
  later WSL run shipped a hooks.json naming an interpreter the distro cannot
  execute. Hook errors do not block, so this only ever wasted a spawn, but it
  accumulated and the branch had no test.
- The comment claimed the read-only child keeps ExitPlanMode "as Claude does
  under plan mode". A --print child is never offered the plan or prompt tools at
  all, so most of both deny lists is inert today. Kept as a guard against a
  version that starts offering them, but the comment now says so.

Also covers "auto" in the gate's non-plan modes, which is a real permission_mode
and the one the child's own Bash classifier runs under.

* Stop the gate failing closed, and bound a wedged child

Second review of this branch, driving real claude 2.1.219 against a mock
endpoint rather than reading.

The gate could fail closed. If plan_gate.py went missing the interpreter exited
2, which Claude treats as a blocking hook error, so the editing tool was denied
in every mode rather than just plan. Running the script through runpy instead of
handing its path to the interpreter turns that into an ordinary traceback, which
is exit 1 and allows. Verified both exit codes directly.

The hook also had no timeout, so a hung one stalled the parent for as long as it
hung, measured past 400s. Bounded at 10s.

The real stall this branch is named for was untouched: run_local_agent polled
communicate() forever, so a local server that accepts and never answers left the
child and the parent blocked indefinitely, measured past 400s. Added a wall-clock
deadline that kills the child and says the server looks wedged.
UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT overrides it, 0 restores the old behaviour.

Also corrected the plan-mode comment. Claude already refuses the editing tool in
plan mode on its own, since it advertises readOnlyHint false; what the hook adds
is a reason naming the read-only tool to call instead. The WSL comment had the
direction backwards: the gate is the Linux path, not the Windows one.

Tests: the hook command's quoting and its behaviour with the gate deleted, both
previously unguarded, plus the timeout and its env override.

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

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

* Keep the gate path out of the shell string

Codex review. The hook command is run by a shell, and the gate path was
interpolated into it, so a session-config root containing shell metacharacters
expanded before Python saw it. Verified on both: sh expands $(..), backticks and
$VAR; cmd expands %VAR%. In every case the path no longer resolves, the gate
exits 1, and because that intentionally fails open the routing message silently
stops appearing.

The path now travels as base64, whose alphabet has no metacharacter in either
shell. Parametrised over all four hostile forms, and the old interpolation makes
those tests fail.

One correction to the report: it says the editing agent becomes callable in plan
mode. It does not. Claude refuses that tool by itself, since it advertises
readOnlyHint false, which was checked earlier by deleting the hook entirely.
What a mangled path costs is the reason naming the read-only agent to call
instead, not the block.

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

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

---------

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>
2026-07-27 04:18:22 -07:00

455 lines
16 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import io
import json
import os
import subprocess
import sys
import time
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": {}},
)
assert initialized["result"]["serverInfo"]["name"] == "unsloth-local-agent"
listed = bridge._response({"jsonrpc": "2.0", "id": 2, "method": "tools/list"})
tool = listed["result"]["tools"][0]
assert tool["name"] == "unsloth_agent"
assert "spawn an Unsloth or local agent" in tool["description"]
assert tool["inputSchema"]["required"] == ["task"]
assert tool["_meta"]["anthropic/maxResultSizeChars"] == 100_000
called = bridge._response(
{
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {"name": "unsloth_agent", "arguments": {"task": " inspect this "}},
},
run_agent = lambda task: f"completed: {task}",
)
assert called["result"] == {
"content": [{"type": "text", "text": "completed: inspect this"}],
"isError": False,
}
def test_protocol_exposes_read_only_agent_for_claude_plan_mode():
listed = bridge._response(
{"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
run_read_only_agent = lambda task: task,
read_only_tool_name = "unsloth_plan_agent",
)
tools = {tool["name"]: tool for tool in listed["result"]["tools"]}
assert tools["unsloth_agent"]["annotations"]["readOnlyHint"] is False
assert tools["unsloth_plan_agent"]["annotations"] == {
"readOnlyHint": True,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": True,
}
called = bridge._response(
{
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "unsloth_plan_agent",
"arguments": {"task": " inspect this "},
},
},
run_agent = lambda task: f"write: {task}",
run_read_only_agent = lambda task: f"plan: {task}",
read_only_tool_name = "unsloth_plan_agent",
)
assert called["result"] == {
"content": [{"type": "text", "text": "plan: inspect this"}],
"isError": False,
}
def test_protocol_returns_tool_errors_to_parent():
response = bridge._response(
{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {"name": "unsloth_agent", "arguments": {"task": "test"}},
},
run_agent = lambda task: (_ for _ in ()).throw(RuntimeError("local failure")),
)
assert response["result"]["isError"] is True
assert response["result"]["content"][0]["text"] == "local failure"
def test_stdio_server_ignores_notifications_and_answers_requests():
requests = "\n".join(
[
json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}),
json.dumps({"jsonrpc": "2.0", "id": 4, "method": "ping"}),
]
)
output = io.StringIO()
bridge.serve(io.StringIO(requests), output)
assert json.loads(output.getvalue()) == {"jsonrpc": "2.0", "id": 4, "result": {}}
def test_stdio_cancellation_reaches_the_running_local_agent():
requests = "\n".join(
[
json.dumps(
{
"jsonrpc": "2.0",
"id": "call-1",
"method": "tools/call",
"params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
}
),
json.dumps(
{
"jsonrpc": "2.0",
"method": "notifications/cancelled",
"params": {"requestId": "call-1", "reason": "user cancelled"},
}
),
]
)
output = io.StringIO()
cancelled = []
def run_agent(task, cancel_event):
assert task == "wait"
assert cancel_event.wait(timeout = 1)
cancelled.append(task)
raise RuntimeError("The local Claude agent was cancelled.")
bridge.serve(io.StringIO(requests), output, run_agent = run_agent)
assert cancelled == ["wait"]
assert output.getvalue() == ""
def test_stdio_sigint_stops_the_running_local_agent(monkeypatch):
request = json.dumps(
{
"jsonrpc": "2.0",
"id": "call-1",
"method": "tools/call",
"params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
}
)
handlers = {}
started = bridge.threading.Event()
cancelled = []
def set_handler(signum, handler):
previous = handlers.get(signum, bridge.signal.SIG_DFL)
handlers[signum] = handler
return previous
monkeypatch.setattr(bridge.signal, "signal", set_handler)
class InterruptingInput:
def __init__(self):
self.sent = False
def __iter__(self):
return self
def __next__(self):
if not self.sent:
self.sent = True
return request + "\n"
assert started.wait(timeout = 1)
handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
raise AssertionError("SIGINT handler must unwind the stdin loop")
def run_agent(task, cancel_event):
assert task == "wait"
started.set()
assert cancel_event.wait(timeout = 1)
# Real Claude Code sends SIGINT twice. The second one must not abort cleanup.
handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
cancelled.append(task)
raise RuntimeError("The local Claude agent was cancelled.")
output = io.StringIO()
bridge.serve(InterruptingInput(), output, run_agent = run_agent)
assert cancelled == ["wait"]
assert output.getvalue() == ""
@pytest.mark.parametrize(
("bypass", "permission"),
[("0", "acceptEdits"), ("1", "bypassPermissions")],
)
def test_local_child_uses_unsloth_without_overwriting_parent_auth(
monkeypatch, tmp_path, bypass, permission
):
captured = {}
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("UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW", "32768")
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", bypass)
monkeypatch.setenv("ANTHROPIC_API_KEY", "cloud-key")
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "cloud-oauth")
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", "{}"])
class Process:
pid = 1234
returncode = 0
def communicate(self, timeout):
captured["timeout"] = timeout
return json.dumps({"is_error": False, "result": "LOCAL_OK"}), ""
def poll(self):
return self.returncode
def popen(command, **kwargs):
captured["command"] = command
captured.update(kwargs)
return Process()
monkeypatch.setattr(bridge.subprocess, "Popen", popen)
assert bridge.run_local_agent("reply exactly LOCAL_OK") == "LOCAL_OK"
command = captured["command"]
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
assert captured["stderr"] is bridge.subprocess.PIPE
if os.name == "nt":
assert captured["creationflags"] == bridge.subprocess.CREATE_NEW_PROCESS_GROUP
else:
assert captured["start_new_session"] is True
child_env = captured["env"]
assert child_env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8888"
assert child_env["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-test"
assert child_env["ANTHROPIC_MODEL"] == "unsloth/model-GGUF:Q4_K_M"
assert child_env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] == "32768"
assert child_env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] == "90"
assert "ANTHROPIC_API_KEY" not in child_env
assert "CLAUDE_CODE_OAUTH_TOKEN" not in child_env
def test_read_only_local_child_uses_plan_mode(monkeypatch, tmp_path):
captured = {}
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("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", "1")
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: [])
class Process:
pid = 1234
returncode = 0
def communicate(self, timeout):
return json.dumps({"is_error": False, "result": "PLAN_OK"}), ""
def poll(self):
return self.returncode
def popen(command, **kwargs):
captured["command"] = command
return Process()
monkeypatch.setattr(bridge.subprocess, "Popen", popen)
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
def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path):
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: [])
cancel_event = bridge.threading.Event()
stopped = []
class Process:
pid = 1234
returncode = None
def communicate(self, timeout):
cancel_event.set()
raise bridge.subprocess.TimeoutExpired("claude", timeout)
def poll(self):
return self.returncode
process = Process()
monkeypatch.setattr(bridge.subprocess, "Popen", lambda *args, **kwargs: process)
def stop(child):
stopped.append(child)
child.returncode = -15
monkeypatch.setattr(bridge, "_stop_child", stop)
with pytest.raises(RuntimeError, match = "cancelled"):
bridge.run_local_agent("wait", cancel_event)
assert stopped == [process]
def test_windows_cancellation_stops_the_child_process_tree(monkeypatch):
monkeypatch.setattr(bridge.os, "name", "nt")
captured = {}
class Process:
pid = 4321
returncode = None
def poll(self):
return self.returncode
def wait(self, timeout = None):
captured["wait_timeout"] = timeout
self.returncode = 1
def terminate(self):
raise AssertionError("taskkill should handle the process tree")
def run(command, **kwargs):
captured["command"] = command
captured.update(kwargs)
return bridge.subprocess.CompletedProcess(command, 0)
monkeypatch.setattr(bridge.subprocess, "run", run)
bridge._stop_child(Process())
assert captured["command"] == ["taskkill", "/PID", "4321", "/T", "/F"]
assert captured["capture_output"] is True
assert captured["check"] is False
assert captured["wait_timeout"] == bridge._CANCEL_GRACE_SECONDS
def test_windows_failed_taskkill_still_terminates_the_child(monkeypatch):
monkeypatch.setattr(bridge.os, "name", "nt")
captured = {}
class Process:
pid = 4321
returncode = None
def poll(self):
return self.returncode
def wait(self, timeout = None):
self.returncode = 1
def terminate(self):
captured["terminated"] = True
self.returncode = 1
monkeypatch.setattr(
bridge.subprocess,
"run",
lambda command, **kwargs: bridge.subprocess.CompletedProcess(command, 1),
)
bridge._stop_child(Process())
assert captured.get("terminated") is True
@pytest.mark.skipif(os.name == "nt", reason = "POSIX process groups")
def test_stop_child_kills_survivors_after_leader_exit(monkeypatch, tmp_path):
monkeypatch.setattr(bridge, "_CANCEL_GRACE_SECONDS", 0.2)
marker = tmp_path / "grandchild-survived"
grandchild = (
"import pathlib, sys, time; time.sleep(1.0); "
"pathlib.Path(sys.argv[1]).write_text('alive')"
)
process = subprocess.Popen(
[
sys.executable,
"-c",
"import subprocess, sys; "
"subprocess.Popen([sys.executable, '-c', sys.argv[1], sys.argv[2]])",
grandchild,
str(marker),
],
start_new_session = True,
)
process.wait()
bridge._stop_child(process)
time.sleep(1.2)
assert not marker.exists()
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