Merge branch 'main' of github.com:microsoft/agent-lightning into feature/collection

agentlightning/emitter/reward.py

@@ -129,12 +129,13 @@ def reward(fn: FnType) -> FnType:
         return wrapper  # type: ignore
 
 
-def emit_reward(reward: float) -> ReadableSpan:
+def emit_reward(reward: float, auto_export: bool = True) -> ReadableSpan:
     """Emit a reward value as an OpenTelemetry span.
 
     Args:
         reward: Numeric reward to record. Integers and booleans are converted to
             floating point numbers for consistency.
+        auto_export: Whether to export the span automatically.
 
     Returns:
         Readable span capturing the recorded reward.
@@ -150,7 +151,7 @@ def emit_reward(reward: float) -> ReadableSpan:
         raise ValueError(f"Reward must be a number, got: {type(reward)}")
 
     # TODO: This should use the tracer from current context by tracer
-    tracer = get_tracer()
+    tracer = get_tracer(use_active_span_processor=auto_export)
     span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
     # Do nothing; it's just a number
     with span:

agentlightning/emitter/utils.py

@@ -2,13 +2,22 @@
 
 """Utilities shared across emitter implementations."""
 
+from typing import cast
+from warnings import filterwarnings
+
 import opentelemetry.trace as trace_api
+from opentelemetry.sdk.trace import SpanLimits, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
 from opentelemetry.trace import get_tracer_provider
 
 
-def get_tracer() -> trace_api.Tracer:
+def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
     """Resolve the OpenTelemetry tracer configured for Agent Lightning.
 
+    Args:
+        use_active_span_processor: Whether to use the active span processor.
+
     Returns:
         OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
 
@@ -18,5 +27,31 @@ def get_tracer() -> trace_api.Tracer:
     if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
         raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
 
-    tracer_provider = get_tracer_provider()
-    return tracer_provider.get_tracer("agentlightning")
+    tracer_provider = cast(TracerProviderImpl, get_tracer_provider())
+
+    if use_active_span_processor:
+        return tracer_provider.get_tracer("agentlightning")
+
+    else:
+        filterwarnings(
+            "ignore",
+            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
+            category=DeprecationWarning,
+            module="opentelemetry.sdk.trace",
+        )
+
+        return Tracer(
+            tracer_provider.sampler,
+            tracer_provider.resource,
+            # We use an empty span processor to avoid emitting spans to the tracer
+            SynchronousMultiSpanProcessor(),
+            tracer_provider.id_generator,
+            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
+            SpanLimits(),
+            InstrumentationScope(
+                "agentlightning",
+                "",
+                "",
+                {},
+            ),
+        )

agentlightning/runner/agent.py

@@ -279,8 +279,9 @@ class LitAgentRunner(Runner[T_task]):
         if isinstance(raw_result, float):
             # Preserve the existing spans before another span is emitted
             trace_spans = list(self._tracer.get_last_trace())
-            # This will emit another span to the tracer
-            reward_span = emit_reward(raw_result)
+            # This will NOT emit another span to the tracer
+            reward_span = emit_reward(raw_result, auto_export=False)
+            # We add it to the store manually
             await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
             trace_spans.append(reward_span)
 

agentlightning/store/__init__.py

@@ -2,6 +2,7 @@
 
 from .base import LightningStore, LightningStoreCapabilities
 from .client_server import LightningStoreClient, LightningStoreServer
+from .collection_based import CollectionBasedLightningStore
 from .memory import InMemoryLightningStore
 from .threading import LightningStoreThreaded
 
@@ -11,5 +12,6 @@ __all__ = [
     "LightningStoreClient",
     "LightningStoreServer",
     "InMemoryLightningStore",
+    "CollectionBasedLightningStore",
     "LightningStoreThreaded",
 ]

agentlightning/store/collection/__init__.py

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

agentlightning/store/collection/base.py

