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

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio
import logging
import random
import re
import time
from contextlib import asynccontextmanager
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    Dict,
    Generic,
    List,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    cast,
)

import aiologic

from agentlightning.utils.metrics import MetricsBackend

if TYPE_CHECKING:
    from typing import Self

from pydantic import BaseModel, TypeAdapter
from pymongo import AsyncMongoClient, ReadPreference, ReturnDocument, WriteConcern
from pymongo.asynchronous.client_session import AsyncClientSession
from pymongo.asynchronous.collection import AsyncCollection
from pymongo.asynchronous.database import AsyncDatabase
from pymongo.errors import (
    BulkWriteError,
    CollectionInvalid,
    ConnectionFailure,
    DuplicateKeyError,
    OperationFailure,
    PyMongoError,
)
from pymongo.read_concern import ReadConcern

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

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

T_model = TypeVar("T_model", bound=BaseModel)

T_generic = TypeVar("T_generic")

T_mapping = TypeVar("T_mapping", bound=Mapping[str, Any])

T_callable = TypeVar("T_callable", bound=Callable[..., Any])

K = TypeVar("K")
V = TypeVar("V")

logger = logging.getLogger(__name__)


def resolve_mongo_error_type(exc: BaseException | None) -> str | None:
    is_transient = isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError")
    if isinstance(exc, OperationFailure):
        if is_transient:
            return f"OperationFailure-{exc.code}-Transient"
        else:
            return f"OperationFailure-{exc.code}"
    if isinstance(exc, DuplicateKeyError):
        return "DuplicateKeyError-Transient" if is_transient else "DuplicateKeyError"
    if isinstance(exc, PyMongoError):
        if is_transient:
            return f"{exc.__class__.__name__}-Transient"
        else:
            return exc.__class__.__name__
    if isinstance(exc, ConnectionFailure):
        return "ConnectionFailure-Transient" if is_transient else "ConnectionFailure"
    if is_transient:
        return "Other-Transient"
    else:
        return None


def _field_ops_to_conditions(field: str, ops: Mapping[str, Any]) -> List[Dict[str, Any]]:
    """Convert a FilterField (ops) into one or more Mongo conditions."""
    conditions: List[Dict[str, Any]] = []

    for op_name, raw_value in ops.items():
        if op_name == "exact":
            if raw_value is None:
                logger.debug(f"Skipping exact filter for field '{field}' with None value")
                continue
            conditions.append({field: raw_value})
        elif op_name == "within":
            if raw_value is None:
                logger.debug(f"Skipping within filter for field '{field}' with None value")
                continue
            try:
                iterable = list(raw_value)
            except TypeError as exc:
                raise ValueError(f"Invalid iterable for within filter for field '{field}': {raw_value!r}") from exc
            conditions.append({field: {"$in": iterable}})
        elif op_name == "contains":
            if raw_value is None:
                logger.debug(f"Skipping contains filter for field '{field}' with None value")
                continue
            value = str(raw_value)
            pattern = f".*{re.escape(value)}.*"
            conditions.append({field: {"$regex": pattern, "$options": "i"}})
        else:
            raise ValueError(f"Unsupported filter operator '{op_name}' for field '{field}'")

    return conditions


def _build_mongo_filter(filter_options: Optional[FilterOptions]) -> Dict[str, Any]:
    """Translate FilterOptions into a MongoDB filter dict."""
    normalized, must_filters, aggregate = normalize_filter_options(filter_options)

    regular_conditions: List[Dict[str, Any]] = []
    must_conditions: List[Dict[str, Any]] = []

    # Normal filters
    if normalized:
        for field_name, ops in normalized.items():
            regular_conditions.extend(_field_ops_to_conditions(field_name, ops))

    # Must filters
    if must_filters:
        for field_name, ops in must_filters.items():
            must_conditions.extend(_field_ops_to_conditions(field_name, ops))

    # No filters at all
    if not regular_conditions and not must_conditions:
        return {}

    # Aggregate logic for regular conditions; _must always ANDs in.
    if aggregate == "and":
        all_conds = regular_conditions + must_conditions
        if len(all_conds) == 1:
            return all_conds[0]
        return {"$and": all_conds}

    # aggregate == "or"
    if regular_conditions and must_conditions:
        # (OR of regular) AND (all must)
        if len(regular_conditions) == 1:
            or_part: Dict[str, Any] = regular_conditions[0]
        else:
            or_part = {"$or": regular_conditions}

        and_parts: List[Dict[str, Any]] = [or_part] + must_conditions
        if len(and_parts) == 1:
            return and_parts[0]
        return {"$and": and_parts}

    if regular_conditions:
        if len(regular_conditions) == 1:
            return regular_conditions[0]
        return {"$or": regular_conditions}

    # Only must conditions
    if len(must_conditions) == 1:
        return must_conditions[0]
    return {"$and": must_conditions}


