diff --git a/openhands-agent-server/openhands/agent_server/conversation_service.py b/openhands-agent-server/openhands/agent_server/conversation_service.py index 32f9e30f59..4f9abac659 100644 --- a/openhands-agent-server/openhands/agent_server/conversation_service.py +++ b/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 diff --git a/openhands-sdk/openhands/sdk/io/local.py b/openhands-sdk/openhands/sdk/io/local.py index 3e2ef4cd25..36179470ea 100644 --- a/openhands-sdk/openhands/sdk/io/local.py +++ b/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: diff --git a/tests/agent_server/test_conversation_service.py b/tests/agent_server/test_conversation_service.py index 4fb3b79784..e867187cf8 100644 --- a/tests/agent_server/test_conversation_service.py +++ b/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."""