Coverage for agentos/swarm/code_sandbox.py: 27%

155 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-08 21:26 +0800

1""" 

2v1.9.5: Isolated Code Sandbox for safe agent code generation & execution. 

3 

4Runs generated code in subprocess with timeout, memory limits, 

5test case validation, and structured error extraction for feedback loops. 

6""" 

7 

8from __future__ import annotations 

9 

10import json 

11import os 

12import subprocess 

13import tempfile 

14import traceback 

15from dataclasses import dataclass, field 

16from typing import Any 

17 

18 

19@dataclass 

20class SandboxResult: 

21 """Result of a sandbox execution.""" 

22 

23 success: bool = False 

24 exit_code: int = -1 

25 stdout: str = "" 

26 stderr: str = "" 

27 exception: str = "" 

28 duration: float = 0.0 

29 max_memory_mb: float = 0.0 

30 test_results: list[dict] = field(default_factory=list) # per-test-case results 

31 all_passed: bool = False 

32 

33 def to_dict(self) -> dict: 

34 return { 

35 "success": self.success, 

36 "exit_code": self.exit_code, 

37 "stdout": self.stdout[:500], 

38 "stderr": self.stderr[:500], 

39 "exception": self.exception[:300], 

40 "duration": f"{self.duration:.2f}s", 

41 "max_memory_mb": f"{self.max_memory_mb:.1f}", 

42 "test_results": self.test_results, 

43 "all_passed": self.all_passed, 

44 } 

45 

46 

47@dataclass 

48class TestCase: 

49 """A single test case for code validation.""" 

50 

51 name: str 

52 input_args: tuple = () 

53 input_kwargs: dict = field(default_factory=dict) 

54 expected_output: Any = None 

55 expected_type: str = "" # e.g. "int", "str", "list" 

56 expected_exception: str = "" # e.g. "ValueError" 

57 weight: float = 1.0 

58 

59 def to_dict(self) -> dict: 

60 d: dict = {"name": self.name, "input_args": str(self.input_args)} 

61 if self.expected_output is not None: 

62 d["expected_output"] = str(self.expected_output)[:100] 

63 if self.expected_exception: 

64 d["expected_exception"] = self.expected_exception 

65 return d 

66 

67 

68class CodeSandbox: 

69 """Isolated execution environment for agent-generated code. 

70 

71 Usage: 

72 sandbox = CodeSandbox(timeout=30, max_memory_mb=256) 

73 result = sandbox.run( 

74 code="def add(a, b): return a + b", 

75 func_name="add", 

76 test_cases=[TestCase(name="1+2", args=(1, 2), expected=3)] 

77 ) 

78 """ 

79 

80 def __init__( 

81 self, 

82 timeout: float = 30.0, 

83 max_memory_mb: int = 256, 

84 allow_imports: bool = True, 

85 forbidden_modules: list[str] | None = None, 

86 ): 

87 self.timeout = timeout 

88 self.max_memory_mb = max_memory_mb 

89 self.allow_imports = allow_imports 

90 self.forbidden_modules = forbidden_modules or [ 

91 "os.system", 

92 "subprocess", 

93 "shutil", 

94 "socket", 

95 "requests", 

96 "urllib", 

97 "http", 

98 "ftp", 

99 ] 

100 

101 def run( 

102 self, 

103 code: str, 

104 func_name: str = "", 

105 test_cases: list[TestCase] | None = None, 

106 setup_code: str = "", 

107 ) -> SandboxResult: 

108 """Execute code in sandbox with test cases. 

109 

110 Args: 

111 code: The code to execute 

112 func_name: Name of the function to test 

113 test_cases: List of test cases to run against func_name 

114 setup_code: Setup code to run before tests (imports, fixtures) 

115 

116 Returns: 

117 SandboxResult with execution details and test outcomes 

118 """ 

119 test_cases = test_cases or [] 

120 

121 # Security check 

122 security_issues = self._check_security(code) 

123 if security_issues: 

