Coverage for agentos/swarm/code_sandbox.py: 27%
154 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +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",
91 "subprocess",
92 "shutil",
93 "socket",
94 "requests",
95 "urllib",
96 "http",
97 "ftp",
98 ]
100 def run(
101 self,
102 code: str,
103 func_name: str = "",
104 test_cases: list[TestCase] | None = None,
105 setup_code: str = "",
106 ) -> SandboxResult:
107 """Execute code in sandbox with test cases.
109 Args:
110 code: The code to execute
111 func_name: Name of the function to test
112 test_cases: List of test cases to run against func_name
113 setup_code: Setup code to run before tests (imports, fixtures)
115 Returns:
116 SandboxResult with execution details and test outcomes
117 """
118 test_cases = test_cases or []
120 # Security check
121 security_issues = self._check_security(code)
122 if security_issues:
123 return SandboxResult(
124 success=False,
125 stderr=f"SECURITY VIOLATION: {security_issues}",
126 exception="Security check failed",
127 )
129 # Syntactic check
130 syntax_ok, syntax_err = self._check_syntax(code)
131 if not syntax_ok:
132 return SandboxResult(
133 success=False,
134 stderr=syntax_err,
135 exception=f"Syntax error: {syntax_err}",
136 )
138 # Write code to temp file and execute
139 with tempfile.NamedTemporaryFile(
140 mode="w", suffix=".py", delete=False, prefix="sandbox_"
141 ) as f:
142 script = self._build_script(code, func_name, test_cases, setup_code)
143 f.write(script)
144 script_path = f.name
146 try:
147 result = self._execute_script(script_path)
148 finally:
149 try:
150 os.unlink(script_path)
151 except OSError:
152 pass
154 return result
156 def _check_security(self, code: str) -> str:
157 """Check for forbidden patterns."""
158 issues = []
159 code_lower = code.lower()
161 # Check forbidden modules
162 for mod in self.forbidden_modules:
163 import_pattern = mod.replace(".", ".")
164 if f"import {import_pattern}" in code_lower:
165 issues.append(f"Forbidden import: {mod}")
166 if f"from {import_pattern}" in code_lower:
167 issues.append(f"Forbidden import: {mod}")
169 # Check dangerous builtins
170 dangerous = [
171 ("eval(", "eval() is forbidden"),
172 ("exec(", "exec() is forbidden"),
173 ("__import__(", "__import__() is forbidden"),
174 ("open(", "File I/O blocked in sandbox"),
175 ]
176 for pattern, msg in dangerous:
177 if pattern in code_lower:
178 issues.append(msg)
180 return "; ".join(issues) if issues else ""
182 def _check_syntax(self, code: str) -> tuple[bool, str]:
183 """Check Python syntax."""
184 try:
185 compile(code, "<sandbox>", "exec")
186 return True, ""
187 except SyntaxError as e:
188 return False, f"Line {e.lineno}: {e.msg}"
190 def _build_script(
191 self,
192 code: str,
193 func_name: str,
194 test_cases: list[TestCase],
195 setup_code: str,
196 ) -> str:
197 """Build the full sandbox execution script."""
199 tc_defs = []
200 tc_names = []
201 for tc in test_cases:
202 tc_names.append(tc.name)
203 tc_defs.append(
204 f' "{tc.name}": {{'
205 f'"args": {list(tc.input_args)}, '
206 f'"kwargs": {tc.input_kwargs}, '
207 f'"expected": {repr(tc.expected_output) if tc.expected_output is not None else None}, '
208 f'"expected_type": "{tc.expected_type}", '
209 f'"expected_exception": "{tc.expected_exception}"'
210 f"}}"
211 )
213 tc_dict = ",\n".join(tc_defs)
214 tc_list = ", ".join(f'"{n}"' for n in tc_names)
216 return f'''#!/usr/bin/env python3
217"""Sandbox execution script — auto-generated by CodeSandbox."""
218import json
219import sys
220import time
221import traceback
222import os
224# Limit memory
225try:
226 import resource
227 resource.setrlimit(resource.RLIMIT_AS, ({self.max_memory_mb} * 1024 * 1024, {self.max_memory_mb} * 1024 * 1024))
228except Exception:
229 pass
231# Setup code
232{setup_code}
234# User code
235{code}
237# Test runner
238test_cases = {{
239{tc_dict}
240}}
242func_name = "{func_name}"
243has_func = func_name and func_name in dir()
245results = []
246total = 0
247passed = 0
249if not has_func and len(test_cases) == 0:
250 print("__SANDBOX_OK__")
251 print(json.dumps({{"status": "executed", "message": "Code executed, no tests"}}))
252 sys.exit(0)
254if not has_func and len(test_cases) > 0:
255 print("__SANDBOX_ERR__")
256 print(json.dumps({{"error": f"Function '{{func_name}}' not found in code"}}))
257 sys.exit(1)
259func = eval(func_name)
260all_test_order = [{tc_list}]
262for tc_name in all_test_order:
263 tc = test_cases[tc_name]
264 total += 1
265 start = time.time()
266 try:
267 output = func(*tc["args"], **tc["kwargs"])
268 elapsed = time.time() - start
270 expected = tc.get("expected")
271 expected_type = tc.get("expected_type")
272 expected_exc = tc.get("expected_exception")
274 if expected_exc:
275 passed_ = False
276 detail = f"Expected exception {{expected_exc}} but none raised"
277 elif expected_type and not isinstance(output, eval(expected_type)):
278 passed_ = False
279 detail = f"Type mismatch: got {{type(output).__name__}}, expected {{expected_type}}"
280 elif expected is not None and output != expected:
281 passed_ = False
282 detail = f"Expected {{repr(expected)}}, got {{repr(output)}}"
283 else:
284 passed_ = True
285 detail = "OK"
287 if passed_:
288 passed += 1
290 results.append({{
291 "name": tc_name,
292 "passed": passed_,
293 "output": repr(output)[:200],
294 "expected": repr(expected)[:200],
295 "detail": detail,
296 "duration_ms": round(elapsed * 1000, 2),
297 }})
298 except Exception as e:
299 exc_name = type(e).__name__
300 expected_exc = tc.get("expected_exception", "")
302 if expected_exc and exc_name == expected_exc:
303 passed_ = True
304 detail = f"Expected exception {{exc_name}} raised"
305 passed += 1
306 else:
307 passed_ = False
308 detail = f"{{exc_name}}: {{str(e)[:200]}}"
310 results.append({{
311 "name": tc_name,
312 "passed": passed_,
313 "output": "",
314 "expected": repr(tc.get("expected"))[:200],
315 "detail": detail,
316 "duration_ms": 0,
317 }})
319print("__SANDBOX_RESULTS__")
320print(json.dumps({{
321 "status": "complete",
322 "total": total,
323 "passed": passed,
324 "failed": total - passed,
325 "all_passed": total > 0 and passed == total,
326 "test_results": results,
327}}))
328'''
330 def _execute_script(self, script_path: str) -> SandboxResult:
331 """Execute the sandbox script as subprocess."""
332 try:
333 start = time_module()
334 proc = subprocess.run(
335 ["python3", script_path],
336 capture_output=True,
337 text=True,
338 timeout=self.timeout,
339 cwd="/tmp",
340 env={
341 **os.environ,
342 "PYTHONDONTWRITEBYTECODE": "1",
343 "PYTHONPATH": "",
344 "SANDBOX_MODE": "1",
345 },
346 )
347 duration = time_module() - start
348 except subprocess.TimeoutExpired as e:
349 return SandboxResult(
350 success=False,
351 exit_code=-1,
352 stderr=str(e.stdout or "") if e.stdout else "",
353 exception=f"Timeout after {self.timeout}s",
354 duration=self.timeout,
355 )
356 except Exception as e:
357 return SandboxResult(
358 success=False,
359 exception=str(e),
360 stderr=traceback.format_exc(),
361 )
363 result = SandboxResult(
364 exit_code=proc.returncode,
365 stdout=proc.stdout,
366 stderr=proc.stderr,
367 duration=duration,
368 )
370 # Parse test results
371 if "__SANDBOX_RESULTS__" in proc.stdout:
372 try:
373 lines = proc.stdout.split("\n")
374 json_start = False
375 json_text = ""
376 for line in lines:
377 if json_start:
378 json_text += line
379 try:
380 data = json.loads(json_text)
381 break
382 except json.JSONDecodeError:
383 continue
384 if "__SANDBOX_RESULTS__" in line:
385 json_start = True
387 if data:
388 result.test_results = data.get("test_results", [])
389 result.all_passed = data.get("all_passed", False)
390 passed = data.get("passed", 0)
391 total = data.get("total", 0)
392 result.success = total > 0 and passed == total
393 except Exception:
394 pass
396 elif proc.returncode == 0 and not proc.stderr:
397 result.success = True
399 # Extract meaningful error info for feedback
400 if not result.success and result.stderr:
401 lines = result.stderr.strip().split("\n")
402 # Get last 3 lines (most relevant error info)
403 result.exception = "\n".join(lines[-5:]) if len(lines) > 5 else result.stderr
405 return result
408def time_module() -> float:
409 """Get current time in seconds."""
410 import time as _time
412 return _time.time()
415class CodeFeedbackExtractor:
416 """Extracts actionable feedback from sandbox failures for retry loops."""
418 ERROR_PATTERNS = {
419 "NameError": "Variable or function not defined. Check spelling and scope.",
420 "TypeError": "Wrong type passed to function. Check argument types.",
421 "ValueError": "Invalid value for operation. Check parameter constraints.",
422 "IndexError": "List index out of range. Check bounds.",
423 "KeyError": "Dictionary key not found. Check key existence.",
424 "AttributeError": "Object has no such attribute. Check method/attribute name.",
425 "ImportError": "Missing import. Add 'import X' or 'from X import Y'.",
426 "SyntaxError": "Python syntax error. Check indentation, brackets, colons.",
427 "ZeroDivisionError": "Division by zero. Add guard for zero denominator.",
428 "RecursionError": "Recursion depth exceeded. Add base case or switch to iteration.",
429 "TimeoutError": "Code timed out. Check for infinite loops or optimize.",
430 }
432 @classmethod
433 def extract(cls, sandbox_result: SandboxResult) -> list[str]:
434 """Extract actionable feedback suggestions from sandbox result."""
435 suggestions: list[str] = []
437 # Check for security issues
438 if "SECURITY" in sandbox_result.stderr:
439 suggestions.append("Code violates sandbox security rules; avoid system calls and I/O.")
440 return suggestions
442 # Check test failures
443 for tc in sandbox_result.test_results:
444 if not tc.get("passed", True):
445 detail = tc.get("detail", "")
446 name = tc.get("name", "unknown")
447 suggestions.append(f"Test '{name}' failed: {detail}")
449 # Check error patterns
450 for error_type, suggestion in cls.ERROR_PATTERNS.items():
451 if error_type in sandbox_result.stderr or error_type in sandbox_result.exception:
452 if suggestion not in suggestions:
453 suggestions.append(suggestion)
455 # Add specific output mismatch advice
456 if not sandbox_result.success and not suggestions:
457 suggestions.append("Code failed to execute. Check logic and edge cases.")
458 if sandbox_result.stderr:
459 suggestions.append(f"Error: {sandbox_result.stderr.strip()[:200]}")
461 return suggestions