Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
4.8 kB · 94 lines
Python
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495from __future__ import annotations
import osimport subprocessimport unittest
from rlm.subprocess_guard import install_subprocess_guards
class SubprocessGuardTests(unittest.TestCase): @classmethod def setUpClass(cls): install_subprocess_guards()
def test_cell_caller_rejects_search_commands(self): # Simulate call from user cell by executing in a frame named __main__ code = compile("import subprocess; subprocess.run(['rg', 'foo'])", "<cell-test-1>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("use search('pattern', path='...')", str(ctx.exception))
def test_cell_caller_rejects_vcs_commands(self): code = compile("import subprocess; subprocess.run(['git', 'status'])", "<cell-test-2>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("git.status(...)", str(ctx.exception))
code_grep = compile("import subprocess; subprocess.run(['git', 'grep', 'foo'])", "<cell-test-2b>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code_grep, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("git.grep(...)", str(ctx.exception))
def test_cell_caller_rejects_bash_lc_with_rg(self): code = compile("import subprocess; subprocess.run(['bash', '-lc', 'rg -l line-height'])", "<cell-test-3>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("bash() rejects `rg`", str(ctx.exception))
def test_cell_caller_rejects_find_commands(self): code = compile("import subprocess; subprocess.run(['find', '.'])", "<cell-test-find>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("search.files", str(ctx.exception))
def test_cell_caller_rejects_ls_commands(self): code = compile("import subprocess; subprocess.run(['ls', '-la'])", "<cell-test-ls>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("search.ls", str(ctx.exception))
def test_cell_caller_rejects_general_shell_commands(self): code = compile("import subprocess; subprocess.run(['whoami'])", "<cell-test-4>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("use await bash(...) instead", str(ctx.exception))
def test_popen_remains_subscriptable_class(self): # mcp.os.win32.utilities evaluates `subprocess.Popen[bytes]` in a # signature at import time; the guard must keep Popen class-like. self.assertTrue(issubclass(subprocess.Popen, subprocess.Popen)) subprocess.Popen[bytes] import mcp.os.win32.utilities # noqa: F401
def test_cell_caller_rejects_nix_commands(self): code_shell = compile("import subprocess; subprocess.run(['nix-shell', '-p', 'cargo'])", "<cell-test-nix-1>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code_shell, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("use `await nix.develop(...)`", str(ctx.exception))
code_dev = compile("import subprocess; subprocess.run(['nix', 'develop'])", "<cell-test-nix-2>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code_dev, {"subprocess": subprocess, "__name__": "__main__"}) self.assertIn("nix.develop", str(ctx.exception))
def test_cell_caller_rejects_private_helper_run(self): # Even if a helper in an exempt module calls subprocess, if it is a private helper # called directly from a cell, it must not bypass guards def _fake_skill_run(cmd): return subprocess.run(cmd)
_fake_skill_run.__module__ = "jj" code = compile("_fake_skill_run(['git', 'status'])", "<cell-test-priv>", "exec") with self.assertRaises(RuntimeError) as ctx: eval(code, {"_fake_skill_run": _fake_skill_run, "__name__": "__main__"}) self.assertIn("git.status(...)", str(ctx.exception))
def test_exempt_caller_allows_subprocess(self): # Calls from unittest (an exempt module) must run normally res = subprocess.run(["echo", "hello"], capture_output=True, text=True) self.assertEqual(res.stdout.strip(), "hello")
if __name__ == "__main__": unittest.main()