Fix: proper task cancellation in scheduler, leakage in defer.py
* **Mechanism**: `EventLoopThread.terminate()` now correctly stops the asyncio loop and joins the thread, removing it from the global registry. * **Cleanup**: `DeferredTask.kill(terminate_thread=True)` now invokes `_drain_event_loop_tasks()`, which runs **inside** the target thread to explicitly cancel and await all pending tasks (including monologue loops) before killing the thread. This prevents "Task was destroyed but it is pending" warnings and ensures clean exits. * **Tracking**: The scheduler now maintains a live registry (`_running_deferred_tasks`) of active `DeferredTask` objects, protected by a reentrant lock. * **State Management**: The `run_task` wrapper uses `asyncio.shield` to ensure that even when a task is cancelled (e.g., by user action), the task state is reliably reset to `IDLE` in the database, preventing tasks from getting stuck in `RUNNING` state. The fix is correctly propagated to all relevant destruction points using `terminate_thread=True`: * **Dedicated Context**: `scheduler_task_delete.py` cancels the specific running task and terminates its thread. * **Shared/Dedicated Context**: Both `chat_remove.py` (Delete Chat) and `chat_reset.py` (Reset Chat) now call `scheduler.cancel_tasks_by_context(...)`. This ensures that if a scheduler task is running in a chat window (monologue), resetting that chat immediately kills the background thread and stops the agent loop.
Rafael Uzarowski committed
Nov 19, 2025 at 12:50 UTC
784fe5589a769ebc61a7a8308dc6d42680c7a068
5 files changed
+127
-33
python/api/chat_remove.py
+3
-1
@@ -8,6 +8,9 @@ class RemoveChat(ApiHandler):
8
async def process(self, input: Input, request: Request) -> Output:
9
ctxid = input.get("context", "")
10
11
+ scheduler = TaskScheduler.get()
12
+ scheduler.cancel_tasks_by_context(ctxid, terminate_thread=True)
13
+
14
context = AgentContext.use(ctxid)
15
if context:
16
# stop processing any tasks
@@ -16,7 +19,6 @@ class RemoveChat(ApiHandler):
19
AgentContext.remove(ctxid)
20
persist_chat.remove_chat(ctxid)
21
19
- scheduler = TaskScheduler.get()
22
await scheduler.reload()
23
24
tasks = scheduler.get_tasks_by_context_id(ctxid)
python/api/chat_reset.py
+4
@@ -2,12 +2,16 @@ from python.helpers.api import ApiHandler, Input, Output, Request, Response
2
3
4
from python.helpers import persist_chat
5
+from python.helpers.task_scheduler import TaskScheduler
6
7
8
class Reset(ApiHandler):
9
async def process(self, input: Input, request: Request) -> Output:
10
ctxid = input.get("context", "")
11
12
+ # attempt to stop any scheduler tasks bound to this context
13
+ TaskScheduler.get().cancel_tasks_by_context(ctxid, terminate_thread=True)
14
+
15
# context instance - get or create
16
context = self.use_context(ctxid)
17
context.reset()
python/api/scheduler_task_delete.py
+1
@@ -34,6 +34,7 @@ class SchedulerTaskDelete(ApiHandler):
34
35
# If the task is running, update its state to IDLE first
36
if task.state == TaskState.RUNNING:
37
+ scheduler.cancel_running_task(task_id, terminate_thread=True)
38
if context:
39
context.reset()
40
# Update the state to IDLE so any ongoing processes know to terminate
python/helpers/defer.py
+57
-27
@@ -6,8 +6,9 @@ from typing import Any, Callable, Optional, Coroutine, TypeVar, Awaitable
6
7
T = TypeVar("T")
8
9
+
10
class EventLoopThread:
10
- _instances = {}
11
+ _instances: dict[str, "EventLoopThread"] = {}
12
_lock = threading.Lock()
13
14
def __init__(self, thread_name: str = "Background") -> None:
@@ -38,8 +39,29 @@ class EventLoopThread:
39
self.loop.run_forever()
40
41
def terminate(self):
41
- if self.loop and self.loop.is_running():
42
- self.loop.stop()
42
+ loop = getattr(self, "loop", None)
43
+ thread = getattr(self, "thread", None)
44
+
45
+ if not loop:
46
+ return
47
+
48
+ if loop.is_running():
49
+ if thread and thread is threading.current_thread():
50
+ loop.stop()
51
+ else:
52
+ loop.call_soon_threadsafe(loop.stop)
53
+ if thread:
54
+ thread.join()
55
+ elif thread and thread.is_alive() and thread is not threading.current_thread():
56
+ thread.join()
57
+
58
+ if not loop.is_closed():
59
+ loop.close()
60
+
61
+ with self.__class__._lock:
62
+ if self.thread_name in self.__class__._instances:
63
+ del self.__class__._instances[self.thread_name]
64
+
65
self.loop = None
66
self.thread = None
67
@@ -79,6 +101,12 @@ class DeferredTask:
101
102
def _start_task(self):
103
self._future = self.event_loop_thread.run_coroutine(self._run())
104
+ if self._future:
105
+ self._future.add_done_callback(self._on_task_done)
106
+
107
+ def _on_task_done(self, _future: Future):
108
+ # Ensure child background tasks are always cleaned up once the parent finishes
109
+ self.kill_children()
110
111
async def _run(self):
112
return await self.func(*self.args, **self.kwargs)
@@ -120,30 +148,16 @@ class DeferredTask:
148
if self._future and not self._future.done():
149
self._future.cancel()
150
123
- if (
124
- terminate_thread
125
- and self.event_loop_thread.loop
126
- and self.event_loop_thread.loop.is_running()
127
- ):
128
-
129
- def cleanup():
130
- tasks = [
131
- t
132
- for t in asyncio.all_tasks(self.event_loop_thread.loop)
133
- if t is not asyncio.current_task(self.event_loop_thread.loop)
134
- ]
135
- for task in tasks:
136
- task.cancel()
137
- try:
138
- # Give tasks a chance to cleanup
139
- if self.event_loop_thread.loop:
140
- self.event_loop_thread.loop.run_until_complete(
141
- asyncio.gather(task, return_exceptions=True)
142
- )
143
- except Exception:
144
- pass # Ignore cleanup errors
145
-
146
- self.event_loop_thread.loop.call_soon_threadsafe(cleanup)
151
+ if terminate_thread and self.event_loop_thread.loop:
152
+ if self.event_loop_thread.loop.is_running():
153
+ try:
154
+ cleanup_future = asyncio.run_coroutine_threadsafe(
155
+ self._drain_event_loop_tasks(), self.event_loop_thread.loop
156
+ )
157
+ cleanup_future.result()
158
+ except Exception:
159
+ pass
160
+
161
self.event_loop_thread.terminate()
162
163
def kill_children(self) -> None:
@@ -196,3 +210,19 @@ class DeferredTask:
210
211
asyncio.run_coroutine_threadsafe(wrapped(), self.event_loop_thread.loop)
212
return asyncio.wrap_future(future)
213
+
214
+ @staticmethod
215
+ async def _drain_event_loop_tasks():
216
+ """Cancel and await all pending tasks on the current event loop."""
217
+ loop = asyncio.get_running_loop()
218
+ current_task = asyncio.current_task(loop=loop)
219
+ pending = [
220
+ task
221
+ for task in asyncio.all_tasks(loop=loop)
222
+ if task is not current_task
223
+ ]
224
+ if not pending:
225
+ return
226
+ for task in pending:
227
+ task.cancel()
228
+ await asyncio.gather(*pending, return_exceptions=True)
python/helpers/task_scheduler.py
+62
-5
@@ -619,6 +619,8 @@ class TaskScheduler:
619
_tasks: SchedulerTaskList
620
_printer: PrintStyle
621
_instance = None
622
+ _running_deferred_tasks: Dict[str, DeferredTask]
623
+ _running_tasks_lock: threading.RLock
624
625
@classmethod
626
def get(cls) -> "TaskScheduler":
@@ -631,8 +633,38 @@ class TaskScheduler:
633
if not hasattr(self, '_initialized'):
634
self._tasks = SchedulerTaskList.get()
635
self._printer = PrintStyle(italic=True, font_color="green", padding=False)
636
+ self._running_deferred_tasks = {}
637
+ self._running_tasks_lock = threading.RLock()
638
self._initialized = True
639
640
+ def _register_running_task(self, task_uuid: str, deferred_task: DeferredTask) -> None:
641
+ with self._running_tasks_lock:
642
+ self._running_deferred_tasks[task_uuid] = deferred_task
643
+
644
+ def _unregister_running_task(self, task_uuid: str) -> None:
645
+ with self._running_tasks_lock:
646
+ self._running_deferred_tasks.pop(task_uuid, None)
647
+
648
+ def cancel_running_task(self, task_uuid: str, terminate_thread: bool = False) -> bool:
649
+ with self._running_tasks_lock:
650
+ deferred_task = self._running_deferred_tasks.get(task_uuid)
651
+ if not deferred_task:
652
+ return False
653
+ self._printer.print(f"Scheduler cancelling task {task_uuid}")
654
+ deferred_task.kill(terminate_thread=terminate_thread)
655
+ return True
656
+
657
+ def cancel_tasks_by_context(self, context_id: str, terminate_thread: bool = False) -> bool:
658
+ cancelled_any = False
659
+ with self._running_tasks_lock:
660
+ running_tasks = list(self._running_deferred_tasks.keys())
661
+ for task_uuid in running_tasks:
662
+ task = self.get_task_by_uuid(task_uuid)
663
+ if task and task.context_id == context_id:
664
+ if self.cancel_running_task(task_uuid, terminate_thread=terminate_thread):
665
+ cancelled_any = True
666
+ return cancelled_any
667
+
668
async def reload(self):
669
await self._tasks.reload()
670
@@ -774,19 +806,23 @@ class TaskScheduler:
806
task_snapshot: Union[ScheduledTask, AdHocTask, PlannedTask] | None = self.get_task_by_uuid(task_uuid)
807
if task_snapshot is None:
808
self._printer.print(f"Scheduler Task with UUID '{task_uuid}' not found")
809
+ self._unregister_running_task(task_uuid)
810
return
811
if task_snapshot.state == TaskState.RUNNING:
812
self._printer.print(f"Scheduler Task '{task_snapshot.name}' already running, skipping")
813
+ self._unregister_running_task(task_uuid)
814
return
815
816
# Atomically fetch and check the task's current state
817
current_task = await self.update_task_checked(task_uuid, lambda task: task.state != TaskState.RUNNING, state=TaskState.RUNNING)
818
if not current_task:
819
self._printer.print(f"Scheduler Task with UUID '{task_uuid}' not found or updated by another process")
820
+ self._unregister_running_task(task_uuid)
821
return
822
if current_task.state != TaskState.RUNNING:
823
# This means the update failed due to state conflict
824
self._printer.print(f"Scheduler Task '{current_task.name}' state is '{current_task.state}', skipping")
825
+ self._unregister_running_task(task_uuid)
826
return
827
828
await current_task.on_run()
@@ -868,6 +904,13 @@ class TaskScheduler:
904
self._printer.print(f"Fixing task state consistency: '{current_task.name}' state is not IDLE after success")
905
await self.update_task(task_uuid, state=TaskState.IDLE)
906
907
+ except asyncio.CancelledError:
908
+ self._printer.print(f"Scheduler Task '{current_task.name}' cancelled by user")
909
+ try:
910
+ await asyncio.shield(self.update_task(task_uuid, state=TaskState.IDLE))
911
+ except Exception:
912
+ pass
913
+ raise
914
except Exception as e:
915
# Error
916
self._printer.print(f"Scheduler Task '{current_task.name}' failed: {e}")
@@ -884,17 +927,31 @@ class TaskScheduler:
927
agent.handle_critical_exception(e)
928
finally:
929
# Call on_finish for task-specific cleanup
887
- await current_task.on_finish()
930
+ try:
931
+ await asyncio.shield(current_task.on_finish())
932
+ except asyncio.CancelledError:
933
+ pass
934
+ except Exception:
935
+ pass
936
937
# Make one final save to ensure all states are persisted
890
- await self._tasks.save()
938
+ try:
939
+ await asyncio.shield(self._tasks.save())
940
+ except asyncio.CancelledError:
941
+ pass
942
+ except Exception:
943
+ pass
944
+
945
+ self._unregister_running_task(task_uuid)
946
947
deferred_task = DeferredTask(thread_name=self.__class__.__name__)
948
+ self._register_running_task(task.uuid, deferred_task)
949
deferred_task.start_task(_run_task_wrapper, task.uuid, task_context)
950
895
- # Ensure background execution doesn't exit immediately on async await, especially in script contexts
896
- # This helps prevent premature exits when running from non-event-loop contexts
897
- asyncio.create_task(asyncio.sleep(0.1))
951
+ # Ensure background execution doesn't exit immediately on async await, especially in script contexts.
952
+ # Yielding briefly keeps callers like CLI scripts alive long enough for the DeferredTask thread to spin up
953
+ # without leaving stray pending tasks that trigger \"Task was destroyed\" warnings when the loop shuts down.
954
+ await asyncio.sleep(0.1)
955
956
def serialize_all_tasks(self) -> list[Dict[str, Any]]:
957
"""