diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 3c0f9eab8..598d6213b 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -9,6 +9,7 @@ import dataclasses import inspect import os import time +import weakref from collections.abc import Coroutine, Iterable, Mapping from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress from dataclasses import dataclass, field @@ -394,7 +395,9 @@ class AgentLoop: self._runtime_context_providers: list[RuntimeContextProvider] = [] self._active_tasks: dict[str, set[asyncio.Task[Any]]] = {} self._background_tasks: set[asyncio.Task[Any]] = set() - self._session_locks: dict[str, asyncio.Lock] = {} + self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( + weakref.WeakValueDictionary() + ) # Per-session pending queues for mid-turn message injection. # When a session has an active task, new messages for that session # are routed here instead of creating a new task. @@ -1206,7 +1209,7 @@ class AgentLoop: session_key = self._effective_session_key(msg) if session_key != msg.session_key: msg = dataclasses.replace(msg, session_key_override=session_key) - lock = self._session_locks.setdefault(session_key, asyncio.Lock()) + lock = self._get_session_lock(session_key) gate = self._concurrency_gate or nullcontext() delivery = self.turn_delivery_factory.unrouted(msg, session_key) @@ -2108,7 +2111,7 @@ class AgentLoop: content=content, media=media or [], metadata=metadata, ) # Share the dispatch lock so direct calls serialize with bus turns. - lock = self._session_locks.setdefault(session_key, asyncio.Lock()) + lock = self._get_session_lock(session_key) try: async with lock: kwargs: dict[str, Any] = { @@ -2139,3 +2142,11 @@ class AgentLoop: finally: await self.runtime_event_publisher.run_status_changed(msg, session_key, "idle") self.runtime_event_publisher.clear_turn(session_key) + + def _get_session_lock(self, session_key: str) -> asyncio.Lock: + """Return the shared lock while allowing idle session entries to expire.""" + lock = self._session_locks.get(session_key) + if lock is None: + lock = asyncio.Lock() + self._session_locks[session_key] = lock + return lock diff --git a/tests/agent/test_session_lock_lifecycle.py b/tests/agent/test_session_lock_lifecycle.py new file mode 100644 index 000000000..43a862175 --- /dev/null +++ b/tests/agent/test_session_lock_lifecycle.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import asyncio +import gc +from unittest.mock import MagicMock + +import pytest + + +def _make_loop(loop_factory): + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + return loop_factory(provider=provider) + + +def test_idle_agent_session_locks_are_released(loop_factory): + loop = _make_loop(loop_factory) + + for index in range(1000): + lock = loop._get_session_lock(f"api:temporary-{index}") + + del lock + gc.collect() + + assert len(loop._session_locks) == 0 + + +@pytest.mark.asyncio +async def test_waiter_keeps_agent_session_lock_alive(loop_factory): + loop = _make_loop(loop_factory) + owner_lock = loop._get_session_lock("api:shared") + await owner_lock.acquire() + waiter_started = asyncio.Event() + waiter_entered = asyncio.Event() + + async def wait_for_lock() -> None: + lock = loop._get_session_lock("api:shared") + waiter_started.set() + async with lock: + waiter_entered.set() + + waiter = asyncio.create_task(wait_for_lock()) + await waiter_started.wait() + + assert loop._get_session_lock("api:shared") is owner_lock + assert not waiter_entered.is_set() + + owner_lock.release() + await waiter + del owner_lock + gc.collect() + + assert "api:shared" not in loop._session_locks