@@ -0,0 +1,329 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import (
+    Any,
+    AsyncContextManager,
+    Generic,
+    Iterable,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    TypedDict,
+    TypeVar,
+    Union,
+)
+
+from pydantic import BaseModel
+
+from agentlightning.types import (
+    Attempt,
+    ResourcesUpdate,
+    Rollout,
+    Span,
+    Worker,
+)
+
+T = TypeVar("T")  # Recommended to be a BaseModel
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class FilterField(TypedDict, total=False):
+    """An operator dict for a single field."""
+
+    exact: Any
+    within: Iterable[Any]
+    contains: str
+
+
+FilterOptions = Mapping[Union[str, Literal["_aggregate"]], Union[FilterField, Literal["and", "or"]]]
+
+"""A mapping of field name -> operator dict.
+
+Each operator dict can contain:
+
+- "exact": value for exact equality.
+- "within": iterable of allowed values.
+- "contains": substring to search for in string fields.
+
+Example:
+
+```json
+{
+    "_aggregate": "or",
+    "status": {"exact": "active"},
+    "id": {"within": [1, 2, 3]},
+    "name": {"contains": "foo"},
+}
+```
+
+The filter can also have a special field called "_aggregate" that can be used to specify the logic
+to combine the results of the filters:
+
+- "and": all conditions must match. This is the default value if not specified.
+- "or": at least one condition must match.
+
+All conditions within a field and between different fields are
+stored in a unified pool and combined using `_aggregate`.
+"""
+
+
+class SortOptions(TypedDict):
+    """Options for sorting the collection."""
+
+    name: str
+    """The name of the field to sort by."""
+    order: Literal["asc", "desc"]
+    """The order to sort by."""
+
+
+class PaginatedResult(BaseModel, Generic[T]):
+    """Result of a paginated query."""
+
+    items: Sequence[T]
+    """Items in the result."""
+    limit: int
+    """Limit of the result."""
+    offset: int
+    """Offset of the result."""
+    total: int
+    """Total number of items in the collection."""
+
+
+class Collection(Generic[T]):
+    """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
+
+    def primary_keys(self) -> Sequence[str]:
+        """Get the primary keys of the collection."""
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the collection."""
+        raise NotImplementedError()
+
+    def size(self) -> int:
+        """Get the number of items in the collection."""
+        raise NotImplementedError()
+
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with the given filters, sort order, and pagination.
+
+        Args:
+            filter:
+                The filters to apply to the collection. See [`FilterOptions`][agentlightning.store.collection.FilterOptions].
+
+            sort:
+                The options for sorting the collection. See [`SortOptions`][agentlightning.store.collection.SortOptions].
+                The field must exist in the model. If field might contain null values, in which case the behavior is undefined
+                (i.e., depending on the implementation).
+
+            limit:
+                Max number of items to return. Use -1 for "no limit".
+
+            offset:
+                Number of items to skip from the start of the *matching* items.
+
+        Returns:
+            PaginatedResult with items, limit, offset, and total matched items.
+        """
+        raise NotImplementedError()
+
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Get the first item that matches the given filters.
+
+        Args:
+            filter: The filters to apply to the collection.
+                See [`FilterOptions`][agentlightning.store.collection.FilterOptions].
+            sort: Sort options. See [`SortOptions`][agentlightning.store.collection.SortOptions].
+
+        Returns:
+            The first item that matches the given filters, or None if no item matches.
+        """
+        raise NotImplementedError()
+
+    async def insert(self, items: Sequence[T]) -> None:
+        """Add the given items to the collection.
+
+        Raises:
+            ValueError: If an item with the same primary key already exists.
+        """
+        raise NotImplementedError()
+
+    async def update(self, items: Sequence[T]) -> None:
+        """Update the given items in the collection.
+
+        Raises:
+            ValueError: If an item with the primary keys does not exist.
+        """
+        raise NotImplementedError()
+
+    async def upsert(self, items: Sequence[T]) -> None:
+        """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.
+        """
+        raise NotImplementedError()
+
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items from the collection.
+
+        Args:
+            items: The items to delete from the collection.
+
+        Raises:
+            ValueError: If the items with the primary keys to be deleted do not exist.
+        """
+        raise NotImplementedError()
+
+
+class Queue(Generic[T]):
+    """Behaves like a deque. Supporting appending items to the end and popping items from the front."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the queue."""
+        raise NotImplementedError()
+
+    async def has(self, item: T) -> bool:
+        """Check if the given item is in the queue."""
+        raise NotImplementedError()
+
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        """Append the given items to the end of the queue.
+
+        Args:
+            items: The items to append to the end of the queue.
+
+        Returns:
+            The items that were appended to the end of the queue.
+        """
+        raise NotImplementedError()
+
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        """Pop the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to pop from the front of the queue.
+
+        Returns:
+            The items that were popped from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        """Peek the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to peek from the front of the queue.
+
+        Returns:
+            The items that were peeked from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    def size(self) -> int:
+        """Get the number of items in the queue."""
+        raise NotImplementedError()
+
+
+class KeyValue(Generic[K, V]):
+    """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__} ({self.size()})>"
+
+    async def has(self, key: K) -> bool:
+        """Check if the given key is in the dictionary."""
+        raise NotImplementedError()
+
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        """Get the value for the given key, or the default value if the key is not found."""
+        raise NotImplementedError()
+
+    async def set(self, key: K, value: V) -> None:
+        """Set the value for the given key."""
+        raise NotImplementedError()
+
+    async def pop(self, key: K, default: V | None = None) -> V | None:
+        """Pop the value for the given key, or the default value if the key is not found."""
+        raise NotImplementedError()
+
+    def size(self) -> int:
+        """Get the number of items in the dictionary."""
+        raise NotImplementedError()
+
+
+class LightningCollections:
+    """Collections of rollouts, attempts, spans, resources, and workers.
+
+    [LightningStore][agentlightning.LightningStore] implementations can use this as a storage base
+    to implement the store API.
+    """
+
+    @property
+    def rollouts(self) -> Collection[Rollout]:
+        """Collections of rollouts."""
+        raise NotImplementedError()
+
+    @property
+    def attempts(self) -> Collection[Attempt]:
+        """Collections of attempts."""
+        raise NotImplementedError()
+
+    @property
+    def spans(self) -> Collection[Span]:
+        """Collections of spans."""
+        raise NotImplementedError()
+
+    @property
+    def resources(self) -> Collection[ResourcesUpdate]:
+        """Collections of resources."""
+        raise NotImplementedError()
+
+    @property
+    def workers(self) -> Collection[Worker]:
+        """Collections of workers."""
+        raise NotImplementedError()
+
+    @property
+    def rollout_queue(self) -> Queue[str]:
+        """Queue of rollouts (tasks)."""
+        raise NotImplementedError()
+
+    @property
+    def span_sequence_ids(self) -> KeyValue[str, int]:
+        """Dictionary (counter) of span sequence IDs."""
+        raise NotImplementedError()
+
+    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[None]:
+        """Perform a atomic operation on the collections.
+
+        Subclass may use args and kwargs to support multiple levels of atomicity.
+
+        Args:
+            *args: Arguments to pass to the operation.
+            **kwargs: Keyword arguments to pass to the operation.
+        """
+        raise NotImplementedError()

agentlightning/store/collection/memory.py

@@ -0,0 +1,731 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import weakref
+from collections import deque
+from contextlib import asynccontextmanager
+from typing import (
+    Any,
+    AsyncGenerator,
+    Deque,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from agentlightning.types import (
+    Attempt,
+    ResourcesUpdate,
+    Rollout,
+    Span,
+    Worker,
+)
+
+from .base import (
+    Collection,
+    FilterField,
+    FilterOptions,
+    KeyValue,
+    LightningCollections,
+    PaginatedResult,
+    Queue,
+    SortOptions,
+)
+
+T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
+K = TypeVar("K")
+V = TypeVar("V")
+
+# Nested structure type:
+# dict[pk1] -> dict[pk2] -> ... -> item
+ListBasedCollectionItemType = Union[
+    Dict[Any, "ListBasedCollectionItemType[T]"],  # intermediate node
+    Dict[Any, T],  # leaf node dictionary
+]
+
+FilterMap = Mapping[str, FilterField]
+MutationMode = Literal["insert", "update", "upsert", "delete"]
+
+
+class ListBasedCollection(Collection[T]):
+    """In-memory implementation of Collection using a nested dict for O(1) primary-key lookup.
+
+    The internal structure is:
+
+        {
+            pk1_value: {
+                pk2_value: {
+                    ...
+                        pkN_value: item
+                }
+            }
+        }
+
+    where the nesting depth equals the number of primary keys.
+
+    Sorting behavior:
+
+    1. If no sort_by is provided, the items are returned in the order of insertion.
+    2. If sort_by is provided, the items are sorted by the value of the sort_by field.
+    3. If the sort_by field is a timestamp, the null values are treated as infinity.
+    4. If the sort_by field is not a timestamp, the null values are treated as empty string
+       if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
+    """
+
+    def __init__(self, items: List[T], item_type: Type[T], primary_keys: Sequence[str]):
+        if not primary_keys:
+            raise ValueError("primary_keys must be non-empty")
+
+        self._items: Dict[Any, Any] = {}
+        self._size: int = 0
+        if issubclass(item_type, dict):
+            raise TypeError(f"Expect item to be not a dict, got {item_type.__name__}")
+        self._item_type: Type[T] = item_type
+        self._primary_keys: Tuple[str, ...] = tuple(primary_keys)
+
+        # Pre-populate the collection with the given items.
+        for item in items or []:
+            self._mutate_single(item, mode="insert")
+
+    def primary_keys(self) -> Sequence[str]:
+        """Return the primary key field names for this collection."""
+        return self._primary_keys
+
+    def item_type(self) -> Type[T]:
+        """Return the Pydantic model type of items stored in this collection."""
+        return self._item_type
+
+    def size(self) -> int:
+        """Return the number of items stored in the collection."""
+        return self._size
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    def _ensure_item_type(self, item: T) -> None:
+        """Validate that the item matches the declared item_type."""
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, " f"got {type(item).__name__}")
+
+    def _extract_primary_key_values(self, item: T) -> Tuple[Any, ...]:
+        """Extract the primary key values from an item.
+
+        Raises:
+            ValueError: If any primary key is missing on the item.
+        """
+        values: List[Any] = []
+        for key in self._primary_keys:
+            if not hasattr(item, key):
+                raise ValueError(f"Item {item} does not have primary key field '{key}'")
+            values.append(getattr(item, key))
+        return tuple(values)
+
+    def _render_key_values(self, key_values: Sequence[Any]) -> str:
+        return ", ".join(f"{name}={value!r}" for name, value in zip(self._primary_keys, key_values))
+
+    def _locate_node(
+        self,
+        key_values: Sequence[Any],
+        create_missing: bool,
+    ) -> Tuple[MutableMapping[Any, Any], Any]:
+        """Locate the parent mapping and final key for an item path.
+
+        Args:
+            key_values: The sequence of primary key values.
+            create_missing: Whether to create intermediate dictionaries as needed.
+
+        Returns:
+            (parent_mapping, final_key)
+
+        Raises:
+            KeyError: If the path does not exist and create_missing is False.
+            ValueError: If the internal structure is corrupted (non-dict where dict is expected).
+        """
+        if not key_values:
+            raise ValueError("key_values must be non-empty")
+
+        current: MutableMapping[Any, Any] = self._items
+        for idx, value in enumerate(key_values):
+            is_last = idx == len(key_values) - 1
+            if is_last:
+                # At the final level, current[value] is the item (or will be).
+                return current, value  # type: ignore
+
+            # Intermediate level: current[value] must be a dict.
+            if value not in current:
+                if not create_missing:
+                    raise KeyError(f"Path does not exist for given primary keys: {self._render_key_values(key_values)}")
+                current[value] = {}
+            next_node = current[value]  # type: ignore
+            if not isinstance(next_node, dict):
+                raise ValueError(f"Internal structure corrupted: expected dict, got {type(next_node)!r}")  # type: ignore
+            current = next_node  # type: ignore
+
+        # We should always return inside the loop.
+        raise RuntimeError("Unreachable")
+
+    def _mutate_single(self, item: T, mode: MutationMode) -> None:
+        """Core mutation logic shared by insert, update, upsert, and delete."""
+        self._ensure_item_type(item)
+        key_values = self._extract_primary_key_values(item)
+
+        if mode in ("insert", "upsert"):
+            parent, final_key = self._locate_node(key_values, create_missing=True)
+            exists = final_key in parent
+
+            if mode == "insert":
+                if exists:
+                    raise ValueError(f"Item already exists with primary key(s): {self._render_key_values(key_values)}")
+                parent[final_key] = item
+                self._size += 1
+            else:  # upsert
+                if not exists:
+                    self._size += 1
+                parent[final_key] = item
+
+        elif mode in ("update", "delete"):
+            # For update/delete we must not create missing paths.
+            try:
+                parent, final_key = self._locate_node(key_values, create_missing=False)
+            except KeyError:
+                raise ValueError(
+                    f"Item does not exist with primary key(s): {self._render_key_values(key_values)}"
+                ) from None
+
+            if final_key not in parent:
+                raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
+
+            if mode == "update":
+                parent[final_key] = item
+            else:  # delete
+                del parent[final_key]
+                self._size -= 1
+        else:
+            raise ValueError(f"Unknown mutation mode: {mode}")
+
+    @staticmethod
+    def _normalize_filter_options(
+        filter_options: Optional[FilterOptions],
+    ) -> Tuple[Optional[FilterMap], Literal["and", "or"]]:
+        """Convert FilterOptions to the internal structure and resolve aggregate logic."""
+        if not filter_options:
+            return None, "and"
+
+        aggregate = cast(Literal["and", "or"], filter_options.get("_aggregate", "and"))
+        if aggregate not in ("and", "or"):
+            raise ValueError(f"Unsupported filter aggregate '{aggregate}'")
+
+        normalized: Dict[str, FilterField] = {}
+        for field_name, ops in filter_options.items():
+            if field_name == "_aggregate":
+                continue
+            normalized[field_name] = cast(FilterField, ops)
+
+        return (normalized or None, aggregate)
+
+    @staticmethod
+    def _resolve_sort_options(sort: Optional[SortOptions]) -> Tuple[Optional[str], Literal["asc", "desc"]]:
+        """Extract sort field/order from the caller-provided SortOptions."""
+        if not sort:
+            return None, "asc"
+
+        sort_name = sort.get("name")
+        if not sort_name:
+            raise ValueError("Sort options must include a 'name' field")
+
+        sort_order = sort.get("order", "asc")
+        if sort_order not in ("asc", "desc"):
+            raise ValueError(f"Unsupported sort order '{sort_order}'")
+
+        return sort_name, sort_order
+
+    def _iter_items(
+        self,
+        root: Optional[Mapping[Any, Any]] = None,
+        filters: Optional[FilterMap] = None,
+        filter_logic: Literal["and", "or"] = "and",
+    ) -> Iterable[T]:
+        """Iterate over all items in the nested dictionary structure, optionally applying filters."""
+        if root is None:
+            root = self._items
+        if not root:
+            return
+        stack: List[Mapping[Any, Any]] = [root]
+        while stack:
+            node = stack.pop()
+            for value in node.values():
+                # Leaf nodes contain items; intermediate nodes are dicts.
+                if isinstance(value, self._item_type):
+                    if self._item_matches_filters(value, filters, filter_logic):
+                        yield value
+                elif isinstance(value, dict):
+                    stack.append(value)  # type: ignore
+                else:
+                    raise ValueError(
+                        f"Internal structure corrupted: expected dict or {self._item_type.__name__}, "
+                        f"got {type(value)!r}"
+                    )
+
+    def _iter_matching_items(
+        self,
+        filters: Optional[FilterMap],
+        filter_logic: Literal["and", "or"],
+    ) -> Iterable[T]:
+        """Efficiently iterate over items matching filters, using primary-key prefix when possible."""
+        # Fast path: no filters or non-AND logic -> just delegate to _iter_items.
+        if not filters or filter_logic != "and":
+            return self._iter_items(filters=filters, filter_logic=filter_logic)
+
+        # Try to derive a primary-key prefix from exact filters.
+        pk_values_prefix: List[Any] = []
+        for pk in self._primary_keys:
+            field_ops = filters.get(pk)  # type: ignore[union-attr]
+            if not field_ops:
+                break
+            # Only allow a pure {"exact": value} constraint.
+            if set(field_ops.keys()) != {"exact"}:
+                break
+            value = field_ops.get("exact")
+            if value is None:
+                break
+            pk_values_prefix.append(value)
+
+        if not pk_values_prefix:
+            return self._iter_items(filters=filters, filter_logic=filter_logic)
+
+        try:
+            if len(pk_values_prefix) == len(self._primary_keys):
+                # All primary keys specified -> at most a single item.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                single_item = parent.get(final_key)
+                if isinstance(single_item, self._item_type) and self._item_matches_filters(
+                    single_item, filters, filter_logic
+                ):
+                    return (single_item,)
+                return ()
+            else:
+                # Prefix of primary keys specified -> iterate only the subtree below that prefix.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                subtree = parent.get(final_key)
+                if isinstance(subtree, dict):
+                    return self._iter_items(subtree, filters=filters, filter_logic=filter_logic)  # type: ignore
+                return ()
+        except KeyError:
+            # No items exist for this primary-key prefix.
+            return ()
+
+    @staticmethod
+    def _item_matches_filters(
+        item: T,
+        filters: Optional[FilterMap],
+        filter_logic: Literal["and", "or"],
+    ) -> bool:
+        """Check whether an item matches the provided filter definition.
+
+        Filter format:
+
+        ```json
+        {
+            "_aggregate": "or",
+            "field_name": {
+                "exact": <value>,
+                "within": <iterable_of_allowed_values>,
+                "contains": <substring_or_element>,
+            },
+            ...
+        }
+        ```
+
+        Operators within the same field are stored in a unified pool and combined using
+        a universal logical operator.
+        """
+        if not filters:
+            return True
+
+        all_conditions_match: List[bool] = []
+
+        for field_name, ops in filters.items():
+            item_value = getattr(item, field_name, None)
+
+            for op_name, expected in ops.items():
+                # Ignore no-op filters
+                if expected is None:
+                    continue
+
+                if op_name == "exact":
+                    all_conditions_match.append(item_value == expected)
+
+                elif op_name == "within":
+                    try:
+                        all_conditions_match.append(item_value in expected)  # type: ignore[arg-type]
+                    except TypeError:
+                        all_conditions_match.append(False)
+
+                elif op_name == "contains":
+                    if item_value is None:
+                        all_conditions_match.append(False)
+                    elif isinstance(item_value, str) and isinstance(expected, str):
+                        all_conditions_match.append(expected in item_value)
+                    else:
+                        # Fallback: treat as generic iterable containment.
+                        try:
+                            all_conditions_match.append(expected in item_value)  # type: ignore[arg-type]
+                        except TypeError:
+                            all_conditions_match.append(False)
+                else:
+                    raise ValueError(f"Unsupported filter operator '{op_name}' for field '{field_name}'")
+
+        return all(all_conditions_match) if filter_logic == "and" else any(all_conditions_match)
+
+    @staticmethod
+    def _get_sort_value(item: T, sort_by: str) -> Any:
+        """Get a sort key for the given item/field.
+
+        - If the field name ends with '_time', values are treated as comparable timestamps.
+        - For other fields we try to infer a safe default from the Pydantic model annotation.
+        """
+        value = getattr(item, sort_by, None)
+
+        if sort_by.endswith("_time"):
+            # For *_time fields, push missing values to the end.
+            return float("inf") if value is None else value
+
+        if value is None:
+            # Introspect model field type to choose a reasonable default for None.
+            model_fields = getattr(item.__class__, "model_fields", {})
+            if sort_by not in model_fields:
+                raise ValueError(
+                    f"Failed to sort items by '{sort_by}': field does not exist " f"on {item.__class__.__name__}"
+                )
+
+            field_type_str = str(model_fields[sort_by].annotation)
+            if "str" in field_type_str or "Literal" in field_type_str:
+                return ""
+            if "int" in field_type_str:
+                return 0
+            if "float" in field_type_str:
+                return 0.0
+            raise ValueError(f"Failed to sort items by '{sort_by}': unsupported field type {field_type_str!r}")
+
+        return value
+
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with filters, sort order, and pagination.
+
+        Args:
+            filter: Mapping of field name to operator dict along with the optional `_aggregate` logic.
+            sort: Options describing which field to sort by and in which order.
+            limit: Max number of items to return. Use -1 for "no limit".
+            offset: Number of items to skip from the start of the *matching* items.
+        """
+        filters, filter_logic = self._normalize_filter_options(filter)
+        sort_by, sort_order = self._resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, filter_logic)
+
+        # No sorting: stream through items and apply pagination on the fly.
+        if not sort_by:
+            matched_items: List[T] = []
+            total_matched = 0
+
+            for item in items_iter:
+                # Count every match for 'total'
+                total_matched += 1
+
+                # Apply offset/limit window
+                if total_matched <= offset:
+                    continue
+                if limit != -1 and len(matched_items) >= limit:
+                    # Still need to finish iteration to get accurate total_matched.
+                    continue
+
+                matched_items.append(item)
+
+            return PaginatedResult(
+                items=matched_items,
+                limit=limit,
+                offset=offset,
+                total=total_matched,
+            )
+
+        # With sorting: we must materialize all matching items to sort them.
+        all_matches: List[T] = list(items_iter)
+
+        total_matched = len(all_matches)
+        reverse = sort_order == "desc"
+        all_matches.sort(key=lambda x: self._get_sort_value(x, sort_by), reverse=reverse)
+
+        if limit == -1:
+            paginated_items = all_matches[offset:]
+        else:
+            paginated_items = all_matches[offset : offset + limit]
+
+        return PaginatedResult(
+            items=paginated_items,
+            limit=limit,
+            offset=offset,
+            total=total_matched,
+        )
+
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Return the first (or best-sorted) item that matches the given filters, or None."""
+        filters, filter_logic = self._normalize_filter_options(filter)
+        sort_by, sort_order = self._resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, filter_logic)
+
+        if not sort_by:
+            # Just return the first matching item, if any.
+            for item in items_iter:
+                return item
+            return None
+
+        # Single-pass min/max according to sort_order.
+        best_item: Optional[T] = None
+        best_key: Any = None
+
+        for item in items_iter:
+            key = self._get_sort_value(item, sort_by)
+            if best_item is None:
+                best_item = item
+                best_key = key
+                continue
+
+            if sort_order == "asc":
+                if key < best_key:
+                    best_item, best_key = item, key
+            else:
+                if key > best_key:
+                    best_item, best_key = item, key
+
+        return best_item
+
+    async def insert(self, items: Sequence[T]) -> None:
+        """Insert the given items.
+
+        Raises:
+            ValueError: If any item with the same primary keys already exists.
+        """
+        for item in items:
+            self._mutate_single(item, mode="insert")
+
+    async def update(self, items: Sequence[T]) -> None:
+        """Update the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        for item in items:
+            self._mutate_single(item, mode="update")
+
+    async def upsert(self, items: Sequence[T]) -> None:
+        """Upsert the given items (insert if missing, otherwise update)."""
+        for item in items:
+            self._mutate_single(item, mode="upsert")
+
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        # We use a two-phase approach to avoid partial deletion if one fails:
+        # first compute key_values to validate, then perform deletions.
+        for item in items:
+            # _mutate_single will validate existence and update size.
+            self._mutate_single(item, mode="delete")
+
+
+class DequeBasedQueue(Queue[T]):
+    """Queue implementation backed by collections.deque.
+
+    Provides O(1) amortized enqueue (append) and dequeue (popleft).
+    """
+
+    def __init__(self, item_type: Type[T], items: Optional[Sequence[T]] = None):
+        self._items: Deque[T] = deque()
+        self._item_type: Type[T] = item_type
+        if items:
+            self._items.extend(items)
+
+    def item_type(self) -> Type[T]:
+        return self._item_type
+
+    async def has(self, item: T) -> bool:
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+        return item in self._items
+
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        for item in items:
+            if not isinstance(item, self._item_type):
+                raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+            self._items.append(item)
+        return items
+
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        out: List[T] = []
+        for _ in range(min(limit, len(self._items))):
+            out.append(self._items.popleft())
+        return out
+
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        result: List[T] = []
+        count = min(limit, len(self._items))
+        for idx, item in enumerate(self._items):
+            if idx >= count:
+                break
+            result.append(item)
+        return result
+
+    def size(self) -> int:
+        return len(self._items)
+
+
+class DictBasedKeyValue(KeyValue[K, V]):
+    """KeyValue implementation backed by a plain dictionary."""
+
+    def __init__(self, data: Optional[Mapping[K, V]] = None):
+        self._values: Dict[K, V] = dict(data) if data else {}
+
+    async def has(self, key: K) -> bool:
+        return key in self._values
+
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        return self._values.get(key, default)
+
+    async def set(self, key: K, value: V) -> None:
+        self._values[key] = value
+
+    async def pop(self, key: K, default: V | None = None) -> V | None:
+        return self._values.pop(key, default)
+
+    def size(self) -> int:
+        return len(self._values)
+
+
+class InMemoryLightningCollections(LightningCollections):
+    """In-memory implementation of LightningCollections using Python data structures.
+
+    Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+    """
+
+    def __init__(self):
+        self._lock = _LoopAwareAsyncLock()
+        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(
+            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"]
+        )
+        self._resources = ListBasedCollection(items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"])
+        self._workers = ListBasedCollection(items=[], item_type=Worker, primary_keys=["worker_id"])
+        self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
+        self._span_sequence_ids = DictBasedKeyValue[str, int](data={})  # rollout_id -> sequence_id
+
+    @property
+    def rollouts(self) -> ListBasedCollection[Rollout]:
+        return self._rollouts
+
+    @property
+    def attempts(self) -> ListBasedCollection[Attempt]:
+        return self._attempts
+
+    @property
+    def spans(self) -> ListBasedCollection[Span]:
+        return self._spans
+
+    @property
+    def resources(self) -> ListBasedCollection[ResourcesUpdate]:
+        return self._resources
+
+    @property
+    def workers(self) -> ListBasedCollection[Worker]:
+        return self._workers
+
+    @property
+    def rollout_queue(self) -> DequeBasedQueue[str]:
+        return self._rollout_queue
+
+    @property
+    def span_sequence_ids(self) -> DictBasedKeyValue[str, int]:
+        return self._span_sequence_ids
+
+    @asynccontextmanager
+    async def atomic(self, *args: Any, **kwargs: Any) -> AsyncGenerator[None, None]:
+        async with self._lock:
+            yield
+
+    async def evict_spans_for_rollout(self, rollout_id: str) -> None:
+        """Evict all spans for a given rollout ID.
+
+        Uses private API for efficiency.
+        """
+        self._spans._items.pop(rollout_id, [])  # pyright: ignore[reportPrivateUsage]
+
+
+class _LoopAwareAsyncLock:
+    """Async lock that transparently rebinds to the current event loop.
+
+    The lock intentionally remains *thread-unsafe*: callers must only use it from
+    one thread at a time. If multiple threads interact with the store, each
+    thread gets its own event loop specific lock.
+    """
+
+    def __init__(self) -> None:
+        self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+    # When serializing and deserializing, we don't need to serialize the locks.
+    # Because another process will have its own set of event loops and its own lock.
+    def __getstate__(self) -> dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self._locks = weakref.WeakKeyDictionary()
+
+    def _get_lock_for_current_loop(self) -> asyncio.Lock:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[loop] = lock
+        return lock
+
+    async def __aenter__(self) -> asyncio.Lock:
+        lock = self._get_lock_for_current_loop()
+        await lock.acquire()
+        return lock
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None or not lock.locked():
+            raise RuntimeError("Lock released without being acquired")
+        lock.release()

agentlightning/store/collection_based.py

@@ -0,0 +1,986 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import hashlib
+import logging
+import time
+import uuid
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+    Worker,
+)
+
+from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset, is_finished, is_queuing
+from .collection import FilterOptions, LightningCollections
+from .utils import healthcheck, propagate_status
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+T_model = TypeVar("T_model", bound=BaseModel)
+T_collections = TypeVar("T_collections", bound=LightningCollections)
+
+
+logger = logging.getLogger(__name__)
+
+
+def _healthcheck_wrapper(func: T_callable) -> T_callable:
+    """
+    Decorator to run the watchdog healthcheck **before** executing the decorated method.
+    Only runs if the store has a watchdog configured.
+    Prevents recursive healthcheck execution using a flag on the store instance.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
+        # Check if healthcheck is already running to prevent recursion
+        if getattr(self, "_healthcheck_running", False):
+            # Skip healthcheck if already running
+            return await func(self, *args, **kwargs)
+
+        # Set flag to prevent recursive healthcheck calls
+        # This flag is not asyncio/thread-safe, but it doesn't matter
+        self._healthcheck_running = True  # type: ignore
+        try:
+            # The following methods should live inside one lock.
+            await self._healthcheck()  # pyright: ignore[reportPrivateUsage]
+        finally:
+            # Always clear the flag, even if healthcheck fails
+            self._healthcheck_running = False  # type: ignore
+
+        # Execute the original method
+        # This should be outside the lock.
+        return await func(self, *args, **kwargs)
+
+    return cast(T_callable, wrapper)
+
+
+def _generate_resources_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "rs-" + short_id
+
+
+def _generate_rollout_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "ro-" + short_id
+
+
+def _generate_attempt_id() -> str:
+    """We don't need that long because attempts are limited to rollouts."""
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
+    return "at-" + short_id
+
+
+class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
+    """It's the standard implementation of LightningStore that uses collections to store data.
+
+    If the store implementation is to use the store's default behavior, it's recommended to
+    inherit from this class and override the methods if needed.
+    Bring your own collection implementation by using a different `collections` argument.
+
+    Args:
+        collections: The collections to use for storage.
+    """
+
+    def __init__(
+        self,
+        *,
+        collections: T_collections,
+    ):
+        # rollouts and spans' storage
+        self.collections = collections
+
+        # Caches the latest resources ID.
+        self._latest_resources_id: Union[str, None, Unset] = UNSET
+
+    async def _ensure_latest_resources_id(self) -> Optional[str]:
+        if isinstance(self._latest_resources_id, Unset):
+            latest_resources = await self.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
+
+    async def _get_or_create_worker(self, worker_id: str) -> Worker:
+        """Create a worker if it doesn't exist.
+
+        This is different from upsert because we don't want to update the worker if it already exists.
+        """
+        worker = await self.collections.workers.get({"worker_id": {"exact": worker_id}})
+        if worker is None:
+            worker = Worker(worker_id=worker_id)
+            await self.collections.workers.insert([worker])
+        return worker
+
+    async def _sync_worker_with_attempt(self, attempt: Attempt) -> None:
+        worker_id = attempt.worker_id
+        if not worker_id:
+            return
+
+        worker = await self._get_or_create_worker(worker_id)
+        now = time.time()
+
+        if attempt.status in ("succeeded", "failed"):
+            if worker.status != "idle":
+                worker.last_idle_time = now
+            worker.status = "idle"
+            worker.current_rollout_id = None
+            worker.current_attempt_id = None
+        elif attempt.status in ("timeout", "unresponsive"):
+            if worker.status != "unknown":
+                worker.last_idle_time = now
+            worker.status = "unknown"
+            worker.current_rollout_id = None
+            worker.current_attempt_id = None
+        else:
+            transitioned = worker.status != "busy" or worker.current_attempt_id != attempt.attempt_id
+            if transitioned:
+                worker.last_busy_time = now
+            worker.status = "busy"
+            worker.current_rollout_id = attempt.rollout_id
+            worker.current_attempt_id = attempt.attempt_id
+
+        # Validate the schema to make sure it's valid.
+        Worker.model_validate(worker.model_dump())
+        await self.collections.workers.update([worker])
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store.
+
+        This store supports no capability. The capability depends on the underlying collections.
+        """
+        return LightningStoreCapabilities()
+
+    @_healthcheck_wrapper
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """Notify the store that I'm about to run a rollout.
+
+        See [`LightningStore.start_rollout()`][agentlightning.LightningStore.start_rollout] for semantics.
+        """
+        async with self.collections.atomic():
+            rollout_id = _generate_rollout_id()
+            current_time = time.time()
+
+            rollout_config = config.model_copy(deep=True) if config is not None else RolloutConfig()
+            rollout_metadata = dict(metadata) if metadata is not None else {}
+
+            resources_id = resources_id if resources_id is not None else await self._ensure_latest_resources_id()
+
+            rollout = Rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                start_time=current_time,
+                status="preparing",
+                config=rollout_config,
+                metadata=rollout_metadata,
+            )
+
+            # Create the initial attempt
+            attempt_id = _generate_attempt_id()
+            attempt = Attempt(
+                rollout_id=rollout.rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=1,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            await self.collections.attempts.insert([attempt])
+            await self.collections.rollouts.insert([rollout])
+
+            # Notify the subclass that the rollout status has changed.
+            await self.on_rollout_update(rollout)
+
+            # Return a rollout with attempt attached.
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        """Adds a new task to the queue with specific metadata and returns the rollout.
+
+        See [`LightningStore.enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout] for semantics.
+        """
+        async with self.collections.atomic():
+            rollout_id = _generate_rollout_id()
+            current_time = time.time()
+
+            rollout_config = config.model_copy(deep=True) if config is not None else RolloutConfig()
+            rollout_metadata = dict(metadata) if metadata is not None else {}
+
+            resources_id = resources_id if resources_id is not None else await self._ensure_latest_resources_id()
+
+            rollout = Rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                start_time=current_time,
+                status="queuing",  # should be queuing
+                config=rollout_config,
+                metadata=rollout_metadata,
+            )
+
+            await self.collections.rollouts.insert([rollout])
+            await self.collections.rollout_queue.enqueue([rollout.rollout_id])  # add it to the end of the queue
+
+            # Notify the subclass that the rollout status has changed.
+            await self.on_rollout_update(rollout)
+
+            # Return the rollout with no attempt attached.
+            return rollout
+
+    @_healthcheck_wrapper
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
+        """Retrieves the next task from the queue without blocking.
+        Returns `None` if the queue is empty.
+
+        Will set the rollout status to preparing and create a new attempt.
+
+        See [`LightningStore.dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] for semantics.
+        """
+        async with self.collections.atomic():
+            if worker_id is not None:
+                worker = await self._get_or_create_worker(worker_id)
+                worker.last_dequeue_time = time.time()
+                worker.status = "idle"
+                await self.collections.workers.update([worker])
+
+            # Keep looking until we find a rollout that's still in queuing status
+            # or the queue is empty
+            while self.collections.rollout_queue.size() > 0:
+                dequeued = await self.collections.rollout_queue.dequeue(1)
+                if not dequeued:
+                    break
+                rollout_id = dequeued[0]
+
+                rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+                if not rollout:
+                    logger.warning(f"Rollout {rollout_id} not found, skipping dequeuing")
+                    continue
+
+                # Check if rollout is still in a queuing state
+                # (it might have been updated to a different status while in queue)
+                if is_queuing(rollout):
+                    # Create a new attempt (could be first attempt or retry)
+                    attempt_id = _generate_attempt_id()
+                    current_time = time.time()
+
+                    # Get existing attempts to determine sequence number
+                    existing_attempts = await self._query_attempts_for_rollout_unlocked(rollout.rollout_id)
+                    sequence_id = len(existing_attempts) + 1
+
+                    attempt = Attempt(
+                        rollout_id=rollout.rollout_id,
+                        attempt_id=attempt_id,
+                        sequence_id=sequence_id,
+                        start_time=current_time,
+                        status="preparing",
+                    )
+
+                    await self.collections.attempts.insert([attempt])
+
+                    # Sync attempt status to rollout
+                    await self._update_rollout_unlocked(rollout.rollout_id, status="preparing")
+
+                    return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+                # If not in queuing state, skip this rollout and continue
+                # (it was updated externally and should not be processed)
+                logger.warning(
+                    f"Rollout {rollout.rollout_id} is not in queuing state: {rollout.status}, skipping dequeuing"
+                )
+
+            # No valid rollouts found
+            return None
+
+    @_healthcheck_wrapper
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """Creates a new attempt for a given rollout ID and return the attempt details.
+
+        See [`LightningStore.start_attempt()`][agentlightning.LightningStore.start_attempt] for semantics.
+        """
+        async with self.collections.atomic():
+            # Get the rollout
+            rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+            if not rollout:
+                raise ValueError(f"Rollout {rollout_id} not found")
+
+            # Get existing attempts to determine sequence number
+            existing_attempts = await self._query_attempts_for_rollout_unlocked(rollout_id)
+            sequence_id = len(existing_attempts) + 1
+
+            # We don't care whether the max attempts have reached or not
+            # This attempt is from user trigger
+
+            # Create new attempt
+            attempt_id = _generate_attempt_id()
+            current_time = time.time()
+
+            attempt = Attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=sequence_id,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            # Add attempt to storage
+            await self.collections.attempts.insert([attempt])
+
+            # Sync attempt status to rollout
+            await self._update_rollout_unlocked(rollout_id, status="preparing")
+
+            # Return the rollout with the new attempt attached.
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        """Retrieves rollouts filtered by their status and rollout ids.
+        If no status is provided, returns all rollouts.
+
+        See [`LightningStore.query_rollouts()`][agentlightning.LightningStore.query_rollouts] for semantics.
+        """
+        # Construct filters condition
+        filters: FilterOptions = {}
+        if rollout_ids is not None:
+            filters["rollout_id"] = {"within": list(rollout_ids)}
+        if status is not None:
+            filters["status"] = {"within": list(status)}
+
+        async with self.collections.atomic():
+            rollouts = await self.collections.rollouts.query(filter=filters or None)
+
+            # Attach the latest attempt to the rollout objects
+            # TODO: Maybe we can use asyncio.gather here to speed up the process?
+            attempted_rollouts = [
+                await self._rollout_to_attempted_rollout_unlocked(rollout) for rollout in rollouts.items
+            ]
+
+            return attempted_rollouts
+
+    async def _query_attempts_for_rollout_unlocked(self, rollout_id: str) -> List[Attempt]:
+        """The unlocked version of `query_attempts_for_rollout`."""
+        result = await self.collections.attempts.query(
+            filter={"rollout_id": {"exact": rollout_id}},
+            sort={"name": "sequence_id", "order": "asc"},
+        )
+        return list(result.items)
+
+    @_healthcheck_wrapper
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
+        """Retrieves a specific rollout by its ID.
+
+        See [`LightningStore.get_rollout_by_id()`][agentlightning.LightningStore.get_rollout_by_id] for semantics.
+
+        If the rollout has been attempted, the latest attempt will also be returned.
+        """
+        async with self.collections.atomic():
+            rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+            if rollout is None:
+                return None
+            return await self._rollout_to_attempted_rollout_unlocked(rollout)
+
+    async def _rollout_to_attempted_rollout_unlocked(self, rollout: Rollout) -> Union[Rollout, AttemptedRollout]:
+        """Query the latest attempt for the rollout, and attach it to the rollout object.
+
+        If the rollout has no attempts, return the rollout object itself.
+        """
+        latest_attempt = await self._get_latest_attempt_unlocked(rollout.rollout_id)
+        if latest_attempt is None:
+            return rollout
+        else:
+            return AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt)
+
+    async def _get_latest_attempt_unlocked(self, rollout_id: str) -> Optional[Attempt]:
+        """The unlocked version of `get_latest_attempt`."""
+        return await self.collections.attempts.get(
+            filter={"rollout_id": {"exact": rollout_id}},
+            sort={"name": "sequence_id", "order": "desc"},
+        )
+
+    @_healthcheck_wrapper
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """Retrieves all attempts associated with a specific rollout ID.
+        Returns an empty list if no attempts are found.
+
+        See [`LightningStore.query_attempts()`][agentlightning.LightningStore.query_attempts] for semantics.
+        """
+        async with self.collections.atomic():
+            result = await self.collections.attempts.query(
+                filter={"rollout_id": {"exact": rollout_id}},
+                sort={"name": "sequence_id", "order": "asc"},
+            )
+            return list(result.items)
+
+    @_healthcheck_wrapper
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """Retrieves the latest attempt for a given rollout ID.
+
+        See [`LightningStore.get_latest_attempt()`][agentlightning.LightningStore.get_latest_attempt] for semantics.
+        """
+        async with self.collections.atomic():
+            return await self._get_latest_attempt_unlocked(rollout_id)
+
+    @_healthcheck_wrapper
+    async def query_resources(self) -> List[ResourcesUpdate]:
+        """Return every stored resource snapshot in insertion order."""
+        async with self.collections.atomic():
+            # No sorting and no pagination by default
+            result = await self.collections.resources.query()
+            return list(result.items)
+
+    @_healthcheck_wrapper
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        """Stores a new version of named resources and sets it as the latest.
+
+        See [`LightningStore.add_resources()`][agentlightning.LightningStore.add_resources] for semantics.
+        """
+        resources_id = _generate_resources_id()
+        async with self.collections.atomic():
+            current_time = time.time()
+            update = ResourcesUpdate(
+                resources_id=resources_id,
+                resources=resources,
+                create_time=current_time,
+                update_time=current_time,
+                version=1,
+            )
+            await self.collections.resources.insert([update])
+            self._latest_resources_id = resources_id
+            return update
+
+    @_healthcheck_wrapper
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+
+        See [`LightningStore.update_resources()`][agentlightning.LightningStore.update_resources] for semantics.
+        """
+        async with self.collections.atomic():
+            current_time = time.time()
+            existing = await self.collections.resources.get({"resources_id": {"exact": resources_id}})
+            if existing is None:
+                update = ResourcesUpdate(
+                    resources_id=resources_id,
+                    resources=resources,
+                    create_time=current_time,
+                    update_time=current_time,
+                    version=1,
+                )
+                await self.collections.resources.insert([update])
+            else:
+                update = existing.model_copy(
+                    update={
+                        "resources": resources,
+                        "update_time": current_time,
+                        "version": existing.version + 1,
+                    }
+                )
+                await self.collections.resources.update([update])
+            self._latest_resources_id = resources_id
+            return update
+
+    @_healthcheck_wrapper
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """Retrieves a specific version of named resources by its ID.
+
+        See [`LightningStore.get_resources_by_id()`][agentlightning.LightningStore.get_resources_by_id] for semantics.
+        """
+        async with self.collections.atomic():
+            return await self.collections.resources.get({"resources_id": {"exact": resources_id}})
+
+    @_healthcheck_wrapper
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """Retrieves the latest version of named resources.
+
+        See [`LightningStore.get_latest_resources()`][agentlightning.LightningStore.get_latest_resources] for semantics.
+        """
+        async with self.collections.atomic():
+            latest_id = await self._ensure_latest_resources_id()
+            if latest_id is None:
+                return None
+            return await self.collections.resources.get({"resources_id": {"exact": latest_id}})
+
+    async def _issue_span_sequence_id_unlocked(self, rollout_id: str) -> int:
+        """Issue a new span sequence ID for a given rollout."""
+        sequence_id = await self.collections.span_sequence_ids.get(rollout_id)
+        if sequence_id is None:
+            sequence_id = 1
+        else:
+            sequence_id += 1
+        await self.collections.span_sequence_ids.set(rollout_id, sequence_id)
+        return sequence_id
+
+    async def _sync_span_sequence_id_unlocked(self, rollout_id: str, sequence_id: int) -> None:
+        """Sync the span sequence ID for a given rollout from the input span sequence ID."""
+        existing_sequence_id = await self.collections.span_sequence_ids.get(rollout_id)
+        if existing_sequence_id is None:
+            existing_sequence_id = 0
+        await self.collections.span_sequence_ids.set(rollout_id, max(existing_sequence_id, sequence_id))
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """Get the next span sequence ID for a given rollout and attempt.
+        The number is strictly increasing for each rollout.
+        The store will not issue the same sequence ID twice.
+
+        See [`LightningStore.get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id] for semantics.
+        """
+        async with self.collections.atomic():
+            return await self._issue_span_sequence_id_unlocked(rollout_id)
+
+    async def add_span(self, span: Span) -> Span:
+        """Persist a pre-converted span.
+
+        See [`LightningStore.add_span()`][agentlightning.LightningStore.add_span] for semantics.
+        """
+        async with self.collections.atomic():
+            # Update the sequence ID to be synced with latest input span
+            await self._sync_span_sequence_id_unlocked(span.rollout_id, span.sequence_id)
+            return await self._add_span_unlocked(span)
+
+    async def add_otel_span(
+        self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
+    ) -> Span:
+        """Add an opentelemetry span to the store.
+
+        See [`LightningStore.add_otel_span()`][agentlightning.LightningStore.add_otel_span] for semantics.
+        """
+        async with self.collections.atomic():
+            if sequence_id is None:
+                # Issue a new sequence ID for the rollout
+                sequence_id = await self._issue_span_sequence_id_unlocked(rollout_id)
+            else:
+                # Comes from a provided sequence ID
+                # Make sure our counter is strictly increasing
+                await self._sync_span_sequence_id_unlocked(rollout_id, sequence_id)
+
+            span = Span.from_opentelemetry(
+                readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id
+            )
+            await self._add_span_unlocked(span)
+            return span
+
+    async def _add_span_unlocked(self, span: Span) -> Span:
+        rollout = await self.collections.rollouts.get({"rollout_id": {"exact": span.rollout_id}})
+        if not rollout:
+            raise ValueError(f"Rollout {span.rollout_id} not found")
+        current_attempt = await self.collections.attempts.get(
+            filter={"rollout_id": {"exact": span.rollout_id}, "attempt_id": {"exact": span.attempt_id}},
+        )
+        latest_attempt = await self.collections.attempts.get(
+            filter={"rollout_id": {"exact": span.rollout_id}},
+            sort={"name": "sequence_id", "order": "desc"},
+        )
+        if not current_attempt:
+            raise ValueError(f"Attempt {span.attempt_id} not found for rollout {span.rollout_id}")
+        if not latest_attempt:
+            raise ValueError(f"No attempts found for rollout {span.rollout_id}")
+
+        try:
+            await self.collections.spans.insert([span])
+        except ValueError as e:
+            if "Item already exists" in str(e):
+                # This is a duplicate span, we warn it
+                logger.error(
+                    f"Duplicate span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
+                )
+                return span
+            raise
+
+        # Update attempt heartbeat and ensure persistence
+        current_attempt.last_heartbeat_time = time.time()
+        if current_attempt.status in ["preparing", "unresponsive"]:
+            current_attempt.status = "running"
+        await self.collections.attempts.update([current_attempt])
+
+        # If the status has already timed out or failed, do not change it (but heartbeat is still recorded)
+
+        # Update rollout status if it's the latest attempt
+        if current_attempt.attempt_id == latest_attempt.attempt_id:
+            if rollout.status == "preparing":
+                rollout.status = "running"
+                await self.collections.rollouts.update([rollout])
+                await self.on_rollout_update(rollout)
+            elif rollout.status in ["queuing", "requeuing"]:
+                rollout.status = "running"
+                await self.collections.rollouts.update([rollout])
+                await self.on_rollout_update(rollout)
+
+        return span
+
+    @_healthcheck_wrapper
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """Wait for specified rollouts to complete with a timeout.
+        Returns the completed rollouts, potentially incomplete if timeout is reached.
+
+        This method does not change the state of the store.
+
+        See [`LightningStore.wait_for_rollouts()`][agentlightning.LightningStore.wait_for_rollouts] for semantics.
+        """
+        # Wait for all rollouts concurrently
+        rollouts = await asyncio.gather(
+            *[self.wait_for_rollout(rid, timeout) for rid in rollout_ids], return_exceptions=True
+        )
+
+        for rollout_id, rollout in zip(rollout_ids, rollouts):
+            if isinstance(rollout, Exception):
+                logger.error(f"Error waiting for rollout {rollout_id}: {rollout}")
+
+        # Filter out the exceptions
+        return [rollout for rollout in rollouts if isinstance(rollout, 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.
+
+        Subclass may use advanced mechanisms like events to accelerate this.
+
+        Returns the completed rollout, or None if timeout is reached.
+        """
+        # First check if already completed
+        async with self.collections.atomic():
+            rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+            if rollout and is_finished(rollout):
+                return rollout
+
+        # No timeout, return immediately
+        if timeout is not None and timeout <= 0:
+            return None
+
+        start_time = time.time()
+        deadline = start_time + timeout if timeout is not None else None
+
+        # If not completed, wait for completion
+        while deadline is None or time.time() < deadline:
+            # Poll every 10 seconds by default
+            rest_time = max(0.01, min(deadline - time.time(), 10.0)) if deadline is not None else 10.0
+            await asyncio.sleep(rest_time)
+            rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+            # check if rollout is finished
+            if rollout and is_finished(rollout):
+                return rollout
+
+        return None
+
+    @_healthcheck_wrapper
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """
+        Query and retrieve all spans associated with a specific rollout ID.
+        Returns an empty list if no spans are found.
+
+        See [`LightningStore.query_spans()`][agentlightning.LightningStore.query_spans] for semantics.
+        """
+        async with self.collections.atomic():
+            if attempt_id is None:
+                spans = await self.collections.spans.query(filter={"rollout_id": {"exact": rollout_id}})
+            elif attempt_id == "latest":
+                latest_attempt = await self.collections.attempts.get(
+                    filter={"rollout_id": {"exact": rollout_id}},
+                    sort={"name": "sequence_id", "order": "desc"},
+                )
+                if not latest_attempt:
+                    logger.debug(f"No attempts found for rollout {rollout_id} when querying latest spans")
+                    return []
+                spans = await self.collections.spans.query(
+                    filter={
+                        "rollout_id": {"exact": rollout_id},
+                        "attempt_id": {"exact": latest_attempt.attempt_id},
+                    }
+                )
+            else:
+                spans = await self.collections.spans.query(
+                    filter={"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
+                )
+            return list(spans.items)
+
+    @_healthcheck_wrapper
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        """Update the rollout status and related metadata.
+
+        See [`LightningStore.update_rollout()`][agentlightning.LightningStore.update_rollout] for semantics.
+        """
+        async with self.collections.atomic():
+            return await self._update_rollout_unlocked(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    @_healthcheck_wrapper
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        """Update a specific or latest attempt for a given rollout.
+
+        See [`LightningStore.update_attempt()`][agentlightning.LightningStore.update_attempt] for semantics.
+        """
+        async with self.collections.atomic():
+            attempt = await self._update_attempt_unlocked(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )
+
+        return attempt
+
+    async def _update_rollout_unlocked(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        # No lock inside this one.
+        rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        # Update fields if they are not UNSET
+        if not isinstance(input, Unset):
+            rollout.input = input
+        if not isinstance(mode, Unset):
+            rollout.mode = mode
+        if not isinstance(resources_id, Unset):
+            rollout.resources_id = resources_id
+        if not isinstance(status, Unset):
+            rollout.status = status
+        if not isinstance(config, Unset):
+            rollout.config = config
+        if not isinstance(metadata, Unset):
+            rollout.metadata = metadata
+
+        # Set end time for finished rollouts
+        # Rollout is only finished when it succeeded or fail with no more retries.
+        if not isinstance(status, Unset) and is_finished(rollout):
+            rollout.end_time = time.time()
+
+        # If requeuing, add back to queue.
+        # Check whether the rollout is already in queue.
+        elif is_queuing(rollout) and not await self.collections.rollout_queue.has(rollout.rollout_id):
+            await self.collections.rollout_queue.enqueue([rollout.rollout_id])
+
+        # We also don't need to remove non-queuing rollouts from the queue, for similar reasons.
+
+        # Re-validate the rollout to ensure legality
+        Rollout.model_validate(rollout.model_dump())
+
+        await self.collections.rollouts.update([rollout])
+        await self.on_rollout_update(rollout)
+
+        return rollout
+
+    async def _update_attempt_unlocked(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        # No lock, but with status propagation.
+        rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        latest_attempt = await self.collections.attempts.get(
+            {"rollout_id": {"exact": rollout_id}},
+            sort={"name": "sequence_id", "order": "desc"},
+        )
+        if not latest_attempt:
+            raise ValueError(f"No attempts found for rollout {rollout_id}")
+
+        # Find the attempt to update
+        if attempt_id == "latest":
+            attempt = latest_attempt
+        else:
+            attempt = await self.collections.attempts.get(
+                {"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
+            )
+            if not attempt:
+                raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
+
+        worker_sync_required = False
+
+        # Update fields if they are not UNSET
+        if not isinstance(worker_id, Unset):
+            attempt.worker_id = worker_id
+            worker_sync_required = worker_sync_required or bool(worker_id)
+        if not isinstance(status, Unset):
+            attempt.status = status
+            # Also update end_time if the status indicates completion
+            if status in ["failed", "succeeded"]:
+                attempt.end_time = time.time()
+            worker_sync_required = worker_sync_required or bool(attempt.worker_id)
+        if not isinstance(last_heartbeat_time, Unset):
+            attempt.last_heartbeat_time = last_heartbeat_time
+        if not isinstance(metadata, Unset):
+            attempt.metadata = metadata
+
+        if worker_sync_required and attempt.worker_id:
+            await self._sync_worker_with_attempt(attempt)
+
+        # Re-validate the attempt to ensure legality
+        Attempt.model_validate(attempt.model_dump())
+        # Update the attempt in storage
+        await self.collections.attempts.update([attempt])
+
+        if attempt.attempt_id == latest_attempt.attempt_id:
+
+            async def _update_status(rollout_id: str, status: RolloutStatus) -> Rollout:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            # Propagate the status to the rollout
+            await propagate_status(
+                _update_status,
+                attempt,
+                rollout.config,
+            )
+
+        return attempt
+
+    @_healthcheck_wrapper
+    async def query_workers(self) -> List[Worker]:
+        """Return the current snapshot of all workers."""
+        async with self.collections.atomic():
+            result = await self.collections.workers.query()
+            return list(result.items)
+
+    @_healthcheck_wrapper
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        async with self.collections.atomic():
+            return await self.collections.workers.get({"worker_id": {"exact": worker_id}})
+
+    @_healthcheck_wrapper
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        """Create or update a worker entry."""
+        async with self.collections.atomic():
+            worker = await self._get_or_create_worker(worker_id)
+            if not isinstance(heartbeat_stats, Unset):
+                worker.heartbeat_stats = dict(heartbeat_stats)
+            worker.last_heartbeat_time = time.time()
+
+            Worker.model_validate(worker.model_dump())
+            await self.collections.workers.update([worker])
+            return worker
+
+    async def on_rollout_update(self, rollout: Rollout) -> None:
+        """Callback for subclasses to implement specific logic when a rollout changes.
+
+        Subclass should not lock this method with `collections.atomic()` because the caller will already hold the lock.
+        """
+        pass
+
+    async def get_running_rollouts(self) -> List[AttemptedRollout]:
+        """Get all running rollouts.
+
+        As this is invoked very frequently (probably at every requests),
+        subclass can implement hacks to make it more efficient.
+        It should also be unlocked and let the caller hold the lock.
+        """
+        running_rollouts: List[AttemptedRollout] = []
+        rollouts = await self.collections.rollouts.query(filter={"status": {"within": ["preparing", "running"]}})
+
+        for rollout in rollouts.items:
+            latest_attempt = await self.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
+
+    async def _healthcheck(self) -> None:
+        """Perform healthcheck against all running rollouts in the store."""
+        async with self.collections.atomic():
+            running_rollouts = await self.get_running_rollouts()
+
+            async def _update_attempt_status(rollout_id: str, attempt_id: str, status: AttemptStatus) -> Attempt:
+                return await self._update_attempt_unlocked(rollout_id, attempt_id, status=status)
+
+            async def _update_rollout_status(rollout_id: str, status: RolloutStatus) -> Rollout:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            await healthcheck(
+                running_rollouts,
+                _update_rollout_status,
+                _update_attempt_status,
+            )

dashboard/test-utils/python-server.py

@@ -37,7 +37,7 @@ from agentlightning.types import (
 )
 
 
-def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) -> None:
+async def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) -> None:
     """
     Inject mock data directly into the InMemoryLightningStore.
 
@@ -217,20 +217,12 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
     )
 
     # Inject rollouts directly into store
-    store._rollouts["ro-story-001"] = rollout1
-    store._rollouts["ro-story-002"] = rollout2
-    store._rollouts["ro-story-003"] = rollout3
-    store._rollouts["ro-story-004"] = rollout4
-    store._rollouts["ro-story-005"] = rollout5
-    store._rollouts["ro-story-006"] = rollout6
+    await store.collections.rollouts.insert([rollout1, rollout2, rollout3, rollout4, rollout5, rollout6])
 
     # Inject attempts directly into store
-    store._attempts["ro-story-001"] = [attempt1]
-    store._attempts["ro-story-002"] = [attempt2_1, attempt2_2]
-    store._attempts["ro-story-003"] = [attempt3_1, attempt3_2, attempt3_3]
-    store._attempts["ro-story-004"] = []  # No attempt for preparing rollout
-    store._attempts["ro-story-005"] = [attempt5]
-    store._attempts["ro-story-006"] = [attempt6]
+    await store.collections.attempts.insert(
+        [attempt1, attempt2_1, attempt2_2, attempt3_1, attempt3_2, attempt3_3, attempt5, attempt6]
+    )
 
     # Create and inject spans with diverse data
     # Spans for ro-story-001 (Running) - Multiple nested spans with ongoing execution
@@ -545,11 +537,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
         ),
     ]
 
