khyari hamza
5d5d3fde2c
Add AsyncResourceTracker (core/async_cleanup.py) providing a central registry for async resources with timeout-bounded close_all(), async context manager support, and a __del__ finalizer that logs leaked resources by name. Enhance LangGraphBridge with cleanup_tasks_async() that awaits in-flight tasks with a deadline instead of fire-and-forget cancel(). Add cancellation_reasons dict to trace why tasks were cancelled. Add StateManager.close() to properly release checkpoint file handles and complete the RxPY BehaviorSubject. Add AcpEventQueue.close() to dispose all subscriptions. Includes 14 Behave scenarios (67 steps), Robot integration tests, ASV benchmarks, and docs/reference/async_safety.md. ISSUES CLOSED: #321
422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""Step definitions for async resource cleanup and leak prevention."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.acp.events import AcpEventQueue
|
|
from cleveragents.core.async_cleanup import AsyncResourceTracker
|
|
from cleveragents.langgraph.state import GraphState, StateManager
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers — lightweight mock async resource (lives in steps, not src/)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class MockAsyncResource:
|
|
"""Test double for an async-closable resource."""
|
|
|
|
def __init__(self, name: str, close_delay: float = 0.0) -> None:
|
|
self.name = name
|
|
self.close_delay = close_delay
|
|
self.closed = False
|
|
self.close_count = 0
|
|
|
|
async def close(self) -> None:
|
|
if self.close_delay > 0:
|
|
await asyncio.sleep(self.close_delay)
|
|
self.closed = True
|
|
self.close_count += 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Background
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given("I have an async resource tracker")
|
|
def step_create_tracker(context):
|
|
context.tracker = AsyncResourceTracker()
|
|
context.resources = {}
|
|
context.log_handler = _CapturingHandler()
|
|
logging.getLogger("cleveragents.core.async_cleanup").addHandler(context.log_handler)
|
|
context.warnings_logged = context.log_handler.records
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Givens
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@given('I have a mock async resource named "{name}"')
|
|
def step_create_mock_resource(context, name):
|
|
resource = MockAsyncResource(name)
|
|
context.resources[name] = resource
|
|
context.current_resource = resource
|
|
context.current_resource_name = name
|
|
|
|
|
|
@given(
|
|
'I have a mock async resource named "{name}" that takes {seconds:g} seconds to close'
|
|
)
|
|
def step_create_slow_resource(context, name, seconds):
|
|
resource = MockAsyncResource(name, close_delay=seconds)
|
|
context.resources[name] = resource
|
|
context.current_resource = resource
|
|
context.current_resource_name = name
|
|
|
|
|
|
@given("I have {count:d} mock async resources with varied close times")
|
|
def step_create_multiple_resources(context, count):
|
|
context.multi_resources = []
|
|
for i in range(count):
|
|
name = f"resource-{i}"
|
|
resource = MockAsyncResource(name, close_delay=0.01 * (i + 1))
|
|
context.resources[name] = resource
|
|
context.multi_resources.append((name, resource))
|
|
context.tracker.register(name, resource)
|
|
|
|
|
|
@given("I register the resource with the tracker")
|
|
def step_register_resource(context):
|
|
context.tracker.register(context.current_resource_name, context.current_resource)
|
|
|
|
|
|
@given("I have an async task tracked by the bridge")
|
|
def step_create_bridge_task(context):
|
|
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
|
|
from cleveragents.reactive.stream_router import ReactiveStreamRouter
|
|
|
|
router = ReactiveStreamRouter()
|
|
context.bridge = RxPyLangGraphBridge(router)
|
|
context.cancellation_reasons = {}
|
|
|
|
async def _long_running():
|
|
try:
|
|
await asyncio.sleep(100)
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
loop = asyncio.new_event_loop()
|
|
|
|
async def _setup():
|
|
task = asyncio.create_task(_long_running())
|
|
context.bridge._active_tasks.add(task)
|
|
context.bridge_task = task
|
|
|
|
loop.run_until_complete(_setup())
|
|
context.bridge_loop = loop
|
|
|
|
|
|
@given("I have a state manager with a temporary checkpoint directory")
|
|
def step_create_state_manager(context):
|
|
context.temp_dir = tempfile.mkdtemp()
|
|
context.checkpoint_dir = Path(context.temp_dir)
|
|
context.state_manager = StateManager(
|
|
initial_state=GraphState(),
|
|
checkpoint_dir=context.checkpoint_dir,
|
|
)
|
|
|
|
|
|
@given("I have an ACP event queue with {count:d} active subscriptions")
|
|
def step_create_event_queue(context, count):
|
|
context.event_queue = AcpEventQueue()
|
|
context.subscription_ids = []
|
|
for _ in range(count):
|
|
sub_id = context.event_queue.subscribe_local(lambda _evt: None)
|
|
context.subscription_ids.append(sub_id)
|
|
|
|
|
|
@given("I have a bridge with {count:d} slow async tasks")
|
|
def step_create_bridge_with_slow_tasks(context, count):
|
|
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
|
|
from cleveragents.reactive.stream_router import ReactiveStreamRouter
|
|
|
|
router = ReactiveStreamRouter()
|
|
context.bridge = RxPyLangGraphBridge(router)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
|
|
async def _setup():
|
|
for _ in range(count):
|
|
|
|
async def _slow():
|
|
try:
|
|
await asyncio.sleep(100)
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
task = asyncio.create_task(_slow())
|
|
context.bridge._active_tasks.add(task)
|
|
|
|
loop.run_until_complete(_setup())
|
|
context.bridge_loop = loop
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Whens
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@when("I register the resource with the tracker")
|
|
def step_when_register_resource(context):
|
|
context.tracker.register(context.current_resource_name, context.current_resource)
|
|
|
|
|
|
@when("I close all tracked resources")
|
|
def step_close_all(context):
|
|
loop = asyncio.new_event_loop()
|
|
loop.run_until_complete(context.tracker.close_all())
|
|
loop.close()
|
|
|
|
|
|
@when("I close all tracked resources with a {timeout:g} second timeout")
|
|
def step_close_all_with_timeout(context, timeout):
|
|
loop = asyncio.new_event_loop()
|
|
loop.run_until_complete(context.tracker.close_all(timeout=timeout))
|
|
loop.close()
|
|
|
|
|
|
@when("the tracker finalizer runs without close_all")
|
|
def step_run_finalizer(context):
|
|
context.tracker._warn_unclosed()
|
|
|
|
|
|
@when("I request graceful cleanup with a {timeout:g} second timeout")
|
|
def step_graceful_bridge_cleanup(context, timeout):
|
|
loop = context.bridge_loop
|
|
|
|
async def _cleanup():
|
|
await context.bridge.cleanup_tasks_async(timeout=timeout)
|
|
|
|
loop.run_until_complete(_cleanup())
|
|
|
|
|
|
@when('I cancel the task with reason "{reason}"')
|
|
def step_cancel_with_reason(context, reason):
|
|
loop = context.bridge_loop
|
|
|
|
async def _cancel():
|
|
await context.bridge.cancel_task_with_reason(context.bridge_task, reason)
|
|
|
|
loop.run_until_complete(_cancel())
|
|
context.cancel_reason = reason
|
|
|
|
|
|
@when("I save a checkpoint and then close the state manager")
|
|
def step_save_and_close_state_manager(context):
|
|
context.state_manager.update_state(
|
|
{"messages": [{"role": "user", "content": "test"}]}
|
|
)
|
|
context.state_manager._save_checkpoint()
|
|
context.state_manager.close()
|
|
|
|
|
|
@when("I close the event queue")
|
|
def step_close_event_queue(context):
|
|
context.event_queue.close()
|
|
|
|
|
|
@when('I try to register another resource named "{name}"')
|
|
def step_try_register_duplicate(context, name):
|
|
try:
|
|
dup_resource = MockAsyncResource(name)
|
|
context.tracker.register(name, dup_resource)
|
|
context.lsp_error = None
|
|
except ValueError as exc:
|
|
context.lsp_error = exc
|
|
|
|
|
|
@when("I try to register a resource with an empty name")
|
|
def step_try_register_empty_name(context):
|
|
try:
|
|
context.tracker.register("", MockAsyncResource("empty"))
|
|
context.lsp_error = None
|
|
except ValueError as exc:
|
|
context.lsp_error = exc
|
|
|
|
|
|
@when('I try to register a None resource with name "{name}"')
|
|
def step_try_register_none_resource(context, name):
|
|
try:
|
|
context.tracker.register(name, None) # type: ignore[arg-type]
|
|
context.lsp_error = None
|
|
except ValueError as exc:
|
|
context.lsp_error = exc
|
|
|
|
|
|
@when("I use the tracker as an async context manager and register the resource")
|
|
def step_use_context_manager(context):
|
|
loop = asyncio.new_event_loop()
|
|
|
|
async def _run():
|
|
async with context.tracker as tracker:
|
|
tracker.register(context.current_resource_name, context.current_resource)
|
|
|
|
loop.run_until_complete(_run())
|
|
loop.close()
|
|
|
|
|
|
@when("I close all tracked resources twice")
|
|
def step_close_all_twice(context):
|
|
loop = asyncio.new_event_loop()
|
|
loop.run_until_complete(context.tracker.close_all())
|
|
loop.run_until_complete(context.tracker.close_all())
|
|
loop.close()
|
|
|
|
|
|
@when("I run enhanced cleanup with a {timeout:g} second timeout")
|
|
def step_enhanced_bridge_cleanup(context, timeout):
|
|
loop = context.bridge_loop
|
|
|
|
async def _cleanup():
|
|
await context.bridge.cleanup_tasks_async(timeout=timeout)
|
|
|
|
loop.run_until_complete(_cleanup())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Thens
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@then('the resource "{name}" should be closed')
|
|
def step_assert_resource_closed(context, name):
|
|
resource = context.resources[name]
|
|
assert resource.closed, f"Resource '{name}' was not closed"
|
|
|
|
|
|
@then("the tracker should have zero open resources")
|
|
def step_assert_zero_open(context):
|
|
assert context.tracker.open_count == 0, (
|
|
f"Expected 0 open resources, got {context.tracker.open_count}"
|
|
)
|
|
|
|
|
|
@then("all {count:d} resources should be closed")
|
|
def step_assert_all_closed(context, count):
|
|
closed_count = sum(1 for _, r in context.multi_resources if r.closed)
|
|
assert closed_count == count, f"Expected {count} closed, got {closed_count}"
|
|
|
|
|
|
@then('a warning should be logged mentioning "{text}"')
|
|
def step_assert_warning_logged(context, text):
|
|
warnings = [r for r in context.log_handler.records if r.levelno >= logging.WARNING]
|
|
matching = [r for r in warnings if text in r.getMessage()]
|
|
assert matching, (
|
|
f"No warning mentioning '{text}' found. "
|
|
f"Warnings: {[r.getMessage() for r in warnings]}"
|
|
)
|
|
|
|
|
|
@then("the task should be cancelled")
|
|
def step_assert_task_cancelled(context):
|
|
task = context.bridge_task
|
|
assert task.done(), "Task should be done after cleanup"
|
|
|
|
|
|
@then("the bridge should have no active tasks")
|
|
def step_assert_no_active_tasks(context):
|
|
assert len(context.bridge._active_tasks) == 0, (
|
|
f"Expected 0 active tasks, got {len(context.bridge._active_tasks)}"
|
|
)
|
|
|
|
|
|
@then("a warning should be logged about forced termination")
|
|
def step_assert_forced_termination_warning(context):
|
|
warnings = [r for r in context.log_handler.records if r.levelno >= logging.WARNING]
|
|
matching = [
|
|
r
|
|
for r in warnings
|
|
if "timeout" in r.getMessage().lower() or "forced" in r.getMessage().lower()
|
|
]
|
|
assert matching, (
|
|
f"No forced termination warning found. "
|
|
f"Warnings: {[r.getMessage() for r in warnings]}"
|
|
)
|
|
|
|
|
|
@then("the tracker should report the timed-out resource")
|
|
def step_assert_timeout_reported(context):
|
|
assert len(context.tracker.timed_out_resources) > 0, (
|
|
"Expected at least one timed-out resource"
|
|
)
|
|
|
|
|
|
@then('the cancellation reason should be "{reason}"')
|
|
def step_assert_cancellation_reason(context, reason):
|
|
assert context.bridge.cancellation_reasons.get(context.bridge_task) == reason, (
|
|
f"Expected reason '{reason}', "
|
|
f"got '{context.bridge.cancellation_reasons.get(context.bridge_task)}'"
|
|
)
|
|
|
|
|
|
@then("the checkpoint file should exist and be readable")
|
|
def step_assert_checkpoint_exists(context):
|
|
checkpoints = list(context.checkpoint_dir.glob("checkpoint_*.json"))
|
|
assert len(checkpoints) >= 1, "No checkpoint files found"
|
|
content = checkpoints[0].read_text(encoding="utf-8")
|
|
assert len(content) > 0, "Checkpoint file is empty"
|
|
|
|
|
|
@then("the state manager should be marked as closed")
|
|
def step_assert_state_manager_closed(context):
|
|
assert context.state_manager.is_closed, "StateManager should be closed"
|
|
|
|
|
|
@then("all subscriptions should be removed")
|
|
def step_assert_subscriptions_removed(context):
|
|
assert len(context.event_queue._subscriptions) == 0, (
|
|
f"Expected 0 subscriptions, got {len(context.event_queue._subscriptions)}"
|
|
)
|
|
|
|
|
|
@then("the subscription count should be zero")
|
|
def step_assert_subscription_count_zero(context):
|
|
assert len(context.event_queue._subscriptions) == 0
|
|
|
|
|
|
@then('the resource "{name}" should be closed after exiting the context')
|
|
def step_assert_closed_after_context(context, name):
|
|
resource = context.resources[name]
|
|
assert resource.closed, f"Resource '{name}' should be closed after context exit"
|
|
|
|
|
|
@then('the resource "{name}" should be closed exactly once')
|
|
def step_assert_closed_once(context, name):
|
|
resource = context.resources[name]
|
|
assert resource.closed, f"Resource '{name}' was not closed"
|
|
assert resource.close_count == 1, (
|
|
f"Resource '{name}' was closed {resource.close_count} times, expected 1"
|
|
)
|
|
|
|
|
|
@then("completed tasks should be logged")
|
|
def step_assert_tasks_logged(context):
|
|
infos = [r for r in context.log_handler.records if r.levelno >= logging.INFO]
|
|
has_cleanup_log = any(
|
|
"cleanup" in r.getMessage().lower() or "task" in r.getMessage().lower()
|
|
for r in infos
|
|
)
|
|
# The cleanup operation should have been logged, and tasks cleared
|
|
assert has_cleanup_log or len(context.bridge._active_tasks) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Logging capture helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _CapturingHandler(logging.Handler):
|
|
"""Captures log records for assertion in tests."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.records: list[logging.LogRecord] = []
|
|
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
self.records.append(record)
|