Coverage for agentos/swarm/code

1"""

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

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

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

6"""

8from __future__ import annotations

10import json

11import os

12import subprocess

13import tempfile

14import traceback

15from dataclasses import dataclass, field

16from typing import Any

19@dataclass

20class SandboxResult:

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

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

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 }

47@dataclass

48class TestCase:

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

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

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

68class CodeSandbox:

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

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 """

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

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

155 statements