-    store._spans["ro-story-001"] = spans_ro1
-    store._spans["ro-story-002"] = spans_ro2_a1 + spans_ro2_a2
-    store._spans["ro-story-003"] = spans_ro3_a3
-    store._spans["ro-story-005"] = spans_ro5
-    store._spans["ro-story-006"] = spans_ro6
+    await store.collections.spans.insert(spans_ro1 + spans_ro2_a1 + spans_ro2_a2 + spans_ro3_a3 + spans_ro5 + spans_ro6)
 
     # Create and inject resources with diverse types
     resource1 = ResourcesUpdate(
@@ -628,11 +616,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
         },
     )
 
-    store._resources["rs-story-001"] = resource1
-    store._resources["rs-story-002"] = resource2
-    store._resources["rs-story-003"] = resource3
-    store._resources["rs-story-004"] = resource4
-    store._resources["rs-story-005"] = resource5
+    await store.collections.resources.insert([resource1, resource2, resource3, resource4, resource5])
     store._latest_resources_id = "rs-story-005"
 
     # Register workers with diverse states and activity windows.
@@ -694,8 +678,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
         ),
     ]
 
-    for worker in workers:
-        store._workers[worker.worker_id] = worker
+    await store.collections.workers.insert(workers)
 
 
 async def main():
