diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index d86b440a7..f10434f1d 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -852,6 +852,22 @@ class AgentLoop: metadata={**metadata, "render_as": "text"}, ) + def _track_active_task(self, key: str, task: asyncio.Task[Any]) -> None: + """Track active session work until its task group becomes empty.""" + tasks = self._active_tasks.setdefault(key, set()) + tasks.add(task) + task.add_done_callback(partial(self._active_task_done, key, tasks)) + + def _active_task_done( + self, + key: str, + tasks: set[asyncio.Task[Any]], + task: asyncio.Task[Any], + ) -> None: + tasks.discard(task) + if not tasks and self._active_tasks.get(key) is tasks: + self._active_tasks.pop(key, None) + async def _cancel_active_tasks(self, key: str) -> int: """Cancel and await all active work for *key*. @@ -1350,12 +1366,7 @@ class AgentLoop: # Compute the effective session key before dispatching # This ensures /stop command can find tasks correctly when unified session is enabled task = asyncio.create_task(self._dispatch(msg)) - active_tasks: set[asyncio.Task[Any]] = self._active_tasks.setdefault( - effective_key, - set(), - ) - active_tasks.add(task) - task.add_done_callback(active_tasks.discard) + self._track_active_task(effective_key, task) finally: await self.aclose() diff --git a/tests/agent/test_task_cancel.py b/tests/agent/test_task_cancel.py index 642772a4b..0e71e8a7f 100644 --- a/tests/agent/test_task_cancel.py +++ b/tests/agent/test_task_cancel.py @@ -42,6 +42,65 @@ def _make_loop(*, tools_config=None): return loop, bus +class TestActiveTaskTracking: + @pytest.mark.asyncio + async def test_completed_task_removes_empty_session_group(self): + loop, _bus = _make_loop() + release = asyncio.Event() + task = asyncio.create_task(release.wait()) + + loop._track_active_task("test:c1", task) + release.set() + await task + await asyncio.sleep(0) + + assert "test:c1" not in loop._active_tasks + + @pytest.mark.asyncio + async def test_session_group_remains_until_last_task_completes(self): + loop, _bus = _make_loop() + releases = [asyncio.Event(), asyncio.Event()] + tasks = [asyncio.create_task(release.wait()) for release in releases] + for task in tasks: + loop._track_active_task("test:c1", task) + + releases[0].set() + await tasks[0] + await asyncio.sleep(0) + + assert loop._active_tasks["test:c1"] == {tasks[1]} + + releases[1].set() + await tasks[1] + await asyncio.sleep(0) + + assert "test:c1" not in loop._active_tasks + + @pytest.mark.asyncio + async def test_old_callback_preserves_replacement_session_group(self): + loop, _bus = _make_loop() + old_release = asyncio.Event() + new_release = asyncio.Event() + old_task = asyncio.create_task(old_release.wait()) + new_task = asyncio.create_task(new_release.wait()) + + loop._track_active_task("test:c1", old_task) + loop._active_tasks.pop("test:c1") + loop._track_active_task("test:c1", new_task) + + old_release.set() + await old_task + await asyncio.sleep(0) + + assert loop._active_tasks["test:c1"] == {new_task} + + new_release.set() + await new_task + await asyncio.sleep(0) + + assert "test:c1" not in loop._active_tasks + + class TestHandleStop: @pytest.mark.asyncio async def test_stop_no_active_task(self):