返回 CodeWhale
test_harness.py
1 from __future__ import annotations
2
3 import importlib
4 import json
5 import sys
6 import types
7 import unittest
8 from dataclasses import dataclass
9 from pathlib import Path
10 from types import SimpleNamespace
11
12
13 def install_verifiers_stubs() -> None:
14 class HarnessConfig:
15 env: dict[str, str] = {}
16 disabled_tools: list[str] | None = None
17
18 def __init__(self, **values):
19 self.env = dict(values.pop("env", {}))
20 self.disabled_tools = values.pop("disabled_tools", None)
21 for cls in reversed(type(self).mro()):
22 for name in getattr(cls, "__annotations__", {}):
23 if name not in self.__dict__ and hasattr(cls, name):
24 setattr(self, name, getattr(cls, name))
25 for name, value in values.items():
26 setattr(self, name, value)
27
28 @property
29 def resolved_env(self):
30 return dict(self.env)
31
32 class Harness:
33 @classmethod
34 def __class_getitem__(cls, _item):
35 return cls
36
37 def __init__(self, config):
38 self.config = config
39
40 def resolve_prompt(self, task):
41 return task.system_prompt, task.prompt
42
43 @dataclass(frozen=True)
44 class ProgramResult:
45 exit_code: int
46 stdout: str
47 stderr: str
48
49 packages = {
50 "verifiers": types.ModuleType("verifiers"),
51 "verifiers.v1": types.ModuleType("verifiers.v1"),
52 "verifiers.v1.clients": types.ModuleType("verifiers.v1.clients"),
53 "verifiers.v1.harness": types.ModuleType("verifiers.v1.harness"),
54 "verifiers.v1.runtimes": types.ModuleType("verifiers.v1.runtimes"),
55 "verifiers.v1.trace": types.ModuleType("verifiers.v1.trace"),
56 }
57 packages["verifiers.v1.clients"].ModelContext = type("ModelContext", (), {})
58 packages["verifiers.v1.harness"].Harness = Harness
59 packages["verifiers.v1.harness"].HarnessConfig = HarnessConfig
60 packages["verifiers.v1.runtimes"].ProgramResult = ProgramResult
61 packages["verifiers.v1.runtimes"].Runtime = type("Runtime", (), {})
62 packages["verifiers.v1.trace"].Trace = type("Trace", (), {})
63 sys.modules.update(packages)
64
65
66 ROOT = Path(__file__).resolve().parents[1]
67 sys.path.insert(0, str(ROOT))
68 install_verifiers_stubs()
69 harness_module = importlib.import_module("codewhale_harness.harness")
70
71 CodewhaleHarness = harness_module.CodewhaleHarness
72 CodewhaleHarnessConfig = harness_module.CodewhaleHarnessConfig
73 ProgramResult = sys.modules["verifiers.v1.runtimes"].ProgramResult
74
75
76 def event(event_type: str, **fields) -> str:
77 return json.dumps(
78 {
79 "schema": "codewhale.exec-stream",
80 "schema_version": 1,
81 "type": event_type,
82 **fields,
83 }
84 )
85
86
87 def successful_stream(model: str = "test/model", sandbox: str = "workspace-write") -> str:
88 sha = "sha256:" + "a" * 64
89 return "\n".join(
90 [
91 event("content", content="sensitive model output"),
92 event(
93 "metadata",
94 meta={
95 "receipt_kind": "terminal",
96 "provider": "openai",
97 "model": model,
98 "approval_posture": "auto_tools",
99 "sandbox_posture": sandbox,
100 "binary_sha256": sha,
101 "prompt_sha256": sha,
102 "status": "completed",
103 "termination_reason": "resolved",
104 "workspace": "/private/runtime/path",
105 "session_id": "private-session-id",
106 },
107 ),
108 event("done"),
109 ]
110 )
111
112
113 class FakeRuntime:
114 def __init__(self, runtime_type="subprocess", result=None):
115 self.type = runtime_type
116 self.result = result or ProgramResult(0, successful_stream(), "")
117 self.runs = []
118 self.programs = []
119 self.writes = []
120
121 async def run(self, argv, env):
122 self.runs.append((argv, env))
123 if argv[-1:] == ["--version"]:
124 return ProgramResult(0, "codewhale 0.9.1 (candidate)", "")
125 return ProgramResult(0, "", "")
126
127 async def run_program(self, argv, env):
128 self.programs.append((argv, env))
129 return self.result
130
131 async def write(self, path, data):
132 self.writes.append((path, data))
133
134
135 def trace(prompt="Do the task", system_prompt="Use evidence"):
136 return SimpleNamespace(
137 id="0123456789abcdef",
138 task=SimpleNamespace(
139 data=SimpleNamespace(prompt=prompt, system_prompt=system_prompt)
140 ),
141 info={},
142 )
143
144
145 class HarnessTests(unittest.IsolatedAsyncioTestCase):
146 async def test_preinstalled_candidate_requires_exact_version(self):
147 config = CodewhaleHarnessConfig(
148 version="0.9.1", binary_path="/candidate/codewhale"
149 )
150 runtime = FakeRuntime()
151
152 await CodewhaleHarness(config).setup(runtime)
153
154 self.assertEqual(runtime.runs, [(["/candidate/codewhale", "--version"], {})])
155
156 async def test_launch_isolates_config_routes_interception_and_keeps_safe_receipt(self):
157 config = CodewhaleHarnessConfig(
158 version="0.9.1",
159 binary_path="/candidate/codewhale",
160 max_turns=17,
161 disabled_tools=["web_run"],
162 env={"CALLER_SETTING": "kept", "OPENAI_API_KEY": "must-not-win"},
163 )
164 runtime = FakeRuntime()
165 rollout = trace()
166 ctx = SimpleNamespace(model="test/model")
167
168 await CodewhaleHarness(config).launch(
169 ctx,
170 rollout,
171 runtime,
172 "http://127.0.0.1:9000/v1/",
173 "vf-session-secret",
174 {"grader": "http://127.0.0.1:8123/mcp"},
175 )
176
177 self.assertEqual(len(runtime.programs), 1)
178 argv, env = runtime.programs[0]
179 self.assertEqual(argv[0], "/candidate/codewhale")
180 self.assertIn("--no-project-config", argv)
181 self.assertIn("--skip-onboarding", argv)
182 self.assertEqual(argv[argv.index("--sandbox") + 1], "workspace-write")
183 self.assertEqual(argv[argv.index("--max-turns") + 1], "17")
184 self.assertEqual(argv[argv.index("--disallowed-tools") + 1], "web_run")
185 self.assertEqual(argv[argv.index("--append-system-prompt") + 1], "Use evidence")
186 self.assertNotIn("vf-session-secret", argv)
187 self.assertEqual(env["OPENAI_API_KEY"], "vf-session-secret")
188 self.assertEqual(env["OPENAI_BASE_URL"], "http://127.0.0.1:9000/v1")
189 self.assertEqual(env["CODEWHALE_PROVIDER"], "openai")
190 self.assertEqual(env["CODEWHALE_MODEL"], "test/model")
191 self.assertEqual(env["CODEWHALE_TELEMETRY"], "false")
192 self.assertEqual(env["CODEWHALE_ALLOW_INSECURE_HTTP"], "1")
193 self.assertEqual(env["CALLER_SETTING"], "kept")
194 self.assertEqual(len(runtime.writes), 1)
195 mcp_path, mcp_bytes = runtime.writes[0]
196 self.assertEqual(mcp_path, env["CODEWHALE_MCP_CONFIG"])
197 self.assertTrue(mcp_path.startswith(".vf-codewhale/"))
198 self.assertNotIn(rollout.id, mcp_path)
199 self.assertEqual(
200 json.loads(mcp_bytes),
201 {"servers": {"grader": {"url": "http://127.0.0.1:8123/mcp"}}},
202 )
203
204 receipt = rollout.info["codewhale"]
205 self.assertEqual(receipt["events"], {"content": 1, "done": 1, "metadata": 1})
206 self.assertNotIn("sensitive model output", json.dumps(receipt))
207 self.assertNotIn("private-session-id", json.dumps(receipt))
208 self.assertNotIn("/private/runtime/path", json.dumps(receipt))
209 self.assertNotIn("vf-session-secret", json.dumps(receipt))
210
211 async def test_isolated_runtime_uses_external_sandbox_without_elevation_flag(self):
212 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
213 runtime = FakeRuntime(
214 runtime_type="prime",
215 result=ProgramResult(
216 0, successful_stream(sandbox="external-sandbox"), ""
217 ),
218 )
219
220 await CodewhaleHarness(config).launch(
221 SimpleNamespace(model="test/model"),
222 trace(),
223 runtime,
224 "https://intercept.example/v1",
225 "secret",
226 {},
227 )
228
229 argv, env = runtime.programs[0]
230 self.assertEqual(argv[argv.index("--sandbox") + 1], "external-sandbox")
231 self.assertNotIn("--allow-sandbox-elevation", argv)
232 self.assertNotIn("--max-turns", argv)
233 self.assertNotIn("CODEWHALE_ALLOW_INSECURE_HTTP", env)
234
235 async def test_success_without_exact_terminal_receipt_fails_closed(self):
236 runtime = FakeRuntime(result=ProgramResult(0, event("content", content="ok"), ""))
237 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
238
239 with self.assertRaisesRegex(RuntimeError, "omitted terminal metadata"):
240 await CodewhaleHarness(config).launch(
241 SimpleNamespace(model="test/model"),
242 trace(),
243 runtime,
244 "https://intercept.example/v1",
245 "secret-not-in-error",
246 {},
247 )
248
249 async def test_non_resolved_terminal_receipt_fails_closed(self):
250 stdout = successful_stream().replace('"resolved"', '"approval_required"')
251 runtime = FakeRuntime(result=ProgramResult(0, stdout, ""))
252 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
253
254 with self.assertRaisesRegex(RuntimeError, "was not resolved"):
255 await CodewhaleHarness(config).launch(
256 SimpleNamespace(model="test/model"),
257 trace(),
258 runtime,
259 "https://intercept.example/v1",
260 "secret-not-in-error",
261 {},
262 )
263
264 async def test_oversized_terminal_receipt_fails_closed_without_copying_content(self):
265 stdout = successful_stream().replace("test/model", "x" * 513)
266 runtime = FakeRuntime(result=ProgramResult(0, stdout, ""))
267 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
268
269 with self.assertRaisesRegex(RuntimeError, "exceeded its string bound"):
270 await CodewhaleHarness(config).launch(
271 SimpleNamespace(model="x" * 513),
272 trace(),
273 runtime,
274 "https://intercept.example/v1",
275 "secret-not-in-error",
276 {},
277 )
278
279 def test_exact_version_rejects_prerelease_and_longer_tokens(self):
280 has_version = harness_module._has_version
281
282 self.assertTrue(has_version("codewhale 0.9.1 (abc123)", "0.9.1"))
283 self.assertFalse(has_version("codewhale 0.9.1-rc.1", "0.9.1"))
284 self.assertFalse(has_version("codewhale 0.9.10", "0.9.1"))
285
286 async def test_invalid_config_rejects_shell_and_tool_injection(self):
287 for config in [
288 CodewhaleHarnessConfig(version="0.9.1; touch /tmp/pwned"),
289 CodewhaleHarnessConfig(disabled_tools=["read_file,exec_shell"]),
290 CodewhaleHarnessConfig(max_turns=0),
291 ]:
292 with self.assertRaises(ValueError):
293 await CodewhaleHarness(config).setup(FakeRuntime())
294
295 async def test_install_is_pinned_and_checksum_verified(self):
296 config = CodewhaleHarnessConfig(version="0.9.1")
297 runtime = FakeRuntime()
298
299 await CodewhaleHarness(config).setup(runtime)
300
301 self.assertEqual(len(runtime.runs), 1)
302 argv, env = runtime.runs[0]
303 self.assertEqual(argv[:2], ["sh", "-c"])
304 self.assertEqual(env, {})
305 command = argv[2]
306 self.assertIn("releases/download", command)
307 self.assertIn("codewhale-artifacts-sha256.txt", command)
308 self.assertIn("sha256sum", command)
309 self.assertIn("flock", command)
310 self.assertIn("util-linux", command)
311
312
313 if __name__ == "__main__":
314 unittest.main()
315
315 lines PYTHON