Coverage for agentos/swarm/code_sandbox.py: 27%
154 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +0800
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 os
11import subprocess
12import tempfile
13import traceback
14from dataclasses import dataclass, field
15from typing import Any
18@dataclass
19class SandboxResult:
20 """Result of a sandbox execution."""
22 success: bool = False
23 exit_code: int = -1
24 stdout: str = ""
25 stderr: str = ""
26 exception: str = ""
27 duration: float = 0.0
28 max_memory_mb: float = 0.0
29 test_results: list[dict] = field(default_factory=list) # per-test-case results
30 all_passed: bool = False
32 def to_dict(self) -> dict:
33 return {
34 "success": self.success,
35 "exit_code": self.exit_code,
36 "stdout": self.stdout[:500],
37 "stderr": self.stderr[:500],
38 "exception": self.exception[:300],
39 "duration": f"{self.duration:.2f}s",
40 "max_memory_mb": f"{self.max_memory_mb:.1f}",
41 "test_results": self.test_results,
42 "all_passed": self.all_passed,
43 }
46@dataclass
47class TestCase:
48 """A single test case for code validation."""
50 name: str
51 input_args: tuple = ()
52 input_kwargs: dict = field(default_factory=dict)
53 expected_output: Any = None
54 expected_type: str = "" # e.g. "int", "str", "list"
55 expected_exception: str = "" # e.g. "ValueError"
56 weight: float = 1.0
58 def to_dict(self) -> dict:
59 d: dict = {"name": self.name, "input_args": str(self.input_args)}
60 if self.expected_output is not None:
61 d["expected_output"] = str(self.expected_output)[:100]
62 if self.expected_exception:
63 d["expected_exception"] = self.expected_exception
64 return d
67class CodeSandbox:
68 """Isolated execution environment for agent-generated code.
70 Usage:
71 sandbox = CodeSandbox(timeout=30, max_memory_mb=256)
72 result = sandbox.run(
73 code="def add(a, b): return a + b",
74 func_name="add",
75 test_cases=[TestCase(name="1+2", args=(1, 2), expected=3)]
76 )
77 """
79 def __init__(
80 self,
81 timeout: float = 30.0,
82 max_memory_mb: int = 256,
83 allow_imports: bool = True,
84 forbidden_modules: list[str] | None = None,
85 ):
86 self.timeout = timeout
87 self.max_memory_mb = max_memory_mb
88 self.allow_imports = allow_imports
89 self.forbidden_modules = forbidden_modules or [
90 "os.system", "subprocess", "shutil", "socket",
91 "requests", "urllib", "http", "ftp",
92 ]
94 def run(
95 self,
96 code: str,
97 func_name: str = "",
98 test_cases: list[TestCase] | None = None,
99 setup_code: str = "",
100 ) -> SandboxResult:
101 """Execute code in sandbox with test cases.
103 Args:
104 code: The code to execute
105 func_name: Name of the function to test
106 test_cases: List of test cases to run against func_name
107 setup_code: Setup code to run before tests (imports, fixtures)
109 Returns:
110 SandboxResult with execution details and test outcomes
111 """
112 test_cases = test_cases or []
114 # Security check
115 security_issues = self._check_security(code)
116 if security_issues:
117 return SandboxResult(
118 success=False,
119 stderr=f"SECURITY VIOLATION: {security_issues}",
120 exception="Security check failed",
121 )
123 # Syntactic check
124 syntax_ok, syntax_err = self._check_syntax(code)
125 if not syntax_ok:
126 return SandboxResult(
127 success=False,
128 stderr=syntax_err,
129 exception=f"Syntax error: {syntax_err}",
130 )
132 # Write code to temp file and execute
133 with tempfile.NamedTemporaryFile(
134 mode="w", suffix=".py", delete=False,
135 prefix="sandbox_"
136 ) as f:
137 script = self._build_script(code, func_name, test_cases, setup_code)
138 f.write(script)
139 script_path = f.name
141 try:
142 result = self._execute_script(script_path)
143 finally:
144 try:
145 os.unlink(script_path)
146 except OSError:
147 pass
149 return result
151 def _check_security(self, code: str) -> str:
152 """Check for forbidden patterns."""
153 issues = []
154 code_lower = code.lower()
156 # Check forbidden modules
157 for mod in self.forbidden_modules:
158 import_pattern = mod.replace(".", ".")
159 if f"import {import_pattern}" in code_lower:
160 issues.append(f"Forbidden import: {mod}")
161 if f"from {import_pattern}" in code_lower:
162 issues.append(f"Forbidden import: {mod}")
164 # Check dangerous builtins
165 dangerous = [
166 ("eval(", "eval() is forbidden"),
167 ("exec(", "exec() is forbidden"),
168 ("__import__(", "__import__() is forbidden"),
169 ("open(", "File I/O blocked in sandbox"),
170 ]
171 for pattern, msg in dangerous:
172 if pattern in code_lower:
173 issues.append(msg)
175 return "; ".join(issues) if issues else ""
177 def _check_syntax(self, code: str) -> tuple[bool, str]:
178 """Check Python syntax."""
179 try:
180 compile(code, "<sandbox>", "exec")
181 return True, ""
182 except SyntaxError as e:
183 return False, f"Line {e.lineno}: {e.msg}"
185 def _build_script(
186 self,
187 code: str,
188 func_name: str,
189 test_cases: list[TestCase],
190 setup_code: str,
191 ) -> str:
192 """Build the full sandbox execution script."""
194 tc_defs = []
195 tc_names = []
196 for tc in test_cases:
197 tc_names.append(tc.name)
198 tc_defs.append(
199 f' "{tc.name}": {{'
200 f'"args": {list(tc.input_args)}, '
201 f'"kwargs": {tc.input_kwargs}, '
202 f'"expected": {repr(tc.expected_output) if tc.expected_output is not None else None}, '
203 f'"expected_type": "{tc.expected_type}", '
204 f'"expected_exception": "{tc.expected_exception}"'
205 f"}}"
206 )
208 tc_dict = ",\n".join(tc_defs)
209 tc_list = ", ".join(f'"{n}"' for n in tc_names)
211 return f'''#!/usr/bin/env python3
212"""Sandbox execution script — auto-generated by CodeSandbox."""
213import json
214import sys
215import time
216import traceback
217import os
219# Limit memory
220try:
221 import resource
222 resource.setrlimit(resource.RLIMIT_AS, ({self.max_memory_mb} * 1024 * 1024, {self.max_memory_mb} * 1024 * 1024))
223except Exception:
224 pass
226# Setup code
227{setup_code}
229# User code
230{code}
232# Test runner
233test_cases = {{
234{tc_dict}
235}}
237func_name = "{func_name}"
238has_func = func_name and func_name in dir()
240results = []
241total = 0
242passed = 0
244if not has_func and len(test_cases) == 0:
245 print("__SANDBOX_OK__")
246 print(json.dumps({{"status": "executed", "message": "Code executed, no tests"}}))
247 sys.exit(0)
249if not has_func and len(test_cases) > 0:
250 print("__SANDBOX_ERR__")
251 print(json.dumps({{"error": f"Function '{{func_name}}' not found in code"}}))
252 sys.exit(1)
254func = eval(func_name)
255all_test_order = [{tc_list}]
257for tc_name in all_test_order:
258 tc = test_cases[tc_name]
259 total += 1
260 start = time.time()
261 try:
262 output = func(*tc["args"], **tc["kwargs"])
263 elapsed = time.time() - start
265 expected = tc.get("expected")
266 expected_type = tc.get("expected_type")
267 expected_exc = tc.get("expected_exception")
269 if expected_exc:
270 passed_ = False
271 detail = f"Expected exception {{expected_exc}} but none raised"
272 elif expected_type and not isinstance(output, eval(expected_type)):
273 passed_ = False
274 detail = f"Type mismatch: got {{type(output).__name__}}, expected {{expected_type}}"
275 elif expected is not None and output != expected:
276 passed_ = False
277 detail = f"Expected {{repr(expected)}}, got {{repr(output)}}"
278 else:
279 passed_ = True
280 detail = "OK"
282 if passed_:
283 passed += 1
285 results.append({{
286 "name": tc_name,
287 "passed": passed_,
288 "output": repr(output)[:200],
289 "expected": repr(expected)[:200],
290 "detail": detail,
291 "duration_ms": round(elapsed * 1000, 2),
292 }})
293 except Exception as e:
294 exc_name = type(e).__name__
295 expected_exc = tc.get("expected_exception", "")
297 if expected_exc and exc_name == expected_exc:
298 passed_ = True
299 detail = f"Expected exception {{exc_name}} raised"
300 passed += 1
301 else:
302 passed_ = False
303 detail = f"{{exc_name}}: {{str(e)[:200]}}"
305 results.append({{
306 "name": tc_name,
307 "passed": passed_,
308 "output": "",
309 "expected": repr(tc.get("expected"))[:200],
310 "detail": detail,
311 "duration_ms": 0,
312 }})
314print("__SANDBOX_RESULTS__")
315print(json.dumps({{
316 "status": "complete",
317 "total": total,
318 "passed": passed,
319 "failed": total - passed,
320 "all_passed": total > 0 and passed == total,
321 "test_results": results,
322}}))
323'''
325 def _execute_script(self, script_path: str) -> SandboxResult:
326 """Execute the sandbox script as subprocess."""
327 try:
328 start = time_module()
329 proc = subprocess.run(
330 ["python3", script_path],
331 capture_output=True,
332 text=True,
333 timeout=self.timeout,
334 cwd="/tmp",
335 env={
336 **os.environ,
337 "PYTHONDONTWRITEBYTECODE": "1",
338 "PYTHONPATH": "",
339 "SANDBOX_MODE": "1",
340 },
341 )
342 duration = time_module() - start
343 except subprocess.TimeoutExpired as e:
344 return SandboxResult(
345 success=False,
346 exit_code=-1,
347 stderr=str(e.stdout or "") if e.stdout else "",
348 exception=f"Timeout after {self.timeout}s",
349 duration=self.timeout,
350 )
351 except Exception as e:
352 return SandboxResult(
353 success=False,
354 exception=str(e),
355 stderr=traceback.format_exc(),
356 )
358 result = SandboxResult(
359 exit_code=proc.returncode,
360 stdout=proc.stdout,
361 stderr=proc.stderr,
362 duration=duration,
363 )
365 # Parse test results
366 if "__SANDBOX_RESULTS__" in proc.stdout:
367 try:
368 lines = proc.stdout.split("\n")
369 json_start = False
370 json_text = ""
371 for line in lines:
372 if json_start:
373 json_text += line
374 try:
375 data = json.loads(json_text)
376 break
377 except json.JSONDecodeError:
378 continue
379 if "__SANDBOX_RESULTS__" in line:
380 json_start = True
382 if data:
383 result.test_results = data.get("test_results", [])
384 result.all_passed = data.get("all_passed", False)
385 passed = data.get("passed", 0)
386 total = data.get("total", 0)
387 result.success = total > 0 and passed == total
388 except Exception:
389 pass
391 elif proc.returncode == 0 and not proc.stderr:
392 result.success = True
394 # Extract meaningful error info for feedback
395 if not result.success and result.stderr:
396 lines = result.stderr.strip().split("\n")
397 # Get last 3 lines (most relevant error info)
398 result.exception = "\n".join(lines[-5:]) if len(lines) > 5 else result.stderr
400 return result
403def time_module() -> float:
404 """Get current time in seconds."""
405 import time as _time
406 return _time.time()
409class CodeFeedbackExtractor:
410 """Extracts actionable feedback from sandbox failures for retry loops."""
412 ERROR_PATTERNS = {
413 "NameError": "Variable or function not defined. Check spelling and scope.",
414 "TypeError": "Wrong type passed to function. Check argument types.",
415 "ValueError": "Invalid value for operation. Check parameter constraints.",
416 "IndexError": "List index out of range. Check bounds.",
417 "KeyError": "Dictionary key not found. Check key existence.",
418 "AttributeError": "Object has no such attribute. Check method/attribute name.",
419 "ImportError": "Missing import. Add 'import X' or 'from X import Y'.",
420 "SyntaxError": "Python syntax error. Check indentation, brackets, colons.",
421 "ZeroDivisionError": "Division by zero. Add guard for zero denominator.",
422 "RecursionError": "Recursion depth exceeded. Add base case or switch to iteration.",
423 "TimeoutError": "Code timed out. Check for infinite loops or optimize.",
424 }
426 @classmethod
427 def extract(cls, sandbox_result: SandboxResult) -> list[str]:
428 """Extract actionable feedback suggestions from sandbox result."""
429 suggestions: list[str] = []
431 # Check for security issues
432 if "SECURITY" in sandbox_result.stderr:
433 suggestions.append("Code violates sandbox security rules; avoid system calls and I/O.")
434 return suggestions
436 # Check test failures
437 for tc in sandbox_result.test_results:
438 if not tc.get("passed", True):
439 detail = tc.get("detail", "")
440 name = tc.get("name", "unknown")
441 suggestions.append(f"Test '{name}' failed: {detail}")
443 # Check error patterns
444 for error_type, suggestion in cls.ERROR_PATTERNS.items():
445 if error_type in sandbox_result.stderr or error_type in sandbox_result.exception:
446 if suggestion not in suggestions:
447 suggestions.append(suggestion)
449 # Add specific output mismatch advice
450 if not sandbox_result.success and not suggestions:
451 suggestions.append("Code failed to execute. Check logic and edge cases.")
452 if sandbox_result.stderr:
453 suggestions.append(f"Error: {sandbox_result.stderr.strip()[:200]}")
455 return suggestions