Python: fix(hosting): avoid retaining per-session locks

python/packages/hosting/agent_framework_hosting/_state.py

@@ -22,6 +22,7 @@
 import asyncio
 import copy
 import inspect
+import weakref
 from collections.abc import Awaitable, Callable, Mapping
 from typing import Any, Generic, Protocol, TypedDict, TypeVar, cast, runtime_checkable
 
@@ -113,7 +114,7 @@ def __init__(
         if not callable(target) and not inspect.isawaitable(target):
             self._cached_target = target
         self._session_store: SessionStore = session_store if session_store is not None else SessionStore()
-        self._session_locks: dict[str, asyncio.Lock] = {}
+        self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
         mark_feature_used(FeatureIndex.HOSTING)
 
     async def get_target(self) -> AgentT:

python/packages/hosting/tests/hosting/test_state.py

@@ -253,6 +253,18 @@ async def set(self, session_id: str, session: AgentSession) -> None:
         assert all(session.session_id == "session-1" for session in sessions)
         assert len(agent.created_sessions) == 1
 
+    async def test_get_or_create_session_does_not_retain_locks_for_completed_sessions(self) -> None:
+        agent = _FakeAgent()
+        state = AgentState(agent)
+
+        for i in range(1000):
+            session_id = f"session-{i}"
+            await state.get_or_create_session(session_id)
+            await state.session_store.delete(session_id)
+
+        assert len(state.session_store._sessions) == 0  # pyright: ignore[reportPrivateUsage]
+        assert len(state._session_locks) == 0  # pyright: ignore[reportPrivateUsage]
+
     async def test_get_or_create_session_reuses_a_session_set_on_the_state(self) -> None:
         agent = _FakeAgent()
         state = AgentState(agent)