124 return SandboxResult( 

125 success=False, 

126 stderr=f"SECURITY VIOLATION: {security_issues}", 

127 exception="Security check failed", 

128 ) 

129 

130 # Syntactic check 

131 syntax_ok, syntax_err = self._check_syntax(code) 

132 if not syntax_ok: 

133 return SandboxResult( 

134 success=False, 

135 stderr=syntax_err, 

136 exception=f"Syntax error: {syntax_err}", 

137 ) 

138 

139 # Write code to temp file and execute 

140 with tempfile.NamedTemporaryFile( 

141 mode="w", suffix=".py", delete=False, prefix="sandbox_" 

142 ) as f: 

143 script = self._build_script(code, func_name, test_cases, setup_code) 

144 f.write(script) 

145 script_path = f.name 

146 

147 try: 

148 result = self._execute_script(script_path) 

149 finally: 

150 try: 

151 os.unlink(script_path) 

152 except OSError: 

153 pass 

154 

155 return result 

156 

157 def _check_security(self, code: str) -> str: 

158 """Check for forbidden patterns.""" 

159 issues = [] 

160 code_lower = code.lower() 

161 

162 # Check forbidden modules 

163 for mod in self.forbidden_modules: 

164 import_pattern = mod.replace(".", ".") 

165 if f"import {import_pattern}" in code_lower: 

166 issues.append(f"Forbidden import: {mod}") 

167 if f"from {import_pattern}" in code_lower: 

168 issues.append(f"Forbidden import: {mod}") 

169 

170 # Check dangerous builtins 

171 dangerous = [ 

172 ("eval(", "eval() is forbidden"), 

173 ("exec(", "exec() is forbidden"), 

174 ("__import__(", "__import__() is forbidden"), 

175 ("open(", "File I/O blocked in sandbox"), 

176 ] 

177 for pattern, msg in dangerous: 

178 if pattern in code_lower: 

179 issues.append(msg) 

180 

181 return "; ".join(issues) if issues else "" 

182 

183 def _check_syntax(self, code: str) -> tuple[bool, str]: 

184 """Check Python syntax.""" 

185 try: 

186 compile(code, "<sandbox>", "exec") 

187 return True, "" 

188 except SyntaxError as e: 

189 return False, f"Line {e.lineno}: {e.msg}" 

190 

191 def _build_script( 

192 self, 

193 code: str, 

194 func_name: str, 

195 test_cases: list[TestCase], 

196 setup_code: str, 

197 ) -> str: 

198 """Build the full sandbox execution script.""" 

199 

200 tc_defs = [] 

201 tc_names = [] 

202 for tc in test_cases: 

203 tc_names.append(tc.name) 

204 tc_defs.append( 

205 f' "{tc.name}": {{' 

206 f'"args": {list(tc.input_args)}, ' 

207 f'"kwargs": {tc.input_kwargs}, ' 

208 f'"expected": {repr(tc.expected_output) if tc.expected_output is not None else None}, ' 

209 f'"expected_type": "{tc.expected_type}", ' 

210 f'"expected_exception": "{tc.expected_exception}"' 

211 f"}}" 

212 ) 

213 

214 tc_dict = ",\n".join(tc_defs) 

215 tc_list = ", ".join(f'"{n}"' for n in tc_names) 

216 

