diff --git a/ms_agent/a2a/session_store.py b/ms_agent/a2a/session_store.py index c74c31dbb..fd7b389bc 100644 --- a/ms_agent/a2a/session_store.py +++ b/ms_agent/a2a/session_store.py @@ -71,16 +71,18 @@ def __init__( async def get_or_create(self, task_id: str) -> A2ATaskEntry: """Return an existing task entry or create a new one.""" - if task_id in self._tasks: - entry = self._tasks[task_id] - entry.touch() - return entry - - if len(self._tasks) >= self.max_tasks: + while True: + if task_id in self._tasks: + entry = self._tasks[task_id] + entry.touch() + return entry + if len(self._tasks) < self.max_tasks: + break evicted_id = self._evict_lru() if evicted_id is None: raise MaxTasksError(self.max_tasks) await self._cleanup_entry(evicted_id) + # Cleanup yields, so recheck task identity and available capacity. if not self.config_path or not os.path.exists(self.config_path): raise ConfigError(f'Config not found: {self.config_path}') diff --git a/ms_agent/acp/session_store.py b/ms_agent/acp/session_store.py index 8746052f6..5f7537cff 100644 --- a/ms_agent/acp/session_store.py +++ b/ms_agent/acp/session_store.py @@ -78,11 +78,12 @@ async def create( meta: dict | None = None, ) -> ACPSessionEntry: """Create a new session backed by a freshly loaded agent.""" - if len(self._sessions) >= self.max_sessions: + while len(self._sessions) >= self.max_sessions: evicted_id = self._evict_lru() if evicted_id is None: raise MaxSessionsError(self.max_sessions) await self._cleanup_session(evicted_id) + # Another request may have filled the slot during cleanup. if not config_path or not os.path.exists(config_path): raise ConfigError(f'Config not found: {config_path}') diff --git a/tests/test_session_store_concurrency.py b/tests/test_session_store_concurrency.py new file mode 100644 index 000000000..35b8f784a --- /dev/null +++ b/tests/test_session_store_concurrency.py @@ -0,0 +1,135 @@ +"""Concurrent admission must preserve protocol-store capacity and identity.""" + +import asyncio +from importlib import import_module +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from omegaconf import OmegaConf + + +@pytest.fixture(params=['a2a', 'acp']) +def store_case(request, tmp_path, monkeypatch): + module = import_module(f'ms_agent.{request.param}.session_store') + config_path = tmp_path / 'agent.yaml' + config_path.write_text('{}', encoding='utf-8') + agents = [] + + def build(**kwargs): + agent = SimpleNamespace(cleanup_tools=AsyncMock()) + agents.append(agent) + return agent + + monkeypatch.setattr(module.Config, 'from_task', + lambda path: OmegaConf.create({})) + monkeypatch.setattr(module.AgentLoader, 'build', build) + if request.param == 'a2a': + store = module.A2AAgentStore(str(config_path), max_tasks=1) + create = store.get_or_create + entries = store._tasks + capacity_error = module.MaxTasksError + else: + store = module.ACPSessionStore(max_sessions=1) + + async def create(key): + return await store.create(str(config_path), str(tmp_path)) + + entries = store._sessions + capacity_error = module.MaxSessionsError + return store, create, entries, agents, capacity_error + + +@pytest.mark.parametrize('store_case', ['a2a'], indirect=True) +@pytest.mark.asyncio +async def test_same_a2a_task_during_eviction_reuses_its_agent(store_case): + store, create, entries, agents, _ = store_case + started, release = asyncio.Event(), asyncio.Event() + + async def cleanup(): + started.set() + await release.wait() + + try: + old = await create('old') + old.agent.cleanup_tools.side_effect = cleanup + first = asyncio.create_task(create('shared')) + await asyncio.wait_for(started.wait(), timeout=2) + second = asyncio.create_task(create('shared')) + await asyncio.sleep(0) + release.set() + first_entry, second_entry = await asyncio.gather(first, second) + assert first_entry is second_entry + assert store.get('shared') is first_entry + assert len(agents) == 2 + finally: + release.set() + await store.close_all() + + +@pytest.mark.asyncio +async def test_concurrent_admission_during_eviction_respects_capacity(store_case): + store, create, entries, agents, capacity_error = store_case + started, release = asyncio.Event(), asyncio.Event() + + async def cleanup(): + started.set() + await release.wait() + + async def admit(key): + entry = await create(key) + entry.is_running = True + return entry + + try: + old = await create('old') + old.agent.cleanup_tools.side_effect = cleanup + first = asyncio.create_task(admit('first')) + await asyncio.wait_for(started.wait(), timeout=2) + second = asyncio.create_task(admit('second')) + await asyncio.sleep(0) + release.set() + results = await asyncio.gather(first, second, return_exceptions=True) + + assert len(entries) == 1 + assert sum(isinstance(r, capacity_error) for r in results) == 1 + assert len(agents) == 2 # the evicted agent and one admitted agent + finally: + release.set() + await store.close_all() + + +@pytest.mark.asyncio +async def test_cancelled_admission_does_not_block_the_next_request(store_case): + store, create, entries, _, _ = store_case + started = asyncio.Event() + + async def cleanup(): + started.set() + await asyncio.Event().wait() + + try: + old = await create('old') + old.agent.cleanup_tools.side_effect = cleanup + pending = asyncio.create_task(create('cancelled')) + await asyncio.wait_for(started.wait(), timeout=2) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + entry = await asyncio.wait_for(create('next'), timeout=2) + assert list(entries.values()) == [entry] + finally: + await store.close_all() + + +@pytest.mark.asyncio +async def test_idle_replacement_cleans_the_evicted_agent(store_case): + store, create, entries, agents, _ = store_case + try: + old = await create('old') + new = await create('new') + assert list(entries.values()) == [new] + old.agent.cleanup_tools.assert_awaited_once() + assert len(agents) == 2 + finally: + await store.close_all()