mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-07-29 10:43:34 +00:00
21c62f6716
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.
115 lines
3.1 KiB
Python
115 lines
3.1 KiB
Python
import asyncio
|
|
import threading
|
|
import uuid
|
|
import weakref
|
|
|
|
import pytest
|
|
|
|
from helpers.defer import DeferredTask
|
|
|
|
|
|
class Owner:
|
|
pass
|
|
|
|
|
|
def make_task() -> DeferredTask:
|
|
return DeferredTask(f"defer-lifecycle-{uuid.uuid4()}")
|
|
|
|
|
|
def test_completed_task_releases_call_references_and_children():
|
|
task = make_task()
|
|
owner = Owner()
|
|
owner_ref = weakref.ref(owner)
|
|
child_killed = threading.Event()
|
|
|
|
class Child:
|
|
def kill(self, terminate_thread: bool = False) -> None:
|
|
assert terminate_thread
|
|
child_killed.set()
|
|
|
|
async def run(captured_owner):
|
|
return "done"
|
|
|
|
try:
|
|
task.add_child_task(Child(), terminate_thread=True) # type: ignore[arg-type]
|
|
task.start_task(run, owner)
|
|
assert task.result_sync(timeout=2) == "done"
|
|
assert child_killed.wait(2)
|
|
assert task.func is None
|
|
assert task.args == ()
|
|
assert task.kwargs == {}
|
|
|
|
del owner
|
|
assert owner_ref() is None
|
|
assert task.result_sync(timeout=2) == "done"
|
|
with pytest.raises(RuntimeError, match="Completed task cannot be restarted"):
|
|
task.restart()
|
|
finally:
|
|
task.kill(terminate_thread=True)
|
|
|
|
|
|
def test_kill_clears_stored_call_without_clearing_running_arguments():
|
|
task = make_task()
|
|
owner = Owner()
|
|
owner_ref = weakref.ref(owner)
|
|
started = threading.Event()
|
|
cancelled = threading.Event()
|
|
finished = threading.Event()
|
|
release: list[asyncio.Event] = []
|
|
|
|
async def run(captured_owner):
|
|
release.append(asyncio.Event())
|
|
started.set()
|
|
try:
|
|
await asyncio.Future()
|
|
except asyncio.CancelledError:
|
|
cancelled.set()
|
|
await release[0].wait()
|
|
finally:
|
|
finished.set()
|
|
|
|
try:
|
|
task.start_task(run, owner)
|
|
assert started.wait(2)
|
|
task.kill()
|
|
assert cancelled.wait(2)
|
|
assert task.func is None
|
|
assert task.args == ()
|
|
assert task.kwargs == {}
|
|
|
|
del owner
|
|
assert owner_ref() is not None
|
|
task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
|
|
assert finished.wait(2)
|
|
asyncio.run_coroutine_threadsafe(
|
|
asyncio.sleep(0), task.event_loop_thread.loop
|
|
).result(2)
|
|
assert owner_ref() is None
|
|
finally:
|
|
if release and task.event_loop_thread.loop:
|
|
task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
|
|
task.kill(terminate_thread=True)
|
|
|
|
|
|
def test_active_task_can_restart_from_its_snapshot():
|
|
task = make_task()
|
|
starts = [threading.Event(), threading.Event()]
|
|
run_count = 0
|
|
|
|
async def run(value):
|
|
nonlocal run_count
|
|
current_run = run_count
|
|
run_count += 1
|
|
assert value == "argument"
|
|
starts[current_run].set()
|
|
await asyncio.Future()
|
|
|
|
try:
|
|
task.start_task(run, "argument")
|
|
assert starts[0].wait(2)
|
|
task.restart()
|
|
assert starts[1].wait(2)
|
|
assert task.func is run
|
|
assert task.args == ("argument",)
|
|
finally:
|
|
task.kill(terminate_thread=True)
|