.

agentlightning/store/base.py

@@ -14,6 +14,7 @@ from agentlightning.types import (
     ResourcesUpdate,
     Rollout,
     RolloutConfig,
+    RolloutMode,
     RolloutStatus,
     Span,
     TaskInput,
@@ -99,6 +100,19 @@ class LightningStoreStatistics(TypedDict, total=False):
     """Memory capacity of the store in bytes."""
 
 
+class _EnqueueRolloutRequestRequired(TypedDict):
+    input: TaskInput
+
+
+class EnqueueRolloutRequest(_EnqueueRolloutRequestRequired, total=False):
+    """Payload describing a rollout to be queued via `enqueue_rollout`."""
+
+    mode: Optional[RolloutMode]
+    resources_id: Optional[str]
+    config: Optional[RolloutConfig]
+    metadata: Optional[Dict[str, Any]]
+
+
 class LightningStore:
     """Contract for the persistent control-plane that coordinates training rollouts.
 
@@ -156,7 +170,7 @@ class LightningStore:
     async def start_rollout(
         self,
         input: TaskInput,
-        mode: Literal["train", "val", "test"] | None = None,
+        mode: RolloutMode | None = None,
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
@@ -227,6 +241,22 @@ class LightningStore:
         """
         raise NotImplementedError()
 
+    async def enqueue_many_rollouts(self, inputs: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        """Persist multiple rollouts in `queuing` state.
+
+        The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
+        per request and preserves the input ordering. Subclasses can override to provide
+        more efficient bulk enqueue semantics.
+
+        Args:
+            inputs: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
+                parameters. Each entry requires `input` and can optionally include other fields.
+
+        Returns:
+            Rollouts enqueued in the same order as `inputs`.
+        """
+        raise NotImplementedError()
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """Claim the oldest queued rollout and transition it to `preparing`.
 
@@ -251,6 +281,29 @@ class LightningStore:
         """
         raise NotImplementedError()
 
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        """Claim up to `limit` queued rollouts without blocking.
+
+        The implementation can repeatedly invokes
+        [`dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] until reaching
+        the requested limit or the queue is empty. Subclasses can override it to fetch
+        multiple rollouts atomically.
+
+        Args:
+            limit: Maximum number of rollouts to claim. Non-positive values return an empty list.
+            worker_id: Optional worker identifier passed through to each dequeue call.
+
+        Returns:
+            Attempted rollouts claimed in FIFO order. May contain fewer than `limit` entries
+            when the queue is exhausted.
+        """
+        raise NotImplementedError()
+
     async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         """Create a manual retry attempt for an existing rollout.
 
@@ -719,7 +772,8 @@ class LightningStore:
 
         When `attempt_id` is `"latest"` the update must target the attempt with the highest
         `sequence_id`; otherwise it must target the specific attempt. Implementations should
-        propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
+        propagate status changes to the rollout (for example
+        via [`rollout_status_from_attempt()`][agentlightning.store.utils.rollout_status_from_attempt])
         once the latest attempt transitions to a terminal state.
 
         Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],

agentlightning/store/client_server.py

@@ -565,6 +565,24 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(results, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/rollouts/search", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
+        async def search_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Rollout)
+            status_in = request.status_in if "status_in" in request.model_fields_set else None
+            rollout_id_in = request.rollout_id_in if "rollout_id_in" in request.model_fields_set else None
+            # Get all rollouts from the underlying store
+            results = await self.query_rollouts(
+                status_in=status_in,
+                rollout_id_in=rollout_id_in,
+                rollout_id_contains=request.rollout_id_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(results, limit=request.limit, offset=request.offset)
+
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Union[AttemptedRollout, Rollout])
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_rollout_by_id(rollout_id)
@@ -600,6 +618,20 @@ class LightningStoreServer(LightningStore):
         async def start_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.start_attempt(rollout_id)
 
+        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/search", response_model=PaginatedResult[Attempt])
+        async def search_attempts(  # pyright: ignore[reportUnusedFunction]
+            rollout_id: str, request: QueryAttemptsRequest
+        ):
+            _validate_paginated_request(request, Attempt)
+            attempts = await self.query_attempts(
+                rollout_id,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(attempts, limit=request.limit, offset=request.offset)
+
         @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
         async def update_attempt(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, attempt_id: str, request: UpdateAttemptRequest = Body(...)
@@ -627,6 +659,21 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(workers, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/workers/search", response_model=PaginatedResult[Worker])
+        async def search_workers(request: QueryWorkersRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Worker)
+            status_in = request.status_in if "status_in" in request.model_fields_set else None
+            workers = await self.query_workers(
+                status_in=status_in,
+                worker_id_contains=request.worker_id_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(workers, limit=request.limit, offset=request.offset)
+
         @api.get(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Optional[Worker])
         async def get_worker(worker_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_worker_by_id(worker_id)
@@ -719,6 +766,28 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(spans, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/spans/search", response_model=PaginatedResult[Span])
+        async def search_spans(request: QuerySpansRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Span)
+            spans = await self.query_spans(
+                request.rollout_id,
+                request.attempt_id,
+                trace_id=request.trace_id,
+                trace_id_contains=request.trace_id_contains,
+                span_id=request.span_id,
+                span_id_contains=request.span_id_contains,
+                parent_id=request.parent_id,
+                parent_id_contains=request.parent_id_contains,
+                name=request.name,
+                name_contains=request.name_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(spans, limit=request.limit, offset=request.offset)
+
         @api.post(API_AGL_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
         async def get_next_span_sequence_id(request: NextSequenceIdRequest):  # pyright: ignore[reportUnusedFunction]
             sequence_id = await self.get_next_span_sequence_id(request.rollout_id, request.attempt_id)
@@ -778,9 +847,13 @@ class LightningStoreServer(LightningStore):
             # Handle "latest" keywords BEFORE generic IDs
             if path.endswith("/attempts/latest") and "/rollouts/" in path:
                 return re.sub(r"rollouts/[^/]+/attempts/latest$", "rollouts/{rollout_id}/attempts/latest", path)
-            elif path.endswith("/resources/latest"):
+            if path.endswith("/attempts/search") and "/rollouts/" in path:
+                return re.sub(r"rollouts/[^/]+/attempts/search$", "rollouts/{rollout_id}/attempts/search", path)
+            if path.endswith("/resources/latest"):
                 return path
-            elif "enqueue" in path or "dequeue" in path:
+            if path.endswith("/search"):
+                return path
+            if "enqueue" in path or "dequeue" in path:
                 return path
 
             # Handle generic IDs
@@ -1526,29 +1599,25 @@ class LightningStoreClient(LightningStore):
         status: Optional[Sequence[RolloutStatus]] = None,
         rollout_ids: Optional[Sequence[str]] = None,
     ) -> PaginatedResult[Union[AttemptedRollout, Rollout]]:
-        params_list: List[Tuple[str, Any]] = []
-
-        def _extend(key: str, values: Sequence[Any]) -> None:
-            for value in values:
-                params_list.append((key, value))
-
         resolved_status = status_in if status_in is not None else status
         resolved_rollout_ids = rollout_id_in if rollout_id_in is not None else rollout_ids
 
+        payload: Dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+        }
         if resolved_status is not None:
-            _extend("status_in", resolved_status)
+            payload["status_in"] = resolved_status
         if resolved_rollout_ids is not None:
-            _extend("rollout_id_in", resolved_rollout_ids)
+            payload["rollout_id_in"] = resolved_rollout_ids
         if rollout_id_contains is not None:
-            params_list.append(("rollout_id_contains", rollout_id_contains))
-        params_list.append(("filter_logic", filter_logic))
+            payload["rollout_id_contains"] = rollout_id_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params_list.append(("sort_by", sort_by))
-            params_list.append(("sort_order", sort_order))
-        params_list.append(("limit", limit))
-        params_list.append(("offset", offset))
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
 
-        data = await self._request_json("get", "/rollouts", params=params_list or None)
+        data = await self._request_json("post", "/rollouts/search", json=payload)
         items = [
             (
                 AttemptedRollout.model_validate(item)
@@ -1568,14 +1637,14 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Attempt]:
-        params: List[Tuple[str, Any]] = [
-            ("limit", limit),
-            ("offset", offset),
-        ]
+        payload: Dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+        }
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
-        data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts", params=params)
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
+        data = await self._request_json("post", f"/rollouts/{rollout_id}/attempts/search", json=payload)
         items = [Attempt.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -1795,32 +1864,30 @@ class LightningStoreClient(LightningStore):
         sort_by: Optional[str] = "sequence_id",
         sort_order: Literal["asc", "desc"] = "asc",
     ) -> PaginatedResult[Span]:
-        params: List[Tuple[str, Any]] = [("rollout_id", rollout_id)]
+        payload: Dict[str, Any] = {"rollout_id": rollout_id, "limit": limit, "offset": offset}
         if attempt_id is not None:
-            params.append(("attempt_id", attempt_id))
+            payload["attempt_id"] = attempt_id
         if trace_id is not None:
-            params.append(("trace_id", trace_id))
+            payload["trace_id"] = trace_id
         if trace_id_contains is not None:
-            params.append(("trace_id_contains", trace_id_contains))
+            payload["trace_id_contains"] = trace_id_contains
         if span_id is not None:
-            params.append(("span_id", span_id))
+            payload["span_id"] = span_id
         if span_id_contains is not None:
-            params.append(("span_id_contains", span_id_contains))
+            payload["span_id_contains"] = span_id_contains
         if parent_id is not None:
-            params.append(("parent_id", parent_id))
+            payload["parent_id"] = parent_id
         if parent_id_contains is not None:
-            params.append(("parent_id_contains", parent_id_contains))
+            payload["parent_id_contains"] = parent_id_contains
         if name is not None:
-            params.append(("name", name))
+            payload["name"] = name
         if name_contains is not None:
-            params.append(("name_contains", name_contains))
-        params.append(("filter_logic", filter_logic))
+            payload["name_contains"] = name_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
-        params.append(("limit", limit))
-        params.append(("offset", offset))
-        data = await self._request_json("get", "/spans", params=params)
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
+        data = await self._request_json("post", "/spans/search", json=payload)
         items = [Span.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -1888,21 +1955,17 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Worker]:
-        params: List[Tuple[str, Any]] = [
-            ("limit", limit),
-            ("offset", offset),
-        ]
+        payload: Dict[str, Any] = {}
         if status_in is not None:
-            for value in status_in:
-                params.append(("status_in", value))
+            payload["status_in"] = status_in
         if worker_id_contains is not None:
-            params.append(("worker_id_contains", worker_id_contains))
-        params.append(("filter_logic", filter_logic))
+            payload["worker_id_contains"] = worker_id_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
 
-        data = await self._request_json("get", "/workers", params=params)
+        data = await self._request_json("post", "/workers/search", json=payload)
         items = [Worker.model_validate(item) for item in data.get("items", [])]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 

agentlightning/store/collection/__init__.py

@@ -1,9 +1,21 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import Collection, FilterOptions, KeyValue, LightningCollections, PaginatedResult, Queue, SortOptions
+from .base import (
+    AtomicLabels,
+    AtomicMode,
+    Collection,
+    FilterOptions,
+    KeyValue,
+    LightningCollections,
+    PaginatedResult,
+    Queue,
+    SortOptions,
+)
 from .memory import DequeBasedQueue, DictBasedKeyValue, InMemoryLightningCollections, ListBasedCollection
 
 __all__ = [
+    "AtomicLabels",
+    "AtomicMode",
     "Collection",
     "Queue",
     "KeyValue",

agentlightning/store/collection/base.py

@@ -41,6 +41,15 @@ T = TypeVar("T")  # Recommended to be a BaseModel
 K = TypeVar("K")
 V = TypeVar("V")
 
+AtomicMode = Literal["r", "w", "rw"]
+"""What is expected within the atomic context. Can be "read", "write", or "read-write"."""
+
+AtomicLabels = Literal["rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids"]
+"""Labels for atomic operations.
+
+These labels are used to identify the collections that are affected by the atomic operation.
+"""
+
 
 class Collection(Generic[T]):
     """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
@@ -114,19 +123,42 @@ class Collection(Generic[T]):
         """
         raise NotImplementedError()
 
-    async def update(self, items: Sequence[T]) -> None:
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Update the given items in the collection.
 
+        Args:
+            items: The items to update in the collection.
+            update_fields: The fields to update. If not provided, all fields in the type will be updated.
+                Only applicable if the item type is a Pydantic BaseModel.
+
         Raises:
             ValueError: If an item with the primary keys does not exist.
+
+        Returns:
+            The items that were updated.
         """
         raise NotImplementedError()
 
-    async def upsert(self, items: Sequence[T]) -> None:
+    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Upsert the given items into the collection.
 
         If the items with the same primary keys already exist, they will be updated.
         Otherwise, they will be inserted.
+
+        The operation has three semantics configurable via `update_fields`:
+
+        - `update_or_insert` via `collection.upsert(items, update_fields=["status", "updated_at"])`.
+          If the item with the same primary keys already exists, only the specified fields will be updated.
+          Otherwise, the item will be inserted.
+        - `get_or_insert` via `collection.upsert(items, update_fields=[])`.
+          If the item with the same primary keys already exists, the item will be left unchanged.
+          Otherwise, the item will be inserted.
+        - `replace_ish` via `collection.upsert(items)`.
+          If the item with the same primary keys already exists, all fields from the item will be set.
+          Otherwise, the item will be inserted.
+
+        Returns:
+            The items that were upserted.
         """
         raise NotImplementedError()
 
@@ -265,20 +297,46 @@ class LightningCollections:
         """Dictionary (counter) of span sequence IDs."""
         raise NotImplementedError()
 
-    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[Self]:
+    def atomic(
+        self,
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
+    ) -> AsyncContextManager[Self]:
         """Perform a atomic operation on the collections.
 
         Subclass may use args and kwargs to support multiple levels of atomicity.
+        The arguments can be seen as tags. They only imply the behavior of the operation, not the implementation.
 
         Args:
-            *args: Arguments to pass to the operation.
+            mode: The mode of atomicity. See [`AtomicMode`][agentlightning.store.collection.AtomicMode].
+            snapshot: Enable read snapshot for repeatable reads. Data consistency is guaranteed. The real behavior is implementation-dependent.
+            commit: Enable commitment for write operations. Unsuccessful operations will be rolled back depending on the implementation.
+                Recommend to use [`execute()`][agentlightning.store.collection.LightningCollections.execute] for this level to enable automatic retries.
+                Remember that the real behavior is implementation-dependent.
+            labels: Labels to add to the atomic operation (commonly used as lock names or collection names).
             **kwargs: Keyword arguments to pass to the operation.
         """
         raise NotImplementedError()
 
-    async def execute(self, callback: Callable[[Self], Awaitable[T]]) -> T:
-        """Execute the given callback within an atomic operation."""
-        async with self.atomic() as collections:
+    async def execute(
+        self,
+        callback: Callable[[Self], Awaitable[T]],
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
+    ) -> T:
+        """Execute the given callback within an atomic operation. Retry on transient errors is implied.
+
+        See [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] for more details.
+        """
+        async with self.atomic(mode=mode, snapshot=snapshot, commit=commit, labels=labels, **kwargs) as collections:
             return await callback(collections)
 
 

agentlightning/store/collection/memory.py

@@ -4,9 +4,10 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import threading
 import weakref
 from collections import deque
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager
 from typing import (
     Any,
     Deque,
@@ -24,6 +25,8 @@ from typing import (
     Union,
 )
 
+from pydantic import BaseModel
+
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -37,6 +40,7 @@ from agentlightning.types import (
 )
 
 from .base import (
+    AtomicMode,
     Collection,
     FilterMap,
     KeyValue,
@@ -282,7 +286,7 @@ class ListBasedCollection(Collection[T]):
         # We should always return inside the loop.
         raise RuntimeError("Unreachable")
 
-    def _mutate_single(self, item: T, mode: MutationMode) -> None:
+    def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
         """Core mutation logic shared by insert, update, upsert, and delete."""
         self._ensure_item_type(item)
         key_values = self._extract_primary_key_values(item)
@@ -299,7 +303,35 @@ class ListBasedCollection(Collection[T]):
             else:  # upsert
                 if not exists:
                     self._size += 1
-                parent[final_key] = item
+                    parent[final_key] = item
+
+                elif update_fields is None:
+                    # update_or_insert: update all fields
+                    parent[final_key] = item
+
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+
+                    # Try to fetch the existing item
+                    existing = parent[final_key]
+                    if not isinstance(existing, self._item_type):
+                        raise ValueError(
+                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
+                        )
+
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+
+            return parent[final_key]
 
         elif mode in ("update", "delete"):
             # For update/delete we must not create missing paths.
@@ -314,7 +346,22 @@ class ListBasedCollection(Collection[T]):
                 raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
 
             if mode == "update":
-                parent[final_key] = item
+                if update_fields is None:
+                    # replace the entire item
+                    parent[final_key] = item
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+                return parent[final_key]
             else:  # delete
                 del parent[final_key]
                 self._size -= 1
@@ -554,19 +601,29 @@ class ListBasedCollection(Collection[T]):
         for item in prepared:
             self._mutate_single(item, mode="insert")
 
-    async def update(self, items: Sequence[T]) -> None:
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Update the given items.
 
         Raises:
             ValueError: If any item with the given primary keys does not exist.
         """
+        updated_items: List[T] = []
         for item in items:
-            self._mutate_single(item, mode="update")
+            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
+            if updated is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            updated_items.append(updated)
+        return updated_items
 
-    async def upsert(self, items: Sequence[T]) -> None:
+    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Upsert the given items (insert if missing, otherwise update)."""
+        upserted_items: List[T] = []
         for item in items:
-            self._mutate_single(item, mode="upsert")
+            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
+            if upserted is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            upserted_items.append(upserted)
+        return upserted_items
 
     async def delete(self, items: Sequence[T]) -> None:
         """Delete the given items.
@@ -662,8 +719,16 @@ class InMemoryLightningCollections(LightningCollections):
     Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
     """
 
-    def __init__(self):
-        self._lock = _LoopAwareAsyncLock()
+    def __init__(self, lock_type: Literal["thread", "asyncio"]):
+        self._lock = {
+            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+        }
         self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
         self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
         self._spans = ListBasedCollection(
@@ -703,9 +768,25 @@ class InMemoryLightningCollections(LightningCollections):
         return self._span_sequence_ids
 
     @asynccontextmanager
-    async def atomic(self, *args: Any, **kwargs: Any):
-        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
-        async with self._lock:
+    async def atomic(
+        self, *, mode: AtomicMode = "rw", snapshot: bool = False, labels: Optional[Sequence[str]] = None, **kwargs: Any
+    ):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
+
+        Skip the locking if mode is "r" and snapshot is False.
+
+        This collection implementation does NOT support rollback / commit.
+        """
+        if mode == "r" and not snapshot:
+            yield self
+            return
+        if not labels:
+            # If no labels are provided, use all locks.
+            labels = list(self._lock.keys())
+        managers = [self._lock[label] for label in labels]
+        async with AsyncExitStack() as stack:
+            for manager in managers:
+                await stack.enter_async_context(manager)
             yield self
 
     async def evict_spans_for_rollout(self, rollout_id: str) -> None:
@@ -754,3 +835,27 @@ class _LoopAwareAsyncLock:
         if lock is None or not lock.locked():
             raise RuntimeError("Lock released without being acquired")
         lock.release()
+
+
+class _ThreadSafeAsyncLock:
+    """A threading.Lock that can be used in both async and sync contexts."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+
+    def __enter__(self):
+        self._lock.acquire()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any):
+        self._lock.release()
+
+    async def __aenter__(self):
+        # We run the blocking .acquire() in a thread pool so we don't block the event loop
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self._lock.acquire)
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any):
+        # .release() is non-blocking, so we can call it directly
+        self._lock.release()

agentlightning/store/collection/mongo.py

@@ -34,7 +34,7 @@ if TYPE_CHECKING:
     from typing import Self
 
 from pydantic import BaseModel, TypeAdapter
-from pymongo import AsyncMongoClient, ReadPreference, WriteConcern
+from pymongo import AsyncMongoClient, ReadPreference, ReturnDocument, WriteConcern
 from pymongo.asynchronous.client_session import AsyncClientSession
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.database import AsyncDatabase
@@ -55,6 +55,7 @@ from agentlightning.types import (
 )
 
 from .base import (
+    AtomicMode,
     Collection,
     KeyValue,
     LightningCollections,
@@ -773,34 +774,104 @@ class MongoBasedCollection(Collection[T_model]):
                 raise ValueError("Duplicate key error while inserting items") from exc
 
     @_mongo_operation("update")
-    async def update(self, items: Sequence[T_model]) -> None:
+    async def update(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
         if not items:
-            return
+            return []
 
+        updated_items: List[T_model] = []
         collection = await self.ensure_collection()
+
         for item in items:
             self._ensure_item_type(item)
             pk_filter = self._pk_filter(item)
             doc = item.model_dump()
             doc["partition_id"] = self._partition_id
-            with self._prometheus_tracker.track("update__replace_one", self._database_name, self._collection_name):
-                result = await collection.replace_one(pk_filter, doc, session=self._session)
-            if result.matched_count == 0:
+
+            updated_doc = None
+
+            # Branch 1: Full Replace
+            if update_fields is None:
+                with self._prometheus_tracker.track(
+                    "update__find_one_and_replace", self._database_name, self._collection_name
+                ):
+                    updated_doc = await collection.find_one_and_replace(
+                        filter=pk_filter,
+                        replacement=doc,
+                        session=self._session,
+                        return_document=ReturnDocument.AFTER,  # Returns the new version
+                    )
+
+            # Branch 2: Partial Update
+            else:
+                update_doc = {field: doc[field] for field in update_fields if field in doc}
+                with self._prometheus_tracker.track(
+                    "update__find_one_and_update", self._database_name, self._collection_name
+                ):
+                    updated_doc = await collection.find_one_and_update(
+                        filter=pk_filter,
+                        update={"$set": update_doc},
+                        session=self._session,
+                        return_document=ReturnDocument.AFTER,  # Returns the new version
+                    )
+
+            # Validation and Reconstruction
+            if updated_doc is None:  # type: ignore
                 raise ValueError(f"Item with primary key(s) {pk_filter} does not exist")
 
-    @_mongo_operation("upsert")
-    async def upsert(self, items: Sequence[T_model]) -> None:
-        if not items:
-            return
+            # Re-instantiate the model from the raw MongoDB dictionary.
+            new_item = self._item_type.model_validate(updated_doc)  # type: ignore[arg-type]
+            updated_items.append(new_item)
 
+        return updated_items
+
+    @_mongo_operation("upsert")
+    async def upsert(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
+        if not items:
+            return []
+
+        upserted_items: List[T_model] = []
         collection = await self.ensure_collection()
+
         for item in items:
             self._ensure_item_type(item)
             pk_filter = self._pk_filter(item)
-            doc = item.model_dump()
-            doc["partition_id"] = self._partition_id
-            with self._prometheus_tracker.track("upsert__replace_one", self._database_name, self._collection_name):
-                await collection.replace_one(pk_filter, doc, upsert=True, session=self._session)
+
+            insert_doc = item.model_dump()
+            insert_doc["partition_id"] = self._partition_id
+
+            # If update_fields is None, we update ALL fields (standard upsert behavior).
+            # Otherwise, we only update specific fields, but insert the full doc if it's new.
+            target_fields = update_fields if update_fields is not None else list(insert_doc.keys())
+
+            # 1. $set: Fields that should be overwritten if the document exists
+            update_subset = {field: insert_doc[field] for field in target_fields if field in insert_doc}
+
+            # 2. $setOnInsert: Fields that are only set if we are creating a NEW document
+            # (Everything in the model that isn't in the update_subset)
+            set_on_insert = {k: v for k, v in insert_doc.items() if k not in update_subset}
+
+            update_spec: Dict[str, Dict[str, Any]] = {}
+            if set_on_insert:
+                update_spec["$setOnInsert"] = set_on_insert
+            if update_subset:
+                update_spec["$set"] = update_subset
+
+            with self._prometheus_tracker.track(
+                "upsert__find_one_and_update", self._database_name, self._collection_name
+            ):
+                result_doc = await collection.find_one_and_update(
+                    filter=pk_filter,
+                    update=update_spec,
+                    upsert=True,
+                    session=self._session,
+                    return_document=ReturnDocument.AFTER,
+                )
+
+            # Because upsert=True, result_doc is guaranteed to be not None
+            new_item = self._item_type.model_validate(result_doc)  # type: ignore[arg-type]
+            upserted_items.append(new_item)
+
+        return upserted_items
 
     @_mongo_operation("delete")
     async def delete(self, items: Sequence[T_model]) -> None:
@@ -1327,49 +1398,52 @@ class MongoLightningCollections(LightningCollections):
         self._collection_ensured = True
 
     @asynccontextmanager
-    async def atomic(self, *args: Any, **kwargs: Any):
+    async def atomic(
+        self, mode: AtomicMode = "rw", snapshot: bool = False, commit: bool = False, *args: Any, **kwargs: Any
+    ):
         """Perform a atomic operation on the collections."""
+        if commit:
+            raise ValueError("Commit should be used with execute() instead.")
         with self._prometheus_tracker.track("atomic", self._database_name, self._collection_name):
             # First step: ensure all collections exist before going into the atomic block
             if not self._collection_ensured:
                 await self._ensure_collections()
-            # One session for one transaction
-            client = await self._client_pool.get_client()
-            async with client.start_session() as session:
-                collection_with_session = self.with_session(session)
-                try:
-                    # Start the transaction now
-                    await session.start_transaction(
-                        write_concern=WriteConcern("majority"),
-                        read_concern=ReadConcern("local"),
-                        read_preference=ReadPreference.PRIMARY,
-                    )
-                    yield collection_with_session
-                    # Commit the transaction
-                    await session.commit_transaction()
-                # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
-                except BaseException as exc:
-                    if session.in_transaction:
-                        await session.abort_transaction()
-                    if isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError"):
-                        # NOTE: Retry is via execute
-                        raise RuntimeError("Transaction failed with transient error") from exc
-                    raise
+            # Execute directly without commit
+            yield self
 
     @_mongo_operation("execute")
-    async def execute(self, callback: Callable[[Self], Awaitable[T_generic]]) -> T_generic:
+    async def execute(
+        self,
+        callback: Callable[[Self], Awaitable[T_generic]],
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        **kwargs: Any,
+    ) -> T_generic:
         """Execute the given callback within an atomic operation, and with retries on transient errors."""
         if not self._collection_ensured:
             await self._ensure_collections()
         client = await self._client_pool.get_client()
 
+        # If commit is not turned on, just execute the callback directly.
+        if not commit:
+            return await callback(self)
+
+        # If snapshot is enabled, use snapshot read concern.
+        read_concern = ReadConcern("snapshot") if snapshot else ReadConcern("local")
+        # If mode is "r", write_concern is not needed.
+        write_concern = WriteConcern("majority") if mode != "r" else None
+
         async with client.start_session() as session:
             collections = self.with_session(session)
             with self._prometheus_tracker.track(
                 "execute__transaction", self._database_name, self._collection_name
             ) as tracker:
                 try:
-                    return await self._with_transaction(session, collections, callback, tracker)
+                    return await self._with_transaction(
+                        session, collections, callback, read_concern, write_concern, tracker
+                    )
                 except (ConnectionFailure, OperationFailure) as exc:
                     # Un-retryable errors.
                     tracker.report_error(exc)
@@ -1380,14 +1454,14 @@ class MongoLightningCollections(LightningCollections):
         session: AsyncClientSession,
         collections: Self,
         callback: Callable[[Self], Awaitable[T_generic]],
+        read_concern: ReadConcern,
+        write_concern: Optional[WriteConcern],
         transaction_tracker: _MongoOperationContext | _DummyOperationContext,
     ) -> T_generic:
         # This will start a transaction, run transaction callback, and commit.
         # It will also transparently retry on some transient errors.
         # Expanded implementation of with_transaction from client_session
         num_attempts = 0
-        read_concern = ReadConcern("local")
-        write_concern = WriteConcern("majority")
         read_preference = ReadPreference.PRIMARY
         transaction_retry_time_limit = 120
         start_time = time.monotonic()

agentlightning/store/memory.py

@@ -19,6 +19,7 @@ from typing import (
     Optional,
     Sequence,
     Set,
+    Tuple,
     TypeVar,
     Union,
     cast,
@@ -26,7 +27,7 @@ from typing import (
 
 from pydantic import BaseModel
 
-from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span
+from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
 
 from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
@@ -82,13 +83,18 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def __init__(
         self,
         *,
+        thread_safe: bool = False,
         eviction_memory_threshold: float | int | None = None,
         safe_memory_threshold: float | int | None = None,
         span_size_estimator: Callable[[Span], int] | None = None,
         prometheus: bool = False,
     ):
-        super().__init__(collections=InMemoryLightningCollections(), prometheus=prometheus)
+        super().__init__(
+            collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio"),
+            prometheus=prometheus,
+        )
 
+        self._thread_safe = thread_safe
         self._start_time_by_rollout: Dict[str, float] = {}
         self._span_bytes_by_rollout: Dict[str, int] = Counter()
         self._total_span_bytes: int = 0
@@ -134,7 +140,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
-            thread_safe=False,
+            thread_safe=self._thread_safe,
             async_safe=True,
             zero_copy=False,
             otlp_traces=False,
@@ -153,7 +159,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     @tracked("wait_for_rollout")
     async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
         """Wait for a specific rollout to complete with a timeout."""
-        async with self.collections.atomic() as collections:
+        async with self.collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["rollouts"]) as collections:
             rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
             if rollout and is_finished(rollout):
                 return rollout
@@ -181,47 +187,71 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
             # If event was set (not timeout), check if rollout is finished
             if result:
-                async with self.collections.atomic() as collections:
+                async with self.collections.atomic(
+                    mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
+                ) as collections:
                     rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
                     if rollout and is_finished(rollout):
                         return rollout
 
         return None
 
-    @tracked("on_rollout_update")
-    async def on_rollout_update(self, rollout: Rollout) -> None:
+    @tracked("add_resources_inmemory")
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        ret = await super().add_resources(resources)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
+            self._latest_resources_id = ret.resources_id
+        return ret
+
+    @tracked("update_resources_inmemory")
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        ret = await super().update_resources(resources_id, resources)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
+            self._latest_resources_id = ret.resources_id
+        return ret
+
+    @tracked("_post_update_rollout_inmemory")
+    async def _post_update_rollout(self, rollouts: Sequence[Tuple[Rollout, Sequence[str]]]) -> None:
         """Update the running rollout ids set when the rollout updates."""
-        if is_running(rollout):
-            self._running_rollout_ids.add(rollout.rollout_id)
-        else:
-            self._running_rollout_ids.discard(rollout.rollout_id)
+        await super()._post_update_rollout(rollouts)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["rollouts"]):
+            for rollout, _ in rollouts:
+                if is_running(rollout):
+                    self._running_rollout_ids.add(rollout.rollout_id)
+                else:
+                    self._running_rollout_ids.discard(rollout.rollout_id)
 
-        if is_finished(rollout):
-            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-            self._completion_events[rollout.rollout_id].set()
-        else:
-            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-        # Rollout status can never transition from finished to running (unlike attempt)
-        # so we don't need to clear the completion event even in case of retrying.
+                if is_finished(rollout):
+                    self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+                    self._completion_events[rollout.rollout_id].set()
+                else:
+                    self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+                # Rollout status can never transition from finished to running (unlike attempt)
+                # so we don't need to clear the completion event even in case of retrying.
 
-        if rollout.rollout_id not in self._start_time_by_rollout:
-            self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
+                if rollout.rollout_id not in self._start_time_by_rollout:
+                    self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
 
-    @tracked("get_running_rollouts")
-    async def get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
-        """Accelerated version of `get_running_rollouts` for in-memory store. Used for healthcheck."""
-        rollouts = await collections.rollouts.query(filter={"rollout_id": {"within": list(self._running_rollout_ids)}})
-        running_rollouts: List[AttemptedRollout] = []
-        for rollout in rollouts.items:
-            latest_attempt = await collections.attempts.get(
-                filter={"rollout_id": {"exact": rollout.rollout_id}},
-                sort={"name": "sequence_id", "order": "desc"},
+    @tracked("_unlocked_get_running_rollouts")
+    async def _unlocked_get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
+        """Accelerated version of `_unlocked_get_running_rollouts` for in-memory store. Used for healthcheck."""
+        async with self.collections.atomic(
+            mode="r", snapshot=self._read_snapshot, labels=["rollouts", "attempts"]
+        ) as collections:
+            rollouts = await collections.rollouts.query(
+                filter={"rollout_id": {"within": list(self._running_rollout_ids)}}
             )
-            if not latest_attempt:
-                # The rollout is running but has no attempts, this should not happen
-                logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
-                continue
-            running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
+            running_rollouts: List[AttemptedRollout] = []
+            for rollout in rollouts.items:
+                latest_attempt = await collections.attempts.get(
+                    filter={"rollout_id": {"exact": rollout.rollout_id}},
+                    sort={"name": "sequence_id", "order": "desc"},
+                )
+                if not latest_attempt:
+                    # The rollout is running but has no attempts, this should not happen
+                    logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
+                    continue
+                running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
         return running_rollouts
 
     @tracked("query_spans_inmemory")  # Since this method calls super, we need to track it separately
@@ -235,28 +265,28 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
         return await super().query_spans(rollout_id, attempt_id, **kwargs)
 
-    @tracked("_add_many_spans_unlocked_inmemory")
-    async def _add_many_spans_unlocked(
-        self, collections: InMemoryLightningCollections, rollout_id: str, attempt_id: str, spans: Sequence[Span]
-    ) -> Sequence[Span]:
+    @tracked("_post_add_spans")
+    async def _post_add_spans(self, spans: Sequence[Span], rollout_id: str, attempt_id: str) -> None:
         """In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""
 
-        inserted = await super()._add_many_spans_unlocked(collections, rollout_id, attempt_id, spans)
-        for span in inserted:
-            await self._account_span_size(span)
-        await self._maybe_evict_spans(collections)
+        await super()._post_add_spans(spans, rollout_id, attempt_id)
+        async with self.collections.atomic(
+            mode="rw", snapshot=self._read_snapshot, labels=["rollouts", "spans"]
+        ) as collections:
+            for span in spans:
+                await self._account_span_size(span)
+            await self._maybe_evict_spans(collections)
 
-        return inserted
-
-    @tracked("_get_latest_resources_id")
-    async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
+    @tracked("_get_latest_resources_inmemory")
+    async def _get_latest_resources(self) -> Optional[ResourcesUpdate]:
         if isinstance(self._latest_resources_id, Unset):
-            latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
-            if latest_resources:
-                self._latest_resources_id = latest_resources.resources_id
-            else:
-                self._latest_resources_id = None
-        return self._latest_resources_id
+            return await super()._get_latest_resources()
+        if self._latest_resources_id is not None:
+            async with self.collections.atomic(
+                mode="r", snapshot=self._read_snapshot, labels=["resources"]
+            ) as collections:
+                return await collections.resources.get(filter={"resources_id": {"exact": self._latest_resources_id}})
+        return None
 
     @staticmethod
     def _resolve_memory_threshold(

agentlightning/store/mongo.py

@@ -115,10 +115,13 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         unfinished_rollout_ids = set(rollout_ids)
 
         while deadline is None or current_time <= deadline:
-            # Query the rollouts that are not finished in a single query
-            rollouts = await self.collections.rollouts.query(
-                filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
-            )
+            async with self.collections.atomic(
+                mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
+            ) as collections:
+                # Query the rollouts that are not finished in a single query
+                rollouts = await collections.rollouts.query(
+                    filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
+                )
             for rollout in rollouts.items:
                 if is_finished(rollout):
                     finished_rollouts[rollout.rollout_id] = rollout
@@ -136,15 +139,16 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         # Reorder the rollouts to match the input order
         return [finished_rollouts[rollout_id] for rollout_id in rollout_ids if rollout_id in finished_rollouts]
 
-    @tracked("_many_rollouts_to_attempted_rollouts_unlocked")
-    async def _many_rollouts_to_attempted_rollouts_unlocked(
+    @tracked("_unlocked_many_rollouts_to_attempted_rollouts")
+    async def _unlocked_many_rollouts_to_attempted_rollouts(
         self, collections: MongoLightningCollections, rollouts: Sequence[Rollout]
     ) -> List[Union[Rollout, AttemptedRollout]]:
         """Query the latest attempts for the rollouts, and attach them to the rollout objects."""
-        attempts = await collections.attempts.query(
-            filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
-            sort={"name": "sequence_id", "order": "desc"},
-        )
+        async with collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["attempts"]) as collections:
+            attempts = await collections.attempts.query(
+                filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
+                sort={"name": "sequence_id", "order": "desc"},
+            )
         latest_attempts: Dict[str, Attempt] = {}
         for attempt in attempts:
             if attempt.rollout_id not in latest_attempts:

agentlightning/store/utils.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import time
-from typing import Awaitable, Callable, List, cast
+from typing import Awaitable, Callable, Dict, List, Tuple
 
 from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
 
@@ -57,66 +57,54 @@ LATENCY_BUCKETS = [
 ]
 
 
-async def propagate_status(
-    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
+async def rollout_status_from_attempt(
     attempt: Attempt,
     config: RolloutConfig,
-) -> Rollout:
+) -> RolloutStatus:
     """
     Propagate the status of an attempt to the rollout.
 
-    The rollout should be made sure in a state to be outdated.
-    Requeue the rollout if it should be retried.
-
-    This operation is completely unlocked. The caller is responsible for locking the store.
+    Returns:
+        The status of the rollout from the perspective of the attempt.
     """
     # Propagate the status directly to the rollout
     if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
-        return await update_rollout_status(
-            attempt.rollout_id,
-            attempt.status,
-        )
+        return attempt.status
 
     if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
         # Check if this status should trigger a retry
         if attempt.status in config.retry_condition:
             # If we haven't exceeded max attempts, retry
             if attempt.sequence_id < config.max_attempts:
-                return await update_rollout_status(
-                    attempt.rollout_id,
-                    "requeuing",
-                )
+                return "requeuing"
 
         # If we can't retry or shouldn't retry, mark as failed
-        return await update_rollout_status(
-            attempt.rollout_id,
-            "failed",
-        )
+        return "failed"
 
     raise ValueError(f"Invalid attempt status: {attempt.status}")
 
 
-async def healthcheck(
+async def scan_unhealthy_rollouts(
     rollouts: List[AttemptedRollout],
-    update_rollout_status: UpdateRolloutStatus,
-    update_attempt_status: UpdateAttemptStatus,
-) -> None:
+) -> Dict[Tuple[str, str], AttemptStatus]:
     """
     Perform health check on all running rollouts in the store.
 
     This method should be called periodically to:
 
-    1. Update rollout status to failed to succeeded when the attempt is done
-    2. Check for unresponsive attempts (no heartbeat or spans for a while)
-    3. Check for timed-out rollouts (running too long since start_time)
-    4. Update attempt/rollout status accordingly
+    1. Check for unresponsive attempts (no heartbeat or spans for a while)
+    2. Check for timed-out rollouts (running too long since start_time)
 
     This operation is completely unlocked. The caller is responsible for locking the store.
 
     Args:
-        store: The LightningStore instance to check rollouts from
+        rollouts: The list of running rollouts to check.
+
+    Returns:
+        A dictionary of updates to the rollouts.
     """
     current_time = time.time()
+    updates: Dict[Tuple[str, str], AttemptStatus] = {}
 
     for rollout in rollouts:
         config = rollout.config  # policy for retry and timeout
@@ -124,52 +112,31 @@ async def healthcheck(
         # Get the latest attempt for this rollout
         latest_attempt = rollout.attempt
         if not latest_attempt:
-            continue
-
-        # Check if the attempt has already failed or succeeded
-        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
-            await propagate_status(update_rollout_status, latest_attempt, config)
+            # This should not happen
             continue
 
         # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
         if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "timeout",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
             continue
 
         # Check for unresponsive condition (based on last heartbeat)
-        if latest_attempt.last_heartbeat_time:
-            if latest_attempt.status == "preparing":
-                # If still preparing, mark it as running
-                latest_attempt = await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "running",
-                )
+        # (1) Haven't received heartbeat for a while
+        if (
+            latest_attempt.last_heartbeat_time
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
+        ):
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
 
-            # Haven't received heartbeat for a while
-            if (
-                config.unresponsive_seconds is not None
-                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
-            ):
-                await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "unresponsive",
-                )
-                continue
-
-        # Check if there's no last heartbeat (no spans) at all
+        # (2) Check if there's no last heartbeat (no spans) at all
         if (
             latest_attempt.last_heartbeat_time is None
             and config.unresponsive_seconds is not None
             and current_time - latest_attempt.start_time > config.unresponsive_seconds
         ):
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "unresponsive",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
+
+    return updates

docs/deep-dive/store.md

@@ -117,7 +117,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
 | ------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------- |
 | N/A                                   | `queuing`                                 | Created by `enqueue_rollout()`.                                                                   |
 | `preparing`                           | `queuing/requeuing` → `preparing`         | Typically `dequeue_rollout()` or `start_rollout()`/`start_attempt()` creates a new attempt.       |
-| `running`                             | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `propagate_status`. |
+| `running`                             | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `rollout_status_from_attempt`. |
 | `succeeded`                           | `*` → `succeeded`                         | Terminal. Rollout `end_time` set.                                                                 |
 | `failed` / `timeout` / `unresponsive` | `*` → `requeuing`                         | **Only if** `status ∈ retry_condition ∧ sequence_id < max_attempts`.                              |
 | `failed` / `timeout` / `unresponsive` | `*` → `failed`                            | Otherwise (no retries left or retries disabled).                                                  |
@@ -125,7 +125,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
 
 !!! note "Why aggregation?"
 
-    In code, we use `propagate_status()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
+    In code, we use `rollout_status_from_attempt()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
 
 ## Spans
 

docs/reference/internal.md

@@ -44,7 +44,9 @@
 
 ::: agentlightning.store.base.UNSET
 
-::: agentlightning.store.utils.propagate_status
+::: agentlightning.store.utils.rollout_status_from_attempt
+
+::: agentlightning.store.utils.scan_unhealthy_rollouts
 
 ## Tracing and OpenTelemetry
 

docs/reference/store.md

@@ -20,6 +20,10 @@
 
 ## Collections and Collection Implementations
 
+::: agentlightning.store.collection.AtomicMode
+
+::: agentlightning.store.collection.AtomicLabels
+
 ::: agentlightning.store.collection.Collection
 
 ::: agentlightning.store.collection.Queue

tests/store/test_collection_memory.py

@@ -3,8 +3,22 @@
 from __future__ import annotations
 
 import asyncio
+import threading
 import time
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Mapping, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    Sequence,
+    Tuple,
+    Union,
+)
 from uuid import uuid4
 
 import pytest
@@ -14,12 +28,16 @@ import agentlightning.store.collection.memory as memory_module
 from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
 from agentlightning.store.collection.base import Collection
 from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
+from agentlightning.store.collection.memory import _LoopAwareAsyncLock  # pyright: ignore[reportPrivateUsage]
+from agentlightning.store.collection.memory import _ThreadSafeAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.types import Rollout
 from tests.store.conftest import QueueItem, SampleItem
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.database import AsyncDatabase
 
+    from agentlightning.store.collection.mongo import MongoLightningCollections
+
 
 def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
     return ListBasedCollection(list(items), SampleItem, ("partition", "index"))
@@ -166,6 +184,126 @@ async def test_list_collection_upsert_updates_when_existing(sample_collection: C
     assert fetched == replacement
 
 
+@pytest.mark.asyncio()
+async def test_list_collection_upsert_get_or_insert_semantics(sample_collection: Collection[SampleItem]) -> None:
+    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 2}}
+    original = await sample_collection.get(filters)
+    assert original is not None
+
+    replacement = SampleItem(
+        partition="beta",
+        index=2,
+        name="replacement",
+        status="queued",
+        tags=["patched"],
+        score=999,
+        rank=999,
+        updated_time=99.0,
+        payload={"priority": 99},
+        metadata="replacement",
+    )
+
+    await sample_collection.upsert([replacement], update_fields=[])
+
+    fetched = await sample_collection.get(filters)
+    assert fetched == original and fetched is not None
+    assert fetched.name == original.name
+    assert fetched.status == original.status
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_upsert_updates_selected_fields(sample_collection: Collection[SampleItem]) -> None:
+    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 1}}
+    original = await sample_collection.get(filters)
+    assert original is not None
+
+    incoming = SampleItem(
+        partition="beta",
+        index=1,
+        name="beta-incoming",
+        status="in-progress",
+        tags=["different"],
+        score=-1.0,
+        rank=42,
+        updated_time=123.45,
+        payload={"priority": -1},
+        metadata="incoming",
+    )
+
+    await sample_collection.upsert([incoming], update_fields=["status", "updated_time"])
+
+    fetched = await sample_collection.get(filters)
+    assert fetched is not None
+    assert fetched.status == incoming.status
+    assert fetched.updated_time == incoming.updated_time
+    # Ensure unspecified fields (e.g. name/tags) remain the same as the original document.
+    assert fetched.name == original.name
+    assert fetched.tags == original.tags
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_update_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
+    replacements = [
+        SampleItem(partition="alpha", index=1, name="alpha-new", status="patched"),
+        SampleItem(partition="delta", index=1, name="delta-new", status="patched"),
+    ]
+
+    returned = await sample_collection.update(replacements)
+    assert list(returned) == replacements
+
+    for expected in replacements:
+        fetched = await sample_collection.get(
+            {"partition": {"exact": expected.partition}, "index": {"exact": expected.index}}
+        )
+        assert fetched == expected
+
+    original_beta = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
+    assert original_beta is not None
+
+    partial_payload = SampleItem(
+        partition="beta", index=1, name="ignored", status="partial", metadata="updated-metadata"
+    )
+    partial_returned = await sample_collection.update([partial_payload], update_fields=["status", "metadata"])
+    assert len(partial_returned) == 1
+
+    fetched_partial = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
+    assert fetched_partial == partial_returned[0]
+    assert fetched_partial is not None
+    assert fetched_partial.status == partial_payload.status
+    assert fetched_partial.metadata == partial_payload.metadata
+    assert fetched_partial.name == original_beta.name
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_upsert_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
+    new_item = SampleItem(partition="omega", index=99, name="omega-new", status="queued")
+    inserted = await sample_collection.upsert([new_item])
+    assert list(inserted) == [new_item]
+
+    fetched_new = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 99}})
+    assert fetched_new == new_item
+
+    original_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
+    assert original_existing is not None
+
+    incoming = SampleItem(
+        partition="beta",
+        index=2,
+        name="beta-incoming-new-name",
+        status="processing",
+        tags=["beta", "patched"],
+    )
+    updated = await sample_collection.upsert([incoming], update_fields=["status", "tags"])
+    assert len(updated) == 1
+
+    fetched_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
+    assert fetched_existing == updated[0]
+    assert fetched_existing is not None
+    assert fetched_existing.status == incoming.status
+    assert fetched_existing.tags == incoming.tags
+    assert fetched_existing.name == original_existing.name
+
+
 @pytest.mark.asyncio()
 async def test_list_collection_delete_multiple_items(sample_collection: Collection[SampleItem]) -> None:
     await sample_collection.delete(
@@ -727,6 +865,139 @@ async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyVa
     assert await dict_key_value.size() == 1
 
 
+def test_thread_safe_async_lock_blocks_threads() -> None:
+    lock = _ThreadSafeAsyncLock()
+    allow_second = threading.Event()
+    second_has_lock = threading.Event()
+    release_first = threading.Event()
+
+    def first() -> None:
+        with lock:
+            allow_second.set()
+            release_first.wait()
+
+    def second() -> None:
+        allow_second.wait()
+        with lock:
+            second_has_lock.set()
+
+    t1 = threading.Thread(target=first)
+    t2 = threading.Thread(target=second)
+    t1.start()
+    t2.start()
+
+    assert allow_second.wait(timeout=1)
+    assert not second_has_lock.wait(0.05)
+
+    release_first.set()
+    t1.join(timeout=1)
+    t2.join(timeout=1)
+
+    assert second_has_lock.is_set()
+
+
+@pytest.mark.asyncio()
+async def test_thread_safe_async_lock_serializes_async_tasks() -> None:
+    lock = _ThreadSafeAsyncLock()
+    first_entered = asyncio.Event()
+    release_first = asyncio.Event()
+    second_acquired = asyncio.Event()
+
+    async def first() -> None:
+        async with lock:
+            first_entered.set()
+            await release_first.wait()
+
+    async def second() -> None:
+        await first_entered.wait()
+        async with lock:
+            second_acquired.set()
+
+    task1 = asyncio.create_task(first())
+    task2 = asyncio.create_task(second())
+
+    await first_entered.wait()
+    await asyncio.sleep(0)
+    assert not second_acquired.is_set()
+
+    release_first.set()
+    await asyncio.gather(task1, task2)
+
+    assert second_acquired.is_set()
+
+
+@pytest.mark.asyncio()
+async def test_loop_aware_async_lock_serializes_tasks() -> None:
+    lock = _LoopAwareAsyncLock()
+    first_entered = asyncio.Event()
+    release_first = asyncio.Event()
+    second_acquired = asyncio.Event()
+
+    async def first() -> None:
+        async with lock:
+            first_entered.set()
+            await release_first.wait()
+
+    async def second() -> None:
+        await first_entered.wait()
+        async with lock:
+            second_acquired.set()
+
+    task1 = asyncio.create_task(first())
+    task2 = asyncio.create_task(second())
+
+    await first_entered.wait()
+    await asyncio.sleep(0)
+    assert not second_acquired.is_set()
+
+    release_first.set()
+    await asyncio.gather(task1, task2)
+
+    assert second_acquired.is_set()
+
+
+@pytest.mark.asyncio()
+async def test_loop_aware_async_lock_reuses_loop_specific_lock() -> None:
+    lock = _LoopAwareAsyncLock()
+    first_lock: asyncio.Lock | None = None
+
+    async with lock as acquired:
+        first_lock = acquired
+        assert first_lock.locked()
+
+    assert first_lock is not None and not first_lock.locked()
+
+    async with lock as acquired_again:
+        assert acquired_again is first_lock
+
+
+@pytest.mark.asyncio()
+async def test_loop_aware_async_lock_distinguishes_event_loops() -> None:
+    lock = _LoopAwareAsyncLock()
+    main_loop_lock: asyncio.Lock | None = None
+
+    async with lock as acquired:
+        main_loop_lock = acquired
+
+    locks_from_threads: List[asyncio.Lock] = []
+
+    def _worker() -> None:
+        async def runner() -> None:
+            async with lock as acquired:
+                locks_from_threads.append(acquired)
+
+        asyncio.run(runner())
+
+    worker = threading.Thread(target=_worker)
+    worker.start()
+    worker.join(timeout=2)
+
+    assert worker.is_alive() is False
+    assert main_loop_lock is not None
+    assert locks_from_threads
+    assert locks_from_threads[0] is not main_loop_lock
+
+
 @pytest.mark.asyncio()
 async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data: Dict[str, int]) -> None:
     key_value = DictBasedKeyValue(dict_key_value_data)
