Update app.py

app.py

@@ -7,8 +7,8 @@ app = FastAPI(title="Agent Sandbox API")
 def run_code(payload: dict = Body(...)):
     """
     Execute Python code in a sandboxed subprocess.
-    Input JSON: {"code": "print(2+2)"}
-    Output JSON: {"output": "4\n"} or {"error": "..."}
+    Input JSON: {"code": "print(2+2)"} or {"code": "2+2"}
+    Output JSON: {"output": "4"} or {"error": "..."}
     """
     code = payload.get("code", "")
     if not code:
@@ -24,9 +24,19 @@ def run_code(payload: dict = Body(...)):
         )
 
         if proc.returncode != 0:
-            return {"error": proc.stderr.strip()}
+            return {"error": proc.stderr.strip() or "Unknown error"}
 
-        return {"output": proc.stdout.strip()}
+        output = proc.stdout.strip()
+
+        # If nothing printed, try evaluating expression directly
+        if not output:
+            try:
+                result = eval(code, {})
+                output = str(result)
+            except Exception:
+                output = "<no output>"
+
+        return {"output": output}
 
     except subprocess.TimeoutExpired:
         return {"error": "Execution timed out"}