update

Author: fengju0213

camel/societies/workforce/worker.py

@@ -45,6 +45,7 @@ def __init__(
     ) -> None:
         super().__init__(description, node_id=node_id)
         self._active_task_ids: Set[str] = set()
+        self._running_tasks: Set[asyncio.Task] = set()
 
     def __repr__(self):
         return f"Worker node {self.node_id} ({self.description})"
@@ -113,15 +114,12 @@ async def _listen_to_channel(self):
         self._running = True
         logger.info(f"{self} started.")
 
-        # Keep track of running task coroutines
-        running_tasks: Set[asyncio.Task] = set()
-
         while self._running:
             try:
                 # Clean up completed tasks
-                completed_tasks = [t for t in running_tasks if t.done()]
+                completed_tasks = [t for t in self._running_tasks if t.done()]
                 for completed_task in completed_tasks:
-                    running_tasks.remove(completed_task)
+                    self._running_tasks.discard(completed_task)
                     # Check for exceptions in completed tasks
                     try:
                         await completed_task
@@ -138,11 +136,11 @@ async def _listen_to_channel(self):
                     task_coroutine = asyncio.create_task(
                         self._process_single_task(task)
                     )
-                    running_tasks.add(task_coroutine)
+                    self._running_tasks.add(task_coroutine)
 
                 except asyncio.TimeoutError:
                     # No tasks available, continue loop
-                    if not running_tasks:
+                    if not self._running_tasks:
                         # No tasks running and none available, short sleep
                         await asyncio.sleep(0.1)
                     continue
@@ -155,12 +153,14 @@ async def _listen_to_channel(self):
                 continue
 
         # Wait for all remaining tasks to complete when stopping
-        if running_tasks:
+        if self._running_tasks:
             logger.info(
-                f"{self} stopping, waiting for {len(running_tasks)} "
+                f"{self} stopping, waiting for {len(self._running_tasks)} "
                 f"tasks to complete..."
             )
-            await asyncio.gather(*running_tasks, return_exceptions=True)
+            await asyncio.gather(
+                *self._running_tasks, return_exceptions=True
+            )
 
         logger.info(f"{self} stopped.")
 
@@ -173,4 +173,8 @@ async def start(self):
     def stop(self):
         r"""Stop the worker."""
         self._running = False
+        # Cancel any in-flight task coroutines
+        for task in list(self._running_tasks):
+            if not task.done():
+                task.cancel()
         return

camel/societies/workforce/workforce.py

@@ -1754,6 +1754,56 @@ def stop_gracefully(self) -> None:
                 f"(event-loop not yet started)."
             )
 
+    async def _async_stop_immediately(self) -> None:
+        r"""Force-stop the workforce without waiting for in-flight tasks."""
+        self._stop_requested = True
+        self._pause_event.set()
+        logger.warning(f"Workforce {self.node_id} force stop requested.")
+
+        # Remove pending tasks and clear channel postings to avoid new work
+        self._pending_tasks.clear()
+        if self._channel:
+            try:
+                in_flight = await self._channel.get_in_flight_tasks(
+                    self.node_id
+                )
+                for task in in_flight:
+                    await self._channel.remove_task(task.id)
+            except Exception as e:
+                logger.error(
+                    f"Failed to clear in-flight tasks during force stop: {e}",
+                    exc_info=True,
+                )
+        self._in_flight_tasks = 0
+
+        # Stop child nodes and cancel their listening tasks immediately
+        for child in self._children:
+            if child._running:
+                child.stop()
+        for task in self._child_listening_tasks:
+            if not task.done():
+                task.cancel()
+
+        self._running = False
+        self._state = WorkforceState.STOPPED
+
+    def stop_immediately(self) -> None:
+        r"""Force-stop without waiting for current tasks to finish."""
+        if self._loop and not self._loop.is_closed():
+            self._submit_coro_to_loop(self._async_stop_immediately())
+        else:
+            # No running loop; best-effort synchronous cleanup
+            self._stop_requested = True
+            self._pause_event.set()
+            self._pending_tasks.clear()
+            self._in_flight_tasks = 0
+            self._running = False
+            self._state = WorkforceState.STOPPED
+            logger.warning(
+                f"Workforce {self.node_id} force stopped "
+                f"(event-loop not yet started)."
+            )
+
     async def _async_skip_gracefully(self) -> None:
         r"""Async implementation of skip_gracefully to run on the event
         loop.
@@ -4990,54 +5040,34 @@ async def start(self) -> None:
 
     @check_if_running(True)
     def stop(self) -> None:
-        r"""Stop all the child nodes under it. The node itself will be stopped
-        by its parent node.
-        """
-        # Stop all child nodes first
-        for child in self._children:
-            if child._running:
-                child.stop()
+        r"""Forcefully stop the workforce and its children immediately.
 
-        # Cancel child listening tasks
-        if self._child_listening_tasks:
-            try:
-                loop = asyncio.get_running_loop()
-                if loop and not loop.is_closed():
-                    # Create graceful cleanup task
-                    async def cleanup():
-                        await asyncio.sleep(0.1)  # Brief grace period
-                        for task in self._child_listening_tasks:
-                            if not task.done():
-                                task.cancel()
-
-                        # Handle both asyncio.Task and concurrent.futures.
-                        # Future
-                        awaitables = []
-                        for task in self._child_listening_tasks:
-                            if isinstance(task, concurrent.futures.Future):
-                                # Convert Future to awaitable
-                                awaitables.append(asyncio.wrap_future(task))
-                            else:
-                                # Already an asyncio.Task
-                                awaitables.append(task)
-
-                        await asyncio.gather(
-                            *awaitables,
-                            return_exceptions=True,
-                        )
-
-                    self._cleanup_task = loop.create_task(cleanup())
-                else:
-                    # No active loop, cancel immediately
-                    for task in self._child_listening_tasks:
-                        task.cancel()
-            except (RuntimeError, Exception) as e:
-                # Fallback: cancel immediately
-                logger.debug(f"Exception during task cleanup: {e}")
-                for task in self._child_listening_tasks:
+        This is now an immediate stop (was previously a graceful lifecycle
+        cleanup). It cancels child listeners, clears pending/in-flight tasks,
+        and sets state to STOPPED without waiting for active work to finish.
+        """
+        if self._loop and not self._loop.is_closed():
+            self._submit_coro_to_loop(self._async_stop_immediately())
+        else:
+            # No running loop; perform synchronous best-effort force stop
+            self._stop_requested = True
+            self._pause_event.set()
+            self._pending_tasks.clear()
+            self._in_flight_tasks = 0
+            # Stop children
+            for child in self._children:
+                if child._running:
+                    child.stop()
+            # Cancel listening tasks
+            for task in self._child_listening_tasks:
+                if not task.done():
                     task.cancel()
-
-        self._running = False
+            self._running = False
+            self._state = WorkforceState.STOPPED
+            logger.warning(
+                f"Workforce {self.node_id} force stopped "
+                f"(event-loop not yet started)."
+            )
 
     def clone(self, with_memory: bool = False) -> 'Workforce':
         r"""Creates a new instance of Workforce with the same configuration.