@@ -735,6 +1006,55 @@ async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data:
     assert dict_key_value_data == {"alpha": 1, "beta": 2}
 
 
+@pytest.mark.asyncio()
+async def test_inmemory_atomic_read_only_skips_lock() -> None:
+    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")
+
+    class FailingLock:
+        async def __aenter__(self) -> None:
+            raise AssertionError("read-only atomic block should not acquire the lock")
+
+        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+            return None
+
+    collections._lock = {"default": FailingLock()}  # type: ignore[attr-defined]
+
+    async with collections.atomic(mode="r", snapshot=False):
+        # Should complete without touching the failing lock.
+        assert collections.rollouts is not None
+
+
+@pytest.mark.asyncio()
+async def test_inmemory_atomic_snapshot_or_write_acquires_lock() -> None:
+    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")
+
+    class RecordingLock:
+        def __init__(self) -> None:
+            self.enter_count = 0
+            self.exit_count = 0
+
+        async def __aenter__(self) -> None:
+            self.enter_count += 1
+
+        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+            self.exit_count += 1
+
+    lock = RecordingLock()
+    collections._lock = {"default": lock}  # type: ignore[attr-defined]
+
+    async with collections.atomic(mode="rw", snapshot=False):
+        assert collections.attempts is not None
+
+    assert (lock.enter_count, lock.exit_count) == (1, 1)
+
+    lock.enter_count = lock.exit_count = 0
+
+    async with collections.atomic(mode="r", snapshot=True):
+        assert collections.spans is not None
+
+    assert (lock.enter_count, lock.exit_count) == (1, 1)
+
+
 @pytest.mark.mongo
 @pytest.mark.asyncio()
 async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[Any]) -> None:
