diff --git a/src/opencode_a2a/execution/session_manager.py b/src/opencode_a2a/execution/session_manager.py index 955c385..a4d562d 100644 --- a/src/opencode_a2a/execution/session_manager.py +++ b/src/opencode_a2a/execution/session_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import weakref from ..invocation import call_with_supported_kwargs from ..server.state_store import MemorySessionStateRepository, SessionStateRepository @@ -24,7 +25,9 @@ def __init__( ) self._lock = asyncio.Lock() self._inflight_session_creates: dict[tuple[str, str], asyncio.Task[str]] = {} - self._session_locks: dict[str, asyncio.Lock] = {} + self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( + weakref.WeakValueDictionary() + ) async def get_or_create_session( self, diff --git a/tests/execution/test_session_lock_lifecycle.py b/tests/execution/test_session_lock_lifecycle.py new file mode 100644 index 0000000..bad45ce --- /dev/null +++ b/tests/execution/test_session_lock_lifecycle.py @@ -0,0 +1,33 @@ +import gc +import weakref +from unittest.mock import AsyncMock + +import pytest + +from opencode_a2a.execution.session_manager import SessionManager +from opencode_a2a.opencode_upstream_client import OpencodeUpstreamClient + + +@pytest.mark.asyncio +async def test_session_manager_reuses_live_lock_for_same_session() -> None: + manager = SessionManager(client=AsyncMock(spec=OpencodeUpstreamClient)) + + lock1 = await manager.get_session_lock("session-1") + lock2 = await manager.get_session_lock("session-1") + + assert lock1 is lock2 + + +@pytest.mark.asyncio +async def test_session_manager_does_not_strongly_retain_idle_locks() -> None: + manager = SessionManager(client=AsyncMock(spec=OpencodeUpstreamClient)) + + lock = await manager.get_session_lock("session-1") + lock_ref = weakref.ref(lock) + assert "session-1" in manager._session_locks + + del lock + gc.collect() + + assert lock_ref() is None + assert "session-1" not in manager._session_locks