217 return f'''#!/usr/bin/env python3 

218"""Sandbox execution script — auto-generated by CodeSandbox.""" 

219import json 

220import sys 

221import time 

222import traceback 

223import os 

224 

225# Limit memory 

226try: 

227 import resource 

228 resource.setrlimit(resource.RLIMIT_AS, ({self.max_memory_mb} * 1024 * 1024, {self.max_memory_mb} * 1024 * 1024)) 

229except Exception: 

230 pass 

231 

232# Setup code 

233{setup_code} 

234 

235# User code 

236{code} 

237 

238# Test runner 

239test_cases = {{ 

240{tc_dict} 

241}} 

242 

243func_name = "{func_name}" 

244has_func = func_name and func_name in dir() 

245 

246results = [] 

247total = 0 

248passed = 0 

249 

250if not has_func and len(test_cases) == 0: 

251 print("__SANDBOX_OK__") 

252 print(json.dumps({{"status": "executed", "message": "Code executed, no tests"}})) 

253 sys.exit(0) 

254 

255if not has_func and len(test_cases) > 0: 

256 print("__SANDBOX_ERR__") 

257 print(json.dumps({{"error": f"Function '{{func_name}}' not found in code"}})) 

258 sys.exit(1) 

259 

260func = eval(func_name) 

261all_test_order = [{tc_list}] 

262 

263for tc_name in all_test_order: 

264 tc = test_cases[tc_name] 

265 total += 1 

266 start = time.time() 

267 try: 

268 output = func(*tc["args"], **tc["kwargs"]) 

269 elapsed = time.time() - start 

270 

271 expected = tc.get("expected") 

272 expected_type = tc.get("expected_type") 

273 expected_exc = tc.get("expected_exception") 

274 

275 if expected_exc: 

276 passed_ = False 

277 detail = f"Expected exception {{expected_exc}} but none raised" 

278 elif expected_type and not isinstance(output, eval(expected_type)): 

279 passed_ = False 

280 detail = f"Type mismatch: got {{type(output).__name__}}, expected {{expected_type}}" 

281 elif expected is not None and output != expected: 

282 passed_ = False 

283 detail = f"Expected {{repr(expected)}}, got {{repr(output)}}" 

284 else: 

285 passed_ = True 

286 detail = "OK" 

287 

288 if passed_: 

289 passed += 1 

290 

291 results.append({{ 

292 "name": tc_name, 

293 "passed": passed_, 

294 "output": repr(output)[:200], 

295 "expected": repr(expected)[:200], 

296 "detail": detail, 

297 "duration_ms": round(elapsed * 1000, 2), 

298 }}) 

299 except Exception as e: 

300 exc_name = type(e).__name__ 

301 expected_exc = tc.get("expected_exception", "") 

302 

303 if expected_exc and exc_name == expected_exc: 

304 passed_ = True 

305 detail = f"Expected exception {{exc_name}} raised" 

306 passed += 1 

307 else: 

308 passed_ = False 

309 detail = f"{{exc_name}}: {{str(e)[:200]}}" 

310 

311 results.append({{ 

312 "name": tc_name, 

313 "passed": passed_, 

314 "output": "", 

315 "expected": repr(tc.get("expected"))[:200], 

316 "detail": detail, 

317 "duration_ms": 0, 

318 }}) 

319 

320print("__SANDBOX_RESULTS__") 

321print(json.dumps({{ 

322 "status": "complete", 

323 "total": total, 

324 "passed": passed, 

325 "failed": total - passed, 

326 "all_passed": total > 0 and passed == total, 

327 "test_results": results, 

328}})) 

329''' 

330 

331 def _execute_script(self, script_path: str) -> SandboxResult: 

332 """Execute the sandbox script as subprocess.""" 

333 try: 

334 start = time_module() 

335 proc = subprocess.run( 

336 ["python3", script_path], 

337 capture_output=True, 

338 text=True, 

339 timeout=self.timeout, 

340 cwd="/tmp", 

341 env={ 

342 **os.environ, 

343 "PYTHONDONTWRITEBYTECODE": "1", 

344 "PYTHONPATH": "", 

345 "SANDBOX_MODE": "1", 

346 }, 

347 ) 

348 duration = time_module() - start 

349 except subprocess.TimeoutExpired as e: 

350 return SandboxResult( 

351 success=False, 

352 exit_code=-1, 

353 stderr=str(e.stdout or "") if e.stdout else "", 

354 exception=f"Timeout after {self.timeout}s", 

355 duration=self.timeout, 

356 ) 

357 except Exception as e: 

358 return SandboxResult( 

359 success=False, 

360 exception=str(e), 

361 stderr=traceback.format_exc(), 

362 ) 

