From 47cf1be181620c3c68aeffdb4d2a79e5ad5f493f Mon Sep 17 00:00:00 2001
From: a7m-1st <Ahmed.jimi.awelkeir500@gmail.com>
Date: Fri, 16 Jan 2026 17:00:21 +0300
Subject: [PATCH] enhance: cleanup env on close session

---
 backend/app/service/task.py                   | 27 ++++++++++++
 backend/app/utils/toolkit/terminal_toolkit.py | 43 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/backend/app/service/task.py b/backend/app/service/task.py
index 7f2e379d..b3f509dc 100644
--- a/backend/app/service/task.py
+++ b/backend/app/service/task.py
@@ -270,6 +270,8 @@ class TaskLock:
     last_accessed: datetime
     background_tasks: set[asyncio.Task]
     """Track all background tasks for cleanup"""
+    registered_toolkits: list[Any]
+    """Track toolkits for cleanup (e.g., TerminalToolkit venvs)"""
 
     # Context management fields
     conversation_history: List[Dict[str, Any]]
@@ -290,6 +292,7 @@ class TaskLock:
         self.created_at = datetime.now()
         self.last_accessed = datetime.now()
         self.background_tasks = set()
+        self.registered_toolkits = []
 
         # Initialize context management fields
         self.conversation_history = []
@@ -339,8 +342,32 @@ class TaskLock:
                 except asyncio.CancelledError:
                     pass
         self.background_tasks.clear()
+        
+        # Clean up registered toolkits (e.g., remove TerminalToolkit venvs)
+        for toolkit in self.registered_toolkits:
+            try:
+                if hasattr(toolkit, 'cleanup'):
+                    toolkit.cleanup()
+                    logger.info("Toolkit cleanup completed", extra={"task_id": self.id, "toolkit": type(toolkit).__name__})
+            except Exception as e:
+                logger.warning(f"Failed to cleanup toolkit: {e}", extra={"task_id": self.id, "toolkit": type(toolkit).__name__})
+        self.registered_toolkits.clear()
+        
         logger.info("Task lock cleanup completed", extra={"task_id": self.id})
 
+    def register_toolkit(self, toolkit: Any) -> None:
+        """Register a toolkit for cleanup when task ends.
+        
+        This is used to track toolkits that create resources (like venvs) that
+        should be cleaned up when the task is complete.
+        """
+        self.registered_toolkits.append(toolkit)
+        logger.debug("Toolkit registered for cleanup", extra={
+            "task_id": self.id, 
+            "toolkit": type(toolkit).__name__,
+            "total_registered": len(self.registered_toolkits)
+        })
+
     def add_conversation(self, role: str, content: str | dict):
         """Add a conversation entry to history"""
         logger.debug("Adding conversation entry", extra={"task_id": self.id, "role": role, "content_length": len(str(content))})
diff --git a/backend/app/utils/toolkit/terminal_toolkit.py b/backend/app/utils/toolkit/terminal_toolkit.py
index 7eff1ccd..2be82317 100644
--- a/backend/app/utils/toolkit/terminal_toolkit.py
+++ b/backend/app/utils/toolkit/terminal_toolkit.py
@@ -73,6 +73,16 @@ class TerminalToolkit(BaseTerminalToolkit, AbstractToolkit):
                 "openpyxl",
             ],
         )
+        
+        # Auto-register with TaskLock for cleanup when task ends
+        from app.service.task import get_task_lock_if_exists
+        task_lock = get_task_lock_if_exists(api_task_id)
+        if task_lock:
+            task_lock.register_toolkit(self)
+            logger.info("TerminalToolkit registered for cleanup", extra={
+                "api_task_id": api_task_id,
+                "working_directory": working_directory
+            })
 
     def _write_to_log(self, log_file: str, content: str) -> None:
         r"""Write content to log file with optional ANSI stripping.
@@ -175,6 +185,39 @@ class TerminalToolkit(BaseTerminalToolkit, AbstractToolkit):
 
         return result
 
+    def cleanup(self, remove_venv: bool = True):
+        """Clean up all active sessions and optionally remove the virtual environment.
+        
+        Args:
+            remove_venv: If True, removes the .venv or .initial_env folder created
+                        by this toolkit. Defaults to True to prevent disk bloat.
+        """
+        # First call parent cleanup to kill all shell sessions
+        super().cleanup()
+        
+        if remove_venv:
+            import shutil
+            
+            # Remove cloned env (.venv) if it exists
+            if self.cloned_env_path and os.path.exists(self.cloned_env_path):
+                try:
+                    shutil.rmtree(self.cloned_env_path)
+                    logger.info(f"Removed cloned venv: {self.cloned_env_path}", extra={
+                        "api_task_id": self.api_task_id
+                    })
+                except Exception as e:
+                    logger.warning(f"Failed to remove cloned venv {self.cloned_env_path}: {e}")
+            
+            # Remove initial env (.initial_env) if it exists
+            if self.initial_env_path and os.path.exists(self.initial_env_path):
+                try:
+                    shutil.rmtree(self.initial_env_path)
+                    logger.info(f"Removed initial env: {self.initial_env_path}", extra={
+                        "api_task_id": self.api_task_id
+                    })
+                except Exception as e:
+                    logger.warning(f"Failed to remove initial env {self.initial_env_path}: {e}")
+
     @classmethod
     def shutdown(cls):
         if cls._thread_pool: