microsoft--agent-lightning/agentlightning/store/collection/memory.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio
import logging
import uuid
import weakref
from collections import deque
from contextlib import AsyncExitStack, asynccontextmanager
from typing import (
    Any,
    Deque,
    Dict,
    Iterable,
    List,
    Literal,
    Mapping,
    MutableMapping,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
)

import aiologic
from pydantic import BaseModel

from agentlightning.types import (
    Attempt,
    FilterField,
    FilterOptions,
    PaginatedResult,
    ResourcesUpdate,
    Rollout,
    SortOptions,
    Span,
    Worker,
)
from agentlightning.utils.metrics import MetricsBackend

from .base import (
    AtomicLabels,
    AtomicMode,
    Collection,
    DuplicatedPrimaryKeyError,
    FilterMap,
    KeyValue,
    LightningCollections,
    Queue,
    ensure_numeric,
    normalize_filter_options,
    resolve_sort_options,
    tracked,
)

T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
K = TypeVar("K")
V = TypeVar("V")

logger = logging.getLogger(__name__)

# Nested structure type:
# dict[pk1] -> dict[pk2] -> ... -> item
ListBasedCollectionItemType = Union[
    Dict[Any, "ListBasedCollectionItemType[T]"],  # intermediate node
    Dict[Any, T],  # leaf node dictionary
]

MutationMode = Literal["insert", "update", "upsert", "delete"]


def _item_matches_filters(
    item: object,
    filters: Optional[FilterMap],
    filter_logic: Literal["and", "or"],
    must_filters: Optional[FilterMap] = None,
) -> 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 must_filters and not _item_matches_filters(item, must_filters, "and"):
        return False

    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)


