fix(agent): release idle session locks

This commit is contained in:
yu-xin-c 2026-07-29 01:03:29 +08:00 committed by Xubin Ren
parent 52680dbe19
commit 9ec4420104
2 changed files with 67 additions and 3 deletions

View File

@ -9,6 +9,7 @@ import dataclasses
import inspect import inspect
import os import os
import time import time
import weakref
from collections.abc import Coroutine, Iterable, Mapping from collections.abc import Coroutine, Iterable, Mapping
from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress
from dataclasses import dataclass, field from dataclasses import dataclass, field
@ -394,7 +395,9 @@ class AgentLoop:
self._runtime_context_providers: list[RuntimeContextProvider] = [] self._runtime_context_providers: list[RuntimeContextProvider] = []
self._active_tasks: dict[str, set[asyncio.Task[Any]]] = {} self._active_tasks: dict[str, set[asyncio.Task[Any]]] = {}
self._background_tasks: set[asyncio.Task[Any]] = set() 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. # Per-session pending queues for mid-turn message injection.
# When a session has an active task, new messages for that session # When a session has an active task, new messages for that session
# are routed here instead of creating a new task. # are routed here instead of creating a new task.
@ -1206,7 +1209,7 @@ class AgentLoop:
session_key = self._effective_session_key(msg) session_key = self._effective_session_key(msg)
if session_key != msg.session_key: if session_key != msg.session_key:
msg = dataclasses.replace(msg, session_key_override=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() gate = self._concurrency_gate or nullcontext()
delivery = self.turn_delivery_factory.unrouted(msg, session_key) delivery = self.turn_delivery_factory.unrouted(msg, session_key)
@ -2108,7 +2111,7 @@ class AgentLoop:
content=content, media=media or [], metadata=metadata, content=content, media=media or [], metadata=metadata,
) )
# Share the dispatch lock so direct calls serialize with bus turns. # 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: try:
async with lock: async with lock:
kwargs: dict[str, Any] = { kwargs: dict[str, Any] = {
@ -2139,3 +2142,11 @@ class AgentLoop:
finally: finally:
await self.runtime_event_publisher.run_status_changed(msg, session_key, "idle") await self.runtime_event_publisher.run_status_changed(msg, session_key, "idle")
self.runtime_event_publisher.clear_turn(session_key) 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

View File

@ -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