async def _ensure_collection(
    db: AsyncDatabase[Mapping[str, Any]],
    collection_name: str,
    primary_keys: Optional[Sequence[str]] = None,
    extra_indexes: Optional[Sequence[Sequence[str]]] = None,
) -> bool:
    """Ensure the backing MongoDB collection exists.

    This method is idempotent and safe to call multiple times.
    """
    # Create collection if it doesn't exist yet
    try:
        await db.create_collection(collection_name)
    except CollectionInvalid as exc:
        # Thrown if collection already exists
        logger.debug(f"Collection '{collection_name}' may have already existed. No need to create it: {exc!r}")
    except OperationFailure as exc:
        logger.debug(f"Failed to create collection '{collection_name}'. Probably already exists: {exc!r}")
        # Some servers use OperationFailure w/ specific codes for "NamespaceExists"
        if exc.code in (48, 68):  # 48: NamespaceExists, 68: already exists on older versions
            pass
        else:
            raise

    # Optionally create a unique index on primary keys (scoped by partition_id)
    if primary_keys:
        # Always include the partition field in the unique index.
        keys = [("partition_id", 1)] + [(pk, 1) for pk in primary_keys]
        try:
            await db[collection_name].create_index(keys, name=f"uniq_partition_{'_'.join(primary_keys)}", unique=True)
        except OperationFailure as exc:
            logger.debug(f"Index for collection '{collection_name}' already exists. No need to create it: {exc!r}")
            # Ignore "index already exists" type errors
            if exc.code in (68, 85):  # IndexOptionsConflict, etc.
                pass
            else:
                raise

    # Optionally create extra indexes
    if extra_indexes:
        for index in extra_indexes:
            try:
                await db[collection_name].create_index(index, name=f"idx_{'_'.join(index)}")
            except OperationFailure as exc:
                logger.debug(f"Index for collection '{collection_name}' already exists. No need to create it: {exc!r}")
                # Ignore "index already exists" type errors
                if exc.code in (68, 85):  # IndexOptionsConflict, etc.
                    pass
                else:
                    raise

    return True


class MongoClientPool(Generic[T_mapping]):
    """A pool of MongoDB clients, each bound to a specific event loop.

    The pool lazily creates `AsyncMongoClient` instances per event loop using the provided
    connection parameters, ensuring we never try to reuse a client across loops.
    """

    def __init__(self, *, mongo_uri: str, mongo_client_kwargs: Mapping[str, Any] | None = None):
        self._get_collection_lock = aiologic.Lock()
        self._get_client_lock = aiologic.Lock()
        self._mongo_uri = mongo_uri
        self._mongo_client_kwargs = dict(mongo_client_kwargs or {})
        self._client_pool: Dict[int, AsyncMongoClient[T_mapping]] = {}
        self._collection_pool: Dict[Tuple[int, str, str], AsyncCollection[T_mapping]] = {}

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
        await self.close()

    async def close(self) -> None:
        """Close all clients currently tracked by the pool."""

        async with self._get_client_lock, self._get_collection_lock:
            clients = list(self._client_pool.values())
            self._client_pool.clear()
            self._collection_pool.clear()

        for client in clients:
            try:
                await client.close()
            except Exception:
                logger.exception("Error closing MongoDB client: %s", client)

    async def get_client(self) -> AsyncMongoClient[T_mapping]:
        loop = asyncio.get_running_loop()
        key = id(loop)

        # If there is already a client specifically for this loop, return it.
        existing = self._client_pool.get(key)
        if existing is not None:
            await existing.aconnect()  # This actually does nothing if the client is already connected.
            return existing

        async with self._get_client_lock:
            # Another coroutine may have already created the client.
            if key in self._client_pool:
                await self._client_pool[key].aconnect()
                return self._client_pool[key]

            # Create a new client for this loop.
            client = AsyncMongoClient[T_mapping](self._mongo_uri, **self._mongo_client_kwargs)
            await client.aconnect()
            self._client_pool[key] = client
            return client

    async def get_collection(self, database_name: str, collection_name: str) -> AsyncCollection[T_mapping]:
        loop = asyncio.get_running_loop()
        key = (id(loop), database_name, collection_name)
        if key in self._collection_pool:
            return self._collection_pool[key]

        async with self._get_collection_lock:
            # Another coroutine may have already created the collection.
            if key in self._collection_pool:
                return self._collection_pool[key]

            # Create a new collection for this loop.
            client = await self.get_client()
            collection = client[database_name][collection_name]
            self._collection_pool.setdefault(key, collection)
            return collection


