Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 import shutil
 import uuid
 from pathlib import Path
+import atexit
 
 # Generate a unique workspace for the user
 session_id = str(uuid.uuid4())
@@ -18,6 +19,12 @@ def run_shell_command(command):
     except subprocess.CalledProcessError as e:
         return f"Error: {e.output}"
 
+# Cleanup workspace when the script exits
+def cleanup_workspace():
+    if os.path.exists(workspace_dir):
+        shutil.rmtree(workspace_dir)
+atexit.register(cleanup_workspace)
+
 # Custom CSS for styling
 st.markdown(
     """
@@ -87,14 +94,6 @@ if st.button("Export Code"):
     else:
         st.error("No code to export!")
 
-# Cleanup workspace on session end
-def cleanup_workspace():
-    if os.path.exists(workspace_dir):
-        shutil.rmtree(workspace_dir)
-
-# Streamlit lifecycle hook for cleanup
-st.on_event("teardown", cleanup_workspace)
-
 st.markdown("---")
 st.info(
     """