363 

364 result = SandboxResult( 

365 exit_code=proc.returncode, 

366 stdout=proc.stdout, 

367 stderr=proc.stderr, 

368 duration=duration, 

369 ) 

370 

371 # Parse test results 

372 if "__SANDBOX_RESULTS__" in proc.stdout: 

373 try: 

374 lines = proc.stdout.split("\n") 

375 json_start = False 

376 json_text = "" 

377 for line in lines: 

378 if json_start: 

379 json_text += line 

380 try: 

381 data = json.loads(json_text) 

382 break 

383 except json.JSONDecodeError: 

384 continue 

385 if "__SANDBOX_RESULTS__" in line: 

386 json_start = True 

387 

388 if data: 

389 result.test_results = data.get("test_results", []) 

390 result.all_passed = data.get("all_passed", False) 

391 passed = data.get("passed", 0) 

392 total = data.get("total", 0) 

393 result.success = total > 0 and passed == total 

394 except Exception: 

395 pass 

396 

397 elif proc.returncode == 0 and not proc.stderr: 

398 result.success = True 

399 

400 # Extract meaningful error info for feedback 

401 if not result.success and result.stderr: 

402 lines = result.stderr.strip().split("\n") 

403 # Get last 3 lines (most relevant error info) 

404 result.exception = "\n".join(lines[-5:]) if len(lines) > 5 else result.stderr 

405 

406 return result 

407 

408 

409def time_module() -> float: 

410 """Get current time in seconds.""" 

411 import time as _time 

412 

413 return _time.time() 

414 

415 

416class CodeFeedbackExtractor: 

417 """Extracts actionable feedback from sandbox failures for retry loops.""" 

418 

419 ERROR_PATTERNS = { 

420 "NameError": "Variable or function not defined. Check spelling and scope.", 

421 "TypeError": "Wrong type passed to function. Check argument types.", 

422 "ValueError": "Invalid value for operation. Check parameter constraints.", 

423 "IndexError": "List index out of range. Check bounds.", 

424 "KeyError": "Dictionary key not found. Check key existence.", 

425 "AttributeError": "Object has no such attribute. Check method/attribute name.", 

426 "ImportError": "Missing import. Add 'import X' or 'from X import Y'.", 

427 "SyntaxError": "Python syntax error. Check indentation, brackets, colons.", 

428 "ZeroDivisionError": "Division by zero. Add guard for zero denominator.", 

429 "RecursionError": "Recursion depth exceeded. Add base case or switch to iteration.", 

430 "TimeoutError": "Code timed out. Check for infinite loops or optimize.", 

431 } 

432 

433 @classmethod 

434 def extract(cls, sandbox_result: SandboxResult) -> list[str]: 

435 """Extract actionable feedback suggestions from sandbox result.""" 

436 suggestions: list[str] = [] 

437 

438 # Check for security issues 

439 if "SECURITY" in sandbox_result.stderr: 

440 suggestions.append("Code violates sandbox security rules; avoid system calls and I/O.") 

441 return suggestions 

442 

443 # Check test failures 

444 for tc in sandbox_result.test_results: 

445 if not tc.get("passed", True): 

446 detail = tc.get("detail", "") 

447 name = tc.get("name", "unknown") 

448 suggestions.append(f"Test '{name}' failed: {detail}") 

449 

450 # Check error patterns 

451 for error_type, suggestion in cls.ERROR_PATTERNS.items(): 

452 if error_type in sandbox_result.stderr or error_type in sandbox_result.exception: 

453 if suggestion not in suggestions: 

454 suggestions.append(suggestion) 

455 

456 # Add specific output mismatch advice 

457 if not sandbox_result.success and not suggestions: 

458 suggestions.append("Code failed to execute. Check logic and edge cases.") 

459 if sandbox_result.stderr: 

460 suggestions.append(f"Error: {sandbox_result.stderr.strip()[:200]}") 

461 

462 return suggestions