class MongoBasedCollection(Collection[T_model]):
    """Mongo-based implementation of Collection.

    Args:
        client_pool: The pool of MongoDB clients.
        database_name: The name of the database.
        collection_name: The name of the collection.
        partition_id: The partition ID. Used to partition the collection into multiple collections.
        primary_keys: The primary keys of the collection.
        item_type: The type of the items in the collection.
        extra_indexes: The extra indexes to create on the collection.
        tracker: The metrics tracker to use.
    """

    def __init__(
        self,
        client_pool: MongoClientPool[Mapping[str, Any]],
        database_name: str,
        collection_name: str,
        partition_id: str,
        primary_keys: Sequence[str],
        item_type: Type[T_model],
        extra_indexes: Sequence[Sequence[str]] = [],
        tracker: MetricsBackend | None = None,
    ):
        super().__init__(tracker=tracker)
        self._client_pool = client_pool
        self._database_name = database_name
        self._collection_name = collection_name
        self._partition_id = partition_id
        self._collection_created = False
        self._extra_indexes = [list(index) for index in extra_indexes]
        self._session: Optional[AsyncClientSession] = None

        if not primary_keys:
            raise ValueError("primary_keys must be non-empty")
        self._primary_keys = list(primary_keys)

        if not issubclass(item_type, BaseModel):  # type: ignore
            raise ValueError(f"item_type must be a subclass of BaseModel, got {item_type.__name__}")
        self._item_type = item_type

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

    @property
    def extra_tracking_labels(self) -> Mapping[str, str]:
        return {
            "database": self._database_name,
        }

    @tracked("ensure_collection")
    async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
        """Ensure the backing MongoDB collection exists (and optionally its indexes).

        This method is idempotent and safe to call multiple times.

        It will also create a unique index across the configured primary key fields.
        """
        if not self._collection_created:
            client = await self._client_pool.get_client()
            self._collection_created = await _ensure_collection(
                client[self._database_name], self._collection_name, self._primary_keys, self._extra_indexes
            )

        return await self._client_pool.get_collection(self._database_name, self._collection_name)

    def with_session(self, session: AsyncClientSession) -> MongoBasedCollection[T_model]:
        """Create a new collection with the same configuration but a new session."""
        collection = MongoBasedCollection(
            client_pool=self._client_pool,
            database_name=self._database_name,
            collection_name=self._collection_name,
            partition_id=self._partition_id,
            primary_keys=self._primary_keys,
            item_type=self._item_type,
            extra_indexes=self._extra_indexes,
            tracker=self._tracker,
        )
        collection._collection_created = self._collection_created
        collection._session = session
        return collection

    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_model]:
        return self._item_type

    @tracked("size")
    async def size(self) -> int:
        collection = await self.ensure_collection()
        return await collection.count_documents({"partition_id": self._partition_id}, session=self._session)

    def _pk_filter(self, item: T_model) -> Dict[str, Any]:
        """Build a Mongo filter for the primary key(s) of a model instance."""
        data = item.model_dump()
        missing = [pk for pk in self._primary_keys if pk not in data]
        if missing:
            raise ValueError(f"Missing primary key fields {missing} on item {item!r}")
        pk_filter: Dict[str, Any] = {"partition_id": self._partition_id}
        pk_filter.update({pk: data[pk] for pk in self._primary_keys})
        return pk_filter

    def _render_pk_values(self, values: Sequence[Any]) -> str:
        return ", ".join(f"{pk}={value!r}" for pk, value in zip(self._primary_keys, values))

    def _ensure_item_type(self, item: T_model) -> None:
        if not isinstance(item, self._item_type):
            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")

    def _inject_partition_filter(self, filter: Optional[FilterOptions]) -> Dict[str, Any]:
        """Ensure every query is scoped to this collection's partition."""
        combined: Dict[str, Any]
        if filter is None:
            combined = {}
        else:
            combined = dict(filter)

        partition_must = {"partition_id": {"exact": self._partition_id}}
        existing_must = combined.get("_must")
        if existing_must is None:
            combined["_must"] = partition_must
            return combined

        if isinstance(existing_must, Mapping):
            combined["_must"] = [existing_must, partition_must]
        elif isinstance(existing_must, Sequence) and not isinstance(existing_must, (str, bytes)):
            combined["_must"] = [*existing_must, partition_must]
        else:
            raise TypeError("`_must` filters must be a mapping or sequence of mappings")

        return combined

    def _model_validate_item(self, raw: Mapping[str, Any]) -> T_model:
        item_type_has_id = "_id" in self._item_type.model_fields
        # Remove _id from the raw document if the item type does not have it.
        if not item_type_has_id:
            raw = {k: v for k, v in raw.items() if k != "_id"}
        # Convert Mongo document to Pydantic model
        return self._item_type.model_validate(raw)  # type: ignore[arg-type]

    @tracked("query")
    async def query(
        self,
        filter: Optional[FilterOptions] = None,
        sort: Optional[SortOptions] = None,
        limit: int = -1,
        offset: int = 0,
    ) -> PaginatedResult[T_model]:
        """Mongo-based implementation of Collection.query.

        The handling of null-values in sorting is different from memory-based implementation.
        In MongoDB, null values are treated as less than non-null values.
        """
        collection = await self.ensure_collection()

        combined = self._inject_partition_filter(filter)
        mongo_filter = _build_mongo_filter(cast(FilterOptions, combined))

        total = await collection.count_documents(mongo_filter, session=self._session)

        if limit == 0:
            return PaginatedResult[T_model](items=[], limit=0, offset=offset, total=total)

        cursor = collection.find(mongo_filter, session=self._session)

        sort_name, sort_order = resolve_sort_options(sort)
        if sort_name is not None:
            model_fields = getattr(self._item_type, "model_fields", {})
            if sort_name not in model_fields:
                raise ValueError(
                    f"Failed to sort items by '{sort_name}': field does not exist on {self._item_type.__name__}"
                )
            direction = 1 if sort_order == "asc" else -1
            cursor = cursor.sort(sort_name, direction)

        if offset > 0:
            cursor = cursor.skip(offset)
        if limit >= 0:
            cursor = cursor.limit(limit)

        items: List[T_model] = []
        async for raw in cursor:
            items.append(self._model_validate_item(raw))

        return PaginatedResult[T_model](items=items, limit=limit, offset=offset, total=total)

    @tracked("get")
    async def get(
        self,
        filter: Optional[FilterOptions] = None,
        sort: Optional[SortOptions] = None,
    ) -> Optional[T_model]:
        collection = await self.ensure_collection()

        combined = self._inject_partition_filter(filter)
        mongo_filter = _build_mongo_filter(cast(FilterOptions, combined))

        sort_name, sort_order = resolve_sort_options(sort)
        mongo_sort: Optional[List[Tuple[str, int]]] = None
        if sort_name is not None:
            model_fields = getattr(self._item_type, "model_fields", {})
            if sort_name not in model_fields:
                raise ValueError(
                    f"Failed to sort items by '{sort_name}': field does not exist on {self._item_type.__name__}"
                )
            direction = 1 if sort_order == "asc" else -1
            mongo_sort = [(sort_name, direction)]

        raw = await collection.find_one(mongo_filter, sort=mongo_sort, session=self._session)

        if raw is None:
            return None

        return self._model_validate_item(raw)

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

        The implementation does NOT do checks for duplicate primary keys,
        neither within the same insert call nor across different insert calls.
        It relies on the database to enforce uniqueness via indexes.
        """
        if not items:
            return

        collection = await self.ensure_collection()
        docs: List[Mapping[str, Any]] = []
        for item in items:
            self._ensure_item_type(item)
            doc = item.model_dump()
            doc["partition_id"] = self._partition_id
            docs.append(doc)

        if not docs:
            return

        try:
            async with self.tracking_context("insert.insert_many", self._collection_name):
                await collection.insert_many(docs, session=self._session)
        except DuplicateKeyError as exc:
            # In case the DB enforces uniqueness via index, normalize to ValueError
            raise DuplicatedPrimaryKeyError("Duplicated primary key(s) while inserting items") from exc
        except BulkWriteError as exc:
            write_errors = exc.details.get("writeErrors", [])
            if write_errors and write_errors[0].get("code") == 11000:
                raise DuplicatedPrimaryKeyError("Duplicated primary key(s) while inserting items") from exc
            raise

    @tracked("update")
    async def update(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
        if not items:
            return []

        updated_items: List[T_model] = []
        collection = await self.ensure_collection()

        for item in items:
            self._ensure_item_type(item)
            pk_filter = self._pk_filter(item)
            doc = item.model_dump()
            doc["partition_id"] = self._partition_id

            updated_doc = None

            # Branch 1: Full Replace
            if update_fields is None:
                async with self.tracking_context("update.find_one_and_replace", self._collection_name):
                    updated_doc = await collection.find_one_and_replace(
                        filter=pk_filter,
                        replacement=doc,
                        session=self._session,
                        return_document=ReturnDocument.AFTER,  # Returns the new version
                    )

            # Branch 2: Partial Update
            else:
                update_doc = {field: doc[field] for field in update_fields if field in doc}
                async with self.tracking_context("update.find_one_and_update", self._collection_name):
                    updated_doc = await collection.find_one_and_update(
                        filter=pk_filter,
                        update={"$set": update_doc},
                        session=self._session,
                        return_document=ReturnDocument.AFTER,  # Returns the new version
                    )

            # Validation and Reconstruction
            if updated_doc is None:  # type: ignore
                raise ValueError(f"Item with primary key(s) {pk_filter} does not exist")

            # Re-instantiate the model from the raw MongoDB dictionary.
            new_item = self._model_validate_item(updated_doc)
            updated_items.append(new_item)

        return updated_items

    @tracked("upsert")
    async def upsert(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
        if not items:
            return []

        upserted_items: List[T_model] = []
        collection = await self.ensure_collection()

        for item in items:
            self._ensure_item_type(item)
            pk_filter = self._pk_filter(item)

            insert_doc = item.model_dump()
            insert_doc["partition_id"] = self._partition_id

            # If update_fields is None, we update ALL fields (standard upsert behavior).
            # Otherwise, we only update specific fields, but insert the full doc if it's new.
            target_fields = update_fields if update_fields is not None else list(insert_doc.keys())

            # 1. $set: Fields that should be overwritten if the document exists
            update_subset = {field: insert_doc[field] for field in target_fields if field in insert_doc}

            # 2. $setOnInsert: Fields that are only set if we are creating a NEW document
            # (Everything in the model that isn't in the update_subset)
            set_on_insert = {k: v for k, v in insert_doc.items() if k not in update_subset}

            update_spec: Dict[str, Dict[str, Any]] = {}
            if set_on_insert:
                update_spec["$setOnInsert"] = set_on_insert
            if update_subset:
                update_spec["$set"] = update_subset

            async with self.tracking_context("upsert.find_one_and_update", self._collection_name):
                result_doc = await collection.find_one_and_update(
                    filter=pk_filter,
                    update=update_spec,
                    upsert=True,
                    session=self._session,
                    return_document=ReturnDocument.AFTER,
                )

            if result_doc is None:  # pyright: ignore[reportUnnecessaryComparison]
                raise RuntimeError(f"Upsert resulted in no document for filter: {pk_filter}")

            # Because upsert=True, result_doc is guaranteed to be not None
            new_item = self._model_validate_item(result_doc)
            upserted_items.append(new_item)

        return upserted_items

    @tracked("delete")
    async def delete(self, items: Sequence[T_model]) -> None:
        if not items:
            return

        collection = await self.ensure_collection()
        for item in items:
            self._ensure_item_type(item)
            pk_filter = self._pk_filter(item)
            async with self.tracking_context("delete.delete_one", self._collection_name):
                result = await collection.delete_one(pk_filter, session=self._session)
            if result.deleted_count == 0:
                raise ValueError(f"Item with primary key(s) {pk_filter} does not exist")


class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
    """Mongo-based implementation of Queue backed by a MongoDB collection.

    Items are stored append-only; dequeue marks items as consumed instead of deleting them.
    """

    def __init__(
        self,
        client_pool: MongoClientPool[Mapping[str, Any]],
        database_name: str,
        collection_name: str,
        partition_id: str,
        item_type: Type[T_generic],
        tracker: MetricsBackend | None = None,
    ) -> None:
        """
        Args:
            client_pool: The pool of MongoDB clients.
            database_name: The name of the database.
            collection_name: The name of the collection backing the queue.
            partition_id: Partition identifier; allows multiple logical queues in one collection.
            item_type: The Python type of queue items (primitive or BaseModel subclass).
        """
        super().__init__(tracker=tracker)
        self._client_pool = client_pool
        self._database_name = database_name
        self._collection_name = collection_name
        self._partition_id = partition_id
        self._item_type = item_type
        self._adapter: TypeAdapter[T_generic] = TypeAdapter(item_type)
        self._collection_created = False

        self._session: Optional[AsyncClientSession] = None

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

    @property
    def extra_tracking_labels(self) -> Mapping[str, str]:
        return {
            "database": self._database_name,
        }

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

    @tracked("ensure_collection")
    async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
        """Ensure the backing collection exists.

        If it already exists, it returns the existing collection.
        """
        if not self._collection_created:
            client = await self._client_pool.get_client()
            self._collection_created = await _ensure_collection(
                client[self._database_name], self._collection_name, primary_keys=["consumed", "_id"]
            )
        return await self._client_pool.get_collection(self._database_name, self._collection_name)

    def with_session(self, session: AsyncClientSession) -> MongoBasedQueue[T_generic]:
        queue = MongoBasedQueue(
            client_pool=self._client_pool,
            database_name=self._database_name,
            collection_name=self._collection_name,
            partition_id=self._partition_id,
            item_type=self._item_type,
            tracker=self._tracker,
        )
        queue._collection_created = self._collection_created
        queue._session = session
        return queue

    @tracked("has")
    async def has(self, item: T_generic) -> bool:
        collection = await self.ensure_collection()
        encoded = self._adapter.dump_python(item, mode="python")
        doc = await collection.find_one(
            {
                "partition_id": self._partition_id,
                "consumed": False,
                "value": encoded,
            },
            session=self._session,
        )
        return doc is not None

    @tracked("enqueue")
    async def enqueue(self, items: Sequence[T_generic]) -> Sequence[T_generic]:
        if not items:
            return []

        collection = await self.ensure_collection()
        docs: List[Mapping[str, Any]] = []
        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__}")
            docs.append(
                {
                    "partition_id": self._partition_id,
                    "value": self._adapter.dump_python(item, mode="python"),
                    "consumed": False,
                    "created_at": datetime.now(),
                }
            )

        async with self.tracking_context("enqueue.insert_many", self.collection_name):
            await collection.insert_many(docs, session=self._session)
        return list(items)

    @tracked("dequeue")
    async def dequeue(self, limit: int = 1) -> Sequence[T_generic]:
        if limit <= 0:
            return []

        collection = await self.ensure_collection()
        results: list[T_generic] = []

        # Atomic claim loop using find_one_and_update
        for _ in range(limit):
            async with self.tracking_context("dequeue.find_one_and_update", self.collection_name):
                doc = await collection.find_one_and_update(
                    {
                        "partition_id": self._partition_id,
                        "consumed": False,
                    },
                    {"$set": {"consumed": True}},
                    sort=[("_id", 1)],  # FIFO using insertion order
                    return_document=True,
                    session=self._session,
                )
            if doc is None:  # type: ignore
                # No more items to dequeue
                break

            raw_value = doc["value"]
            item = self._adapter.validate_python(raw_value)
            results.append(item)

        return results

    @tracked("peek")
    async def peek(self, limit: int = 1) -> Sequence[T_generic]:
        if limit <= 0:
            return []

        collection = await self.ensure_collection()
        async with self.tracking_context("peek.find", self.collection_name):
            cursor = (
                collection.find(
                    {
                        "partition_id": self._partition_id,
                        "consumed": False,
                    },
                    session=self._session,
                )
                .sort("_id", 1)
                .limit(limit)
            )

        items: list[T_generic] = []
        async for doc in cursor:
            raw_value = doc["value"]
            items.append(self._adapter.validate_python(raw_value))

        return items

    @tracked("size")
    async def size(self) -> int:
        collection = await self.ensure_collection()
        return await collection.count_documents(
            {
                "partition_id": self._partition_id,
                "consumed": False,
            },
            session=self._session,
        )


class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
    """Mongo-based implementation of KeyValue."""

    def __init__(
        self,
        client_pool: MongoClientPool[Mapping[str, Any]],
        database_name: str,
        collection_name: str,
        partition_id: str,
        key_type: Type[K],
        value_type: Type[V],
        tracker: MetricsBackend | None = None,
    ) -> None:
        """
        Args:
            client_pool: The pool of MongoDB clients.
            database_name: The name of the database.
            collection_name: The name of the collection backing the key-value store.
            partition_id: Partition identifier; allows multiple logical maps in one collection.
            key_type: The Python type of keys (primitive or BaseModel).
            value_type: The Python type of values (primitive or BaseModel).
            tracker: The metrics tracker to use.
        """
        super().__init__(tracker=tracker)
        self._client_pool = client_pool
        self._database_name = database_name
        self._collection_name = collection_name
        self._partition_id = partition_id
        self._key_type = key_type
        self._value_type = value_type
        self._key_adapter: TypeAdapter[K] = TypeAdapter(key_type)
        self._value_adapter: TypeAdapter[V] = TypeAdapter(value_type)
        self._collection_created = False

        self._session: Optional[AsyncClientSession] = None

    @property
    def extra_tracking_labels(self) -> Mapping[str, str]:
        return {
            "database": self._database_name,
        }

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

    @tracked("ensure_collection")
    async def ensure_collection(self, *, create_indexes: bool = True) -> AsyncCollection[Mapping[str, Any]]:
        """Ensure the backing collection exists (and optionally its indexes)."""
        if not self._collection_created:
            client = await self._client_pool.get_client()
            self._collection_created = await _ensure_collection(
                client[self._database_name], self._collection_name, primary_keys=["key"]
            )
        return await self._client_pool.get_collection(self._database_name, self._collection_name)

    def with_session(self, session: AsyncClientSession) -> MongoBasedKeyValue[K, V]:
        key_value = MongoBasedKeyValue(
            client_pool=self._client_pool,
            database_name=self._database_name,
            collection_name=self._collection_name,
            partition_id=self._partition_id,
            key_type=self._key_type,
            value_type=self._value_type,
            tracker=self._tracker,
        )
        key_value._collection_created = self._collection_created
        key_value._session = session

        return key_value

    @tracked("has")
    async def has(self, key: K) -> bool:
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        doc = await collection.find_one(
            {
                "partition_id": self._partition_id,
                "key": encoded_key,
            },
            session=self._session,
        )
        return doc is not None

    @tracked("get")
    async def get(self, key: K, default: V | None = None) -> V | None:
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        doc = await collection.find_one(
            {
                "partition_id": self._partition_id,
                "key": encoded_key,
            },
            session=self._session,
        )
        if doc is None:
            return default

        raw_value = doc["value"]
        return self._value_adapter.validate_python(raw_value)

    @tracked("set")
    async def set(self, key: K, value: V) -> None:
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        encoded_value = self._value_adapter.dump_python(value, mode="python")
        try:
            async with self.tracking_context("set.replace_one", self.collection_name):
                await collection.replace_one(
                    {
                        "partition_id": self._partition_id,
                        "key": encoded_key,
                    },
                    {
                        "partition_id": self._partition_id,
                        "key": encoded_key,
                        "value": encoded_value,
                    },
                    upsert=True,
                    session=self._session,
                )
        except DuplicateKeyError as exc:
            # Very unlikely with replace_one+upsert, but normalize anyway.
            raise DuplicatedPrimaryKeyError("Duplicate key error while setting key-value item") from exc

    @tracked("inc")
    async def inc(self, key: K, amount: V) -> V:
        assert ensure_numeric(amount, description="amount")
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        encoded_amount = self._value_adapter.dump_python(amount, mode="python")
        try:
            async with self.tracking_context("inc.find_one_and_update", self.collection_name):
                doc = await collection.find_one_and_update(
                    {
                        "partition_id": self._partition_id,
                        "key": encoded_key,
                    },
                    {
                        "$inc": {"value": encoded_amount},
                    },
                    upsert=True,
                    return_document=ReturnDocument.AFTER,
                    session=self._session,
                )
        except OperationFailure as exc:
            if exc.code == 14 or "Cannot apply $inc" in str(exc):
                raise TypeError(f"value for key {key!r} is not numeric") from exc
            raise
        if doc is None:  # type: ignore
            raise RuntimeError("Failed to increment value; MongoDB did not return a document")
        raw_value = doc["value"]
        return self._value_adapter.validate_python(raw_value)

    @tracked("chmax")
    async def chmax(self, key: K, value: V) -> V:
        assert ensure_numeric(value, description="value")
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        encoded_value = self._value_adapter.dump_python(value, mode="python")
        try:
            async with self.tracking_context("chmax.find_one_and_update", self.collection_name):
                doc = await collection.find_one_and_update(
                    {
                        "partition_id": self._partition_id,
                        "key": encoded_key,
                    },
                    {
                        "$max": {"value": encoded_value},
                    },
                    upsert=True,
                    return_document=ReturnDocument.AFTER,
                    session=self._session,
                )
        except OperationFailure as exc:
            if exc.code == 14 or "Cannot apply $max" in str(exc):
                raise TypeError(f"value for key {key!r} is not numeric") from exc
            raise
        if doc is None:  # type: ignore
            raise RuntimeError("Failed to update value; MongoDB did not return a document")
        raw_value = doc["value"]
        return self._value_adapter.validate_python(raw_value)

    @tracked("pop")
    async def pop(self, key: K, default: V | None = None) -> V | None:
        collection = await self.ensure_collection()
        encoded_key = self._key_adapter.dump_python(key, mode="python")
        doc = await collection.find_one_and_delete(
            {
                "partition_id": self._partition_id,
                "key": encoded_key,
            },
            session=self._session,
        )
        if doc is None:  # type: ignore
            return default

        raw_value = doc["value"]
        return self._value_adapter.validate_python(raw_value)

    @tracked("size")
    async def size(self) -> int:
        collection = await self.ensure_collection()
        return await collection.count_documents(
            {
                "partition_id": self._partition_id,
            },
            session=self._session,
        )


class MongoLightningCollections(LightningCollections):
    """Mongo implementation of LightningCollections using MongoDB collections.

    Serves as the storage base for [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore].
    """

    def __init__(
        self,
        client_pool: MongoClientPool[Mapping[str, Any]],
        database_name: str,
        partition_id: str,
        rollouts: Optional[MongoBasedCollection[Rollout]] = None,
        attempts: Optional[MongoBasedCollection[Attempt]] = None,
        spans: Optional[MongoBasedCollection[Span]] = None,
        resources: Optional[MongoBasedCollection[ResourcesUpdate]] = None,
        workers: Optional[MongoBasedCollection[Worker]] = None,
        rollout_queue: Optional[MongoBasedQueue[str]] = None,
        span_sequence_ids: Optional[MongoBasedKeyValue[str, int]] = None,
        tracker: MetricsBackend | None = None,
    ):
        super().__init__(tracker=tracker, extra_labels=["database"])
        self._client_pool = client_pool
        self._database_name = database_name
        self._partition_id = partition_id
        self._collection_ensured = False
        self._lock = aiologic.Lock()  # used for generic atomic operations like scan debounce seconds
        self._rollouts = (
            rollouts
            if rollouts is not None
            else MongoBasedCollection(
                self._client_pool,
                self._database_name,
                "rollouts",
                self._partition_id,
                ["rollout_id"],
                Rollout,
                [["status"]],
                tracker=self._tracker,
            )
        )
        self._attempts = (
            attempts
            if attempts is not None
            else MongoBasedCollection(
                self._client_pool,
                self._database_name,
                "attempts",
                self._partition_id,
                ["rollout_id", "attempt_id"],
                Attempt,
                [["status"], ["sequence_id"]],
                tracker=self._tracker,
            )
        )
        self._spans = (
            spans
            if spans is not None
            else MongoBasedCollection(
                self._client_pool,
                self._database_name,
                "spans",
                self._partition_id,
                ["rollout_id", "attempt_id", "span_id"],
                Span,
                [["sequence_id"]],
                tracker=self._tracker,
            )
        )
        self._resources = (
            resources
            if resources is not None
            else MongoBasedCollection(
                self._client_pool,
                self._database_name,
                "resources",
                self._partition_id,
                ["resources_id"],
                ResourcesUpdate,
                ["update_time"],
                tracker=self._tracker,
            )
        )
        self._workers = (
            workers
            if workers is not None
            else MongoBasedCollection(
                self._client_pool,
                self._database_name,
                "workers",
                self._partition_id,
                ["worker_id"],
                Worker,
                ["status"],
                tracker=self._tracker,
            )
        )
        self._rollout_queue = (
            rollout_queue
            if rollout_queue is not None
            else MongoBasedQueue(
                self._client_pool,
                self._database_name,
                "rollout_queue",
                self._partition_id,
                str,
                tracker=self._tracker,
            )
        )
        self._span_sequence_ids = (
            span_sequence_ids
            if span_sequence_ids is not None
            else MongoBasedKeyValue(
                self._client_pool,
                self._database_name,
                "span_sequence_ids",
                self._partition_id,
                str,
                int,
                tracker=self._tracker,
            )
        )

    @property
    def collection_name(self) -> str:
        return "router"  # Special collection name for tracking transactions

    @property
    def extra_tracking_labels(self) -> Mapping[str, str]:
        return {
            "database": self._database_name,
        }

    def with_session(self, session: AsyncClientSession) -> Self:
        instance = self.__class__(
            client_pool=self._client_pool,
            database_name=self._database_name,
            partition_id=self._partition_id,
            rollouts=self._rollouts.with_session(session),
            attempts=self._attempts.with_session(session),
            spans=self._spans.with_session(session),
            resources=self._resources.with_session(session),
            workers=self._workers.with_session(session),
            rollout_queue=self._rollout_queue.with_session(session),
            span_sequence_ids=self._span_sequence_ids.with_session(session),
            tracker=self._tracker,
        )
        instance._collection_ensured = self._collection_ensured
        return instance

    @property
    def rollouts(self) -> MongoBasedCollection[Rollout]:
        return self._rollouts

    @property
    def attempts(self) -> MongoBasedCollection[Attempt]:
        return self._attempts

    @property
    def spans(self) -> MongoBasedCollection[Span]:
        return self._spans

    @property
    def resources(self) -> MongoBasedCollection[ResourcesUpdate]:
        return self._resources

    @property
    def workers(self) -> MongoBasedCollection[Worker]:
        return self._workers

    @property
    def rollout_queue(self) -> MongoBasedQueue[str]:
        return self._rollout_queue

    @property
    def span_sequence_ids(self) -> MongoBasedKeyValue[str, int]:
        return self._span_sequence_ids

    @tracked("ensure_collections")
    async def _ensure_collections(self) -> None:
        """Ensure all collections exist."""
        if self._collection_ensured:
            return
        await self._rollouts.ensure_collection()
        await self._attempts.ensure_collection()
        await self._spans.ensure_collection()
        await self._resources.ensure_collection()
        await self._workers.ensure_collection()
        await self._rollout_queue.ensure_collection()
        await self._span_sequence_ids.ensure_collection()
        self._collection_ensured = True

    @asynccontextmanager
    async def _lock_manager(self, labels: Optional[Sequence[AtomicLabels]]):
        if labels is None or "generic" not in labels:
            yield

        else:
            # Only lock the generic label.
            try:
                async with self.tracking_context("lock", self.collection_name):
                    await self._lock.async_acquire()
                yield
            finally:
                self._lock.async_release()

    @asynccontextmanager
    async def atomic(
        self,
        mode: AtomicMode = "rw",
        snapshot: bool = False,
        commit: bool = False,
        labels: Optional[Sequence[AtomicLabels]] = None,
        *args: Any,
        **kwargs: Any,
    ):
        """Perform a atomic operation on the collections."""
        if commit:
            raise ValueError("Commit should be used with execute() instead.")
        async with self._lock_manager(labels):
            async with self.tracking_context("atomic", self.collection_name):
                # First step: ensure all collections exist before going into the atomic block
                if not self._collection_ensured:
                    await self._ensure_collections()
                # Execute directly without commit
                yield self

    @tracked("execute")
    async def execute(
        self,
        callback: Callable[[Self], Awaitable[T_generic]],
        *,
        mode: AtomicMode = "rw",
        snapshot: bool = False,
        commit: bool = False,
        labels: Optional[Sequence[AtomicLabels]] = None,
        **kwargs: Any,
    ) -> T_generic:
        """Execute the given callback within an atomic operation, and with retries on transient errors."""
        if not self._collection_ensured:
            await self._ensure_collections()
        client = await self._client_pool.get_client()

        # If commit is not turned on, just execute the callback directly.
        if not commit:
            async with self._lock_manager(labels):
                return await callback(self)

        # If snapshot is enabled, use snapshot read concern.
        read_concern = ReadConcern("snapshot") if snapshot else ReadConcern("local")
        # If mode is "r", write_concern is not needed.
        write_concern = WriteConcern("majority") if mode != "r" else None

        async with client.start_session() as session:
            collections = self.with_session(session)
            try:
                async with self._lock_manager(labels):
                    return await self.with_transaction(session, collections, callback, read_concern, write_concern)
            except (ConnectionFailure, OperationFailure) as exc:
                # Un-retryable errors.
                raise RuntimeError("Transaction failed with connection or operation error") from exc

    @tracked("with_transaction")
    async def with_transaction(
        self,
        session: AsyncClientSession,
        collections: Self,
        callback: Callable[[Self], Awaitable[T_generic]],
        read_concern: ReadConcern,
        write_concern: Optional[WriteConcern],
    ) -> T_generic:
        # This will start a transaction, run transaction callback, and commit.
        # It will also transparently retry on some transient errors.
        # Expanded implementation of with_transaction from client_session
        read_preference = ReadPreference.PRIMARY
        transaction_retry_time_limit = 120
        start_time = time.monotonic()

        def _within_time_limit() -> bool:
            return time.monotonic() - start_time < transaction_retry_time_limit

        async def _jitter_before_retry() -> None:
            async with self.tracking_context("execute.jitter", self.collection_name):
                await asyncio.sleep(random.uniform(0, 0.05))

        while True:
            await session.start_transaction(read_concern, write_concern, read_preference)

            try:
                # The _session is always the same within one transaction,
                # so we can use the same collections object.
                async with self.tracking_context("execute.callback", self.collection_name):
                    ret = await callback(collections)
            # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
            except BaseException as exc:
                if session.in_transaction:
                    await session.abort_transaction()
                if (
                    isinstance(exc, PyMongoError)
                    and exc.has_error_label("TransientTransactionError")
                    and _within_time_limit()
                ):
                    # Retry the entire transaction.
                    await _jitter_before_retry()
                    continue
                raise

            if not session.in_transaction:
                # Assume callback intentionally ended the transaction.
                return ret

            # Tracks the commit operation.
            async with self.tracking_context("execute.commit", self.collection_name):
                # Loop until the commit succeeds or we hit the time limit.
                while True:
                    # Tracks the commit attempt.
                    try:
                        async with self.tracking_context("execute.commit_once", self.collection_name):
                            await session.commit_transaction()
                    except PyMongoError as exc:
                        if (
                            exc.has_error_label("UnknownTransactionCommitResult")
                            and _within_time_limit()
                            and not (isinstance(exc, OperationFailure) and exc.code == 50)  # max_time_expired_error
                        ):
                            # Retry the commit.
                            await _jitter_before_retry()
                            continue

                        if exc.has_error_label("TransientTransactionError") and _within_time_limit():
                            # Retry the entire transaction.
                            await _jitter_before_retry()
                            break
                        raise

                    # Commit succeeded.
                    return ret