def _get_sort_value(item: object, 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


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],
        id: Optional[str] = None,
        tracker: Optional[MetricsBackend] = None,
    ):
        super().__init__(tracker=tracker)
        if not primary_keys:
            raise ValueError("primary_keys must be non-empty")

        self._id = id if id is not None else str(uuid.uuid4())
        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")

    @property
    def collection_name(self) -> str:
        return self._id

    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

    async 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, update_fields: Sequence[str] | None = None) -> Optional[T]:
        """Core mutation logic shared by insert, update, upsert, and delete."""
        self._ensure_item_type(item)
        key_values = self._extract_primary_key_values(item)

        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 DuplicatedPrimaryKeyError(
                        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 update_fields is None:
                    # update_or_insert: update all fields
                    parent[final_key] = item

                else:
                    if not issubclass(self._item_type, BaseModel):
                        raise TypeError(
                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
                        )

                    # Try to fetch the existing item
                    existing = parent[final_key]
                    if not isinstance(existing, self._item_type):
                        raise ValueError(
                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
                        )

                    if not isinstance(item, self._item_type):
                        raise TypeError(
                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
                        )

                    parent[final_key] = parent[final_key].model_copy(
                        update={field: getattr(item, field) for field in update_fields}
                    )

            return parent[final_key]

        elif mode in ("update", "delete"):
            # For update/delete we must not create missing paths.
            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":
                if update_fields is None:
                    # replace the entire item
                    parent[final_key] = item
                else:
                    if not issubclass(self._item_type, BaseModel):
                        raise TypeError(
                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
                        )
                    if not isinstance(item, self._item_type):
                        raise TypeError(
                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
                        )
                    parent[final_key] = parent[final_key].model_copy(
                        update={field: getattr(item, field) for field in update_fields}
                    )
                return parent[final_key]
            else:  # delete
                del parent[final_key]
                self._size -= 1
        else:
            raise ValueError(f"Unknown mutation mode: {mode}")

    def _iter_items(
        self,
        root: Optional[Mapping[Any, Any]] = None,
        filters: Optional[FilterMap] = None,
        must_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 _item_matches_filters(value, filters, filter_logic, must_filters):
                        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],
        must_filters: Optional[FilterMap],
        filter_logic: Literal["and", "or"],
    ) -> Iterable[T]:
        """Efficiently iterate over items matching filters, using primary-key prefix when possible."""
        # Fast path: when optional filters can't form a prefix, fall back to scanning.
        if filter_logic != "and" and must_filters is None:
            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)

        # Try to derive a primary-key prefix from exact filters.
        pk_values_prefix: List[Any] = []
        prefix_sources: List[FilterMap] = []
        if must_filters:
            prefix_sources.append(must_filters)
        if filter_logic == "and" and filters:
            prefix_sources.append(filters)

        for pk in self._primary_keys:
            # combined_ops are: [{"exact": value}, {"within": [...]}, ...]
            combined_ops: List[FilterField] = []
            for source in prefix_sources:
                field_ops = source.get(pk)  # type: ignore[union-attr]
                if field_ops:
                    combined_ops.append(field_ops)
            if not combined_ops:
                break
            # Only allow a pure {"exact": value} constraint.
            exact_value: Any | None = None
            allow_prefix = True
            for ops in combined_ops:
                if set(ops.keys()) != {"exact"}:
                    allow_prefix = False
                    break
                candidate = ops.get("exact")
                if candidate is None:
                    allow_prefix = False
                    break
                if exact_value is not None and candidate != exact_value:
                    # Contradictory exact filters mean no items can match.
                    logger.warning(f"Contradictory exact filters for field '{pk}': {exact_value} != {candidate}")
                    return ()
                exact_value = candidate

            if not allow_prefix:
                break

            value = exact_value
            if value is None:
                break
            pk_values_prefix.append(value)

        if not pk_values_prefix:
            return self._iter_items(filters=filters, must_filters=must_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 _item_matches_filters(
                    single_item,
                    filters,
                    filter_logic,
                    must_filters,
                ):
                    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,  # type: ignore
                        filters=filters,
                        must_filters=must_filters,
                        filter_logic=filter_logic,
                    )
                return ()
        except KeyError:
            # No items exist for this primary-key prefix.
            return ()

    @tracked("query")
    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, must_filters, filter_logic = normalize_filter_options(filter)
        sort_by, sort_order = resolve_sort_options(sort)
        items_iter: Iterable[T] = self._iter_matching_items(filters, must_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: _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,
        )

    @tracked("get")
    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, must_filters, filter_logic = normalize_filter_options(filter)
        sort_by, sort_order = resolve_sort_options(sort)
        items_iter: Iterable[T] = self._iter_matching_items(filters, must_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 = _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

    @tracked("insert")
    async def insert(self, items: Sequence[T]) -> None:
        """Insert the given items.

        Raises:
            DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
        """
        seen_keys: set[Tuple[Any, ...]] = set()
        prepared: List[T] = []
        for item in items:
            self._ensure_item_type(item)
            key_values = self._extract_primary_key_values(item)
            if key_values in seen_keys:
                raise DuplicatedPrimaryKeyError(
                    f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
                )
            seen_keys.add(key_values)
            prepared.append(item)

        for item in prepared:
            self._mutate_single(item, mode="insert")

    @tracked("update")
    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
        """Update the given items.

        Raises:
            ValueError: If any item with the given primary keys does not exist.
        """
        updated_items: List[T] = []
        for item in items:
            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
            if updated is None:
                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
            updated_items.append(updated)
        return updated_items

    @tracked("upsert")
    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
        """Upsert the given items (insert if missing, otherwise update)."""
        upserted_items: List[T] = []
        for item in items:
            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
            if upserted is None:
                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
            upserted_items.append(upserted)
        return upserted_items

    @tracked("delete")
    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,
        id: Optional[str] = None,
        tracker: Optional[MetricsBackend] = None,
    ):
        super().__init__(tracker=tracker)
        self._items: Deque[T] = deque()
        self._item_type: Type[T] = item_type
        self._id = id if id is not None else str(uuid.uuid4())
        if items:
            self._items.extend(items)

    def item_type(self) -> Type[T]:
        return self._item_type

    @property
    def collection_name(self) -> str:
        return self._id

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"

    @tracked("has")
    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

    @tracked("enqueue")
    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

    @tracked("dequeue")
    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

    @tracked("peek")
    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

    @tracked("size")
    async 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, id: Optional[str] = None, tracker: Optional[MetricsBackend] = None
    ):
        super().__init__(tracker=tracker)
        self._values: Dict[K, V] = dict(data) if data else {}
        self._id = id if id is not None else str(uuid.uuid4())

    @property
    def collection_name(self) -> str:
        return self._id

    @tracked("has")
    async def has(self, key: K) -> bool:
        return key in self._values

    @tracked("get")
    async def get(self, key: K, default: V | None = None) -> V | None:
        return self._values.get(key, default)

    @tracked("set")
    async def set(self, key: K, value: V) -> None:
        self._values[key] = value

    @tracked("inc")
    async def inc(self, key: K, amount: V) -> V:
        assert ensure_numeric(amount, description="amount")
        if key in self._values:
            current_value = self._values[key]
            assert ensure_numeric(current_value, description=f"value for key {key!r}")
            new_value = cast(V, current_value + amount)
            self._values[key] = new_value
        else:
            new_value = amount
            self._values[key] = new_value
        return new_value

    @tracked("chmax")
    async def chmax(self, key: K, value: V) -> V:
        assert ensure_numeric(value, description="value")
        if key in self._values:
            current_value = self._values[key]
            assert ensure_numeric(current_value, description=f"value for key {key!r}")
            if value > current_value:
                self._values[key] = value
                return value
            return current_value
        else:
            self._values[key] = value
            return value

    @tracked("pop")
    async def pop(self, key: K, default: V | None = None) -> V | None:
        return self._values.pop(key, default)

    @tracked("size")
    async 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, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
        super().__init__(tracker=tracker)
        self._lock: Mapping[AtomicLabels, _LoopAwareAsyncLock | _ThreadSafeAsyncLock] = {
            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
        }
        self._rollouts = ListBasedCollection(
            items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
        )
        self._attempts = ListBasedCollection(
            items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"], id="attempts", tracker=tracker
        )
        self._spans = ListBasedCollection(
            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"], id="spans", tracker=tracker
        )
        self._resources = ListBasedCollection(
            items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"], id="resources", tracker=tracker
        )
        self._workers = ListBasedCollection(
            items=[], item_type=Worker, primary_keys=["worker_id"], id="workers", tracker=tracker
        )
        self._rollout_queue = DequeBasedQueue(items=[], item_type=str, id="rollout_queue", tracker=tracker)
        self._span_sequence_ids = DictBasedKeyValue[str, int](
            data={}, id="span_sequence_ids", tracker=tracker
        )  # rollout_id -> sequence_id

    @property
    def collection_name(self) -> str:
        return "router"

    @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,
        *,
        mode: AtomicMode = "rw",
        snapshot: bool = False,
        labels: Optional[Sequence[AtomicLabels]] = None,
        **kwargs: Any,
    ):
        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.

        Skip the locking if mode is "r" and snapshot is False.

        This collection implementation does NOT support rollback / commit.
        """
        if mode == "r" and not snapshot:
            yield self
            return
        if not labels:
            # If no labels are provided, use all locks.
            labels = list(self._lock.keys())

        # IMPORTANT: Sort the labels to ensure consistent locking order.
        # This is necessary to avoid deadlocks when multiple threads/coroutines
        # are trying to acquire the same locks in different orders.
        labels = sorted(labels)

        async with self.tracking_context(operation="atomic", collection=self.collection_name):
            managers = [(label, self._lock[label]) for label in labels]
            async with AsyncExitStack() as stack:
                for label, manager in managers:
                    async with self.tracking_context(operation="lock", collection=label):
                        await stack.enter_async_context(manager)
                yield self

    @tracked("evict_spans_for_rollout")
    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()


class _ThreadSafeAsyncLock:
    """A thread lock powered by aiologic that can be used in both async and sync contexts.

    aiologic claims itself to be a thread-safe asyncio lock.
    """

    def __init__(self):
        self._lock = aiologic.Lock()

    async def __aenter__(self):
        await self._lock.async_acquire()
        return self

    async def __aexit__(self, *args: Any, **kwargs: Any):
        # .release() is non-blocking, so we can call it directly
        self._lock.async_release()