fix(agent-server): replace global _lifecycle_lock with per-conversation locks

openhands-agent-server/openhands/agent_server/conversation_service.py

@@ -5,11 +5,12 @@
 import os
 from collections.abc import Awaitable, Callable
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import suppress
+from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 from uuid import UUID, uuid4
+from weakref import WeakValueDictionary
 
 import httpx
 from pydantic import BaseModel
@@ -645,6 +646,14 @@ class ConversationService:
         default_factory=dict, init=False
     )
     _lifecycle_lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
+    _lifecycle_condition: asyncio.Condition = field(
+        default_factory=asyncio.Condition, init=False
+    )
+    _active_lifecycle_operations: int = field(default=0, init=False)
+    _exclusive_lifecycle_pending: bool = field(default=False, init=False)
+    _conversation_locks: WeakValueDictionary[UUID, asyncio.Lock] = field(
+        default_factory=WeakValueDictionary, init=False
+    )
     _conversation_webhook_subscribers: list["ConversationWebhookSubscriber"] = field(
         default_factory=list, init=False
     )
@@ -751,7 +760,7 @@ async def activate_credential_binding(
         secret_name: str,
         binding: VersionedCredentialBinding,
     ) -> None:
-        async with self._lifecycle_lock:
+        async with self._conversation_lifecycle(conversation_id):
             event_services = self._event_services
             event_service = (
                 event_services.get(conversation_id)
@@ -770,7 +779,7 @@ async def activate_credential_binding(
             )
 
     async def prepare_for_sandbox_pause(self) -> None:
-        async with self._lifecycle_lock:
+        async with self._exclusive_lifecycle():
             event_services = self._event_services
             if event_services is None:
                 raise ValueError("inactive_service")
@@ -1047,10 +1056,56 @@ def _prepare_persisted_runtime(self, stored: StoredConversation) -> None:
                 context=f"resuming conversation {stored.id}",
             )
 
+    def _get_conversation_lock(self, conversation_id: UUID) -> asyncio.Lock:
+        lock = self._conversation_locks.get(conversation_id)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._conversation_locks[conversation_id] = lock
+        return lock
+
+    @asynccontextmanager
+    async def _conversation_lifecycle(self, conversation_id: UUID):
+        async with self._lifecycle_condition:
+            await self._lifecycle_condition.wait_for(
+                lambda: not self._exclusive_lifecycle_pending
+            )
+            self._active_lifecycle_operations += 1
+        try:
+            async with self._get_conversation_lock(conversation_id):
+                yield
+        finally:
+            async with self._lifecycle_condition:
+                self._active_lifecycle_operations -= 1
+                if self._active_lifecycle_operations == 0:
+                    self._lifecycle_condition.notify_all()
+
+    @asynccontextmanager
+    async def _exclusive_lifecycle(self):
+        async with self._lifecycle_lock:
+            try:
+                async with self._lifecycle_condition:
+                    self._exclusive_lifecycle_pending = True
+                    await self._lifecycle_condition.wait_for(
+                        lambda: self._active_lifecycle_operations == 0
+                    )
+                yield
+            finally:
+                async with self._lifecycle_condition:
+                    self._exclusive_lifecycle_pending = False
+                    self._lifecycle_condition.notify_all()
+
     async def _get_or_load_event_service(
         self, conversation_id: UUID
     ) -> EventService | None:
