fix: release deferred task resources
Clear DeferredTask callables and arguments after completion or cancellation while running each invocation from a safe private snapshot. Preserve result retrieval and active restart behavior, clean up child tasks, and add best-effort local and SSH session destructors with focused lifecycle regression coverage.
Alessandro committed
Jul 20, 2026 at 21:15 UTC
21c62f6716d0d6ed76e080c2e34e3d23d302a7d6
6 files changed
+159
-6
helpers/defer.py
+26
-6
@@ -88,6 +88,9 @@ class DeferredTask:
88
self.event_loop_thread = EventLoopThread(thread_name)
89
self._future: Optional[Future] = None
90
self.children: list[ChildTask] = []
91
+ self.func: Optional[Callable[..., Coroutine[Any, Any, Any]]] = None
92
+ self.args: tuple[Any, ...] = ()
93
+ self.kwargs: dict[str, Any] = {}
94
95
def start_task(
96
self, func: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any
@@ -102,16 +105,29 @@ class DeferredTask:
105
self.kill()
106
107
def _start_task(self):
105
- self._future = self.event_loop_thread.run_coroutine(self._run())
108
+ if self.func is None:
109
+ raise RuntimeError("Task callable is no longer available")
110
+
111
+ self._future = self.event_loop_thread.run_coroutine(
112
+ self._run(self.func, self.args, self.kwargs)
113
+ )
114
if self._future:
115
self._future.add_done_callback(self._on_task_done)
116
109
- def _on_task_done(self, _future: Future):
117
+ def _on_task_done(self, future: Future):
118
# Ensure child background tasks are always cleaned up once the parent finishes
111
- self.kill_children()
119
+ if future is self._future:
120
+ self.kill_children()
121
+ self._clear_call()
122
113
- async def _run(self):
114
- return await self.func(*self.args, **self.kwargs)
123
+ def _clear_call(self) -> None:
124
+ self.func = None
125
+ self.args = ()
126
+ self.kwargs = {}
127
+
128
+ @staticmethod
129
+ async def _run(func, args, kwargs):
130
+ return await func(*args, **kwargs)
131
132
def is_ready(self) -> bool:
133
return self._future.done() if self._future else False
@@ -149,6 +165,7 @@ class DeferredTask:
165
self.kill_children()
166
if self._future and not self._future.done():
167
self._future.cancel()
168
+ self._clear_call()
169
170
if terminate_thread and self.event_loop_thread.loop:
171
if self.event_loop_thread.loop.is_running():
@@ -171,8 +188,11 @@ class DeferredTask:
188
return self._future and not self._future.done() # type: ignore
189
190
def restart(self, terminate_thread: bool = False) -> None:
191
+ if self.func is None:
192
+ raise RuntimeError("Completed task cannot be restarted")
193
+ func, args, kwargs = self.func, self.args, self.kwargs
194
self.kill(terminate_thread=terminate_thread)
175
- self._start_task()
195
+ self.start_task(func, *args, **kwargs)
196
197
def add_child_task(
198
self, task: "DeferredTask", terminate_thread: bool = False
helpers/defer.py.dox.md
+2
@@ -29,6 +29,8 @@
29
## Runtime Contracts
30
31
- Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together.
32
+- `DeferredTask` retains its callable and arguments only while an invocation is active; completion and `kill()` clear those references after the running coroutine has taken its own snapshot.
33
+- Task results remain available after completion. `restart()` can restart an active invocation, but a completed invocation has no retained call recipe and must be started again explicitly.
34
- Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.
35
- Observed side-effect areas: scheduler state.
36
- Imported dependency areas include: `asyncio`, `concurrent.futures`, `dataclasses`, `threading`, `typing`.
plugins/_code_execution/AGENTS.md
+1
@@ -17,6 +17,7 @@
17
- Execute multi-line terminal input as one current-shell compound so intermediate prompts cannot mark queued work complete; preserve `cd`, exports, and other shell state.
18
- Treat local process exit and SSH channel termination as definitive command completion even when no final prompt is emitted; recreate terminated sessions before their next command.
19
- Terminal reset/close must not hang on foreground commands or shells that ignore SIGTERM.
20
+- Local and SSH session wrappers must synchronously release their owned process or connection resources when discarded.
21
- Explicitly target local versus SSH execution runtimes.
22
- Do not hardcode secrets, SSH credentials, or local user paths.
23
plugins/_code_execution/helpers/shell_local.py
+7
@@ -31,6 +31,13 @@ class LocalInteractiveSession:
31
self.full_output = ''
32
self.cwd = cwd
33
34
+ def __del__(self):
35
+ try:
36
+ if self.session:
37
+ self.session.kill()
38
+ except Exception:
39
+ pass
40
+
41
async def connect(self):
42
self.session = tty_session.TTYSession(
43
runtime.get_terminal_executable(),
plugins/_code_execution/helpers/shell_ssh.py
+8
@@ -36,6 +36,14 @@ class SSHInteractiveSession:
36
self.cwd = cwd
37
self._exit_code: int | None = None
38
39
+ def __del__(self):
40
+ for resource in (getattr(self, "shell", None), getattr(self, "client", None)):
41
+ try:
42
+ if resource:
43
+ resource.close()
44
+ except Exception:
45
+ pass
46
+
47
async def connect(self, keepalive_interval: int = 5):
48
"""
49
Establish the SSH connection and start an interactive shell.
tests/test_defer_lifecycle.py
new
+115
@@ -0,0 +1,115 @@
1
+import asyncio
2
+import threading
3
+import uuid
4
+import weakref
5
+
6
+import pytest
7
+
8
+from helpers.defer import DeferredTask
9
+
10
+
11
+class Owner:
12
+ pass
13
+
14
+
15
+def make_task() -> DeferredTask:
16
+ return DeferredTask(f"defer-lifecycle-{uuid.uuid4()}")
17
+
18
+
19
+def test_completed_task_releases_call_references_and_children():
20
+ task = make_task()
21
+ owner = Owner()
22
+ owner_ref = weakref.ref(owner)
23
+ child_killed = threading.Event()
24
+
25
+ class Child:
26
+ def kill(self, terminate_thread: bool = False) -> None:
27
+ assert terminate_thread
28
+ child_killed.set()
29
+
30
+ async def run(captured_owner):
31
+ return "done"
32
+
33
+ try:
34
+ task.add_child_task(Child(), terminate_thread=True) # type: ignore[arg-type]
35
+ task.start_task(run, owner)
36
+ assert task.result_sync(timeout=2) == "done"
37
+ assert child_killed.wait(2)
38
+ assert task.func is None
39
+ assert task.args == ()
40
+ assert task.kwargs == {}
41
+
42
+ del owner
43
+ assert owner_ref() is None
44
+ assert task.result_sync(timeout=2) == "done"
45
+ with pytest.raises(RuntimeError, match="Completed task cannot be restarted"):
46
+ task.restart()
47
+ finally:
48
+ task.kill(terminate_thread=True)
49
+
50
+
51
+def test_kill_clears_stored_call_without_clearing_running_arguments():
52
+ task = make_task()
53
+ owner = Owner()
54
+ owner_ref = weakref.ref(owner)
55
+ started = threading.Event()
56
+ cancelled = threading.Event()
57
+ finished = threading.Event()
58
+ release: list[asyncio.Event] = []
59
+
60
+ async def run(captured_owner):
61
+ release.append(asyncio.Event())
62
+ started.set()
63
+ try:
64
+ await asyncio.Future()
65
+ except asyncio.CancelledError:
66
+ cancelled.set()
67
+ await release[0].wait()
68
+ finally:
69
+ finished.set()
70
+
71
+ try:
72
+ task.start_task(run, owner)
73
+ assert started.wait(2)
74
+ task.kill()
75
+ assert cancelled.wait(2)
76
+ assert task.func is None
77
+ assert task.args == ()
78
+ assert task.kwargs == {}
79
+
80
+ del owner
81
+ assert owner_ref() is not None
82
+ task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
83
+ assert finished.wait(2)
84
+ asyncio.run_coroutine_threadsafe(
85
+ asyncio.sleep(0), task.event_loop_thread.loop
86
+ ).result(2)
87
+ assert owner_ref() is None
88
+ finally:
89
+ if release and task.event_loop_thread.loop:
90
+ task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
91
+ task.kill(terminate_thread=True)
92
+
93
+
94
+def test_active_task_can_restart_from_its_snapshot():
95
+ task = make_task()
96
+ starts = [threading.Event(), threading.Event()]
97
+ run_count = 0
98
+
99
+ async def run(value):
100
+ nonlocal run_count
101
+ current_run = run_count
102
+ run_count += 1
103
+ assert value == "argument"
104
+ starts[current_run].set()
105
+ await asyncio.Future()
106
+
107
+ try:
108
+ task.start_task(run, "argument")
109
+ assert starts[0].wait(2)
110
+ task.restart()
111
+ assert starts[1].wait(2)
112
+ assert task.func is run
113
+ assert task.args == ("argument",)
114
+ finally:
115
+ task.kill(terminate_thread=True)