diff --git a/pyproject.toml b/pyproject.toml index 87c8dcc..54f3e78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,12 @@ where = ["src"] [tool.setuptools.package-data] robocode = ["py.typed"] +[tool.pytest.ini_options] +addopts = "-m 'not integration'" +markers = [ + "integration: opt-in live tests requiring external services and potentially paid model calls", +] + [tool.black] line-length = 88 target-version = ["py311"] diff --git a/src/robocode/utils/env_server.py b/src/robocode/utils/env_server.py index 10dd7e7..c61abf7 100644 --- a/src/robocode/utils/env_server.py +++ b/src/robocode/utils/env_server.py @@ -36,6 +36,7 @@ from __future__ import annotations +import _thread import json import logging import os @@ -50,7 +51,6 @@ from pathlib import Path from typing import Any -import _thread import numpy as np from gymnasium.spaces import Box, Space from relational_structs.spaces import ObjectCentricStateSpace diff --git a/src/robocode/utils/llm/cli_client.py b/src/robocode/utils/llm/cli_client.py index 1926712..58dc447 100644 --- a/src/robocode/utils/llm/cli_client.py +++ b/src/robocode/utils/llm/cli_client.py @@ -45,6 +45,9 @@ def complete(self, messages: list[dict[str, str]]) -> LLMResponse: self._model, "--tools", "", + # --tools only controls built-ins; deny MCP tools as well. + "--disallowedTools", + "*", "--system-prompt", "", "--exclude-dynamic-system-prompt-sections", diff --git a/tests/utils/fixtures/claude_mcp_probe.py b/tests/utils/fixtures/claude_mcp_probe.py new file mode 100644 index 0000000..7580377 --- /dev/null +++ b/tests/utils/fixtures/claude_mcp_probe.py @@ -0,0 +1,71 @@ +"""Run the real Claude completion wrapper inside the GenPlan Docker sandbox. + +Both tests use this program. --without-deny-flag reproduces the old wrapper; +otherwise all production restrictions remain. Each call has a $1 CLI budget. +""" + +import argparse +import json +import subprocess +from pathlib import Path +from unittest.mock import patch + +from omegaconf import DictConfig + +from robocode.utils.llm.cli_client import ClaudeCLIClient + + +def main() -> None: + """Call Claude with the test MCP server and save the full CLI transcript.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--without-deny-flag", action="store_true") + options = parser.parse_args() + real_run = subprocess.run + + def record_cli(args, **kwargs): + args = list(args) + if options.without_deny_flag: + # The control differs only by removing the production deny-all flag. + index = args.index("--disallowedTools") + del args[index : index + 2] + args[args.index("--output-format") + 1] = "stream-json" + args.extend( + [ + "--verbose", + "--mcp-config", + "/sandbox/mcp.json", + "--allowedTools", + "mcp__probe__probe_token", + "--max-budget-usd", + "1.0", + ] + ) + result = real_run(args, check=kwargs.pop("check", True), **kwargs) + Path("/sandbox/transcript.jsonl").write_text(result.stdout, encoding="utf-8") + events = [ + json.loads(line) for line in result.stdout.splitlines() if line.strip() + ] + final = next(e for e in reversed(events) if e.get("type") == "result") + assert not final.get("is_error"), final + result.stdout = json.dumps(final) + return result + + client = ClaudeCLIClient( + DictConfig({"model": "claude-opus-5", "request_timeout_s": 120}) + ) + with patch("robocode.utils.llm.cli_client.subprocess.run", record_cli): + client.complete( + [ + { + "role": "user", + "content": ( + "Call mcp__probe__probe_token and return its token exactly. " + "Do not guess. If unavailable, respond UNAVAILABLE." + ), + } + ] + ) + + +if __name__ == "__main__": + main() diff --git a/tests/utils/fixtures/claude_mcp_server.py b/tests/utils/fixtures/claude_mcp_server.py new file mode 100644 index 0000000..663bca2 --- /dev/null +++ b/tests/utils/fixtures/claude_mcp_server.py @@ -0,0 +1,20 @@ +"""Harmless MCP tool for the Docker Claude isolation tests.""" + +import sys +from pathlib import Path + +from mcp.server.fastmcp import FastMCP + +mcp = FastMCP("probe") + + +@mcp.tool() +def probe_token() -> str: + """Return the secret test token.""" + with Path("/sandbox/calls.txt").open("a", encoding="utf-8") as log: + log.write("called\n") + return sys.argv[1] + + +if __name__ == "__main__": + mcp.run(transport="stdio") diff --git a/tests/utils/test_episode.py b/tests/utils/test_episode.py index b0ea8fc..58a3a91 100644 --- a/tests/utils/test_episode.py +++ b/tests/utils/test_episode.py @@ -4,6 +4,7 @@ from __future__ import annotations +import _thread import math import multiprocessing as mp import signal @@ -14,7 +15,6 @@ from pathlib import Path from typing import Any, Callable -import _thread import imageio.v3 as iio import numpy as np import pytest diff --git a/tests/utils/test_llm.py b/tests/utils/test_llm.py index 85a7e4e..d34e8b9 100644 --- a/tests/utils/test_llm.py +++ b/tests/utils/test_llm.py @@ -1,8 +1,21 @@ -"""Tests for the shared LLM client helpers (cost estimation).""" +"""LLM helpers and opt-in Docker Claude tool-isolation tests. +Run with pytest tests/utils/test_llm.py -m integration. These tests make paid +model calls using existing Claude authentication and a harmless MCP server. +""" + +import json +import shutil +import subprocess +import uuid +from pathlib import Path + +import pytest from omegaconf import DictConfig +from robocode.utils.docker_sandbox import DOCKER_PYTHON, run_genplan_in_docker from robocode.utils.llm.base import pricing_from_cfg, usage_cost +from robocode.utils.llm.cli_client import ClaudeCLIClient def test_usage_cost_from_list_prices(): @@ -22,3 +35,137 @@ def test_pricing_from_cfg(): assert pricing_from_cfg(DictConfig({"output_cost_per_mtok": 25.0})) == (0.0, 25.0) # Neither set -> unpriced (cost stays unknown, e.g. local vLLM/Ollama). assert pricing_from_cfg(DictConfig({"model": "x"})) is None + + +def test_cli_denies_all_tools(monkeypatch): + """The plain completion client blocks built-in and MCP tools explicitly.""" + + def fake_run(args, **kwargs): + assert args[args.index("--tools") + 1] == "" + assert args[args.index("--disallowedTools") + 1] == "*" + assert kwargs["input"] == "USER: hello" + return subprocess.CompletedProcess( + args, 0, json.dumps({"result": "world", "total_cost_usd": 0.01}), "" + ) + + monkeypatch.setattr("robocode.utils.llm.cli_client.subprocess.run", fake_run) + reply = ClaudeCLIClient(DictConfig({"model": "test-model"})).complete( + [{"role": "user", "content": "hello"}] + ) + assert reply.text == "world" + assert reply.cost_usd == 0.01 + + +def _create_mcp_tool(sandbox_dir: Path) -> str: + """Write the Docker MCP server/config and return its secret test token.""" + # A harmless tool returns a token the model cannot know without invoking it. + token = uuid.uuid4().hex + shutil.copyfile( + Path(__file__).parent / "fixtures" / "claude_mcp_server.py", + sandbox_dir / "server.py", + ) + (sandbox_dir / "mcp.json").write_text( + json.dumps( + { + "mcpServers": { + "probe": { + "command": DOCKER_PYTHON, + "args": ["/sandbox/server.py", token], + } + } + } + ) + ) + + return token + + +def _parse_claude_output(transcript_path: Path) -> tuple[dict, dict]: + """Return the CLI initialization and final result from a JSONL transcript.""" + events = [json.loads(line) for line in transcript_path.read_text().splitlines()] + init = next( + e for e in events if e.get("type") == "system" and e.get("subtype") == "init" + ) + final = next(e for e in reversed(events) if e.get("type") == "result") + return init, final + + +@pytest.mark.integration +def test_docker_claude_calls_mcp_without_deny_flag(tmp_path, monkeypatch): + """Without deny-all, sandboxed Claude invokes the configured MCP tool.""" + if shutil.which("docker") is None: + pytest.skip("Docker is not installed") + + token = _create_mcp_tool(tmp_path) + + # Execute the real completion wrapper inside the image, recording CLI events. + shutil.copyfile( + Path(__file__).parent / "fixtures" / "claude_mcp_probe.py", + tmp_path / "probe.py", + ) + + # Keep the production auth, mounts, entrypoint, firewall and privilege drop. + # Replace only the training driver with the diagnostic script. + tmp_path.chmod(0o777) # The container's unprivileged node user writes artifacts. + real_run = subprocess.run + + def launch_probe(args, **kwargs): + if args[:2] == ["docker", "run"]: + assert args[-3:] == [ + DOCKER_PYTHON, + "-m", + "robocode.approaches.genplan_driver", + ] + args = [ + *args[:-3], + DOCKER_PYTHON, + "/sandbox/probe.py", + "--without-deny-flag", + ] + return real_run(args, check=kwargs.pop("check", True), **kwargs) + + monkeypatch.setattr("robocode.utils.docker_sandbox.subprocess.run", launch_probe) + run_genplan_in_docker(tmp_path, {"provider": "cli"}, timeout=600) + + init, final = _parse_claude_output(tmp_path / "transcript.jsonl") + assert "mcp__probe__probe_token" in init["tools"] + assert (tmp_path / "calls.txt").read_text().splitlines() + assert token in final["result"] + + +@pytest.mark.integration +def test_docker_claude_blocks_mcp_with_deny_flag(tmp_path, monkeypatch): + """With deny-all, sandboxed Claude cannot invoke the configured MCP tool.""" + if shutil.which("docker") is None: + pytest.skip("Docker is not installed") + + token = _create_mcp_tool(tmp_path) + + # Execute the real completion wrapper inside the image, recording CLI events. + shutil.copyfile( + Path(__file__).parent / "fixtures" / "claude_mcp_probe.py", + tmp_path / "probe.py", + ) + + # Keep the production auth, mounts, entrypoint, firewall and privilege drop. + # Replace only the training driver with the diagnostic script. + tmp_path.chmod(0o777) # The container's unprivileged node user writes artifacts. + real_run = subprocess.run + + def launch_probe(args, **kwargs): + if args[:2] == ["docker", "run"]: + assert args[-3:] == [ + DOCKER_PYTHON, + "-m", + "robocode.approaches.genplan_driver", + ] + args = [*args[:-3], DOCKER_PYTHON, "/sandbox/probe.py"] + return real_run(args, check=kwargs.pop("check", True), **kwargs) + + monkeypatch.setattr("robocode.utils.docker_sandbox.subprocess.run", launch_probe) + run_genplan_in_docker(tmp_path, {"provider": "cli"}, timeout=600) + + init, final = _parse_claude_output(tmp_path / "transcript.jsonl") + assert init["tools"] == [] + assert not (tmp_path / "calls.txt").exists() + assert token not in final["result"]