-        async with self._lifecycle_lock:
+        event_services = self._event_services
+        if event_services is None:
+            raise ValueError("inactive_service")
+        if (
+            conversation_id not in event_services
+            and conversation_id not in self._conversation_records
+        ):
+            return None
+        async with self._conversation_lifecycle(conversation_id):
             return await self._get_or_load_event_service_locked(conversation_id)
 
     async def _get_or_load_event_service_locked(
@@ -1318,7 +1373,7 @@ async def _start_conversation(
         if existing_record is not None or (
             existing_event_service is not None and existing_event_service.is_open()
         ):
-            async with self._lifecycle_lock:
+            async with self._conversation_lifecycle(conversation_id):
                 existing_event_service = self._event_services.get(conversation_id)
                 if (
                     existing_event_service is not None
@@ -1625,7 +1680,7 @@ async def _start_conversation(
                 launched_agent_profile=launched_agent_profile,
                 **request_data,
             )
-        async with self._lifecycle_lock:
+        async with self._conversation_lifecycle(conversation_id):
             # New conversation: the agent is written to base_state.json (its
             # single source of truth), not to meta.json. Pass it explicitly.
             # ``new_agent`` is ``request.agent`` (decrypted when the request was
@@ -1683,10 +1738,15 @@ async def resume_conversation(self, conversation_id: UUID) -> bool:
         return bool(await self._get_or_load_event_service(conversation_id))
 
     async def delete_conversation(self, conversation_id: UUID) -> bool:
-        async with self._lifecycle_lock:
-            event_services = self._event_services
-            if event_services is None:
-                raise ValueError("inactive_service")
+        event_services = self._event_services
+        if event_services is None:
+            raise ValueError("inactive_service")
+        if (
+            conversation_id not in event_services
+            and conversation_id not in self._conversation_records
+        ):
+            return False
+        async with self._conversation_lifecycle(conversation_id):
             event_service = await self._get_or_load_event_service_locked(
                 conversation_id,
                 require_runtime_bindings=False,
@@ -1911,7 +1971,7 @@ async def fork_conversation(
         # directory so we don't leave stale state on disk.
         fork_dir = self.conversations_dir / fork_conv_id.hex
         try:
-            async with self._lifecycle_lock:
+            async with self._conversation_lifecycle(fork_conv_id):
                 fork_event_service = await self._start_event_service(
                     fork_stored, is_new_conversation=True, agent=fork_agent
                 )
@@ -2036,7 +2096,7 @@ async def _evict_idle_conversations(self, ttl_seconds: float) -> None:
 
         Running or externally-subscribed conversations are skipped.
         """
-        async with self._lifecycle_lock:
+        async with self._exclusive_lifecycle():
             event_services = self._event_services
             if event_services is None:
                 return
@@ -2092,7 +2152,7 @@ async def __aexit__(self, exc_type, exc_value, traceback):
                 await self._lease_renewal_task
             self._lease_renewal_task = None
 
-        async with self._lifecycle_lock:
+        async with self._exclusive_lifecycle():
             event_services = self._event_services
             if event_services is None:
                 return

openhands-sdk/openhands/sdk/io/local.py

@@ -2,11 +2,13 @@
 import shutil
 from collections.abc import Iterator
 from contextlib import contextmanager
+from pathlib import Path
 
 from filelock import FileLock, Timeout
 
 from openhands.sdk.io.cache import MemoryLRUCache
 from openhands.sdk.logger import get_logger
+from openhands.sdk.utils.files import atomic_write_text
 from openhands.sdk.utils.path import to_posix_path
 
 from .base import FileStore
@@ -62,8 +64,7 @@ def write(self, path: str, contents: str | bytes) -> None:
         full_path = self.get_full_path(path)
         os.makedirs(os.path.dirname(full_path), exist_ok=True)
         if isinstance(contents, str):
-            with open(full_path, "w", encoding="utf-8") as f:
-                f.write(contents)
+            atomic_write_text(Path(full_path), contents)
             self.cache[full_path] = contents
         else:
             with open(full_path, "wb") as f:

tests/agent_server/test_conversation_service.py

@@ -590,6 +590,80 @@ async def test_prepare_for_sandbox_pause_drains_active_services(tmp_path):
     await service.__aexit__(None, None, None)
 
 
+@pytest.mark.asyncio
+async def test_conversation_lifecycle_serializes_only_matching_ids(tmp_path):
+    service = ConversationService(conversations_dir=tmp_path / "conversations")
+    first_id = uuid4()
+    second_id = uuid4()
+    same_id_entered = asyncio.Event()
+    other_id_entered = asyncio.Event()
+
+    async def enter_lifecycle(conversation_id: UUID, entered: asyncio.Event):
+        async with service._conversation_lifecycle(conversation_id):
+            entered.set()
+
+    async with service._conversation_lifecycle(first_id):
+        same_id_task = asyncio.create_task(enter_lifecycle(first_id, same_id_entered))
+        other_id_task = asyncio.create_task(
+            enter_lifecycle(second_id, other_id_entered)
+        )
+        await asyncio.wait_for(other_id_entered.wait(), timeout=1)
+        assert not same_id_entered.is_set()
+
+    await asyncio.gather(same_id_task, other_id_task)
+    assert same_id_entered.is_set()
+
+
+@pytest.mark.asyncio
+async def test_prepare_for_sandbox_pause_blocks_new_hydration(
+    persisted_conversation,
+):
+    conversations_dir, conversation_id = persisted_conversation
+    service = ConversationService(conversations_dir=conversations_dir)
+    await service.__aenter__()
+    existing_id = uuid4()
+    existing = AsyncMock(spec=EventService)
+    close_entered = asyncio.Event()
+    finish_close = asyncio.Event()
+    hydration_entered = asyncio.Event()
+    record = service._conversation_records[conversation_id]
+    hydrated = AsyncMock(spec=EventService)
+    hydrated.stored = record.stored
+
+    async def close(*_args):
+        close_entered.set()
+        await finish_close.wait()
+
+    async def hydrate(*_args, **_kwargs):
+        hydration_entered.set()
+        assert service._event_services is not None
+        service._event_services[conversation_id] = hydrated
+        return hydrated
+
+    existing.__aexit__.side_effect = close
+    assert service._event_services is not None
+    service._event_services[existing_id] = existing
+
+    try:
+        with patch.object(service, "_start_event_service", side_effect=hydrate):
+            pause_task = asyncio.create_task(service.prepare_for_sandbox_pause())
+            await asyncio.wait_for(close_entered.wait(), timeout=1)
+            hydration_task = asyncio.create_task(
+                service.get_event_service(conversation_id)
+            )
+            await asyncio.sleep(0)
+            assert not hydration_entered.is_set()
+
+            finish_close.set()
+            await asyncio.wait_for(pause_task, timeout=1)
+            assert service._event_services == {}
+            assert not hydration_entered.is_set()
+            assert await asyncio.wait_for(hydration_task, timeout=1) is hydrated
+    finally:
+        finish_close.set()
+        await service.__aexit__(None, None, None)
+
+
 @pytest.mark.asyncio
 async def test_prepare_for_sandbox_pause_closes_services_concurrently(tmp_path):
     service = ConversationService(conversations_dir=tmp_path / "conversations")
@@ -926,7 +1000,9 @@ async def test_waiting_hydration_cannot_restore_deleted_conversation(
         replacement_runtime = AsyncMock(spec=EventService)
         replacement_runtime.stored = record.stored
 
-        async def publish_replacement(_stored: StoredConversation) -> EventService:
+        async def publish_replacement(
+            _stored: StoredConversation, *, agent: AgentBase | None = None
+        ) -> EventService:
             assert service._event_services is not None
             service._event_services[conversation_id] = replacement_runtime
             service._conversation_records[conversation_id] = record
@@ -940,7 +1016,8 @@ async def publish_replacement(_stored: StoredConversation) -> EventService:
             ) as start_event_service,
             patch("openhands.agent_server.conversation_service.safe_rmtree"),
         ):
-            await service._lifecycle_lock.acquire()
+            conversation_lock = service._get_conversation_lock(conversation_id)
+            await conversation_lock.acquire()
             try:
                 getter_task = asyncio.create_task(
                     service.get_event_service(conversation_id)
@@ -953,7 +1030,7 @@ async def publish_replacement(_stored: StoredConversation) -> EventService:
                 )
                 await asyncio.sleep(0)
             finally:
-                service._lifecycle_lock.release()
+                conversation_lock.release()
 
             getter_result, deleted = await asyncio.gather(getter_task, delete_task)
 
@@ -2627,6 +2704,22 @@ async def test_delete_conversation_not_found(self, conversation_service):
         result = await conversation_service.delete_conversation(uuid4())
         assert result is False
 
+    @pytest.mark.asyncio
+    async def test_missing_conversations_do_not_accumulate_locks(
+        self, conversation_service
+    ):
+        for _ in range(100):
+            conversation_id = uuid4()
+            assert await conversation_service.get_event_service(conversation_id) is None
+            assert (
+                await conversation_service.resume_conversation(conversation_id) is False
+            )
+            assert (
+                await conversation_service.delete_conversation(conversation_id) is False
+            )
+
+        assert len(conversation_service._conversation_locks) == 0
+
     @pytest.mark.asyncio
     async def test_delete_conversation_success(self, conversation_service):
         """Test successful conversation deletion."""