@@ -891,3 +1211,110 @@ async def test_mongo_ensure_collection_repeats_without_altering_indexes(
                 unique_indexes.append((index["name"], list(index["key"].items())))  # type: ignore
 
         assert unique_indexes == [("uniq_partition_index", [("partition_id", 1), ("index", 1)])]
+
+
+async def _with_mongo_collections(
+    db: AsyncDatabase[Any],
+    callback: Callable[[MongoLightningCollections], Awaitable[Any]],
+) -> Any:
+    from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections
+
+    async with MongoClientPool(db.client) as client_pool:
+        collections = MongoLightningCollections(
+            client_pool=client_pool,
+            database_name=db.name,
+            partition_id=f"partition-{uuid4().hex}",
+        )
+        return await callback(collections)
+
+
+async def _initialize_counter(collections: MongoLightningCollections, key: str) -> None:
+    async def _init(coll: MongoLightningCollections) -> None:
+        await coll.span_sequence_ids.set(key, 0)
+
+    await collections.execute(_init, commit=False)
+
+
+async def _read_counter(collections: MongoLightningCollections, key: str) -> int:
+    async def _read(coll: MongoLightningCollections) -> int:
+        value = await coll.span_sequence_ids.get(key)
+        assert value is not None
+        return value
+
+    return await collections.execute(_read, commit=False)
+
+
+async def _contention_run(
+    collections: MongoLightningCollections,
+    *,
+    key: str,
+    commit: bool,
+    concurrency: int,
+) -> int:
+    read_lock = asyncio.Lock()
+    ready = asyncio.Event()
+    readers_seen = 0
+
+    async def _barrier() -> None:
+        nonlocal readers_seen
+        async with read_lock:
+            readers_seen += 1
+            if readers_seen == concurrency:
+                ready.set()
+        await ready.wait()
+
+    async def worker(_: int) -> None:
+        first_attempt = True
+
+        async def callback(coll: MongoLightningCollections) -> None:
+            nonlocal first_attempt
+            value = await coll.span_sequence_ids.get(key)
+            assert value is not None
+            if first_attempt:
+                await _barrier()
+                first_attempt = False
+            await asyncio.sleep(0)
+            await coll.span_sequence_ids.set(key, value + 1)
+
+        await collections.execute(callback, commit=commit, snapshot=True, mode="rw")
+
+    await asyncio.gather(*(worker(i) for i in range(concurrency)))
+    return await _read_counter(collections, key)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_execute_without_commit_allows_lost_updates(
+    temporary_mongo_database: AsyncDatabase[Any],
+) -> None:
+    async def scenario(collections: MongoLightningCollections) -> None:
+        counter_key = f"counter-{uuid4().hex}"
+        await _initialize_counter(collections, counter_key)
+        final_value = await _contention_run(
+            collections,
+            key=counter_key,
+            commit=False,
+            concurrency=6,
+        )
+        assert final_value == 1
+
+    await _with_mongo_collections(temporary_mongo_database, scenario)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_execute_with_commit_retries_until_success(
+    temporary_mongo_database: AsyncDatabase[Any],
+) -> None:
+    async def scenario(collections: MongoLightningCollections) -> None:
+        counter_key = f"counter-{uuid4().hex}"
+        await _initialize_counter(collections, counter_key)
+        final_value = await _contention_run(
+            collections,
+            key=counter_key,
+            commit=True,
+            concurrency=6,
+        )
+        assert final_value == 6
+
+    await _with_mongo_collections(temporary_mongo_database, scenario)

tests/store/test_implementation.py

@@ -29,6 +29,7 @@ from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
     LLM,
+    Attempt,
     AttemptedRollout,
     Event,
     Link,
@@ -1255,6 +1256,35 @@ async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_r
     assert updated_attempt.last_heartbeat_time is not None
 
 
+@pytest.mark.asyncio
+async def test_spans_promote_preparing_attempt_with_heartbeat(
+    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Spans should set heartbeat time and promote preparing attempts/rollouts to running."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "preparing-heartbeat"})
+    dequeued = await store_fixture.dequeue_rollout()
+    assert dequeued is not None
+
+    attempts_before = await store_fixture.query_attempts(rollout.rollout_id)
+    assert attempts_before
+    attempt_id = attempts_before[0].attempt_id
+    assert attempts_before[0].status == "preparing"
+    assert attempts_before[0].last_heartbeat_time is None
+
+    heartbeat_time = 1234.5
+    monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: heartbeat_time)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.status == "running"
+    assert attempt_after.last_heartbeat_time == heartbeat_time
+
+    rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id)
+    assert rollout_after is not None
+    assert rollout_after.status == "running"
+
+
 @pytest.mark.asyncio
 async def test_unresponsive_attempt_recovers_after_span(
     store_fixture: LightningStore, mock_readable_span: Mock
@@ -1303,6 +1333,77 @@ async def test_running_attempt_updates_heartbeat(
     assert attempt_after_second.last_heartbeat_time == first_heartbeat + 100.0
 
 
+@pytest.mark.asyncio
+async def test_span_post_add_preserves_concurrent_updates(
+    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Concurrent attempt updates should not be clobbered when spans record heartbeats."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "span-concurrency"})
+    dequeued = await store_fixture.dequeue_rollout()
+    assert dequeued is not None
+    attempt_id = dequeued.attempt.attempt_id
+
+    original_post = store_fixture._post_add_spans  # type: ignore
+
+    async def patched_post(spans: List[Span], rollout_id: str, mutated_attempt_id: str) -> None:
+        await store_fixture.update_attempt(rollout_id, mutated_attempt_id, metadata={"concurrent": True})
+        await original_post(spans, rollout_id, mutated_attempt_id)
+
+    monkeypatch.setattr(store_fixture, "_post_add_spans", patched_post)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.metadata == {"concurrent": True}
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_marks_unresponsive_and_updates_worker(
+    store_fixture: LightningStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Healthcheck should mark attempts unresponsive and sync worker state via the helper."""
+
+    class TimeStub:
+        def __init__(self, value: float):
+            self.value = value
+
+        def __call__(self) -> float:
+            return self.value
+
+    time_stub = TimeStub(100.0)
+    monkeypatch.setattr("agentlightning.store.collection_based.time.time", time_stub)
+
+    rollout = await store_fixture.enqueue_rollout(
+        input={"test": "healthcheck-worker"},
+        config=RolloutConfig(unresponsive_seconds=1.0),
+    )
+    dequeued = await store_fixture.dequeue_rollout(worker_id="worker-sync")
+    assert dequeued is not None
+    attempt_id = dequeued.attempt.attempt_id
+    await store_fixture.update_attempt(rollout.rollout_id, attempt_id, worker_id="worker-sync")
+
+    original_sync = store_fixture._sync_workers_with_attempts  # type: ignore
+    sync_calls: List[str] = []
+
+    async def tracking_sync(attempts: Sequence[Attempt]) -> None:
+        for attempt in attempts:
+            sync_calls.append(attempt.attempt_id)
+        await original_sync(attempts)
+
+    monkeypatch.setattr(store_fixture, "_sync_workers_with_attempts", tracking_sync)
+
+    time_stub.value = 105.0
+    await store_fixture.get_rollout_by_id(rollout.rollout_id)
+
+    worker = await store_fixture.get_worker_by_id("worker-sync")
+    assert worker is not None
+    assert worker.status == "unknown"
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.status == "unresponsive"
+    assert sync_calls == [attempt_after.attempt_id]
+
+
 @pytest.mark.asyncio
 async def test_duplicate_span_id_error(
     store_fixture: LightningStore, mock_readable_span: Mock, caplog: pytest.LogCaptureFixture
@@ -2329,7 +2430,7 @@ async def test_concurrent_resource_updates(store_fixture: LightningStore) -> Non
 @pytest.mark.asyncio
 async def test_update_nonexistent_rollout(store_fixture: LightningStore) -> None:
     """Test updating non-existent rollout raises error."""
-    with pytest.raises(ValueError, match="Rollout nonexistent not found"):
+    with pytest.raises(ValueError, match=r"Item.*does not exist"):
         await store_fixture.update_rollout(rollout_id="nonexistent", status="failed")
 
 

tests/store/test_restful.py

@@ -12,7 +12,7 @@ Test categories:
 
 import asyncio
 import contextlib
-from typing import AsyncGenerator, List, Tuple
+from typing import Any, AsyncGenerator, Dict, List, Tuple
 
 import aiohttp
 import pytest
@@ -69,6 +69,77 @@ async def _run_server_with_cors(cors_origins: List[str] | str | None = None):
         await server.stop()
 
 
+async def _request_rollouts_page(
+    session: aiohttp.ClientSession, api_endpoint: str, method: str, payload: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for rollouts and return parsed payload."""
+
+    if method == "get":
+        async with session.get(f"{api_endpoint}/rollouts", params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{api_endpoint}/rollouts/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_attempts_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    rollout_id: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for attempts and return parsed payload."""
+
+    base = f"{api_endpoint}/rollouts/{rollout_id}/attempts"
+    if method == "get":
+        async with session.get(base, params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{base}/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_spans_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    rollout_id: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for spans and return parsed payload."""
+
+    base_params = {"rollout_id": rollout_id}
+    base_params.update(payload)
+    if method == "get":
+        async with session.get(f"{api_endpoint}/spans", params=base_params) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{api_endpoint}/spans/search", json=base_params) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_workers_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for workers and return parsed payload."""
+
+    base = f"{api_endpoint}/workers"
+    if method == "get":
+        async with session.get(base, params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{base}/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
 @pytest_asyncio.fixture
 async def server_client(
     store_fixture: LightningStore,
@@ -207,6 +278,24 @@ async def test_rollouts_pagination_disabled(
         assert len(data["items"]) == 15
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_rollouts_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    """Ensure both GET and POST /search endpoints behave the same for rollouts."""
+
+    server, _client, session, api_endpoint = server_client
+    for i in range(3):
+        await server.enqueue_rollout(input={"index": i})
+
+    data = await _request_rollouts_page(session, api_endpoint, method, {"limit": 2, "offset": 0})
+    assert data["total"] == 3
+    assert data["limit"] == 2
+    assert data["offset"] == 0
+    assert len(data["items"]) == 2
+
+
 @pytest.mark.asyncio
 async def test_rollouts_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -481,6 +570,22 @@ async def test_attempts_pagination_basic(
         assert len(data["items"]) == 5
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_attempts_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    rollout = await server.enqueue_rollout(input={"test": "attempt-search"})
+    await server.start_attempt(rollout.rollout_id)
+
+    data = await _request_attempts_page(session, api_endpoint, rollout.rollout_id, method, {"limit": 1, "offset": 0})
+    assert data["total"] == 1
+    assert data["limit"] == 1
+    assert len(data["items"]) == 1
+
+
 @pytest.mark.asyncio
 async def test_attempts_sorting(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -746,6 +851,30 @@ async def test_spans_pagination_basic(
         assert len(data["items"]) == 10
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_spans_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    attempted = await server.start_rollout(input={"test": "span-search"})
+    attempt_id = attempted.attempt.attempt_id
+    for seq in range(3):
+        await server.add_span(_make_span(attempted.rollout_id, attempt_id, seq + 1, f"span-{seq}"))
+
+    data = await _request_spans_page(
+        session,
+        api_endpoint,
+        attempted.rollout_id,
+        method,
+        {"attempt_id": attempt_id, "limit": 2, "offset": 0},
+    )
+    assert data["total"] == 3
+    assert data["limit"] == 2
+    assert len(data["items"]) == 2
+
+
 @pytest.mark.asyncio
 async def test_spans_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -1116,6 +1245,50 @@ async def test_request_json_spans_returns_pagination_metadata(
     assert len(data["items"]) == 1
 
 
+# Update semantics tests
+
+
+@pytest.mark.asyncio
+async def test_update_rollout_distinguishes_unset_fields(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    rollout = await server.enqueue_rollout(input={"id": "unset-check"})
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
+        json={"metadata": {"foo": "bar"}},
+    ) as resp:
+        assert resp.status == 200
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
+        json={"status": None},
+    ) as resp:
+        assert resp.status == 400
+
+
+@pytest.mark.asyncio
+async def test_update_attempt_distinguishes_unset_fields(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    attempted = await server.start_rollout(input={"id": "attempt-unset"})
+    attempt_id = attempted.attempt.attempt_id
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
+        json={"status": "running"},
+    ) as resp:
+        assert resp.status == 200
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
+        json={"worker_id": None},
+    ) as resp:
+        assert resp.status == 400
+
+
 # Client Compatibility Tests
 
 
@@ -1162,6 +1335,27 @@ async def test_client_query_with_filters(
     assert rollouts[0].rollout_id == r2.rollout_id
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_workers_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    await server.update_worker("worker-1", heartbeat_stats={"cpu": 0.5})
+    await server.update_worker("worker-2", heartbeat_stats={"cpu": 0.7})
+
+    data = await _request_workers_page(
+        session,
+        api_endpoint,
+        method,
+        {"limit": 1, "offset": 0, "sort_by": "worker_id", "sort_order": "asc"},
+    )
+    assert data["total"] == 2
+    assert data["limit"] == 1
+    assert len(data["items"]) == 1
+
+
 @pytest.mark.asyncio
 async def test_workers_endpoint_supports_updates(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],

tests/store/test_utils.py

@@ -2,11 +2,11 @@
 
 import time
 from typing import List, Optional, cast
-from unittest.mock import AsyncMock, patch
+from unittest.mock import patch
 
 import pytest
 
-from agentlightning.store.utils import healthcheck, propagate_status
+from agentlightning.store.utils import rollout_status_from_attempt, scan_unhealthy_rollouts
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -14,11 +14,11 @@ from agentlightning.types import (
     RolloutConfig,
 )
 
-# Tests for propagate_status function
+# Tests for rollout_status_from_attempt function
 
 
 @pytest.mark.parametrize(
-    "status,expected_call",
+    "status,expected_status",
     [
         ("preparing", "preparing"),
         ("running", "running"),
@@ -26,21 +26,22 @@ from agentlightning.types import (
     ],
 )
 @pytest.mark.asyncio
-async def test_propagate_status_direct_statuses(status: AttemptStatus, expected_call: AttemptStatus) -> None:
-    """Test propagate_status directly propagates preparing/running/succeeded statuses."""
+async def test_rollout_status_from_attempt_direct_statuses(
+    status: AttemptStatus, expected_status: AttemptStatus
+) -> None:
+    """Test rollout_status_from_attempt directly propagates preparing/running/succeeded statuses."""
     attempt = Attempt(
         rollout_id="test-rollout", attempt_id="test-attempt", sequence_id=1, start_time=time.time(), status=status
     )
     config = RolloutConfig()
-    update_rollout_mock = AsyncMock()
 
-    await propagate_status(update_rollout_mock, attempt, config)
+    result = await rollout_status_from_attempt(attempt, config)
 
-    update_rollout_mock.assert_called_once_with("test-rollout", expected_call)
+    assert result == expected_status
 
 
 @pytest.mark.parametrize(
-    "status,in_retry_condition,sequence_id,max_attempts,expected_call",
+    "status,in_retry_condition,sequence_id,max_attempts,expected_status",
     [
         ("failed", True, 1, 3, "requeuing"),  # Should retry
         ("failed", True, 3, 3, "failed"),  # Max attempts reached
@@ -51,10 +52,10 @@ async def test_propagate_status_direct_statuses(status: AttemptStatus, expected_
     ],
 )
 @pytest.mark.asyncio
-async def test_propagate_status_retry_logic(
-    status: AttemptStatus, in_retry_condition: bool, sequence_id: int, max_attempts: int, expected_call: AttemptStatus
+async def test_rollout_status_from_attempt_retry_logic(
+    status: AttemptStatus, in_retry_condition: bool, sequence_id: int, max_attempts: int, expected_status: AttemptStatus
 ) -> None:
-    """Test propagate_status retry logic for different combinations."""
+    """Test rollout_status_from_attempt retry logic for different combinations."""
     attempt = Attempt(
         rollout_id="test-rollout",
         attempt_id="test-attempt",
@@ -65,16 +66,15 @@ async def test_propagate_status_retry_logic(
 
     retry_condition: List[AttemptStatus] = [status] if in_retry_condition else []
     config = RolloutConfig(max_attempts=max_attempts, retry_condition=retry_condition)
-    update_rollout_mock = AsyncMock()
 
-    await propagate_status(update_rollout_mock, attempt, config)
+    result = await rollout_status_from_attempt(attempt, config)
 
-    update_rollout_mock.assert_called_once_with("test-rollout", expected_call)
+    assert result == expected_status
 
 
 @pytest.mark.asyncio
-async def test_propagate_status_invalid_status() -> None:
-    """Test propagate_status raises error for invalid status."""
+async def test_rollout_status_from_attempt_invalid_status() -> None:
+    """Test rollout_status_from_attempt raises error for invalid status."""
     # Create a valid attempt first, then modify its status
     attempt = Attempt(
         rollout_id="test-rollout", attempt_id="test-attempt", sequence_id=1, start_time=time.time(), status="failed"
@@ -83,31 +83,25 @@ async def test_propagate_status_invalid_status() -> None:
     attempt.status = cast(AttemptStatus, "invalid_status")  # Invalid status
 
     config = RolloutConfig()
-    update_rollout_mock = AsyncMock()
 
     with pytest.raises(ValueError, match="Invalid attempt status: invalid_status"):
-        await propagate_status(update_rollout_mock, attempt, config)
+        await rollout_status_from_attempt(attempt, config)
 
 
-# Tests for healthcheck function
+# Tests for scan_unhealthy_rollouts function
 
 
 @pytest.mark.asyncio
-async def test_healthcheck_empty_rollouts_list() -> None:
-    """Test healthcheck handles empty rollouts list gracefully."""
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
+async def test_scan_unhealthy_rollouts_empty_list() -> None:
+    """Test scan_unhealthy_rollouts handles empty rollouts list gracefully."""
+    updates = await scan_unhealthy_rollouts([])
 
-    await healthcheck([], update_rollout_mock, update_attempt_mock)
-
-    # Should not call any updates
-    update_rollout_mock.assert_not_called()
-    update_attempt_mock.assert_not_called()
+    assert updates == {}
 
 
 @pytest.mark.asyncio
-async def test_healthcheck_multiple_rollouts_different_timeouts() -> None:
-    """Test healthcheck handles multiple rollouts with different timeout configs."""
+async def test_scan_unhealthy_rollouts_multiple_rollouts_different_timeouts() -> None:
+    """Test scan_unhealthy_rollouts handles multiple rollouts with different timeout configs."""
     current_time = time.time()
 
     # Rollout 1: Short timeout, should timeout
@@ -146,14 +140,11 @@ async def test_healthcheck_multiple_rollouts_different_timeouts() -> None:
         attempt=attempt2,
     )
 
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
-
     with patch("time.time", return_value=current_time):
-        await healthcheck([rollout1, rollout2], update_rollout_mock, update_attempt_mock)
+        updates = await scan_unhealthy_rollouts([rollout1, rollout2])
 
     # Only rollout1 should be marked as timeout
-    update_attempt_mock.assert_called_once_with("rollout-1", "attempt-1", "timeout")
+    assert updates == {("rollout-1", "attempt-1"): "timeout"}
 
 
 @pytest.mark.parametrize(
@@ -167,13 +158,13 @@ async def test_healthcheck_multiple_rollouts_different_timeouts() -> None:
     ],
 )
 @pytest.mark.asyncio
-async def test_healthcheck_timeout_configurations(
+async def test_scan_unhealthy_rollouts_timeout_configurations(
     timeout_seconds: Optional[float],
     unresponsive_seconds: Optional[float],
     should_timeout: bool,
     should_unresponsive: bool,
 ) -> None:
-    """Test healthcheck with various timeout configurations."""
+    """Test scan_unhealthy_rollouts with various timeout configurations."""
     current_time = time.time()
 
     config = RolloutConfig(timeout_seconds=timeout_seconds, unresponsive_seconds=unresponsive_seconds)
@@ -196,22 +187,20 @@ async def test_healthcheck_timeout_configurations(
         attempt=attempt,
     )
 
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
-
     with patch("time.time", return_value=current_time):
-        await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
+        updates = await scan_unhealthy_rollouts([rollout])
 
+    expected_updates = {}
     if should_timeout:
-        update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "timeout")
+        expected_updates[("test-rollout", "test-attempt")] = "timeout"
     elif should_unresponsive:
-        update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "unresponsive")
-    else:
-        update_attempt_mock.assert_not_called()
+        expected_updates[("test-rollout", "test-attempt")] = "unresponsive"
+
+    assert updates == expected_updates
 
 
 @pytest.mark.asyncio
-async def test_healthcheck_unresponsive_with_heartbeat_timing() -> None:
+async def test_scan_unhealthy_rollouts_unresponsive_with_heartbeat_timing() -> None:
     """Test unresponsive detection considers heartbeat timing correctly."""
     current_time = time.time()
     config = RolloutConfig(unresponsive_seconds=1.0)
@@ -252,51 +241,16 @@ async def test_healthcheck_unresponsive_with_heartbeat_timing() -> None:
         attempt=attempt_old,
     )
 
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
-
     with patch("time.time", return_value=current_time):
-        await healthcheck([rollout_recent, rollout_old], update_rollout_mock, update_attempt_mock)
+        updates = await scan_unhealthy_rollouts([rollout_recent, rollout_old])
 
     # Only the old heartbeat should trigger unresponsive
-    update_attempt_mock.assert_called_once_with("rollout-old", "attempt-old", "unresponsive")
+    assert updates == {("rollout-old", "attempt-old"): "unresponsive"}
 
 
 @pytest.mark.asyncio
-async def test_healthcheck_preparing_with_heartbeat_promotion() -> None:
-    """Test healthcheck promotes preparing attempts with heartbeat to running."""
-    current_time = time.time()
-
-    config = RolloutConfig()
-    attempt = Attempt(
-        rollout_id="test-rollout",
-        attempt_id="test-attempt",
-        sequence_id=1,
-        start_time=current_time,
-        status="preparing",
-        last_heartbeat_time=current_time,  # Has heartbeat
-    )
-    rollout = AttemptedRollout(
-        rollout_id="test-rollout",
-        input={"test": 1},
-        status="preparing",
-        start_time=current_time,
-        config=config,
-        attempt=attempt,
-    )
-
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
-
-    await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
-
-    # Should promote to running
-    update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "running")
-
-
-@pytest.mark.asyncio
-async def test_healthcheck_skips_rollouts_without_attempts() -> None:
-    """Test healthcheck gracefully skips rollouts with no attempts."""
+async def test_scan_unhealthy_rollouts_skips_rollouts_without_attempts() -> None:
+    """Test scan_unhealthy_rollouts gracefully skips rollouts with no attempts."""
     config = RolloutConfig()
 
     # Create a valid attempt first, then set it to None
@@ -315,11 +269,7 @@ async def test_healthcheck_skips_rollouts_without_attempts() -> None:
     # Bypass Pydantic validation by directly setting the attribute
     rollout.attempt = cast(Attempt, None)  # No attempt
 
-    update_rollout_mock = AsyncMock()
-    update_attempt_mock = AsyncMock()
+    updates = await scan_unhealthy_rollouts([rollout])
 
-    await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
-
-    # Should not call any updates
-    update_rollout_mock.assert_not_called()
-    update_attempt_mock.assert_not_called()
+    # Should not include rollout without attempts
+    assert updates == {}