fix(executor): Fix the issue of cleaning up sessions (#186)

Author: Micro66

executor/agents/agno/agno_agent.py

@@ -337,7 +337,6 @@ async def _async_execute(self) -> TaskStatus:
             # Checkpoint 1: Check cancellation before execution starts
             if self.task_state_manager.is_cancelled(self.task_id):
                 logger.info(f"Task {self.task_id} cancelled before execution")
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             progress = 65
@@ -379,7 +378,6 @@ async def _async_execute(self) -> TaskStatus:
             # Checkpoint 2: Check cancellation after team/agent creation
             if self.task_state_manager.is_cancelled(self.task_id):
                 logger.info(f"Task {self.task_id} cancelled after team/agent creation")
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             # Prepare prompt
@@ -701,7 +699,6 @@ async def _run_agent_streaming_async(self, prompt: str) -> TaskStatus:
                 # Checkpoint: Check cancellation during streaming
                 if self.task_state_manager.is_cancelled(self.task_id):
                     logger.info(f"Task {self.task_id} cancelled during agent streaming")
-                    await self._cleanup_resources()
                     return TaskStatus.COMPLETED
 
                 result_content = await self._handle_agent_streaming_event(
@@ -710,7 +707,6 @@ async def _run_agent_streaming_async(self, prompt: str) -> TaskStatus:
 
             # Check if task was cancelled
             if self.task_state_manager.is_cancelled(self.task_id):
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             return self._handle_execution_result(result_content, "agent streaming execution")
@@ -813,7 +809,6 @@ async def _run_team_streaming_async(self, prompt: str) -> TaskStatus:
                 # Checkpoint: Check cancellation during streaming
                 if self.task_state_manager.is_cancelled(self.task_id):
                     logger.info(f"Task {self.task_id} cancelled during team streaming")
-                    await self._cleanup_resources()
                     return TaskStatus.COMPLETED
 
                 result_content, reasoning = await self._handle_team_streaming_event(
@@ -824,7 +819,6 @@ async def _run_team_streaming_async(self, prompt: str) -> TaskStatus:
 
             # Check if task was cancelled
             if self.task_state_manager.is_cancelled(self.task_id):
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             return self._handle_execution_result(result_content, "team streaming execution")
@@ -1118,16 +1112,6 @@ def cancel_run(self) -> bool:
             self.task_state_manager.set_state(self.task_id, TaskState.CANCELLED)
             return False
 
-    async def _cleanup_resources(self) -> None:
-        """
-        Clean up resources when task is cancelled
-        """
-        try:
-            logger.info(f"Cleaning up resources for task {self.task_id}")
-            await self.resource_manager.cleanup_task_resources(self.task_id)
-        except Exception as e:
-            logger.exception(f"Error cleaning up resources for task {self.task_id}: {str(e)}")
-
     async def cleanup(self) -> None:
         """
         Clean up resources used by the agent

executor/agents/claude_code/claude_code_agent.py

@@ -693,7 +693,6 @@ async def _async_execute(self) -> TaskStatus:
             # Check if task was cancelled before execution
             if self.task_state_manager.is_cancelled(self.task_id):
                 logger.info(f"Task {self.task_id} was cancelled before execution")
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             progress = 65
@@ -768,14 +767,12 @@ async def _async_execute(self) -> TaskStatus:
                 self.resource_manager.register_resource(
                     task_id=self.task_id,
                     resource_id=f"claude_client_{self.session_id}",
-                    cleanup_func=self._cleanup_client,
                     is_async=True
                 )
 
             # Check cancellation again before proceeding
             if self.task_state_manager.is_cancelled(self.task_id):
                 logger.info(f"Task {self.task_id} cancelled during client setup")
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             # Prepare prompt
@@ -790,7 +787,6 @@ async def _async_execute(self) -> TaskStatus:
             # Check cancellation before sending query
             if self.task_state_manager.is_cancelled(self.task_id):
                 logger.info(f"Task {self.task_id} cancelled before sending query")
-                await self._cleanup_resources()
                 return TaskStatus.COMPLETED
 
             # Use session_id to send messages, ensuring messages are in the same session
@@ -1012,32 +1008,7 @@ async def _async_cancel_run(self) -> None:
         except Exception as e:
             logger.exception(f"Error during async interrupt for session_id {self.session_id}: {str(e)}")
 
-    async def _cleanup_client(self) -> None:
-        """Clean up client connection and remove from cache"""
-        try:
-            if self.client and hasattr(self.client, 'close'):
-                await self.client.close()
-                logger.info(f"Closed client for task {self.task_id}")
 
-            # Remove client from cache to prevent reuse of terminated client
-            if self.session_id in self._clients:
-                del self._clients[self.session_id]
-                logger.info(f"Removed client from cache for session_id: {self.session_id}")
-        except Exception as e:
-            logger.exception(f"Error closing client for task {self.task_id}: {e}")
-    
-    async def _cleanup_resources(self) -> None:
-        """Clean up all task resources"""
-        try:
-            # Use resource manager to clean up all registered resources
-            await self.resource_manager.cleanup_task_resources(self.task_id)
-            
-            # Clean up task state
-            self.task_state_manager.cleanup(self.task_id)
-            
-            logger.info(f"Completed resource cleanup for task {self.task_id}")
-        except Exception as e:
-            logger.exception(f"Error in resource cleanup for task {self.task_id}: {e}")
     def _setup_claudecode_dir(self, project_path: str, custom_rules: Dict[str, str]) -> None:
         """
         Setup .claudecode directory with custom instruction files for Claude Code compatibility

executor/tasks/resource_manager.py

@@ -24,9 +24,6 @@
 class ResourceHandle:
     """Resource handle"""
     resource_id: str
-    cleanup_func: Callable
-    cleanup_args: tuple = field(default_factory=tuple)
-    cleanup_kwargs: dict = field(default_factory=dict)
     is_async: bool = False
 
 
@@ -53,9 +50,6 @@ def register_resource(
         self, 
         task_id: int, 
         resource_id: str,
-        cleanup_func: Callable,
-        cleanup_args: tuple = (),
-        cleanup_kwargs: Optional[dict] = None,
         is_async: bool = False
     ) -> None:
         """
@@ -64,23 +58,15 @@ def register_resource(
         Args:
             task_id: Task ID
             resource_id: Unique resource identifier
-            cleanup_func: Cleanup function
-            cleanup_args: Positional arguments for cleanup function
-            cleanup_kwargs: Keyword arguments for cleanup function
             is_async: Whether cleanup function is asynchronous
         """
-        if cleanup_kwargs is None:
-            cleanup_kwargs = {}
 
         with self._resource_lock:
             if task_id not in self._resources:
                 self._resources[task_id] = []
 
             handle = ResourceHandle(
                 resource_id=resource_id,
-                cleanup_func=cleanup_func,
-                cleanup_args=cleanup_args,
-                cleanup_kwargs=cleanup_kwargs,
                 is_async=is_async
             )
             self._resources[task_id].append(handle)
@@ -104,55 +90,6 @@ def unregister_resource(self, task_id: int, resource_id: str) -> None:
                 if len(self._resources[task_id]) < original_count:
                     logger.debug(f"Unregistered resource '{resource_id}' for task {task_id}")
 
-    async def cleanup_task_resources(self, task_id: int) -> None:
-        """
-        Clean up all resources for a task (async version)
-        
-        Args:
-            task_id: Task ID
-        """
-        resources = []
-        with self._resource_lock:
-            resources = self._resources.pop(task_id, [])
-        
-        if not resources:
-            logger.debug(f"No resources to cleanup for task {task_id}")
-            return
-        
-        logger.info(f"Cleaning up {len(resources)} resources for task {task_id}")
-        
-        # Cleanup in reverse order, last registered cleaned up first
-        for handle in reversed(resources):
-            try:
-                if handle.is_async:
-                    await handle.cleanup_func(*handle.cleanup_args, **handle.cleanup_kwargs)
-                else:
-                    handle.cleanup_func(*handle.cleanup_args, **handle.cleanup_kwargs)
-                logger.debug(f"Cleaned up resource '{handle.resource_id}' for task {task_id}")
-            except Exception as e:
-                logger.exception(f"Error cleaning up resource '{handle.resource_id}' for task {task_id}: {e}")
-    
-    def cleanup_task_resources_sync(self, task_id: int) -> None:
-        """
-        Clean up all resources for a task (sync version)
-        
-        Args:
-            task_id: Task ID
-        """
-        try:
-            # Try to get current event loop
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                # If in async context, create task
-                asyncio.create_task(self.cleanup_task_resources(task_id))
-                logger.debug(f"Created async cleanup task for task {task_id}")
-            else:
-                # Otherwise run directly
-                loop.run_until_complete(self.cleanup_task_resources(task_id))
-        except RuntimeError:
-            # No event loop, create new one
-            asyncio.run(self.cleanup_task_resources(task_id))
-    
     def get_resource_count(self, task_id: int) -> int:
         """
         Get count of registered resources for a task