From 0b25ee2e5e25c1a409623e2ac5d93c0997d15bb0 Mon Sep 17 00:00:00 2001 From: suhasagg Date: Wed, 7 Oct 2026 21:17:30 +0530 Subject: [PATCH] fix(hosting): avoid retaining per-session locks --- .../hosting/agent_framework_hosting/_state.py | 3 ++- python/packages/hosting/tests/hosting/test_state.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index cbebd07031a..34f02cb3afd 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/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: diff --git a/python/packages/hosting/tests/hosting/test_state.py b/python/packages/hosting/tests/hosting/test_state.py index f4c991fca50..ab6223dde17 100644 --- a/python/packages/hosting/tests/hosting/test_state.py +++ b/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)