Update app.py

app.py

@@ -5,10 +5,15 @@ import resource
 import multiprocessing
 import signal
 import time
+import threading
+import shutil
 import gradio as gr
 
+# -------------------------
 # Store running sandboxes
+# -------------------------
 sandboxes = {}
+SANDBOX_TTL = 2 * 3600  # 2 hours
 
 # -------------------------
 # Sandbox worker function
@@ -33,7 +38,7 @@ def sandbox_worker(code: str, sandbox_id: str):
             ["python3", code_file],
             capture_output=True,
             text=True,
-            timeout=7200  # 2h wall-clock
+            timeout=SANDBOX_TTL
         )
         # Save logs
         with open(os.path.join(user_dir, "stdout.log"), "w") as f:
@@ -44,6 +49,26 @@ def sandbox_worker(code: str, sandbox_id: str):
         with open(os.path.join(user_dir, "stderr.log"), "w") as f:
             f.write("Execution timed out (2h limit reached).")
 
+# -------------------------
+# Cleanup thread
+# -------------------------
+def cleanup_sandboxes():
+    while True:
+        now = time.time()
+        for sid, info in list(sandboxes.items()):
+            if now - info["start_time"] > SANDBOX_TTL:
+                try:
+                    os.kill(info["pid"], signal.SIGKILL)
+                except:
+                    pass
+                folder = f"/tmp/sandbox_{sid}"
+                if os.path.exists(folder):
+                    shutil.rmtree(folder)
+                del sandboxes[sid]
+        time.sleep(60)  # check every minute
+
+threading.Thread(target=cleanup_sandboxes, daemon=True).start()
+
 # -------------------------
 # Gradio interface functions
 # -------------------------
@@ -70,6 +95,9 @@ def kill_sandbox(sandbox_id):
     pid = sandboxes[sandbox_id]["pid"]
     try:
         os.kill(pid, signal.SIGKILL)
+        folder = f"/tmp/sandbox_{sandbox_id}"
+        if os.path.exists(folder):
+            shutil.rmtree(folder)
         del sandboxes[sandbox_id]
         return f"Sandbox {sandbox_id} killed."
     except Exception as e:
@@ -85,27 +113,23 @@ def status_sandbox(sandbox_id):
 # Gradio UI
 # -------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# Sandbox Backend")
-    
+    gr.Markdown("# Sandbox Backend with Auto-Cleanup")
+
     code_input = gr.Textbox(label="Python Code", lines=15, placeholder="Write your code here...")
     launch_btn = gr.Button("Launch Sandbox")
     sandbox_id_output = gr.Textbox(label="Sandbox ID")
-    
+
     logs_btn = gr.Button("Fetch Logs")
     logs_output = gr.Textbox(label="Logs", lines=10)
-    
+
     kill_btn = gr.Button("Kill Sandbox")
     kill_output = gr.Textbox(label="Kill Status")
-    
+
     status_btn = gr.Button("Check Status")
     status_output = gr.Textbox(label="Sandbox Status")
-    
-    # Temporary storage for sandbox ID
+
     sandbox_id_storage = gr.State(value=None)
-    
-    # -------------------------
-    # Button actions
-    # -------------------------
+
     launch_btn.click(
         launch_sandbox,
         inputs=code_input,
@@ -127,7 +151,4 @@ with gr.Blocks() as demo:
         outputs=status_output
     )
 
-# -------------------------
-# Launch Gradio app
-# -------------------------
 demo.launch()