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.
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 ]
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.
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)
116 Returns:
117 SandboxResult with execution details and test outcomes
118 """
119 test_cases = test_cases or []
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 )
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 )
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
147 try:
148 result = self._execute_script(script_path)
149 finally:
150 try:
151 os.unlink(script_path)
152 except OSError:
153 pass
155 return result
157 def _check_security(self, code: str) -> str:
158 """Check for forbidden patterns."""
159 issues = []
160 code_lower = code.lower()
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}")
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)
181 return "; ".join(issues) if issues else ""
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}"
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."""
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 )
214 tc_dict = ",\n".join(tc_defs)
215 tc_list = ", ".join(f'"{n}"' for n in tc_names)
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
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
232# Setup code
233{setup_code}
235# User code
236{code}
238# Test runner
239test_cases = {{
240{tc_dict}
241}}
243func_name = "{func_name}"
244has_func = func_name and func_name in dir()
246results = []
247total = 0
248passed = 0
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)
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)
260func = eval(func_name)
261all_test_order = [{tc_list}]
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
271 expected = tc.get("expected")
272 expected_type = tc.get("expected_type")
273 expected_exc = tc.get("expected_exception")
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"
288 if passed_:
289 passed += 1
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", "")
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]}}"
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 }})
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'''
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 )
364 result = SandboxResult(
365 exit_code=proc.returncode,
366 stdout=proc.stdout,
367 stderr=proc.stderr,
368 duration=duration,
369 )
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
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
397 elif proc.returncode == 0 and not proc.stderr:
398 result.success = True
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
406 return result
409def time_module() -> float:
410 """Get current time in seconds."""
411 import time as _time
413 return _time.time()
416class CodeFeedbackExtractor:
417 """Extracts actionable feedback from sandbox failures for retry loops."""
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 }
433 @classmethod
434 def extract(cls, sandbox_result: SandboxResult) -> list[str]:
435 """Extract actionable feedback suggestions from sandbox result."""
436 suggestions: list[str] = []
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
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}")
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)
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]}")
462 return suggestions