properly cleanup agents and conversation

Author: Shreyas-Microsoft

src/backend/sql_agents/convert_script.py

@@ -44,24 +44,24 @@ async def convert_script(
     logger.info("Migrating query: %s\n", source_script)
 
     # Setup the group chat for the agents
-    comms_manager = CommsManager(
+    async with CommsManager(
         sql_agents.idx_agents,
-        max_retries=5,          # Retry up to 5 times for rate limits
-        initial_delay=1.0,      # Start with 1 second delay
-        backoff_factor=2.0,     # Double delay each retry
-    )
+        max_retries=5,
+        initial_delay=1.0,
+        backoff_factor=2.0,
+    ) as comms_manager:
 
     # send websocket notification that file processing has started
-    send_status_update(
-        status=FileProcessUpdate(
-            file.batch_id,
-            file.file_id,
-            ProcessStatus.IN_PROGRESS,
-            AgentType.ALL,
-            "File processing started",
-            file_result=FileResult.INFO,
-        ),
-    )
+        send_status_update(
+            status=FileProcessUpdate(
+                file.batch_id,
+                file.file_id,
+                ProcessStatus.IN_PROGRESS,
+                AgentType.ALL,
+                "File processing started",
+                file_result=FileResult.INFO,
+            ),
+        )
 
     # orchestrate the chat
     current_migration = "No migration"

src/backend/sql_agents/helpers/comms_manager.py

@@ -156,113 +156,151 @@ def __init__(
 
     async def invoke_async(self):
         """Invoke the group chat with the given agents (original method maintained for compatibility)."""
-        return self.group_chat.invoke()
+        try:
+            return self.group_chat.invoke()
+        finally:
+            # Ensure cleanup happens even if invoke fails
+            await self.cleanup()
 
     async def async_invoke(self) -> AsyncIterable[ChatMessageContent]:
         """Invoke the group chat with retry logic and error handling."""
         attempt = 0
         current_delay = self.initial_delay
 
-        while attempt < self.max_retries:
-            try:
-                # Grab a snapshot of the history of the group chat
-                # Using "SHALLOW" copy to avoid getting a reference to the original list
-                history_snap = copy.copy(self.group_chat.history)
-
-                self.logger.debug(
-                    "History before invoke: %s",
-                    [msg.name for msg in self.group_chat.history],
-                )
-
-                # Get a fresh iterator from the function
-                async_iter = self.group_chat.invoke()
-
-                # If simple truncation is set, truncate the history
-                if (
-                    self.simple_truncation
-                    and len(self.group_chat.history) > self.simple_truncation
-                ):
-                    # Truncate the history to the last n messages
-                    self.group_chat.history = history_snap[-self.simple_truncation :]
-
-                # Yield each item from the iterator
-                async for item in async_iter:
-                    yield item
-
-                # If we get here without exception, we're done
-                break
-
-            except AgentInvokeException as aie:
-                attempt += 1
-                if attempt >= self.max_retries:
-                    self.logger.error(
-                        "Function invoke failed after %d attempts. Final error: %s. Consider increasing the models rate limit.",
-                        self.max_retries,
-                        str(aie),
+        try:
+            while attempt < self.max_retries:
+                try:
+                    # Grab a snapshot of the history of the group chat
+                    # Using "SHALLOW" copy to avoid getting a reference to the original list
+                    history_snap = copy.copy(self.group_chat.history)
+
+                    self.logger.debug(
+                        "History before invoke: %s",
+                        [msg.name for msg in self.group_chat.history],
                     )
-                    # Re-raise the last exception if all retries failed
-                    raise
 
-                # Return history state for retry
-                self.group_chat.history = history_snap
+                    # Get a fresh iterator from the function
+                    async_iter = self.group_chat.invoke()
 
-                try:
-                    # Try to extract wait time from error message
-                    wait_time_match = re.search(self._EXTRACT_WAIT_TIME, str(aie))
-                    if wait_time_match:
-                        # If regex is found, set the delay to the value in seconds
-                        current_delay = int(wait_time_match.group(1))
-                    else:
+                    # If simple truncation is set, truncate the history
+                    if (
+                        self.simple_truncation
+                        and len(self.group_chat.history) > self.simple_truncation
+                    ):
+                        # Truncate the history to the last n messages
+                        self.group_chat.history = history_snap[-self.simple_truncation :]
+
+                    # Yield each item from the iterator
+                    async for item in async_iter:
+                        yield item
+
+                    # If we get here without exception, we're done
+                    break
+
+                except AgentInvokeException as aie:
+                    attempt += 1
+                    if attempt >= self.max_retries:
+                        self.logger.error(
+                            "Function invoke failed after %d attempts. Final error: %s. Consider increasing the models rate limit.",
+                            self.max_retries,
+                            str(aie),
+                        )
+                        # Re-raise the last exception if all retries failed
+                        raise
+
+                    # Return history state for retry
+                    self.group_chat.history = history_snap
+
+                    try:
+                        # Try to extract wait time from error message
+                        wait_time_match = re.search(self._EXTRACT_WAIT_TIME, str(aie))
+                        if wait_time_match:
+                            # If regex is found, set the delay to the value in seconds
+                            current_delay = int(wait_time_match.group(1))
+                        else:
+                            current_delay = self.initial_delay
+
+                        self.logger.warning(
+                            "Attempt %d/%d for function invoke failed: %s. Retrying in %.2f seconds...",
+                            attempt,
+                            self.max_retries,
+                            str(aie),
+                            current_delay,
+                        )
+
+                        # Wait before retrying
+                        await asyncio.sleep(current_delay)
+
+                        if not wait_time_match:
+                            # Increase delay for next attempt using backoff factor
+                            current_delay *= self.backoff_factor
+
+                    except Exception as ex:
+                        self.logger.error(
+                            "Retry error: %s. Using default delay.",
+                            ex,
+                        )
                         current_delay = self.initial_delay
 
+                except self.exception_types as e:
+                    attempt += 1
+                    if attempt >= self.max_retries:
+                        self.logger.error(
+                            "Function invoke failed after %d attempts. Final error: %s",
+                            self.max_retries,
+                            str(e),
+                        )
+                        raise
+
                     self.logger.warning(
-                        "Attempt %d/%d for function invoke failed: %s. Retrying in %.2f seconds...",
+                        "Attempt %d/%d failed with %s: %s. Retrying in %.2f seconds...",
                         attempt,
                         self.max_retries,
-                        str(aie),
+                        type(e).__name__,
+                        str(e),
                         current_delay,
                     )
 
-                    # Wait before retrying
                     await asyncio.sleep(current_delay)
+                    current_delay *= self.backoff_factor
 
-                    if not wait_time_match:
-                        # Increase delay for next attempt using backoff factor
-                        current_delay *= self.backoff_factor
-
-                except Exception as ex:
-                    self.logger.error(
-                        "Retry error: %s. Using default delay.",
-                        ex,
-                    )
-                    current_delay = self.initial_delay
+        finally:
+            # Ensure cleanup happens regardless of success or failure
+            await self.cleanup()
 
-            except self.exception_types as e:
-                attempt += 1
-                if attempt >= self.max_retries:
-                    self.logger.error(
-                        "Function invoke failed after %d attempts. Final error: %s",
-                        self.max_retries,
-                        str(e),
-                    )
-                    raise
-
-                self.logger.warning(
-                    "Attempt %d/%d failed with %s: %s. Retrying in %.2f seconds...",
-                    attempt,
-                    self.max_retries,
-                    type(e).__name__,
-                    str(e),
-                    current_delay,
-                )
-
-                await asyncio.sleep(current_delay)
-                current_delay *= self.backoff_factor
-
-    async def close(self):
-        """Clean up the thread after agent execution is complete."""
+    async def cleanup(self):
+        """Clean up all resources including internal threads."""
+        try:
+            if self.group_chat is not None:
+                self.logger.debug("Cleaning up AgentGroupChat resources...")
+                
+                # Reset the group chat - this clears conversation state and deletes remote threads
+                await self.group_chat.reset()
+                
+                self.logger.debug("AgentGroupChat cleanup completed successfully")
+                
+        except Exception as e:
+            self.logger.error("Error during cleanup: %s", str(e))
+            # Don't re-raise cleanup errors to avoid masking original exceptions
+
+    async def __aenter__(self):
+        """Context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit with cleanup."""
+        await self.cleanup()
+
+    def __del__(self):
+        """Destructor to ensure cleanup if not explicitly called."""
         try:
-            await self.group_chat.delete()
-            self.logger.info("Deleted group chat thread successfully.")
-        except Exception as ex:
-            self.logger.warning("Failed to delete group chat thread: %s", str(ex))
+            # Only attempt cleanup if there's an active event loop
+            loop = asyncio.get_running_loop()
+            if loop and not loop.is_closed():
+                # Schedule cleanup as a task
+                loop.create_task(self.cleanup())
+        except RuntimeError:
+            # No event loop running, can't clean up asynchronously
+            self.logger.warning("No event loop available for cleanup in destructor")
+        except Exception as e:
+            self.logger.error("Error in destructor cleanup: %s", str(e))