@@ -704,7 +687,7 @@ async def main():
     args = parser.parse_args()
 
     store = InMemoryLightningStore()
-    inject_mock_data(store, now=args.now)
+    await inject_mock_data(store, now=args.now)
 
     # Start server
     server = LightningStoreServer(store, "127.0.0.1", 8765, "*")

docs/reference/store.md

@@ -8,6 +8,8 @@
 
 ::: agentlightning.InMemoryLightningStore
 
+::: agentlightning.CollectionBasedLightningStore
+
 ## Client-Server and Thread-safe Wrappers
 
 ::: agentlightning.LightningStoreServer
@@ -15,3 +17,27 @@
 ::: agentlightning.LightningStoreClient
 
 ::: agentlightning.LightningStoreThreaded
+
+## Collections and Collection Implementations
+
+::: agentlightning.store.collection.Collection
+
+::: agentlightning.store.collection.Queue
+
+::: agentlightning.store.collection.KeyValue
+
+::: agentlightning.store.collection.FilterOptions
+
+::: agentlightning.store.collection.SortOptions
+
+::: agentlightning.store.collection.PaginatedResult
+
+::: agentlightning.store.collection.LightningCollections
+
+::: agentlightning.store.collection.ListBasedCollection
+
+::: agentlightning.store.collection.DequeBasedQueue
+
+::: agentlightning.store.collection.DictBasedKeyValue
+
+::: agentlightning.store.collection.InMemoryLightningCollections

