返回 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("CODEWHALE_ALLOW_INSECURE_HTTP", env)
233
234 async def test_success_without_exact_terminal_receipt_fails_closed(self):
235 runtime = FakeRuntime(result=ProgramResult(0, event("content", content="ok"), ""))
236 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
237
238 with self.assertRaisesRegex(RuntimeError, "omitted terminal metadata"):
239 await CodewhaleHarness(config).launch(
240 SimpleNamespace(model="test/model"),
241 trace(),
242 runtime,
243 "https://intercept.example/v1",
244 "secret-not-in-error",
245 {},
246 )
247
248 async def test_non_resolved_terminal_receipt_fails_closed(self):
249 stdout = successful_stream().replace('"resolved"', '"approval_required"')
250 runtime = FakeRuntime(result=ProgramResult(0, stdout, ""))
251 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
252
253 with self.assertRaisesRegex(RuntimeError, "was not resolved"):
254 await CodewhaleHarness(config).launch(
255 SimpleNamespace(model="test/model"),
256 trace(),
257 runtime,
258 "https://intercept.example/v1",
259 "secret-not-in-error",
260 {},
261 )
262
263 async def test_oversized_terminal_receipt_fails_closed_without_copying_content(self):
264 stdout = successful_stream().replace("test/model", "x" * 513)
265 runtime = FakeRuntime(result=ProgramResult(0, stdout, ""))
266 config = CodewhaleHarnessConfig(binary_path="/candidate/codewhale")
267
268 with self.assertRaisesRegex(RuntimeError, "exceeded its string bound"):
269 await CodewhaleHarness(config).launch(
270 SimpleNamespace(model="x" * 513),
271 trace(),
272 runtime,
273 "https://intercept.example/v1",
274 "secret-not-in-error",
275 {},
276 )
277
278 def test_exact_version_rejects_prerelease_and_longer_tokens(self):
279 has_version = harness_module._has_version
280
281 self.assertTrue(has_version("codewhale 0.9.1 (abc123)", "0.9.1"))
282 self.assertFalse(has_version("codewhale 0.9.1-rc.1", "0.9.1"))
283 self.assertFalse(has_version("codewhale 0.9.10", "0.9.1"))
284
285 async def test_invalid_config_rejects_shell_and_tool_injection(self):
286 for config in [
287 CodewhaleHarnessConfig(version="0.9.1; touch /tmp/pwned"),
288 CodewhaleHarnessConfig(disabled_tools=["read_file,exec_shell"]),
289 CodewhaleHarnessConfig(max_turns=0),
290 ]:
291 with self.assertRaises(ValueError):
292 await CodewhaleHarness(config).setup(FakeRuntime())
293
294 async def test_install_is_pinned_and_checksum_verified(self):
295 config = CodewhaleHarnessConfig(version="0.9.1")
296 runtime = FakeRuntime()
297
298 await CodewhaleHarness(config).setup(runtime)
299
300 self.assertEqual(len(runtime.runs), 1)
301 argv, env = runtime.runs[0]
302 self.assertEqual(argv[:2], ["sh", "-c"])
303 self.assertEqual(env, {})
304 command = argv[2]
305 self.assertIn("releases/download", command)
306 self.assertIn("codewhale-artifacts-sha256.txt", command)
307 self.assertIn("sha256sum", command)
308 self.assertIn("flock", command)
309 self.assertIn("util-linux", command)
310
311
312 if __name__ == "__main__":
313 unittest.main()
314
314 lines PYTHON