update rag workflow

.github/workflows/badge-claude-code.yml

@@ -0,0 +1,29 @@
+name: Badge - Claude Code
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Claude Code
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-claude-code.yml', label: 'claude-code', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-examples.yml

@@ -9,6 +9,7 @@ on:
       - Examples - Unsloth
       - Examples - Tinker
       - Examples - Azure
+      - Examples - Claude Code
     types: [completed]
 
   workflow_dispatch:
@@ -35,5 +36,6 @@ jobs:
               { workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
               { workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
               { workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
+              { workflow: 'examples-claude-code.yml', label: 'examples-claude-code.stable', variants: ['stable'] },
             ];
             await badgeAggregation({ github, context, core, dependencies });

.github/workflows/benchmark.yml

@@ -127,6 +127,7 @@ jobs:
             sleep 1
           done
           echo "Store did not become ready in time" >&2
+          docker compose -f "$COMPOSE_FILE" logs app
           exit 1
 
       - name: Prepare artifact directory
@@ -177,6 +178,9 @@ jobs:
           if [ -d docker/data/prometheus ]; then
             tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-${SCENARIO_ID}-${BACKEND_ID}.tar.gz" prometheus
           fi
+          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
+            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-${SCENARIO_ID}-${BACKEND_ID}.log" || true
+          fi
 
       - name: Upload benchmark artifacts
         if: ${{ always() }}
@@ -185,3 +189,146 @@ jobs:
           name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
           path: ${{ env.ARTIFACT_DIR }}
           if-no-files-found: error
+
+  micro-benchmark:
+    name: Micro-benchmark (${{ matrix.backend.id }}, ${{ matrix.mode.display }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        backend:
+          - id: memory
+            compose_file: compose.prometheus-memory-store.yml
+          - id: mongo
+            compose_file: compose.prometheus-mongo-store.yml
+        mode:
+          - id: worker
+            display: Update worker throughput
+            cli: worker
+          - id: dequeue-empty
+            display: Dequeue empty throughput
+            cli: dequeue-empty
+          - id: rollout
+            display: Rollout + span throughput
+            cli: rollout
+    env:
+      STORE_URL: http://localhost:4747
+      STORE_API_URL: http://localhost:4747/v1/agl
+      PROM_URL: http://localhost:9090
+      BACKEND_ID: ${{ matrix.backend.id }}
+      MODE_ID: ${{ matrix.mode.id }}
+      ARTIFACT_DIR: artifacts/micro-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
+      AGL_STORE_N_WORKERS: 8
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - name: Sync dependencies
+        run: uv sync --frozen --extra mongo --group core-stable --group dev
+
+      - name: Reset benchmark data directories
+        run: |
+          set -euo pipefail
+          cd docker
+          rm -rf data
+          bash setup.sh
+
+      - name: Launch ${{ matrix.backend.id }} Prometheus stack
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+          docker compose -f "$COMPOSE_FILE" up -d --quiet-pull
+
+      - name: Wait for store readiness
+        run: |
+          set -euo pipefail
+          for attempt in {1..60}; do
+            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+              exit 0
+            fi
+            sleep 1
+          done
+          echo "Store did not become ready in time" >&2
+          cd docker && docker compose -f "$COMPOSE_FILE" logs app
+          exit 1
+
+      - name: Prepare artifact directory
+        run: mkdir -p "$ARTIFACT_DIR"
+
+      - name: Record micro benchmark start
+        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run ${{ matrix.mode.display }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          uv run --locked --no-sync python -m tests.benchmark.micro_benchmark \
+            --store-url "$STORE_URL" \
+            --summary-file "$ARTIFACT_DIR/summary-${MODE_ID}.txt" \
+            "${{ matrix.mode.cli }}" | tee "$ARTIFACT_DIR/micro-${MODE_ID}.txt"
+
+      - name: Record micro benchmark end
+        if: ${{ always() }}
+        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run micro benchmark analysis
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
+            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+            exit 1
+          fi
+          uv run --locked --no-sync python -m tests.benchmark.analysis \
+            --prom-url "$PROM_URL" \
+            --store-url "$STORE_API_URL" \
+            --start "$BENCHMARK_START" \
+            --end "$BENCHMARK_END" \
+            | tee "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+
+      - name: Show micro benchmark summary
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          summary_file="$ARTIFACT_DIR/summary-${MODE_ID}.txt"
+          if [ -f "$summary_file" ]; then
+            echo "Micro benchmark summary ($MODE_ID/$BACKEND_ID):"
+            cat "$summary_file"
+          else
+            echo "Summary file not found: $summary_file"
+          fi
+
+      - name: Stop ${{ matrix.backend.id }} Prometheus stack
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+
+      - name: Archive Prometheus metrics
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          if [ -d docker/data/prometheus ]; then
+            tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-micro-${MODE_ID}-${BACKEND_ID}.tar.gz" prometheus
+          fi
+          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
+            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-micro-${MODE_ID}-${BACKEND_ID}.log" || true
+          fi
+
+      - name: Upload micro benchmark artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: micro-benchmark-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error

.github/workflows/examples-rag.yml

@@ -0,0 +1,179 @@
+name: Examples - RAG
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: '0 22 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-rag, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'RAG - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('RAG - {0}', github.event_name) }}
+
+jobs:
+  rag:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-rag' ||
+      github.event.action == 'ci-all'
+    name: RAG (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare RAG dataset
+        run: |
+          set -euo pipefail
+          cd examples/rag
+          mkdir -p data
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" -O data/dataset_tiny.parquet
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" -O data/chunks_candidate_tiny.pkl
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" -O data/index_hnsw_faiss_n32e40_tiny.index
+
+      - name: Run WIKI Retriever MCP Server
+        run: |
+          set -euo pipefail
+          cd examples/rag
+          uv run python wiki_retriever_mcp.py &
+          for i in {1..20}; do
+            sleep 5
+            if nc -z localhost 8099; then
+              echo "MCP server is up!"
+              exit 0
+            else
+              echo "Waiting for MCP server to start..."
+            fi
+          done
+          echo "MCP server failed to start within expected time."
+          exit 1
+
+      - name: Run vLLM Server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen2.5-1.5B-Instruct \
+            --enable-auto-tool-choice \
+            --tool-call-parser hermes \
+            --port 8000 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:8000/v1/models > /dev/null 2>&1; then
+              echo "vLLM server is ready!"
+              VLLM_READY=1
+              break
+            fi
+            echo "Waiting for vLLM server to be ready... (${i})"
+            sleep 5
+          done
+          if [[ "$VLLM_READY" != "1" ]]; then
+            echo "vLLM server failed to start!"
+            exit 1
+          fi
+
+      - name: Run RAG Sanity check
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/rag
+          uv run python rag_agent.py
+        shell: bash
+
+      - name: Stop vLLM Server
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+
+      - name: RAG training
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/rag
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_rag.py fast
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: rag_train
+
+      - name: Validate RAG training
+        run: |
+          set -ex
+          # Allow up to 5 rollouts to fail to produce rewards
+          uv run scripts/validate_example_wandb.py ${{ steps.rag_train.outputs.project_name }} ${{ steps.rag_train.outputs.run_name }} --reward-tolerance 5
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

.github/workflows/tests-full.yml

@@ -320,3 +320,24 @@ jobs:
             echo "Waiting for llm_proxy.py to finish..."
             sleep 5
           done
+
+      - name: MultiMetrics backend example
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_metrics.py --duration 8 --prom-port 9105 --prom-host 0.0.0.0 2>&1 | tee metrics.log &
+          pid=$!
+
+          for attempt in $(seq 1 20); do
+            if curl -sSf http://localhost:9105/metrics | grep -q minimal_requests_total; then
+              echo "Metrics endpoint responding"
+              wait $pid
+              cat metrics.log
+              exit 0
+            fi
+            sleep 1
+          done
+
+          echo "Metrics endpoint did not respond"
+          exit 1

agentlightning/cli/store.py

@@ -64,7 +64,9 @@ def main(argv: Iterable[str] | None = None) -> int:
     setup_logging(args.log_level)
 
     if args.backend == "memory":
-        store = InMemoryLightningStore(prometheus=args.prometheus)
+        store = InMemoryLightningStore(
+            prometheus=args.prometheus, thread_safe=True
+        )  # Using thread_safe store for server
     elif args.backend == "mongo":
         from agentlightning.store.mongo import MongoLightningStore
 

agentlightning/emitter/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .annotation import emit_annotation
+from .annotation import emit_annotation, operation
 from .exception import emit_exception
 from .message import emit_message, get_message_value
 from .object import emit_object, get_object_value
@@ -16,6 +16,7 @@ from .reward import (
 
 __all__ = [
     "reward",
+    "operation",
     "emit_reward",
     "get_reward_value",
     "get_rewards_from_span",

agentlightning/emitter/annotation.py

@@ -1,15 +1,36 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Helpers for emitting annotation spans."""
+"""Helpers for emitting annotation/operation spans."""
 
+import asyncio
+import functools
+import inspect
+import json
 import logging
-from typing import Any, Dict
+from types import TracebackType
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
+from opentelemetry import trace
 from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import Status, StatusCode
 
-from agentlightning.semconv import AGL_ANNOTATION
+from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
 from agentlightning.utils.otel import flatten_attributes, get_tracer
 
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])
+
 logger = logging.getLogger(__name__)
 
 
@@ -46,3 +67,298 @@ def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> Reada
         raise ValueError(f"Span is not a ReadableSpan: {span}")
 
     return span
+
+
+def _safe_json_dump(obj: Any) -> str:
+    """Serialize an object to JSON, falling back to ``str(obj)`` if needed.
+
+    Args:
+        obj: Object to be serialized.
+
+    Returns:
+        The JSON-encoded string representation of the object, or its string
+        representation if JSON encoding fails.
+    """
+    try:
+        return json.dumps(obj, default=str, ensure_ascii=False)
+    except Exception:
+        return str(obj)
+
+
+class OperationContext:
+    """Context manager and decorator for tracing operations.
+
+    This class manages an OpenTelemetry span for a logical unit of work. It can
+    be used either:
+
+    * As a decorator, in which case inputs and outputs are inferred
+      automatically from the wrapped function's signature.
+    * As a context manager, in which case inputs and outputs can be recorded
+      explicitly via :meth:`set_input` and :meth:`set_output`.
+
+    Attributes:
+        name: Human-readable span name.
+        initial_attributes: Attributes applied when the span is created.
+        tracer: OpenTelemetry tracer used to create spans.
+        span: The currently active span, if any.
+    """
+
+    def __init__(self, name: str, attributes: Dict[str, Any], *, propagate: bool = True) -> None:
+        """Initialize a new operation context.
+
+        Args:
+            name: Human-readable name of the span.
+            attributes: Initial attributes attached to the span. Values are
+                JSON-serialized where necessary.
+            propagate: Whether the span should be sent to active exporters.
+        """
+        self.name: str = name
+        self.initial_attributes: Dict[str, Any] = attributes
+        self.propagate: bool = propagate
+        self.tracer: trace.Tracer = get_tracer(use_active_span_processor=propagate)
+        self.span: Optional[trace.Span] = None
+        self._ctx_token: Optional[ContextManager[Any]] = None
+
+    def __enter__(self) -> "OperationContext":
+        """Enter the context manager and start a new span.
+
+        Returns:
+            The current :class:`OperationContext` instance with an active span.
+        """
+        # 1. Start the span with initial attributes (JSON serialized)
+        sanitized_attrs = {
+            k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
+            for k, v in self.initial_attributes.items()
+        }
+
+        self.span = self.tracer.start_span(self.name, attributes=sanitized_attrs)
+        self._ctx_token = trace.use_span(self.span, end_on_exit=True)
+        self._ctx_token.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Exit the context manager and finish the span.
+
+        Any exception raised inside the context is recorded on the span and the
+        span status is set to error.
+
+        Args:
+            exc_type: Exception type, if an exception occurred.
+            exc_val: Exception instance, if an exception occurred.
+            exc_tb: Traceback object, if an exception occurred.
+        """
+        # 1. Record Exception if present
+        if exc_val and self.span:
+            self.span.record_exception(exc_val)
+            self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
+
+        # 2. Close span
+        if self._ctx_token:
+            self._ctx_token.__exit__(exc_type, exc_val, exc_tb)
+
+    def set_input(self, *args: Any, **kwargs: Any) -> None:
+        """Record input arguments on the current span.
+
+        Positional arguments are stored under the ``input.args`` attribute,
+        and keyword arguments are stored under ``input.<name>`` attributes.
+
+        This is intended for use inside a ``with operation(...) as op`` block.
+
+        Args:
+            *args: Positional arguments to record.
+            **kwargs: Keyword arguments to record.
+        """
+        if not self.span:
+            return
+
+        if args:
+            self.span.set_attribute("input.args", _safe_json_dump(args))
+        if kwargs:
+            for k, v in kwargs.items():
+                self.span.set_attribute(f"input.{k}", _safe_json_dump(v))
+
+    def set_output(self, output: Any) -> None:
+        """Record the output value on the current span.
+
+        This is intended for use inside a ``with operation(...) as op`` block.
+
+        Args:
+            output: The output value to record.
+        """
+        if not self.span:
+            return
+        self.span.set_attribute("output", _safe_json_dump(output))
+
+    def __call__(self, fn: _FnType) -> _FnType:
+        """Wrap a callable so its execution is traced in a span.
+
+        When used as a decorator, a new span is created for each call to
+        the wrapped function. The bound arguments are recorded as input
+        attributes, the return value is recorded as an output attribute,
+        and any exception is recorded and marks the span as an error.
+
+        Args:
+            fn: The function or coroutine function to wrap.
+
+        Returns:
+            The wrapped callable.
+        """
+        function_name = fn.__name__
+
+        sig = inspect.signature(fn)
+
+        def _record_auto_inputs(span: trace.Span, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None:
+            """Bind arguments to signature and log them on the span.
+
+            Args:
+                span: Span on which to record attributes.
+                args: Positional arguments passed to the wrapped callable.
+                kwargs: Keyword arguments passed to the wrapped callable.
+            """
+            try:
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                for k, v in bound.arguments.items():
+                    span.set_attribute(
+                        f"{LightningSpanAttributes.OPERATION_INPUT.value}.{k}",
+                        _safe_json_dump(v),
+                    )
+            except Exception:
+                span.set_attribute(
+                    f"{LightningSpanAttributes.OPERATION_INPUT.value}.args",
+                    _safe_json_dump(args),
+                )
+                span.set_attribute(
+                    f"{LightningSpanAttributes.OPERATION_INPUT.value}.kwargs",
+                    _safe_json_dump(kwargs),
+                )
+
+        if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Async wrapper that traces the wrapped coroutine."""
+                # Reuse __enter__ logic via 'with self' would share state incorrectly
+                # across concurrent calls. We must create a new span per call.
+                # So we manually reimplement the span logic for the wrapper here.
+
+                sanitized_attrs = {
+                    k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
+                    for k, v in self.initial_attributes.items()
+                }
+
+                with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
+                    span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
+                    _record_auto_inputs(span, args, kwargs)
+                    try:
+                        result = await fn(*args, **kwargs)
+                        span.set_attribute(
+                            LightningSpanAttributes.OPERATION_OUTPUT.value,
+                            _safe_json_dump(result),
+                        )
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+
+            return cast(_FnType, async_wrapper)
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Sync wrapper that traces the wrapped callable."""
+                sanitized_attrs = {
+                    k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
+                    for k, v in self.initial_attributes.items()
+                }
+
+                with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
+                    span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
+                    _record_auto_inputs(span, args, kwargs)
+                    try:
+                        result = fn(*args, **kwargs)
+                        span.set_attribute(
+                            LightningSpanAttributes.OPERATION_OUTPUT.value,
+                            _safe_json_dump(result),
+                        )
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+
+            return cast(_FnType, sync_wrapper)
+
+
+@overload
+def operation(fn: _FnType, *, propagate: bool = True, **additional_attributes: Any) -> _FnType: ...
+
+
+@overload
+def operation(*, propagate: bool = True, **additional_attributes: Any) -> OperationContext: ...
+
+
+def operation(
+    fn: Optional[_FnType] = None,
+    *,
+    propagate: bool = True,
+    **additional_attributes: Any,
+) -> Union[_FnType, OperationContext]:
+    """Entry point for tracking operations.
+
+    This helper can be used either as a decorator or as a context manager.
+    The span name is fixed to [`AGL_OPERATION`][agentlightning.semconv.AGL_OPERATION];
+    custom span names are not supported. Any keyword arguments are recorded as span attributes.
+
+    Usage as a decorator:
+
+    ```python
+    @operation
+    def func(...):
+        ...
+
+    @operation(category="compute")
+    def func(...):
+        ...
+    ```
+
+    Usage as a context manager:
+
+    ```python
+    with operation(user_id=123) as op:
+        op.set_input(data=data)
+        # ... do work ...
+        op.set_output(result)
+    ```
+
+    Args:
+        fn: When used as `@operation`, this is the wrapped function.
+            When used as `operation(**attrs)`, this should be omitted (or
+            left as `None`) and only keyword attributes are provided.
+        propagate: Whether spans should use the active span processor. When False,
+            spans will stay local and not be exported.
+        **additional_attributes: Additional span attributes to attach at
+            creation time.
+
+    Returns:
+        Either a wrapped callable (when used as a decorator) or an
+        [`OperationContext`][agentlightning.emitter.annotation.OperationContext]
+        (when used as a context manager factory).
+    """
+    # Case 1: Used as @operation (bare decorator or with attributes)
+    if callable(fn):
+        # Create context with fixed name, then immediately wrap the function
+        return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn)
+
+    # Case 2: Used as operation(...) / with operation(...)
+    # Custom span names are intentionally not supported; use AGL_OPERATION.
+    if fn is not None:
+        raise ValueError("Custom span names are intentionally not supported when used as a context manager.")
+    return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)

agentlightning/runner/agent.py

@@ -33,8 +33,8 @@ from opentelemetry.sdk.trace import ReadableSpan
 from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.store.base import LightningStore
-from agentlightning.tracer.agentops import AgentOpsTracer
 from agentlightning.tracer.base import Tracer
+from agentlightning.tracer.otel import OtelTracer
 from agentlightning.types import (
     AttemptedRollout,
     Hook,
@@ -73,7 +73,7 @@ class LitAgentRunner(Runner[T_task]):
         max_rollouts: Optional[int] = None,
         poll_interval: float = 5.0,
         heartbeat_interval: float = 10.0,
-        interval_jitter: float = 0.1,
+        interval_jitter: float = 0.5,
         heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
     ) -> None:
         """Initialize the agent runner.
@@ -277,13 +277,21 @@ class LitAgentRunner(Runner[T_task]):
         store = self.get_store()
 
         trace_spans: list[ReadableSpan] | list[Span] = []
+        result_recognized: bool = False
 
         # Case 0: result is None
         if raw_result is None:
             trace_spans = self._tracer.get_last_trace()
+            result_recognized = True
 
         # Case 1: result is a float (final reward)
-        if isinstance(raw_result, float):
+        if isinstance(raw_result, (bool, int, float)):
+            if isinstance(raw_result, (bool, int)):
+                logger.warning(
+                    f"{self._log_prefix(rollout.rollout_id)} Reward is not a number, got: {type(raw_result)}. "
+                    "Auto converting to float."
+                )
+                raw_result = float(raw_result)
             # Preserve the existing spans before another span is emitted
             trace_spans = list(self._tracer.get_last_trace())
             # This will NOT emit another span to the tracer
@@ -291,7 +299,9 @@ class LitAgentRunner(Runner[T_task]):
             # 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)
+            result_recognized = True
 
+        # Case 2-3: result is a list
         if isinstance(raw_result, list):
             # For rollout methods that return a list, we assume that the returned spans
             # are the complete span set from the whole rollout
@@ -299,10 +309,7 @@ class LitAgentRunner(Runner[T_task]):
 
             # Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
             if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
-
-                if not isinstance(
-                    self._tracer, AgentOpsTracer
-                ):  # TODO: this should be replaced with general OpenTelemetry tracer in next version
+                if not isinstance(self._tracer, OtelTracer):
                     for span in raw_result:
                         await store.add_otel_span(
                             rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
@@ -313,6 +320,7 @@ class LitAgentRunner(Runner[T_task]):
                         "The traces should have already been added to the store. "
                         "No need to return anything from rollout."
                     )
+                result_recognized = True
 
             # Case 3: result is a list of Span (agentlightning spans)
             elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
@@ -320,6 +328,7 @@ class LitAgentRunner(Runner[T_task]):
                 for span in raw_result:
                     await store.add_span(cast(Span, span))
                 trace_spans = raw_result
+                result_recognized = True
 
             # Left over cases for list
             elif len(raw_result) == 0:
@@ -328,6 +337,7 @@ class LitAgentRunner(Runner[T_task]):
                     "Please check your rollout implementation."
                 )
                 trace_spans = raw_result
+                result_recognized = True
 
             else:
                 types = [type(t).__name__ for t in raw_result][:10]
@@ -336,6 +346,12 @@ class LitAgentRunner(Runner[T_task]):
                     f"but got: {', '.join(types)}..."
                 )
 
+        if not result_recognized:
+            raise TypeError(
+                f"Invalid raw result type. It's expected to be none, float, or a list of ReadableSpan or Span, "
+                f"but got: {type(raw_result).__name__}..."
+            )
+
         return trace_spans
 
     async def _emit_heartbeat(self, store: LightningStore) -> None:
@@ -577,16 +593,6 @@ class LitAgentRunner(Runner[T_task]):
                 if next_rollout is None:
                     return
 
-                try:
-                    # Claim the rollout but updating the current worker id
-                    await store.update_attempt(
-                        next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
-                    )
-                except Exception:
-                    # This exception could happen if the rollout is dequeued and the other end died for some reason
-                    logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
-                    continue
-
                 # Execute the step
                 await self._step_impl(next_rollout)
 
@@ -640,12 +646,8 @@ class LitAgentRunner(Runner[T_task]):
         else:
             resources_id = None
 
-        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
-        # Register the attempt as running by the current worker
-        await self.get_store().update_attempt(
-            attempted_rollout.rollout_id,
-            attempted_rollout.attempt.attempt_id,
-            worker_id=self.get_worker_id(),
+        attempted_rollout = await self.get_store().start_rollout(
+            input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
         )
         rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
 

agentlightning/semconv.py

@@ -29,6 +29,11 @@ AGL_EXCEPTION = "agentlightning.exception"
 Used by the exception emitter to record exception details.
 """
 
+AGL_OPERATION = "agentlightning.operation"
+"""Agent-lightning's standard span name for functions.
+Wrap function or code-blocks as operations.
+"""
+
 AGL_VIRTUAL = "agentlightning.virtual"
 """Agent-lightning's standard span name for virtual operations.
 
@@ -84,6 +89,15 @@ class LightningSpanAttributes(Enum):
     OBJECT_JSON = "agentlightning.object.json"
     """Attribute name for object serialized value (JSON) in object spans."""
 
+    OPERATION_NAME = "agentlightning.operation.name"
+    """Attribute name for operation name in operation spans, normally the function name."""
+
+    OPERATION_INPUT = "agentlightning.operation.input"
+    """Attribute name for operation input in operation spans."""
+
+    OPERATION_OUTPUT = "agentlightning.operation.output"
+    """Attribute name for operation output in operation spans."""
+
 
 class RewardAttributes(Enum):
     """Multi-dimensional reward attributes will look like:

agentlightning/store/base.py

@@ -10,10 +10,12 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
     RolloutConfig,
+    RolloutMode,
     RolloutStatus,
     Span,
     TaskInput,
@@ -156,10 +158,11 @@ class LightningStore:
     async def start_rollout(
         self,
         input: TaskInput,
-        mode: Literal["train", "val", "test"] | None = None,
+        mode: RolloutMode | None = None,
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: str | None = None,
     ) -> AttemptedRollout:
         """Register a rollout and immediately create its first attempt.
 
@@ -182,6 +185,7 @@ class LightningStore:
             resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
             config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
             metadata: Free-form metadata persisted verbatim with the rollout.
+            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
@@ -227,6 +231,22 @@ class LightningStore:
         """
         raise NotImplementedError()
 
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        """Persist multiple rollouts in `queuing` state.
+
+        The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
+        per request and preserves the input ordering. Subclasses can override to provide
+        more efficient bulk enqueue semantics.
+
+        Args:
+            rollouts: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
+                parameters. Each entry requires `input` and can optionally include other fields.
+
+        Returns:
+            Rollouts enqueued in the same order as `rollouts`.
+        """
+        raise NotImplementedError()
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """Claim the oldest queued rollout and transition it to `preparing`.
 
@@ -243,6 +263,9 @@ class LightningStore:
         * Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
           (e.g., `last_dequeue_time`) when `worker_id` is provided.
 
+        Args:
+            worker_id: Optional worker identifier to associate the claimed attempt with.
+
         Returns:
             The next attempt to execute, or `None` when no eligible rollouts are queued.
 
@@ -251,7 +274,30 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        """Claim up to `limit` queued rollouts without blocking.
+
+        The implementation can repeatedly invokes
+        [`dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] until reaching
+        the requested limit or the queue is empty. Subclasses can override it to fetch
+        multiple rollouts atomically.
+
+        Args:
+            limit: Maximum number of rollouts to claim. Non-positive values return an empty list.
+            worker_id: Optional worker identifier passed through to each dequeue call.
+
+        Returns:
+            Attempted rollouts claimed in FIFO order. May contain fewer than `limit` entries
+            when the queue is exhausted.
+        """
+        raise NotImplementedError()
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
         """Create a manual retry attempt for an existing rollout.
 
         This is typically invoked by runners that wish to retry outside of the
@@ -262,6 +308,7 @@ class LightningStore:
 
         Args:
             rollout_id: Unique identifier of the rollout receiving a new attempt.
+            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The rollout paired with its newly-created attempt.
@@ -719,7 +766,8 @@ class LightningStore:
 
         When `attempt_id` is `"latest"` the update must target the attempt with the highest
         `sequence_id`; otherwise it must target the specific attempt. Implementations should
-        propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
+        propagate status changes to the rollout (for example
+        via [`rollout_status_from_attempt()`][agentlightning.store.utils.rollout_status_from_attempt])
         once the latest attempt transitions to a terminal state.
 
         Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],

agentlightning/store/client_server.py

@@ -47,6 +47,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     PaginatedResult,
     ResourcesUpdate,
@@ -81,12 +82,26 @@ class RolloutRequest(BaseModel):
     resources_id: Optional[str] = None
     config: Optional[RolloutConfig] = None
     metadata: Optional[Dict[str, Any]] = None
+    worker_id: Optional[str] = None
 
 
 class DequeueRolloutRequest(BaseModel):
     worker_id: Optional[str] = None
 
 
+class StartAttemptRequest(BaseModel):
+    worker_id: Optional[str] = None
+
+
+class EnqueueManyRolloutsRequest(BaseModel):
+    rollouts: List[EnqueueRolloutRequest]
+
+
+class DequeueManyRolloutsRequest(BaseModel):
+    limit: int = 1
+    worker_id: Optional[str] = None
+
+
 class QueryRolloutsRequest(BaseModel):
     status_in: Optional[List[RolloutStatus]] = Field(FastAPIQuery(default=None))
     rollout_id_in: Optional[List[str]] = Field(FastAPIQuery(default=None))
@@ -423,6 +438,7 @@ class LightningStoreServer(LightningStore):
         if self._prometheus:
             self._setup_prometheus(api=api, app=self.app)
 
+        # TODO: This should only be enabled in development mode.
         @self.app.middleware("http")
         async def _app_exception_handler(  # pyright: ignore[reportUnusedFunction]
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
@@ -458,7 +474,9 @@ class LightningStoreServer(LightningStore):
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
         ):
             # If not API request, just pass through
-            if not request.url.path.startswith(API_V1_AGL_PREFIX):
+            if not request.url.path.startswith(API_V1_AGL_PREFIX) and not request.url.path.startswith(
+                API_V1_PREFIX + "/traces"
+            ):
                 return await call_next(request)
 
             start = time.perf_counter()
@@ -522,22 +540,38 @@ class LightningStoreServer(LightningStore):
         async def health():  # pyright: ignore[reportUnusedFunction]
             return {"status": "ok"}
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=Rollout)
-        async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
-            return await self.enqueue_rollout(
-                input=request.input,
-                mode=request.mode,
-                resources_id=request.resources_id,
-                config=request.config,
-                metadata=request.metadata,
-            )
+        @api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=List[Rollout])
+        async def enqueue_rollouts(  # pyright: ignore[reportUnusedFunction]
+            request: EnqueueManyRolloutsRequest,
+        ) -> List[Rollout]:
+            enqueue_requests = request.rollouts
+            if not enqueue_requests:
+                return []
+            if len(enqueue_requests) == 1:
+                single = enqueue_requests[0]
+                rollout = await self.enqueue_rollout(
+                    input=single.input,
+                    mode=single.mode,
+                    resources_id=single.resources_id,
+                    config=single.config,
+                    metadata=single.metadata,
+                )
+                return [rollout]
+            rollouts = await self.enqueue_many_rollouts(enqueue_requests)
+            return list(rollouts)
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
-        async def dequeue_rollout(  # pyright: ignore[reportUnusedFunction]
-            request: DequeueRolloutRequest | None = Body(None),
-        ):
-            worker_id = request.worker_id if request else None
-            return await self.dequeue_rollout(worker_id=worker_id)
+        @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=List[AttemptedRollout])
+        async def dequeue_rollouts(  # pyright: ignore[reportUnusedFunction]
+            request: DequeueManyRolloutsRequest | None = Body(None),
+        ) -> List[AttemptedRollout]:
+            payload = request or DequeueManyRolloutsRequest()
+            if payload.limit <= 0:
+                return []
+            if payload.limit == 1:
+                single = await self.dequeue_rollout(worker_id=payload.worker_id)
+                return [single] if single else []
+            rollouts = await self.dequeue_many_rollouts(limit=payload.limit, worker_id=payload.worker_id)
+            return list(rollouts)
 
         @api.post(API_AGL_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
@@ -547,6 +581,7 @@ class LightningStoreServer(LightningStore):
                 resources_id=request.resources_id,
                 config=request.config,
                 metadata=request.metadata,
+                worker_id=request.worker_id,
             )
 
         @api.get(API_AGL_PREFIX + "/rollouts", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
@@ -565,6 +600,24 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(results, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/rollouts/search", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
+        async def search_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Rollout)
+            status_in = request.status_in if "status_in" in request.model_fields_set else None
+            rollout_id_in = request.rollout_id_in if "rollout_id_in" in request.model_fields_set else None
+            # Get all rollouts from the underlying store
+            results = await self.query_rollouts(
+                status_in=status_in,
+                rollout_id_in=rollout_id_in,
+                rollout_id_contains=request.rollout_id_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(results, limit=request.limit, offset=request.offset)
+
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Union[AttemptedRollout, Rollout])
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_rollout_by_id(rollout_id)
@@ -597,8 +650,25 @@ class LightningStoreServer(LightningStore):
             )
 
         @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", status_code=201, response_model=AttemptedRollout)
-        async def start_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
-            return await self.start_attempt(rollout_id)
+        async def start_attempt(  # pyright: ignore[reportUnusedFunction]
+            rollout_id: str, request: StartAttemptRequest | None = Body(None)
+        ):
+            worker_id = request.worker_id if request else None
+            return await self.start_attempt(rollout_id, worker_id=worker_id)
+
+        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/search", response_model=PaginatedResult[Attempt])
+        async def search_attempts(  # pyright: ignore[reportUnusedFunction]
+            rollout_id: str, request: QueryAttemptsRequest
+        ):
+            _validate_paginated_request(request, Attempt)
+            attempts = await self.query_attempts(
+                rollout_id,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(attempts, limit=request.limit, offset=request.offset)
 
         @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
         async def update_attempt(  # pyright: ignore[reportUnusedFunction]
@@ -627,6 +697,21 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(workers, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/workers/search", response_model=PaginatedResult[Worker])
+        async def search_workers(request: QueryWorkersRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Worker)
+            status_in = request.status_in if "status_in" in request.model_fields_set else None
+            workers = await self.query_workers(
+                status_in=status_in,
+                worker_id_contains=request.worker_id_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(workers, limit=request.limit, offset=request.offset)
+
         @api.get(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Optional[Worker])
         async def get_worker(worker_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_worker_by_id(worker_id)
@@ -719,6 +804,28 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(spans, limit=params.limit, offset=params.offset)
 
+        @api.post(API_AGL_PREFIX + "/spans/search", response_model=PaginatedResult[Span])
+        async def search_spans(request: QuerySpansRequest):  # pyright: ignore[reportUnusedFunction]
+            _validate_paginated_request(request, Span)
+            spans = await self.query_spans(
+                request.rollout_id,
+                request.attempt_id,
+                trace_id=request.trace_id,
+                trace_id_contains=request.trace_id_contains,
+                span_id=request.span_id,
+                span_id_contains=request.span_id_contains,
+                parent_id=request.parent_id,
+                parent_id_contains=request.parent_id_contains,
+                name=request.name,
+                name_contains=request.name_contains,
+                filter_logic=request.filter_logic,
+                sort_by=request.sort_by,
+                sort_order=request.sort_order,
+                limit=request.limit,
+                offset=request.offset,
+            )
+            return _build_paginated_response(spans, limit=request.limit, offset=request.offset)
+
         @api.post(API_AGL_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
         async def get_next_span_sequence_id(request: NextSequenceIdRequest):  # pyright: ignore[reportUnusedFunction]
             sequence_id = await self.get_next_span_sequence_id(request.rollout_id, request.attempt_id)
@@ -770,7 +877,7 @@ class LightningStoreServer(LightningStore):
         HTTP_LATENCY = Histogram(
             "http_request_duration_seconds",
             "Latency of HTTP requests",
-            ["method", "path"],
+            ["method", "path", "status_code"],
             buckets=LATENCY_BUCKETS,
         )
 
@@ -778,9 +885,13 @@ class LightningStoreServer(LightningStore):
             # Handle "latest" keywords BEFORE generic IDs
             if path.endswith("/attempts/latest") and "/rollouts/" in path:
                 return re.sub(r"rollouts/[^/]+/attempts/latest$", "rollouts/{rollout_id}/attempts/latest", path)
-            elif path.endswith("/resources/latest"):
+            if path.endswith("/attempts/search") and "/rollouts/" in path:
+                return re.sub(r"rollouts/[^/]+/attempts/search$", "rollouts/{rollout_id}/attempts/search", path)
+            if path.endswith("/resources/latest"):
                 return path
-            elif "enqueue" in path or "dequeue" in path:
+            if path.endswith("/search"):
+                return path
+            if "enqueue" in path or "dequeue" in path:
                 return path
 
             # Handle generic IDs
@@ -797,18 +908,30 @@ class LightningStoreServer(LightningStore):
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
         ) -> Response:
             start = time.perf_counter()
-            response = await call_next(request)
-            elapsed = time.perf_counter() - start
+            status = 520  # Default to 520 if things crash hard
 
-            # Strip the ID-specific URL parts
-            path = get_template_path(request.url.path)
-            method = request.method
-            status = response.status_code
+            try:
+                response = await call_next(request)
+                status = response.status_code
+                return response
+            except asyncio.CancelledError:
+                # Client disconnected (Timeout)
+                status = 499  # Standard Nginx code for "Client Closed Request"
+                raise  # Re-raise to let Uvicorn handle the cleanup
+            except Exception:
+                # TODO: Record the error type
+                status = 500
+                raise
+            finally:
+                # This block executes NO MATTER WHAT happens above
+                elapsed = time.perf_counter() - start
 
-            HTTP_REQUESTS.labels(method, path, status).inc()
-            HTTP_LATENCY.labels(method, path).observe(elapsed)
+                # Strip the ID-specific URL parts
+                path = get_template_path(request.url.path)
+                method = request.method
 
-            return response
+                HTTP_REQUESTS.labels(method, path, status).inc()
+                HTTP_LATENCY.labels(method, path, status).observe(elapsed)
 
         metrics_app = make_asgi_app(registry=registry)  # type: ignore
 
@@ -934,6 +1057,7 @@ class LightningStoreServer(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         return await self._call_store_method(
             "start_rollout",
@@ -942,6 +1066,7 @@ class LightningStoreServer(LightningStore):
             resources_id,
             config,
             metadata,
+            worker_id,
         )
 
     async def enqueue_rollout(
@@ -961,11 +1086,22 @@ class LightningStoreServer(LightningStore):
             metadata,
         )
 
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        return await self._call_store_method("enqueue_many_rollouts", rollouts)
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         return await self._call_store_method("dequeue_rollout", worker_id)
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        return await self._call_store_method("start_attempt", rollout_id)
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        return await self._call_store_method("dequeue_many_rollouts", limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        return await self._call_store_method("start_attempt", rollout_id, worker_id)
 
     async def query_rollouts(
         self,
@@ -1439,6 +1575,7 @@ class LightningStoreClient(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         data = await self._request_json(
             "post",
@@ -1449,6 +1586,7 @@ class LightningStoreClient(LightningStore):
                 resources_id=resources_id,
                 config=config,
                 metadata=metadata,
+                worker_id=worker_id,
             ).model_dump(exclude_none=False),
         )
         return AttemptedRollout.model_validate(data)
@@ -1461,18 +1599,64 @@ class LightningStoreClient(LightningStore):
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> Rollout:
+        request_body = EnqueueManyRolloutsRequest(
+            rollouts=[
+                EnqueueRolloutRequest(
+                    input=input,
+                    mode=mode,
+                    resources_id=resources_id,
+                    config=config,
+                    metadata=metadata,
+                )
+            ]
+        ).model_dump(exclude_none=False)
         data = await self._request_json(
             "post",
             "/queues/rollouts/enqueue",
-            json=RolloutRequest(
-                input=input,
-                mode=mode,
-                resources_id=resources_id,
-                config=config,
-                metadata=metadata,
-            ).model_dump(exclude_none=False),
+            json=request_body,
         )
-        return Rollout.model_validate(data)
+        if not data:
+            raise RuntimeError("enqueue_rollout returned no rollouts")
+        return Rollout.model_validate(data[0])
+
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        if not rollouts:
+            return []
+        request_body = EnqueueManyRolloutsRequest(rollouts=list(rollouts)).model_dump(exclude_none=False)
+        data = await self._request_json(
+            "post",
+            "/queues/rollouts/enqueue",
+            json=request_body,
+        )
+        return [Rollout.model_validate(entry) for entry in data]
+
+    async def _dequeue_batch(
+        self,
+        *,
+        limit: int,
+        worker_id: Optional[str],
+    ) -> List[AttemptedRollout]:
+        if limit <= 0:
+            return []
+        session = await self._get_session()
+        url = f"{self.server_address}/queues/rollouts/dequeue"
+        payload: Dict[str, Any] = {"limit": limit}
+        if worker_id is not None:
+            payload["worker_id"] = worker_id
+        try:
+            async with session.post(url, json=payload) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                self._dequeue_was_successful = True
+                return [AttemptedRollout.model_validate(item) for item in data]
+        except Exception as e:
+            if self._dequeue_was_successful:
+                if self._dequeue_first_unsuccessful:
+                    client_logger.warning(f"dequeue_rollout failed with exception: {e}")
+                    self._dequeue_first_unsuccessful = False
+            client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
+            # Else ignore the exception because the server is not ready yet
+            return []
 
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """
@@ -1485,30 +1669,23 @@ class LightningStoreClient(LightningStore):
             This method does NOT retry on failures. If any exception occurs (network error,
             server error, etc.), it logs the error and returns None immediately.
         """
-        session = await self._get_session()
-        url = f"{self.server_address}/queues/rollouts/dequeue"
-        request_kwargs: Dict[str, Any] = {}
-        if worker_id is not None:
-            request_kwargs["json"] = {"worker_id": worker_id}
-        try:
-            async with session.post(url, **request_kwargs) as resp:
-                resp.raise_for_status()
-                data = await resp.json()
-                self._dequeue_was_successful = True
-                return AttemptedRollout.model_validate(data) if data else None
-        except Exception as e:
-            if self._dequeue_was_successful:
-                if self._dequeue_first_unsuccessful:
-                    client_logger.warning(f"dequeue_rollout failed with exception: {e}")
-                    self._dequeue_first_unsuccessful = False
-            client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
-            # Else ignore the exception because the server is not ready yet
-            return None
+        attempts = await self._dequeue_batch(limit=1, worker_id=worker_id)
+        return attempts[0] if attempts else None
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        return await self._dequeue_batch(limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        payload = {"worker_id": worker_id} if worker_id is not None else None
         data = await self._request_json(
             "post",
             f"/rollouts/{rollout_id}/attempts",
+            json=payload,
         )
         return AttemptedRollout.model_validate(data)
 
@@ -1526,29 +1703,25 @@ class LightningStoreClient(LightningStore):
         status: Optional[Sequence[RolloutStatus]] = None,
         rollout_ids: Optional[Sequence[str]] = None,
     ) -> PaginatedResult[Union[AttemptedRollout, Rollout]]:
-        params_list: List[Tuple[str, Any]] = []
-
-        def _extend(key: str, values: Sequence[Any]) -> None:
-            for value in values:
-                params_list.append((key, value))
-
         resolved_status = status_in if status_in is not None else status
         resolved_rollout_ids = rollout_id_in if rollout_id_in is not None else rollout_ids
 
+        payload: Dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+        }
         if resolved_status is not None:
-            _extend("status_in", resolved_status)
+            payload["status_in"] = resolved_status
         if resolved_rollout_ids is not None:
-            _extend("rollout_id_in", resolved_rollout_ids)
+            payload["rollout_id_in"] = resolved_rollout_ids
         if rollout_id_contains is not None:
-            params_list.append(("rollout_id_contains", rollout_id_contains))
-        params_list.append(("filter_logic", filter_logic))
+            payload["rollout_id_contains"] = rollout_id_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params_list.append(("sort_by", sort_by))
-            params_list.append(("sort_order", sort_order))
-        params_list.append(("limit", limit))
-        params_list.append(("offset", offset))
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
 
-        data = await self._request_json("get", "/rollouts", params=params_list or None)
+        data = await self._request_json("post", "/rollouts/search", json=payload)
         items = [
             (
                 AttemptedRollout.model_validate(item)
@@ -1568,14 +1741,14 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Attempt]:
-        params: List[Tuple[str, Any]] = [
-            ("limit", limit),
-            ("offset", offset),
-        ]
+        payload: Dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+        }
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
-        data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts", params=params)
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
+        data = await self._request_json("post", f"/rollouts/{rollout_id}/attempts/search", json=payload)
         items = [Attempt.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -1795,32 +1968,30 @@ class LightningStoreClient(LightningStore):
         sort_by: Optional[str] = "sequence_id",
         sort_order: Literal["asc", "desc"] = "asc",
     ) -> PaginatedResult[Span]:
-        params: List[Tuple[str, Any]] = [("rollout_id", rollout_id)]
+        payload: Dict[str, Any] = {"rollout_id": rollout_id, "limit": limit, "offset": offset}
         if attempt_id is not None:
-            params.append(("attempt_id", attempt_id))
+            payload["attempt_id"] = attempt_id
         if trace_id is not None:
-            params.append(("trace_id", trace_id))
+            payload["trace_id"] = trace_id
         if trace_id_contains is not None:
-            params.append(("trace_id_contains", trace_id_contains))
+            payload["trace_id_contains"] = trace_id_contains
         if span_id is not None:
-            params.append(("span_id", span_id))
+            payload["span_id"] = span_id
         if span_id_contains is not None:
-            params.append(("span_id_contains", span_id_contains))
+            payload["span_id_contains"] = span_id_contains
         if parent_id is not None:
-            params.append(("parent_id", parent_id))
+            payload["parent_id"] = parent_id
         if parent_id_contains is not None:
-            params.append(("parent_id_contains", parent_id_contains))
+            payload["parent_id_contains"] = parent_id_contains
         if name is not None:
-            params.append(("name", name))
+            payload["name"] = name
         if name_contains is not None:
-            params.append(("name_contains", name_contains))
-        params.append(("filter_logic", filter_logic))
+            payload["name_contains"] = name_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
-        params.append(("limit", limit))
-        params.append(("offset", offset))
-        data = await self._request_json("get", "/spans", params=params)
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
+        data = await self._request_json("post", "/spans/search", json=payload)
         items = [Span.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -1888,21 +2059,17 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Worker]:
-        params: List[Tuple[str, Any]] = [
-            ("limit", limit),
-            ("offset", offset),
-        ]
+        payload: Dict[str, Any] = {}
         if status_in is not None:
-            for value in status_in:
-                params.append(("status_in", value))
+            payload["status_in"] = status_in
         if worker_id_contains is not None:
-            params.append(("worker_id_contains", worker_id_contains))
-        params.append(("filter_logic", filter_logic))
+            payload["worker_id_contains"] = worker_id_contains
+        payload["filter_logic"] = filter_logic
         if sort_by is not None:
-            params.append(("sort_by", sort_by))
-            params.append(("sort_order", sort_order))
+            payload["sort_by"] = sort_by
+            payload["sort_order"] = sort_order
 
-        data = await self._request_json("get", "/workers", params=params)
+        data = await self._request_json("post", "/workers/search", json=payload)
         items = [Worker.model_validate(item) for item in data.get("items", [])]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 

agentlightning/store/collection/__init__.py

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

agentlightning/store/collection/base.py

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

agentlightning/store/collection/memory.py

@@ -4,9 +4,10 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import time
 import weakref
 from collections import deque
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager
 from typing import (
     Any,
     Deque,
@@ -24,6 +25,10 @@ from typing import (
     Union,
 )
 
+import aiologic
+from pydantic import BaseModel
+
+from agentlightning.store.utils import LATENCY_BUCKETS
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -37,6 +42,7 @@ from agentlightning.types import (
 )
 
 from .base import (
+    AtomicMode,
     Collection,
     FilterMap,
     KeyValue,
@@ -282,7 +288,7 @@ class ListBasedCollection(Collection[T]):
         # We should always return inside the loop.
         raise RuntimeError("Unreachable")
 
-    def _mutate_single(self, item: T, mode: MutationMode) -> None:
+    def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
         """Core mutation logic shared by insert, update, upsert, and delete."""
         self._ensure_item_type(item)
         key_values = self._extract_primary_key_values(item)
@@ -299,7 +305,35 @@ class ListBasedCollection(Collection[T]):
             else:  # upsert
                 if not exists:
                     self._size += 1
-                parent[final_key] = item
+                    parent[final_key] = item
+
+                elif update_fields is None:
+                    # update_or_insert: update all fields
+                    parent[final_key] = item
+
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+
+                    # Try to fetch the existing item
+                    existing = parent[final_key]
+                    if not isinstance(existing, self._item_type):
+                        raise ValueError(
+                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
+                        )
+
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+
+            return parent[final_key]
 
         elif mode in ("update", "delete"):
             # For update/delete we must not create missing paths.
@@ -314,7 +348,22 @@ class ListBasedCollection(Collection[T]):
                 raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
 
             if mode == "update":
-                parent[final_key] = item
+                if update_fields is None:
+                    # replace the entire item
+                    parent[final_key] = item
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+                return parent[final_key]
             else:  # delete
                 del parent[final_key]
                 self._size -= 1
@@ -554,19 +603,29 @@ class ListBasedCollection(Collection[T]):
         for item in prepared:
             self._mutate_single(item, mode="insert")
 
-    async def update(self, items: Sequence[T]) -> None:
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Update the given items.
 
         Raises:
             ValueError: If any item with the given primary keys does not exist.
         """
+        updated_items: List[T] = []
         for item in items:
-            self._mutate_single(item, mode="update")
+            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
+            if updated is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            updated_items.append(updated)
+        return updated_items
 
-    async def upsert(self, items: Sequence[T]) -> None:
+    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Upsert the given items (insert if missing, otherwise update)."""
+        upserted_items: List[T] = []
         for item in items:
-            self._mutate_single(item, mode="upsert")
+            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
+            if upserted is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            upserted_items.append(upserted)
+        return upserted_items
 
     async def delete(self, items: Sequence[T]) -> None:
         """Delete the given items.
@@ -662,8 +721,16 @@ class InMemoryLightningCollections(LightningCollections):
     Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
     """
 
-    def __init__(self):
-        self._lock = _LoopAwareAsyncLock()
+    def __init__(self, lock_type: Literal["thread", "asyncio"], prometheus: bool = False):
+        self._lock = {
+            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+        }
         self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
         self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
         self._spans = ListBasedCollection(
@@ -674,6 +741,22 @@ class InMemoryLightningCollections(LightningCollections):
         self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
         self._span_sequence_ids = DictBasedKeyValue[str, int](data={})  # rollout_id -> sequence_id
 
+        self._prometheus = prometheus
+        if self._prometheus:
+            from prometheus_client import Counter, Histogram
+
+            self._rate_metric = Counter(
+                "memory_collection_lock_rate",
+                "Rate of memory collection locks",
+                ["collection"],
+            )
+            self._latency_metric = Histogram(
+                "memory_collection_lock_latency_seconds",
+                "Latency of memory collection locks",
+                ["collection"],
+                buckets=LATENCY_BUCKETS,
+            )
+
     @property
     def rollouts(self) -> ListBasedCollection[Rollout]:
         return self._rollouts
@@ -703,9 +786,36 @@ class InMemoryLightningCollections(LightningCollections):
         return self._span_sequence_ids
 
     @asynccontextmanager
-    async def atomic(self, *args: Any, **kwargs: Any):
-        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
-        async with self._lock:
+    async def atomic(
+        self, *, mode: AtomicMode = "rw", snapshot: bool = False, labels: Optional[Sequence[str]] = None, **kwargs: Any
+    ):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
+
+        Skip the locking if mode is "r" and snapshot is False.
+
+        This collection implementation does NOT support rollback / commit.
+        """
+        if mode == "r" and not snapshot:
+            yield self
+            return
+        if not labels:
+            # If no labels are provided, use all locks.
+            labels = list(self._lock.keys())
+
+        # 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)
+
+        managers = [(label, self._lock[label]) for label in labels]
+        async with AsyncExitStack() as stack:
+            for label, manager in managers:
+                start_time = time.perf_counter()
+                await stack.enter_async_context(manager)
+                elapsed = time.perf_counter() - start_time
+                if self._prometheus:
+                    self._rate_metric.labels(collection=label).inc()
+                    self._latency_metric.labels(collection=label).observe(elapsed)
             yield self
 
     async def evict_spans_for_rollout(self, rollout_id: str) -> None:
@@ -754,3 +864,21 @@ class _LoopAwareAsyncLock:
         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()

agentlightning/store/collection/mongo.py

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

agentlightning/store/memory.py

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

agentlightning/store/mongo.py

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

agentlightning/store/threading.py

@@ -11,6 +11,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -59,9 +60,17 @@ class LightningStoreThreaded(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         with self._lock:
-            return await self.store.start_rollout(input, mode, resources_id, config, metadata)
+            return await self.store.start_rollout(
+                input,
+                mode,
+                resources_id,
+                config,
+                metadata,
+                worker_id,
+            )
 
     async def enqueue_rollout(
         self,
@@ -74,13 +83,26 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
 
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        with self._lock:
+            return await self.store.enqueue_many_rollouts(rollouts)
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         with self._lock:
             return await self.store.dequeue_rollout(worker_id=worker_id)
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
         with self._lock:
-            return await self.store.start_attempt(rollout_id)
+            return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id, worker_id)
 
     async def query_rollouts(
         self,

agentlightning/store/utils.py

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

agentlightning/trainer/trainer.py

@@ -152,6 +152,13 @@ class Trainer(TrainerLegacy):
         # super().__init__() will call TrainerLegacy's initialization, which is not intended.
         self.worker_id: Optional[int] = None
 
+        if dev:
+            logger.warning(
+                "Trainer(dev=True) is deprecated and will be removed in future versions. "
+                "Please use Trainer.dev(...) instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         self._dev = dev
         self.daemon = daemon
         self._client: AgentLightningClient | None = None  # Will be initialized in fit or fit_v0
@@ -213,10 +220,6 @@ class Trainer(TrainerLegacy):
         # We might be able to support a list of resources in future.
         self.initial_resources = initial_resources
 
-        # The active store for the current execution context
-        self.store = self._make_store(store)
-        self.runner = self._make_runner(runner)
-
         self.port = port
 
         self.strategy = self._make_strategy(
@@ -224,6 +227,11 @@ class Trainer(TrainerLegacy):
             n_runners=self.n_runners,
             port=port,
         )
+
+        # The active store for the current execution context
+        self.store = self._make_store(store, self.strategy)
+        self.runner = self._make_runner(runner)
+
         if hasattr(self.strategy, "n_runners"):
             strategy_runners = getattr(self.strategy, "n_runners")
             if isinstance(strategy_runners, int) and strategy_runners > 0:
@@ -282,13 +290,19 @@ class Trainer(TrainerLegacy):
             type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
         )
 
-    def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
-        """Resolve the store implementation backing rollouts, attempts, spans, and resources."""
+    def _make_store(self, store: ComponentSpec[LightningStore], strategy: ExecutionStrategy) -> LightningStore:
+        """Resolve the store implementation backing rollouts, attempts, spans, and resources.
+
+        By default, it's always a in-memory store. If using a client/server execution strategy,
+        the in-memory store will be initialized in a thread-safe manner.
+        """
+        is_client_server = isinstance(strategy, ClientServerExecutionStrategy)
+        default_store_factory = lambda: InMemoryLightningStore(thread_safe=is_client_server)
         return build_component(
             store,
             expected_type=LightningStore,
             spec_name="store",
-            default_factory=InMemoryLightningStore,
+            default_factory=default_store_factory,
             invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
             type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
         )

agentlightning/types/core.py

@@ -53,6 +53,7 @@ __all__ = [
     "Rollout",
     "Attempt",
     "AttemptedRollout",
+    "EnqueueRolloutRequest",
     "Hook",
     "Worker",
     "WorkerStatus",
@@ -211,6 +212,24 @@ class AttemptedRollout(Rollout):
         return self
 
 
+class EnqueueRolloutRequest(BaseModel):
+    """Payload describing a rollout to be queued via [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout].
+
+    A subset of fields from [`Rollout`][agentlightning.Rollout] used for queuing new rollouts.
+    """
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+    config: Optional[RolloutConfig] = None
+    """Retry and timeout configuration associated with the rollout."""
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
 WorkerStatus = Literal["idle", "busy", "unknown"]
 
 

agentlightning/utils/metrics.py

@@ -0,0 +1,873 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Metrics abstraction with explicit registration and several backends.
+
+It provides:
+
+- MetricsBackend: Abstract interface for registering and recording metrics.
+- ConsoleMetricsBackend: In-process backend with sliding-window
+  aggregations (rate, P50, P95, P99) logged to stdout.
+- PrometheusMetricsBackend: Thin wrapper around prometheus_client.
+- MultiMetricsBackend: Fan-out backend that forwards calls to multiple underlying backends.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import tempfile
+import threading
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple
+
+if TYPE_CHECKING:
+    from prometheus_client import CollectorRegistry
+
+LabelDict = Dict[str, str]
+LabelKey = Tuple[Tuple[str, str], ...]  # normalized, sorted (key, value) pairs
+
+logger = logging.getLogger(__name__)
+
+
+def _validate_labels(
+    kind: str,
+    metric_name: str,
+    labels: LabelDict,
+    expected_names: Tuple[str, ...],
+) -> LabelKey:
+    """Validates label keys against the metric definition.
+
+    Args:
+        kind: Metric kind for error messages ("counter" or "histogram").
+        metric_name: Metric name.
+        labels: Provided label dictionary.
+        expected_names: Expected label names as a tuple.
+
+    Returns:
+        A tuple of (key, value) pairs sorted by registered label order.
+
+    Raises:
+        ValueError: If label keys do not match expected_names.
+    """
+
+    label_items: List[Tuple[str, str]] = []
+    for label_name in expected_names:
+        if label_name not in labels:
+            raise ValueError(f"Label '{label_name}' is required for {kind.capitalize()} '{metric_name}'.")
+        label_items.append((label_name, labels[label_name]))
+
+    return tuple(label_items)
+
+
+def _normalize_label_names(label_names: Optional[Sequence[str]]) -> Tuple[str, ...]:
+    """Normalizes label names into a canonical tuple.
+
+    Args:
+        label_names: Iterable of label names or None.
+
+    Returns:
+        A tuple of label names sorted alphabetically.
+    """
+    if not label_names:
+        return ()
+    return tuple(sorted(label_names))
+
+
+@dataclass(frozen=True)
+class _CounterDef:
+    """Definition of a registered counter metric."""
+
+    name: str
+    label_names: Tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class _HistogramDef:
+    """Definition of a registered histogram metric."""
+
+    name: str
+    label_names: Tuple[str, ...]
+    buckets: Tuple[float, ...]
+
+
+@dataclass
+class _CounterState:
+    """Runtime state of a counter metric group (for console backend)."""
+
+    timestamps: List[float]
+    amounts: List[float]
+
+
+@dataclass
+class _HistogramState:
+    """Runtime state of a histogram metric group (for console backend)."""
+
+    timestamps: List[float]
+    values: List[float]
+
+
+class MetricsBackend:
+    """Abstract base class for metrics backends."""
+
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Registers a counter metric.
+
+        Args:
+            name: Metric name.
+            label_names: List of label names. Order is not important.
+
+        Raises:
+            ValueError: If the metric is already registered with a different
+                type or label set.
+        """
+        raise NotImplementedError()
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+    ) -> None:
+        """Registers a histogram metric.
+
+        Args:
+            name: Metric name.
+            label_names: List of label names. Order is not important.
+            buckets: Bucket boundaries (exclusive upper bounds). If None, the
+                backend may choose defaults.
+
+        Raises:
+            ValueError: If the metric is already registered with a different
+                type or label set.
+        """
+        raise NotImplementedError()
+
+    def inc_counter(
+        self,
+        name: str,
+        amount: float = 1.0,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Increments a registered counter.
+
+        Args:
+            name: Metric name (must be registered as a counter).
+            amount: Increment amount.
+            labels: Label values.
+
+        Raises:
+            ValueError: If the metric is not registered, has the wrong type,
+                or label keys do not match the registered label names.
+        """
+        raise NotImplementedError()
+
+    def observe_histogram(
+        self,
+        name: str,
+        value: float,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Records an observation for a registered histogram.
+
+        Args:
+            name: Metric name (must be registered as a histogram).
+            value: Observed value.
+            labels: Label values.
+
+        Raises:
+            ValueError: If the metric is not registered, has the wrong type,
+                or label keys do not match the registered label names.
+        """
+        raise NotImplementedError()
+
+
+class ConsoleMetricsBackend(MetricsBackend):
+    """Console backend with sliding-window aggregations and label grouping.
+
+    This backend:
+
+    * Requires explicit metric registration.
+    * Stores timestamped events per (metric_name, labels) key.
+    * Computes rate and percentiles (P50, P95, P99) over a sliding time window.
+    * Uses a single global logging decision: when logging is triggered, it
+      logs all metric groups, not just the one being updated.
+
+    Rate is always per second.
+
+    Label grouping: When logging, labels are truncated to the first `group_level` label
+    pairs (according to sorted label key order). For example:
+
+        labels = {"method": "GET", "path": "/", "status": "200"}
+        group_level = 2 -> logged labels {"method": "GET", "path": "/"}
+
+    If `group_level` is None or < 1, all labels are logged.
+
+    Thread-safety: A single lock protects shared state mutation, pruning, and snapshotting.
+    Percentile computation, formatting, and printing are done after releasing the lock.
+    """
+
+    def __init__(
+        self,
+        window_seconds: Optional[float] = 60.0,
+        log_interval_seconds: float = 5.0,
+        group_level: Optional[int] = None,
+    ) -> None:
+        """Initializes ConsoleMetricsBackend.
+
+        Args:
+            window_seconds: Sliding window size (in seconds) used when computing
+                rate and percentiles. If None, all in-memory events are used.
+            log_interval_seconds: Minimum time (in seconds) between log bursts.
+                When the interval elapses, the next metric event triggers a
+                snapshot and logging of all metrics.
+            group_level: Label grouping depth. When logging, only the first
+                `group_level` labels (sorted by key) are included. If None or
+                < 1, all labels are included.
+        """
+        self.window_seconds = window_seconds
+        self.log_interval_seconds = log_interval_seconds
+        self.group_level = group_level
+
+        self._counters: Dict[str, _CounterDef] = {}
+        self._histograms: Dict[str, _HistogramDef] = {}
+
+        # Runtime state keyed by (metric_name, label_key)
+        self._counter_state: Dict[Tuple[str, LabelKey], _CounterState] = {}
+        self._hist_state: Dict[Tuple[str, LabelKey], _HistogramState] = {}
+
+        # Global last log time (for all metrics)
+        self._last_log_time: Optional[float] = None
+
+        self._lock = threading.Lock()
+
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Registers a counter metric.
+
+        See base class for argument documentation.
+        """
+        label_tuple = _normalize_label_names(label_names)
+        with self._lock:
+            existing_counter = self._counters.get(name)
+            existing_hist = self._histograms.get(name)
+
+            if existing_hist is not None:
+                raise ValueError(f"Metric '{name}' already registered as histogram.")
+
+            if existing_counter is not None:
+                if existing_counter.label_names != label_tuple:
+                    raise ValueError(
+                        f"Counter '{name}' already registered with labels "
+                        f"{existing_counter.label_names}, got {label_tuple}."
+                    )
+                return
+
+            self._counters[name] = _CounterDef(name=name, label_names=label_tuple)
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+    ) -> None:
+        """Registers a histogram metric.
+
+        See base class for argument documentation.
+        """
+        label_tuple = _normalize_label_names(label_names)
+        if buckets is None:
+            bucket_tuple: Tuple[float, ...] = (0.1, 0.2, 0.5, 1.0, 2.0)
+        else:
+            bucket_tuple = tuple(buckets)
+
+        with self._lock:
+            existing_counter = self._counters.get(name)
+            existing_hist = self._histograms.get(name)
+
+            if existing_counter is not None:
+                raise ValueError(f"Metric '{name}' already registered as counter.")
+
+            if existing_hist is not None:
+                if existing_hist.label_names != label_tuple or existing_hist.buckets != bucket_tuple:
+                    raise ValueError(
+                        f"Histogram '{name}' already registered with "
+                        f"labels={existing_hist.label_names}, "
+                        f"buckets={existing_hist.buckets}."
+                    )
+                return
+
+            self._histograms[name] = _HistogramDef(
+                name=name,
+                label_names=label_tuple,
+                buckets=bucket_tuple,
+            )
+
+    def inc_counter(
+        self,
+        name: str,
+        amount: float = 1.0,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Increments a registered counter metric.
+
+        See base class for behavior and error conditions.
+        """
+        now = time.time()
+        labels = labels or {}
+
+        definition = self._counters.get(name)
+        if definition is None:
+            raise ValueError(f"Counter '{name}' is not registered.")
+
+        label_key = _validate_labels("counter", name, labels, definition.label_names)
+        state_key = (name, label_key)
+
+        with self._lock:
+            state = self._counter_state.get(state_key)
+            if state is None:
+                state = _CounterState(timestamps=[], amounts=[])
+                self._counter_state[state_key] = state
+
+            state.timestamps.append(now)
+            state.amounts.append(amount)
+            self._prune_events(state.timestamps, state.amounts, now)
+
+            should_log = self._should_log_locked(now)
+            if should_log:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
+                snapshot_time = now
+            else:
+                counter_snaps = hist_snaps = []
+                snapshot_time = now
+
+        if should_log and (counter_snaps or hist_snaps):
+            self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)
+
+    def observe_histogram(
+        self,
+        name: str,
+        value: float,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Records an observation for a registered histogram metric.
+
+        See base class for behavior and error conditions.
+        """
+        now = time.time()
+        labels = labels or {}
+
+        definition = self._histograms.get(name)
+        if definition is None:
+            raise ValueError(f"Histogram '{name}' is not registered.")
+
+        label_key = _validate_labels("histogram", name, labels, definition.label_names)
+        state_key = (name, label_key)
+
+        with self._lock:
+            state = self._hist_state.get(state_key)
+            if state is None:
+                state = _HistogramState(timestamps=[], values=[])
+                self._hist_state[state_key] = state
+
+            state.timestamps.append(now)
+            state.values.append(value)
+            self._prune_events(state.timestamps, state.values, now)
+
+            should_log = self._should_log_locked(now)
+            if should_log:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
+                snapshot_time = now
+            else:
+                counter_snaps = hist_snaps = []
+                snapshot_time = now
+
+        if should_log and (counter_snaps or hist_snaps):
+            self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)
+
+    def _prune_events(
+        self,
+        timestamps: List[float],
+        values: List[float],
+        now: float,
+    ) -> None:
+        """Prunes events older than the sliding window.
+
+        Args:
+            timestamps: List of event timestamps (ascending).
+            values: List of corresponding values or amounts.
+            now: Current time.
+        """
+        if self.window_seconds is None or not timestamps:
+            return
+        cutoff = now - self.window_seconds
+        idx = 0
+        for i, ts in enumerate(timestamps):
+            if ts >= cutoff:
+                idx = i
+                break
+        else:
+            idx = len(timestamps)
+        if idx > 0:
+            del timestamps[:idx]
+            del values[:idx]
+
+    def _should_log_locked(self, now: float) -> bool:
+        """Determines whether to emit a log snapshot (lock must be held).
+
+        This decision is global: if it returns True, all metrics will be
+        logged based on a snapshot taken at this time.
+
+        Args:
+            now: Current timestamp.
+
+        Returns:
+            True if enough time has elapsed since the last log; False otherwise.
+        """
+        last = self._last_log_time
+        if last is None or now - last >= self.log_interval_seconds:
+            self._last_log_time = now
+            return True
+        return False
+
+    def _snapshot_locked(
+        self,
+        now: float,
+    ) -> Tuple[
+        List[Tuple[str, LabelDict, List[float], List[float]]],
+        List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]],
+    ]:
+        """Creates a snapshot of all metric state (lock must be held).
+
+        Args:
+            now: Current timestamp.
+
+        Returns:
+            A tuple (counter_snapshots, histogram_snapshots) where:
+              - counter_snapshots: list of (metric_name, labels, timestamps, amounts)
+              - histogram_snapshots: list of (metric_name, labels, values, buckets)
+        """
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+
+        # Prune and snapshot counters.
+        for (name, label_key), state in self._counter_state.items():
+            self._prune_events(state.timestamps, state.amounts, now)
+            if not state.timestamps:
+                continue
+            labels = dict(label_key)
+            counter_snaps.append(
+                (
+                    name,
+                    labels,
+                    list(state.timestamps),
+                    list(state.amounts),
+                )
+            )
+
+        # Prune and snapshot histograms.
+        for (name, label_key), state in self._hist_state.items():
+            self._prune_events(state.timestamps, state.values, now)
+            if not state.values:
+                continue
+            labels = dict(label_key)
+            buckets = self._histograms[name].buckets
+            hist_snaps.append(
+                (
+                    name,
+                    labels,
+                    list(state.values),
+                    buckets,
+                )
+            )
+
+        return counter_snaps, hist_snaps
+
+    def _truncate_labels_for_logging(self, labels: LabelDict) -> LabelDict:
+        """Returns a label dict truncated to the configured group depth.
+
+        Args:
+            labels: Original label dictionary.
+
+        Returns:
+            A new dictionary containing at most `group_level` label pairs,
+            chosen by sorted key order. If group_level is None or < 1, returns
+            a shallow copy of the original labels.
+        """
+        if self.group_level is None or self.group_level < 1:
+            return dict(labels)
+        items = sorted(labels.items())
+        return dict(items[: self.group_level])
+
+    def _log(self, message: str) -> None:
+        """Logs a message via the module logger."""
+        logger.info(message)
+
+    def _log_snapshot(
+        self,
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]],
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]],
+        snapshot_time: float,
+    ) -> None:
+        """Logs all metrics from a snapshot.
+
+        Args:
+            counter_snaps: Counter snapshot list.
+            hist_snaps: Histogram snapshot list.
+        """
+        entries: List[str] = []
+        for name, labels, timestamps, amounts in counter_snaps:
+            truncated_labels = self._truncate_labels_for_logging(labels)
+            line = self._log_counter(name, truncated_labels, timestamps, amounts, snapshot_time)
+            if line:
+                entries.append(line)
+
+        for name, labels, values, buckets in hist_snaps:
+            truncated_labels = self._truncate_labels_for_logging(labels)
+            line = self._log_histogram(name, truncated_labels, values, buckets, snapshot_time)
+            if line:
+                entries.append(line)
+
+        if entries:
+            self._log("  ".join(entries))
+
+    def _log_counter(
+        self,
+        name: str,
+        labels: LabelDict,
+        timestamps: List[float],
+        amounts: List[float],
+        snapshot_time: float,
+    ) -> Optional[str]:
+        """Computes counter stats and returns formatted line."""
+        if not timestamps:
+            return None
+
+        total = sum(amounts)
+        window_start = timestamps[0]
+        if self.window_seconds is not None:
+            window_start = max(window_start, snapshot_time - self.window_seconds)
+        min_duration = self.log_interval_seconds if self.log_interval_seconds > 0 else 1e-3
+        duration = max(snapshot_time - window_start, min_duration)
+        rate = total / duration
+
+        label_str = _format_label_string(labels)
+        return f"{name}{label_str}={rate:.2f}/s"
+
+    def _log_histogram(
+        self,
+        name: str,
+        labels: LabelDict,
+        values: List[float],
+        buckets: Tuple[float, ...],
+        snapshot_time: float,
+    ) -> Optional[str]:
+        """Computes histogram stats and returns formatted line."""
+        if not values:
+            return None
+
+        sorted_vals = sorted(values)
+        n = len(sorted_vals)
+
+        def percentile(p: float) -> float:
+            if n == 1:
+                return sorted_vals[0]
+            pos = (p / 100.0) * (n - 1)
+            lo = int(pos)
+            hi = min(lo + 1, n - 1)
+            if lo == hi:
+                return sorted_vals[lo]
+            w = pos - lo
+            return sorted_vals[lo] * (1 - w) + sorted_vals[hi] * w
+
+        p50 = percentile(50.0)
+        p95 = percentile(95.0)
+        p99 = percentile(99.0)
+
+        label_str = _format_label_string(labels)
+        formatted = ",".join([_format_duration(p50), _format_duration(p95), _format_duration(p99)])
+        return f"{name}{label_str}={formatted}"
+
+
+def _format_label_string(labels: LabelDict) -> str:
+    if not labels:
+        return "{}"
+    ordered = ",".join(f"{key}={value}" for key, value in sorted(labels.items()))
+    return f"{{{ordered}}}"
+
+
+def _format_duration(value: float) -> str:
+    abs_value = abs(value)
+    if abs_value >= 1.0:
+        return f"{value:.2f}s"
+    if abs_value >= 1e-3:
+        return f"{value * 1_000:.2f}ms"
+    if abs_value >= 1e-6:
+        return f"{value * 1_000_000:.2f}µs"
+    return f"{value * 1_000_000_000:.2f}ns"
+
+
+class PrometheusMetricsBackend(MetricsBackend):
+    """Metrics backend that forwards events to prometheus_client.
+
+    All metrics must be registered before use. This backend does not compute
+    any aggregations; it only updates Prometheus metrics.
+
+    Thread-safety: Registration is protected by a lock. Metric updates assume metrics
+    are registered during initialization and then remain stable.
+    """
+
+    def __init__(self) -> None:
+        """Initializes PrometheusMetricsBackend.
+
+        Raises:
+            ImportError: If prometheus_client is not installed.
+        """
+        try:
+            import prometheus_client  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "prometheus_client is not installed. Please either install it or use ConsoleMetricsBackend instead."
+            )
+
+        self._counters: Dict[str, _CounterDef] = {}
+        self._histograms: Dict[str, _HistogramDef] = {}
+        self._prom_counters: Dict[str, Any] = {}
+        self._prom_histograms: Dict[str, Any] = {}
+
+        self._lock = threading.Lock()
+
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Registers a Prometheus counter metric."""
+        from prometheus_client import Counter as PromCounter
+
+        label_tuple = _normalize_label_names(label_names)
+
+        with self._lock:
+            if name in self._histograms:
+                raise ValueError(f"Metric '{name}' already registered as histogram.")
+
+            existing = self._counters.get(name)
+            if existing is not None:
+                if existing.label_names != label_tuple:
+                    raise ValueError(
+                        f"Counter '{name}' already registered with labels "
+                        f"{existing.label_names}, got {label_tuple}."
+                    )
+                return
+
+            self._counters[name] = _CounterDef(name=name, label_names=label_tuple)
+
+            prom_counter = PromCounter(
+                name,
+                f"Counter {name}",
+                labelnames=label_tuple,
+            )
+            self._prom_counters[name] = prom_counter
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+    ) -> None:
+        """Registers a Prometheus histogram metric."""
+        from prometheus_client import Histogram as PromHistogram
+
+        label_tuple = _normalize_label_names(label_names)
+        bucket_tuple = tuple(buckets) if buckets is not None else ()
+
+        with self._lock:
+            if name in self._counters:
+                raise ValueError(f"Metric '{name}' already registered as counter.")
+
+            existing = self._histograms.get(name)
+            if existing is not None:
+                if existing.label_names != label_tuple or existing.buckets != bucket_tuple:
+                    raise ValueError(
+                        f"Histogram '{name}' already registered with "
+                        f"labels={existing.label_names}, "
+                        f"buckets={existing.buckets}."
+                    )
+                return
+
+            self._histograms[name] = _HistogramDef(
+                name=name,
+                label_names=label_tuple,
+                buckets=bucket_tuple,
+            )
+
+            if bucket_tuple:
+                prom_hist = PromHistogram(
+                    name,
+                    f"Histogram {name}",
+                    labelnames=label_tuple,
+                    buckets=bucket_tuple,
+                )
+            else:
+                prom_hist = PromHistogram(
+                    name,
+                    f"Histogram {name}",
+                    labelnames=label_tuple,
+                )
+
+            self._prom_histograms[name] = prom_hist
+
+    def inc_counter(
+        self,
+        name: str,
+        amount: float = 1.0,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Increments a registered Prometheus counter."""
+        labels = labels or {}
+        definition = self._counters.get(name)
+        if definition is None:
+            raise ValueError(f"Counter '{name}' is not registered.")
+
+        prom_counter = self._prom_counters[name]
+        if definition.label_names:
+            label_key = _validate_labels("counter", name, labels, definition.label_names)
+            prom_counter.labels(**dict(label_key)).inc(amount)
+        else:
+            prom_counter.inc(amount)
+
+    def observe_histogram(
+        self,
+        name: str,
+        value: float,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Records an observation for a registered Prometheus histogram."""
+        labels = labels or {}
+        definition = self._histograms.get(name)
+        if definition is None:
+            raise ValueError(f"Histogram '{name}' is not registered.")
+
+        prom_hist = self._prom_histograms[name]
+        if definition.label_names:
+            label_key = _validate_labels("histogram", name, labels, definition.label_names)
+            prom_hist.labels(**dict(label_key)).observe(value)
+        else:
+            prom_hist.observe(value)
+
+
+class MultiMetricsBackend(MetricsBackend):
+    """Metrics backend that forwards calls to multiple underlying backends."""
+
+    def __init__(self, backends: Sequence[MetricsBackend]) -> None:
+        """Initializes MultiMetricsBackend.
+
+        Args:
+            backends: Sequence of underlying backends.
+
+        Raises:
+            ValueError: If no backends are provided.
+        """
+        if not backends:
+            raise ValueError("MultiMetricsBackend requires at least one backend.")
+        self._backends = list(backends)
+
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Registers a counter metric in all underlying backends."""
+        for backend in self._backends:
+            backend.register_counter(name, label_names=label_names)
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+    ) -> None:
+        """Registers a histogram metric in all underlying backends."""
+        for backend in self._backends:
+            backend.register_histogram(
+                name,
+                label_names=label_names,
+                buckets=buckets,
+            )
+
+    def inc_counter(
+        self,
+        name: str,
+        amount: float = 1.0,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Increments a counter metric in all underlying backends."""
+        for backend in self._backends:
+            backend.inc_counter(name, amount=amount, labels=labels)
+
+    def observe_histogram(
+        self,
+        name: str,
+        value: float,
+        labels: Optional[LabelDict] = None,
+    ) -> None:
+        """Records a histogram observation in all underlying backends."""
+        for backend in self._backends:
+            backend.observe_histogram(name, value=value, labels=labels)
+
+
+_prometheus_multiproc_dir: tempfile.TemporaryDirectory[str] | None = None
+
+
+def setup_multiprocess_prometheus():
+    """Set up prometheus multiprocessing directory if not already configured."""
+
+    global _prometheus_multiproc_dir
+
+    if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
+        # Make TemporaryDirectory for prometheus multiprocessing
+        # Note: global TemporaryDirectory will be automatically
+        # cleaned up upon exit.
+        _prometheus_multiproc_dir = tempfile.TemporaryDirectory()
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = _prometheus_multiproc_dir.name
+        logger.debug("Created PROMETHEUS_MULTIPROC_DIR at %s", _prometheus_multiproc_dir.name)
+    else:
+        logger.warning(
+            "Found PROMETHEUS_MULTIPROC_DIR was set by user. " "This directory must be wiped between multiple runs."
+        )
+
+
+def get_prometheus_registry() -> CollectorRegistry:
+    """Get the appropriate prometheus registry based on multiprocessing configuration."""
+    from prometheus_client import REGISTRY, CollectorRegistry, multiprocess
+
+    if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None:
+        logger.debug("Using multiprocess registry for prometheus metrics")
+        registry = CollectorRegistry()
+        multiprocess.MultiProcessCollector(registry)
+        return registry
+
+    return REGISTRY
+
+
+def shutdown_metrics():
+    """Shutdown prometheus metrics."""
+
+    from prometheus_client import multiprocess
+
+    path = _prometheus_multiproc_dir
+    if path is None:
+        return
+    try:
+        pid = os.getpid()
+        multiprocess.mark_process_dead(pid, path.name)  # type: ignore
+        logger.debug("Marked Prometheus metrics for process %d as dead", pid)
+    except Exception as e:
+        logger.error("Error during metrics cleanup: %s", str(e))

agentlightning/verl/daemon.py

@@ -9,7 +9,7 @@ import time
 import uuid
 from collections import defaultdict
 from collections.abc import Mapping
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple, cast
 
 import numpy as np
 import requests
@@ -22,7 +22,7 @@ from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLeg
 from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store.base import LightningStore
-from agentlightning.types import Rollout, RolloutConfig, Task
+from agentlightning.types import EnqueueRolloutRequest, Rollout, RolloutConfig, Task
 
 __all__ = [
     "AgentModeDaemon",
@@ -377,42 +377,57 @@ class AgentModeDaemon:
         num_samples = len(data[keys[0]])
         rollouts_per_sample = self.train_rollout_n if is_train else 1
 
+        enqueue_rollout_requests: List[EnqueueRolloutRequest] = []
+        data_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
+
         for i in range(num_samples):
             data_id = str(uuid.uuid4())
             original_sample = {key: data[key][i] for key in keys}
             original_sample["data_id"] = data_id
+            data_id_to_original_sample[data_id] = original_sample
 
             # For training, each sample is rolled out multiple times
+            # Data ID is different from Rollout ID, as one data can have multiple rollouts.
             for _ in range(rollouts_per_sample):
                 task_metadata = {"data_id": data_id, "is_train": is_train}
-
-                # Data ID is different from Rollout ID, as one data can have multiple rollouts.
                 if self.mode == "v0":
+                    # Queue immediately
                     rollout_id = await self.server.queue_task(
                         sample=_to_native(original_sample),
                         mode="train" if is_train else "val",
                         resources_id=resources_id,
                         metadata=task_metadata,
                     )
-                else:
-                    rollout = await self.store.enqueue_rollout(
-                        input=_to_native(original_sample),
-                        mode="train" if is_train else "val",
-                        resources_id=resources_id,
-                        metadata=task_metadata,
-                    )
-                    await self.store.update_rollout(
-                        rollout_id=rollout.rollout_id,
-                        config=RolloutConfig(
-                            unresponsive_seconds=self.llm_timeout_seconds,
-                            timeout_seconds=self.llm_timeout_seconds,
-                        ),
-                    )
-                    rollout_id = rollout.rollout_id
 
-                # Store original sample data to reconstruct batch information later
-                self._task_id_to_original_sample[rollout_id] = original_sample
-                self._total_tasks_queued += 1
+                    # Store original sample data to reconstruct batch information later
+                    self._task_id_to_original_sample[rollout_id] = original_sample
+                    self._total_tasks_queued += 1
+                else:
+                    # Collect tasks to enqueue in batch and queue them later
+                    enqueue_rollout_requests.append(
+                        EnqueueRolloutRequest(
+                            input=_to_native(original_sample),
+                            mode="train" if is_train else "val",
+                            resources_id=resources_id,
+                            config=RolloutConfig(
+                                unresponsive_seconds=self.llm_timeout_seconds,
+                                timeout_seconds=self.llm_timeout_seconds,
+                            ),
+                            metadata=task_metadata,
+                        )
+                    )
+
+        if self.mode == "v1":
+            # Enqueue all the tasks in a single batch
+            rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
+            self._task_id_to_original_sample.update(
+                {
+                    # Recover the original data and store it for later use.
+                    rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
+                    for rollout in rollouts
+                }
+            )
+            self._total_tasks_queued += len(rollouts)
 
     def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Synchronous wrapper for setting up data and server resources."""

docker/compose.grafana.yml

@@ -0,0 +1,29 @@
+services:
+  prometheus:
+    image: prom/prometheus:latest
+    command:
+      - "--config.file=/etc/prometheus/prometheus.yml"
+      - "--storage.tsdb.path=/prometheus"
+    volumes:
+      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ${AGL_MONITORING_DATA_PATH:?Set AGL_MONITORING_DATA_PATH to the metrics directory}/prometheus:/prometheus
+    ports:
+      - "9090:9090"
+
+  grafana:
+    image: grafana/grafana:latest
+    depends_on:
+      - prometheus
+    ports:
+      - "9091:3000"
+    volumes:
+      - ./data/grafana:/var/lib/grafana
+      - ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
+      - ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
+      - ./grafana/dashboards:/var/lib/grafana/dashboards
+    environment:
+      - GF_INSTALL_PLUGINS=grafana-piechart-panel
+      - GF_AUTH_ANONYMOUS_ENABLED=true
+      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
+      - GF_AUTH_DISABLE_LOGIN_FORM=true
+      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json

docker/grafana/dashboards/agentlightning.json

@@ -377,7 +377,7 @@
       }
     },
     {
-      "title": "add_span: Mongo Ops per Call (by Operation)",
+      "title": "add_many_spans: Mongo Ops per Call (by Operation)",
       "type": "timeseries",
       "gridPos": {
         "h": 8,
@@ -387,7 +387,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (operation, collection) (rate(mongo_operation_total{operation!='ensure_collection', store_method='add_span'}[1m])) / scalar(clamp_min(sum(rate(collection_store_total{method='add_span'}[1m])), 1e-9))",
+          "expr": "sum by (operation, collection) (rate(mongo_operation_total{operation!='ensure_collection', store_method='add_many_spans'}[1m])) / scalar(clamp_min(sum(rate(collection_store_total{method='add_many_spans'}[1m])), 1e-9))",
           "legendFormat": "{{operation}} - {{collection}}"
         }
       ],
@@ -398,7 +398,7 @@
       }
     },
     {
-      "title": "add_span: Mongo Latency (P50)",
+      "title": "add_many_spans: Mongo Latency (P50)",
       "type": "timeseries",
       "gridPos": {
         "h": 8,
@@ -408,7 +408,7 @@
       },
       "targets": [
         {
-          "expr": "histogram_quantile(0.50, sum by (le, operation, collection) (rate(mongo_operation_duration_seconds_bucket{operation!='ensure_collection', store_method='add_span'}[1m])))",
+          "expr": "histogram_quantile(0.50, sum by (le, operation, collection) (rate(mongo_operation_duration_seconds_bucket{operation!='ensure_collection', store_method='add_many_spans'}[1m])))",
           "legendFormat": "{{operation}} - {{collection}}"
         }
       ],
@@ -419,7 +419,7 @@
       }
     },
     {
-      "title": "add_span: Mongo Time per Call (P50)",
+      "title": "add_many_spans: Mongo Time per Call (P50)",
       "type": "timeseries",
       "gridPos": {
         "h": 8,
@@ -429,7 +429,7 @@
       },
       "targets": [
         {
-          "expr": "(histogram_quantile(0.50, sum by (le, operation, collection) (rate(mongo_operation_duration_seconds_bucket{operation!='ensure_collection', store_method='add_span'}[1m])))) * (sum by (operation, collection) (rate(mongo_operation_total{operation!='ensure_collection', store_method='add_span'}[1m])) / scalar(clamp_min(sum(rate(collection_store_total{method='add_span'}[1m])), 1e-9)))",
+          "expr": "(histogram_quantile(0.50, sum by (le, operation, collection) (rate(mongo_operation_duration_seconds_bucket{operation!='ensure_collection', store_method='add_many_spans'}[1m])))) * (sum by (operation, collection) (rate(mongo_operation_total{operation!='ensure_collection', store_method='add_many_spans'}[1m])) / scalar(clamp_min(sum(rate(collection_store_total{method='add_many_spans'}[1m])), 1e-9)))",
           "legendFormat": "{{operation}} - {{collection}}"
         }
       ],
@@ -1017,6 +1017,123 @@
       ],
       "title": "Rollout Completion Rate (by Status)",
       "type": "timeseries"
+    },
+    {
+      "collapsed": false,
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 140
+      },
+      "id": 304,
+      "panels": [],
+      "title": "In-memory Collection Locks",
+      "type": "row"
+    },
+    {
+      "fieldConfig": {
+        "defaults": {
+          "unit": "ops/s"
+        }
+      },
+      "gridPos": {
+        "h": 10,
+        "w": 12,
+        "x": 0,
+        "y": 141
+      },
+      "id": 305,
+      "targets": [
+        {
+          "expr": "sum by (collection) (rate(memory_collection_lock_rate_total[1m]))",
+          "legendFormat": "{{collection}}"
+        }
+      ],
+      "title": "Memory Lock Rate (by Collection)",
+      "type": "timeseries"
+    },
+    {
+      "fieldConfig": {
+        "defaults": {
+          "unit": "s"
+        }
+      },
+      "gridPos": {
+        "h": 10,
+        "w": 12,
+        "x": 12,
+        "y": 141
+      },
+      "id": 306,
+      "targets": [
+        {
+          "expr": "histogram_quantile(0.50, sum by (le, collection) (rate(memory_collection_lock_latency_seconds_bucket[1m])))",
+          "legendFormat": "P50 {{collection}}"
+        },
+        {
+          "expr": "histogram_quantile(0.95, sum by (le, collection) (rate(memory_collection_lock_latency_seconds_bucket[1m])))",
+          "legendFormat": "P95 {{collection}}"
+        },
+        {
+          "expr": "histogram_quantile(0.99, sum by (le, collection) (rate(memory_collection_lock_latency_seconds_bucket[1m])))",
+          "legendFormat": "P99 {{collection}}"
+        }
+      ],
+      "title": "Memory Lock Latency (P50, P95, P99)",
+      "type": "timeseries"
+    },
+    {
+      "collapsed": false,
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 151
+      },
+      "id": 307,
+      "panels": [],
+      "title": "HTTP Status Insights",
+      "type": "row"
+    },
+    {
+      "gridPos": {
+        "h": 10,
+        "w": 12,
+        "x": 0,
+        "y": 152
+      },
+      "id": 308,
+      "targets": [
+        {
+          "expr": "sum by (method, path, status_code) (rate(http_requests_total{path!='/v1/prometheus/'}[1m]))",
+          "legendFormat": "{{method}} {{path}} {{status_code}}"
+        }
+      ],
+      "title": "HTTP Requests / Sec (by Path & Status)",
+      "type": "timeseries"
+    },
+    {
+      "fieldConfig": {
+        "defaults": {
+          "unit": "s"
+        }
+      },
+      "gridPos": {
+        "h": 10,
+        "w": 12,
+        "x": 12,
+        "y": 152
+      },
+      "id": 309,
+      "targets": [
+        {
+          "expr": "histogram_quantile(0.95, sum by (le, method, path, status_code) (rate(http_request_duration_seconds_bucket{path!='/v1/prometheus/'}[1m])))",
+          "legendFormat": "{{method}} {{path}} {{status_code}}"
+        }
+      ],
+      "title": "HTTP Latency P95 (by Path & Status)",
+      "type": "timeseries"
     }
   ],
   "schemaVersion": 36,

docs/community/contributing.md

@@ -28,7 +28,7 @@ Documentation improvements are the easiest way to get started. You can find more
 
 Bug fixes are the fastest way to get familiar with the codebase. To get started, you can:
 
-- Browse the ["good first issue"](https://github.com/microsoft/agent-lightning/labels/good%20first%20issue) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
+- Browse the ["help wanted"](https://github.com/microsoft/agent-lightning/labels/help%20wanted) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
 - For fresh bugs, open an issue with reproduction steps, logs, and expected behavior before submitting a fix.
 - Keep each pull request focused, ideally avoiding breaking API changes. Larger refactors should be discussed via RFC or maintainer sync.
 

docs/deep-dive/store.md

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

docs/how-to/examples-catalog.md

@@ -30,6 +30,14 @@
 
     [:octicons-repo-24: Browse source]({{ src("examples/calc_x") }})
 
+-   :material-code-braces:{ .lg .middle } __Claude Code SWE-bench__
+
+    ---
+
+    Instrumented driver that runs Anthropic's Claude Code workflow on SWE-bench instances while streaming traces through Agent-lightning—supports hosted vLLM, official Anthropic, or any OpenAI-compatible backend and emits datasets for downstream tuning.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/claude_code") }})
+
 -   :material-view-grid:{ .lg .middle } __Minimal building blocks__
 
     ---

docs/reference/agent.md

@@ -22,6 +22,8 @@
 
 ## Emitter
 
+::: agentlightning.operation
+
 ::: agentlightning.emit_annotation
 
 ::: agentlightning.emit_reward

docs/reference/internal.md

@@ -24,6 +24,8 @@
 
 ::: agentlightning.litagent.decorator.prompt_rollout
 
+::: agentlightning.emitter.annotation.OperationContext
+
 ## LLM Proxy
 
 ::: agentlightning.llm_proxy.ModelConfig
@@ -44,7 +46,9 @@
 
 ::: agentlightning.store.base.UNSET
 
-::: agentlightning.store.utils.propagate_status
+::: agentlightning.store.utils.rollout_status_from_attempt
+
+::: agentlightning.store.utils.scan_unhealthy_rollouts
 
 ## Tracing and OpenTelemetry
 

docs/reference/store.md

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

docs/reference/types.md

@@ -22,6 +22,8 @@
 
 ::: agentlightning.Rollout
 
+::: agentlightning.EnqueueRolloutRequest
+
 ::: agentlightning.Attempt
 
 ::: agentlightning.AttemptedRollout

docs/tutorials/write-agents.md

@@ -211,6 +211,8 @@ While returning a single float for the final reward is sufficient for many algor
 
 Agent-lightning provides an **emitter** module that allows you to record custom spans from within your agent's logic. Like many common operations (like LLM calls) that are automatically instrumented by [Tracer][agentlightning.Tracer], the emitter will also send a [Span][agentlightning.Span] that records an Agent-lightning-specific operation. Then algorithms can query and read those spans later. See [Working with Traces](./traces.md) for more details.
 
+For multi-step routines (function calls, tools, or adapters) you can wrap code with [`operation`][agentlightning.operation], either as a decorator or a context manager,to capture inputs, outputs, and metadata on a dedicated `"agentlightning.operation"` span. This makes it easier to correlate downstream annotations (like rewards or messages) with the higher-level work that produced them.
+
 You can find the emitter functions from [agentlightning.emitter](../reference/agent.md).
 
 ### Emitting Rewards, Messages, and More
@@ -221,7 +223,6 @@ Here are the primary emitter functions:
 * [`emit_message(message: str)`][agentlightning.emit_message]: Records a simple log message as a span.
 * [`emit_exception(exception: BaseException)`][agentlightning.emit_exception]: Records a Python exception, including its type, message, and stack trace.
 * [`emit_object(obj: Any)`][agentlightning.emit_object]: Records any JSON-serializable object, perfect for structured data.
-
 Let's see an example of an agent using these emitters to provide detailed feedback.
 
 ```python
@@ -256,3 +257,55 @@ def multi_step_agent(task: dict, prompt_template: PromptTemplate) -> float:
 ```
 
 By using the emitter, you create a rich, detailed trace of your agent's execution. This data can be invaluable for debugging and is essential for advanced algorithms that can learn from more than just a single final score.
+
+### Linking to Other Spans
+
+Sometimes a span should explicitly point back to another span that produced the input it is working on (for example, linking a reward annotation to the `"agentlightning.operation"` span that generated a response). Agent-lightning encodes these relationships through flattened link attributes. The helper [`make_link_attributes`][agentlightning.utils.otel.make_link_attributes] converts a dictionary of keys—such as `trace_id`, `span_id`, or any custom attribute—into the `"agentlightning.link.*"` fields expected by the backend. Later on, [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] can be used to recover the original span(s) from those link descriptors.
+
+```python
+import opentelemetry.trace as trace_api
+from agentlightning import emit_annotation, operation
+from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
+
+with operation(conversation_id="chat-42") as op:
+    # ... perform the work ...
+    span_ctx = op.span.get_span_context()
+    link_attrs = make_link_attributes({
+        "conversation_id": "chat-42",
+    })
+
+    emit_annotation(
+        {
+            **link_attrs,
+            **make_tag_attributes(["reward", "good"]),
+        }
+    )
+```
+
+When analyzing in adapters, pass the extracted link models to [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] to retrieve the matching span(s):
+
+```python
+from agentlightning.utils.otel import extract_links_from_attributes, query_linked_spans
+
+annotation_span = ...  # Span from your trace store
+operation_spans = [...]  # list of spans you want to search
+
+link_models = extract_links_from_attributes(annotation_span.attributes)
+matches = query_linked_spans(operation_spans, link_models)
+assert matches  # Contains the original operation span
+```
+
+!!! tip "Correlating Rewards with LLM Requests"
+
+    [Tracer](./traces.md) instruments each request/response as its own span. You can link to the [`gen_ai.response.id`](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/) attribute, which comes from the LLM response ID.
+
+    ```python
+    from agentlightning import emit_reward
+    from agentlightning.utils.otel import make_link_attributes
+
+    result = call_llm(prompt)
+    reward_links = make_link_attributes({"gen_ai.response.id": result.id})
+    emit_reward(0.9, attributes=reward_links)
+    ```
+
+    Later, use the same `gen_ai.response.id` key inside `query_linked_spans` to find the reward(s) that reference that specific LLM request span.

examples/.gitignore

@@ -12,3 +12,6 @@ unsloth/unsloth_training_checkpoints/
 apo/pomltrace/
 tinker/logs/
 tinker/crewai_*.html
+rag/dataset_tiny.parquet
+rag/chunks_candidate_tiny.pkl
+rag/index_hnsw_faiss_n32e40_tiny.index

examples/README.md

@@ -7,6 +7,7 @@ This catalog highlights the examples shipped with Agent-lightning.
 | [apo](./apo) | Automatic Prompt Optimization tutorials covering built-in, custom, and debugging workflows. | [![apo workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml) |
 | [azure](./azure) | Supervised fine-tuning with Azure OpenAI. | [![azure workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-azure.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml) |
 | [calc_x](./calc_x) | VERL-powered math reasoning agent training that uses AutoGen with an MCP calculator tool. | [![calc_x workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-calc-x.yml) |
+| [claude_code](./claude_code) | Claude Code SWE-bench harness that records Agent-lightning traces across Anthropic, vLLM, and OpenAI-compatible backends. | [![claude_code workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-claude-code.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml) |
 | [minimal](./minimal) | Bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation. | [![minimal workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
 | [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | **Unmaintained** — last verified with Agent-lightning v0.1.1 |
 | [search_r1](./search_r1) | Framework-free Search-R1 reinforcement learning training workflow with a retrieval backend. | **Unmaintained** — last verified with Agent-lightning v0.1.2 |

examples/azure/README.md

@@ -1,5 +1,7 @@
 # Supervised Fine-tuning with Azure OpenAI
 
+[![azure CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml)
+
 This example walks through an end-to-end supervised fine-tuning loop on Azure OpenAI. The trainer runs a toy capital-lookup agent, collects traces with rewards, submits fine-tuning jobs using those traces, and deploys every successful checkpoint as a new Azure OpenAI deployment.
 
 **NOTE: The example is tested and compatible with Agent-lightning v0.2.x, but it's not yet maintained on CI due to the difficulty of maintaining a logged-in status in the testing environment.**

examples/claude_code/README.md

@@ -1,46 +1,54 @@
 # Training Claude Code with Agent-lightning
 
-This example demonstrates how to train a Claude Code agent with Agent-lightning. **The example is still under development.**
+[![claude-code CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml)
 
-It wraps Claude Code as the agent to:
+This example shows how to wrap Anthropic's Claude Code experience with Agent-lightning instrumentation to solve SWE-bench tasks, collect spans/logs, and optionally convert those traces into HuggingFace datasets.
 
-1. collect traces from agent execution on coding tasks;
-2. train a hosted LLM with the traces ***🔨 Under development***
+**NOTE:** This example only shows how to integrate Claude Code as an agent in Agent-lightning. The training part is still under development and welcoming contributions!
+
+## Overview
+
+`claude_code_agent.py` spins up a Lightning Store, an LLM proxy, and the Claude Code controller. Each SWE-bench instance is executed inside the official container image so you can either prompt-tune against Anthropic's hosted models or point Claude Code at a self-hosted OpenAI-compatible backend such as vLLM. When a backend surfaces token IDs/logprobs (e.g., vLLM), the traces are turned into triplets that downstream fine-tuning pipelines can consume.
 
 ## Requirements
 
-1. Install agentlightning following [installation instructions](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/);
-2. `(uv) pip install swebench` for evaluation.
+First, install Agent-lightning following the [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/). Then install the SWE-bench harness plus utilities used by this example:
+
+```bash
+(uv) pip install swebench transformers datasets python-dotenv
+```
+
+Docker must be available because each SWE-bench instance is executed in a container via `swebench_utils`.
+
+Finally, set API credentials depending on backend:
+
+- `ANTHROPIC_API_KEY` for the official Claude Code path.
+- `OPENAI_API_KEY` (or another OpenAI-compatible key) for the `openai` backend.
+- A running OpenAI-compatible server (e.g., vLLM) when using the `vllm` backend.
 
 ## Dataset
 
-We provide a small dataset `swebench_samples.jsonl` which is a subset of [SWE-bench](https://huggingface.co/datasets/SWE-bench/SWE-bench) for sanity check.
-
-The instruction to prepare the full dataset is still underway.
+`swebench_samples.jsonl` contains a handful of SWE-bench issues for smoke testing. For full-scale benchmarks load `princeton-nlp/SWE-bench` via `load_swebench_dataset` or point `--dataset-path` to your own JSONL file.
 
 ## Included Files
 
-| Filename                        | Description |
-|--------------------------------|-------------|
-| `cc_agent.py`                   | Main entry point for running Claude Code agent on coding tasks with trace collection capabilities |
-| `claude_code_controller.py`     | Controller implementation for managing Claude Code agent interactions and execution |
-| `custom_adapter.py`             | Custom adapter for integrating with Claude Code's interface and communication protocols |
-| `custom_callbacks.py`           | Callback handlers for customizing agent behavior and responses during execution |
-| `handle_hook.template.sh`       | Template script for handling hooks during agent execution |
-| `settings.template.json`        | Template configuration file with default settings for Claude Code agent |
-| `swe_debug.jsonl`               | Debug dataset containing a subset of SWE-bench samples for testing and verification |
-| `swebench_utils/`               | Utility module with helper functions for SWE-bench dataset containerized exeuction and evaluation |
+| File/Directory | Description |
+|----------------|-------------|
+| `claude_code_agent.py` | CLI entry point that launches the Lightning store, LLM proxy, and Claude Code agent |
+| `claude_code_controller.py` | Manages the SWE-bench Docker runtime and translates model outputs into git patches |
+| `extended_adapter.py` | Adapter that converts LLM proxy spans into triplets with token IDs, logprobs, and chat history |
+| `swebench_samples.jsonl` | Mini SWE-bench subset for quick validation |
+| `swebench_utils/` | Utilities for running/evaluating SWE-bench instances inside containers |
+| `templates/handle_hook.template.sh` | Helper script injected into containers for hook handling |
+| `templates/settings.template.json` | Base configuration consumed by Claude Code CLI |
 
-## Trace collection
+## Running the Example
 
-We support running Claude Code via two ways:
+All commands are issued from `examples/claude_code`. Inspect the module-level docstring in `claude_code_agent.py` for the full CLI reference.
 
-- Hosted LLM servers (i.e., vLLM), useful for fine-tuning the LLM;
-- Official Claude Code (i.e., via Anthropic API), useful for prompt tuning.
+### Hosted vLLM (open-source models)
 
-### From Hosted LLM server
-
-1. Prepare an OpenAI-compatible server:
+First, launch your model behind an OpenAI-compatible endpoint, for example:
 
 ```bash
 vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
@@ -49,37 +57,56 @@ vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
     --tool-call-parser qwen3_coder
 ```
 
-2. Sanity check:
+Run the Agent-lightning harness and point it at the server:
 
 ```bash
-# Suppose the vllm server is running at localhost:8000
-python cc_agent \
-    --model_name_or_path Qwen/Qwen3-Coder-30B-A3B-Instruct \
-    --server_address http://localhost:8000/v1 \
-    --dataset_path swe_debug.jsonl \
-    --max_step 32 \
-    --output_dir data_debug
+python claude_code_agent.py vllm \
+    --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
+    --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
+    --frontend-model-high claude-sonnet-4-5-20250929 \
+    --frontend-model-low claude-haiku-4-5-20251001 \
+    --base-url http://localhost:8000/v1 \
+    --dataset-path swebench_samples.jsonl \
+    --output-dir data_debug \
+    --max-turns 5 \
+    --limit 2
 ```
 
-The above commands will generate a `data_debug` dir, which contains two targets: (1) a Huggingface Dataset named `dataset-<instance_id>` and (2) a trace file named `stream_<instance_id>.jsonl`, where `instance_id` is a unique key of the SWE-bench samples.
-The dataset showcases the versatile customization capability of agent-lightning. In particular, we support extracting **prompt/response ids**, **logprobs** from the vllm server.
-The trace file is the conversation logs for claude code to tackle the SWE-bench instance.
+The backend model names must match what the server exposes. Because this mode surfaces token IDs/logprobs, the script saves both raw span logs and HuggingFace datasets per instance.
 
-In addition, there will be a `logs` dir, which is the output of the docker container executing agent calls.
+### Official Claude Code (Anthropic API)
 
-### From official Claude Code
-1. Prepare ANTHROPIC_API_KEY
 ```bash
-export ANTHROPIC_API_KEY=sk-<your private key>
+export ANTHROPIC_API_KEY=sk-...
+python claude_code_agent.py anthropic \
+    --dataset-path swebench_samples.jsonl \
+    --output-dir data_anthropic \
+    --frontend-model-high claude-sonnet-4-5-20250929 \
+    --frontend-model-low claude-haiku-4-5-20251001
 ```
 
-2. Sanity check
+Backend model flags are optional here because the Anthropic API strings match the frontend names. This path is ideal for validating prompts against the hosted experience (trace outputs do not contain token IDs or logprobs).
+
+### OpenAI-Compatible Providers
+
 ```bash
-cd examples/cc
-python cc_agent \
-    --official \
-    --dataset_path swe_debug.jsonl \
-    --max_step 32 \
-    --output_dir data_debug
+export OPENAI_API_KEY=sk-...
+python claude_code_agent.py openai \
+    --backend-model-high gpt-4.1 \
+    --backend-model-low gpt-4o-mini \
+    --dataset-path swebench_samples.jsonl \
+    --output-dir data_openai
 ```
-As the underlying model is provided by Anthropic, we cannot obtain prompt/response ids and logprobs. However, we can still obtain a trace file named `<instance_id>.json` under `data_debug`.
+
+Use this mode whenever Claude Code should talk to Azure OpenAI, OpenAI, or another compatible provider. `--base-url` is optional—pass it if your endpoint differs from the public OpenAI URL.
+
+Adjust `--max-turns`, `--cooldown-seconds`, and `--limit` to control runtime and rate limits regardless of backend.
+
+## Outputs and Trace Collection
+
+- `output_dir/stream_<instance_id>.json` contains the complete span stream captured from the Lightning Store for each rollout.
+- When running with `backend_type=vllm`, `output_dir/dataset-<instance_id>/` stores a HuggingFace dataset with token IDs, logprobs, prompts, and metadata produced by `ExtendedLlmProxyTraceToTriplet`.
+- `logs/<instance_id>/` is created by the SWE-bench runtime and mirrors the console output from the container.
+- Return values from the agent are also evaluated via `swebench_utils.evaluation.evaluate`, so `data_debug` (or your chosen folder) will contain evaluation reports alongside traces.
+
+Use these artifacts to fine-tune models, debug Claude Code behavior, or replay rollouts in downstream Agent-lightning workflows.

examples/claude_code/claude_code_agent.py

@@ -1,16 +1,54 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Main module for the Claude Code Agent implementation.
+"""Instrumented driver for running Claude Code on SWE-bench with Agent-lightning.
 
-This module provides the core functionality for running Claude Code agent experiments
-on SWE-bench datasets. It includes the ClaudeCodeAgent class that implements the agent logic,
-functions for loading datasets, and asynchronous execution functions for running experiments.
+This script wires together the Lightning Store, LLM proxy, and Claude Code controller so
+that every SWE-bench instance is executed inside the official Claude container while
+capturing full Agent-lightning traces. It supports three backend modes:
 
-Key components:
+- `vllm`: wrap an OpenAI-compatible endpoint (e.g., vLLM) for hosted OSS models while
+  collecting prompt/response token ids and logprobs.
+- `anthropic`: call the official Claude Code API via `ANTHROPIC_API_KEY` for prompt
+  tuning. Backend model defaults to the provided frontend names.
+- `openai`: route through any OpenAI-compatible provider using `OPENAI_API_KEY`.
 
-- Dataset loading utilities
-- ClaudeCodeAgent: Main agent implementation that handles rollout logic
-- Asynchronous execution functions for dry runs and full datasets
+Typical usage: hosted vLLM (requires model paths and --base-url)
+
+```bash
+# Run vLLM in background
+vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+  --max-model-len 131072 \
+  --enable-auto-tool-choice \
+  --tool-call-parser qwen3_coder \
+  --port 45993 &
+
+python claude_code_agent.py vllm \
+  --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
+  --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
+  --base-url http://localhost:45993/v1 \
+  --dataset-path swebench_samples.jsonl \
+```
+
+Official Claude Code via Anthropic:
+
+```bash
+export ANTHROPIC_API_KEY=sk-...
+python claude_code_agent.py anthropic \
+  --dataset-path swebench_samples.jsonl \
+  --output-dir data_anthropic
+```
+
+Any OpenAI-compatible backend:
+
+```bash
+export OPENAI_API_KEY=sk-...
+python claude_code_agent.py openai \
+  --backend-model-high gpt-5.1-codex-mini \
+  --backend-model-low gpt-4.1-mini \
+  --dataset-path swebench_samples.jsonl
+```
+
+Use `--debug` to enable debug loggings.
 """
 
 import asyncio

examples/minimal/README.md

@@ -1,5 +1,7 @@
 # Minimal Component Showcase
 
+[![minimal CI status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
+
 `examples/minimal` provides bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation.
 
 Each module have been documented with its own CLI usage in the module-level docstring. Use this directory as a reference when wiring the same pieces into a larger system.
@@ -9,6 +11,7 @@ Each module have been documented with its own CLI usage in the module-level docs
 | Component | Demonstrated In | Highlights |
 | --- | --- | --- |
 | LightningStore + OTLP ingestion | `write_traces.py` | Shows how `OtelTracer` and `AgentOpsTracer` open rollouts, emit spans, and optionally forward them to a remote store client. |
+| MultiMetrics backend | `write_metrics.py` | Emits counters/histograms through `ConsoleMetricsBackend` and `PrometheusMetricsBackend` simultaneously, exposing `/metrics` for scraping. |
 | LLM proxying | `llm_proxy.py` | Guards either OpenAI or a local vLLM deployment with `LLMProxy`, proving how requests are routed through `/rollout/<id>/attempt/<id>` namespaces and captured in the store. |
 | vLLM lifecycle | `vllm_server.py` | Minimal context manager that shells out to `vllm serve`, monitors readiness, and tears down the process safely. |
 

examples/minimal/write_metrics.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Demonstrate `MultiMetricsBackend` by emitting metrics to both console and Prometheus.
+
+Usage:
+    python write_metrics.py --duration 10 --prom-port 8000
+
+The script registers a counter and a histogram, pushes events through both the
+`ConsoleMetricsBackend` (for immediate feedback) and the `PrometheusMetricsBackend`
+for scraping via `/metrics`.
+
+Run a Prometheus server (for example via `docker/compose.prometheus-memory-store.yml`)
+and add the host running this script as a scrape target. By default the metrics
+endpoint binds to `0.0.0.0:9105`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+import signal
+import sys
+import time
+from typing import Sequence
+
+from prometheus_client import start_http_server
+
+from agentlightning import setup_logging
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+)
+
+
+def _register_metrics(backend: MetricsBackend) -> None:
+    backend.register_counter("minimal_requests_total", ["operation", "status"])
+    backend.register_histogram(
+        "minimal_latency_seconds",
+        ["operation"],
+        buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0],
+    )
+
+
+def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
+    statuses = ["200", "404", "500"]
+    end_time = time.time() + duration
+    random.seed(1337)
+    while time.time() < end_time:
+        operation = random.choice(operations)
+        status = random.choices(statuses, weights=[0.9, 0.05, 0.05], k=1)[0]
+        latency = random.lognormvariate(-4.0, 0.5)
+        backend.inc_counter("minimal_requests_total", labels={"operation": operation, "status": status})
+        backend.observe_histogram("minimal_latency_seconds", value=latency, labels={"operation": operation})
+        time.sleep(0.25)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--duration", type=float, default=10.0, help="Seconds to emit metrics before shutting down.")
+    parser.add_argument("--prom-port", type=int, default=9105, help="Port for the /metrics endpoint.")
+    parser.add_argument("--prom-host", default="0.0.0.0", help="Host/IP for the /metrics endpoint.")
+    parser.add_argument("--group-level", type=int, default=2, help="ConsoleMetricsBackend label grouping depth.")
+    args = parser.parse_args()
+
+    setup_logging()
+
+    console_backend = ConsoleMetricsBackend(window_seconds=15.0, log_interval_seconds=2.0, group_level=args.group_level)
+    prom_backend = PrometheusMetricsBackend()
+    backend = MultiMetricsBackend([console_backend, prom_backend])
+    _register_metrics(backend)
+
+    start_http_server(args.prom_port, addr=args.prom_host)
+    print(f"Prometheus metrics exposed on http://{args.prom_host}:{args.prom_port}/metrics")
+    print(f"Emitting demo metrics for {args.duration:.1f}s ...")
+
+    # Handle CTRL+C gracefully
+    interrupted = False
+
+    def _handle_interrupt(signum: int, frame: object | None) -> None:  # pragma: no cover - signal handler
+        nonlocal interrupted
+        print(f"Received signal {signum}, stopping...")
+        interrupted = True
+
+    original_handler = signal.signal(signal.SIGINT, _handle_interrupt)
+    try:
+        _emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"])
+    finally:
+        signal.signal(signal.SIGINT, original_handler)
+
+    if interrupted:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/rag/README.md

@@ -1,62 +1,84 @@
 # RAG Agent Example
 
-This example demonstrates training a Retrieval-Augmented Generation (RAG) agent using Agent-Lightning with Wikipedia retrieval capabilities. The agent answers multi-hop questions from the MuSiQue dataset by retrieving and reasoning over Wikipedia passages. **It's tested and compatible with Agent-lightning v0.1.x**.
+This example demonstrates training a Retrieval-Augmented Generation (RAG) agent using Agent-Lightning with retrieval capabilities. The agent answers multi-hop questions from a tiny MuSiQue dataset by retrieving and reasoning over Wikipedia passages.
 
 ## Overview
 
-This example originally runs on a single node with four GPUs, each requiring at least 40GB of memory.
+This example can run on a single GPU for demonstration purposes.
 
-1. Prepare the RAG dataset in the wiki_retriever_mcp folder. Wiki chunks (`nq_list.pkl`) and Faiss index (`nq_hnsw_faiss_n32e40.index`) are required. (Full wiki dump files are huge, additional information will be provided later)
-2. Prepare the training data in the `data` folder. Download from [here](https://drive.google.com/drive/folders/1hEqOY4EbplUB5ew-8UPFhV_5QU2j7WCN?usp=drive_link). `musique_train.parquet` and `musique_dev_128.parquet` are required.
-3. Set up the environment for wiki retriever MCP: `bash wiki_retriever_install.sh`. This will install the required packages and set up the environment for the wiki retriever MCP.
-4. Start the wiki retriever MCP: `python wiki_retriever_mcp.py`. This will start the wiki retriever MCP server.
-5. Start Ray: `bash ../../scripts/restart_ray.sh`. To use Wandb, you need to set the WANDB_API_KEY environment variable before starting Ray.
-6. Run the agent: `python rag_agent.py`. This automatically launches 12 agent workers by default.
-7. In another terminal, launch the training server: `bash train.sh`.
+**Step 1:** Set up the environment. It is recommended to setup with uv and activate the virtual environment with:
+
+```bash
+uv sync --frozen --extra apo --group agents --group torch-gpu-stable --extra verl --group rag
+source .venv/bin/activate
+```
+
+**Step 2:** Prepare the tiny dataset.
+
+```bash
+pip install gdown
+
+# tiny training dataset
+cd examples/rag
+gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" \
+    -O dataset_tiny.parquet
+
+# chunks_candidate_tiny.pkl
+gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" \
+    -O chunks_candidate_tiny.pkl
+
+# index_hnsw_faiss_n32e40_tiny.index
+gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" \
+    -O index_hnsw_faiss_n32e40_tiny.index
+```
+
+**Step 3:** Start the MCP server. Open a terminal and run:
+
+```bash
+python wiki_retriever_mcp.py
+```
+
+**Step 4:** Start training. Open another terminal and run:
+
+```bash
+python train_rag.py
+```
 
 ## Included Files
 
 | File/Directory | Description |
 |----------------|-------------|
-| `rag_agent.py` | Entry point for running the Agent-Lightning RAG training pipeline |
-| `train.sh` | Starts the GRPO training server that updates the agent |
-| `utils.py` | Scoring utilities for exact match, F1, and response parsing |
-| `wiki_retriever_mcp/` | Setup scripts and MCP server (`wiki_retriever_install.sh`, `wiki_retriever_mcp.py`) for Wikipedia retrieval |
+| `rag_agent.py` | RAG agent example using the OpenAI Agents SDK, with debugging utils |
+| `train_rag.py` | Initiates the GRPO training process |
+| `metric_utils.py` | Scoring utilities for exact match, F1 score, and response parsing |
+| `wiki_retriever_mcp.py` | MCP server for Wikipedia retrieval |
 
-## Preparing the Retrieval Corpus
+## How to Prepare the Retrieval Corpus Yourself
 
-To enable semantic retrieval with this mcp server, we need two files:
+To enable semantic retrieval with this MCP server, you need two files:
 
 1. **FAISS index file** (`.index`)
 2. **Chunk list file** (`.pkl`)
 
 These two files work together: the FAISS index stores the vector embeddings and their mapping to integer IDs, while the pickle file stores the actual text chunks. The integer IDs in the index correspond exactly to the positions in the chunk list.
 
----
+### Step 1: Collecting Text Chunks
 
-### Step 1. Collecting Text Chunks
+First, you need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` from the [FlashRAG_dataset](https://huggingface.co/datasets/FlashRAG) or use other pre-split corpora.
 
-You first need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` in the [FlashRAG_dataset](https://huggingface.co/datasets/FlashRAG) or use other pre-split corpora.
-
----
-
-### Step 2. Creating the FAISS Index (`nq_hnsw_faiss_n32e40.index`)
+### Step 2: Creating the FAISS Index (`nq_hnsw_faiss_n32e40.index`)
 
 - Use a sentence embedding model (e.g., `BAAI/bge-large-en-v1.5`) to encode each chunk into a vector.
 - Build a FAISS index from these vectors.
 - In this example, we use an **HNSW index** (Hierarchical Navigable Small World graph), which supports efficient approximate nearest-neighbor search.
-- The index only stores embeddings and integer IDs (no raw text).
+- The index stores only embeddings and integer IDs (no raw text).
 
----
-
-### Step 3. Creating the Chunk List (`nq_list.pkl`)
+### Step 3: Creating the Chunk List (`nq_list.pkl`)
 
 - Store the raw text chunks in a Python list.
 - Save this list with `pickle`.
 - The index ID returned by FAISS corresponds to the list index in this file. For example, if FAISS search returns `I[0][i] = 12345`, then the corresponding text chunk is `chunks[12345]`.
 
----
-
 ### Example Schema
 
 - **`nq_hnsw_faiss_n32e40.index`**
@@ -78,10 +100,9 @@ You first need a collection of text passages (chunks). For example, you can down
     ]
     ```
 
----
+### Step 4: Code Example - Building Index and Chunk List
 
-### Step 4. Code Example: Building Index and Chunk List
-Warning: The following example only demonstrates a small-scale workflow. In practice, if the dataset is large, you should encode the text in batches and incrementally add them to the index.
+**Warning:** The following example demonstrates a small-scale workflow only. In practice, for large datasets, you should encode the text in batches and incrementally add them to the index.
 
 ```python
 import faiss
@@ -117,8 +138,3 @@ with open("nq_list.pkl", "wb") as f:
 
 print("Index and chunk list saved successfully.")
 ```
-
-
-## Evaluation
-
-Results are coming soon.

examples/rag/metric_utils.py

@@ -109,10 +109,10 @@ def split_response(text: str) -> Tuple[str, str]:
 def extract_recall_chunk(prompt: str, response: str) -> Tuple[Set[str], Set[str]]:
     import re
 
-    # 正则表达式，匹配每个search_step内1.和2.后面的内容
+    # Regular expression to match content after 1. and 2. within each search_step
     pattern = r"Retrieved sentences:\s*1\.\s*(.*?)\s*2\.\s*(.*?)(?:\n\s*\d+\.|\n\n|$)"
 
-    # 使用re.findall 提取所有的(s1, s2)
+    # Use re.findall to extract all (s1, s2) pairs
     origin_recall = re.findall(pattern, prompt, re.DOTALL)
     sequential_recall = re.findall(pattern, response, re.DOTALL)
     origin_recall_set = set(s for pair in origin_recall for s in pair)
@@ -121,14 +121,11 @@ def extract_recall_chunk(prompt: str, response: str) -> Tuple[Set[str], Set[str]
     return origin_recall_set, sequential_recall_set
 
 
-import re
-
-
 def extract_retrieved_paragraphs(log_text: str) -> List[str]:
-    # 正则表达式匹配 "Retrieved paragraph:" 后的内容
+    # Regular expression to match content after "Retrieved paragraph:"
     pattern = re.compile(r"Retrieved paragraph:\s*(.*?)\n", re.DOTALL)
 
-    # 提取匹配的段落
+    # Extract matched paragraphs
     matches = pattern.findall(log_text)
     matches = list(set(matches))
     return matches

examples/rag/rag_agent.py

@@ -2,23 +2,19 @@
 
 from __future__ import annotations
 
-from typing import Any, cast
+import logging
+from typing import Any, Dict, List, cast
 
+import pandas as pd
 from agents import Agent, Runner
 from agents.extensions.models.litellm_model import LitellmModel
 from agents.mcp import MCPServerSse
 from agents.model_settings import ModelSettings
-from utils import compute_scores
+from metric_utils import compute_scores
 
-from agentlightning import (
-    LLM,
-    LitAgent,
-    NamedResources,
-    Trainer,
-    setup_logging,
-)
+import agentlightning as agl
 
-setup_logging()
+logger = logging.getLogger("rag_agent")
 
 agent_prompt = """You are an assistant who answers questions using Wikipedia retriever. Answer the question using only the retrieved passages. Verify your answer directly against the text.
 
@@ -32,22 +28,36 @@ After each search:
 Repeat as needed. When done, wrap your final, concise answer in <answer> tags."""
 
 
-class RAGAgent(LitAgent[Any]):
-    def __init__(self, trained_agents: str | None = None) -> None:
-        super().__init__(trained_agents=trained_agents)
+class RAGAgent(agl.LitAgent[Dict[str, Any]]):
+    """RAGAgent is an agent that relies on a MCP-based retriever to answer questions."""
+
+    def __init__(self) -> None:
+        super().__init__()
         self.mcp_server_url = "http://127.0.0.1:8099/sse"
 
-    async def training_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
-        llm: LLM = cast(LLM, resources.get("main_llm"))
-        print("Training with model:", llm.model, "on endpoint:", llm.endpoint)
+    async def training_rollout_async(
+        self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout
+    ) -> float | None:
+        # llm resources
+        llm = cast(agl.LLM, resources["main_llm"])
+
+        # The rollout should carry an attempt inside
+        rollout = cast(agl.AttemptedRollout, rollout)
+        base_url = llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id)
+
+        logger.info(f"Training with model: {llm.model} on endpoint: {base_url}")
+
         async with MCPServerSse(
             name="wiki_retriever_mcp",
             params={"url": self.mcp_server_url},
         ) as server:
             agent = Agent(
-                model=LitellmModel(model="hosted_vllm/" + llm.model, base_url=llm.endpoint),
+                model=LitellmModel(
+                    model="hosted_vllm/" + llm.model,
+                    base_url=base_url,
+                ),
                 model_settings=ModelSettings(
-                    max_tokens=4096,
+                    max_tokens=2048,
                     temperature=0.7,
                 ),
                 name="Assistant",
@@ -55,26 +65,68 @@ class RAGAgent(LitAgent[Any]):
                 mcp_servers=[server],
             )
             result = await Runner.run(agent, task["question"])
-            answer = result.final_output  # type: ignore
-            reward = compute_scores(answer, str(task["answer"]))
-            print(
-                "question:{} answer: {} ground_truth: {} reward: {}".format(
-                    task["question"], answer, task["answer"], reward
-                )
-            )
-            return reward
+            answer = result.final_output
 
-    async def validation_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
-        llm: LLM = cast(LLM, resources.get("main_llm"))
-        resources = {
-            "main_llm": LLM(
-                endpoint=llm.endpoint,
+            # reward
+            reward = compute_scores(answer, str(task["answer"]))
+
+            logger.info(
+                "Question: %s\nAnswer: %s\nGround truth: %s\nReward: %s",
+                task["question"],
+                answer,
+                task["answer"],
+                reward,
+            )
+            return float(reward)  # Convert to float for compatibility with the Runner
+
+    async def validation_rollout_async(
+        self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout
+    ) -> float | None:
+        """Validation rollout will share the same logic as the training rollout."""
+        # Same as training rollout, but with different temperature
+        llm = cast(agl.LLM, resources["main_llm"])
+        rollout = cast(agl.AttemptedRollout, rollout)
+
+        # set temperature
+        val_resources: agl.NamedResources = {
+            "main_llm": agl.LLM(
+                endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
                 model=llm.model,
                 sampling_parameters={"temperature": 0.7},
             )
         }
-        return await self.training_rollout_async(task, rollout_id, resources)
+
+        # reuse training rollout for validation
+        return await self.training_rollout_async(task, val_resources, rollout)
+
+
+def debug():
+    """Debug the RAGAgent."""
+
+    agl.setup_logging("DEBUG", apply_to=[logger.name])
+
+    # 1. loading dataset
+    dataset_path = "data/dataset_tiny.parquet"
+    df: pd.DataFrame = pd.read_parquet(dataset_path)  # type: ignore
+    data: List[Dict[str, Any]] = df.head(5).to_dict(orient="records")  # type: ignore
+    # NOTE: The following dummy data can also be used if you don't have the dataset.
+    # data: List[Dict[str, Any]] = [{"question": "What is the capital of France?", "answer": "Paris"}]
+
+    # 2. configuring resources (LLM)
+    # Note: You need to start a local service compatible with the OpenAI API (such as vLLM)
+    # For example: python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-1.5B-Instruct --port 8000
+    resources: dict[str, agl.ResourceUnion] = {
+        "main_llm": agl.LLM(
+            endpoint="http://localhost:8000/v1",  # Replace with your actual vLLM address
+            model="Qwen/Qwen2.5-1.5B-Instruct",  # Replace with your actual loaded model name
+            sampling_parameters={"temperature": 0.0},
+        )
+    }
+
+    # 3. run agent
+    trainer = agl.Trainer(initial_resources=resources)
+    trainer.dev(RAGAgent(), train_dataset=data)  # type: ignore
 
 
 if __name__ == "__main__":
-    Trainer(n_workers=12).fit_v0(RAGAgent(), "http://localhost:9999/")
+    debug()

examples/rag/train.sh

@@ -1,53 +0,0 @@
-#!/bin/bash
-
-set -e
-
-export N_GPUS=1
-export BASE_MODEL=Qwen/Qwen3-1.7B
-export DATA_DIR=data
-export ROLLOUT_TP_SIZE=1
-export EXPERIMENT_NAME=rag_agent
-export PROJECT_NAME=AgentLightning
-
-echo "Starting training script..."
-
-python -m agentlightning.verl \
-    algorithm.adv_estimator=grpo \
-    data.train_files=${DATA_DIR}/musique_train.parquet \
-    data.val_files=${DATA_DIR}/musique_dev_128.parquet \
-    actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP_SIZE \
-    trainer.n_gpus_per_node=${N_GPUS} \
-    data.train_batch_size=32 \
-    actor_rollout_ref.rollout.n=4 \
-    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
-    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
-    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
-    actor_rollout_ref.rollout.multi_turn.format=hermes \
-    actor_rollout_ref.model.path=${BASE_MODEL} \
-    data.max_prompt_length=4096 \
-    data.max_response_length=2048 \
-    data.truncation='error' \
-    trainer.val_before_train=True \
-    actor_rollout_ref.actor.optim.lr=1e-6 \
-    actor_rollout_ref.model.use_remove_padding=True \
-    actor_rollout_ref.actor.use_kl_loss=False \
-    actor_rollout_ref.actor.kl_loss_coef=0.000 \
-    actor_rollout_ref.actor.entropy_coeff=0 \
-    actor_rollout_ref.actor.clip_ratio_low=0.2 \
-    actor_rollout_ref.actor.clip_ratio_high=0.3 \
-    actor_rollout_ref.model.enable_gradient_checkpointing=True \
-    actor_rollout_ref.actor.fsdp_config.param_offload=True \
-    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
-    actor_rollout_ref.rollout.name=vllm \
-    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
-    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
-    actor_rollout_ref.ref.fsdp_config.param_offload=True \
-    algorithm.use_kl_in_reward=False \
-    trainer.critic_warmup=0 \
-    trainer.logger=['console','wandb'] \
-    trainer.project_name=${PROJECT_NAME} \
-    trainer.experiment_name=${EXPERIMENT_NAME} \
-    trainer.nnodes=1 \
-    trainer.save_freq=40 \
-    trainer.test_freq=20 \
-    trainer.total_epochs=2 $@

examples/rag/train_rag.py

@@ -0,0 +1,200 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train a RAG agent using Agent-lightning.
+
+Usage:
+    python train_rag.py fast        # Fast training for CI/testing
+    python train_rag.py single_gpu  # Optimized for Single GPU (1.5B/7B models)
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from rag_agent import RAGAgent  # Make sure to import your RAGAgent class
+
+import agentlightning as agl
+
+# Base configuration (default configuration, can be overridden)
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",  # Use GRPO algorithm
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_batch_size": 16,  # Default configuration for multi-GPU
+        "max_prompt_length": 8192,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,  # Generate 4 responses per sampling
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},  # Ensure using template format matching the model
+            "name": "vllm",
+            "gpu_memory_utilization": 0.6,  # vLLM GPU memory utilization
+            "engine_kwargs": {
+                "vllm": {
+                    "enable_auto_tool_choice": True,
+                    "tool_call_parser": "hermes",
+                }
+            },
+        },
+        "actor": {
+            "ppo_mini_batch_size": 16,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6},
+            "use_kl_loss": False,
+            "kl_loss_coef": 0.0,
+            "entropy_coeff": 0,
+            "clip_ratio_low": 0.2,
+            "clip_ratio_high": 0.3,
+            "fsdp_config": {
+                "param_offload": True,  # Enable parameter offloading to save GPU memory
+                "optimizer_offload": True,
+            },
+        },
+        "ref": {
+            "log_prob_micro_batch_size_per_gpu": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-1.5B-Instruct",  # Default model
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console"],  # Disable wandb for easier local debugging, add back when needed
+        "project_name": "AgentLightning",
+        "experiment_name": "rag_agent",
+        "nnodes": 1,
+        "test_freq": 10,
+        "total_epochs": 200,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """Fast training configuration for CI/testing"""
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    random_suffix = uuid.uuid4().hex[:8]
+    EXPERIMENT_NAME = f"rag_fast_{timestamp}_{random_suffix}"
+
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+        print("Set environment variables:")
+        print(f"PROJECT_NAME={PROJECT_NAME}")
+        print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+
+    # Keep it tiny/light without adding new knobs
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
+    config["trainer"]["total_epochs"] = 2
+    config["trainer"]["test_freq"] = 5
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["logger"] = ["console", "wandb"]
+    return config
+
+
+def config_train_single_gpu() -> Dict[str, Any]:
+    """Single GPU training optimized configuration (optimized for 24GB GPU memory)"""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+
+    # 1. Reduce vLLM memory usage to leave space for training
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.4
+
+    # 2. Reduce Batch Size to prevent OOM
+    config["data"]["train_batch_size"] = 4
+    config["actor_rollout_ref"]["actor"]["ppo_mini_batch_size"] = 4
+    config["actor_rollout_ref"]["actor"]["ppo_micro_batch_size_per_gpu"] = 1
+    config["actor_rollout_ref"]["rollout"]["log_prob_micro_batch_size_per_gpu"] = 2
+
+    # 3. Ensure Offload is enabled
+    config["actor_rollout_ref"]["actor"]["fsdp_config"]["param_offload"] = True
+    config["actor_rollout_ref"]["actor"]["fsdp_config"]["optimizer_offload"] = True
+
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the RAG agent with the given configuration."""
+
+    # 1. Instantiate your Agent
+    agent = RAGAgent()
+
+    # 2. Initialize algorithm (VERL)
+    algorithm = agl.VERL(config)
+
+    # 3. Initialize Trainer
+    # n_runners=4 means 4 concurrent rollout runners (can be reduced if insufficient memory, or managed internally by VERL)
+    trainer = agl.Trainer(n_runners=4, algorithm=algorithm, adapter={"agent_match": active_agent})
+
+    # 4. Load data
+    # NOTE: Fill in the path to your previously converted parquet file here
+    # For demo purposes, we use the same dataset for training and validation,
+    # which should be avoided in production.
+    train_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore
+    val_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore
+
+    # Keep the rest of the code unchanged
+    train_data: List[Dict[str, Any]] = train_df.to_dict(orient="records")  # type: ignore
+    val_data: List[Dict[str, Any]] = val_df.to_dict(orient="records")  # type: ignore
+
+    # 5. Start training
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train a RAG agent using different configurations")
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "single_gpu"],
+        default="single_gpu",
+        nargs="?",
+        help="Training configuration name",
+    )
+
+    parser.add_argument("--active-agent", type=str, help="Override the active agent name")
+
+    args = parser.parse_args()
+
+    config_functions = {
+        "fast": config_train_fast,
+        "single_gpu": config_train_single_gpu,
+    }
+    config = config_functions[args.config]()
+
+    # Print key information for confirmation
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Model: {config['actor_rollout_ref']['model']['path']}")
+    print(f"Batch Size: {config['data']['train_batch_size']}")
+    print(f"GPU Mem Util: {config['actor_rollout_ref']['rollout']['gpu_memory_utilization']}")
+
+    train(config, args.active_agent)
+
+
+if __name__ == "__main__":
+    main()

examples/rag/wiki_retriever_mcp.py

@@ -8,15 +8,14 @@ import faiss
 from fastmcp import FastMCP
 from sentence_transformers import SentenceTransformer
 
-# index = faiss.read_index("/mnt/input/agent_lightning/nq_hnsw_faiss_n32e40.index")
-index = faiss.read_index("nq_hnsw_faiss_n32e40.index")
+index = faiss.read_index("data/index_hnsw_faiss_n32e40_tiny.index")
 print("Index loaded successfully.")
 
 model = SentenceTransformer("BAAI/bge-large-en-v1.5")
 print("Model loaded successfully.")
 
 # with open('/mnt/input/agent_lightning/nq_list.pkl', 'rb') as f:
-with open("nq_list.pkl", "rb") as f:
+with open("data/chunks_candidate_tiny.pkl", "rb") as f:
     chunks = pickle.load(f)
 print("Chunks loaded successfully.")
 
@@ -37,7 +36,7 @@ def retrieve(query: str) -> list:
     Returns:
         list: A list of dictionaries containing the retrieved chunks and their metadata.
     """
-    top_k = 4  # Number of top results to return
+    top_k = 1  # Number of top results to return
     embedding = model.encode([query], normalize_embeddings=True)
     D, I = index.search(embedding, top_k)
 
@@ -45,7 +44,7 @@ def retrieve(query: str) -> list:
     for i in range(top_k):
         if I[0][i] != -1:
             chunk = chunks[I[0][i]]
-            results.append({"chunk": chunk, "chunk_id": I[0][i], "distance": D[0][i]})
+            results.append({"chunk": chunk, "chunk_id": int(I[0][i]), "distance": float(D[0][i])})
     return results
 
 

examples/rag/wiki_retriever_mcp/wiki_retriever_install.sh

@@ -1,3 +0,0 @@
-conda create -n mcp_server python=3.12 -y
-conda activate mcp_server
-pip install faiss-cpu==1.11.0 fastmcp==2.5.1 sentence-transformers==4.1.0

examples/tinker/README.md

@@ -1,8 +1,8 @@
 # Tinker + Agent-lightning Integration
 
-This example shows how to use [Tinker's reinforcement-learning infrastructure](https://tinker-docs.thinkingmachines.ai/) as a fine-tuning backend for agents written against Agent-lightning. You author the agent exactly the way you would for deployment, while the bridge code reconstructs Tinker-compatible trajectories from Agent-lightning traces.
+[![tinker CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml)
 
-**NOTE: The example is tested and compatible with Agent-lightning v0.2.x, but it's not yet maintained on CI due to the cost of running the Tinker training service.**
+This example shows how to use [Tinker's reinforcement-learning infrastructure](https://tinker-docs.thinkingmachines.ai/) as a fine-tuning backend for agents written against Agent-lightning. You author the agent exactly the way you would for deployment, while the bridge code reconstructs Tinker-compatible trajectories from Agent-lightning traces.
 
 ## How this differs from the original Tinker Cookbook RL recipe
 

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
   "portpicker",
   "gunicorn",
   "uvicorn_worker",
+  "aiologic",
 ]
 
 [project.optional-dependencies]
@@ -102,8 +103,8 @@ torch-stable = [
   "transformers>=4.55.0,!=4.57.2",
   # vLLM 0.11.1 requires PyTorch 2.9.0, which is incompatible with flash-attn
   # https://github.com/Dao-AILab/flash-attention/issues/1967
-  # Similar issues with vLLM 0.11.2
-  "vllm>=0.10.2,!=0.11.1,!=0.11.2",
+  # Similar issues with vLLM 0.11.2 and 0.12.0
+  "vllm>=0.10.2,!=0.11.1,!=0.11.2,!=0.12.0",
   # LiteLLM can then be upgraded with new vLLM
   "litellm[proxy]>=1.78",
 ]
@@ -191,6 +192,11 @@ sql = [
   "sqlparse",
   "nltk",
 ]
+rag = [
+  "fastmcp>=2.13.1",
+  "faiss-cpu>=1.11.0",
+  "sentence-transformers>=4.1.0",
+]
 crewai = [
   # https://github.com/crewAIInc/crewAI/issues/3959
   "crewai[tools]>=1.2.0,!=1.2.1,!=1.3.0,!=1.4.0,!=1.4.1,!=1.5.0",
@@ -241,6 +247,9 @@ override-dependencies = [
   "uvicorn>=0.38.0",
   # Conflicts between packaging dependency of pyvers (dependency of tensordict) and agentops.
   "packaging>=24.0",
+  # Conflicts between litellm and fastmcp
+  "websockets>=15.0.1",
+  "rich>=13.9.4",
 ]
 
 [tool.uv.sources]

tests/benchmark/analysis.py

@@ -153,8 +153,13 @@ def format_window(seconds: float) -> str:
     return f"{seconds}s"
 
 
-def compute_rate_window(duration_seconds: float) -> str:
-    return format_window(min(duration_seconds, 60.0))
+def clamp_window_seconds(duration_seconds: float) -> int:
+    return max(int(duration_seconds), 1)
+
+
+def compute_peak_window(duration_seconds: float) -> str:
+    peak_seconds = max(min(int(duration_seconds), 60), 1)
+    return f"{peak_seconds}s"
 
 
 def compute_subquery_step(duration_seconds: float) -> str:
@@ -355,20 +360,19 @@ class RolloutOutcomeStats:
 def gather_store_methods(
     client: PrometheusClient,
     window: str,
-    rate_window: str,
+    window_seconds: int,
+    peak_window: str,
     subquery_step: str,
 ) -> Tuple[List[StoreMethodStats], StatsSummary]:
-    ops_expr = f"sum by (method)(rate(collection_store_total[{rate_window}]))"
-    ops_mean = vector_to_map(
-        safe_vector(client, f"avg_over_time(({ops_expr})[{window}:{subquery_step}])"),
-        ("method",),
-    )
+    mean_expr = f"(sum by (method)(increase(collection_store_total[{window}]))) / {window_seconds}"
+    ops_mean = vector_to_map(safe_vector(client, mean_expr), ("method",))
+    peak_expr = f"sum by (method)(irate(collection_store_total[{peak_window}]))"
     ops_max = vector_to_map(
-        safe_vector(client, f"max_over_time(({ops_expr})[{window}:{subquery_step}])"),
+        safe_vector(client, f"max_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method",),
     )
     ops_min = vector_to_map(
-        safe_vector(client, f"min_over_time(({ops_expr})[{window}:{subquery_step}])"),
+        safe_vector(client, f"min_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method",),
     )
     p50 = vector_to_map(
@@ -412,16 +416,16 @@ def gather_store_methods(
     overall: StatsSummary = {
         "ops_mean": safe_scalar(
             client,
-            f"avg_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+            f"(sum(increase(collection_store_total[{window}]))) / {window_seconds}",
         )
         or 0.0,
         "ops_max": safe_scalar(
             client,
-            f"max_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+            f"max_over_time(((sum(irate(collection_store_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "ops_min": safe_scalar(
             client,
-            f"min_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+            f"min_over_time(((sum(irate(collection_store_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "p50": safe_scalar(
             client,
@@ -445,17 +449,20 @@ def gather_store_methods(
 def gather_rollout_outcomes(
     client: PrometheusClient,
     window: str,
-    rate_window: str,
+    window_seconds: int,
 ) -> List[RolloutOutcomeStats]:
     rate_map = vector_to_map(
-        safe_vector(client, f"sum by (status)(rate(collection_store_rollout_total[{rate_window}]))"),
+        safe_vector(
+            client,
+            f"(sum by (status)(increase(collection_store_rollout_total[{window}]))) / {window_seconds}",
+        ),
         ("status",),
     )
     p25_map = vector_to_map(
         safe_vector(
             client,
             f"histogram_quantile(0.25, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
         ),
         ("status",),
     )
@@ -463,7 +470,7 @@ def gather_rollout_outcomes(
         safe_vector(
             client,
             f"histogram_quantile(0.50, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
         ),
         ("status",),
     )
@@ -471,7 +478,7 @@ def gather_rollout_outcomes(
         safe_vector(
             client,
             f"histogram_quantile(0.75, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
         ),
         ("status",),
     )
@@ -479,7 +486,7 @@ def gather_rollout_outcomes(
         safe_vector(
             client,
             f"histogram_quantile(1.00, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
         ),
         ("status",),
     )
@@ -516,43 +523,58 @@ class HttpPathStats:
     p99: Optional[float]
 
 
+@dataclass
+class HttpPathStatusStats:
+    method: str
+    path: str
+    status_code: str
+    qps_mean: float
+    qps_max: Optional[float]
+    qps_min: Optional[float]
+    p50: Optional[float]
+    p95: Optional[float]
+    p99: Optional[float]
+
+
 def gather_http_paths(
     client: PrometheusClient,
     window: str,
-    rate_window: str,
+    window_seconds: int,
+    peak_window: str,
     subquery_step: str,
 ) -> Tuple[List[HttpPathStats], StatsSummary]:
-    qps_expr = f"sum by (method, path)(rate(http_requests_total[{rate_window}]))"
+    mean_expr = f"(sum by (method, path)(increase(http_requests_total[{window}]))) / {window_seconds}"
     qps_mean = vector_to_map(
-        safe_vector(client, f"avg_over_time(({qps_expr})[{window}:{subquery_step}])"),
+        safe_vector(client, mean_expr),
         ("method", "path"),
     )
+    peak_expr = f"sum by (method, path)(irate(http_requests_total[{peak_window}]))"
     qps_max = vector_to_map(
-        safe_vector(client, f"max_over_time(({qps_expr})[{window}:{subquery_step}])"),
+        safe_vector(client, f"max_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method", "path"),
     )
     qps_min = vector_to_map(
-        safe_vector(client, f"min_over_time(({qps_expr})[{window}:{subquery_step}])"),
+        safe_vector(client, f"min_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method", "path"),
     )
     p50 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.50, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.50, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
         ),
         ("method", "path"),
     )
     p95 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.95, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.95, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
         ),
         ("method", "path"),
     )
     p99 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.99, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.99, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
         ),
         ("method", "path"),
     )
@@ -594,33 +616,105 @@ def gather_http_paths(
     overall: StatsSummary = {
         "qps_mean": safe_scalar(
             client,
-            f"avg_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+            f"(sum(increase(http_requests_total[{window}]))) / {window_seconds}",
         )
         or 0.0,
         "qps_max": safe_scalar(
             client,
-            f"max_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+            f"max_over_time(((sum(irate(http_requests_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "qps_min": safe_scalar(
             client,
-            f"min_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+            f"min_over_time(((sum(irate(http_requests_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "p50": safe_scalar(
-            client, f"histogram_quantile(0.50, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.50, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
         )
         or 0.0,
         "p95": safe_scalar(
-            client, f"histogram_quantile(0.95, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.95, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
         )
         or 0.0,
         "p99": safe_scalar(
-            client, f"histogram_quantile(0.99, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.99, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
         )
         or 0.0,
     }
     return path_stats, overall
 
 
+def gather_http_paths_with_status(
+    client: PrometheusClient,
+    window: str,
+    window_seconds: int,
+    peak_window: str,
+    subquery_step: str,
+) -> List[HttpPathStatusStats]:
+    qps_expr = f"sum by (method, path, status_code)(irate(http_requests_total[{peak_window}]))"
+
+    def fetch_status_metric(expr: str) -> Dict[Tuple[str, str, str], Optional[float]]:
+        samples = safe_vector(client, expr)
+        typed: Dict[Tuple[str, str, str], Optional[float]] = {}
+        if not samples:
+            return typed
+        for sample in samples:
+            metric_obj = sample.get("metric", {})
+            metric_map: Mapping[str, Any]
+            if isinstance(metric_obj, Mapping):
+                metric_map = cast(Mapping[str, Any], metric_obj)
+            else:
+                metric_map = {}
+            method_raw = metric_map.get("method")
+            path_raw = metric_map.get("path")
+            status_raw = metric_map.get("status_code")
+            key = (
+                _normalize_label(method_raw),
+                _normalize_label(path_raw),
+                _normalize_label(status_raw),
+            )
+            typed[key] = _sample_value(sample)
+        return typed
+
+    def _normalize_label(value: Any) -> str:
+        if value is None:
+            return "-"
+        text = str(value)
+        return text if text else "-"
+
+    qps_mean_norm = fetch_status_metric(
+        f"(sum by (method, path, status_code)(increase(http_requests_total[{window}]))) / {window_seconds}"
+    )
+    qps_max_norm = fetch_status_metric(f"max_over_time(({qps_expr})[{window}:{subquery_step}])")
+    qps_min_norm = fetch_status_metric(f"min_over_time(({qps_expr})[{window}:{subquery_step}])")
+    p50_norm = fetch_status_metric(
+        f"histogram_quantile(0.50, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+    )
+    p95_norm = fetch_status_metric(
+        f"histogram_quantile(0.95, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+    )
+    p99_norm = fetch_status_metric(
+        f"histogram_quantile(0.99, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+    )
+
+    stats: List[HttpPathStatusStats] = []
+    for key in sorted(qps_mean_norm.keys()):
+        method_label, path_label, status_label = key
+        stats.append(
+            HttpPathStatusStats(
+                method_label,
+                path_label,
+                status_label,
+                qps_mean_norm.get(key, 0.0) or 0.0,
+                qps_max_norm.get(key),
+                qps_min_norm.get(key),
+                p50_norm.get(key),
+                p95_norm.get(key),
+                p99_norm.get(key),
+            )
+        )
+    return stats
+
+
 # ---------------------------------------------------------------------------
 # Part 4 – diagnostics
 # ---------------------------------------------------------------------------
@@ -635,6 +729,27 @@ def gather_diagnostics(client: PrometheusClient, window: str) -> Dict[str, Any]:
         ),
         ("operation",),
     )
+    diagnostics["mongo_latency_p50"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.50, sum by (le, operation)(rate(mongo_operation_duration_seconds_bucket{{operation!='ensure_collection'}}[{window}])))",
+        ),
+        ("operation",),
+    )
+    diagnostics["mongo_latency_p95"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.95, sum by (le, operation)(rate(mongo_operation_duration_seconds_bucket{{operation!='ensure_collection'}}[{window}])))",
+        ),
+        ("operation",),
+    )
+    diagnostics["mongo_latency_p99"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.99, sum by (le, operation)(rate(mongo_operation_duration_seconds_bucket{{operation!='ensure_collection'}}[{window}])))",
+        ),
+        ("operation",),
+    )
     opcounters_samples = safe_vector(client, f"sum by (legacy_op_type)(rate(mongodb_ss_opcounters[{window}]))")
     mongo_opcounters: Dict[str, float] = {}
     if opcounters_samples:
@@ -651,6 +766,31 @@ def gather_diagnostics(client: PrometheusClient, window: str) -> Dict[str, Any]:
                 mongo_opcounters[str(label or "-")] = value
     diagnostics["mongo_opcounters"] = mongo_opcounters
     diagnostics["mongo_connections"] = safe_scalar(client, "avg(mongodb_ss_connections{conn_type='current'})")
+    diagnostics["memory_lock_rate"] = vector_to_map(
+        safe_vector(client, f"sum by (collection)(rate(memory_collection_lock_rate_total[{window}]))"),
+        ("collection",),
+    )
+    diagnostics["memory_lock_p50"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.50, sum by (le, collection)(rate(memory_collection_lock_latency_seconds_bucket[{window}])))",
+        ),
+        ("collection",),
+    )
+    diagnostics["memory_lock_p95"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.95, sum by (le, collection)(rate(memory_collection_lock_latency_seconds_bucket[{window}])))",
+        ),
+        ("collection",),
+    )
+    diagnostics["memory_lock_p99"] = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.99, sum by (le, collection)(rate(memory_collection_lock_latency_seconds_bucket[{window}])))",
+        ),
+        ("collection",),
+    )
     diagnostics["cpu_usage"] = safe_scalar(client, f"1 - avg(rate(node_cpu_seconds_total{{mode='idle'}}[{window}]))")
     diagnostics["memory_total"] = safe_scalar(client, "avg(node_memory_MemTotal_bytes)")
     diagnostics["memory_available"] = safe_scalar(client, "avg(node_memory_MemAvailable_bytes)")
@@ -749,8 +889,9 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
         start = end - duration
     assert start is not None
     duration_seconds = max((end - start).total_seconds(), 1.0)
+    window_seconds = clamp_window_seconds(duration_seconds)
     window = format_window(duration_seconds)
-    rate_window = compute_rate_window(duration_seconds)
+    peak_window = compute_peak_window(duration_seconds)
     subquery_step = compute_subquery_step(duration_seconds)
 
     client = PrometheusClient(args.prom_url, timeout=args.timeout, default_time=end)
@@ -794,7 +935,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     )
 
     # Store internals
-    store_stats, store_overall = gather_store_methods(client, window, rate_window, subquery_step)
+    store_stats, store_overall = gather_store_methods(client, window, window_seconds, peak_window, subquery_step)
     store_rows: List[List[str]] = [
         [
             stat.method,
@@ -821,7 +962,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     )
     lines.extend(section("CollectionBasedLightningStore", store_lines))
 
-    rollout_outcomes = gather_rollout_outcomes(client, window, rate_window)
+    rollout_outcomes = gather_rollout_outcomes(client, window, window_seconds)
     rollout_rows = [
         [
             stat.status,
@@ -838,7 +979,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     )
 
     # HTTP traffic
-    http_paths, http_overall = gather_http_paths(client, window, rate_window, subquery_step)
+    http_paths, http_overall = gather_http_paths(client, window, window_seconds, peak_window, subquery_step)
     http_rows: List[List[str]] = [
         [
             stat.method,
@@ -865,16 +1006,60 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     )
     lines.extend(section("HTTP Endpoints", http_lines))
 
+    http_status_stats = gather_http_paths_with_status(client, window, window_seconds, peak_window, subquery_step)
+    http_status_rows: List[List[str]] = []
+    for stat in http_status_stats:
+        http_status_rows.append(
+            [
+                stat.method,
+                stat.path,
+                stat.status_code,
+                fmt_rate(stat.qps_mean),
+                fmt_rate(stat.qps_max),
+                fmt_rate(stat.qps_min),
+                fmt_latency(stat.p50),
+                fmt_latency(stat.p95),
+                fmt_latency(stat.p99),
+            ]
+        )
+    lines.extend(
+        section(
+            "HTTP Endpoints (by Status)",
+            render_table(
+                ["Method", "Path", "Status Code", "Mean Req/s", "Max Req/s", "Min Req/s", "P50", "P95", "P99"],
+                http_status_rows,
+            ),
+        )
+    )
+
     # Diagnostics
     diag = gather_diagnostics(client, window)
     diagnostics_blocks: List[List[str]] = []
 
     mongo_ops = cast(Dict[str, float], diag.get("mongo_ops", {}))
+    mongo_latency_p50 = cast(Dict[str, float], diag.get("mongo_latency_p50", {}))
+    mongo_latency_p95 = cast(Dict[str, float], diag.get("mongo_latency_p95", {}))
+    mongo_latency_p99 = cast(Dict[str, float], diag.get("mongo_latency_p99", {}))
+    mongo_op_keys = sorted(
+        {
+            *mongo_ops.keys(),
+            *mongo_latency_p50.keys(),
+            *mongo_latency_p95.keys(),
+            *mongo_latency_p99.keys(),
+        },
+        key=str,
+    )
     mongo_ops_rows = [
-        [operation or "-", fmt_rate(rate)]
-        for operation, rate in sorted(mongo_ops.items(), key=lambda item: str(item[0]))
+        [
+            op or "-",
+            fmt_rate(mongo_ops.get(op)),
+            fmt_latency(mongo_latency_p50.get(op)),
+            fmt_latency(mongo_latency_p95.get(op)),
+            fmt_latency(mongo_latency_p99.get(op)),
+        ]
+        for op in mongo_op_keys
     ]
-    diagnostics_blocks.append(render_table(["Mongo Operation", "Ops/s"], mongo_ops_rows))
+    diagnostics_blocks.append(render_table(["Mongo Operation", "Ops/s", "P50", "P95", "P99"], mongo_ops_rows))
 
     mongo_opcounters = cast(Dict[str, float], diag.get("mongo_opcounters", {}))
     mongo_opcounters_rows = [
@@ -883,6 +1068,31 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     ]
     diagnostics_blocks.append(render_table(["MongoDB Opcounter", "Ops/s"], mongo_opcounters_rows))
 
+    memory_lock_rate = cast(Dict[str, float], diag.get("memory_lock_rate", {}))
+    memory_lock_p50 = cast(Dict[str, float], diag.get("memory_lock_p50", {}))
+    memory_lock_p95 = cast(Dict[str, float], diag.get("memory_lock_p95", {}))
+    memory_lock_p99 = cast(Dict[str, float], diag.get("memory_lock_p99", {}))
+    memory_collections = sorted(
+        {
+            *memory_lock_rate.keys(),
+            *memory_lock_p50.keys(),
+            *memory_lock_p95.keys(),
+            *memory_lock_p99.keys(),
+        },
+        key=str,
+    )
+    memory_lock_rows = [
+        [
+            collection or "-",
+            fmt_rate(memory_lock_rate.get(collection)),
+            fmt_latency(memory_lock_p50.get(collection)),
+            fmt_latency(memory_lock_p95.get(collection)),
+            fmt_latency(memory_lock_p99.get(collection)),
+        ]
+        for collection in memory_collections
+    ]
+    diagnostics_blocks.append(render_table(["Memory Collection", "Locks/s", "P50", "P95", "P99"], memory_lock_rows))
+
     mongo_misc_rows: List[List[str]] = []
     if diag.get("mongo_connections") is not None:
         mongo_misc_rows.append(["MongoDB connections (avg)", f"{diag['mongo_connections']:.2f}"])

tests/benchmark/micro_benchmark.py

@@ -0,0 +1,238 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Micro benchmarks for the store."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import multiprocessing
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Sequence
+
+from rich.console import Console
+
+import agentlightning as agl
+from agentlightning.types.tracer import OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.utils.system_snapshot import system_snapshot
+
+console = Console()
+
+
+def _close_store_client(store: agl.LightningStoreClient) -> None:
+    try:
+        asyncio.run(store.close())
+    except Exception:
+        pass
+
+
+def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
+    trace_hex = f"{sequence_id:032x}"
+    span_hex = f"{sequence_id:016x}"
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=trace_hex,
+        span_id=span_hex,
+        parent_id=None,
+        name=name,
+        status=TraceStatus(status_code="OK"),
+        attributes={},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}),
+        parent=None,
+        resource=OtelResource(attributes={}, schema_url=""),
+    )
+
+
+@dataclass
+class BenchmarkSummary:
+    mode: str
+    total_tasks: int
+    successes: int
+    duration: float
+
+    @property
+    def success_rate(self) -> float:
+        if self.total_tasks == 0:
+            return 0.0
+        return self.successes / self.total_tasks
+
+    @property
+    def throughput(self) -> float:
+        if self.duration <= 0:
+            return 0.0
+        return self.successes / self.duration
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Micro benchmarks for the store.")
+    parser.add_argument("--store-url", default="http://localhost:4747", help="Lightning Store endpoint base URL.")
+    parser.add_argument("--summary-file", help="File to append final benchmark summary.")
+    parser.add_argument(
+        "mode",
+        choices=("worker", "dequeue-empty", "rollout"),
+        help="Mode to exercise different operations.",
+    )
+    args = parser.parse_args(argv)
+    return args
+
+
+def _update_worker_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"Updating worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+    try:
+        asyncio.run(store.update_worker(worker_id, system_snapshot()))
+        return True
+    except Exception as e:
+        console.print(f"Error updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_many_update_workers(store_url: str) -> BenchmarkSummary:
+    """Simulate many update workers."""
+
+    start_time = time.time()
+
+    # Use a multiprocessing pool to update workers.
+    worker_ids = [(f"Worker-{i % 1024}", f"Task-{j}") for i in range(1024) for j in range(10)]
+    with multiprocessing.get_context("fork").Pool(processes=1024) as pool:
+        successful_tasks = pool.map(_update_worker_task, [(store_url, *worker_id) for worker_id in worker_ids])
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Success rate: {successes / len(worker_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} workers/second")
+    return BenchmarkSummary(mode="worker", total_tasks=len(worker_ids), successes=successes, duration=duration)
+
+
+def _dequeue_empty_and_update_workers_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"Dequeueing empty and updating worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> None:
+        await store.dequeue_rollout(worker_id=worker_id)
+        await store.update_worker(worker_id, system_snapshot())
+
+    try:
+        asyncio.run(_async_task())
+        return True
+    except Exception as e:
+        console.print(f"Error dequeueing empty and updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_dequeue_empty_and_update_workers(store_url: str) -> BenchmarkSummary:
+    """Simulate dequeue empty and update workers."""
+    start_time = time.time()
+
+    worker_ids = [(f"Worker-{i % 1024}", f"Task-{j}") for i in range(1024) for j in range(10)]
+    with multiprocessing.get_context("fork").Pool(processes=1024) as pool:
+        successful_tasks = pool.map(
+            _dequeue_empty_and_update_workers_task, [(store_url, *worker_id) for worker_id in worker_ids]
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Success rate: {successes / len(worker_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} workers/second")
+    return BenchmarkSummary(mode="dequeue-empty", total_tasks=len(worker_ids), successes=successes, duration=duration)
+
+
+def _rollout_flow_task(args: tuple[str, int, int]) -> bool:
+    store_url, task_id, spans_per_attempt = args
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> None:
+        console.print(f"Starting rollout for task {task_id} with {spans_per_attempt} spans")
+        attempted = await store.start_rollout(input={"task": task_id})
+        rollout_id = attempted.rollout_id
+        attempt_id = attempted.attempt.attempt_id
+        for seq in range(1, spans_per_attempt + 1):
+            console.print(f"Adding span {seq} for task {task_id} with {spans_per_attempt} spans")
+            span = _make_span(
+                rollout_id,
+                attempt_id,
+                task_id * spans_per_attempt + seq,
+                f"micro-span-{seq}",
+            )
+            await store.add_span(span)
+        console.print(f"Updating attempt {attempt_id} for task {task_id} with {spans_per_attempt} spans")
+        await store.update_attempt(rollout_id, attempt_id, status="succeeded")
+
+    try:
+        asyncio.run(_async_task())
+        return True
+    except Exception as e:
+        console.print(f"Error running rollout task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_rollout_with_spans(store_url: str, spans_per_attempt: int = 4) -> BenchmarkSummary:
+    """Simulate full rollout lifecycle with spans."""
+    start_time = time.time()
+    task_ids = list(range(1024 * 4))
+    with multiprocessing.get_context("fork").Pool(processes=256) as pool:
+        successful_tasks = pool.map(
+            _rollout_flow_task, [(store_url, task_id, spans_per_attempt) for task_id in task_ids]
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Rollout success rate: {successes / len(task_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(mode="rollout", total_tasks=len(task_ids), successes=successes, duration=duration)
+
+
+def record_summary(summary: BenchmarkSummary, summary_file: Optional[str]) -> None:
+    message = (
+        f"[summary] mode={summary.mode} success_rate={summary.success_rate:.3f} "
+        f"throughput={summary.throughput:.3f} ops/s duration={summary.duration:.3f}s "
+        f"success={summary.successes}/{summary.total_tasks}"
+    )
+    console.print(message)
+    if summary_file:
+        path = Path(summary_file)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", encoding="utf-8") as fh:
+            fh.write(message + "\n")
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    if args.mode == "worker":
+        summary = simulate_many_update_workers(args.store_url)
+    elif args.mode == "dequeue-empty":
+        summary = simulate_dequeue_empty_and_update_workers(args.store_url)
+    elif args.mode == "rollout":
+        summary = simulate_rollout_with_spans(args.store_url)
+    else:
+        raise ValueError(f"Invalid mode: {args.mode}")
+    record_summary(summary, args.summary_file)
+
+
+if __name__ == "__main__":
+    main()

tests/emitter/test_operation.py

@@ -0,0 +1,396 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+import opentelemetry.trace as trace_api
+import pytest
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import Status, StatusCode
+
+import agentlightning.emitter.annotation as annotation_module
+from agentlightning.emitter.annotation import _safe_json_dump  # pyright: ignore[reportPrivateUsage]
+from agentlightning.emitter.annotation import (
+    OperationContext,
+    emit_annotation,
+    operation,
+)
+from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
+from agentlightning.utils.otel import extract_links_from_attributes, make_link_attributes, query_linked_spans
+
+
+class RecordingSpan:
+    def __init__(self) -> None:
+        self.attributes: Dict[str, Any] = {}
+        self.recorded_exceptions: List[BaseException] = []
+        self.statuses: List[Status] = []
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        self.attributes[key] = value
+
+    def record_exception(self, exc: BaseException) -> None:
+        self.recorded_exceptions.append(exc)
+
+    def set_status(self, status: Status) -> None:
+        self.statuses.append(status)
+
+
+class DummySpanContextManager:
+    def __init__(self, span: RecordingSpan) -> None:
+        self.span = span
+        self.exit_calls: List[
+            Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
+        ] = []
+
+    def __enter__(self) -> RecordingSpan:
+        return self.span
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> bool:
+        self.exit_calls.append((exc_type, exc_val, exc_tb))
+        return False
+
+
+class DummyTracer:
+    def __init__(self, start_span_instance: Optional[RecordingSpan] = None) -> None:
+        self._start_span_instance = start_span_instance
+        self.start_span_calls: List[Tuple[str, Dict[str, Any]]] = []
+        self.start_as_current_span_calls: List[Tuple[str, Dict[str, Any], RecordingSpan]] = []
+
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> RecordingSpan:
+        span = self._start_span_instance or RecordingSpan()
+        self.start_span_calls.append((name, dict(attributes or {})))
+        return span
+
+    def start_as_current_span(
+        self,
+        name: str,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> DummySpanContextManager:
+        span = RecordingSpan()
+        self.start_as_current_span_calls.append((name, dict(attributes or {}), span))
+        return DummySpanContextManager(span)
+
+
+class DummyUseSpan:
+    def __init__(self) -> None:
+        self.calls: List[Tuple[RecordingSpan, bool]] = []
+        self.exit_calls: List[
+            Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
+        ] = []
+
+    def __call__(self, span: RecordingSpan, end_on_exit: bool) -> DummyUseSpan:
+        self.calls.append((span, end_on_exit))
+        self._span = span
+        return self
+
+    def __enter__(self) -> None:
+        return None
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> bool:
+        self.exit_calls.append((exc_type, exc_val, exc_tb))
+        return False
+
+
+@dataclass
+class ComplexResult:
+    values: Tuple[int, ...]
+    marker: str
+
+
+def test_safe_json_dump_handles_recursive_structures() -> None:
+    payload: List[Any] = []
+    payload.append(payload)
+
+    assert _safe_json_dump(payload) == "[[...]]"
+
+
+def test_operation_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    ctx = OperationContext("custom-span", {"meta": {"foo": 1}, "count": 2})
+
+    with ctx as op:
+        op.set_input({"payload": 1}, flag=True)
+        op.set_output({"success": True})
+
+    assert tracer.start_span_calls
+    start_name, start_attributes = tracer.start_span_calls[0]
+    assert start_name == "custom-span"
+    assert json.loads(start_attributes["meta"]) == {"foo": 1}
+    assert start_attributes["count"] == 2
+
+    assert json.loads(span.attributes["input.args"]) == [{"payload": 1}]
+    assert span.attributes["input.flag"] == "true"
+    assert json.loads(span.attributes["output"]) == {"success": True}
+    assert use_span.calls == [(span, True)]
+
+
+def test_operation_context_set_input_supports_multiple_values(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    ctx = OperationContext("ctx", {})
+
+    with ctx as op:
+        op.set_input(1, 2, data={"foo": ["bar"]}, flags=[True, False])
+
+    assert json.loads(span.attributes["input.args"]) == [1, 2]
+    assert json.loads(span.attributes["input.data"]) == {"foo": ["bar"]}
+    assert json.loads(span.attributes["input.flags"]) == [True, False]
+
+
+def test_operation_context_records_non_serializable_output(monkeypatch: pytest.MonkeyPatch) -> None:
+    class Unserializable:
+        def __str__(self) -> str:
+            return "<Unserializable>"
+
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    ctx = OperationContext("ctx", {})
+
+    with ctx as op:
+        op.set_output(Unserializable())
+
+    assert json.loads(span.attributes["output"]) == "<Unserializable>"
+
+
+def test_operation_context_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    ctx = OperationContext("custom-span", {})
+
+    with pytest.raises(RuntimeError):
+        with ctx:
+            raise RuntimeError("boom")
+
+    assert isinstance(span.recorded_exceptions[0], RuntimeError)
+    status = span.statuses[-1]
+    assert status.status_code == StatusCode.ERROR
+    assert status.description == "boom"
+    assert use_span.exit_calls[-1][1].args == ("boom",)  # type: ignore
+
+
+def test_operation_factory_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    with operation(tags=["one", "two"]) as ctx:
+        ctx.set_input("alpha", meta={"score": 0.5})
+        ctx.set_output(["beta", "gamma"])
+
+    start_name, attrs = tracer.start_span_calls[0]
+    assert start_name == AGL_OPERATION
+    assert json.loads(attrs["tags"]) == ["one", "two"]
+    assert json.loads(span.attributes["input.args"]) == ["alpha"]
+    assert json.loads(span.attributes["input.meta"]) == {"score": 0.5}
+    assert json.loads(span.attributes["output"]) == ["beta", "gamma"]
+
+
+def test_operation_factory_uses_standard_span_name(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = RecordingSpan()
+    tracer = DummyTracer(start_span_instance=span)
+    use_span = DummyUseSpan()
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    with operation(user={"id": 5}) as ctx:
+        ctx.set_output("done")
+
+    assert tracer.start_span_calls
+    start_name, attrs = tracer.start_span_calls[0]
+    assert start_name == AGL_OPERATION
+    assert json.loads(attrs["user"]) == {"id": 5}
+
+
+def test_operation_rejects_custom_span_names() -> None:
+    with pytest.raises(ValueError):
+        operation("custom-name")  # type: ignore
+
+
+def test_operation_decorator_sync_records_span_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
+    tracer = DummyTracer()
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+
+    @operation(category={"kind": "combine"})
+    def combine(data: Dict[str, int], *, meta: Dict[str, str]) -> Dict[str, Any]:
+        return {"joined": {**data, **meta}}
+
+    result = combine({"value": 1}, meta={"source": "unit"})
+
+    assert result == {"joined": {"value": 1, "source": "unit"}}
+    assert tracer.start_as_current_span_calls
+    span_name, span_attributes, span = tracer.start_as_current_span_calls[0]
+    assert span_name == AGL_OPERATION
+    assert json.loads(span_attributes["category"]) == {"kind": "combine"}
+
+    input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
+    assert json.loads(span.attributes[f"{input_prefix}.data"]) == {"value": 1}
+    assert json.loads(span.attributes[f"{input_prefix}.meta"]) == {"source": "unit"}
+    assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "combine"
+    assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
+
+
+def test_operation_decorator_handles_complex_signature(monkeypatch: pytest.MonkeyPatch) -> None:
+    tracer = DummyTracer()
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+
+    @operation()
+    def complicated(
+        first: int,
+        /,
+        required: str,
+        default: int = 5,
+        *extra: int,
+        kwonly: str,
+        kwdefault: str = "fallback",
+        **rest: Any,
+    ) -> ComplexResult:
+        return ComplexResult(values=(first, len(extra), len(rest)), marker=kwonly + kwdefault + required)
+
+    result = complicated(1, "req", 7, 8, 9, kwonly="x", kwdefault="y", tag="value")
+
+    span = tracer.start_as_current_span_calls[0][2]
+    input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
+
+    assert json.loads(span.attributes[f"{input_prefix}.first"]) == 1
+    assert json.loads(span.attributes[f"{input_prefix}.required"]) == "req"
+    assert json.loads(span.attributes[f"{input_prefix}.default"]) == 7
+    assert json.loads(span.attributes[f"{input_prefix}.extra"]) == [8, 9]
+    assert json.loads(span.attributes[f"{input_prefix}.kwonly"]) == "x"
+    assert json.loads(span.attributes[f"{input_prefix}.kwdefault"]) == "y"
+    assert json.loads(span.attributes[f"{input_prefix}.rest"]) == {"tag": "value"}
+    assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "complicated"
+    assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == str(result)
+
+
+def test_operation_decorator_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
+    tracer = DummyTracer()
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+
+    @operation()
+    def fail(value: int) -> int:
+        raise ValueError("bad input")
+
+    with pytest.raises(ValueError):
+        fail(1)
+
+    span = tracer.start_as_current_span_calls[0][2]
+    assert isinstance(span.recorded_exceptions[0], ValueError)
+    status = span.statuses[-1]
+    assert status.status_code == StatusCode.ERROR
+    assert status.description == "bad input"
+
+
+@pytest.mark.asyncio()
+async def test_operation_async_wrapper_records_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
+    tracer = DummyTracer()
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+
+    @operation()
+    async def echo(payload: Dict[str, Any]) -> Dict[str, Any]:
+        return {"payload": payload}
+
+    result = await echo({"value": 3})
+
+    assert result == {"payload": {"value": 3}}
+    span = tracer.start_as_current_span_calls[0][2]
+    prefix = LightningSpanAttributes.OPERATION_INPUT.value
+    assert json.loads(span.attributes[f"{prefix}.payload"]) == {"value": 3}
+    assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
+
+
+def test_operation_span_can_be_resolved_via_annotation_links(monkeypatch: pytest.MonkeyPatch) -> None:
+    provider = TracerProvider()
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = provider.get_tracer(__name__)
+
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+    monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
+
+    @operation(conversation_id="conv-1")
+    def decorated(value: int) -> int:
+        return value + 1
+
+    assert decorated(41) == 42
+
+    spans = exporter.get_finished_spans()
+    operation_span = next(span for span in spans if span.name == AGL_OPERATION)
+    assert operation_span.attributes["conversation_id"] == "conv-1"  # type: ignore
+
+    trace_id_hex = trace_api.format_trace_id(operation_span.context.trace_id)  # type: ignore
+    span_id_hex = trace_api.format_span_id(operation_span.context.span_id)  # type: ignore
+    link_attrs = make_link_attributes({"trace_id": trace_id_hex, "span_id": span_id_hex})
+
+    emit_annotation({**link_attrs, "note": "operation-follow-up"})
+
+    spans = exporter.get_finished_spans()
+    annotation_span = next(span for span in spans if span.name == AGL_ANNOTATION)
+    annotation_links = extract_links_from_attributes(dict(annotation_span.attributes or {}))
+
+    matches = query_linked_spans([operation_span], annotation_links)
+    assert matches == [operation_span]
+
+
+def test_operation_honors_propagate_flag(monkeypatch: pytest.MonkeyPatch) -> None:
+    tracer = DummyTracer()
+    flags: List[bool] = []
+    use_span = DummyUseSpan()
+
+    def fake_get_tracer(use_active_span_processor: bool = True) -> DummyTracer:
+        flags.append(use_active_span_processor)
+        return tracer
+
+    monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
+    monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
+
+    @operation(propagate=False)
+    def decorated(value: int) -> int:
+        return value
+
+    assert decorated(7) == 7
+
+    with operation(propagate=False):
+        pass
+
+    assert flags == [False, False]

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -38,6 +38,8 @@ try:
     GPU_AVAILABLE = torch.cuda.is_available()
 except Exception:
     GPU_AVAILABLE = False  # type: ignore
+
+if not GPU_AVAILABLE:
     pytest.skip(reason="GPU not available", allow_module_level=True)
 
 

tests/runner/test_agent_runner.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import logging
 import random
 from contextlib import asynccontextmanager
 from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Sequence, Tuple, cast
@@ -264,6 +265,89 @@ async def test_step_emits_reward_for_float_result() -> None:
     assert rewards == [0.75]
 
 
+@pytest.mark.asyncio
+async def test_step_emits_reward_for_bool_result(caplog: pytest.LogCaptureFixture) -> None:
+    class BoolRewardAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> bool:
+            return True
+
+    agent = BoolRewardAgent()
+    runner, store, _ = await setup_runner(agent)
+    caplog.set_level(logging.WARNING)
+    try:
+        await runner.step({"prompt": "hello"})
+    finally:
+        teardown_runner(runner)
+
+    assert "Reward is not a number" in caplog.text
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    rewards = [span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION]
+    assert rewards == [1.0]
+
+
+@pytest.mark.asyncio
+async def test_step_raises_for_invalid_result_type() -> None:
+    class InvalidResultAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(  # type: ignore[reportIncompatibleMethodOverride]
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> Dict[str, Any]:
+            return {"unexpected": True}
+
+    agent = InvalidResultAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        with pytest.raises(TypeError, match="Invalid raw result type"):
+            await runner.step({"task": "bad-result"})
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert len(rollouts) == 1
+    attempts = await store.query_attempts(rollouts[0].rollout_id)
+    assert attempts[-1].status == "failed"
+
+
+@pytest.mark.asyncio
+async def test_readable_spans_return_skip_store_when_tracer_is_otel(monkeypatch: pytest.MonkeyPatch) -> None:
+    class DummyTracerWithOtel(DummyTracer):
+        pass
+
+    class RecordingStore(InMemoryLightningStore):
+        def __init__(self) -> None:
+            super().__init__()
+            self.add_otel_span_calls = 0
+
+        async def add_otel_span(
+            self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
+        ) -> Span | None:
+            self.add_otel_span_calls += 1
+            return await super().add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+
+    monkeypatch.setattr("agentlightning.runner.agent.OtelTracer", DummyTracerWithOtel)
+
+    store = RecordingStore()
+    await store.update_resources("default", {"llm": LLM(endpoint="http://localhost", model="dummy")})
+
+    tracer = DummyTracerWithOtel()
+    runner = LitAgentRunner[Any](tracer=tracer)
+    agent = HeartbeatAgent()
+    runner.init(agent)
+    runner.init_worker(worker_id=0, store=store)
+
+    attempted_rollout = await store.start_rollout(input={"task": "otel"}, mode="val")
+    spans = [create_readable_span("otel-span")]
+
+    result_spans = await runner._post_process_rollout_result(  # pyright: ignore[reportPrivateUsage]
+        attempted_rollout, spans
+    )
+
+    assert result_spans == spans
+    assert store.add_otel_span_calls == 0
+
+    teardown_runner(runner)
+
+
 @pytest.mark.asyncio
 async def test_step_handles_non_llm_resource() -> None:
     class PromptAgent(LitAgent[str]):
@@ -706,6 +790,72 @@ async def test_step_with_custom_resources_returns_rollout() -> None:
     assert result.resources_id is not None
 
 
+@pytest.mark.asyncio
+async def test_step_registers_worker_id_on_start_rollout(monkeypatch: pytest.MonkeyPatch) -> None:
+    """runner.step should pass the formatted worker ID down to the store."""
+
+    class WorkerAwareAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
+            return 1.0
+
+    agent = WorkerAwareAgent()
+    runner, store, _ = await setup_runner(agent)
+
+    expected_worker_label = runner.get_worker_id()
+    captured: Dict[str, Optional[str]] = {}
+    original_start_rollout = store.start_rollout
+
+    async def wrapped_start_rollout(*args: Any, **kwargs: Any):
+        captured["worker_id"] = kwargs.get("worker_id")
+        return await original_start_rollout(*args, **kwargs)
+
+    monkeypatch.setattr(store, "start_rollout", wrapped_start_rollout)
+
+    try:
+        await runner.step({"task": "worker-aware"})
+    finally:
+        teardown_runner(runner)
+
+    assert captured["worker_id"] == expected_worker_label
+
+
+@pytest.mark.asyncio
+async def test_iter_passes_worker_id_to_dequeue(monkeypatch: pytest.MonkeyPatch) -> None:
+    """iter() should poll the store with the formatted worker identifier."""
+
+    class IdleAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> float:  # pragma: no cover - not invoked
+            return 0.0
+
+    agent = IdleAgent()
+    runner, store, _ = await setup_runner(agent, poll_interval=0.01)
+
+    expected_worker_label = runner.get_worker_id()
+    captured: Dict[str, Optional[str]] = {}
+    event = ThreadingEvent()
+
+    async def fake_dequeue(*, worker_id: Optional[str] = None):
+        captured["worker_id"] = worker_id
+        event.set()
+        return None
+
+    async def fast_sleep(self: LitAgentRunner[Any], event: Optional[ExecutionEvent] = None) -> None:
+        if event is not None:
+            event.set()
+
+    monkeypatch.setattr(store, "dequeue_rollout", fake_dequeue)
+    monkeypatch.setattr(LitAgentRunner, "_sleep_until_next_poll", fast_sleep)
+
+    try:
+        await runner.iter(event=event)
+    finally:
+        teardown_runner(runner)
+
+    assert captured["worker_id"] == expected_worker_label
+
+
 @pytest.mark.asyncio
 async def test_emit_heartbeat_updates_worker_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
     snapshot = {"cpu_pct": 42.0, "mem_pct": 10.5}

tests/store/dummy_store.py

@@ -10,6 +10,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -42,8 +43,9 @@ class DummyLightningStore(LightningStore):
         resources_id: Optional[str] = None,
         config: Optional[RolloutConfig] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
-        self.calls.append(("start_rollout", (input, mode, resources_id, config, metadata), {}))
+        self.calls.append(("start_rollout", (input, mode, resources_id, config, metadata, worker_id), {}))
         return self.return_values["start_rollout"]
 
     async def enqueue_rollout(
@@ -57,12 +59,25 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("enqueue_rollout", (input, mode, resources_id, config, metadata), {}))
         return self.return_values["enqueue_rollout"]
 
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        self.calls.append(("enqueue_many_rollouts", (rollouts,), {}))
+        return self.return_values["enqueue_many_rollouts"]
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         self.calls.append(("dequeue_rollout", (worker_id,), {}))
         return self.return_values["dequeue_rollout"]
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        self.calls.append(("start_attempt", (rollout_id,), {}))
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        self.calls.append(("dequeue_many_rollouts", (), {"limit": limit, "worker_id": worker_id}))
+        return self.return_values["dequeue_many_rollouts"]
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        self.calls.append(("start_attempt", (rollout_id, worker_id), {}))
         return self.return_values["start_attempt"]
 
     async def query_rollouts(self, *args: Any, **kwargs: Any) -> List[Rollout]:

tests/store/test_client_server.py

@@ -18,7 +18,16 @@ from yarl import URL
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.types import LLM, OtelResource, PaginatedResult, PromptTemplate, RolloutConfig, Span, TraceStatus
+from agentlightning.types import (
+    LLM,
+    EnqueueRolloutRequest,
+    OtelResource,
+    PaginatedResult,
+    PromptTemplate,
+    RolloutConfig,
+    Span,
+    TraceStatus,
+)
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncherArgs
 
 
@@ -159,6 +168,188 @@ async def test_server_client_statistics_match(server_client: Tuple[LightningStor
     assert server_stats["total_rollouts"] >= 1  # type: ignore
 
 
+@pytest.mark.asyncio
+async def test_client_start_rollout_propagates_worker_id(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    server, client = server_client
+    attempt = await client.start_rollout(input={"source": "remote-worker"}, worker_id="client-worker-start")
+
+    assert attempt.attempt.worker_id == "client-worker-start"
+    worker = await server.get_worker_by_id("client-worker-start")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == attempt.rollout_id
+    assert worker.current_attempt_id == attempt.attempt.attempt_id
+
+
+@pytest.mark.asyncio
+async def test_client_start_attempt_propagates_worker_id(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    server, client = server_client
+    initial = await client.start_rollout(input={"source": "retry-worker"})
+    retry = await client.start_attempt(initial.rollout_id, worker_id="client-worker-retry")
+
+    assert retry.attempt.sequence_id == 2
+    assert retry.attempt.worker_id == "client-worker-retry"
+    worker = await server.get_worker_by_id("client-worker-retry")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == retry.rollout_id
+    assert worker.current_attempt_id == retry.attempt.attempt_id
+
+
+@pytest.mark.asyncio
+async def test_client_enqueue_many_rollouts_uses_batch_payload(monkeypatch: MonkeyPatch) -> None:
+    client = LightningStoreClient("http://localhost:9000")
+    captured: Dict[str, Any] = {}
+
+    async def fake_request_json(_, method: str, path: str, *, json: Any = None, params: Any = None) -> Any:
+        captured.update({"method": method, "path": path, "json": json})
+        count = len(json["rollouts"]) if json and "rollouts" in json else 0  # type: ignore[index]
+        return [{"rollout_id": f"bulk-{idx}", "input": {"idx": idx}, "start_time": float(idx)} for idx in range(count)]
+
+    monkeypatch.setattr(LightningStoreClient, "_request_json", fake_request_json, raising=False)  # type: ignore
+
+    requests = [
+        EnqueueRolloutRequest(input={"idx": 0}, mode="train", metadata={"batch": "left"}),
+        EnqueueRolloutRequest(input={"idx": 1}, resources_id="resources-1"),
+    ]
+    rollouts = await client.enqueue_many_rollouts(requests)
+
+    assert captured["method"] == "post"
+    assert captured["path"] == "/queues/rollouts/enqueue"
+    assert len(captured["json"]["rollouts"]) == 2  # type: ignore[index]
+    assert captured["json"]["rollouts"][0]["mode"] == "train"  # type: ignore[index]
+    assert captured["json"]["rollouts"][1]["resources_id"] == "resources-1"  # type: ignore[index]
+    assert len(rollouts) == 2
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_client_dequeue_methods_share_batch_logic(monkeypatch: MonkeyPatch) -> None:
+    client = LightningStoreClient("http://localhost:9001")
+
+    def attempt_payload(idx: int) -> Dict[str, Any]:
+        attempt_id = f"attempt-{idx}"
+        rollout_id = f"rollout-{idx}"
+        return {
+            "rollout_id": rollout_id,
+            "input": {"idx": idx},
+            "start_time": float(idx),
+            "status": "preparing",
+            "attempt": {
+                "rollout_id": rollout_id,
+                "attempt_id": attempt_id,
+                "sequence_id": 1,
+                "start_time": float(idx),
+                "status": "preparing",
+                "worker_id": "batch-worker",
+            },
+        }
+
+    payload_queue = [
+        [attempt_payload(0), attempt_payload(1)],
+        [attempt_payload(0)],
+    ]
+
+    class FakeResponse:
+        def __init__(self, body: Any):
+            self._body = body
+            self.status = 200
+
+        async def __aenter__(self) -> "FakeResponse":
+            return self
+
+        async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+            return None
+
+        def raise_for_status(self) -> None:
+            return None
+
+        async def json(self) -> Any:
+            return self._body
+
+    class RecordingSession:
+        def __init__(self) -> None:
+            self.calls: list[Dict[str, Any]] = []
+
+        def post(self, url: str, json: Dict[str, Any]) -> FakeResponse:
+            self.calls.append({"url": url, "json": json})
+            body = payload_queue.pop(0)
+            return FakeResponse(body)
+
+    session = RecordingSession()
+
+    async def fake_get_session() -> RecordingSession:
+        return session
+
+    monkeypatch.setattr(client, "_get_session", fake_get_session)
+
+    batch = await client.dequeue_many_rollouts(limit=2, worker_id="batch-worker")
+    assert len(batch) == 2
+    single = await client.dequeue_rollout(worker_id="batch-worker")
+    assert single is not None
+    assert [call["json"]["limit"] for call in session.calls] == [2, 1]
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_client_dequeue_many_rollouts_skips_network_for_non_positive_limit(monkeypatch: MonkeyPatch) -> None:
+    client = LightningStoreClient("http://localhost:9002")
+
+    async def fail_get_session() -> None:
+        pytest.fail("Client should not request a session when limit <= 0")
+
+    monkeypatch.setattr(client, "_get_session", fail_get_session)
+
+    assert await client.dequeue_many_rollouts(limit=0, worker_id="idle") == []
+    assert await client.dequeue_many_rollouts(limit=-5) == []
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_client_concurrent_enqueue_many_rollouts(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    _, client = server_client
+
+    async def enqueue_batch(batch_idx: int) -> list[str]:
+        requests = [EnqueueRolloutRequest(input={"batch": batch_idx, "idx": item}) for item in range(3)]
+        rollouts = await client.enqueue_many_rollouts(requests)
+        return [rollout.rollout_id for rollout in rollouts]
+
+    batches = await asyncio.gather(*(enqueue_batch(batch_idx) for batch_idx in range(5)))
+    all_ids = {rollout_id for batch in batches for rollout_id in batch}
+    assert len(all_ids) == 15
+
+    queried = await client.query_rollouts(limit=-1)
+    assert isinstance(queried, PaginatedResult)
+    assert queried.total >= 15
+
+
+@pytest.mark.asyncio
+async def test_client_concurrent_dequeue_many_rollouts(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    server, client = server_client
+    requests = [EnqueueRolloutRequest(input={"idx": idx}) for idx in range(6)]
+    # Seed queue from the server to avoid races with background processing
+    await asyncio.gather(*(server.enqueue_rollout(**req.model_dump()) for req in requests))
+
+    async def consume(limit: int, worker: str):
+        return await client.dequeue_many_rollouts(limit=limit, worker_id=worker)
+
+    batches = await asyncio.gather(
+        consume(3, "worker-a"),
+        consume(3, "worker-b"),
+    )
+    claimed_ids = {attempt.rollout_id for batch in batches for attempt in batch}
+    assert len(claimed_ids) == 6
+    assert await client.dequeue_many_rollouts(limit=1) == []
+
+
 @pytest.mark.asyncio
 async def test_add_resources_via_server(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
     """Test that add_resources works correctly via server."""
@@ -313,7 +504,7 @@ async def test_client_server_end_to_end(
     dequeued = await server.dequeue_rollout(worker_id=server_worker_id)
     server_worker_after_dequeue = await server.get_worker_by_id(server_worker_id)
     assert server_worker_after_dequeue is not None
-    assert server_worker_after_dequeue.status == "idle"
+    assert server_worker_after_dequeue.status == "busy"  # should be busy after dequeue
     assert server_worker_after_dequeue.last_dequeue_time is not None
     dequeue_time = server_worker_after_dequeue.last_dequeue_time
     started_attempt = await server.start_attempt(queued_rollout.rollout_id)
@@ -391,7 +582,7 @@ async def test_client_server_end_to_end(
     assert dequeued_client is not None
     client_worker_after_dequeue = await client.get_worker_by_id(client_worker_id)
     assert client_worker_after_dequeue is not None
-    assert client_worker_after_dequeue.status == "idle"
+    assert client_worker_after_dequeue.status == "busy"  # should be busy after dequeue
     assert client_worker_after_dequeue.last_dequeue_time is not None
     client_dequeue_time = client_worker_after_dequeue.last_dequeue_time
     started_client_attempt = await client.start_attempt(dequeued_client.rollout_id)

tests/store/test_collection_memory.py

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

tests/store/test_implementation.py

@@ -29,7 +29,9 @@ from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
     LLM,
+    Attempt,
     AttemptedRollout,
+    EnqueueRolloutRequest,
     Event,
     Link,
     OtelResource,
@@ -674,6 +676,49 @@ async def test_requeue_mechanism(store_fixture: LightningStore) -> None:
     assert latest_attempt.sequence_id == 2
 
 
+@pytest.mark.asyncio
+async def test_enqueue_many_rollouts_preserves_order(store_fixture: LightningStore) -> None:
+    """enqueue_many_rollouts should enqueue tasks in the provided order with matching metadata."""
+
+    requests = [
+        EnqueueRolloutRequest(input={"idx": 0}, metadata={"batch": "a"}),
+        EnqueueRolloutRequest(input={"idx": 1}, mode="train"),
+        EnqueueRolloutRequest(input={"idx": 2}, config=RolloutConfig(timeout_seconds=3.5)),
+    ]
+
+    rollouts = await store_fixture.enqueue_many_rollouts(requests)
+
+    assert [rollout.input["idx"] for rollout in rollouts] == [0, 1, 2]
+    assert all(rollout.status == "queuing" for rollout in rollouts)
+    assert rollouts[0].metadata == {"batch": "a"}
+    assert rollouts[1].mode == "train"
+    assert rollouts[2].config.timeout_seconds == 3.5
+
+    for expected_idx in range(3):
+        dequeued = await store_fixture.dequeue_rollout()
+        assert dequeued is not None
+        assert dequeued.input["idx"] == expected_idx
+
+
+@pytest.mark.asyncio
+async def test_dequeue_many_rollouts_with_limit(store_fixture: LightningStore) -> None:
+    """dequeue_many_rollouts should honor limit and propagate worker IDs to attempts."""
+
+    requests = [EnqueueRolloutRequest(input={"idx": idx}) for idx in range(4)]
+    await store_fixture.enqueue_many_rollouts(requests)
+
+    first_batch = await store_fixture.dequeue_many_rollouts(limit=2, worker_id="bulk-worker")
+    assert len(first_batch) == 2
+    assert [attempt.input["idx"] for attempt in first_batch] == [0, 1]
+    assert all(attempt.attempt.worker_id == "bulk-worker" for attempt in first_batch)
+
+    second_batch = await store_fixture.dequeue_many_rollouts(limit=5, worker_id="bulk-worker")
+    assert len(second_batch) == 2
+    assert [attempt.input["idx"] for attempt in second_batch] == [2, 3]
+
+    assert await store_fixture.dequeue_many_rollouts(limit=1) == []
+
+
 @pytest.mark.asyncio
 async def test_update_and_query_workers(store_fixture: LightningStore) -> None:
     """Workers can be created, heartbeats recorded, and telemetry auto-updated."""
@@ -717,6 +762,70 @@ async def test_update_and_query_workers(store_fixture: LightningStore) -> None:
         await store_fixture.update_worker("worker-1", heartbeat_stats=None)  # type: ignore[arg-type]
 
 
+@pytest.mark.asyncio
+async def test_worker_sync_preserves_existing_fields(store_fixture: LightningStore) -> None:
+    """Worker heartbeats should survive status syncs triggered by attempts."""
+    worker_id = "worker-sync-preserve"
+    initial = await store_fixture.update_worker(worker_id, heartbeat_stats={"cpu": 0.42})
+    assert initial.last_heartbeat_time is not None
+
+    attempted = await store_fixture.start_rollout(input={"task": "preserve"}, worker_id=worker_id)
+
+    busy = await store_fixture.get_worker_by_id(worker_id)
+    assert busy is not None
+    assert busy.heartbeat_stats == {"cpu": 0.42}
+    assert busy.last_heartbeat_time == initial.last_heartbeat_time
+
+    await store_fixture.update_attempt(attempted.rollout_id, attempted.attempt.attempt_id, status="succeeded")
+    idle = await store_fixture.get_worker_by_id(worker_id)
+    assert idle is not None
+    assert idle.heartbeat_stats == {"cpu": 0.42}
+    assert idle.last_heartbeat_time == initial.last_heartbeat_time
+
+
+@pytest.mark.asyncio
+async def test_start_rollout_assigns_worker(store_fixture: LightningStore) -> None:
+    """start_rollout should immediately associate attempts with the provided worker."""
+    attempted = await store_fixture.start_rollout(input={"task": "direct"}, worker_id="worker-direct")
+
+    assert attempted.attempt.worker_id == "worker-direct"
+    worker = await store_fixture.get_worker_by_id("worker-direct")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == attempted.rollout_id
+    assert worker.current_attempt_id == attempted.attempt.attempt_id
+
+
+@pytest.mark.asyncio
+async def test_dequeue_rollout_assigns_worker(store_fixture: LightningStore) -> None:
+    """dequeue_rollout should stamp attempts and worker telemetry with worker_id."""
+    await store_fixture.enqueue_rollout(input={"task": "queued"})
+    dequeued = await store_fixture.dequeue_rollout(worker_id="worker-dequeue")
+
+    assert dequeued is not None
+    assert dequeued.attempt.worker_id == "worker-dequeue"
+    worker = await store_fixture.get_worker_by_id("worker-dequeue")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == dequeued.rollout_id
+    assert worker.current_attempt_id == dequeued.attempt.attempt_id
+
+
+@pytest.mark.asyncio
+async def test_start_attempt_assigns_worker(store_fixture: LightningStore) -> None:
+    """Manual retries should also update worker state when worker_id is provided."""
+    initial = await store_fixture.start_rollout(input={"task": "retry-seed"})
+    retry = await store_fixture.start_attempt(initial.rollout_id, worker_id="worker-retry")
+
+    assert retry.attempt.sequence_id == 2
+    assert retry.attempt.worker_id == "worker-retry"
+    worker = await store_fixture.get_worker_by_id("worker-retry")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == retry.rollout_id
+    assert worker.current_attempt_id == retry.attempt.attempt_id
+
+
 @pytest.mark.asyncio
 async def test_query_workers_supports_filters(store_fixture: LightningStore) -> None:
     """Worker queries should support filtering, sorting, and pagination."""
@@ -1255,6 +1364,35 @@ async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_r
     assert updated_attempt.last_heartbeat_time is not None
 
 
+@pytest.mark.asyncio
+async def test_spans_promote_preparing_attempt_with_heartbeat(
+    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Spans should set heartbeat time and promote preparing attempts/rollouts to running."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "preparing-heartbeat"})
+    dequeued = await store_fixture.dequeue_rollout()
+    assert dequeued is not None
+
+    attempts_before = await store_fixture.query_attempts(rollout.rollout_id)
+    assert attempts_before
+    attempt_id = attempts_before[0].attempt_id
+    assert attempts_before[0].status == "preparing"
+    assert attempts_before[0].last_heartbeat_time is None
+
+    heartbeat_time = 1234.5
+    monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: heartbeat_time)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.status == "running"
+    assert attempt_after.last_heartbeat_time == heartbeat_time
+
+    rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id)
+    assert rollout_after is not None
+    assert rollout_after.status == "running"
+
+
 @pytest.mark.asyncio
 async def test_unresponsive_attempt_recovers_after_span(
     store_fixture: LightningStore, mock_readable_span: Mock
@@ -1303,6 +1441,77 @@ async def test_running_attempt_updates_heartbeat(
     assert attempt_after_second.last_heartbeat_time == first_heartbeat + 100.0
 
 
+@pytest.mark.asyncio
+async def test_span_post_add_preserves_concurrent_updates(
+    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Concurrent attempt updates should not be clobbered when spans record heartbeats."""
+    rollout = await store_fixture.enqueue_rollout(input={"test": "span-concurrency"})
+    dequeued = await store_fixture.dequeue_rollout()
+    assert dequeued is not None
+    attempt_id = dequeued.attempt.attempt_id
+
+    original_post = store_fixture._post_add_spans  # type: ignore
+
+    async def patched_post(spans: List[Span], rollout_id: str, mutated_attempt_id: str) -> None:
+        await store_fixture.update_attempt(rollout_id, mutated_attempt_id, metadata={"concurrent": True})
+        await original_post(spans, rollout_id, mutated_attempt_id)
+
+    monkeypatch.setattr(store_fixture, "_post_add_spans", patched_post)
+
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.metadata == {"concurrent": True}
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_marks_unresponsive_and_updates_worker(
+    store_fixture: LightningStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Healthcheck should mark attempts unresponsive and sync worker state via the helper."""
+
+    class TimeStub:
+        def __init__(self, value: float):
+            self.value = value
+
+        def __call__(self) -> float:
+            return self.value
+
+    time_stub = TimeStub(100.0)
+    monkeypatch.setattr("agentlightning.store.collection_based.time.time", time_stub)
+
+    rollout = await store_fixture.enqueue_rollout(
+        input={"test": "healthcheck-worker"},
+        config=RolloutConfig(unresponsive_seconds=1.0),
+    )
+    dequeued = await store_fixture.dequeue_rollout(worker_id="worker-sync")
+    assert dequeued is not None
+    attempt_id = dequeued.attempt.attempt_id
+    await store_fixture.update_attempt(rollout.rollout_id, attempt_id, worker_id="worker-sync")
+
+    original_sync = store_fixture._sync_workers_with_attempts  # type: ignore
+    sync_calls: List[str] = []
+
+    async def tracking_sync(attempts: Sequence[Attempt]) -> None:
+        for attempt in attempts:
+            sync_calls.append(attempt.attempt_id)
+        await original_sync(attempts)
+
+    monkeypatch.setattr(store_fixture, "_sync_workers_with_attempts", tracking_sync)
+
+    time_stub.value = 105.0
+    await store_fixture.get_rollout_by_id(rollout.rollout_id)
+
+    worker = await store_fixture.get_worker_by_id("worker-sync")
+    assert worker is not None
+    assert worker.status == "unknown"
+
+    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
+    assert attempt_after.status == "unresponsive"
+    assert sync_calls == [attempt_after.attempt_id]
+
+
 @pytest.mark.asyncio
 async def test_duplicate_span_id_error(
     store_fixture: LightningStore, mock_readable_span: Mock, caplog: pytest.LogCaptureFixture
@@ -2329,7 +2538,7 @@ async def test_concurrent_resource_updates(store_fixture: LightningStore) -> Non
 @pytest.mark.asyncio
 async def test_update_nonexistent_rollout(store_fixture: LightningStore) -> None:
     """Test updating non-existent rollout raises error."""
-    with pytest.raises(ValueError, match="Rollout nonexistent not found"):
+    with pytest.raises(ValueError, match=r"Item.*does not exist"):
         await store_fixture.update_rollout(rollout_id="nonexistent", status="failed")
 
 

tests/store/test_restful.py

@@ -12,7 +12,7 @@ Test categories:
 
 import asyncio
 import contextlib
-from typing import AsyncGenerator, List, Tuple
+from typing import Any, AsyncGenerator, Dict, List, Tuple
 
 import aiohttp
 import pytest
@@ -69,6 +69,77 @@ async def _run_server_with_cors(cors_origins: List[str] | str | None = None):
         await server.stop()
 
 
+async def _request_rollouts_page(
+    session: aiohttp.ClientSession, api_endpoint: str, method: str, payload: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for rollouts and return parsed payload."""
+
+    if method == "get":
+        async with session.get(f"{api_endpoint}/rollouts", params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{api_endpoint}/rollouts/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_attempts_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    rollout_id: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for attempts and return parsed payload."""
+
+    base = f"{api_endpoint}/rollouts/{rollout_id}/attempts"
+    if method == "get":
+        async with session.get(base, params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{base}/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_spans_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    rollout_id: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for spans and return parsed payload."""
+
+    base_params = {"rollout_id": rollout_id}
+    base_params.update(payload)
+    if method == "get":
+        async with session.get(f"{api_endpoint}/spans", params=base_params) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{api_endpoint}/spans/search", json=base_params) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
+async def _request_workers_page(
+    session: aiohttp.ClientSession,
+    api_endpoint: str,
+    method: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Send either GET or POST /search request for workers and return parsed payload."""
+
+    base = f"{api_endpoint}/workers"
+    if method == "get":
+        async with session.get(base, params=payload) as resp:
+            assert resp.status == 200
+            return await resp.json()
+    async with session.post(f"{base}/search", json=payload) as resp:
+        assert resp.status == 200
+        return await resp.json()
+
+
 @pytest_asyncio.fixture
 async def server_client(
     store_fixture: LightningStore,
@@ -137,6 +208,89 @@ async def test_cors_allows_wildcard_origin() -> None:
                 assert allow_credentials == "true"
 
 
+@pytest.mark.asyncio
+async def test_enqueue_endpoint_batches_payloads(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    _server, _client, session, api_endpoint = server_client
+
+    single_payload = {"rollouts": [{"input": {"task": "single"}}]}
+    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=single_payload) as resp:
+        assert resp.status == 201
+        body = await resp.json()
+        assert isinstance(body, list)
+        assert len(body) == 1  # type: ignore
+        assert body[0]["input"] == {"task": "single"}
+
+    batch_payload = {
+        "rollouts": [
+            {"input": {"task": "batch-1"}, "metadata": {"batch": 1}},
+            {"input": {"task": "batch-2"}},
+        ]
+    }
+    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=batch_payload) as resp:
+        assert resp.status == 201
+        body = await resp.json()
+        assert len(body) == 2
+        assert [item["input"]["task"] for item in body] == ["batch-1", "batch-2"]
+        assert body[0]["metadata"] == {"batch": 1}
+
+
+@pytest.mark.asyncio
+async def test_dequeue_endpoint_returns_batches(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    for idx in range(3):
+        await server.enqueue_rollout(input={"idx": idx})
+
+    async with session.post(
+        f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 2, "worker_id": "rest-worker"}
+    ) as resp:
+        assert resp.status == 200
+        body = await resp.json()
+        assert len(body) == 2
+        assert all(item["attempt"]["worker_id"] == "rest-worker" for item in body)
+
+    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp:
+        assert resp.status == 200
+        body = await resp.json()
+        assert len(body) == 1
+
+    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp:
+        assert resp.status == 200
+        body = await resp.json()
+        assert body == []
+
+
+@pytest.mark.asyncio
+async def test_enqueue_endpoint_requires_rollouts_field(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    _server, _client, session, api_endpoint = server_client
+
+    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json={}) as resp:
+        assert resp.status == 422
+
+
+@pytest.mark.asyncio
+async def test_dequeue_endpoint_zero_limit_returns_empty(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    await server.enqueue_rollout(input={"idx": 0})
+
+    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 0}) as resp:
+        assert resp.status == 200
+        assert await resp.json() == []
+
+    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 1}) as resp:
+        assert resp.status == 200
+        body = await resp.json()
+        assert len(body) == 1
+
+
 @pytest.mark.asyncio
 async def test_statistics_endpoint_returns_counts(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -154,6 +308,52 @@ async def test_statistics_endpoint_returns_counts(
     assert payload["total_rollouts"] >= 1
 
 
+@pytest.mark.asyncio
+async def test_rest_start_rollout_propagates_worker_id(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    payload = {"input": {"source": "rest-worker"}, "worker_id": "rest-start-worker"}
+
+    async with session.post(f"{api_endpoint}/rollouts", json=payload) as resp:
+        assert resp.status == 201
+        data = await resp.json()
+
+    assert data["attempt"]["worker_id"] == "rest-start-worker"
+    worker = await server.get_worker_by_id("rest-start-worker")
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == data["rollout_id"]
+    assert worker.current_attempt_id == data["attempt"]["attempt_id"]
+
+
+@pytest.mark.asyncio
+async def test_rest_start_attempt_propagates_worker_id(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    async with session.post(f"{api_endpoint}/rollouts", json={"input": {"source": "rest-retry"}}) as resp:
+        assert resp.status == 201
+        base_rollout = await resp.json()
+
+    attempt_worker = "rest-attempt-worker"
+    async with session.post(
+        f"{api_endpoint}/rollouts/{base_rollout['rollout_id']}/attempts",
+        json={"worker_id": attempt_worker},
+    ) as resp:
+        assert resp.status == 201
+        retry_payload = await resp.json()
+
+    assert retry_payload["attempt"]["sequence_id"] == 2
+    assert retry_payload["attempt"]["worker_id"] == attempt_worker
+    worker = await server.get_worker_by_id(attempt_worker)
+    assert worker is not None
+    assert worker.status == "busy"
+    assert worker.current_rollout_id == retry_payload["rollout_id"]
+    assert worker.current_attempt_id == retry_payload["attempt"]["attempt_id"]
+
+
 # Rollouts Pagination, Sorting, and Filtering Tests
 
 
@@ -207,6 +407,24 @@ async def test_rollouts_pagination_disabled(
         assert len(data["items"]) == 15
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_rollouts_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    """Ensure both GET and POST /search endpoints behave the same for rollouts."""
+
+    server, _client, session, api_endpoint = server_client
+    for i in range(3):
+        await server.enqueue_rollout(input={"index": i})
+
+    data = await _request_rollouts_page(session, api_endpoint, method, {"limit": 2, "offset": 0})
+    assert data["total"] == 3
+    assert data["limit"] == 2
+    assert data["offset"] == 0
+    assert len(data["items"]) == 2
+
+
 @pytest.mark.asyncio
 async def test_rollouts_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -481,6 +699,22 @@ async def test_attempts_pagination_basic(
         assert len(data["items"]) == 5
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_attempts_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    rollout = await server.enqueue_rollout(input={"test": "attempt-search"})
+    await server.start_attempt(rollout.rollout_id)
+
+    data = await _request_attempts_page(session, api_endpoint, rollout.rollout_id, method, {"limit": 1, "offset": 0})
+    assert data["total"] == 1
+    assert data["limit"] == 1
+    assert len(data["items"]) == 1
+
+
 @pytest.mark.asyncio
 async def test_attempts_sorting(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -746,6 +980,30 @@ async def test_spans_pagination_basic(
         assert len(data["items"]) == 10
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_spans_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    attempted = await server.start_rollout(input={"test": "span-search"})
+    attempt_id = attempted.attempt.attempt_id
+    for seq in range(3):
+        await server.add_span(_make_span(attempted.rollout_id, attempt_id, seq + 1, f"span-{seq}"))
+
+    data = await _request_spans_page(
+        session,
+        api_endpoint,
+        attempted.rollout_id,
+        method,
+        {"attempt_id": attempt_id, "limit": 2, "offset": 0},
+    )
+    assert data["total"] == 3
+    assert data["limit"] == 2
+    assert len(data["items"]) == 2
+
+
 @pytest.mark.asyncio
 async def test_spans_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -1116,6 +1374,50 @@ async def test_request_json_spans_returns_pagination_metadata(
     assert len(data["items"]) == 1
 
 
+# Update semantics tests
+
+
+@pytest.mark.asyncio
+async def test_update_rollout_distinguishes_unset_fields(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    rollout = await server.enqueue_rollout(input={"id": "unset-check"})
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
+        json={"metadata": {"foo": "bar"}},
+    ) as resp:
+        assert resp.status == 200
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
+        json={"status": None},
+    ) as resp:
+        assert resp.status == 400
+
+
+@pytest.mark.asyncio
+async def test_update_attempt_distinguishes_unset_fields(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+    attempted = await server.start_rollout(input={"id": "attempt-unset"})
+    attempt_id = attempted.attempt.attempt_id
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
+        json={"status": "running"},
+    ) as resp:
+        assert resp.status == 200
+
+    async with session.post(
+        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
+        json={"worker_id": None},
+    ) as resp:
+        assert resp.status == 400
+
+
 # Client Compatibility Tests
 
 
@@ -1162,6 +1464,27 @@ async def test_client_query_with_filters(
     assert rollouts[0].rollout_id == r2.rollout_id
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post"])
+async def test_workers_search_supports_get_and_post(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    await server.update_worker("worker-1", heartbeat_stats={"cpu": 0.5})
+    await server.update_worker("worker-2", heartbeat_stats={"cpu": 0.7})
+
+    data = await _request_workers_page(
+        session,
+        api_endpoint,
+        method,
+        {"limit": 1, "offset": 0, "sort_by": "worker_id", "sort_order": "asc"},
+    )
+    assert data["total"] == 2
+    assert data["limit"] == 1
+    assert len(data["items"]) == 1
+
+
 @pytest.mark.asyncio
 async def test_workers_endpoint_supports_updates(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],

tests/store/test_threading.py

@@ -17,6 +17,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     OtelResource,
     ResourcesUpdate,
@@ -273,6 +274,42 @@ async def test_threaded_store_delegates_all_methods() -> None:
     assert [name for name, *_ in dummy_store.calls] == expected_order
 
 
+@pytest.mark.asyncio
+async def test_threaded_store_enqueue_many_rollouts_delegates() -> None:
+    requests = [
+        EnqueueRolloutRequest(input={"idx": 0}, mode="train", metadata={"batch": "left"}),
+        EnqueueRolloutRequest(input={"idx": 1}, mode=None, resources_id="resources-1"),
+    ]
+    rollouts = [
+        Rollout(rollout_id="bulk-0", input={"idx": 0}, start_time=0.0),
+        Rollout(rollout_id="bulk-1", input={"idx": 1}, start_time=1.0),
+    ]
+    dummy_store = DummyLightningStore({"enqueue_many_rollouts": rollouts})
+    threaded_store = LightningStoreThreaded(dummy_store)
+
+    result = await threaded_store.enqueue_many_rollouts(requests)
+    assert result == rollouts
+    assert dummy_store.calls[-1][0] == "enqueue_many_rollouts"
+    assert dummy_store.calls[-1][1][0] == requests
+
+
+@pytest.mark.asyncio
+async def test_threaded_store_dequeue_many_rollouts_delegates() -> None:
+    attempt_a = Attempt(rollout_id="bulk-0", attempt_id="attempt-0", sequence_id=1, start_time=0.0)
+    attempt_b = Attempt(rollout_id="bulk-1", attempt_id="attempt-1", sequence_id=1, start_time=0.0)
+    attempts = [
+        AttemptedRollout(rollout_id="bulk-0", input={"idx": 0}, start_time=0.0, attempt=attempt_a),
+        AttemptedRollout(rollout_id="bulk-1", input={"idx": 1}, start_time=0.0, attempt=attempt_b),
+    ]
+    dummy_store = DummyLightningStore({"dequeue_many_rollouts": attempts})
+    threaded_store = LightningStoreThreaded(dummy_store)
+
+    result = await threaded_store.dequeue_many_rollouts(limit=2, worker_id="thread-worker")
+    assert result == attempts
+    assert dummy_store.calls[-1][0] == "dequeue_many_rollouts"
+    assert dummy_store.calls[-1][2] == {"limit": 2, "worker_id": "thread-worker"}
+
+
 def test_threaded_store_serializes_update_attempt_calls() -> None:
     store = SlowAttemptStore()
     threaded_store = LightningStoreThreaded(store)

tests/store/test_utils.py

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

tests/utils/test_metrics.py

@@ -0,0 +1,340 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+import logging
+import sys
+import types
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import pytest
+
+import agentlightning.utils.metrics as metrics_module
+from agentlightning.utils.metrics import ConsoleMetricsBackend, MetricsBackend, MultiMetricsBackend
+
+
+def test_validate_labels_reports_missing_label_with_metric_name() -> None:
+    labels = {"method": "GET", "status": "200"}
+    result = metrics_module._validate_labels("counter", "requests", labels, ("method", "status"))
+    assert result == (("method", "GET"), ("status", "200"))
+
+    with pytest.raises(ValueError) as excinfo:
+        metrics_module._validate_labels("counter", "requests", {"method": "POST"}, ("method", "status"))
+
+    message = str(excinfo.value)
+    assert "Counter 'requests'" in message
+    assert "'status' is required" in message
+
+
+def test_normalize_label_names_is_sorted() -> None:
+    assert metrics_module._normalize_label_names(["b", "a", "b"]) == ("a", "b", "b")
+    assert metrics_module._normalize_label_names(None) == ()
+
+
+def test_console_backend_logs_counters_with_sorted_labels() -> None:
+    backend = ConsoleMetricsBackend()
+    line = backend._log_counter(
+        "rate",
+        {"group2": "b", "group1": "a"},
+        timestamps=[0.0, 1.0],
+        amounts=[1.0, 1.0],
+        snapshot_time=2.0,
+    )
+
+    assert line == "rate{group1=a,group2=b}=0.40/s"
+
+
+def test_console_backend_logs_histograms_with_human_units() -> None:
+    backend = ConsoleMetricsBackend()
+    line = backend._log_histogram(
+        "latency",
+        {"group2": "b", "group1": "a"},
+        values=[0.00395, 0.0168, 3.5],
+        buckets=(0.5,),
+        snapshot_time=1.0,
+    )
+
+    assert line is not None and line.startswith("latency{group1=a,group2=b}=")
+    payload = line.rsplit("=", 1)[1]
+    p50, p95, p99 = payload.split(",", 2)
+    assert p50.endswith("ms")
+    assert p95.endswith("s")
+    assert p99.endswith("s")
+
+
+def test_console_backend_respects_group_level_limit():
+    backend = ConsoleMetricsBackend(group_level=2)
+    truncated = backend._truncate_labels_for_logging({"group3": "c", "group1": "a", "group2": "b"})
+    line = backend._log_counter("metric", truncated, [0.0, 2.0], [1.0, 1.0], snapshot_time=3.0)
+
+    assert line == "metric{group1=a,group2=b}=0.40/s"
+
+
+def test_console_backend_log_uses_logger(caplog: pytest.LogCaptureFixture) -> None:
+    backend = ConsoleMetricsBackend()
+    caplog.set_level(logging.INFO, logger="agentlightning.utils.metrics")
+    backend._log("hello metrics")
+    assert any(record.message == "hello metrics" for record in caplog.records)
+
+
+class _RecordingBackend(MetricsBackend):
+    def __init__(self) -> None:
+        self.calls: List[Tuple[str, Tuple[Any, ...]]] = []
+
+    def register_counter(self, name: str, label_names: Optional[Sequence[str]] = None) -> None:
+        self.calls.append(("register_counter", (name, tuple(label_names or ()))))
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+    ) -> None:
+        self.calls.append(("register_histogram", (name, tuple(label_names or ()), tuple(buckets or ()))))
+
+    def inc_counter(self, name: str, amount: float = 1.0, labels: Optional[Dict[str, str]] = None) -> None:
+        self.calls.append(("inc_counter", (name, amount, labels or {})))
+
+    def observe_histogram(self, name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
+        self.calls.append(("observe_histogram", (name, value, labels or {})))
+
+
+def test_multi_metrics_backend_fans_out_calls() -> None:
+    backend_a = _RecordingBackend()
+    backend_b = _RecordingBackend()
+    multi = MultiMetricsBackend([backend_a, backend_b])
+
+    multi.register_counter("hits", label_names=["method"])
+    multi.register_histogram("latency", label_names=["method"], buckets=[0.1, 1.0])
+    multi.inc_counter("hits", amount=2.5, labels={"method": "GET"})
+    multi.observe_histogram("latency", value=0.4, labels={"method": "GET"})
+
+    expected_calls: List[Tuple[str, Tuple[Any, ...]]] = [
+        ("register_counter", ("hits", ("method",))),
+        ("register_histogram", ("latency", ("method",), (0.1, 1.0))),
+        ("inc_counter", ("hits", 2.5, {"method": "GET"})),
+        ("observe_histogram", ("latency", 0.4, {"method": "GET"})),
+    ]
+
+    assert backend_a.calls == expected_calls
+    assert backend_b.calls == expected_calls
+
+
+class _PrometheusStub(types.ModuleType):
+    def __init__(self) -> None:
+        super().__init__("prometheus_client")
+        self.counter_instances: List["_PromCounter"] = []
+        self.histogram_instances: List["_PromHistogram"] = []
+
+        self_owner = self
+
+        class _CounterChild:
+            def __init__(self) -> None:
+                self.value = 0.0
+
+            def inc(self, amount: float = 1.0) -> None:
+                self.value += amount
+
+        class _HistogramChild:
+            def __init__(self) -> None:
+                self.values: List[float] = []
+
+            def observe(self, value: float) -> None:
+                self.values.append(value)
+
+        class _PromCounter:
+            def __init__(self, name: str, doc: str, labelnames: Sequence[str]) -> None:
+                self.name = name
+                self.doc = doc
+                self.labelnames = tuple(labelnames)
+                self.default = _CounterChild()
+                self.children: Dict[Tuple[Tuple[str, str], ...], _CounterChild] = {}
+                self._register()
+
+            def _register(self) -> None:
+                self_owner.counter_instances.append(self)
+
+            def labels(self, **kwargs: str) -> _CounterChild:
+                key = tuple(sorted(kwargs.items()))
+                return self.children.setdefault(key, _CounterChild())
+
+            def inc(self, amount: float = 1.0) -> None:
+                self.default.inc(amount)
+
+        class _PromHistogram:
+            def __init__(self, name: str, doc: str, labelnames: Sequence[str], buckets: Sequence[float]) -> None:
+                self.name = name
+                self.doc = doc
+                self.labelnames = tuple(labelnames)
+                self.buckets = tuple(buckets)
+                self.default = _HistogramChild()
+                self.children: Dict[Tuple[Tuple[str, str], ...], _HistogramChild] = {}
+                self._register()
+
+            def _register(self) -> None:
+                self_owner.histogram_instances.append(self)
+
+            def labels(self, **kwargs: str) -> _HistogramChild:
+                key = tuple(sorted(kwargs.items()))
+                return self.children.setdefault(key, _HistogramChild())
+
+            def observe(self, value: float) -> None:
+                self.default.observe(value)
+
+        self.Counter = _PromCounter
+        self.Histogram = _PromHistogram
+
+        class CollectorRegistry:
+            pass
+
+        self.CollectorRegistry = CollectorRegistry
+        self.REGISTRY = CollectorRegistry()
+
+        class _Multiprocess:
+            def __init__(self) -> None:
+                self.registry: Optional[CollectorRegistry] = None
+
+            def MultiProcessCollector(self, registry: CollectorRegistry) -> None:
+                self.registry = registry
+
+        self.multiprocess = _Multiprocess()
+
+
+def _make_prometheus_stub() -> _PrometheusStub:
+    return _PrometheusStub()
+
+
+def test_prometheus_backend_binds_stubbed_prometheus(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = _make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("hits", ["method"])
+    backend.register_histogram("latency", ["method"], buckets=[0.5])
+
+    backend.inc_counter("hits", amount=2.0, labels={"method": "GET"})
+    backend.observe_histogram("latency", value=0.1, labels={"method": "GET"})
+
+    counter_instance = stub.counter_instances[0]
+    histogram_instance = stub.histogram_instances[0]
+    assert counter_instance.children[(("method", "GET"),)].value == 2.0
+    assert histogram_instance.children[(("method", "GET"),)].values == [0.1]
+
+
+def _split_segments(log_lines: Sequence[str]) -> List[str]:
+    segments: List[str] = []
+    for line in log_lines:
+        segments.extend(part.strip() for part in line.split("  "))
+    return [seg for seg in segments if seg]
+
+
+def test_console_backend_sliding_window_rate_and_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=5.0, log_interval_seconds=0.0)
+    backend.register_counter("requests", ["group", "path"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 6.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    labels = {"group": "api", "path": "/list"}
+    backend.inc_counter("requests", labels=labels)
+    backend.inc_counter("requests", labels=labels)
+    backend.inc_counter("requests", labels=labels)
+
+    segments = [seg for seg in _split_segments(logged) if seg.startswith("requests{group=api,path=/list}")]
+    assert len(segments) == 3
+
+    latest = segments[-1]
+    rate_str = latest.rsplit("=", 1)[1].rstrip("/s")
+    rate = float(rate_str)
+    assert abs(rate - 0.40) < 1e-2
+
+
+def _duration_to_seconds(payload: str) -> float:
+    if payload.endswith("ms"):
+        return float(payload[:-2]) / 1_000
+    if payload.endswith("µs"):
+        return float(payload[:-2]) / 1_000_000
+    if payload.endswith("ns"):
+        return float(payload[:-2]) / 1_000_000_000
+    if payload.endswith("s"):
+        return float(payload[:-1])
+    raise AssertionError(f"Unknown duration format: {payload}")
+
+
+def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=10.0, log_interval_seconds=0.0, group_level=2)
+    backend.register_histogram("latency", ["service", "endpoint", "status"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 2.0, 3.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    svc_a_labels = {"service": "svcA", "endpoint": "/search", "status": "200"}
+    svc_b_labels = {"service": "svcB", "endpoint": "/chat", "status": "500"}
+
+    backend.observe_histogram("latency", value=0.01, labels=svc_a_labels)
+    backend.observe_histogram("latency", value=0.02, labels=svc_a_labels)
+    backend.observe_histogram("latency", value=0.03, labels=svc_a_labels)
+    backend.observe_histogram("latency", value=2.0, labels=svc_b_labels)
+
+    svc_a_segments = [
+        seg for seg in _split_segments(logged) if seg.startswith("latency{endpoint=/search,service=svcA}")
+    ]
+    assert svc_a_segments, "expected log entries for service A"
+    latest = svc_a_segments[-1]
+    assert "status" not in latest
+
+    payload = latest.rsplit("=", 1)[1]
+    p50_str, p95_str, p99_str = payload.split(",", 2)
+    p50 = _duration_to_seconds(p50_str)
+    p95 = _duration_to_seconds(p95_str)
+    p99 = _duration_to_seconds(p99_str)
+
+    assert abs(p50 - 0.02) < 1e-3
+    assert abs(p95 - 0.029) < 5e-3
+    assert abs(p99 - 0.0298) < 5e-3
+
+
+def test_console_backend_logs_all_metric_groups(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_counter("requests", ["group"])
+    backend.register_counter("errors", ["group"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    backend.inc_counter("requests", labels={"group": "api"})
+    backend.inc_counter("errors", labels={"group": "api"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{group=api}" in last_line
+    assert "errors{group=api}" in last_line
+
+
+def test_console_backend_snapshot_logs_single_line() -> None:
+    backend = ConsoleMetricsBackend()
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("counter", {"g": "1"}, [0.0, 1.0], [1.0, 1.0]),
+        ],
+        [
+            ("latency", {"g": "1"}, [0.1, 0.2], (0.5,)),
+        ],
+        snapshot_time=1.5,
+    )
+
+    assert logged and "counter{g=1}" in logged[0] and "latency{g=1}" in logged[0]