tests/store/conftest.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import time
+from itertools import count
 from unittest.mock import Mock
 
 import pytest
@@ -43,14 +44,20 @@ def mock_readable_span() -> ReadableSpan:
     """Create a mock ReadableSpan for testing."""
     span = Mock()
     span.name = "test_span"
+    context_counter = count(1)
+
+    def _make_context() -> Mock:
+        """Generate a distinct span context each time it is requested."""
+        index = next(context_counter)
+        context = Mock()
+        context.trace_id = 111111
+        context.span_id = 222222 + index
+        context.is_remote = False
+        context.trace_state = {}
+        return context
 
     # Mock context
-    context = Mock()
-    context.trace_id = 111111
-    context.span_id = 222222
-    context.is_remote = False
-    context.trace_state = {}  # Make it an empty dict instead of Mock
-    span.get_span_context = Mock(return_value=context)
+    span.get_span_context = Mock(side_effect=_make_context)
 
     # Mock other attributes
     span.parent = None

tests/store/test_collection_memory.py

@@ -0,0 +1,746 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, List, Literal, Mapping, Sequence, Tuple
+
+import pytest
+from pydantic import BaseModel, Field
+
+from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
+
+
+class SampleItem(BaseModel):
+    partition: str
+    index: int
+    name: str
+    status: str
+    tags: List[str] = Field(default_factory=list)
+    score: float | None = None
+    rank: int | None = None
+    updated_time: float | None = None
+    payload: Dict[str, int] = Field(default_factory=dict)
+    metadata: str | None = None
+
+
+def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
+    return ListBasedCollection(list(items), SampleItem, ("partition", "index"))
+
+
+@pytest.fixture()
+def sample_items() -> List[SampleItem]:
+    return [
+        SampleItem(
+            partition="alpha",
+            index=1,
+            name="urgent-phase-one",
+            status="new",
+            tags=["core", "urgent"],
+            score=10.5,
+            rank=3,
+            updated_time=12.0,
+            payload={"priority": 10},
+            metadata="alpha-start",
+        ),
+        SampleItem(
+            partition="alpha",
+            index=2,
+            name="phase-two",
+            status="running",
+            tags=["core"],
+            score=5.0,
+            rank=2,
+            updated_time=None,
+            payload={"priority": 5},
+            metadata=None,
+        ),
+        SampleItem(
+            partition="alpha",
+            index=3,
+            name="delayed-phase",
+            status="blocked",
+            tags=["delayed"],
+            score=None,
+            rank=5,
+            updated_time=15.1,
+            payload={"priority": 8},
+            metadata="delayed-phase",
+        ),
+        SampleItem(
+            partition="beta",
+            index=1,
+            name="beta-critical",
+            status="new",
+            tags=["beta", "urgent"],
+            score=8.0,
+            rank=1,
+            updated_time=7.0,
+            payload={"priority": 7},
+            metadata="beta critical",
+        ),
+        SampleItem(
+            partition="beta",
+            index=2,
+            name="beta optional",
+            status="done",
+            tags=["beta"],
+            score=3.0,
+            rank=None,
+            updated_time=2.0,
+            payload={"priority": 1},
+            metadata="optional path",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=1,
+            name="gamma-phase",
+            status="running",
+            tags=[],
+            score=9.5,
+            rank=4,
+            updated_time=None,
+            payload={"priority": 9},
+            metadata="gamma-phase data",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=2,
+            name="gamma-late",
+            status="done",
+            tags=["late", "core"],
+            score=1.0,
+            rank=6,
+            updated_time=20.0,
+            payload={"priority": 2},
+            metadata="gamma late entry",
+        ),
+        SampleItem(
+            partition="delta",
+            index=1,
+            name="delta misc",
+            status="archived",
+            tags=["misc"],
+            score=4.2,
+            rank=7,
+            updated_time=11.0,
+            payload={"priority": 3},
+            metadata="delta misc block",
+        ),
+    ]
+
+
+BASE_KEY_ORDER: List[Tuple[str, int]] = [
+    ("alpha", 1),
+    ("alpha", 2),
+    ("alpha", 3),
+    ("beta", 1),
+    ("beta", 2),
+    ("gamma", 1),
+    ("gamma", 2),
+    ("delta", 1),
+]
+
+
+@pytest.fixture()
+def sample_collection(sample_items: Sequence[SampleItem]) -> ListBasedCollection[SampleItem]:
+    return _build_collection(sample_items)
+
+
+def _key_pairs(items: Sequence[SampleItem]) -> List[Tuple[str, int]]:
+    return [(item.partition, item.index) for item in items]
+
+
+def _sorted_pairs(items: Sequence[SampleItem]) -> List[Tuple[str, int]]:
+    return sorted(_key_pairs(items))
+
+
+def test_list_collection_requires_primary_keys(sample_items: Sequence[SampleItem]) -> None:
+    with pytest.raises(ValueError):
+        ListBasedCollection(list(sample_items), SampleItem, ())
+
+
+def test_list_collection_primary_keys(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    assert tuple(sample_collection.primary_keys()) == ("partition", "index")
+
+
+def test_list_collection_item_type(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    assert sample_collection.item_type() is SampleItem
+
+
+def test_list_collection_initial_size(
+    sample_collection: ListBasedCollection[SampleItem], sample_items: Sequence[SampleItem]
+) -> None:
+    assert sample_collection.size() == len(sample_items)
+
+
+def test_list_collection_repr_contains_model_info(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    result = repr(sample_collection)
+    assert "ListBasedCollection" in result and "SampleItem" in result and str(sample_collection.size()) in result
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_insert_adds_item(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    new_item = SampleItem(partition="omega", index=1, name="omega", status="new")
+    await sample_collection.insert([new_item])
+    assert sample_collection.size() == 9
+    result = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 1}})
+    assert result == new_item
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_insert_duplicate_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    duplicate = SampleItem(partition="alpha", index=1, name="dup", status="new")
+    with pytest.raises(ValueError):
+        await sample_collection.insert([duplicate])
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_insert_wrong_type(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    class Another(BaseModel):
+        partition: str
+        index: int
+
+    wrong = Another(partition="omega", index=5)
+    with pytest.raises(TypeError):
+        await sample_collection.insert([wrong])  # type: ignore[arg-type]
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_update_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    updated = SampleItem(partition="alpha", index=1, name="updated", status="new")
+    await sample_collection.update([updated])
+    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
+    assert result == updated
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_update_missing_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    missing = SampleItem(partition="omega", index=99, name="missing", status="lost")
+    with pytest.raises(ValueError):
+        await sample_collection.update([missing])
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_delete_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    target = SampleItem(partition="alpha", index=1, name="ignored", status="new")
+    await sample_collection.delete([target])
+    assert sample_collection.size() == 7
+    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
+    assert result is None
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_delete_missing_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    missing = SampleItem(partition="omega", index=3, name="x", status="y")
+    with pytest.raises(ValueError):
+        await sample_collection.delete([missing])
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_upsert_inserts_when_missing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    created = SampleItem(partition="omega", index=4, name="new", status="queued")
+    await sample_collection.upsert([created])
+    assert sample_collection.size() == 9
+    fetched = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 4}})
+    assert fetched == created
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_upsert_updates_when_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    replacement = SampleItem(partition="beta", index=2, name="replacement", status="done")
+    await sample_collection.upsert([replacement])
+    assert sample_collection.size() == 8
+    fetched = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
+    assert fetched == replacement
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_delete_multiple_items(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    await sample_collection.delete(
+        [
+            SampleItem(partition="alpha", index=1, name="", status=""),
+            SampleItem(partition="beta", index=1, name="", status=""),
+        ]
+    )
+    assert sample_collection.size() == 6
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_insert_accepts_tuple_sequence(
+    sample_collection: ListBasedCollection[SampleItem],
+) -> None:
+    extra = (
+        SampleItem(partition="tuple", index=1, name="a", status="pending"),
+        SampleItem(partition="tuple", index=2, name="b", status="pending"),
+    )
+    await sample_collection.insert(extra)
+    assert sample_collection.size() == 10
+    fetched = await sample_collection.query(filter={"partition": {"exact": "tuple"}})
+    assert _sorted_pairs(fetched.items) == [("tuple", 1), ("tuple", 2)]
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_query_without_filters_returns_all(
+    sample_collection: ListBasedCollection[SampleItem],
+) -> None:
+    result = await sample_collection.query()
+    assert result.total == 8
+    assert len(result.items) == 8
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize(
+    ("filters", "expected"),
+    [
+        pytest.param({"status": {"exact": "new"}}, [("alpha", 1), ("beta", 1)], id="exact-single-field"),
+        pytest.param(
+            {"partition": {"exact": "alpha"}, "index": {"exact": 2}},
+            [("alpha", 2)],
+            id="exact-multiple-fields",
+        ),
+        pytest.param(
+            {"status": {"within": {"running", "blocked"}}},
+            [("alpha", 2), ("alpha", 3), ("gamma", 1)],
+            id="within-set",
+        ),
+        pytest.param(
+            {"partition": {"within": ["gamma", "delta"]}},
+            [("gamma", 1), ("gamma", 2), ("delta", 1)],
+            id="within-list",
+        ),
+        pytest.param(
+            {"name": {"contains": "phase"}},
+            [("alpha", 1), ("alpha", 2), ("alpha", 3), ("gamma", 1)],
+            id="contains-substring",
+        ),
+        pytest.param(
+            {"tags": {"contains": "urgent"}},
+            [("alpha", 1), ("beta", 1)],
+            id="contains-list",
+        ),
+        pytest.param({"metadata": {"contains": "phase"}}, [("alpha", 3), ("gamma", 1)], id="contains-with-none-values"),
+        pytest.param({"tags": {"contains": "missing"}}, [], id="contains-no-match"),
+        pytest.param({"partition": {"exact": "delta"}}, [("delta", 1)], id="single-exact-match"),
+        pytest.param({"missing": {"exact": "value"}}, [], id="exact-missing-field"),
+        pytest.param({"score": {"contains": "phase"}}, [], id="contains-typeerror"),
+        pytest.param({"name": {"contains": None}}, list(BASE_KEY_ORDER), id="contains-null-check"),
+        pytest.param({"status": {"exact": None}}, list(BASE_KEY_ORDER), id="exact-null-no-filter"),
+        pytest.param({"status": {"within": 1}}, [], id="within-non-iterable"),
+    ],
+)
+async def test_list_collection_query_filters(
+    sample_collection: ListBasedCollection[SampleItem],
+    filters: Dict[str, Dict[str, object]],
+    expected: Sequence[Tuple[str, int]],
+) -> None:
+    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
+    assert _sorted_pairs(result.items) == sorted(expected)
+    assert result.total == len(expected)
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize(
+    ("filters", "filter_logic", "expected"),
+    [
+        (
+            {"status": {"exact": "new"}, "tags": {"contains": "beta"}},
+            "and",
+            [("beta", 1)],
+        ),
+        (
+            {"status": {"exact": "new"}, "tags": {"contains": "beta"}},
+            "or",
+            [("alpha", 1), ("beta", 1), ("beta", 2)],
+        ),
+        (
+            {"status": {"exact": "done"}, "tags": {"contains": "core"}},
+            "and",
+            [("gamma", 2)],
+        ),
+        (
+            {"status": {"exact": "done"}, "tags": {"contains": "core"}},
+            "or",
+            [("alpha", 1), ("alpha", 2), ("beta", 2), ("gamma", 2)],
+        ),
+    ],
+)
+async def test_list_collection_filter_logic(
+    sample_collection: ListBasedCollection[SampleItem],
+    filters: Dict[str, Dict[str, object]],
+    filter_logic: Literal["and", "or"],
+    expected: Sequence[Tuple[str, int]],
+) -> None:
+    filter_payload = dict(filters)
+    filter_payload["_aggregate"] = filter_logic  # type: ignore[index]
+    result = await sample_collection.query(filter=filter_payload)  # type: ignore[arg-type]
+    assert _sorted_pairs(result.items) == sorted(expected)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_primary_key_prefix_limits_filter_checks(
+    sample_items: Sequence[SampleItem],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    collection = _build_collection(sample_items)
+    seen: List[Tuple[str, int]] = []
+    original = (  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType,reportUnknownVariableType]
+        ListBasedCollection._item_matches_filters  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType]
+    )
+
+    def tracking(item: SampleItem, filters: object, filter_logic: str) -> bool:
+        seen.append((item.partition, item.index))
+        return original(item, filters, filter_logic)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(ListBasedCollection, "_item_matches_filters", staticmethod(tracking))  # type: ignore[arg-type]
+
+    filters = {"partition": {"exact": "alpha"}, "index": {"within": {1, 2}}}
+    result = await collection.query(filter=filters)  # type: ignore[arg-type]
+    assert _sorted_pairs(result.items) == [("alpha", 1), ("alpha", 2)]
+    assert set(seen) == {("alpha", 1), ("alpha", 2), ("alpha", 3)}
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_full_primary_key_avoids_tree_scan(
+    sample_collection: ListBasedCollection[SampleItem],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    call_count = 0
+    original_iter_items = (  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType,reportUnknownVariableType]
+        ListBasedCollection._iter_items  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType]
+    )
+
+    def tracking(
+        self: ListBasedCollection[SampleItem],
+        root: Mapping[str, object] | None = None,
+        filters: object | None = None,
+        filter_logic: str = "and",
+    ) -> Iterable[SampleItem]:
+        nonlocal call_count
+        call_count += 1
+        return original_iter_items(self, root, filters, filter_logic)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(ListBasedCollection, "_iter_items", tracking)
+
+    filters = {"partition": {"exact": "beta"}, "index": {"exact": 2}}
+    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
+    assert _sorted_pairs(result.items) == [("beta", 2)]
+    assert call_count == 0
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize(
+    ("sort_by", "sort_order", "limit", "expected"),
+    [
+        ("name", "asc", 4, [("beta", 2), ("beta", 1), ("alpha", 3), ("delta", 1)]),
+        ("name", "desc", 4, [("alpha", 1), ("alpha", 2), ("gamma", 1), ("gamma", 2)]),
+        ("rank", "asc", 4, [("beta", 2), ("beta", 1), ("alpha", 2), ("alpha", 1)]),
+        ("rank", "desc", 4, [("delta", 1), ("gamma", 2), ("alpha", 3), ("gamma", 1)]),
+        ("score", "asc", 4, [("alpha", 3), ("gamma", 2), ("beta", 2), ("delta", 1)]),
+        ("score", "desc", 4, [("alpha", 1), ("gamma", 1), ("beta", 1), ("alpha", 2)]),
+        ("updated_time", "asc", 4, [("beta", 2), ("beta", 1), ("delta", 1), ("alpha", 1)]),
+        ("updated_time", "desc", 4, [("gamma", 1), ("alpha", 2), ("gamma", 2), ("alpha", 3)]),
+    ],
+)
+async def test_list_collection_sorting(
+    sample_collection: ListBasedCollection[SampleItem],
+    sort_by: str,
+    sort_order: str,
+    limit: int,
+    expected: Sequence[Tuple[str, int]],
+) -> None:
+    result = await sample_collection.query(sort={"name": sort_by, "order": sort_order}, limit=limit)  # type: ignore[arg-type]
+    assert _key_pairs(result.items) == list(expected)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_sort_by_missing_field_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    with pytest.raises(ValueError):
+        await sample_collection.query(sort={"name": "does_not_exist", "order": "asc"})
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize(
+    ("limit", "offset", "expected", "total"),
+    [
+        (1, 0, [("alpha", 1)], 3),
+        (2, 1, [("alpha", 2), ("alpha", 3)], 3),
+        (-1, 1, [("alpha", 2), ("alpha", 3)], 3),
+        (10, 0, [("alpha", 1), ("alpha", 2), ("alpha", 3)], 3),
+        (0, 0, [], 3),
+        (1, 10, [], 3),
+    ],
+)
+async def test_list_collection_pagination_without_sort(
+    sample_collection: ListBasedCollection[SampleItem],
+    limit: int,
+    offset: int,
+    expected: Sequence[Tuple[str, int]],
+    total: int,
+) -> None:
+    result = await sample_collection.query(filter={"partition": {"exact": "alpha"}}, limit=limit, offset=offset)
+    assert _key_pairs(result.items) == list(expected)
+    assert result.total == total
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_pagination_with_sort(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=2, offset=3)
+    assert _key_pairs(result.items) == [("delta", 1), ("gamma", 2)]
+    assert result.total == 8
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_limit_unbounded_with_sort(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=-1, offset=6)
+    assert _key_pairs(result.items) == [("alpha", 2), ("alpha", 1)]
+    assert result.total == 8
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_limit_zero_reports_total(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    result = await sample_collection.query(filter={"status": {"exact": "done"}}, limit=0)
+    assert result.items == []
+    assert result.total == 2
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_offset_beyond_total_returns_empty(
+    sample_collection: ListBasedCollection[SampleItem],
+) -> None:
+    result = await sample_collection.query(filter={"status": {"exact": "done"}}, offset=10)
+    assert result.items == []
+    assert result.total == 2
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_query_reports_total_with_limit(
+    sample_collection: ListBasedCollection[SampleItem],
+) -> None:
+    result = await sample_collection.query(filter={"partition": {"exact": "alpha"}}, limit=1)
+    assert result.total == 3
+    assert len(result.items) == 1
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_get_returns_first_match(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    item = await sample_collection.get({"status": {"exact": "new"}})
+    assert item is not None
+    assert (item.partition, item.index) == ("beta", 1)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_get_returns_none(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    result = await sample_collection.get({"partition": {"exact": "missing"}})
+    assert result is None
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_get_respects_filter_logic(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    filters = {"status": {"exact": "done"}, "tags": {"contains": "urgent"}, "_aggregate": "or"}
+    item = await sample_collection.get(filters)  # type: ignore[arg-type]
+    assert item is not None
+    assert (item.partition, item.index) == ("gamma", 2)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_get_honors_sort_by(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    filters = {"partition": {"exact": "alpha"}}
+    item = await sample_collection.get(filters, sort={"name": "rank", "order": "asc"})  # type: ignore[arg-type]
+    assert item is not None
+    assert (item.partition, item.index) == ("alpha", 2)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_get_honors_sort_order(sample_collection: ListBasedCollection[SampleItem]) -> None:
+    filters = {"partition": {"exact": "alpha"}}
+    item = await sample_collection.get(filters, sort={"name": "rank", "order": "desc"})  # type: ignore[arg-type]
+    assert item is not None
+    assert (item.partition, item.index) == ("alpha", 3)
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_query_handles_large_dataset() -> None:
+    bulk_items = [
+        SampleItem(
+            partition=f"partition-{i % 5}",
+            index=i,
+            name=f"bulk-{i}",
+            status="bulk",
+            score=float(i),
+            rank=i,
+            updated_time=float(i),
+        )
+        for i in range(1500)
+    ]
+    collection = _build_collection(bulk_items)
+    result = await collection.query(sort={"name": "index", "order": "asc"}, limit=50, offset=100)
+    assert result.total == 1500
+    assert len(result.items) == 50
+    assert result.items[0].index == 100
+    assert result.items[-1].index == 149
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_bulk_delete_and_size() -> None:
+    items = [SampleItem(partition="bulk", index=i, name=f"item-{i}", status="bulk") for i in range(40)]
+    collection = _build_collection(items)
+    await collection.delete(items[:20])
+    assert collection.size() == 20
+    await collection.delete(items[20:])
+    assert collection.size() == 0
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_query_rejects_unknown_operator(
+    sample_collection: ListBasedCollection[SampleItem],
+) -> None:
+    with pytest.raises(ValueError):
+        await sample_collection.query(filter={"status": {"invalid": "x"}})  # type: ignore[arg-type]
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_query_result_type() -> None:
+    collection = _build_collection([])
+    result = await collection.query(filter=None)
+    assert result.items == []
+    assert result.offset == 0
+
+
+class QueueItem(BaseModel):
+    idx: int
+
+
+@pytest.fixture()
+def deque_queue() -> DequeBasedQueue[QueueItem]:
+    return DequeBasedQueue(QueueItem, [QueueItem(idx=i) for i in range(3)])
+
+
+def test_deque_queue_initial_size(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert deque_queue.size() == 3
+
+
+def test_deque_queue_item_type(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert deque_queue.item_type() is QueueItem
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_has_detects_members(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert await deque_queue.has(QueueItem(idx=1))
+    assert not await deque_queue.has(QueueItem(idx=99))
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_enqueue_appends_items(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    items = [QueueItem(idx=3), QueueItem(idx=4)]
+    returned = await deque_queue.enqueue(items)
+    assert returned == items
+    assert deque_queue.size() == 5
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_enqueue_rejects_wrong_type(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    class Wrong(BaseModel):
+        idx: int
+
+    with pytest.raises(TypeError):
+        await deque_queue.enqueue([Wrong(idx=9)])  # type: ignore[arg-type]
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("limit", [1, 2, 5])
+async def test_deque_queue_dequeue_respects_limit(deque_queue: DequeBasedQueue[QueueItem], limit: int) -> None:
+    result = await deque_queue.dequeue(limit)
+    assert len(result) == min(limit, 3)
+    assert deque_queue.size() == 3 - min(limit, 3)
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_dequeue_zero_returns_empty(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert await deque_queue.dequeue(0) == []
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_dequeue_more_than_available(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    result = await deque_queue.dequeue(10)
+    assert len(result) == 3
+    assert deque_queue.size() == 0
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_peek_preserves_items(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    snapshot = await deque_queue.peek(2)
+    assert [item.idx for item in snapshot] == [0, 1]
+    assert deque_queue.size() == 3
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_peek_zero_returns_empty(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert await deque_queue.peek(0) == []
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_peek_after_partial_dequeue(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    await deque_queue.dequeue(1)
+    snapshot = await deque_queue.peek(2)
+    assert [item.idx for item in snapshot] == [1, 2]
+
+
+@pytest.mark.asyncio()
+async def test_deque_queue_handles_large_volume() -> None:
+    queue = DequeBasedQueue(QueueItem)
+    items = [QueueItem(idx=i) for i in range(2000)]
+    await queue.enqueue(items)
+    assert queue.size() == 2000
+    drained = await queue.dequeue(1500)
+    assert len(drained) == 1500
+    assert queue.size() == 500
+
+
+@pytest.fixture()
+def dict_key_value_data() -> Dict[str, int]:
+    return {"alpha": 1, "beta": 2}
+
+
+@pytest.fixture()
+def dict_key_value(dict_key_value_data: Dict[str, int]) -> DictBasedKeyValue[str, int]:
+    return DictBasedKeyValue(dict_key_value_data)
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_initial_state(dict_key_value: DictBasedKeyValue[str, int]) -> None:
+    assert dict_key_value.size() == 2
+    assert await dict_key_value.get("alpha") == 1
+    assert await dict_key_value.get("missing") is None
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_has_handles_presence(dict_key_value: DictBasedKeyValue[str, int]) -> None:
+    assert await dict_key_value.has("alpha")
+    assert not await dict_key_value.has("gamma")
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_set_updates_and_expands(dict_key_value: DictBasedKeyValue[str, int]) -> None:
+    await dict_key_value.set("gamma", 3)
+    assert dict_key_value.size() == 3
+    await dict_key_value.set("alpha", 99)
+    assert await dict_key_value.get("alpha") == 99
+    assert dict_key_value.size() == 3
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyValue[str, int]) -> None:
+    result = await dict_key_value.pop("beta")
+    assert result == 2
+    assert dict_key_value.size() == 1
+    result = await dict_key_value.pop("missing", 42)
+    assert result == 42
+    assert dict_key_value.size() == 1
+
+
+@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)
+    await key_value.set("gamma", 3)  # type: ignore[arg-type]
+    await key_value.pop("alpha")  # type: ignore[arg-type]
+    assert dict_key_value_data == {"alpha": 1, "beta": 2}

tests/store/test_implementation.py

@@ -16,6 +16,7 @@ It should work for multiple store implementations (InMemory, SQL, etc.).
 """
 
 import asyncio
+import logging
 import sys
 import time
 from typing import List, Optional, cast
@@ -24,7 +25,7 @@ from unittest.mock import Mock
 import pytest
 from pydantic import BaseModel
 
-from agentlightning.store.base import LightningStore
+from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
     LLM,
@@ -629,6 +630,26 @@ async def test_resource_lifecycle(store_fixture: LightningStore) -> None:
     assert old.resources["main_llm"].model == "test-model-v1"
 
 
+@pytest.mark.asyncio
+async def test_latest_resources_rehydrates_cache(store_fixture: LightningStore) -> None:
+    """get_latest_resources should consult storage even if the cache is unset."""
+    llm = LLM(
+        resource_type="llm",
+        endpoint="http://localhost:8080/v1",
+        model="cache-model",
+        sampling_parameters={"temperature": 0.1},
+    )
+    update = await store_fixture.update_resources("cache-test", {"main_llm": llm})
+
+    # Simulate a fresh process by clearing the cache.
+    store_fixture._latest_resources_id = UNSET  # type: ignore[attr-defined]
+
+    latest = await store_fixture.get_latest_resources()
+    assert latest is not None
+    assert latest.resources_id == update.resources_id
+    assert latest.resources["main_llm"].model == "cache-model"  # type: ignore
+
+
 @pytest.mark.asyncio
 async def test_task_inherits_latest_resources(store_fixture: LightningStore) -> None:
     """Test that new tasks inherit latest resources_id if not specified."""
@@ -698,6 +719,94 @@ async def test_span_sequence_generation(store_fixture: LightningStore, mock_read
     assert seq_id == 5
 
 
+@pytest.mark.asyncio
+async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
+    """Adding a span should persist attempt heartbeat and transition status to running."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "attempt-status"})
+    await store_fixture.dequeue_rollout()
+
+    attempts = await store_fixture.query_attempts(rollout.rollout_id)
+    assert attempts
+    attempt_id = attempts[0].attempt_id
+    assert attempts[0].status == "preparing"
+    assert attempts[0].last_heartbeat_time is None
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    updated_attempt = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert updated_attempt.status == "running"
+    assert updated_attempt.last_heartbeat_time is not None
+
+
+@pytest.mark.asyncio
+async def test_unresponsive_attempt_recovers_after_span(
+    store_fixture: LightningStore, mock_readable_span: Mock
+) -> None:
+    """Spans arriving for an unresponsive attempt should mark it running again."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "unresponsive"})
+    dequeued = await store_fixture.dequeue_rollout()
+    assert dequeued is not None
+    attempt_id = dequeued.attempt.attempt_id
+
+    await store_fixture.update_attempt(
+        rollout_id=rollout.rollout_id,
+        attempt_id=attempt_id,
+        status="unresponsive",
+    )
+
+    attempt_before = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_before.status == "unresponsive"
+
+    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 is not None
+
+
+@pytest.mark.asyncio
+async def test_running_attempt_updates_heartbeat(
+    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Adding spans to an already running attempt should advance its heartbeat."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "running-heartbeat"})
+    await store_fixture.dequeue_rollout()
+    attempt_id = (await store_fixture.query_attempts(rollout.rollout_id))[0].attempt_id
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+    attempt_after_first = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after_first.status == "running"
+    first_heartbeat = attempt_after_first.last_heartbeat_time
+    assert first_heartbeat is not None
+
+    monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: first_heartbeat + 100.0)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+    attempt_after_second = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after_second.last_heartbeat_time == first_heartbeat + 100.0
+
+
+@pytest.mark.asyncio
+async def test_duplicate_span_id_error(
+    store_fixture: LightningStore, mock_readable_span: Mock, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Adding two spans with the same span_id should raise a ValueError."""
+    caplog.set_level(logging.ERROR)
+    rollout = await store_fixture.enqueue_rollout(input={"test": "data"})
+    await store_fixture.dequeue_rollout()
+    attempts = await store_fixture.query_attempts(rollout.rollout_id)
+    attempt_id = attempts[0].attempt_id
+
+    # Force the mock to reuse the same span context for every call.
+    fixed_context = mock_readable_span.get_span_context()
+    mock_readable_span.get_span_context = Mock(return_value=fixed_context)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+    assert "Duplicate span added" in caplog.text
+
+
 @pytest.mark.asyncio
 async def test_span_with_explicit_sequence_id(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
     """Test providing explicit sequence_id to spans."""