update catalog

.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

@@ -1,6 +1,3 @@
-# This workflow is used to benchmark the performance of the project.
-# It's kept as a placeholder for now.
-
 name: Benchmark
 permissions:
   contents: read
@@ -9,11 +6,182 @@ on:
 
 jobs:
   benchmark:
-    name: Benchmark
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    name: Benchmark (${{ matrix.backend.id }}, ${{ matrix.scenario.display }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
     timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      matrix:
+        backend:
+          - id: memory
+            compose_file: compose.prometheus-memory-store.yml
+          - id: mongo
+            compose_file: compose.prometheus-mongo-store.yml
+        scenario:
+          - id: minimal-production
+            display: Minimal production scale
+            store_workers: 4
+            args: >-
+              --mode batch
+              --total-tasks 4096
+              --batch-size 256
+              --n-runners 32
+              --max-rounds 6
+              --sleep-seconds 0.5
+          - id: medium-production
+            display: Medium production scale
+            store_workers: 16
+            args: >-
+              --mode batch
+              --total-tasks 10000
+              --batch-size 1000
+              --n-runners 100
+              --max-rounds 10
+              --sleep-seconds 0.1
+          - id: large-batch
+            display: Large batch waves
+            store_workers: 32
+            args: >-
+              --mode batch
+              --total-tasks 100000
+              --batch-size 8192
+              --n-runners 256
+              --max-rounds 6
+              --sleep-seconds 0.1
+          - id: long-queues
+            display: Long rollout queues
+            store_workers: 32
+            args: >-
+              --mode batch_partial
+              --total-tasks 100000
+              --batch-size 1024
+              --n-runners 256
+              --remaining-tasks 4096
+              --max-rounds 4
+              --sleep-seconds 0.1
+          - id: high-concurrency
+            display: High-throughput concurrent requests
+            store_workers: 32
+            args: >-
+              --mode single
+              --total-tasks 100000
+              --concurrency 2048
+              --n-runners 256
+              --max-rounds 2
+              --sleep-seconds 0.1
+          - id: heavy-traces
+            display: Heavy rollouts with deep traces
+            store_workers: 64
+            args: >-
+              --mode batch_partial
+              --total-tasks 10000
+              --batch-size 1024
+              --remaining-tasks 256
+              --n-runners 512
+              --max-rounds 20
+              --sleep-seconds 1.0
+    env:
+      STORE_URL: http://localhost:4747
+      STORE_API_URL: http://localhost:4747/v1/agl
+      PROM_URL: http://localhost:9090
+      SCENARIO_ID: ${{ matrix.scenario.id }}
+      BACKEND_ID: ${{ matrix.backend.id }}
+      ARTIFACT_DIR: artifacts/${{ matrix.scenario.id }}-${{ matrix.backend.id }}
+      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
+      AGL_STORE_N_WORKERS: ${{ matrix.scenario.store_workers }}
     steps:
-      - name: Check GPU status
-        run: nvidia-smi
+      - 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: Check disk space
         run: df -h
+
+      - 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
+          exit 1
+
+      - name: Prepare artifact directory
+        run: mkdir -p "$ARTIFACT_DIR"
+
+      - name: Record benchmark start
+        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run ${{ matrix.scenario.display }} workload
+        run: |
+          set -euo pipefail
+          uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
+            --store-url "$STORE_URL" \
+            ${{ matrix.scenario.args }}
+
+      - name: Record benchmark end
+        if: ${{ always() }}
+        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run 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.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.txt"
+
+      - 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-${SCENARIO_ID}-${BACKEND_ID}.tar.gz" prometheus
+          fi
+
+      - name: Upload benchmark artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error

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

@@ -0,0 +1,151 @@
+name: Examples - Claude Code
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: "0 20 * * *"
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-claude-code, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Claude Code - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Claude Code - {0}', github.event_name) }}
+
+jobs:
+  claude-code:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-claude-code' ||
+      github.event.action == 'ci-all'
+    name: Claude Code (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.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
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-gpu-stable
+      - 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-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Download model
+        run: |
+          source .venv/bin/activate
+          python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
+
+      - name: Launch vLLM server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+            --max-model-len 131072 \
+            --enable-auto-tool-choice \
+            --tool-call-parser qwen3_coder \
+            --port 45993 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:45993/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: Claude Code sanity check with vLLM models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          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 --debug
+        shell: bash
+
+      - name: Upload sanity check artifacts for vLLM
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error
+
+      - name: Cleanup vLLM
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+          rm -rf examples/claude_code/data/
+          rm -rf examples/claude_code/logs/
+
+      - name: Claude Code sanity check with OpenAI models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+
+      - name: Upload sanity check artifacts for OpenAI
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error

agentlightning/__init__.py

@@ -7,6 +7,7 @@ from .algorithm import *
 from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
 from .config import *
 from .emitter import *
+from .env_var import *
 from .execution import *
 from .litagent import *
 from .llm_proxy import *

agentlightning/adapter/triplet.py

@@ -11,7 +11,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
 
-from agentlightning.types import Span, SpanNames, Triplet
+from agentlightning.emitter.reward import get_reward_value
+from agentlightning.types import Span, Triplet
 
 from .base import TraceAdapter
 
@@ -313,24 +314,11 @@ class TraceTree:
         Returns:
             Dictionary containing reward metadata, or an empty dictionary when no reward is found.
         """
-        for key in [
-            "agentops.task.output",  # newer versions of agentops
-            "agentops.entity.output",
-        ]:
-            output = self.span.attributes.get(key)  # type: ignore
-            if output:
-                if isinstance(output, dict):
-                    return output
-                elif isinstance(output, str):
-                    try:
-                        return json.loads(output)
-                    except json.JSONDecodeError:
-                        return {}
-
-        # Latest emit reward format
-        if self.span.name == SpanNames.REWARD.value and self.span.attributes:
-            return {"type": "reward", "value": self.span.attributes.get("reward", None)}
-        return {}
+        reward_value = get_reward_value(self.span)
+        if reward_value is not None:
+            return {"type": "reward", "value": reward_value}
+        else:
+            return {}
 
     def is_reward_span(self) -> bool:
         """Return whether the span explicitly encodes a reward.
@@ -776,24 +764,7 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
 
     def _maybe_reward_value(self, span: Span) -> Optional[float]:
         """Parse reward from typical AgentOps payloads or explicit reward spans."""
-        attrs = span.attributes or {}
-
-        # AgentOps new/old keys
-        for k in ("agentops.task.output", "agentops.entity.output"):
-            v = attrs.get(k)
-            v = self._literal_eval_maybe(v)
-            if isinstance(v, dict) and cast(Dict[str, Any], v).get("type") == "reward":
-                rv = cast(Dict[str, Any], v).get("value", None)
-                if rv is None or isinstance(rv, (int, float)):
-                    return None if rv is None else float(rv)
-
-        # Explicit reward span
-        if span.name == SpanNames.REWARD.value:
-            rv = attrs.get("reward", None)
-            if rv is None or isinstance(rv, (int, float)):
-                return None if rv is None else float(rv)
-
-        return None
+        return get_reward_value(span)
 
     def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
         # Prefer OpenAI-like id if present, else proxy raw id.

agentlightning/cli/store.py

@@ -64,11 +64,11 @@ def main(argv: Iterable[str] | None = None) -> int:
     setup_logging(args.log_level)
 
     if args.backend == "memory":
-        store = InMemoryLightningStore()
+        store = InMemoryLightningStore(prometheus=args.prometheus)
     elif args.backend == "mongo":
         from agentlightning.store.mongo import MongoLightningStore
 
-        store = MongoLightningStore(client=args.mongo_uri)
+        store = MongoLightningStore(client=args.mongo_uri, prometheus=args.prometheus)
     else:
         raise ValueError(f"Invalid backend: {args.backend}")
 

agentlightning/emitter/__init__.py

@@ -1,13 +1,15 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from .annotation import emit_annotation
 from .exception import emit_exception
-from .message import emit_message
-from .object import emit_object
+from .message import emit_message, get_message_value
+from .object import emit_object, get_object_value
 from .reward import (
     emit_reward,
     find_final_reward,
     find_reward_spans,
     get_reward_value,
+    get_rewards_from_span,
     is_reward_span,
     reward,
 )
@@ -16,10 +18,14 @@ __all__ = [
     "reward",
     "emit_reward",
     "get_reward_value",
+    "get_rewards_from_span",
     "is_reward_span",
     "find_reward_spans",
     "find_final_reward",
     "emit_message",
     "emit_object",
     "emit_exception",
+    "emit_annotation",
+    "get_message_value",
+    "get_object_value",
 ]

agentlightning/emitter/annotation.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helpers for emitting annotation spans."""
+
+import logging
+from typing import Any, Dict
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.semconv import AGL_ANNOTATION
+from agentlightning.utils.otel import flatten_attributes, get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> ReadableSpan:
+    """Emit a new annotation span.
+
+    This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward].
+
+    Annotation spans are used to annotate a specific event or a part of rollout.
+    See [semconv][agentlightning.semconv] for conventional annotation keys in Agent-lightning.
+
+    If annotations contain nested dicts, they will be flattened before emitting.
+    Complex objects will lead to emitting failures.
+
+    Args:
+        annotation: Dictionary containing annotation key-value pairs.
+            Representatives are rewards, tags, and metadata.
+        propagate: Whether to propagate the span to exporters automatically.
+    """
+    annotation_attributes = flatten_attributes(annotation)
+    if any(not isinstance(v, (str, int, float, bool, bytes)) for v in annotation_attributes.values()):
+        raise TypeError("All annotation attributes must be primitive types (str, int, float, bool, bytes)")
+
+    # TODO: this should use a tracer from current context rather than the singleton
+    tracer = get_tracer(use_active_span_processor=propagate)
+    span = tracer.start_span(
+        AGL_ANNOTATION,
+        attributes=annotation_attributes,
+    )
+    logger.debug("Emitting annotation span with keys %s", annotation_attributes)
+    with span:
+        pass
+    if not isinstance(span, ReadableSpan):
+        raise ValueError(f"Span is not a ReadableSpan: {span}")
+
+    return span

agentlightning/emitter/exception.py

@@ -2,43 +2,53 @@
 
 import logging
 import traceback
+from typing import Any, Dict, Optional
 
 from opentelemetry.semconv.attributes import exception_attributes
 
-from agentlightning.types import SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_EXCEPTION
+from agentlightning.utils.otel import get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_exception(exception: BaseException) -> None:
+def emit_exception(
+    exception: BaseException, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True
+) -> None:
     """Record an exception with OpenTelemetry metadata.
 
+    Classic OpenTelemetry records exceptions in a dedicated logging service.
+    We simplify the model and use trace spans to record exceptions as well.
+
     Args:
         exception: Raised exception instance to serialize into telemetry attributes.
+        attributes: Additional attributes to attach to the exception span.
+        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
-        The helper validates its input. Non-exception values are ignored to prevent
-        noisy telemetry and indicate programming mistakes via the logger.
+
+        The helper validates its input. If a non-exception value is provided,
+        a TypeError is raised to indicate a programming mistake.
     """
     if not isinstance(exception, BaseException):  # type: ignore
-        logger.error(f"Expected an BaseException instance, got: {type(exception)}. Skip emit_exception.")
-        return
+        raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
 
-    tracer = get_tracer()
+    tracer = get_tracer(use_active_span_processor=propagate)
     stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
-    attributes = {
+    span_attributes = {
         exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
         exception_attributes.EXCEPTION_MESSAGE: str(exception),
         exception_attributes.EXCEPTION_ESCAPED: True,
     }
     if stacktrace.strip():
-        attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+        span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+
+    if attributes:
+        span_attributes.update(attributes)
 
     span = tracer.start_span(
-        SpanNames.EXCEPTION.value,
-        attributes=attributes,
+        AGL_EXCEPTION,
+        attributes=span_attributes,
     )
     logger.debug("Emitting exception span for %s", type(exception).__name__)
     with span:

agentlightning/emitter/message.py

@@ -1,33 +1,55 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import logging
+from typing import Any, Dict, Optional
 
-from agentlightning.types import SpanAttributeNames, SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
+from agentlightning.types import SpanLike
+from agentlightning.utils.otel import get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_message(message: str) -> None:
+def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
     """Emit a textual message as an OpenTelemetry span.
 
+    Commonly used for sending debugging and logging messages.
+
     Args:
         message: Human readable message to attach as a span attribute.
+        attributes: Additional attributes to attach to the message span.
+        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
         OpenTelemetry distinguishes between logs and spans. Emitting the message as a
         span keeps all Agent Lightning telemetry in a single data store for analysis.
     """
     if not isinstance(message, str):  # type: ignore
-        logger.error(f"Message must be a string, got: {type(message)}. Skip emit_message.")
-        return
+        raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
 
-    tracer = get_tracer()
+    tracer = get_tracer(use_active_span_processor=propagate)
+    span_attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
+    if attributes:
+        span_attributes.update(attributes)
     span = tracer.start_span(
-        SpanNames.MESSAGE.value,
-        attributes={SpanAttributeNames.MESSAGE.value: message},
+        AGL_MESSAGE,
+        attributes=span_attributes,
     )
     logger.debug("Emitting message span with message: %s", message)
     with span:
         pass
+
+
+def get_message_value(span: SpanLike) -> Optional[str]:
+    """Extract the message string from a message span.
+
+    Args:
+        span: Span-like object to extract the message from.
+    """
+    span_attributes = span.attributes or {}
+    if LightningSpanAttributes.MESSAGE_BODY.value not in span_attributes:
+        return None
+    message = span_attributes[LightningSpanAttributes.MESSAGE_BODY.value]
+    if isinstance(message, str):
+        return message
+    raise TypeError(f"Message must be a string, got: {type(message)}.")

agentlightning/emitter/object.py

@@ -1,37 +1,106 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import base64
 import json
 import logging
-from typing import Any
+from typing import Any, Dict, Optional
 
-from agentlightning.types import SpanAttributeNames, SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
+from agentlightning.types import SpanLike
+from agentlightning.utils.otel import full_qualified_name, get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_object(object: Any) -> None:
+def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
     """Emit an object's serialized representation as an OpenTelemetry span.
 
     Args:
         object: Data structure to encode as JSON and attach to the span payload.
+        attributes: Additional attributes to attach to the object span.
+        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
-        The payload must be JSON serializable. Non-serializable objects are ignored and
-        an error is logged to aid debugging.
+        The payload must be JSON serializable. Non-serializable objects will lead to a RuntimeError.
     """
-    try:
-        serialized = json.dumps(object)
-    except (TypeError, ValueError):
-        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
-        return
-
-    tracer = get_tracer()
+    span_attributes = encode_object(object)
+    if attributes:
+        span_attributes.update(attributes)
+    tracer = get_tracer(use_active_span_processor=propagate)
     span = tracer.start_span(
-        SpanNames.OBJECT.value,
-        attributes={SpanAttributeNames.OBJECT.value: serialized},
+        AGL_OBJECT,
+        attributes=span_attributes,
     )
-    logger.debug("Emitting object span with payload size %d characters", len(serialized))
+    attr_length = 0
+    if LightningSpanAttributes.OBJECT_JSON.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_JSON.value])
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_LITERAL.value])
+    logger.debug("Emitting object span with payload size %d characters", attr_length)
     with span:
         pass
+
+
+def encode_object(object: Any) -> Dict[str, Any]:
+    """Encode an object as span attributes.
+
+    Args:
+        object: Data structure to encode as JSON.
+    """
+    span_attributes = {}
+    if isinstance(object, (str, int, float, bool)):
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: type(object).__name__,
+            LightningSpanAttributes.OBJECT_LITERAL.value: str(object),
+        }
+    elif isinstance(object, bytes):
+        b64_encoded = base64.b64encode(object).decode("utf-8")
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
+            LightningSpanAttributes.OBJECT_LITERAL.value: b64_encoded,
+        }
+    else:
+        try:
+            serialized = json.dumps(object)
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError(f"Object must be JSON serializable, got: {type(object)}.") from exc
+
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: full_qualified_name(type(object)),  # type: ignore
+            LightningSpanAttributes.OBJECT_JSON.value: serialized,
+        }
+
+    return span_attributes
+
+
+def get_object_value(span: SpanLike) -> Any:
+    """Extract the object payload from an object span.
+
+    Args:
+        span: Span object produced by Agent Lightning emitters.
+    """
+    attributes = span.attributes or {}
+    if LightningSpanAttributes.OBJECT_JSON.value in attributes:
+        serialized = attributes[LightningSpanAttributes.OBJECT_JSON.value]
+        try:
+            return json.loads(serialized)  # type: ignore
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError("Failed to deserialize object JSON from span.") from exc
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in attributes:
+        literal = attributes[LightningSpanAttributes.OBJECT_LITERAL.value]
+        obj_type = attributes.get(LightningSpanAttributes.OBJECT_TYPE.value, "str")
+        if obj_type == "str":
+            return literal
+        elif obj_type == "int":
+            # Let it raise errors if there are any
+            return int(literal)  # type: ignore
+        elif obj_type == "float":
+            return float(literal)  # type: ignore
+        elif obj_type == "bool":
+            return literal.lower() == "true"  # type: ignore
+        elif obj_type == "bytes":
+            return base64.b64decode(literal.encode("utf-8"))  # type: ignore
+        else:
+            raise RuntimeError(f"Unsupported object type for literal deserialization: {obj_type}")
+    else:
+        return None

agentlightning/emitter/reward.py

@@ -23,10 +23,13 @@ from typing import (
 import agentops
 from agentops.sdk.decorators import operation
 from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import TypeAdapter
 
-from agentlightning.types import SpanLike, SpanNames
+from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
+from agentlightning.types import SpanLike
+from agentlightning.utils.otel import filter_and_unflatten_attributes
 
-from .utils import get_tracer
+from .annotation import emit_annotation
 
 logger = logging.getLogger(__name__)
 
@@ -34,18 +37,26 @@ __all__ = [
     "reward",
     "emit_reward",
     "get_reward_value",
+    "get_rewards_from_span",
     "is_reward_span",
     "find_reward_spans",
     "find_final_reward",
 ]
 
 
-class RewardSpanData(TypedDict):
+class RewardDimension(TypedDict):
+    """Type representing a single dimension in a multi-dimensional reward."""
+
+    name: str
+    value: float
+
+
+class _RewardSpanData(TypedDict):
     type: Literal["reward"]
     value: Optional[float]
 
 
-FnType = TypeVar("FnType", bound=Callable[..., Any])
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])
 
 
 def _agentops_initialized() -> bool:
@@ -53,7 +64,7 @@ def _agentops_initialized() -> bool:
     return agentops.get_client().initialized
 
 
-def reward(fn: FnType) -> FnType:
+def reward(fn: _FnType) -> _FnType:
     """Decorate a reward function so its outputs are tracked as spans.
 
     The decorator integrates with AgentOps when it is available and falls back to
@@ -70,7 +81,7 @@ def reward(fn: FnType) -> FnType:
         Wrapped callable that preserves the original signature.
     """
 
-    def wrap_result(result: Optional[float]) -> RewardSpanData:
+    def wrap_result(result: Optional[float]) -> _RewardSpanData:
         """Normalize the reward value into the span payload format."""
         if result is None:
             return {"type": "reward", "value": None}
@@ -94,7 +105,7 @@ def reward(fn: FnType) -> FnType:
             result: Optional[float] = None
 
             @operation
-            async def agentops_reward_operation() -> RewardSpanData:
+            async def agentops_reward_operation() -> _RewardSpanData:
                 # The reward function we are interested in tracing
                 # It takes zero inputs and return a formatted dict
                 nonlocal result
@@ -118,7 +129,7 @@ def reward(fn: FnType) -> FnType:
             result: Optional[float] = None
 
             @operation
-            def agentops_reward_operation() -> RewardSpanData:
+            def agentops_reward_operation() -> _RewardSpanData:
                 nonlocal result
                 result = fn(*args, **kwargs)
                 return wrap_result(result)
@@ -129,13 +140,36 @@ def reward(fn: FnType) -> FnType:
         return wrapper  # type: ignore
 
 
-def emit_reward(reward: float, auto_export: bool = True) -> ReadableSpan:
+def emit_reward(
+    reward: float | Dict[str, Any],
+    *,
+    primary_key: str | None = None,
+    attributes: Dict[str, Any] | None = None,
+    propagate: bool = True,
+) -> ReadableSpan:
     """Emit a reward value as an OpenTelemetry span.
 
+    Examples:
+        Emit a single-dimensional reward:
+        >>> emit_reward(1.0)
+
+        Emit multi-dimensional rewards:
+        >>> emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
+
+        Emit a reward with additional attributes (for example linking to another response span):
+        >>> from agentlightning.utils.otel import make_link_attributes
+        >>> emit_reward(0.5, attributes=make_link_attributes({"gen_ai.response.id": "response-123"}))
+
+        Or adding tags onto the reward span:
+        >>> from agentlightning.utils.otel import make_tag_attributes
+        >>> emit_reward(0.7, attributes=make_tag_attributes(["fast", "reliable"]))
+
     Args:
         reward: Numeric reward to record. Integers and booleans are converted to
             floating point numbers for consistency.
-        auto_export: Whether to export the span automatically.
+            Use a dictionary to represent a multi-dimensional reward.
+        attributes: Other optional span attributes.
+        propagate: Whether to propagate the span to exporters automatically.
 
     Returns:
         Readable span capturing the recorded reward.
@@ -145,20 +179,34 @@ def emit_reward(reward: float, auto_export: bool = True) -> ReadableSpan:
             resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
     """
     logger.debug(f"Emitting reward: {reward}")
-    if isinstance(reward, (int, bool)):
-        reward = float(reward)
-    if not isinstance(reward, float):
-        raise ValueError(f"Reward must be a number, got: {type(reward)}")
+    reward_dimensions: List[RewardDimension] = []
+    if isinstance(reward, dict):
+        reward_dict: Dict[str, float] = {}
+        for k, v in reward.items():
+            if isinstance(v, (int, bool)):
+                reward_dict[k] = float(v)
+            elif isinstance(v, float):
+                reward_dict[k] = v
+            else:
+                raise ValueError(f"Reward value must be a number, got: {type(v)} for key {k}")
+        if primary_key is None:
+            raise ValueError("When emitting a multi-dimensional reward as a dict, primary_key must be provided.")
+        if primary_key not in reward_dict:
+            raise ValueError(f"Primary key '{primary_key}' not found in reward dict keys: {list(reward_dict.keys())}")
+        reward_dimensions.append(RewardDimension(name=primary_key, value=reward_dict[primary_key]))
+        for k, v in reward_dict.items():
+            if k != primary_key:
+                reward_dimensions.append(RewardDimension(name=k, value=v))
+    else:
+        if isinstance(reward, (int, bool)):
+            reward = float(reward)
+        elif not isinstance(reward, float):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise TypeError(f"Reward must be a number, got: {type(reward)}")
+        reward_dimensions.append(RewardDimension(name="primary", value=reward))
 
-    # TODO: This should use the tracer from current context by tracer
-    tracer = get_tracer(use_active_span_processor=auto_export)
-    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
-    # Do nothing; it's just a number
-    with span:
-        pass
-    if not isinstance(span, ReadableSpan):
-        raise ValueError(f"Span is not a ReadableSpan: {span}")
-    return span
+    return emit_annotation(
+        {LightningSpanAttributes.REWARD.value: reward_dimensions, **(attributes or {})}, propagate=propagate
+    )
 
 
 def get_reward_value(span: SpanLike) -> Optional[float]:
@@ -168,8 +216,14 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
         span: Span object produced by AgentOps or Agent Lightning emitters.
 
     Returns:
-        The reward encoded in the span or `None` when the span does not represent a reward.
+        The primary reward encoded in the span or `None` when the span does not represent a reward.
     """
+    # v0.3+ emit reward format
+    reward_list = get_rewards_from_span(span)
+    if reward_list:
+        # Reward list is ordered and the first element is the primary reward
+        return reward_list[0].value
+
     for key in [
         "agentops.task.output",  # newer versions of agentops
         "agentops.entity.output",
@@ -192,19 +246,45 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
                 return None
             if not isinstance(reward_value, float):
                 logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+            logger.warning(
+                f"Extracted reward {reward_value} from AgentOps. This format is deprecated, please migrate to using `emit_reward`."
+            )
             return cast(float, reward_value)
 
-    # Latest emit reward format
-    if span.name == SpanNames.REWARD.value and span.attributes:
+    # v0.2 emit reward format
+    if span.name == AGL_ANNOTATION and span.attributes:
         reward_value = span.attributes.get("reward", None)
         if reward_value is None:
             return None
         if not isinstance(reward_value, float):
             logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+        logger.warning(
+            f"Extracted reward {reward_value} from a legacy version of reward span. You might have inconsistent agent-lightning versions."
+        )
         return cast(float, reward_value)
+
     return None
 
 
+def get_rewards_from_span(span: SpanLike) -> List[RewardPydanticModel]:
+    """Extract the reward as a list from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or Agent Lightning emitters.
+
+    Returns:
+        A list of reward dimensions encoded in the span or an empty list when the span does not represent a reward.
+    """
+    if span.attributes and any(key.startswith(LightningSpanAttributes.REWARD.value) for key in span.attributes):
+        reward_attr = filter_and_unflatten_attributes(
+            cast(Any, span.attributes or {}), LightningSpanAttributes.REWARD.value
+        )
+        recovered_rewards = TypeAdapter(List[RewardPydanticModel]).validate_python(reward_attr)
+        return recovered_rewards
+    else:
+        return []
+
+
 def is_reward_span(span: SpanLike) -> bool:
     """Return ``True`` when the provided span encodes a reward value."""
     maybe_reward = get_reward_value(span)

agentlightning/emitter/utils.py

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

agentlightning/env_var.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Environment variable managements."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+from typing import overload
+
+__all__ = [
+    "LightningEnvVar",
+    "resolve_bool_env_var",
+    "resolve_int_env_var",
+    "resolve_str_env_var",
+]
+
+
+class LightningEnvVar(Enum):
+    """Environment variables for Agent Lightning."""
+
+    AGL_EMITTER_DEBUG = "AGL_EMITTER_DEBUG"
+    """Enable debug logging for the emitter."""
+
+    AGL_MANAGED_STORE = "AGL_MANAGED_STORE"
+    """If yes, the [`ExecutionStrategy`][agentlightning.ExecutionStrategy]
+    constructs LightningStore wrappers automatically. When `False` the provided
+    `store` is passed directly to the bundles, allowing callers to manage
+    store wrappers manually."""
+
+    AGL_CURRENT_ROLE = "AGL_CURRENT_ROLE"
+    """Which side(s) to run in this process. Used in
+    [`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]."""
+
+    AGL_SERVER_HOST = "AGL_SERVER_HOST"
+    """Interface the [`LightningStoreServer`][agentlightning.LightningStoreServer]
+    binds to when running the algorithm bundle locally."""
+
+    AGL_SERVER_PORT = "AGL_SERVER_PORT"
+    """Port the [`LightningStoreServer`][agentlightning.LightningStoreServer] listens to."""
+
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+_FALSY_VALUES = {"0", "false", "no", "off"}
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, override: bool, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, *, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None: ...
+
+
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None:
+    """Resolve a boolean environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    normalized = env_value.strip().lower()
+    if normalized in _TRUTHY_VALUES:
+        return True
+    if normalized in _FALSY_VALUES:
+        return False
+
+    raise ValueError(f"{env_var.value} must be one of {_TRUTHY_VALUES} or {_FALSY_VALUES}")
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, override: int, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, *, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None: ...
+
+
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None:
+    """Resolve an integer environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    try:
+        return int(env_value)
+    except ValueError:
+        raise ValueError(f"{env_var.value} must be an integer")
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, override: str, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, *, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None: ...
+
+
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None:
+    """Resolve a string environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    return env_value

agentlightning/execution/base.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from typing import Protocol
 
 from agentlightning.store.base import LightningStore
@@ -13,47 +12,6 @@ from .events import ExecutionEvent
 logger = logging.getLogger(__name__)
 
 
-_TRUTHY_VALUES = {"1", "true", "yes", "on"}
-_FALSY_VALUES = {"0", "false", "no", "off"}
-
-
-def resolve_managed_store_flag(value: bool | None) -> bool:
-    """Determine whether execution helpers should wrap the provided store.
-
-    The helper first honours an explicit `value`. When `None` it falls back
-    to the `AGL_MANAGED_STORE` environment variable, accepting a variety
-    of truthy and falsy spellings. Missing environment configuration defaults to
-    `True` so that higher-level strategies create the appropriate client or
-    server wrappers automatically.
-
-    Args:
-        value: Optional override supplied by the caller.
-
-    Returns:
-        `True` when a managed store should be created around the provided
-        instance, otherwise `False`.
-
-    Raises:
-        ValueError: If `AGL_MANAGED_STORE` is set to an unsupported
-            value.
-    """
-
-    if value is not None:
-        return value
-
-    env_value = os.getenv("AGL_MANAGED_STORE")
-    if env_value is None:
-        return True
-
-    normalized = env_value.strip().lower()
-    if normalized in _TRUTHY_VALUES:
-        return True
-    if normalized in _FALSY_VALUES:
-        return False
-
-    raise ValueError("AGL_MANAGED_STORE must be one of 1, 0, true, false, yes, no, on, or off")
-
-
 class AlgorithmBundle(Protocol):
     """Callable bundle produced by [`Trainer`][agentlightning.Trainer].
 

agentlightning/execution/client_server.py

@@ -9,10 +9,11 @@ import time
 from multiprocessing.context import BaseContext
 from typing import Callable, Iterable, Literal, cast
 
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_int_env_var, resolve_str_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 
-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
 from .events import ExecutionEvent, MultiprocessingEvent
 
 logger = logging.getLogger(__name__)
@@ -99,44 +100,28 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                 By default, runner can exit gracefully with code 0 or terminated
                 by SIGTERM (-15).
         """
-        if role is None:
-            role_env = os.getenv("AGL_CURRENT_ROLE")
-            if role_env is None:
-                # Use both if not specified via env var or argument
-                role = "both"
-            elif role_env not in ("algorithm", "runner", "both"):
-                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
-            else:
-                role = role_env
-
-        if server_host is None:
-            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
-
-        if server_port is None:
-            server_port_env = os.getenv("AGL_SERVER_PORT")
-            if server_port_env is None:
-                server_port = 4747
-            else:
-                try:
-                    server_port = int(server_port_env)
-                except ValueError as exc:
-                    raise ValueError("AGL_SERVER_PORT must be an integer") from exc
-
-        self.role = role
+        resolved_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, override=role, fallback="both")
+        if resolved_role not in ("algorithm", "runner", "both"):
+            raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
+        self.role: Literal["algorithm", "runner", "both"] = resolved_role
         self.n_runners = n_runners
-        self.server_host = server_host
-        self.server_port = server_port
+        self.server_host = resolve_str_env_var(
+            LightningEnvVar.AGL_SERVER_HOST, override=server_host, fallback="localhost"
+        )
+        self.server_port = resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, override=server_port, fallback=4747)
         self.graceful_timeout = graceful_timeout
         self.terminate_timeout = terminate_timeout
         if main_process not in ("algorithm", "runner"):
             raise ValueError("main_process must be 'algorithm' or 'runner'")
         if main_process == "runner":
-            if role != "both":
+            if self.role != "both":
                 raise ValueError("main_process='runner' is only supported when role='both'")
             if n_runners != 1:
                 raise ValueError("main_process='runner' requires n_runners to be 1")
         self.main_process = main_process
-        self.managed_store = resolve_managed_store_flag(managed_store)
+        self.managed_store = resolve_bool_env_var(
+            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
+        )
         self.allowed_exit_codes = tuple(allowed_exit_codes)
 
     async def _execute_algorithm(

agentlightning/execution/shared_memory.py

@@ -7,10 +7,11 @@ from contextlib import suppress
 from queue import SimpleQueue
 from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
 
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.threading import LightningStoreThreaded
 
-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
 from .events import ExecutionEvent, ThreadingEvent
 
 logger = logging.getLogger(__name__)
@@ -62,7 +63,9 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         self.join_timeout = join_timeout
         self.graceful_delay = graceful_delay
         self.poll_interval = poll_interval
-        self.managed_store = resolve_managed_store_flag(managed_store)
+        self.managed_store = resolve_bool_env_var(
+            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
+        )
 
     async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
         """Run `coro` until it finishes or a cooperative stop is requested.

agentlightning/llm_proxy.py

@@ -47,7 +47,8 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.types import Scope
 
-from agentlightning.types import LLM, ProxyLLM, SpanNames
+from agentlightning.semconv import LightningResourceAttributes
+from agentlightning.types import LLM, ProxyLLM
 from agentlightning.utils.server_launcher import (
     LaunchMode,
     PythonServerLauncher,
@@ -174,6 +175,24 @@ class AddReturnTokenIds(CustomLogger):
         return {**data, "return_token_ids": True}
 
 
+class AddLogprobs(CustomLogger):
+    """LiteLLM logger hook to request logprobs from vLLM.
+
+    This mutates the outgoing request payload to include `logprobs=1`
+    for backends that support logprobs return (e.g., vLLM).
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload."""
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure logprobs are requested from the backend when supported.
+        return {**data, "logprobs": 1}
+
+
 class LightningSpanExporter(SpanExporter):
     """Buffered OTEL span exporter with subtree flushing and training-store sink.
 
@@ -396,9 +415,9 @@ class LightningSpanExporter(SpanExporter):
                     span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
                         Resource.create(
                             {
-                                SpanNames.ROLLOUT_ID: rollout_id,
-                                SpanNames.ATTEMPT_ID: attempt_id,
-                                SpanNames.SPAN_SEQUENCE_ID: sequence_id_decimal,
+                                LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                                LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                                LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id_decimal,
                             }
                         )
                     )
@@ -980,6 +999,7 @@ _MIDDLEWARE_REGISTRY: Dict[str, Type[BaseHTTPMiddleware]] = {
 
 _CALLBACK_REGISTRY = {
     "return_token_ids": AddReturnTokenIds,
+    "logprobs": AddLogprobs,
     "opentelemetry": LightningOpenTelemetry,
 }
 
@@ -1038,7 +1058,7 @@ class LLMProxy:
             Middlewares are the **first layer** of request processing. They are applied to all requests before the LiteLLM proxy.
         callbacks: List of LiteLLM callback classes or strings to register. You can specify the class aliases or classes that have been imported.
             If not provided, the default callbacks (AddReturnTokenIds and LightningOpenTelemetry) will be used.
-            Available callback aliases are: "return_token_ids", "opentelemetry".
+            Available callback aliases are: "return_token_ids", "opentelemetry", "logprobs".
     """
 
     def __init__(
@@ -1052,8 +1072,8 @@ class LLMProxy:
         num_workers: int = 1,
         launch_mode: LaunchMode = "mp",
         launcher_args: PythonServerLauncherArgs | None = None,
-        middlewares: List[Union[Type[BaseHTTPMiddleware], str]] | None = None,
-        callbacks: List[Union[Type[CustomLogger], str]] | None = None,
+        middlewares: Sequence[Union[Type[BaseHTTPMiddleware], str]] | None = None,
+        callbacks: Sequence[Union[Type[CustomLogger], str]] | None = None,
     ):
         self.store = store
 
@@ -1159,6 +1179,9 @@ class LLMProxy:
         if _global_llm_proxy is not None:
             logger.warning("A global LLMProxy is already set. Overwriting it with the new instance.")
 
+        # Patch for LiteLLM v1.80.6+: https://github.com/BerriAI/litellm/issues/17243
+        os.environ["USE_OTEL_LITELLM_REQUEST_SPAN"] = "true"
+
         # Set the global LLMProxy reference for middleware/exporter access.
         set_active_llm_proxy(self)
 

agentlightning/runner/agent.py

@@ -287,7 +287,7 @@ class LitAgentRunner(Runner[T_task]):
             # 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
-            reward_span = emit_reward(raw_result, auto_export=False)
+            reward_span = emit_reward(raw_result, propagate=False)
             # We add it to the store manually
             await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
             trace_spans.append(reward_span)

agentlightning/semconv.py

@@ -0,0 +1,144 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Semantic conventions for Agent-lightning spans.
+
+Conventions in this file are added on demand. We generally DO NOT add
+new semantic conventions unless it's absolutely needed for certain algorithms or scenarios.
+"""
+
+from enum import Enum
+
+from pydantic import BaseModel
+
+AGL_ANNOTATION = "agentlightning.annotation"
+"""Agent-lightning's standard span name for annotations.
+
+Annotations are minimal span units for rewards, tags, and metadatas.
+They are used to "annotate" a specific event or a part of rollout.
+"""
+
+AGL_MESSAGE = "agentlightning.message"
+"""Agent-lightning's standard span name for messages and logs."""
+
+AGL_OBJECT = "agentlightning.object"
+"""Agent-lightning's standard span name for customized objects."""
+
+AGL_EXCEPTION = "agentlightning.exception"
+"""Agent-lightning's standard span name for exceptions.
+
+Used by the exception emitter to record exception details.
+"""
+
+AGL_VIRTUAL = "agentlightning.virtual"
+"""Agent-lightning's standard span name for virtual operations.
+
+Mostly used in adapter when needing to represent the root or intermediate operations.
+"""
+
+
+class LightningResourceAttributes(Enum):
+    """Resource attribute names used in Agent-lightning spans."""
+
+    ROLLOUT_ID = "agentlightning.rollout_id"
+    """Resource name for rollout ID in Agent-lightning spans."""
+
+    ATTEMPT_ID = "agentlightning.attempt_id"
+    """Resource name for attempt ID in Agent-lightning spans."""
+
+    SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
+    """Resource name for span sequence ID in Agent-lightning spans."""
+
+
+class LightningSpanAttributes(Enum):
+    """Attribute names that commonly appear in Agent-lightning spans.
+
+    Exception types can't be found here because they are defined in OpenTelemetry's official semantic conventions.
+    """
+
+    REWARD = "agentlightning.reward"
+    """Attribute prefix for rewards-related data in reward spans.
+
+    It should be used as a prefix. For example, "agentlightning.reward.0.value" can
+    be used to track a specific metric. See [RewardAttributes][agentlightning.semconv.RewardAttributes].
+    """
+
+    LINK = "agentlightning.link"
+    """Attribute name for linking the current span to another span or other objects like requests/responses."""
+
+    TAG = "agentlightning.tag"
+    """Attribute name for tagging spans with customized strings."""
+
+    MESSAGE_BODY = "agentlightning.message.body"
+    """Attribute name for message text in message spans."""
+
+    OBJECT_TYPE = "agentlightning.object.type"
+    """Attribute name for object type (full qualified name) in object spans.
+
+    I think builtin types like str, int, bool, list, dict are self-explanatory and
+    should also be qualified to use here.
+    """
+
+    OBJECT_LITERAL = "agentlightning.object.literal"
+    """Attribute name for object literal value in object spans (for str, int, bool, ...)."""
+
+    OBJECT_JSON = "agentlightning.object.json"
+    """Attribute name for object serialized value (JSON) in object spans."""
+
+
+class RewardAttributes(Enum):
+    """Multi-dimensional reward attributes will look like:
+
+    ```json
+    {"agentlightning.reward.0.name": "efficiency", "agentlightning.reward.0.value": 0.75}
+    ```
+
+    The first reward in the reward list will automatically be the primary reward.
+    If the reward list has greater than 1, it shall be a multi-dimensional case.
+    """
+
+    REWARD_NAME = "name"
+    """Key for each dimension in multi-dimensional reward spans."""
+
+    REWARD_VALUE = "value"
+    """Value for each dimension in multi-dimensional reward spans."""
+
+
+class RewardPydanticModel(BaseModel):
+    """A stricter implementation of RewardAttributes used in otel helpers."""
+
+    name: str
+    """Name of the reward dimension."""
+
+    value: float
+    """Value of the reward dimension."""
+
+
+class LinkAttributes(Enum):
+    """Standard link types used in Agent-lightning spans.
+
+    The link is more powerful than [OpenTelemetry link](https://opentelemetry.io/docs/specs/otel/trace/api/#link)
+    in that it supports linking to a queryset of spans.
+    It can even link to span object that hasn't been emitted yet.
+    """
+
+    KEY_MATCH = "key_match"
+    """Linking to spans with matching attribute keys.
+
+    `trace_id` and `span_id` are reserved and will be used to link to specific spans directly.
+
+    For example, it can be `gen_ai.response.id` if intended to be link to a chat completion response span.
+    Or it can be `span_id` to link to a specific span by its ID.
+    """
+
+    VALUE_MATCH = "value_match"
+    """Linking to spans with corresponding attribute values on those keys."""
+
+
+class LinkPydanticModel(BaseModel):
+    """A stricter implementation of LinkAttributes used in otel helpers."""
+
+    key_match: str
+    """The attribute key to match on the target spans."""
+
+    value_match: str
+    """The attribute value to match on the target spans."""

agentlightning/store/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import LightningStore, LightningStoreCapabilities
+from .base import LightningStore, LightningStoreCapabilities, LightningStoreStatistics
 from .client_server import LightningStoreClient, LightningStoreServer
 from .collection_based import CollectionBasedLightningStore
 from .memory import InMemoryLightningStore
@@ -9,6 +9,7 @@ from .threading import LightningStoreThreaded
 __all__ = [
     "LightningStore",
     "LightningStoreCapabilities",
+    "LightningStoreStatistics",
     "LightningStoreClient",
     "LightningStoreServer",
     "InMemoryLightningStore",

agentlightning/store/base.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal, Optional, Sequence, TypedDict
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -70,6 +70,35 @@ class LightningStoreCapabilities(TypedDict, total=False):
     """Whether the store supports OTLP/HTTP traces."""
 
 
+class LightningStoreStatistics(TypedDict, total=False):
+    """Statistics of a LightningStore implementation."""
+
+    name: str
+    """Name of the store implementation."""
+    total_rollouts: int
+    """Total number of rollouts in the store."""
+    total_attempts: int
+    """Total number of attempts in the store."""
+    total_spans: int
+    """Total number of spans in the store."""
+    total_resources: int
+    """Total number of resources in the store."""
+    total_workers: int
+    """Total number of workers in the store."""
+    uptime: float
+    """Uptime of since the store has been started."""
+
+    # Memory-related statistics
+    total_span_bytes: int
+    """Total number of bytes of spans in the store."""
+    eviction_threshold_bytes: int
+    """Eviction threshold for spans in bytes."""
+    safe_threshold_bytes: int
+    """Safe threshold for spans in bytes."""
+    memory_capacity_bytes: int
+    """Memory capacity of the store in bytes."""
+
+
 class LightningStore:
     """Contract for the persistent control-plane that coordinates training rollouts.
 
@@ -102,6 +131,12 @@ class LightningStore:
             otlp_traces=False,
         )
 
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        return {
+            "name": self.__class__.__name__,
+        }
+
     def otlp_traces_endpoint(self) -> str:
         """Return the OTLP/HTTP traces endpoint of the store.
 
@@ -237,7 +272,15 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def add_span(self, span: Span) -> Span:
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        """Persist a sequence of pre-constructed spans emitted during rollout execution.
+
+        Implementations can simply delegate to [`add_span()`][agentlightning.LightningStore.add_span] for each span.
+        However, if the store supports bulk insertion, it can implement this method to improve performance.
+        """
+        raise NotImplementedError()
+
+    async def add_span(self, span: Span) -> Optional[Span]:
         """Persist a pre-constructed span emitted during rollout execution.
 
         The provided [`Span`][agentlightning.Span] must already contain the `rollout_id`,
@@ -254,6 +297,7 @@ class LightningStore:
 
         Returns:
             The stored span record (implementations may return a copy).
+            Return `None` if the span was not added due to a duplicate.
 
         Raises:
             NotImplementedError: Subclasses must implement span persistence.
@@ -267,7 +311,7 @@ class LightningStore:
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         """Convert and persist an OpenTelemetry span for a particular attempt.
 
         Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
@@ -284,7 +328,7 @@ class LightningStore:
                 automatically.
 
         Returns:
-            The stored span record.
+            The stored span record. Return `None` if the span was not added due to a duplicate.
 
         Raises:
             NotImplementedError: Subclasses must implement span persistence.
@@ -485,6 +529,20 @@ class LightningStore:
         """
         raise NotImplementedError()
 
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        """Bulk allocate the next strictly increasing sequence number used to order spans.
+
+        Implementations may delegate to [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
+        for each rollout and attempt.
+
+        Args:
+            rollout_attempt_ids: List of tuples of rollout and attempt identifiers.
+
+        Returns:
+            List of sequence numbers.
+        """
+        raise NotImplementedError()
+
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """Block until the targeted rollouts reach a terminal status or the timeout expires.
 

agentlightning/store/client_server.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import os
+import re
 import threading
 import time
 import traceback
@@ -23,6 +24,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 import aiohttp
@@ -59,7 +61,8 @@ from agentlightning.types import (
 from agentlightning.utils.otlp import handle_otlp_export, spans_from_proto
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncher, PythonServerLauncherArgs
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
+from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
+from .utils import LATENCY_BUCKETS
 
 server_logger = logging.getLogger("agentlightning.store.server")
 client_logger = logging.getLogger("agentlightning.store.client")
@@ -637,6 +640,10 @@ class LightningStoreServer(LightningStore):
                 heartbeat_stats=_get_mandatory_field_or_unset(request, "heartbeat_stats"),
             )
 
+        @api.get(API_AGL_PREFIX + "/statistics", response_model=Dict[str, Any])
+        async def get_statistics():  # pyright: ignore[reportUnusedFunction]
+            return await self.statistics()
+
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", response_model=PaginatedResult[Attempt])
         async def query_attempts(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, params: QueryAttemptsRequest = Depends()
@@ -686,7 +693,7 @@ class LightningStoreServer(LightningStore):
         async def get_resources_by_id(resources_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_resources_by_id(resources_id)
 
-        @api.post(API_AGL_PREFIX + "/spans", status_code=201, response_model=Span)
+        @api.post(API_AGL_PREFIX + "/spans", status_code=201, response_model=Optional[Span])
         async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
             return await self.add_span(span)
 
@@ -733,31 +740,58 @@ class LightningStoreServer(LightningStore):
     def _setup_prometheus(self, api: APIRouter, app: FastAPI):
         """Setup Prometheus metrics endpoints."""
         try:
+            from prometheus_client import make_asgi_app  # type: ignore
             from prometheus_client import (
-                CONTENT_TYPE_LATEST,
+                REGISTRY,
+                CollectorRegistry,
                 Counter,
                 Histogram,
-                generate_latest,
+                multiprocess,
             )
         except ImportError:
             raise ImportError(
                 "Prometheus client is not installed. Please either install it or set prometheus to False."
             )
 
+        # Multi-process mode: https://prometheus.github.io/client_python/multiprocess/
+        is_multiprocess = self.launcher_args.launch_mode == "mp" and self.launcher_args.n_workers > 1
+        if is_multiprocess:
+            registry = CollectorRegistry()
+            multiprocess.MultiProcessCollector(registry)
+        else:
+            registry = REGISTRY
+
         HTTP_REQUESTS = Counter(
             "http_requests_total",
             "Total HTTP requests",
             ["method", "path", "status_code"],
         )
 
-        # TODO: For multi-process scenarios, should use prometheus_client.multiprocess mode.
         HTTP_LATENCY = Histogram(
             "http_request_duration_seconds",
             "Latency of HTTP requests",
             ["method", "path"],
-            buckets=[0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10],
+            buckets=LATENCY_BUCKETS,
         )
 
+        def get_template_path(path: str) -> str:
+            # 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"):
+                return path
+            elif "enqueue" in path or "dequeue" in path:
+                return path
+
+            # Handle generic IDs
+            # (Order matters: longest paths first or lookaheads)
+            path = re.sub(r"/attempts/[^/]+$", "/attempts/{attempt_id}", path)
+            path = re.sub(r"/rollouts/[^/]+", "/rollouts/{rollout_id}", path)  # Handles root and middle
+            path = re.sub(r"/resources/[^/]+$", "/resources/{resources_id}", path)
+            path = re.sub(r"/workers/[^/]+$", "/workers/{worker_id}", path)
+
+            return path
+
         @app.middleware("http")
         async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
@@ -766,7 +800,8 @@ class LightningStoreServer(LightningStore):
             response = await call_next(request)
             elapsed = time.perf_counter() - start
 
-            path = request.url.path
+            # Strip the ID-specific URL parts
+            path = get_template_path(request.url.path)
             method = request.method
             status = response.status_code
 
@@ -775,24 +810,22 @@ class LightningStoreServer(LightningStore):
 
             return response
 
-        @api.get("/prometheus")
-        async def prometheus_metrics():  # pyright: ignore[reportUnusedFunction]
-            return Response(
-                content=generate_latest(),
-                media_type=CONTENT_TYPE_LATEST,
-            )
+        metrics_app = make_asgi_app(registry=registry)  # type: ignore
+
+        # This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
+        app.mount(api.prefix + "/prometheus", metrics_app)  # pyright: ignore[reportUnknownArgumentType]
 
     def _setup_otlp(self, api: APIRouter):
         """Setup OTLP endpoints."""
 
         async def _trace_handler(request: PbExportTraceServiceRequest) -> None:
-            spans = await spans_from_proto(request, self)
+            spans = await spans_from_proto(request, self.get_many_span_sequence_ids)
             server_logger.debug(f"Received {len(spans)} OTLP spans: {', '.join([span.name for span in spans])}")
-            for span in spans:
-                await self.add_span(span)
+            await self.add_many_spans(spans)
 
         # Reserved methods for OTEL traces
         # https://opentelemetry.io/docs/specs/otlp/#otlphttp-request
+        # This is currently the recommended path for Otel compatibility and bulk-insertion support.
         @api.post("/traces")
         async def otlp_traces(request: Request):  # pyright: ignore[reportUnusedFunction]
             return await handle_otlp_export(
@@ -844,6 +877,8 @@ class LightningStoreServer(LightningStore):
 
         @self.app.get("/{full_path:path}", include_in_schema=False)
         def spa_fallback(full_path: str):  # pyright: ignore[reportUnusedFunction]
+            if full_path.startswith("v1/"):
+                raise HTTPException(status_code=404, detail="Not Found")
             # Let the frontend router handle it
             return FileResponse(index_file)
 
@@ -889,6 +924,9 @@ class LightningStoreServer(LightningStore):
             self._client = LightningStoreClient(self.endpoint)
         return await getattr(self._client, method_name)(*args, **kwargs)
 
+    async def statistics(self) -> LightningStoreStatistics:
+        return await self._call_store_method("statistics")
+
     async def start_rollout(
         self,
         input: TaskInput,
@@ -1013,19 +1051,25 @@ class LightningStoreServer(LightningStore):
     async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
         return await self._call_store_method("get_latest_resources")
 
-    async def add_span(self, span: Span) -> Span:
+    async def add_span(self, span: Span) -> Optional[Span]:
         return await self._call_store_method("add_span", span)
 
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        return await self._call_store_method("add_many_spans", spans)
+
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
         return await self._call_store_method("get_next_span_sequence_id", rollout_id, attempt_id)
 
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        return await self._call_store_method("get_many_span_sequence_ids", rollout_attempt_ids)
+
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         return await self._call_store_method(
             "add_otel_span",
             rollout_id,
@@ -1208,6 +1252,10 @@ class LightningStoreClient(LightningStore):
         """Return the OTLP/HTTP traces endpoint of the store."""
         return f"{self.server_address_root}/v1/traces"
 
+    async def statistics(self) -> LightningStoreStatistics:
+        payload = await self._request_json("get", "/statistics")
+        return cast(LightningStoreStatistics, payload)
+
     def __getstate__(self):
         """
         When LightningStoreClient is pickled (e.g., passed to a subprocess), we only
@@ -1570,7 +1618,9 @@ class LightningStoreClient(LightningStore):
         """
         try:
             data = await self._request_json("get", f"/rollouts/{rollout_id}")
-            if isinstance(data, dict) and "attempt" in data:
+            if data is None:
+                return None
+            elif isinstance(data, dict) and "attempt" in data:
                 return AttemptedRollout.model_validate(data)
             else:
                 return Rollout.model_validate(data)
@@ -1660,9 +1710,17 @@ class LightningStoreClient(LightningStore):
             client_logger.error(f"get_latest_resources failed after all retries: {e}", exc_info=True)
             return None
 
-    async def add_span(self, span: Span) -> Span:
+    async def add_span(self, span: Span) -> Optional[Span]:
         data = await self._request_json("post", "/spans", json=span.model_dump(mode="json"))
-        return Span.model_validate(data)
+        return Span.model_validate(data) if data is not None else None
+
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        result: List[Span] = []
+        for span in spans:
+            ret = await self.add_span(span)
+            if ret is not None:
+                result.append(ret)
+        return result
 
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
         data = await self._request_json(
@@ -1673,13 +1731,19 @@ class LightningStoreClient(LightningStore):
         response = NextSequenceIdResponse.model_validate(data)
         return response.sequence_id
 
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        return [
+            await self.get_next_span_sequence_id(rollout_id, attempt_id)
+            for rollout_id, attempt_id in rollout_attempt_ids
+        ]
+
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         # unchanged logic, now benefits from retries inside add_span/get_next_span_sequence_id
         if sequence_id is None:
             sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
@@ -1689,9 +1753,7 @@ class LightningStoreClient(LightningStore):
             attempt_id=attempt_id,
             sequence_id=sequence_id,
         )
-        print("created span", span)
-        await self.add_span(span)
-        return span
+        return await self.add_span(span)
 
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """Wait for rollouts to complete.

agentlightning/store/collection/memory.py

@@ -539,7 +539,19 @@ class ListBasedCollection(Collection[T]):
         Raises:
             ValueError: If any item with the same primary keys already exists.
         """
+        seen_keys: set[Tuple[Any, ...]] = set()
+        prepared: List[T] = []
         for item in items:
+            self._ensure_item_type(item)
+            key_values = self._extract_primary_key_values(item)
+            if key_values in seen_keys:
+                raise ValueError(
+                    f"Insert payload contains duplicate primary key(s): {self._render_key_values(key_values)}"
+                )
+            seen_keys.add(key_values)
+            prepared.append(item)
+
+        for item in prepared:
             self._mutate_single(item, mode="insert")
 
     async def update(self, items: Sequence[T]) -> None:

agentlightning/store/collection/mongo.py

@@ -3,9 +3,14 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
+import functools
+import inspect
 import logging
+import random
 import re
 import threading
+import time
 from contextlib import asynccontextmanager
 from datetime import datetime
 from typing import (
@@ -36,6 +41,8 @@ from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.errors import CollectionInvalid, ConnectionFailure, DuplicateKeyError, OperationFailure, PyMongoError
 from pymongo.read_concern import ReadConcern
 
+from agentlightning.store.base import LightningStore
+from agentlightning.store.utils import LATENCY_BUCKETS
 from agentlightning.types import (
     Attempt,
     FilterOptions,
@@ -47,7 +54,14 @@ from agentlightning.types import (
     Worker,
 )
 
-from .base import Collection, KeyValue, LightningCollections, Queue, normalize_filter_options, resolve_sort_options
+from .base import (
+    Collection,
+    KeyValue,
+    LightningCollections,
+    Queue,
+    normalize_filter_options,
+    resolve_sort_options,
+)
 
 T_model = TypeVar("T_model", bound=BaseModel)
 
@@ -55,11 +69,244 @@ T_generic = TypeVar("T_generic")
 
 T_mapping = TypeVar("T_mapping", bound=Mapping[str, Any])
 
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
 K = TypeVar("K")
 V = TypeVar("V")
 
 logger = logging.getLogger(__name__)
 
+_OPERATION_CONTEXT: contextvars.ContextVar["_MongoOperationContext | None"] = contextvars.ContextVar(
+    "_mongo_operation_context", default=None
+)
+
+_LIGHTNING_STORE_PUBLIC_METHODS = frozenset(
+    [name for name, value in LightningStore.__dict__.items() if not name.startswith("_") and callable(value)]
+    + ["_healthcheck"]
+)
+
+_UNKNOWN_STORE_METHOD = "unknown"
+
+
+def _nearest_lightning_store_method_from_stack() -> str:
+    """Stack introspection so that we capture the nearest public API method from the
+    call stack whenever metrics are recorded."""
+    frame = inspect.currentframe()
+    try:
+        if frame is None:
+            return _UNKNOWN_STORE_METHOD
+        frame = frame.f_back
+        while frame is not None:
+            self_obj = frame.f_locals.get("self")
+            method_name = frame.f_locals.get("method_name")
+            if method_name in _LIGHTNING_STORE_PUBLIC_METHODS and isinstance(self_obj, LightningStore):
+                return method_name
+            frame = frame.f_back
+        return _UNKNOWN_STORE_METHOD
+    except Exception:
+        return _UNKNOWN_STORE_METHOD
+    finally:
+        del frame
+
+
+class MongoOperationPrometheusTracker:
+    """A tracker for MongoDB operations metrics.
+
+    All classes should share one single instance of this tracker.
+    """
+
+    def __init__(self, enabled: bool):
+        self._enabled = enabled
+
+        if enabled:
+            from prometheus_client import Counter, Histogram
+
+            base_labels = ["operation", "database", "collection", "store_method"]
+            self._latency_metric = Histogram(
+                "mongo_operation_duration_seconds",
+                "Latency of MongoDB operations",
+                base_labels,
+                buckets=LATENCY_BUCKETS,
+            )
+            self._total_metric = Counter(
+                "mongo_operation_total",
+                "Total MongoDB operations",
+                base_labels + ["status"],
+            )
+            self._error_metric = Counter(
+                "mongo_operation_errors_total",
+                "Total MongoDB operations that failed",
+                base_labels + ["error_type"],
+            )
+            self._num_attempts_metric = Histogram(
+                "mongo_operation_num_attempts",
+                "Number of attempts for MongoDB operations",
+                base_labels,
+                buckets=list(range(10)) + list(range(10, 100, 5)),
+            )
+
+    def track(self, operation: str, database: str, collection: str) -> _MongoOperationContext | _DummyOperationContext:
+        if not self._enabled:
+            return _DummyOperationContext()
+        return _MongoOperationContext(self, operation, database, collection)
+
+    @staticmethod
+    def classify_error(exc: BaseException | None) -> str:
+        if exc is None:
+            return "Other"
+        is_transient = isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError")
+        if isinstance(exc, OperationFailure):
+            if is_transient:
+                return f"OperationFailure-{exc.code}-Transient"
+            else:
+                return f"OperationFailure-{exc.code}"
+        if isinstance(exc, DuplicateKeyError):
+            return "DuplicateKeyError-Transient" if is_transient else "DuplicateKeyError"
+        if isinstance(exc, PyMongoError):
+            if is_transient:
+                return f"{exc.__class__.__name__}-Transient"
+            else:
+                return exc.__class__.__name__
+        if isinstance(exc, ConnectionFailure):
+            return "ConnectionFailure-Transient" if is_transient else "ConnectionFailure"
+        return "Other-Transient" if is_transient else "Other"
+
+    def observe(
+        self,
+        *,
+        operation: str,
+        database: str,
+        collection: str,
+        elapsed: float,
+        status: str,
+        error_type: str | None,
+        num_attempts: int | None = None,
+    ) -> None:
+        if not self._enabled:
+            return
+
+        store_method = _nearest_lightning_store_method_from_stack()
+        self._total_metric.labels(operation, database, collection, store_method, status).inc()
+        self._latency_metric.labels(operation, database, collection, store_method).observe(elapsed)
+        if status == "error" and error_type:
+            self._error_metric.labels(operation, database, collection, store_method, error_type).inc()
+        if num_attempts is not None:
+            self._num_attempts_metric.labels(operation, database, collection, store_method).observe(num_attempts)
+
+
+class _MongoOperationContext:
+    """A context manager for tracking MongoDB operations.
+
+    Used via:
+
+    ```python
+    with self.tracker.track("insert", "database", "collection") as track_context:
+        try:
+            await collection.insert_one({})
+        except Exception as exc:
+            # For errors that can be ignored, report the error to the tracker.
+            track_context.report_error(exc)
+    ```
+    """
+
+    def __init__(self, tracker: MongoOperationPrometheusTracker, operation: str, database: str, collection: str):
+        self._tracker = tracker
+        self._operation = operation
+        self._database = database
+        self._collection = collection
+        self._start: float = 0.0
+        self._active: bool = False
+        self._error_type: str | None = None
+        self._num_attempts: int | None = None
+
+    def __enter__(self) -> "_MongoOperationContext":
+        self._active = True
+        self._start = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> bool:
+        if not self._active:
+            return False
+
+        elapsed = time.perf_counter() - self._start
+
+        # Try to classify the error
+        if self._error_type is not None:
+            error_type = self._error_type
+        elif exc is not None:
+            error_type = self._tracker.classify_error(exc)
+        else:
+            error_type = None
+
+        status = "ok" if error_type is None else "error"
+        self._tracker.observe(
+            operation=self._operation,
+            database=self._database,
+            collection=self._collection,
+            elapsed=elapsed,
+            status=status,
+            error_type=error_type,
+            num_attempts=self._num_attempts,
+        )
+        return False
+
+    def report_error(self, exc: BaseException) -> None:
+        """Used to report errors that occurred in the middle of an operation."""
+        self._error_type = self._tracker.classify_error(exc)
+
+    def report_num_attempts(self, num_attempts: int) -> None:
+        """Used to report the number of attempts that occurred in the middle of an operation."""
+        self._num_attempts = num_attempts
+
+
+class _DummyOperationContext:
+    """A dummy context manager that does nothing, but compatible with _MongoOperationContext."""
+
+    def __init__(self):
+        pass
+
+    def __enter__(self) -> "_DummyOperationContext":
+        return self
+
+    def __exit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> bool:
+        return False
+
+    def report_error(self, exc: BaseException) -> None:
+        pass
+
+    def report_num_attempts(self, num_attempts: int) -> None:
+        pass
+
+
+def _mongo_operation(operation: str) -> Callable[[T_callable], T_callable]:
+    def decorator(func: T_callable) -> T_callable:
+        if not asyncio.iscoroutinefunction(func):
+            raise TypeError(f"_mongo_operation decorator requires coroutine functions, got {func.__name__}")
+
+        @functools.wraps(func)
+        async def wrapper(
+            self: (
+                MongoBasedCollection[T_model]
+                | MongoBasedQueue[T_generic]
+                | MongoBasedKeyValue[K, V]
+                | MongoLightningCollections
+            ),
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
+            tracker = self._prometheus_tracker  # pyright: ignore[reportPrivateUsage]
+            if not tracker._enabled:  # pyright: ignore[reportPrivateUsage]
+                # Skip the tracking because tracking is not configured
+                return await func(self, *args, **kwargs)
+            with tracker.track(
+                operation, self._database_name, self._collection_name  # pyright: ignore[reportPrivateUsage]
+            ):
+                return await func(self, *args, **kwargs)
+
+        return wrapper  # type: ignore
+
+    return decorator
+
 
 def _field_ops_to_conditions(field: str, ops: Mapping[str, Any]) -> List[Dict[str, Any]]:
     """Convert a FilterField (ops) into one or more Mongo conditions."""
@@ -308,6 +555,7 @@ class MongoBasedCollection(Collection[T_model]):
         primary_keys: Sequence[str],
         item_type: Type[T_model],
         extra_indexes: Sequence[Sequence[str]] = [],
+        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
     ):
         if isinstance(client_pool, AsyncMongoClient):
             self._client_pool = MongoClientPool(client_pool)
@@ -319,6 +567,9 @@ class MongoBasedCollection(Collection[T_model]):
         self._collection_created = False
         self._extra_indexes = [list(index) for index in extra_indexes]
         self._session: Optional[AsyncClientSession] = None
+        self._prometheus_tracker = (
+            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
+        )
 
         if not primary_keys:
             raise ValueError("primary_keys must be non-empty")
@@ -328,6 +579,7 @@ class MongoBasedCollection(Collection[T_model]):
             raise ValueError(f"item_type must be a subclass of BaseModel, got {item_type.__name__}")
         self._item_type = item_type
 
+    @_mongo_operation("ensure_collection")
     async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing MongoDB collection exists (and optionally its indexes).
 
@@ -353,6 +605,7 @@ class MongoBasedCollection(Collection[T_model]):
             primary_keys=self._primary_keys,
             item_type=self._item_type,
             extra_indexes=self._extra_indexes,
+            prometheus_tracker=self._prometheus_tracker,
         )
         collection._collection_created = self._collection_created
         collection._session = session
@@ -365,6 +618,7 @@ class MongoBasedCollection(Collection[T_model]):
     def item_type(self) -> Type[T_model]:
         return self._item_type
 
+    @_mongo_operation("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents({"partition_id": self._partition_id}, session=self._session)
@@ -379,6 +633,9 @@ class MongoBasedCollection(Collection[T_model]):
         pk_filter.update({pk: data[pk] for pk in self._primary_keys})
         return pk_filter
 
+    def _render_pk_values(self, values: Sequence[Any]) -> str:
+        return ", ".join(f"{pk}={value!r}" for pk, value in zip(self._primary_keys, values))
+
     def _ensure_item_type(self, item: T_model) -> None:
         if not isinstance(item, self._item_type):
             raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
@@ -406,6 +663,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return combined
 
+    @_mongo_operation("query")
     async def query(
         self,
         filter: Optional[FilterOptions] = None,
@@ -457,6 +715,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return PaginatedResult[T_model](items=items, limit=limit, offset=offset, total=total)
 
+    @_mongo_operation("get")
     async def get(
         self,
         filter: Optional[FilterOptions] = None,
@@ -465,19 +724,25 @@ class MongoBasedCollection(Collection[T_model]):
         result = await self.query(filter=filter, sort=sort, limit=1, offset=0)
         return result.items[0] if result.items else None
 
+    @_mongo_operation("insert")
     async def insert(self, items: Sequence[T_model]) -> None:
         if not items:
             return
 
         collection = await self.ensure_collection()
         docs: List[Mapping[str, Any]] = []
+        pk_conditions: List[Dict[str, Any]] = []
+        seen_primary_keys: set[Tuple[Any, ...]] = set()
         for item in items:
             self._ensure_item_type(item)
-            # Pre-check for existence to provide a clearer ValueError
             pk_filter = self._pk_filter(item)
-            existing = await collection.find_one(pk_filter, session=self._session)
-            if existing is not None:
-                raise ValueError(f"Item with primary key(s) {pk_filter} already exists")
+            pk_values = tuple(pk_filter[pk] for pk in self._primary_keys)
+            if pk_values in seen_primary_keys:
+                raise ValueError(
+                    f"Insert payload contains duplicate primary key(s): {self._render_pk_values(pk_values)}"
+                )
+            seen_primary_keys.add(pk_values)
+            pk_conditions.append({pk: pk_filter[pk] for pk in self._primary_keys})
 
             doc = item.model_dump()
             doc["partition_id"] = self._partition_id
@@ -486,12 +751,28 @@ class MongoBasedCollection(Collection[T_model]):
         if not docs:
             return
 
-        try:
-            await collection.insert_many(docs, session=self._session)
-        except DuplicateKeyError as exc:
-            # In case the DB enforces uniqueness via index, normalize to ValueError
-            raise ValueError("Duplicate key error while inserting items") from exc
+        if len(pk_conditions) == 1:
+            existing_filter: Dict[str, Any] = {"partition_id": self._partition_id, **pk_conditions[0]}
+        else:
+            existing_filter = {"partition_id": self._partition_id, "$or": pk_conditions}
 
+        with self._prometheus_tracker.track("insert__find_existing", self._database_name, self._collection_name):
+            existing = await collection.find_one(existing_filter, session=self._session)
+            if existing is not None:
+                existing_values = tuple(existing.get(pk) for pk in self._primary_keys)
+                raise ValueError(f"Item with primary key(s) {self._render_pk_values(existing_values)} already exists")
+
+        with self._prometheus_tracker.track(
+            "insert__insert_many", self._database_name, self._collection_name
+        ) as tracker:
+            try:
+                await collection.insert_many(docs, session=self._session)
+            except DuplicateKeyError as exc:
+                # In case the DB enforces uniqueness via index, normalize to ValueError
+                tracker.report_error(exc)
+                raise ValueError("Duplicate key error while inserting items") from exc
+
+    @_mongo_operation("update")
     async def update(self, items: Sequence[T_model]) -> None:
         if not items:
             return
@@ -502,10 +783,12 @@ class MongoBasedCollection(Collection[T_model]):
             pk_filter = self._pk_filter(item)
             doc = item.model_dump()
             doc["partition_id"] = self._partition_id
-            result = await collection.replace_one(pk_filter, doc, session=self._session)
+            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:
                 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
@@ -516,8 +799,10 @@ class MongoBasedCollection(Collection[T_model]):
             pk_filter = self._pk_filter(item)
             doc = item.model_dump()
             doc["partition_id"] = self._partition_id
-            await collection.replace_one(pk_filter, doc, upsert=True, session=self._session)
+            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)
 
+    @_mongo_operation("delete")
     async def delete(self, items: Sequence[T_model]) -> None:
         if not items:
             return
@@ -526,7 +811,8 @@ class MongoBasedCollection(Collection[T_model]):
         for item in items:
             self._ensure_item_type(item)
             pk_filter = self._pk_filter(item)
-            result = await collection.delete_one(pk_filter, session=self._session)
+            with self._prometheus_tracker.track("delete__delete_one", self._database_name, self._collection_name):
+                result = await collection.delete_one(pk_filter, session=self._session)
             if result.deleted_count == 0:
                 raise ValueError(f"Item with primary key(s) {pk_filter} does not exist")
 
@@ -544,6 +830,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         collection_name: str,
         partition_id: str,
         item_type: Type[T_generic],
+        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
     ) -> None:
         """
         Args:
@@ -565,10 +852,14 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         self._collection_created = False
 
         self._session: Optional[AsyncClientSession] = None
+        self._prometheus_tracker = (
+            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
+        )
 
     def item_type(self) -> Type[T_generic]:
         return self._item_type
 
+    @_mongo_operation("ensure_collection")
     async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing collection exists.
 
@@ -588,11 +879,13 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
             collection_name=self._collection_name,
             partition_id=self._partition_id,
             item_type=self._item_type,
+            prometheus_tracker=self._prometheus_tracker,
         )
         queue._collection_created = self._collection_created
         queue._session = session
         return queue
 
+    @_mongo_operation("has")
     async def has(self, item: T_generic) -> bool:
         collection = await self.ensure_collection()
         encoded = self._adapter.dump_python(item, mode="python")
@@ -606,6 +899,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         )
         return doc is not None
 
+    @_mongo_operation("enqueue")
     async def enqueue(self, items: Sequence[T_generic]) -> Sequence[T_generic]:
         if not items:
             return []
@@ -624,9 +918,11 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
                 }
             )
 
-        await collection.insert_many(docs, session=self._session)
+        with self._prometheus_tracker.track("enqueue__insert_many", self._database_name, self._collection_name):
+            await collection.insert_many(docs, session=self._session)
         return list(items)
 
+    @_mongo_operation("dequeue")
     async def dequeue(self, limit: int = 1) -> Sequence[T_generic]:
         if limit <= 0:
             return []
@@ -636,16 +932,19 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         # Atomic claim loop using find_one_and_update
         for _ in range(limit):
-            doc = await collection.find_one_and_update(
-                {
-                    "partition_id": self._partition_id,
-                    "consumed": False,
-                },
-                {"$set": {"consumed": True}},
-                sort=[("_id", 1)],  # FIFO using insertion order
-                return_document=True,
-                session=self._session,
-            )
+            with self._prometheus_tracker.track(
+                "dequeue__find_one_and_update", self._database_name, self._collection_name
+            ):
+                doc = await collection.find_one_and_update(
+                    {
+                        "partition_id": self._partition_id,
+                        "consumed": False,
+                    },
+                    {"$set": {"consumed": True}},
+                    sort=[("_id", 1)],  # FIFO using insertion order
+                    return_document=True,
+                    session=self._session,
+                )
             if doc is None:  # type: ignore
                 # No more items to dequeue
                 break
@@ -656,22 +955,24 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         return results
 
+    @_mongo_operation("peek")
     async def peek(self, limit: int = 1) -> Sequence[T_generic]:
         if limit <= 0:
             return []
 
         collection = await self.ensure_collection()
-        cursor = (
-            collection.find(
-                {
-                    "partition_id": self._partition_id,
-                    "consumed": False,
-                },
-                session=self._session,
+        with self._prometheus_tracker.track("peek__find", self._database_name, self._collection_name):
+            cursor = (
+                collection.find(
+                    {
+                        "partition_id": self._partition_id,
+                        "consumed": False,
+                    },
+                    session=self._session,
+                )
+                .sort("_id", 1)
+                .limit(limit)
             )
-            .sort("_id", 1)
-            .limit(limit)
-        )
 
         items: list[T_generic] = []
         async for doc in cursor:
@@ -680,6 +981,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         return items
 
+    @_mongo_operation("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents(
@@ -702,6 +1004,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         partition_id: str,
         key_type: Type[K],
         value_type: Type[V],
+        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
     ) -> None:
         """
         Args:
@@ -726,7 +1029,11 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         self._collection_created = False
 
         self._session: Optional[AsyncClientSession] = None
+        self._prometheus_tracker = (
+            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
+        )
 
+    @_mongo_operation("ensure_collection")
     async def ensure_collection(self, *, create_indexes: bool = True) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing collection exists (and optionally its indexes)."""
         if not self._collection_created:
@@ -744,12 +1051,14 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
             partition_id=self._partition_id,
             key_type=self._key_type,
             value_type=self._value_type,
+            prometheus_tracker=self._prometheus_tracker,
         )
         key_value._collection_created = self._collection_created
         key_value._session = session
 
         return key_value
 
+    @_mongo_operation("has")
     async def has(self, key: K) -> bool:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -762,6 +1071,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         )
         return doc is not None
 
+    @_mongo_operation("get")
     async def get(self, key: K, default: V | None = None) -> V | None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -778,28 +1088,34 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         raw_value = doc["value"]
         return self._value_adapter.validate_python(raw_value)
 
+    @_mongo_operation("set")
     async def set(self, key: K, value: V) -> None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
         encoded_value = self._value_adapter.dump_python(value, mode="python")
-        try:
-            await collection.replace_one(
-                {
-                    "partition_id": self._partition_id,
-                    "key": encoded_key,
-                },
-                {
-                    "partition_id": self._partition_id,
-                    "key": encoded_key,
-                    "value": encoded_value,
-                },
-                upsert=True,
-                session=self._session,
-            )
-        except DuplicateKeyError as exc:
-            # Very unlikely with replace_one+upsert, but normalize anyway.
-            raise ValueError("Duplicate key error while setting key-value item") from exc
+        with self._prometheus_tracker.track(
+            "upsert__replace_one", self._database_name, self._collection_name
+        ) as tracer:
+            try:
+                await collection.replace_one(
+                    {
+                        "partition_id": self._partition_id,
+                        "key": encoded_key,
+                    },
+                    {
+                        "partition_id": self._partition_id,
+                        "key": encoded_key,
+                        "value": encoded_value,
+                    },
+                    upsert=True,
+                    session=self._session,
+                )
+            except DuplicateKeyError as exc:
+                # Very unlikely with replace_one+upsert, but normalize anyway.
+                tracer.report_error(exc)
+                raise ValueError("Duplicate key error while setting key-value item") from exc
 
+    @_mongo_operation("pop")
     async def pop(self, key: K, default: V | None = None) -> V | None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -816,6 +1132,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         raw_value = doc["value"]
         return self._value_adapter.validate_python(raw_value)
 
+    @_mongo_operation("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents(
@@ -844,10 +1161,16 @@ class MongoLightningCollections(LightningCollections):
         workers: Optional[MongoBasedCollection[Worker]] = None,
         rollout_queue: Optional[MongoBasedQueue[str]] = None,
         span_sequence_ids: Optional[MongoBasedKeyValue[str, int]] = None,
+        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
     ):
         self._client_pool = client_pool
         self._database_name = database_name
+        self._collection_name = "collections"  # Special collection name for tracking transactions
         self._partition_id = partition_id
+        self._prometheus_tracker = (
+            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
+        )
+        self._collection_ensured = False
         self._rollouts = (
             rollouts
             if rollouts is not None
@@ -859,6 +1182,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id"],
                 Rollout,
                 [["status"]],
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
         self._attempts = (
@@ -872,6 +1196,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id", "attempt_id"],
                 Attempt,
                 [["status"], ["sequence_id"]],
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
         self._spans = (
@@ -885,6 +1210,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id", "attempt_id", "span_id"],
                 Span,
                 [["sequence_id"]],
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
         self._resources = (
@@ -898,30 +1224,51 @@ class MongoLightningCollections(LightningCollections):
                 ["resources_id"],
                 ResourcesUpdate,
                 ["update_time"],
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
         self._workers = (
             workers
             if workers is not None
             else MongoBasedCollection(
-                self._client_pool, self._database_name, "workers", self._partition_id, ["worker_id"], Worker, ["status"]
+                self._client_pool,
+                self._database_name,
+                "workers",
+                self._partition_id,
+                ["worker_id"],
+                Worker,
+                ["status"],
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
         self._rollout_queue = (
             rollout_queue
             if rollout_queue is not None
-            else MongoBasedQueue(self._client_pool, self._database_name, "rollout_queue", self._partition_id, str)
+            else MongoBasedQueue(
+                self._client_pool,
+                self._database_name,
+                "rollout_queue",
+                self._partition_id,
+                str,
+                prometheus_tracker=self._prometheus_tracker,
+            )
         )
         self._span_sequence_ids = (
             span_sequence_ids
             if span_sequence_ids is not None
             else MongoBasedKeyValue(
-                self._client_pool, self._database_name, "span_sequence_ids", self._partition_id, str, int
+                self._client_pool,
+                self._database_name,
+                "span_sequence_ids",
+                self._partition_id,
+                str,
+                int,
+                prometheus_tracker=self._prometheus_tracker,
             )
         )
 
     def with_session(self, session: AsyncClientSession) -> Self:
-        return self.__class__(
+        instance = self.__class__(
             client_pool=self._client_pool,
             database_name=self._database_name,
             partition_id=self._partition_id,
@@ -932,7 +1279,10 @@ class MongoLightningCollections(LightningCollections):
             workers=self._workers.with_session(session),
             rollout_queue=self._rollout_queue.with_session(session),
             span_sequence_ids=self._span_sequence_ids.with_session(session),
+            prometheus_tracker=self._prometheus_tracker,
         )
+        instance._collection_ensured = self._collection_ensured
+        return instance
 
     @property
     def rollouts(self) -> MongoBasedCollection[Rollout]:
@@ -962,8 +1312,11 @@ class MongoLightningCollections(LightningCollections):
     def span_sequence_ids(self) -> MongoBasedKeyValue[str, int]:
         return self._span_sequence_ids
 
+    @_mongo_operation("ensure_collections")
     async def _ensure_collections(self) -> None:
         """Ensure all collections exist."""
+        if self._collection_ensured:
+            return
         await self._rollouts.ensure_collection()
         await self._attempts.ensure_collection()
         await self._spans.ensure_collection()
@@ -971,56 +1324,144 @@ class MongoLightningCollections(LightningCollections):
         await self._workers.ensure_collection()
         await self._rollout_queue.ensure_collection()
         await self._span_sequence_ids.ensure_collection()
+        self._collection_ensured = True
 
     @asynccontextmanager
     async def atomic(self, *args: Any, **kwargs: Any):
         """Perform a atomic operation on the collections."""
-        # First step: ensure all collections exist before going into the atomic block
-        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
+        with self._prometheus_tracker.track("atomic", self._database_name, self._collection_name):
+            # First step: ensure all collections exist before going into the atomic block
+            if not self._collection_ensured:
+                await self._ensure_collections()
+            # One session for one transaction
+            client = await self._client_pool.get_client()
+            async with client.start_session() as session:
+                collection_with_session = self.with_session(session)
+                try:
+                    # Start the transaction now
+                    await session.start_transaction(
+                        write_concern=WriteConcern("majority"),
+                        read_concern=ReadConcern("local"),
+                        read_preference=ReadPreference.PRIMARY,
+                    )
+                    yield collection_with_session
+                    # Commit the transaction
+                    await session.commit_transaction()
+                # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
+                except BaseException as exc:
+                    if session.in_transaction:
+                        await session.abort_transaction()
+                    if isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError"):
+                        # NOTE: Retry is via execute
+                        raise RuntimeError("Transaction failed with transient error") from exc
+                    raise
 
+    @_mongo_operation("execute")
     async def execute(self, callback: Callable[[Self], Awaitable[T_generic]]) -> T_generic:
         """Execute the given callback within an atomic operation, and with retries on transient errors."""
-        await self._ensure_collections()
+        if not self._collection_ensured:
+            await self._ensure_collections()
         client = await self._client_pool.get_client()
 
         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)
+                except (ConnectionFailure, OperationFailure) as exc:
+                    # Un-retryable errors.
+                    tracker.report_error(exc)
+                    raise RuntimeError("Transaction failed with connection or operation error") from exc
 
-            async def _txn_callback(_session: AsyncClientSession) -> T_generic:
-                # The _session is always the same within one transaction,
-                # so we can use the same collections object.
-                return await callback(collections)
+    async def _with_transaction(
+        self,
+        session: AsyncClientSession,
+        collections: Self,
+        callback: Callable[[Self], Awaitable[T_generic]],
+        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()
 
-            try:
-                # This will start a transaction, run transaction callback, and commit.
-                # It will also transparently retry on some transient errors.
-                return await session.with_transaction(
-                    _txn_callback,
-                    read_concern=ReadConcern("local"),
-                    write_concern=WriteConcern("majority"),
-                    read_preference=ReadPreference.PRIMARY,
-                )
-            except (ConnectionFailure, OperationFailure) as exc:
-                raise RuntimeError("Transaction failed with connection or operation error") from exc
+        def _within_time_limit() -> bool:
+            return time.monotonic() - start_time < transaction_retry_time_limit
+
+        async def _jitter_before_retry() -> None:
+            with self._prometheus_tracker.track("execute__jitter", self._database_name, self._collection_name):
+                await asyncio.sleep(random.uniform(0, 0.05))
+
+        while True:
+            await session.start_transaction(read_concern, write_concern, read_preference)
+
+            with self._prometheus_tracker.track(
+                "execute__callback", self._database_name, self._collection_name
+            ) as callback_tracker:
+                try:
+                    num_attempts += 1
+                    transaction_tracker.report_num_attempts(num_attempts)
+                    # The _session is always the same within one transaction,
+                    # so we can use the same collections object.
+                    ret = await callback(collections)
+                # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
+                except BaseException as exc:
+                    callback_tracker.report_error(exc)
+                    if session.in_transaction:
+                        await session.abort_transaction()
+                    if (
+                        isinstance(exc, PyMongoError)
+                        and exc.has_error_label("TransientTransactionError")
+                        and _within_time_limit()
+                    ):
+                        # Retry the entire transaction.
+                        await _jitter_before_retry()
+                        continue
+                    raise
+
+            if not session.in_transaction:
+                # Assume callback intentionally ended the transaction.
+                return ret
+
+            commit_num_attempts = 0
+
+            # Tracks the commit operation.
+            with self._prometheus_tracker.track(
+                "execute__commit", self._database_name, self._collection_name
+            ) as commit_tracker:
+                # Loop until the commit succeeds or we hit the time limit.
+                while True:
+                    # Tracks the commit attempt.
+                    with self._prometheus_tracker.track(
+                        "execute__commit_attempt", self._database_name, self._collection_name
+                    ) as commit_attempt_tracker:
+                        try:
+                            commit_num_attempts += 1
+                            commit_tracker.report_num_attempts(commit_num_attempts)
+                            await session.commit_transaction()
+                        except PyMongoError as exc:
+                            commit_attempt_tracker.report_error(exc)
+                            if (
+                                exc.has_error_label("UnknownTransactionCommitResult")
+                                and _within_time_limit()
+                                and not (isinstance(exc, OperationFailure) and exc.code == 50)  # max_time_expired_error
+                            ):
+                                # Retry the commit.
+                                await _jitter_before_retry()
+                                continue
+
+                            if exc.has_error_label("TransientTransactionError") and _within_time_limit():
+                                # Retry the entire transaction.
+                                await _jitter_before_retry()
+                                break
+                            raise
+
+                        # Commit succeeded.
+                        return ret

agentlightning/store/collection_based.py

@@ -9,6 +9,7 @@ import logging
 import time
 import uuid
 import warnings
+from collections import defaultdict
 from types import CoroutineType
 from typing import (
     Any,
@@ -20,6 +21,7 @@ from typing import (
     Optional,
     ParamSpec,
     Sequence,
+    Tuple,
     TypeVar,
     Union,
     cast,
@@ -47,9 +49,17 @@ from agentlightning.types import (
     WorkerStatus,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset, is_finished, is_queuing
+from .base import (
+    UNSET,
+    LightningStore,
+    LightningStoreCapabilities,
+    LightningStoreStatistics,
+    Unset,
+    is_finished,
+    is_queuing,
+)
 from .collection import FilterOptions, LightningCollections
-from .utils import healthcheck, propagate_status
+from .utils import LATENCY_BUCKETS, healthcheck, propagate_status
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 T_model = TypeVar("T_model", bound=BaseModel)
@@ -83,7 +93,44 @@ def _with_collections_execute(
     return wrapper
 
 
-def _healthcheck_wrapper(func: T_callable) -> T_callable:
+def tracked(name: str):
+    """Decorator to track the execution of the decorated method with Prometheus."""
+
+    _public_methods = frozenset([name for name in LightningStore.__dict__ if not name.startswith("_")])
+
+    def decorator(func: T_callable) -> T_callable:
+
+        @functools.wraps(func)
+        async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
+            # For backtracking in Mongo-collection methods.
+            # Only track the public methods (+healthcheck)
+            if name in _public_methods or name == "_healthcheck":
+                method_name = name  # pyright: ignore[reportUnusedVariable]
+            else:
+                method_name = None  # pyright: ignore[reportUnusedVariable]
+
+            if not self._prometheus:  # pyright: ignore[reportPrivateUsage]
+                # Skip the tracking because tracking is not configured
+                return await func(self, *args, **kwargs)
+
+            start_time = time.perf_counter()
+            try:
+                ret = await func(self, *args, **kwargs)
+                self._total_metric.labels(name, "OK").inc()  # pyright: ignore[reportPrivateUsage]
+                return ret
+            except Exception as exc:
+                self._total_metric.labels(name, exc.__class__.__name__).inc()  # pyright: ignore[reportPrivateUsage]
+                raise
+            finally:
+                elapsed = time.perf_counter() - start_time
+                self._latency_metric.labels(name).observe(elapsed)  # pyright: ignore[reportPrivateUsage]
+
+        return cast(T_callable, wrapper)
+
+    return decorator
+
+
+def healthcheck_before(func: T_callable) -> T_callable:
     """
     Decorator to run the watchdog healthcheck **before** executing the decorated method.
     Only runs if the store has a watchdog configured.
@@ -144,10 +191,52 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         collections: The collections to use for storage.
     """
 
-    def __init__(self, collections: T_collections):
+    def __init__(self, collections: T_collections, prometheus: bool = False):
         # rollouts and spans' storage
         self.collections = collections
+        self._prometheus = prometheus
+        self._launch_time = time.time()
 
+        if prometheus:
+            from prometheus_client import Counter, Histogram
+
+            self._latency_metric = Histogram(
+                "collection_store_latency_seconds",
+                "Latency of CollectionBasedLightningStore methods",
+                ["method"],
+                buckets=LATENCY_BUCKETS,
+            )
+            self._total_metric = Counter(
+                "collection_store_total",
+                "Total MongoDB operations",
+                ["method", "error_type"],
+            )
+            self._rollout_counter = Counter(
+                "collection_store_rollout_total",
+                "Total rollouts",
+                ["status", "mode"],
+            )
+            self._rollout_duration_metric = Histogram(
+                "collection_store_rollout_duration_seconds",
+                "Duration of rollouts",
+                ["status", "mode"],
+                buckets=LATENCY_BUCKETS,
+            )
+
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        current_time = time.time()
+        return {
+            "name": self.__class__.__name__,
+            "total_rollouts": await self.collections.rollouts.size(),
+            "total_attempts": await self.collections.attempts.size(),
+            "total_spans": await self.collections.spans.size(),
+            "total_resources": await self.collections.resources.size(),
+            "total_workers": await self.collections.workers.size(),
+            "uptime": current_time - self._launch_time,
+        }
+
+    @tracked("_get_latest_resources_id")
     async def _get_latest_resources_id(self, collections: T_collections) -> Optional[str]:
         """Get the latest resources ID from the collections. Returns `None` if no resources are found."""
         latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
@@ -155,6 +244,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             return latest_resources.resources_id
         return None
 
+    @tracked("_get_or_create_worker")
     async def _get_or_create_worker(self, collections: T_collections, worker_id: str) -> Worker:
         """Create a worker if it doesn't exist.
 
@@ -166,6 +256,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             await collections.workers.insert([worker])
         return worker
 
+    @tracked("_sync_worker_with_attempt")
     async def _sync_worker_with_attempt(self, collections: T_collections, attempt: Attempt) -> None:
         worker_id = attempt.worker_id
         if not worker_id:
@@ -206,7 +297,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         """
         return LightningStoreCapabilities()
 
-    @_healthcheck_wrapper
+    @tracked("start_rollout")
+    @healthcheck_before
     @_with_collections_execute
     async def start_rollout(
         self,
@@ -259,7 +351,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # Return a rollout with attempt attached.
         return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
 
-    @_healthcheck_wrapper
+    @tracked("enqueue_rollout")
+    @healthcheck_before
     @_with_collections_execute
     async def enqueue_rollout(
         self,
@@ -302,7 +395,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # Return the rollout with no attempt attached.
         return rollout
 
-    @_healthcheck_wrapper
+    @tracked("dequeue_rollout")
+    @healthcheck_before
     @_with_collections_execute
     async def dequeue_rollout(
         self, collections: T_collections, worker_id: Optional[str] = None
@@ -368,7 +462,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # No valid rollouts found
         return None
 
-    @_healthcheck_wrapper
+    @tracked("start_attempt")
+    @healthcheck_before
     @_with_collections_execute
     async def start_attempt(self, collections: T_collections, rollout_id: str) -> AttemptedRollout:
         """Creates a new attempt for a given rollout ID and return the attempt details.
@@ -408,7 +503,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # Return the rollout with the new attempt attached.
         return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
 
-    @_healthcheck_wrapper
+    @tracked("query_rollouts")
+    @healthcheck_before
     @_with_collections_execute
     async def query_rollouts(
         self,
@@ -465,10 +561,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         )
 
         # Attach the latest attempt to the rollout objects
-        # TODO: Maybe we can use asyncio.gather here to speed up the process?
-        attempted_rollouts = [
-            await self._rollout_to_attempted_rollout_unlocked(collections, rollout) for rollout in rollouts.items
-        ]
+        attempted_rollouts = await self._many_rollouts_to_attempted_rollouts_unlocked(collections, rollouts.items)
 
         return PaginatedResult(
             items=attempted_rollouts, limit=rollouts.limit, offset=rollouts.offset, total=rollouts.total
@@ -482,7 +575,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         )
         return list(result.items)
 
-    @_healthcheck_wrapper
+    @tracked("get_rollout_by_id")
+    @healthcheck_before
     @_with_collections_execute
     async def get_rollout_by_id(
         self, collections: T_collections, rollout_id: str
@@ -498,6 +592,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             return None
         return await self._rollout_to_attempted_rollout_unlocked(collections, rollout)
 
+    @tracked("_rollout_to_attempted_rollout")
     async def _rollout_to_attempted_rollout_unlocked(
         self, collections: T_collections, rollout: Rollout
     ) -> Union[Rollout, AttemptedRollout]:
@@ -511,6 +606,15 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         else:
             return AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt)
 
+    @tracked("_many_rollouts_to_attempted_rollouts_unlocked")
+    async def _many_rollouts_to_attempted_rollouts_unlocked(
+        self, collections: T_collections, rollouts: Sequence[Rollout]
+    ) -> List[Union[Rollout, AttemptedRollout]]:
+        """Query the latest attempts for the rollouts, and attach them to the rollout objects."""
+        # TODO: Maybe we can use asyncio.gather here to speed up the process?
+        return [await self._rollout_to_attempted_rollout_unlocked(collections, rollout) for rollout in rollouts]
+
+    @tracked("_get_latest_attempt")
     async def _get_latest_attempt_unlocked(self, collections: T_collections, rollout_id: str) -> Optional[Attempt]:
         """The unlocked version of `get_latest_attempt`."""
         return await collections.attempts.get(
@@ -518,7 +622,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             sort={"name": "sequence_id", "order": "desc"},
         )
 
-    @_healthcheck_wrapper
+    @tracked("query_attempts")
+    @healthcheck_before
     @_with_collections_execute
     async def query_attempts(
         self,
@@ -538,7 +643,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             offset=offset,
         )
 
-    @_healthcheck_wrapper
+    @tracked("get_latest_attempt")
+    @healthcheck_before
     @_with_collections_execute
     async def get_latest_attempt(self, collections: T_collections, rollout_id: str) -> Optional[Attempt]:
         """Retrieves the latest attempt for a given rollout ID.
@@ -547,7 +653,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         """
         return await self._get_latest_attempt_unlocked(collections, rollout_id)
 
-    @_healthcheck_wrapper
+    @tracked("query_resources")
+    @healthcheck_before
     @_with_collections_execute
     async def query_resources(
         self,
@@ -576,7 +683,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             offset=offset,
         )
 
-    @_healthcheck_wrapper
+    @tracked("add_resources")
+    @healthcheck_before
     @_with_collections_execute
     async def add_resources(self, collections: T_collections, resources: NamedResources) -> ResourcesUpdate:
         """Stores a new version of named resources and sets it as the latest.
@@ -596,7 +704,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         self._latest_resources_id = resources_id
         return update
 
-    @_healthcheck_wrapper
+    @tracked("update_resources")
+    @healthcheck_before
     @_with_collections_execute
     async def update_resources(
         self, collections: T_collections, resources_id: str, resources: NamedResources
@@ -629,7 +738,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         self._latest_resources_id = resources_id
         return update
 
-    @_healthcheck_wrapper
+    @tracked("get_resources_by_id")
+    @healthcheck_before
     @_with_collections_execute
     async def get_resources_by_id(self, collections: T_collections, resources_id: str) -> Optional[ResourcesUpdate]:
         """Retrieves a specific version of named resources by its ID.
@@ -638,7 +748,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         """
         return await collections.resources.get({"resources_id": {"exact": resources_id}})
 
-    @_healthcheck_wrapper
+    @tracked("get_latest_resources")
+    @healthcheck_before
     @_with_collections_execute
     async def get_latest_resources(self, collections: T_collections) -> Optional[ResourcesUpdate]:
         """Retrieves the latest version of named resources.
@@ -650,16 +761,32 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             return None
         return await collections.resources.get({"resources_id": {"exact": latest_id}})
 
-    async def _issue_span_sequence_id_unlocked(self, collections: T_collections, rollout_id: str) -> int:
+    @tracked("_issue_many_span_sequence_ids")
+    async def _issue_many_span_sequence_ids_unlocked(
+        self, collections: T_collections, rollout_ids: List[str]
+    ) -> List[int]:
         """Issue a new span sequence ID for a given rollout."""
-        sequence_id = await collections.span_sequence_ids.get(rollout_id)
-        if sequence_id is None:
-            sequence_id = 1
-        else:
-            sequence_id += 1
-        await collections.span_sequence_ids.set(rollout_id, sequence_id)
-        return sequence_id
+        # Cache the next sequence IDs for the rollouts (for both RW)
+        next_sequence_ids_cache: Dict[str, int] = {}
+        result: List[int] = []
+        for rollout_id in rollout_ids:
+            if rollout_id not in next_sequence_ids_cache:
+                retrieved_id = await collections.span_sequence_ids.get(rollout_id)
+                if retrieved_id is None:
+                    retrieved_id = 0
+                next_sequence_ids_cache[rollout_id] = retrieved_id
 
+            # Increment the sequence ID for the rollout
+            next_sequence_ids_cache[rollout_id] += 1
+            result.append(next_sequence_ids_cache[rollout_id])
+
+        # Propagate the cache to storage
+        for rollout_id, sequence_id in next_sequence_ids_cache.items():
+            await collections.span_sequence_ids.set(rollout_id, sequence_id)
+
+        return result
+
+    @tracked("_sync_span_sequence_id")
     async def _sync_span_sequence_id_unlocked(
         self, collections: T_collections, rollout_id: str, sequence_id: int
     ) -> None:
@@ -669,6 +796,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             existing_sequence_id = 0
         await collections.span_sequence_ids.set(rollout_id, max(existing_sequence_id, sequence_id))
 
+    @tracked("get_next_span_sequence_id")
     @_with_collections_execute
     async def get_next_span_sequence_id(self, collections: T_collections, rollout_id: str, attempt_id: str) -> int:
         """Get the next span sequence ID for a given rollout and attempt.
@@ -677,18 +805,52 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
 
         See [`LightningStore.get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id] for semantics.
         """
-        return await self._issue_span_sequence_id_unlocked(collections, rollout_id)
+        ret = await self._issue_many_span_sequence_ids_unlocked(collections, [rollout_id])
+        return ret[0]
 
+    @tracked("get_many_span_sequence_ids")
     @_with_collections_execute
-    async def add_span(self, collections: T_collections, span: Span) -> Span:
+    async def get_many_span_sequence_ids(
+        self, collections: T_collections, rollout_attempt_ids: Sequence[Tuple[str, str]]
+    ) -> Sequence[int]:
+        """Get the next span sequence IDs for a given list of rollout and attempt identifiers."""
+        return await self._issue_many_span_sequence_ids_unlocked(
+            collections, [rollout_id for rollout_id, _ in rollout_attempt_ids]
+        )
+
+    @tracked("add_span")
+    @_with_collections_execute
+    async def add_span(self, collections: T_collections, span: Span) -> Optional[Span]:
         """Persist a pre-converted span.
 
         See [`LightningStore.add_span()`][agentlightning.LightningStore.add_span] for semantics.
         """
         # Update the sequence ID to be synced with latest input span
         await self._sync_span_sequence_id_unlocked(collections, span.rollout_id, span.sequence_id)
-        return await self._add_span_unlocked(collections, span)
+        ret = await self._add_many_spans_unlocked(collections, span.rollout_id, span.attempt_id, [span])
+        return ret[0] if len(ret) > 0 else None
 
+    @tracked("add_many_spans")
+    @_with_collections_execute
+    async def add_many_spans(self, collections: T_collections, spans: Sequence[Span]) -> Sequence[Span]:
+        """Persist a sequence of pre-converted spans.
+
+        See [`LightningStore.add_many_spans()`][agentlightning.LightningStore.add_many_spans] for semantics.
+        """
+        # Group spans by rollout and attempt
+        spans_by_rollout_attempt: Dict[Tuple[str, str], List[Span]] = defaultdict(list)
+        for span in spans:
+            spans_by_rollout_attempt[(span.rollout_id, span.attempt_id)].append(span)
+
+        # Bulk add spans for each rollout and attempt
+        successful_spans: List[Span] = []
+        for (rollout_id, attempt_id), spans in spans_by_rollout_attempt.items():
+            await self._sync_span_sequence_id_unlocked(collections, rollout_id, max(span.sequence_id for span in spans))
+            ret = await self._add_many_spans_unlocked(collections, rollout_id, attempt_id, spans)
+            successful_spans.extend(ret)
+        return successful_spans
+
+    @tracked("add_otel_span")
     @_with_collections_execute
     async def add_otel_span(
         self,
@@ -697,14 +859,14 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         """Add an opentelemetry span to the store.
 
         See [`LightningStore.add_otel_span()`][agentlightning.LightningStore.add_otel_span] for semantics.
         """
         if sequence_id is None:
             # Issue a new sequence ID for the rollout
-            sequence_id = await self._issue_span_sequence_id_unlocked(collections, rollout_id)
+            sequence_id = (await self._issue_many_span_sequence_ids_unlocked(collections, [rollout_id]))[0]
         else:
             # Comes from a provided sequence ID
             # Make sure our counter is strictly increasing
@@ -713,35 +875,55 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         span = Span.from_opentelemetry(
             readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id
         )
-        await self._add_span_unlocked(collections, span)
-        return span
+        ret = await self._add_many_spans_unlocked(collections, rollout_id, attempt_id, [span])
+        return ret[0] if len(ret) > 0 else None
 
-    async def _add_span_unlocked(self, collections: T_collections, span: Span) -> Span:
-        rollout = await collections.rollouts.get({"rollout_id": {"exact": span.rollout_id}})
+    @tracked("_add_many_spans_unlocked")
+    async def _add_many_spans_unlocked(
+        self, collections: T_collections, rollout_id: str, attempt_id: str, spans: Sequence[Span]
+    ) -> Sequence[Span]:
+        """All spans must be for the same rollout and attempt."""
+        rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
         if not rollout:
-            raise ValueError(f"Rollout {span.rollout_id} not found")
+            raise ValueError(f"Rollout {rollout_id} not found")
         current_attempt = await collections.attempts.get(
-            filter={"rollout_id": {"exact": span.rollout_id}, "attempt_id": {"exact": span.attempt_id}},
-        )
-        latest_attempt = await collections.attempts.get(
-            filter={"rollout_id": {"exact": span.rollout_id}},
-            sort={"name": "sequence_id", "order": "desc"},
+            filter={"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}},
         )
+        latest_attempt = await self._get_latest_attempt_unlocked(collections, rollout_id)
         if not current_attempt:
-            raise ValueError(f"Attempt {span.attempt_id} not found for rollout {span.rollout_id}")
+            raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
         if not latest_attempt:
-            raise ValueError(f"No attempts found for rollout {span.rollout_id}")
+            raise ValueError(f"No attempts found for rollout {rollout_id}")
 
+        async def _add_span_fallback(span: Span) -> bool:
+            try:
+                await collections.spans.insert([span])
+                return True
+            except ValueError as e:
+                if "already exists" in str(e) or "contains duplicate" in str(e):
+                    logger.error(
+                        f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
+                    )
+                    return False
+                raise
+
+        successful_spans: List[Span] = []
         try:
-            await collections.spans.insert([span])
+            await collections.spans.insert(spans)
+            successful_spans.extend(spans)
         except ValueError as e:
-            if "already exists" in str(e):
-                # This is a duplicate span, we warn it
-                logger.error(
-                    f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
-                )
-                return span
-            raise
+            if "already exists" in str(e) or "contains duplicate" in str(e):
+                # There is a duplicate span, we warn it
+                # We fallback to adding the spans one by one
+                for span in spans:
+                    if await _add_span_fallback(span):
+                        successful_spans.append(span)
+            else:
+                raise
+
+        if not successful_spans:
+            # No spans were added, skip the rest.
+            return []
 
         # Update attempt heartbeat and ensure persistence
         current_attempt.last_heartbeat_time = time.time()
@@ -762,9 +944,10 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                 await collections.rollouts.update([rollout])
                 await self.on_rollout_update(rollout)
 
-        return span
+        return successful_spans
 
-    @_healthcheck_wrapper
+    @tracked("wait_for_rollouts")
+    @healthcheck_before
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """Wait for specified rollouts to complete with a timeout.
         Returns the completed rollouts, potentially incomplete if timeout is reached.
@@ -785,6 +968,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # Filter out the exceptions
         return [rollout for rollout in rollouts if isinstance(rollout, Rollout)]
 
+    @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.
 
@@ -823,7 +1007,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
 
         return None
 
-    @_healthcheck_wrapper
+    @tracked("query_spans")
+    @healthcheck_before
     @_with_collections_execute
     async def query_spans(
         self,
@@ -856,10 +1041,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         if attempt_id is None:
             resolved_attempt_id = None
         elif attempt_id == "latest":
-            latest_attempt = await collections.attempts.get(
-                filter={"rollout_id": {"exact": rollout_id}},
-                sort={"name": "sequence_id", "order": "desc"},
-            )
+            latest_attempt = await self._get_latest_attempt_unlocked(collections, rollout_id)
             if not latest_attempt:
                 logger.debug(f"No attempts found for rollout {rollout_id} when querying latest spans")
                 return PaginatedResult(items=[], limit=limit, offset=offset, total=0)
@@ -896,7 +1078,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             offset=offset,
         )
 
-    @_healthcheck_wrapper
+    @tracked("update_rollout")
+    @healthcheck_before
     @_with_collections_execute
     async def update_rollout(
         self,
@@ -924,7 +1107,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             metadata=metadata,
         )
 
-    @_healthcheck_wrapper
+    @tracked("update_attempt")
+    @healthcheck_before
     @_with_collections_execute
     async def update_attempt(
         self,
@@ -950,6 +1134,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             metadata=metadata,
         )
 
+    @tracked("_update_rollout_unlocked")
     async def _update_rollout_unlocked(
         self,
         collections: T_collections,
@@ -984,6 +1169,11 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         # Rollout is only finished when it succeeded or fail with no more retries.
         if not isinstance(status, Unset) and is_finished(rollout):
             rollout.end_time = time.time()
+            if self._prometheus:
+                self._rollout_counter.labels(rollout.status, rollout.mode).inc()
+                self._rollout_duration_metric.labels(rollout.status, rollout.mode).observe(
+                    rollout.end_time - rollout.start_time
+                )
 
         # If requeuing, add back to queue.
         # Check whether the rollout is already in queue.
@@ -1000,6 +1190,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
 
         return rollout
 
+    @tracked("_update_attempt_unlocked")
     async def _update_attempt_unlocked(
         self,
         collections: T_collections,
@@ -1015,10 +1206,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         if not rollout:
             raise ValueError(f"Rollout {rollout_id} not found")
 
-        latest_attempt = await collections.attempts.get(
-            {"rollout_id": {"exact": rollout_id}},
-            sort={"name": "sequence_id", "order": "desc"},
-        )
+        latest_attempt = await self._get_latest_attempt_unlocked(collections, rollout_id)
         if not latest_attempt:
             raise ValueError(f"No attempts found for rollout {rollout_id}")
 
@@ -1071,7 +1259,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
 
         return attempt
 
-    @_healthcheck_wrapper
+    @tracked("query_workers")
+    @healthcheck_before
     @_with_collections_execute
     async def query_workers(
         self,
@@ -1100,12 +1289,14 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             offset=offset,
         )
 
-    @_healthcheck_wrapper
+    @tracked("get_worker_by_id")
+    @healthcheck_before
     @_with_collections_execute
     async def get_worker_by_id(self, collections: T_collections, worker_id: str) -> Optional[Worker]:
         return await collections.workers.get({"worker_id": {"exact": worker_id}})
 
-    @_healthcheck_wrapper
+    @tracked("update_worker")
+    @healthcheck_before
     @_with_collections_execute
     async def update_worker(
         self,
@@ -1123,6 +1314,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         await collections.workers.update([worker])
         return worker
 
+    @tracked("on_rollout_update")
     async def on_rollout_update(self, rollout: Rollout) -> None:
         """Callback for subclasses to implement specific logic when a rollout changes.
 
@@ -1130,6 +1322,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         """
         pass
 
+    @tracked("get_running_rollouts")
     async def get_running_rollouts(self, collections: T_collections) -> List[AttemptedRollout]:
         """Get all running rollouts.
 
@@ -1137,21 +1330,19 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         subclass can implement hacks to make it more efficient.
         It should also be unlocked and let the caller hold the lock.
         """
-        running_rollouts: List[AttemptedRollout] = []
-        rollouts = await collections.rollouts.query(filter={"status": {"within": ["preparing", "running"]}})
+        filtered_rollouts = await collections.rollouts.query(filter={"status": {"within": ["preparing", "running"]}})
+        running_rollouts = await self._many_rollouts_to_attempted_rollouts_unlocked(collections, filtered_rollouts)
 
-        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
+        running_attempted_rollouts: List[AttemptedRollout] = []
+        for rollout in running_rollouts:
+            if not isinstance(rollout, AttemptedRollout):
                 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
+            running_attempted_rollouts.append(rollout)
 
+        return running_attempted_rollouts
+
+    @tracked("_healthcheck")
     @_with_collections_execute
     async def _healthcheck(self, collections: T_collections) -> None:
         """Perform healthcheck against all running rollouts in the store."""

agentlightning/store/memory.py

@@ -17,6 +17,7 @@ from typing import (
     Literal,
     Mapping,
     Optional,
+    Sequence,
     Set,
     TypeVar,
     Union,
@@ -27,9 +28,9 @@ from pydantic import BaseModel
 
 from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span
 
-from .base import UNSET, LightningStoreCapabilities, Unset, is_finished, is_running
+from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
-from .collection_based import CollectionBasedLightningStore
+from .collection_based import CollectionBasedLightningStore, tracked
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 
@@ -84,8 +85,9 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         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())
+        super().__init__(collections=InMemoryLightningCollections(), prometheus=prometheus)
 
         self._start_time_by_rollout: Dict[str, float] = {}
         self._span_bytes_by_rollout: Dict[str, int] = Counter()
@@ -138,6 +140,17 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             otlp_traces=False,
         )
 
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        return {
+            **(await super().statistics()),
+            "total_span_bytes": self._total_span_bytes,
+            "eviction_threshold_bytes": self._eviction_threshold_bytes,
+            "safe_threshold_bytes": self._safe_threshold_bytes,
+            "memory_capacity_bytes": self._memory_capacity_bytes,
+        }
+
+    @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:
@@ -175,6 +188,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
         return None
 
+    @tracked("on_rollout_update")
     async def on_rollout_update(self, rollout: Rollout) -> None:
         """Update the running rollout ids set when the rollout updates."""
         if is_running(rollout):
@@ -193,6 +207,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         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)}})
@@ -209,6 +224,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             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
     async def query_spans(
         self,
         rollout_id: str,
@@ -219,15 +235,20 @@ 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)
 
-    async def _add_span_unlocked(self, collections: InMemoryLightningCollections, span: Span) -> Span:
+    @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]:
         """In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""
 
-        await super()._add_span_unlocked(collections, span)
-        self._account_span_size(span)
+        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)
 
-        return span
+        return inserted
 
+    @tracked("_get_latest_resources_id")
     async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
         if isinstance(self._latest_resources_id, Unset):
             latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
@@ -267,7 +288,8 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
         return resolved
 
-    def _account_span_size(self, span: Span) -> int:
+    @tracked("_account_span_size")
+    async def _account_span_size(self, span: Span) -> int:
         if self._custom_span_size_estimator is not None:
             size = max(int(self._custom_span_size_estimator(span)), 0)
         else:
@@ -277,6 +299,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         self._total_span_bytes += size
         return size
 
+    @tracked("_maybe_evict_spans")
     async def _maybe_evict_spans(self, collections: InMemoryLightningCollections) -> None:
         if self._total_span_bytes <= self._eviction_threshold_bytes:
             return
@@ -299,6 +322,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             await self._evict_spans_for_rollout(collections, rollout_id)
         logger.info(f"Freed up {memory_consumed_before - self._total_span_bytes} bytes of memory")
 
+    @tracked("_evict_spans_for_rollout")
     async def _evict_spans_for_rollout(self, collections: InMemoryLightningCollections, rollout_id: str) -> None:
         await collections.evict_spans_for_rollout(rollout_id)
         removed_bytes = self._span_bytes_by_rollout.pop(rollout_id, 0)

agentlightning/store/mongo.py

@@ -2,21 +2,30 @@
 
 from __future__ import annotations
 
+import asyncio
 import hashlib
 import logging
+import time
 import uuid
 from typing import (
     Any,
     Callable,
+    Dict,
+    List,
     Mapping,
+    Optional,
+    Sequence,
     TypeVar,
+    Union,
 )
 
 from pymongo import AsyncMongoClient
 
-from .base import LightningStoreCapabilities
-from .collection.mongo import MongoClientPool, MongoLightningCollections
-from .collection_based import CollectionBasedLightningStore
+from agentlightning.types import Attempt, AttemptedRollout, Rollout
+
+from .base import LightningStoreCapabilities, is_finished
+from .collection.mongo import MongoClientPool, MongoLightningCollections, MongoOperationPrometheusTracker
+from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 
@@ -45,7 +54,9 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         client: AsyncMongoClient[Mapping[str, Any]] | str,
         database_name: str | None = None,
         partition_id: str | None = None,
+        prometheus: bool = False,
     ) -> None:
+        self._enable_prometheus = prometheus
         self._auto_created_client = False
         if isinstance(client, str):
             self._client = AsyncMongoClient[Mapping[str, Any]](client)
@@ -62,7 +73,15 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
 
         self._client_pool = MongoClientPool(self._client)
 
-        super().__init__(collections=MongoLightningCollections(self._client_pool, database_name, partition_id))
+        super().__init__(
+            collections=MongoLightningCollections(
+                self._client_pool,
+                database_name,
+                partition_id,
+                prometheus_tracker=MongoOperationPrometheusTracker(enabled=self._enable_prometheus),
+            ),
+            prometheus=self._enable_prometheus,
+        )
 
     @property
     def capabilities(self) -> LightningStoreCapabilities:
@@ -80,3 +99,63 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         # If I created the client, I should close it too.
         if self._auto_created_client:
             await self._client.close()
+
+    @tracked("wait_for_rollouts")
+    @healthcheck_before
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """Wait for specified rollouts to complete with a timeout.
+
+        Concurrently wait for all rollouts to complete with a timeout.
+        """
+        start_time = time.time()
+        current_time = start_time
+        deadline = start_time + timeout if timeout is not None else None
+
+        finished_rollouts: Dict[str, Rollout] = {}
+        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)}}
+            )
+            for rollout in rollouts.items:
+                if is_finished(rollout):
+                    finished_rollouts[rollout.rollout_id] = rollout
+                    unfinished_rollout_ids.remove(rollout.rollout_id)
+
+            if not unfinished_rollout_ids:
+                break
+
+            # Poll every 10 seconds by default
+            # Minus 0.1 to make sure the time is still sufficient for another call
+            rest_time = max(0.01, min(deadline - time.time() - 0.1, 10.0)) if deadline is not None else 10.0
+            await asyncio.sleep(rest_time)
+            current_time = time.time()
+
+        # 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(
+        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"},
+        )
+        latest_attempts: Dict[str, Attempt] = {}
+        for attempt in attempts:
+            if attempt.rollout_id not in latest_attempts:
+                latest_attempts[attempt.rollout_id] = attempt
+            # Otherwise we ignore the attempt because there's already a newer attempt
+
+        return [
+            (
+                AttemptedRollout(**rollout.model_dump(), attempt=latest_attempts[rollout.rollout_id])
+                if rollout.rollout_id in latest_attempts
+                else rollout
+            )
+            for rollout in rollouts
+        ]

agentlightning/store/threading.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import threading
-from typing import Any, Dict, List, Literal, Optional, Sequence
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -22,7 +22,7 @@ from agentlightning.types import (
     WorkerStatus,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
+from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
 
 
 class LightningStoreThreaded(LightningStore):
@@ -47,6 +47,11 @@ class LightningStoreThreaded(LightningStore):
             "thread_safe": True,
         }
 
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        with self._lock:
+            return await self.store.statistics()
+
     async def start_rollout(
         self,
         input: TaskInput,
@@ -167,7 +172,11 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.get_latest_resources()
 
-    async def add_span(self, span: Span) -> Span:
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        with self._lock:
+            return await self.store.add_many_spans(spans)
+
+    async def add_span(self, span: Span) -> Optional[Span]:
         with self._lock:
             return await self.store.add_span(span)
 
@@ -177,7 +186,7 @@ class LightningStoreThreaded(LightningStore):
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         with self._lock:
             return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
 
@@ -189,6 +198,10 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
 
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        with self._lock:
+            return await self.store.get_many_span_sequence_ids(rollout_attempt_ids)
+
     async def query_spans(
         self,
         rollout_id: str,

agentlightning/store/utils.py

@@ -9,6 +9,54 @@ UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
 UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
 
 
+LATENCY_BUCKETS = [
+    0.000001,
+    0.000002,
+    0.000005,
+    0.00001,
+    0.00002,
+    0.00005,
+    0.0001,
+    0.0002,
+    0.0005,
+    0.001,
+    0.002,
+    0.003,
+    0.005,
+    0.007,
+    0.01,
+    0.015,
+    0.02,
+    0.03,
+    0.05,
+    0.07,
+    0.1,
+    0.2,
+    0.3,
+    0.5,
+    0.7,
+    1.0,
+    2.0,
+    3.0,
+    5.0,
+    7.0,
+    10.0,
+    12.0,
+    15.0,
+    20.0,
+    25.0,
+    30.0,
+    40.0,
+    50.0,
+    60.0,
+    90.0,
+    120.0,
+    180.0,
+    240.0,
+    300.0,
+]
+
+
 async def propagate_status(
     update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
     attempt: Attempt,

agentlightning/tracer/otel.py

@@ -18,8 +18,9 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 
+from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.store.base import LightningStore
-from agentlightning.types.tracer import SpanNames
+from agentlightning.utils.otel import get_tracer_provider
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
 from .base import Tracer
@@ -51,7 +52,16 @@ class OtelTracer(Tracer):
         logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")
 
         if self._initialized:
-            logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
+            logger.info(f"[Worker {worker_id}] Tracer provider is already initialized. Skipping initialization.")
+            return
+
+        try:
+            get_tracer_provider()
+            logger.error(
+                f"[Worker {worker_id}] Tracer provider is already initialized but not by OtelTracer. OpenTelemetry may not work as expected."
+            )
+        except RuntimeError:
+            logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
 
         self._tracer_provider = TracerProvider()
         trace_api.set_tracer_provider(self._tracer_provider)
@@ -66,8 +76,7 @@ class OtelTracer(Tracer):
 
     def teardown_worker(self, worker_id: int):
         super().teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
-        self._tracer_provider = None
+        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")
 
     @asynccontextmanager
     async def trace_context(
@@ -144,8 +153,8 @@ class OtelTracer(Tracer):
         tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
             Resource.create(
                 {
-                    SpanNames.ROLLOUT_ID: rollout_id,
-                    SpanNames.ATTEMPT_ID: attempt_id,
+                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
                 }
             )
         )
@@ -182,8 +191,8 @@ class OtelTracer(Tracer):
         tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
             Resource.create(
                 {
-                    SpanNames.ROLLOUT_ID: "",
-                    SpanNames.ATTEMPT_ID: "",
+                    LightningResourceAttributes.ROLLOUT_ID.value: "",
+                    LightningResourceAttributes.ATTEMPT_ID.value: "",
                 }
             )
         )  # reset resource
@@ -219,6 +228,30 @@ class LightningSpanProcessor(SpanProcessor):
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._loop_thread: Optional[threading.Thread] = None
 
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            + f"disable_store_submission={self.disable_store_submission}, "
+            + f"store={self.store!r}, "
+            + f"rollout_id={self.rollout_id!r}, "
+            + f"attempt_id={self.attempt_id!r})"
+        )
+
+    @property
+    def store(self) -> Optional[LightningStore]:
+        """The store to submit the spans to."""
+        return self._store
+
+    @property
+    def rollout_id(self) -> Optional[str]:
+        """The rollout ID to submit the spans to."""
+        return self._rollout_id
+
+    @property
+    def attempt_id(self) -> Optional[str]:
+        """The attempt ID to submit the spans to."""
+        return self._attempt_id
+
     @property
     def disable_store_submission(self) -> bool:
         """Whether to disable submitting spans to the store."""

agentlightning/types/tracer.py

@@ -16,6 +16,8 @@ from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace.status import Status as OtelStatus
 from pydantic import BaseModel, ConfigDict
 
+from agentlightning.semconv import AGL_VIRTUAL
+
 __all__ = [
     "AttributeValue",
     "Attributes",
@@ -379,7 +381,7 @@ class Span(BaseModel):
                 is_remote=False,
                 trace_state={},
             ),
-            name=name or SpanNames.VIRTUAL.value,
+            name=name or AGL_VIRTUAL,
             resource=resource or OtelResource(attributes={}, schema_url=""),
             attributes=attributes,
             status=TraceStatus(status_code="OK"),
@@ -399,7 +401,7 @@ class Span(BaseModel):
 
 
 class SpanNames(str, Enum):
-    """Enumerated span names recognised by Agent-lightning."""
+    """Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""
 
     REWARD = "agentlightning.reward"
     """The name of the reward span."""
@@ -420,7 +422,7 @@ class SpanNames(str, Enum):
 
 
 class SpanAttributeNames(str, Enum):
-    """Canonical attribute names written by Agent Lightning emitters."""
+    """Canonical attribute names written by Agent Lightning emitters. Deprecated in favor of [semconv][agentlightning.semconv]."""
 
     MESSAGE = "message"
     """The name of the message attribute."""

agentlightning/utils/otel.py

@@ -0,0 +1,401 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities shared for OpenTelemetry span (attributes) support."""
+
+import logging
+from typing import Any, Dict, List, Sequence, Union, cast
+from warnings import filterwarnings
+
+import opentelemetry.trace as trace_api
+from agentops.sdk.exporters import OTLPSpanExporter
+from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
+from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
+from opentelemetry.trace import get_tracer_provider as otel_get_tracer_provider
+from pydantic import TypeAdapter
+
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
+from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
+from agentlightning.types import SpanLike
+from agentlightning.utils.otlp import LightningStoreOTLPExporter
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "full_qualified_name",
+    "get_tracer_provider",
+    "get_tracer",
+    "make_tag_attributes",
+    "extract_tags_from_attributes",
+    "make_link_attributes",
+    "query_linked_spans",
+    "extract_links_from_attributes",
+    "filter_attributes",
+    "filter_and_unflatten_attributes",
+    "flatten_attributes",
+    "unflatten_attributes",
+]
+
+
+def full_qualified_name(obj: type) -> str:
+    if str(obj.__module__) == "builtins":
+        return obj.__qualname__
+    return f"{obj.__module__}.{obj.__qualname__}"
+
+
+def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
+    """Get the OpenTelemetry tracer provider configured for Agent Lightning.
+
+    Args:
+        inspect: Whether to inspect the tracer provider and log its configuration.
+            When it's on, make sure you also set the logger level to DEBUG to see the logs.
+    """
+    from agentlightning.tracer.otel import LightningSpanProcessor
+
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+    tracer_provider = otel_get_tracer_provider()
+    if not isinstance(tracer_provider, TracerProviderImpl):
+        logger.error(
+            "Tracer provider is expected to be an instance of opentelemetry.sdk.trace.TracerProvider, found: %s",
+            full_qualified_name(type(tracer_provider)),
+        )
+        return cast(TracerProviderImpl, tracer_provider)
+
+    if not inspect:
+        return tracer_provider
+
+    emitter_debug = resolve_bool_env_var(LightningEnvVar.AGL_EMITTER_DEBUG, fallback=None)
+    logger_effective_level = logger.getEffectiveLevel()
+    if emitter_debug is True and logger_effective_level > logging.DEBUG:
+        logger.warning(
+            "Emitter debug logging is enabled but logging level is not set to DEBUG. Nothing will be logged."
+        )
+
+    if emitter_debug is None:
+        # Set to true by default if the logging level is lower than DEBUG
+        emitter_debug = logging.DEBUG >= logger_effective_level
+
+    if emitter_debug:
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+        processors: List[str] = []
+        active_span_processor_cls = active_span_processor.__class__.__name__
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # The legacy case for tracers without OTLP support.
+                processors.append(f"{active_span_processor_cls} - {processor!r}")
+            elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+                processor_cls = processor.__class__.__name__
+                if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
+                    # This should be the main path now.
+                    processors.append(f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter!r}")
+                elif isinstance(processor.span_exporter, OTLPSpanExporter):
+                    # You need to be careful if the code goes into this path.
+                    endpoint = processor.span_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
+                    processors.append(
+                        f"{active_span_processor_cls} - {processor_cls} - "
+                        f"{processor.span_exporter.__class__.__name__}(endpoint={endpoint!r})"
+                    )
+                else:
+                    # Other cases like Console Span Exporter.
+                    processors.append(
+                        f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter.__class__.__name__}"
+                    )
+            else:
+                processors.append(f"{active_span_processor_cls} - {processor.__class__.__name__}")
+
+        logger.debug(f"Tracer provider: {tracer_provider!r}. Active span processors:")
+        for processor in processors:
+            logger.debug("  * " + processor)
+
+    return tracer_provider
+
+
+def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
+    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
+
+    Args:
+        use_active_span_processor: Whether to use the active span processor.
+
+    Returns:
+        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
+
+    Raises:
+        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
+    """
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = get_tracer_provider(inspect=True)  # inspection is on by default
+
+    if use_active_span_processor:
+        return tracer_provider.get_tracer("agentlightning")
+
+    else:
+        filterwarnings(
+            "ignore",
+            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
+            category=DeprecationWarning,
+            module="opentelemetry.sdk.trace",
+        )
+
+        return Tracer(
+            tracer_provider.sampler,
+            tracer_provider.resource,
+            # We use an empty span processor to avoid emitting spans to the tracer
+            SynchronousMultiSpanProcessor(),
+            tracer_provider.id_generator,
+            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
+            SpanLimits(),
+            InstrumentationScope(
+                "agentlightning",
+                "",
+                "",
+                {},
+            ),
+        )
+
+
+def make_tag_attributes(tags: List[str]) -> Dict[str, Any]:
+    """Convert a list of tags into flattened attributes for span tagging.
+
+    There is no syntax enforced for tags, they are just strings. For example:
+
+    ```python
+    ["gen_ai.model:gpt-4", "reward.extrinsic"]
+    ```
+    """
+    return flatten_attributes({LightningSpanAttributes.TAG.value: tags})
+
+
+def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
+    """Extract tag attributes from flattened span attributes.
+
+    Args:
+        attributes: A dictionary of flattened span attributes.
+    """
+    maybe_tag_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.TAG.value)
+    return TypeAdapter(List[str]).validate_python(maybe_tag_list)
+
+
+def make_link_attributes(links: Dict[str, str]) -> Dict[str, Any]:
+    """Convert a dictionary of links into flattened attributes for span linking.
+
+    Links example:
+
+    ```python
+    {
+        "gen_ai.response.id": "response-123",
+        "span_id": "abcd-efgh-ijkl",
+    }
+    ```
+    """
+    link_list: List[Dict[str, str]] = []
+    for key, value in links.items():
+        if not isinstance(value, str):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise ValueError(f"Link value must be a string, got {type(value)} for key '{key}'")
+        link_list.append({LinkAttributes.KEY_MATCH.value: key, LinkAttributes.VALUE_MATCH.value: value})
+    return flatten_attributes({LightningSpanAttributes.LINK.value: link_list})
+
+
+def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]) -> List[SpanLike]:
+    """Query spans that are linked by the given link attributes.
+
+    Args:
+        spans: A sequence of spans to search.
+        links: A list of link attributes to match.
+
+    Returns:
+        A list of spans that match the given link attributes.
+    """
+    matched_spans: List[SpanLike] = []
+
+    for span in spans:
+        span_attributes = span.attributes or {}
+        is_match = True
+        for link in links:
+            # trace_id and span_id must be full match.
+            if link.key_match == "trace_id":
+                if isinstance(span, ReadableSpan):
+                    trace_id = trace_api.format_trace_id(span.context.trace_id) if span.context else None
+                else:
+                    trace_id = span.trace_id
+                if trace_id != link.value_match:
+                    is_match = False
+                    break
+
+            elif link.key_match == "span_id":
+                if isinstance(span, ReadableSpan):
+                    span_id = trace_api.format_span_id(span.context.span_id) if span.context else None
+                else:
+                    span_id = span.span_id
+                if span_id != link.value_match:
+                    is_match = False
+                    break
+
+            else:
+                attribute = span_attributes.get(link.key_match)
+                # attributes must also be a full match currently.
+                if attribute != link.value_match:
+                    is_match = False
+                    break
+
+        if is_match:
+            matched_spans.append(span)
+
+    return matched_spans
+
+
+def extract_links_from_attributes(attributes: Dict[str, Any]) -> List[LinkPydanticModel]:
+    """Extract link attributes from flattened span attributes.
+
+    Args:
+        attributes: A dictionary of flattened span attributes.
+    """
+    maybe_link_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+    return TypeAdapter(List[LinkPydanticModel]).validate_python(maybe_link_list)
+
+
+def filter_attributes(attributes: Dict[str, Any], prefix: str) -> Dict[str, Any]:
+    """Filter attributes that start with the given prefix.
+
+    The attribute must start with `prefix.` or be exactly `prefix` to be included.
+
+    Args:
+        attributes: A dictionary of span attributes.
+        prefix: The prefix to filter by.
+
+    Returns:
+        A dictionary of attributes that start with the given prefix.
+    """
+    return {k: v for k, v in attributes.items() if k.startswith(prefix + ".") or k == prefix}
+
+
+def filter_and_unflatten_attributes(attributes: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
+    """Filter attributes that start with the given prefix and unflatten them.
+    The prefix will be removed during unflattening.
+
+    Args:
+        attributes: A dictionary of span attributes.
+        prefix: The prefix to filter by.
+
+    Returns:
+        A nested dictionary or list of attributes that start with the given prefix.
+    """
+    filtered_attributes = filter_attributes(attributes, prefix)
+    stripped_attributes: Dict[str, Any] = {}
+    for k, v in filtered_attributes.items():
+        if k == prefix:
+            raise ValueError(f"Cannot unflatten attribute with key exactly equal to prefix: {prefix}")
+        else:
+            stripped_key = k[len(prefix) + 1 :]  # +1 to remove the dot
+        stripped_attributes[stripped_key] = v
+    return unflatten_attributes(stripped_attributes)
+
+
+def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[str, Any]:
+    """Flatten a nested dictionary or list into a flat dictionary with dotted keys.
+
+    This function recursively traverses dictionaries and lists, producing a flat
+    key-value mapping where nested paths are represented via dot-separated keys.
+    Lists are indexed numerically.
+
+    Example:
+
+        >>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}})
+        {"a.b": 1, "a.c.0": 2, "a.c.1": 3}
+
+    Args:
+        nested_data: A nested structure composed of dictionaries, lists, or
+            primitive values.
+
+    Returns:
+        A flat dictionary mapping dotted-string paths to primitive values.
+    """
+
+    flat: Dict[str, Any] = {}
+
+    def _walk(value: Any, prefix: str = "") -> None:
+        if isinstance(value, dict):
+            for k, v in cast(Dict[Any, Any], value).items():
+                if not isinstance(k, str):
+                    raise ValueError(
+                        f"Only string keys are supported in dictionaries, got '{k}' of type {type(k)} in {prefix}"
+                    )
+                new_prefix = f"{prefix}.{k}" if prefix else k
+                _walk(v, new_prefix)
+        elif isinstance(value, list):
+            for idx, item in enumerate(cast(List[Any], value)):
+                new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
+                _walk(item, new_prefix)
+        else:
+            flat[prefix] = value
+
+    _walk(nested_data)
+    return flat
+
+
+def unflatten_attributes(flat_data: Dict[str, Any]) -> Union[Dict[str, Any], List[Any]]:
+    """Reconstruct a nested dictionary/list structure from a flat dictionary.
+
+    Keys are dot-separated paths. Segments that are digit strings will only
+    become list indices if *all* keys in that dict form a consecutive
+    0..n-1 range. Otherwise they remain dict keys.
+
+    Example:
+
+        >>> unflatten_attributes({"a.b": 1, "a.c.0": 2, "a.c.1": 3})
+        {"a": {"b": 1, "c": [2, 3]}}
+
+    Args:
+        flat_data: A dictionary whose keys are dot-separated paths and whose
+            values are primitive data elements.
+
+    Returns:
+        A nested dictionary (and lists where appropriate) corresponding to
+        the flattened structure.
+    """
+    # 1) Build a pure dict tree first (no lists yet)
+    root: Dict[str, Any] = {}
+
+    for flat_key, value in flat_data.items():
+        parts = flat_key.split(".")
+        curr: Dict[str, Any] = root
+
+        for part in parts[:-1]:
+            # Ensure intermediate node is a dict
+            if part not in curr or not isinstance(curr[part], dict):
+                curr[part] = {}
+            curr = curr[part]  # type: ignore[assignment]
+
+        curr[parts[-1]] = value
+
+    # 2) Recursively convert dicts-with-consecutive-numeric-keys into lists
+    def convert(node: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
+        if isinstance(node, dict):
+            # First convert children
+            for k, v in list(node.items()):
+                node[k] = convert(v)
+
+            if not node:
+                # empty dict stays dict
+                return node
+
+            # Check if keys are all numeric strings
+            keys = list(node.keys())
+            if all(isinstance(k, str) and k.isdigit() for k in keys):  # pyright: ignore[reportUnnecessaryIsInstance]
+                indices = sorted(int(k) for k in keys)
+                # Must be exactly 0..n-1
+                if indices == list(range(len(indices))):
+                    return [node[str(i)] for i in range(len(indices))]
+
+            return node
+
+        if isinstance(node, list):  # pyright: ignore[reportUnnecessaryIsInstance]
+            return [convert(v) for v in node]
+
+        # Keep as is
+        return node
+
+    return convert(root)

agentlightning/utils/otlp.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import gzip
 import logging
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
@@ -29,7 +31,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.util.types import AttributeValue
 
-from agentlightning.store.base import LightningStore
+from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.types.tracer import (
     Attributes,
     Event,
@@ -37,7 +39,6 @@ from agentlightning.types.tracer import (
     OtelResource,
     Span,
     SpanContext,
-    SpanNames,
     TraceStatus,
     convert_timestamp,
 )
@@ -108,7 +109,10 @@ async def handle_otlp_export(
     )
 
 
-async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningStore) -> List[Span]:
+async def spans_from_proto(
+    request: ExportTraceServiceRequest,
+    sequence_id_bulk_issuer: Callable[[Sequence[Tuple[str, str]]], Awaitable[Sequence[int]]],
+) -> List[Span]:
     """Parse an OTLP proto payload into List[Span].
 
     A store is needed here for generating a sequence ID for each span.
@@ -119,11 +123,11 @@ async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningS
         # Resource-level attributes & IDs
         resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
         # rollout_id, attempt_id from resource attributes when present.
-        rollout_id_resource = resource_attrs.get(SpanNames.ROLLOUT_ID)
-        attempt_id_resource = resource_attrs.get(SpanNames.ATTEMPT_ID)
+        rollout_id_resource = resource_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
+        attempt_id_resource = resource_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
         # If sequence id is provided, all the spans will share the same sequence ID.
         # unless otherwise overridden by span-level attributes.
-        sequence_id_resource = resource_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
+        sequence_id_resource = resource_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
 
         otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
 
@@ -154,9 +158,9 @@ async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningS
 
                 # Try to get if span attributes contain something like rollout_id or attempt_id
                 # Override the resource-level attributes with the span-level attributes if present.
-                rollout_id_span = span_attrs.get(SpanNames.ROLLOUT_ID)
-                attempt_id_span = span_attrs.get(SpanNames.ATTEMPT_ID)
-                sequence_id_span = span_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
+                rollout_id_span = span_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
+                attempt_id_span = span_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
+                sequence_id_span = span_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
 
                 # Normalize to regular strings and ints
                 rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
@@ -178,9 +182,13 @@ async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningS
 
                 # Generate a new sequence ID if not provided
                 if sequence_id is None:
-                    current_sequence_id = await store.get_next_span_sequence_id(
-                        rollout_id=rollout_id, attempt_id=attempt_id
+                    current_sequence_id = -1
+                elif sequence_id < 0:
+                    logger.error(
+                        "Invalid sequence_id value in resource attributes: %r. Must be a positive integer. Regenerating one.",
+                        sequence_id,
                     )
+                    current_sequence_id = -1
                 else:
                     current_sequence_id = sequence_id
 
@@ -206,6 +214,14 @@ async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningS
 
                 output_spans.append(span)
 
+    # Finalize the sequence IDs
+    bulk_issue_requests = [(span.rollout_id, span.attempt_id) for span in output_spans if span.sequence_id < 0]
+    bulk_sequence_ids = await sequence_id_bulk_issuer(bulk_issue_requests)
+    for span, sequence_id in zip(
+        [span for span in output_spans if span.sequence_id < 0], bulk_sequence_ids, strict=True
+    ):
+        span.sequence_id = sequence_id
+
     return output_spans
 
 
@@ -226,6 +242,36 @@ class LightningStoreOTLPExporter(OTLPSpanExporter):
     _rollout_id: Optional[str] = None
     _attempt_id: Optional[str] = None
 
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            + f"endpoint={self.endpoint!r}, "
+            + f"rollout_id={self.rollout_id!r}, "
+            + f"attempt_id={self.attempt_id!r}, "
+            + f"should_bypass={self.should_bypass()!r})"
+        )
+
+    @property
+    def endpoint(self) -> Optional[str]:
+        """The endpoint to submit the spans to."""
+        if hasattr(self, "_endpoint"):
+            return self._endpoint
+        return None
+
+    @property
+    def rollout_id(self) -> Optional[str]:
+        """The rollout ID to submit the spans to."""
+        if hasattr(self, "_rollout_id"):
+            return self._rollout_id
+        return None
+
+    @property
+    def attempt_id(self) -> Optional[str]:
+        """The attempt ID to submit the spans to."""
+        if hasattr(self, "_attempt_id"):
+            return self._attempt_id
+        return None
+
     def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
         """Enable storing OTLP data to a specific LightningStore rollout/attempt."""
         self._rollout_id = rollout_id
@@ -254,8 +300,8 @@ class LightningStoreOTLPExporter(OTLPSpanExporter):
                 span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
                     Resource.create(
                         {
-                            SpanNames.ROLLOUT_ID: self._rollout_id,
-                            SpanNames.ATTEMPT_ID: self._attempt_id,
+                            LightningResourceAttributes.ROLLOUT_ID.value: self._rollout_id,
+                            LightningResourceAttributes.ATTEMPT_ID.value: self._attempt_id,
                         }
                     )
                 )

agentlightning/utils/server_launcher.py

@@ -6,6 +6,7 @@ import asyncio
 import inspect
 import logging
 import multiprocessing
+import os
 import queue
 import signal
 import socket
@@ -938,6 +939,11 @@ class PythonServerLauncher:
                     self.args.process_join_timeout / 2
                 ),  # Allow half the timeout for graceful shutdown
             }
+            if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+                from prometheus_client import multiprocess
+
+                options["child_exit"] = lambda server, worker: multiprocess.mark_process_dead(worker.pid)  # type: ignore
+
             self._gunicorn_app = GunicornApp(self.app, options)
 
             self._proc = ctx.Process(

agentlightning/verl/daemon.py

@@ -561,7 +561,7 @@ class AgentModeDaemon:
             final_reward = self._fillna_reward(rollout)
             if not rollout.triplets:
                 print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
-                sample_stat_list.append({"reward": final_reward})
+                sample_stat_list.append({"reward": final_reward, "has_reward": final_reward_raw is not None})
                 continue
             response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
 

docker/compose.mongo.yml

@@ -2,9 +2,8 @@
 # It's used to test the MongoDB store implementation.
 
 services:
-
   mongo:
-    image: mongo:latest
+    image: mongo:8.2
     ulimits:
       nofile:
         soft: 65535

docker/compose.prometheus-memory-store.yml

@@ -1,5 +1,4 @@
 services:
-
   app:
     extends:
       file: compose.store.yml
@@ -11,18 +10,17 @@ services:
     image: prom/node-exporter:latest
     # In CI you might not have full /proc, but this is OK for container-level stats
     pid: "host"
-    network_mode: "service:app" # share network with app for simplicity
     command:
-      - '--path.rootfs=/host'
+      - "--path.rootfs=/host"
     volumes:
-      - '/:/host:ro,rslave'
+      - "/:/host:ro,rslave"
 
   prometheus:
     image: prom/prometheus:latest
     command:
-      - '--config.file=/etc/prometheus/prometheus.yml'
-      - '--storage.tsdb.path=/prometheus'
-      - '--storage.tsdb.retention.time=1h'
+      - "--config.file=/etc/prometheus/prometheus.yml"
+      - "--storage.tsdb.path=/prometheus"
+      - "--storage.tsdb.retention.time=1h"
     volumes:
       - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
@@ -31,3 +29,25 @@ services:
       - node-exporter
     ports:
       - "9090:9090"
+
+  grafana:
+    image: grafana/grafana:latest
+    ports:
+      - "9091:3000"
+    depends_on:
+      - prometheus
+    volumes:
+      - ./data/grafana:/var/lib/grafana
+      # 1. Mount the Datasource Config
+      - ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
+      # 2. Mount the Dashboard Provider Config
+      - ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
+      # 3. Mount the folder containing the actual JSON files
+      - ./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/compose.prometheus-mongo-store.yml

@@ -1,5 +1,4 @@
 services:
-
   mongo:
     extends:
       file: compose.mongo.yml
@@ -23,14 +22,24 @@ services:
     depends_on:
       - mongo
 
-    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend mongo --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 --n-workers 4
+    command:
+      - /bin/bash
+      - -c
+      - |
+          mkdir -p /tmp/prometheus &&
+          agl store --host 0.0.0.0 --port 4747 \
+            --prometheus --backend mongo \
+            --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
+            --n-workers ${AGL_STORE_N_WORKERS:-32}
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus
 
   mongodb-exporter:
     image: percona/mongodb_exporter:0.47.1
     command:
-      - '--mongodb.uri=mongodb://mongo:27017/'
-      - '--collect-all'
-      - '--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids'
+      - "--mongodb.uri=mongodb://mongo:27017/"
+      - "--collect-all"
+      - "--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids"
     depends_on:
       - mongo
     ports:
@@ -41,16 +50,16 @@ services:
     # In CI you might not have full /proc, but this is OK for container-level stats
     pid: "host"
     command:
-      - '--path.rootfs=/host'
+      - "--path.rootfs=/host"
     volumes:
-      - '/:/host:ro,rslave'
+      - "/:/host:ro,rslave"
 
   prometheus:
     image: prom/prometheus:latest
     command:
-      - '--config.file=/etc/prometheus/prometheus.yml'
-      - '--storage.tsdb.path=/prometheus'
-      - '--storage.tsdb.retention.time=1h'
+      - "--config.file=/etc/prometheus/prometheus.yml"
+      - "--storage.tsdb.path=/prometheus"
+      - "--storage.tsdb.retention.time=1h"
     volumes:
       - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
@@ -60,3 +69,25 @@ services:
       - node-exporter
     ports:
       - "9090:9090"
+
+  grafana:
+    image: grafana/grafana:latest
+    ports:
+      - "9091:3000"
+    depends_on:
+      - prometheus
+    volumes:
+      - ./data/grafana:/var/lib/grafana
+      # 1. Mount the Datasource Config
+      - ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
+      # 2. Mount the Dashboard Provider Config
+      - ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
+      # 3. Mount the folder containing the actual JSON files
+      - ./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/compose.store.yml

@@ -1,5 +1,4 @@
 services:
-
   app:
     build:
       context: ../

docker/grafana/dashboard-provider.yml

@@ -0,0 +1,12 @@
+apiVersion: 1
+
+providers:
+  - name: "default"
+    orgId: 1
+    folder: ""
+    type: file
+    disableDeletion: false
+    updateIntervalSeconds: 10
+    options:
+      # This tells Grafana to look for JSON files in this directory inside the container
+      path: /var/lib/grafana/dashboards

docker/grafana/datasource.yml

@@ -0,0 +1,8 @@
+apiVersion: 1
+
+datasources:
+  - name: Prometheus
+    type: prometheus
+    access: proxy
+    url: http://prometheus:9090
+    isDefault: true

docker/prometheus.memory-store.yml

@@ -6,11 +6,7 @@ scrape_configs:
   - job_name: app
     static_configs:
       - targets: ["app:4747"]
-    metrics_path: /v1/prometheus
-
-  - job_name: mongodb
-    static_configs:
-      - targets: ["mongodb-exporter:9216"]
+    metrics_path: /v1/prometheus/
 
   - job_name: node
     static_configs:

docker/prometheus.mongo-store.yml

@@ -6,8 +6,12 @@ scrape_configs:
   - job_name: app
     static_configs:
       - targets: ["app:4747"]
-    metrics_path: /v1/prometheus
+    metrics_path: /v1/prometheus/
 
   - job_name: node
     static_configs:
       - targets: ["node-exporter:9100"]
+
+  - job_name: mongodb
+    static_configs:
+      - targets: ["mongodb-exporter:9216"]

docker/setup.sh

@@ -3,7 +3,7 @@
 set -euo pipefail
 
 # Create data directories
-mkdir -p data/prometheus data/mongo-container data/mongo-host
+mkdir -p data/prometheus data/mongo-container data/mongo-host data/grafana
 
 # Change permissions
-chmod 777 data/prometheus data/mongo-container data/mongo-host
+chmod 777 data/prometheus data/mongo-container data/mongo-host data/grafana

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.emit_annotation
+
 ::: agentlightning.emit_reward
 
 ::: agentlightning.emit_message
@@ -30,7 +32,11 @@
 
 ::: agentlightning.emit_exception
 
-## Reward Helpers
+## Emitter Helpers
+
+::: agentlightning.get_message_value
+
+::: agentlightning.get_object_value
 
 ::: agentlightning.find_final_reward
 
@@ -38,8 +44,6 @@
 
 ::: agentlightning.get_reward_value
 
+::: agentlightning.get_rewards_from_span
+
 ::: agentlightning.is_reward_span
-
-## Legacy Emitter Decorators
-
-::: agentlightning.reward.reward

docs/reference/internal.md

@@ -4,6 +4,8 @@
 
     The following APIs should be used with extra caution because they are very likely to change in the future.
 
+## Algorithms and Adapters
+
 ::: agentlightning.adapter.messages.OpenAIMessages
 
 ::: agentlightning.adapter.triplet.TraceTree
@@ -14,12 +16,16 @@
 
 ::: agentlightning.algorithm.decorator.FunctionalAlgorithm
 
+## LitAgent
+
 ::: agentlightning.litagent.decorator.FunctionalLitAgent
 
 ::: agentlightning.litagent.decorator.llm_rollout
 
 ::: agentlightning.litagent.decorator.prompt_rollout
 
+## LLM Proxy
+
 ::: agentlightning.llm_proxy.ModelConfig
 
 ::: agentlightning.llm_proxy.LightningSpanExporter
@@ -34,24 +40,56 @@
 
 ::: agentlightning.llm_proxy.RolloutAttemptMiddleware
 
+## Store
+
 ::: agentlightning.store.base.UNSET
 
 ::: agentlightning.store.utils.propagate_status
 
+## Tracing and OpenTelemetry
+
 ::: agentlightning.tracer.otel.LightningSpanProcessor
 
+## Utilities
+
 ::: agentlightning.utils.server_launcher.PythonServerLauncher
 
 ::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
 
 ::: agentlightning.utils.server_launcher.LaunchMode
 
+::: agentlightning.utils.otel.full_qualified_name
+
+::: agentlightning.utils.otel.get_tracer_provider
+
+::: agentlightning.utils.otel.get_tracer
+
+::: agentlightning.utils.otel.make_tag_attributes
+
+::: agentlightning.utils.otel.extract_tags_from_attributes
+
+::: agentlightning.utils.otel.make_link_attributes
+
+::: agentlightning.utils.otel.query_linked_spans
+
+::: agentlightning.utils.otel.extract_links_from_attributes
+
+::: agentlightning.utils.otel.filter_attributes
+
+::: agentlightning.utils.otel.filter_and_unflatten_attributes
+
+::: agentlightning.utils.otel.flatten_attributes
+
+::: agentlightning.utils.otel.unflatten_attributes
+
 ::: agentlightning.utils.otlp.handle_otlp_export
 
 ::: agentlightning.utils.otlp.spans_from_proto
 
 ## Deprecated APIs
 
+::: agentlightning.emitter.reward.reward
+
 ::: agentlightning.server.AgentLightningServer
 
 ::: agentlightning.server.ServerDataStore

docs/reference/types.md

@@ -76,8 +76,20 @@
 
 ::: agentlightning.Span
 
-::: agentlightning.SpanNames
-
 ::: agentlightning.SpanAttributeNames
 
 ::: agentlightning.SpanLike
+
+## Semantic Conventions
+
+::: agentlightning.semconv
+
+## Environment Variables
+
+::: agentlightning.LightningEnvVar
+
+::: agentlightning.resolve_bool_env_var
+
+::: agentlightning.resolve_int_env_var
+
+::: agentlightning.resolve_str_env_var

docs/stylesheets/extra.css

@@ -22,6 +22,10 @@
   font-size: 1.15em;
 }
 
+.md-typeset h5, .md-typeset h6 {
+  font-size: 1em;
+}
+
 /* Increase spacing between API references */
 .doc-class, .doc-function, .doc-attribute {
   padding-bottom: 2em;

docs/tutorials/debug.md

@@ -43,7 +43,7 @@ Example output (with a reward span captured):
 
 ```python
 [Rollout(rollout_id='ro-519769241af8', input='Explain why the sky appears blue using principles of light scattering in 100 words.', start_time=1760706315.6996238, ..., status='succeeded')]
-[Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='agentlightning.reward', attributes={'reward': 0.95}, ...)]
+[Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='agentlightning.annotation', attributes={'agentlightning.reward.0.value': 0.95}, ...)]
 ```
 
 Swap in an [`AgentOpsTracer`][agentlightning.AgentOpsTracer] instead of [`OtelTracer`][agentlightning.OtelTracer] to see the underlying LLM spans alongside reward information:
@@ -52,7 +52,7 @@ Swap in an [`AgentOpsTracer`][agentlightning.AgentOpsTracer] instead of [`OtelTr
 [
     Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='openai.chat.completion', attributes={..., 'gen_ai.prompt.0.role': 'user', 'gen_ai.prompt.0.content': 'You are a helpful assistant. Explain why the sky appears blue using principles of light scattering in 100 words.', ...}),
     Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=2, ..., name='openai.chat.completion', attributes={..., 'gen_ai.prompt.0.role': 'user', 'gen_ai.prompt.0.content': 'Evaluate how well the output fulfills the task...', ...}),
-    Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=3, ..., name='agentlightning.reward', attributes={'reward': 0.95}, ...)
+    Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=3, ..., name='agentlightning.annotation', attributes={'agentlightning.reward.0.value': 0.95}, ...)
 ]
 ```
 
@@ -220,7 +220,7 @@ Just like [`Runner.run_context`][agentlightning.Runner.run_context], [`Trainer.d
 21:20:35 [Rollout ro-302fb202bd85 | Attempt 1] ID: at-f84ad21c. Status: succeeded. Worker: Worker-0
 21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 3a286a856af6bea8] #1 (openai.chat.completion) ... 1.95 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
 21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span e2f44b775e058dd6] #2 (openai.chat.completion) ... 1.24 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
-21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 45ee3c94fa1070ec] #3 (agentlightning.reward) ... 0.00 seconds. Attribute keys: ['reward']
+21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 45ee3c94fa1070ec] #3 (agentlightning.annotation) ... 0.00 seconds. Attribute keys: ['agentlightning.reward.0.value']
 21:20:35 [Rollout ro-302fb202bd85] Adapted data: [Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=None, metadata={'response_id': '...', 'agent_name': ''}), Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=0.95, metadata={'response_id': '...', 'agent_name': ''})]
 21:20:35 Finished 1 rollouts.
 21:20:35 [Rollout ro-e65a3ffaa540] Status changed to preparing.
@@ -228,7 +228,7 @@ Just like [`Runner.run_context`][agentlightning.Runner.run_context], [`Trainer.d
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt 1] ID: at-eaefa5d4. Status: succeeded. Worker: Worker-0
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 901dd6acc0f50147] #1 (openai.chat.completion) ... 1.30 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 52e0aa63e02be611] #2 (openai.chat.completion) ... 1.26 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
-21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 6c452de193fbffd3] #3 (agentlightning.reward) ... 0.00 seconds. Attribute keys: ['reward']
+21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 6c452de193fbffd3] #3 (agentlightning.annotation) ... 0.00 seconds. Attribute keys: ['agentlightning.reward.0.value']
 21:20:40 [Rollout ro-e65a3ffaa540] Adapted data: [Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=None, metadata={'response_id': '...', 'agent_name': ''}), Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=1.0, metadata={'response_id': '...', 'agent_name': ''})]
 21:20:40 Finished 2 rollouts.
 ```

docs/tutorials/traces.md

@@ -110,7 +110,7 @@ You can also customize an [`Adapter`][agentlightning.Adapter] by extending the i
 
 ### Reading Rewards
 
-Rewards are recorded as dedicated spans named [`agentlightning.reward`][agentlightning.SpanNames.REWARD]. Emitting a reward through [`emit_reward`][agentlightning.emit_reward] or the [`@reward` decorator][agentlightning.reward.reward] ensures the value is stored in the span’s `attributes["reward"]`. To audit rewards, fetch spans from the store and use the helper utilities in [`agentlightning.emitter`](../reference/agent.md):
+Rewards are recorded as dedicated spans named [`agentlightning.annotation`][agentlightning.semconv.AGL_ANNOTATION]. Emitting a reward through [`emit_reward`][agentlightning.emit_reward] or [`emit_annotation`][agentlightning.emit_annotation] ensures the value is stored in the span’s `attributes`. To audit rewards, fetch spans from the store and use the helper utilities in [`agentlightning.emitter`](../reference/agent.md):
 
 ```python
 from agentlightning.emitter import find_final_reward

docs/tutorials/write-agents.md

@@ -115,7 +115,7 @@ The value your agent function returns (i.e., the return value of the function de
 
 !!! important "Emitting the Final Reward"
 
-    When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Emitter section][using-emitter] below) or by wrapping your reward calculation function with the [`@reward`][agentlightning.reward.reward] decorator.
+    When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Emitter section][using-emitter] below). Wrapping your reward calculation function with the `@reward` decorator is NOT the recommended approach any more.
 
 * **`list[ReadableSpan]`** or **`list[Span]`**: For advanced use cases, you can manually construct and return a complete list of all spans for the rollout. This gives you full control over the trace data. You can return either a list of OpenTelemetry `ReadableSpan` objects or Agent-lightning's native `Span` objects.
 

examples/.gitignore

@@ -4,6 +4,7 @@ outputs/
 checkpoints/
 calc-x-data.zip
 spider-data.zip
+claude_code/logs/
 agentops.log
 unsloth/models/
 unsloth/unsloth_compiled_cache/

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/calc_x/train_calc_agent.py

@@ -38,6 +38,7 @@ from calc_agent import MathProblem, calc_agent
 from datasets import Dataset as HuggingFaceDataset
 
 import agentlightning as agl
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_str_env_var
 
 
 def verl_default_config() -> Dict[str, Any]:
@@ -153,7 +154,7 @@ def train(
         PROJECT_NAME = "AgentLightningCI"
 
         # Skip this step if AGL_CURRENT_ROLE is runner
-        agl_current_role = os.getenv("AGL_CURRENT_ROLE")
+        agl_current_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE)
 
         if agl_current_role != "runner":
             # Simulate writing to $GITHUB_OUTPUT if it’s set
@@ -222,7 +223,7 @@ def main():
 
     if args.external_store_address:
         print(f"Connecting to external store at: {args.external_store_address}")
-        if not os.getenv("AGL_MANAGED_STORE"):
+        if resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=True):
             raise ValueError(
                 "When using an external store, please set the environment variable AGL_MANAGED_STORE=0. "
                 "Otherwise the trainer will still try to manage the store lifecycle for you!"

examples/claude_code/README.md

@@ -0,0 +1,112 @@
+# Training Claude Code with Agent-lightning
+
+[![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)
+
+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.
+
+**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
+
+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
+
+`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
+
+| 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 |
+
+## Running the Example
+
+All commands are issued from `examples/claude_code`. Inspect the module-level docstring in `claude_code_agent.py` for the full CLI reference.
+
+### Hosted vLLM (open-source models)
+
+First, launch your model behind an OpenAI-compatible endpoint, for example:
+
+```bash
+vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+    --max-model-len 131072 \
+    --enable-auto-tool-choice \
+    --tool-call-parser qwen3_coder
+```
+
+Run the Agent-lightning harness and point it at the server:
+
+```bash
+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 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.
+
+### Official Claude Code (Anthropic API)
+
+```bash
+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
+```
+
+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
+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
+```
+
+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

@@ -0,0 +1,540 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Instrumented driver for running Claude Code on SWE-bench with Agent-lightning.
+
+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:
+
+- `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`.
+
+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
+import json
+import logging
+import os
+import resource
+from argparse import ArgumentParser
+from typing import Any, Dict, List, Literal, Optional, Sequence, cast
+
+from claude_code_controller import ClaudeController
+from datasets import Dataset
+from extended_adapter import ExtendedLlmProxyTraceToTriplet
+from swebench.harness.constants import SWEbenchInstance
+from swebench.harness.utils import load_swebench_dataset  # pyright: ignore[reportUnknownVariableType]
+from swebench_utils.evaluation import evaluate
+from swebench_utils.logging import log_for_evaluation
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from agentlightning import (
+    InMemoryLightningStore,
+    LightningStoreServer,
+    LitAgentRunner,
+    OtelTracer,
+    setup_logging,
+    setup_module_logging,
+)
+from agentlightning.litagent import LitAgent
+from agentlightning.llm_proxy import LLMProxy, ModelConfig
+from agentlightning.store import LightningStore
+from agentlightning.types import AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span
+
+logger = logging.getLogger("claude_code_agent")
+
+
+def _load_dataset(path: str, epoch: int = 0, limit: Optional[int] = None) -> List[SWEbenchInstance]:
+    instances: List[SWEbenchInstance] = []
+    with open(path) as f:
+        for line in f:
+            instance = json.loads(line)
+            instance["epoch"] = epoch
+            instances.append(instance)
+
+    if limit is not None:
+        instances = instances[:limit]
+    return instances
+
+
+def _flatten_messages(messages: List[Any]) -> List[Dict[str, str]]:
+    flattened: List[Dict[str, str]] = []
+    for msg in messages:
+        if msg["role"] in ["system", "user"] and isinstance(msg["content"], list):
+            msg_content: List[str] = []
+            for content in msg["content"]:
+                msg_content.append(content["text"])
+
+            msg["content"] = "".join(msg_content)
+        elif msg["role"] == "assistant" and "tool_calls" in msg:
+            # NOTE:
+            # Tool calls are list of dict, though in most case only one tool call is made per call
+            # We serialize it as json string here to avoid nested structure
+            msg["tool_calls"] = json.dumps(msg["tool_calls"])
+
+        for k in msg:
+            assert isinstance(msg[k], str), f"\n>>> {msg}"
+        flattened.append(msg)
+    return flattened
+
+
+class ClaudeCodeAgent(LitAgent[SWEbenchInstance]):
+    """Claude Code Agent implementation.
+
+    This agent is a wrapper of the Claude Code controller,
+    and it should be used to run the Claude Code agent on SWE-bench datasets.
+    """
+
+    def __init__(
+        self,
+        namespace: Literal["swebench", "starryzhang"] = "swebench",
+        max_turns: int = 5,
+        run_method: Literal["python", "cli"] = "cli",
+        open_file_limit: int = 4096,
+        cache_level: str = "env",  # ["none", "base", "env", "instance"]
+        clean: bool = False,
+        force_rebuild: bool = False,
+        timeout: int = 1_800,  # in sec
+        instance_image_tag: str = "latest",
+        rewrite_reports: bool = False,
+        swebench_full_dataset: Optional[List[SWEbenchInstance]] = None,
+    ) -> None:
+        super().__init__()
+        self.namespace = namespace
+        self.max_turns = max_turns
+        self.run_method = run_method
+
+        self.cache_level = cache_level
+        self.clean = clean
+        self.force_rebuild = force_rebuild
+        self.timeout = timeout
+        self.instance_image_tag = instance_image_tag
+        self.rewrite_reports = rewrite_reports
+
+        self.swebench_full_dataset = (
+            {each["instance_id"]: each for each in swebench_full_dataset} if swebench_full_dataset is not None else {}
+        )
+
+        # Set the maximum number of open files to the specified limit.
+        resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
+
+    async def rollout_async(
+        self, task: SWEbenchInstance, resources: NamedResources, rollout: Rollout
+    ) -> RolloutRawResult:
+        if not isinstance(rollout, AttemptedRollout):
+            # Technically, rollout should be an AttemptedRollout here.
+            # but the API is not stabilized yet.
+            raise ValueError("Rollout is not an AttemptedRollout.")
+
+        run_id = f"epoch_{task.get('epoch', 0)}"
+        image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_")
+
+        llm = cast(ProxyLLM, resources["llm"])
+
+        try:
+            # 1. init container
+            controller = ClaudeController(
+                image,
+                task,
+                run_id,
+                llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+                llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy"),
+            )
+            # 2. execute task
+            prediction = controller.run_instance(
+                task, max_turns=self.max_turns, run_method=cast(Literal["python", "cli"], self.run_method)
+            )
+            del controller
+        except Exception as e:
+            log_for_evaluation(run_id, task["instance_id"], f"Exception during rollout: {e}")
+            return 0.0
+
+        # 3. obtain rewards (evaluation result)
+        reward = 0.0
+        # empty patch
+        if prediction["model_patch"] in ["", None]:
+            return reward
+
+        instance_id = prediction["instance_id"]
+
+        result = evaluate(
+            cast(Any, prediction),
+            self.swebench_full_dataset[instance_id],
+            self.cache_level,
+            self.clean,
+            self.force_rebuild,
+            run_id,
+            self.timeout,
+            namespace=self.namespace,
+            instance_image_tag=self.instance_image_tag,
+            rewrite_reports=self.rewrite_reports,
+        )
+
+        # error patch
+        if result is None:
+            return reward
+
+        report = result[1]
+        # resolved/unresolved patch
+        if report[instance_id]["resolved"]:
+            reward = 1.0
+        return reward
+
+
+def sanity_check_spans(spans: Sequence[Span]) -> None:
+    assert len(spans) > 1, f"At least two spans are expected for a valid rollout. Found {len(spans)} spans."
+    assert any(span.name == "raw_gen_ai_request" for span in spans), "raw_gen_ai_request span not found"
+    assert any(span.name == "agentlightning.annotation" for span in spans), "agentlightning.annotation span not found"
+
+
+async def run_instance_async(
+    instance: SWEbenchInstance,
+    agent: ClaudeCodeAgent,
+    runner: LitAgentRunner[SWEbenchInstance],
+    store: LightningStore,
+    output_dir: Optional[str],
+    adapter: Optional[ExtendedLlmProxyTraceToTriplet],
+    tokenizer: Optional[PreTrainedTokenizerBase],
+) -> None:
+    """Runs the agent on a specific SWE-bench instance.
+
+    Running on specific SWE-bench instance and queries the traced spans.
+    It then extracts the triplets and saves the dataset.
+    """
+
+    instance_id = instance["instance_id"]
+    logger.info(f"Starting to run instance: {instance_id}")
+
+    # Run the agent and query the traced spans.
+    with runner.run_context(agent=agent, store=store):
+        rollout = await runner.step(instance)
+
+    logger.info(f"Finished running instance: {instance_id}")
+
+    spans = await store.query_spans(rollout.rollout_id)
+
+    if output_dir is None:
+        logger.info(f"Generated {len(spans)} spans for {instance_id}")
+        return
+
+    # 1. Dump raw spans (Common for both types)
+    raw_path = os.path.join(output_dir, f"stream_{instance_id}.json")
+    with open(raw_path, "w") as f:
+        for span in spans:
+            f.write(json.dumps(span.model_dump()) + "\n")
+    logger.info(f"Dumped {len(spans)} spans to {raw_path}")
+
+    # 2. Extract Triplets and Save Dataset (vLLM specific)
+    if adapter is not None and tokenizer is not None:
+        try:
+            triplets = adapter.adapt(cast(List[Span], spans))
+            logger.info(f"Extracted {len(triplets)} triplets for {instance_id}")
+
+            all_triplets: List[Dict[str, Any]] = []
+            recent_reward: Optional[float] = None
+
+            # Process in reverse to propagate rewards if necessary/logic dictates
+            for triplet in reversed(triplets):
+                if triplet.reward is not None:
+                    recent_reward = triplet.reward
+
+                prompt_text = tokenizer.decode(triplet.prompt["token_ids"])  # type: ignore
+                all_triplets.append(
+                    {
+                        "repo": instance.get("repo", ""),
+                        "instance_id": instance_id,
+                        "turn": triplet.metadata["sequence_id"],
+                        "prompt_ids": triplet.prompt["token_ids"],
+                        "gold_completion_ids": triplet.response["token_ids"],
+                        "logprobs": triplet.response["logprobs"],
+                        "reward": recent_reward,
+                        "prompt": prompt_text,
+                        "messages": _flatten_messages(triplet.metadata["messages"]),
+                    }
+                )
+
+            if all_triplets:
+                ds = Dataset.from_list(all_triplets)  # type: ignore
+                save_path = os.path.join(output_dir, f"dataset-{instance_id}")
+                ds.save_to_disk(save_path)  # type: ignore
+                logger.info(f"Saved HuggingFace dataset to {save_path}")
+
+        except Exception as e:
+            logger.error(f"Failed to extract triplets for {instance_id}: {e}")
+
+    logger.info(f"Finished extracting spans and traces for instance: {instance_id}")
+
+    # Quickly sanity check the spans
+    sanity_check_spans(spans)
+    logger.info(f"Sanity check passed for instance: {instance_id}")
+
+
+async def dry_run_claude_code(
+    *,
+    dataset_path: str,
+    haiku_frontend_name: str,
+    haiku_backend_name: str,
+    sonnet_frontend_name: str,
+    sonnet_backend_name: str,
+    backend_type: Literal["vllm", "anthropic", "openai"],
+    api_base_url: Optional[str],
+    output_dir: Optional[str],
+    max_turns: int,
+    limit: Optional[int],
+    cooldown_seconds: float,
+) -> None:
+    """Executes a dry run of the Claude Code agent on a dataset.
+
+    This function handles both 'official' runs (interacting with Anthropic APIs)
+    and 'hosted' runs (interacting with vLLM or compatible servers). It manages
+    initialization of the Lightning Store, LLM Proxy, and the execution loop.
+
+    If running in 'vllm' mode, it will also attempt to extract triplets using
+    the provided backend name as the tokenizer path and save a HuggingFace Dataset.
+
+    Args:
+        dataset_path: Path to the JSONL dataset file.
+        haiku_frontend_name: The model name used in the code to request the 'fast' model.
+        haiku_backend_name: The actual model name/path on the backend.
+        sonnet_frontend_name: The model name used in the code to request the 'strong' model.
+        sonnet_backend_name: The actual model name/path on the backend.
+        backend_type: The type of backend to configure ("vllm", "anthropic" or "openai").
+        api_base_url: Base URL for the API. Required for "vllm" or "openai".
+        output_dir: Directory to save logs, spans, and datasets.
+        max_turns: Maximum number of steps the agent can take per instance.
+        limit: Optional limit on the number of instances to process.
+    """
+    dataset = _load_dataset(dataset_path, limit=limit)
+
+    # Initialize Infrastructure
+    tracer = OtelTracer()
+    runner = LitAgentRunner[SWEbenchInstance](tracer)
+    store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
+    await store.start()
+
+    # Enable callbacks for training data extraction if using vLLM
+    callbacks = ["return_token_ids", "opentelemetry", "logprobs"] if backend_type == "vllm" else ["opentelemetry"]
+    llm_proxy = LLMProxy(port=12358, store=store, callbacks=callbacks)
+
+    # Configure Models based on backend type
+    model_configs: List[ModelConfig] = []
+    model_params: Dict[str, Any] = {}
+
+    if backend_type == "vllm":
+        model_namespace = "hosted_vllm"
+        if api_base_url:
+            model_params["api_base"] = api_base_url
+        else:
+            raise ValueError("api_base_url is required for vllm backend")
+    elif backend_type == "anthropic":
+        model_namespace = "anthropic"
+        model_params["api_key"] = "os.environ/ANTHROPIC_API_KEY"
+        if api_base_url:
+            model_params["api_base"] = api_base_url
+    elif backend_type == "openai":
+        model_namespace = "openai"
+        model_params["api_key"] = "os.environ/OPENAI_API_KEY"
+        if api_base_url:
+            # Users can still override this via environment variables,
+            # even if they don't pass it in as an argument.
+            model_params["api_base"] = api_base_url
+
+    model_configs.extend(
+        [
+            ModelConfig(
+                model_name=sonnet_frontend_name,
+                litellm_params={
+                    "model": f"{model_namespace}/{sonnet_backend_name}",
+                    **model_params,
+                },
+            ),
+            ModelConfig(
+                model_name=haiku_frontend_name,
+                litellm_params={
+                    "model": f"{model_namespace}/{haiku_backend_name}",
+                    **model_params,
+                },
+            ),
+        ]
+    )
+
+    logger.info(f"Updating model list: {model_configs}")
+
+    llm_proxy.update_model_list(model_configs)
+    await llm_proxy.start()
+
+    try:
+        # Add the LLM proxy as a resource to the store
+        await store.add_resources({"llm": llm_proxy.as_resource(model="local")})
+
+        # Prepare for triplet extraction if vllm
+        adapter = ExtendedLlmProxyTraceToTriplet() if backend_type == "vllm" else None
+        tokenizer = None
+        if backend_type == "vllm":
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(sonnet_backend_name)  # type: ignore
+            except Exception as e:
+                logger.warning(f"Could not load tokenizer for {sonnet_backend_name}: {e}")
+
+        # Load full swebench dataset. Mainly for evaluation purposes.
+        swebench_full_dataset = load_swebench_dataset("princeton-nlp/SWE-bench", split="test")
+
+        # Initialize Claude Code Agent
+        claude_code_agent = ClaudeCodeAgent(swebench_full_dataset=swebench_full_dataset, max_turns=max_turns)
+
+        # Execution Loop
+        for instance in dataset:
+            await run_instance_async(
+                instance,
+                claude_code_agent,
+                runner,
+                store,
+                output_dir,
+                adapter,
+                cast(PreTrainedTokenizerBase, tokenizer),
+            )
+            # Basic sleep to allow resource cleanup or rate limit cooling
+            await asyncio.sleep(cooldown_seconds)
+
+    finally:
+        await llm_proxy.stop()
+        await store.stop()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Run Claude Code Agent experiments.")
+
+    # Backend Selection
+    parser.add_argument(
+        "backend_type",
+        type=str,
+        choices=["vllm", "anthropic", "openai"],
+        help="Backend type: 'vllm' for hosted models, 'anthropic' for official API, 'openai' for OpenAI API.",
+    )
+
+    # Model Configuration
+    parser.add_argument(
+        "--backend-model-high",
+        type=str,
+        default=None,
+        help="Backend model path/name for expensive model usages (used as vLLM model name / OpenAI model name).",
+    )
+    parser.add_argument(
+        "--backend-model-low",
+        type=str,
+        default=None,
+        help="Backend model path/name for low-price model usages (used as vLLM model name / OpenAI model name).",
+    )
+    parser.add_argument(
+        "--base-url", type=str, default="http://localhost:8000/v1", help="LLM server address (required for vllm)."
+    )
+
+    # Frontend/Agent Configuration
+    parser.add_argument(
+        "--frontend-model-high",
+        type=str,
+        default="claude-sonnet-4-5-20250929",
+        help="The frontend high-price model name provided to Claude Code.",
+    )
+    parser.add_argument(
+        "--frontend-model-low",
+        type=str,
+        default="claude-haiku-4-5-20251001",
+        help="The frontend low-price model name provided to Claude Code.",
+    )
+
+    # Execution Configuration
+    parser.add_argument("--dataset-path", type=str, default="swebench_samples.jsonl", help="Path to the dataset.")
+    parser.add_argument("--max-turns", type=int, default=5, help="Maximum turns per instance.")
+    parser.add_argument("--output-dir", type=str, default="data", help="Directory to save output logs.")
+    parser.add_argument("--limit", type=int, default=None, help="Limit the number of instances to run (for debugging).")
+    parser.add_argument("--cooldown-seconds", type=float, default=2.0, help="Cooldown seconds between instances.")
+    parser.add_argument("--debug", action="store_true", help="Enable debug loggings.")
+
+    args = parser.parse_args()
+
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    if args.debug:
+        setup_logging()
+        setup_module_logging("DEBUG", name="claude_code_agent")
+    else:
+        setup_logging(apply_to=[logger.name])
+
+    # Map backend_type to the appropriate args
+    backend_mode = cast(Literal["vllm", "anthropic", "openai"], args.backend_type)
+
+    # If using anthropic, the backend name usually matches the frontend or is specific API string.
+    # Otherwise, the backend name is the model name/path (e.g., Qwen/...) and must be provided.
+    if args.backend_model_high is None:
+        if args.backend_type == "anthropic":
+            backend_model_high = args.frontend_model_high
+        else:
+            raise ValueError("--backend-model-high is required for non-anthropic backends")
+    else:
+        backend_model_high = args.backend_model_high
+
+    if args.backend_model_low is None:
+        if args.backend_type == "anthropic":
+            backend_model_low = args.frontend_model_low
+        else:
+            raise ValueError("--backend-model-low is required for non-anthropic backends")
+    else:
+        backend_model_low = args.backend_model_low
+
+    asyncio.run(
+        dry_run_claude_code(
+            dataset_path=args.dataset_path,
+            haiku_frontend_name=args.frontend_model_low,
+            haiku_backend_name=backend_model_low,
+            sonnet_frontend_name=args.frontend_model_high,
+            sonnet_backend_name=backend_model_high,
+            backend_type=backend_mode,
+            api_base_url=args.base_url if backend_mode == "vllm" else None,
+            output_dir=args.output_dir,
+            max_turns=args.max_turns,
+            limit=args.limit,
+            cooldown_seconds=args.cooldown_seconds,
+        )
+    )

examples/claude_code/claude_code_controller.py

@@ -0,0 +1,227 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Controller module for managing Claude Code executions in containerized environments.
+
+This module provides the ClaudeController class that manages the execution of Claude Code
+within Docker containers. It handles container initialization, command execution, and
+patch application for SWE-bench evaluation tasks.
+"""
+
+import logging
+from functools import partial
+from typing import Literal, TypedDict
+
+import dotenv
+from swebench.harness.constants import SWEbenchInstance
+from swebench_utils.docker_runtime import Runtime
+from swebench_utils.logging import log_for_evaluation
+
+SWEBENCH_EXTRA_SYSTEM_PROMPT = """
+ You are an expert software engineer solving swebench bug fixing tasks.
+"""
+
+SWEBENCH_USER_PROMPT = """
+You are given a code repository in the current directory (/testbed).
+The bug description is:
+{description}
+=================================================
+You task is to fix the bug with the following steps:
+(1) write test cases to reproduce the bug.
+(2) explore the source codes to locate the bug.
+(3) edit the source codes to fix the bug.
+(4) rerun your written test cases to validate that the bug is fixed. If not, go back to explore the source codes and fix the codes again.
+(5) remember to delete the test cases you write at last.
+Please do not commit your edits. We will do it later.
+"""
+
+logger = logging.getLogger("claude_code_agent")
+
+
+class RunInstanceResult(TypedDict):
+    instance_id: str
+    model_patch: str
+    model_name_or_path: str
+
+
+class ClaudeController:
+    """Manages the execution of Claude Code within a Docker runtime.
+
+    This controller handles the lifecycle of a SWE-bench task execution, including
+    environment setup, tool installation, agent execution (via CLI or Python SDK),
+    and result extraction.
+
+    Attributes:
+        container: The active Docker runtime session.
+    """
+
+    def __init__(self, image: str, instance: SWEbenchInstance, run_id: str, endpoint: str, api_key: str) -> None:
+        """Initialize the ClaudeController.
+
+        Args:
+            image: The Docker image tag.
+            instance: The dataset instance containing the problem statement and ID.
+            run_id: The identifier for the evaluation run.
+            endpoint: The API endpoint URL.
+            api_key: The API authentication key.
+        """
+        self.image = image
+        self.instance = instance
+        self.run_id = run_id
+        self.endpoint = endpoint
+        self.api_key = api_key
+        self.container: Runtime = self.init_container(self.image, self.instance)
+
+    def init_container(self, image: str, instance: SWEbenchInstance) -> Runtime:
+        """Initializes the Docker container and sets up the Claude Code environment.
+
+        This method starts the container session, installs the Claude CLI,
+        configures environment variables for authentication and sandbox mode.
+
+        Args:
+            image: The Docker image tag to start.
+            instance: The dataset instance to load into the environment.
+
+        Returns:
+            An initialized and configured Docker runtime object.
+        """
+        container = Runtime.start_session(
+            image,
+            instance,
+            log_function=partial(log_for_evaluation, run_id=self.run_id, instance_id=instance["instance_id"]),
+        )
+        # Install Claude CLI
+        container.send_command("curl -fsSL https://claude.ai/install.sh | bash")
+        container.send_command('alias claude="$HOME/.local/bin/claude"')
+
+        # Configure Environment
+        dotenv.load_dotenv()
+        container.send_command(f"export ANTHROPIC_BASE_URL={self.endpoint}")
+        container.send_command(f"export ANTHROPIC_AUTH_TOKEN={self.api_key}")
+        container.send_command("export IS_SANDBOX=1")
+
+        return container
+
+    def _run_cli(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
+        """Executes Claude Code using the Command Line Interface.
+
+        Constructs a safe heredoc for the prompt to avoid shell interpolation issues
+        and executes the `claude` binary directly.
+
+        Args:
+            instance: The problem instance containing the problem statement.
+            max_turns: The maximum number of interaction turns allowed.
+            time_limit: The execution time limit in minutes.
+        """
+        # Prepare prompt safely: write it to a file inside the container using a single-quoted heredoc
+        # directly applying prompt for heredoc may raise error for windows line ending \r\n
+        prompt_text = SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
+
+        # Choose a simple filename and a heredoc delimiter unlikely to collide
+        heredoc_cmd = "cat > /tmp/cc_prompt.txt <<'CC_PROMPT'\n" + prompt_text + "\nCC_PROMPT\n"
+        self.container.send_command(heredoc_cmd)
+
+        # Run claude reading the prompt from the file
+        claude_cmd = (
+            f'claude -p "$(cat /tmp/cc_prompt.txt)" '
+            f'--append-system-prompt "{SWEBENCH_EXTRA_SYSTEM_PROMPT}" '
+            f"--max-turns {max_turns} "
+            f"--dangerously-skip-permissions "
+            f"--output-format json --verbose"
+        )
+        logger.info(f"Running Claude Code CLI command: {claude_cmd}")
+        self.container.send_command(claude_cmd, time_limit * 60)
+        logger.info(f"Claude Code CLI command completed")
+
+    def _run_python_sdk(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
+        """Executes Claude Code using the Python SDK wrapper.
+
+        Installs the Python SDK if necessary, hydrates a template script with the
+        problem prompt, and executes the generated Python script.
+
+        Note:
+            This path is still under development and not yet stable.
+
+        Args:
+            instance: The problem instance containing the problem statement.
+            max_turns: The maximum number of interaction turns allowed.
+            time_limit: The execution time limit in minutes.
+        """
+        # Ensure Python 3.12 is available
+        self.container.send_command(
+            f"""
+if ! command -v python3 &> /dev/null; then
+    echo "Python is not installed. Installing Python 3.12..."
+    sudo apt-get update -qq && sudo apt-get install -y -qq python3.12
+else
+    echo "Python is already installed."
+fi
+"""
+        )
+        self.container.send_command("python3 -m pip install claude-code-sdk")
+
+        # Load and fill the execution template
+        with open("src/agent/cc/claude_code_main.py.template") as f:
+            entrance_template = f.read()
+
+        script_content = (
+            entrance_template.replace("SYS_PROMPT", SWEBENCH_EXTRA_SYSTEM_PROMPT)
+            .replace(
+                "PROMPT", SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
+            )
+            .replace("MAX_STEP", str(max_turns))
+        )
+
+        # Write the script to the container and execute
+        self.container.send_command(f"cat > /tmp/claude_code_main.py <<'CC_MAIN'\n{script_content}\nCC_MAIN\n")
+        self.container.send_command("python3 /tmp/claude_code_main.py", time_limit * 60)
+        return
+
+    def run_instance(
+        self,
+        instance: SWEbenchInstance,
+        max_turns: int = 40,
+        time_limit: int = 30,
+        run_method: Literal["python", "cli"] = "python",
+    ) -> RunInstanceResult:
+        """Runs the agent on a specific SWE-bench instance.
+
+        This method orchestrates the agent execution via the specified method (CLI or Python),
+        and extracts the generated git diff (patch) upon completion.
+
+        Args:
+            instance: The dataset instance dictionary.
+            max_turns: Maximum conversation turns allowed for the agent. Defaults to 40.
+            time_limit: Time limit for the execution in minutes. Defaults to 30.
+            run_method: The execution method, either "python" (SDK) or "cli". Defaults to "python".
+
+        Returns:
+            A dictionary containing the result:
+
+            - instance_id: The ID of the processed instance.
+            - model_patch: The git diff generated by the agent.
+            - model_name_or_path: Hardcoded to "cc" (Claude Code).
+
+        Raises:
+            ValueError: If `run_method` is not "python" or "cli".
+        """
+        if run_method == "python":
+            logger.warning("Running Claude Code using Python SDK is still under development and not yet stable.")
+            self._run_python_sdk(instance, max_turns, time_limit)
+        elif run_method == "cli":
+            self._run_cli(instance, max_turns, time_limit)
+        else:
+            raise ValueError(f"Wrong run_method '{run_method}', run_method should be in ['python', 'cli']")
+
+        result = self.container.send_command("git --no-pager diff HEAD")
+        git_diff = result.output.replace("git --no-pager diff HEAD\n", "")
+
+        return {
+            "instance_id": instance["instance_id"],
+            "model_patch": git_diff,
+            "model_name_or_path": "cc",
+        }
+
+    def __del__(self) -> None:
+        """Destructor to ensure container resources are cleaned up."""
+        if hasattr(self, "container"):
+            self.container.cleanup()

examples/claude_code/extended_adapter.py

@@ -0,0 +1,163 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Custom adapter module for converting LLM proxy traces to augmented trajectories.
+
+This module provides an augmented LlmProxyTraceToTriplet adapter that converts
+LLM proxy spans into augmented trajectories for analysis and evaluation.
+It extends the base LlmProxyTraceToTriplet to include additional metadata like chat messages,
+log probabilities, and sequence IDs.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from agentlightning.adapter.triplet import LlmProxyTraceToTriplet
+from agentlightning.types import Span, Triplet
+
+logger = logging.getLogger(__name__)
+
+
+class ExtendedLlmProxyTraceToTriplet(LlmProxyTraceToTriplet):
+    """Convert LLM Proxy spans into trajectories with logprobs and customized metadata.
+
+    Augmented fields include:
+
+    - chat messages history from [`llm.hosted_vllm.messages`], saved to `Triplet.metadata['messages']`
+    - logprobs from [`llm.hosted_vllm.choices`], saved to `Triplet.response['logprobs']`
+    - sequence_id from [`Span.sequence_id`] to locate the order of the span (conversation turn), saved to `Triplet.metadata['sequence_id']`
+    """
+
+    def _extract_tokens_from_raw(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int], List[float]]:  # type: ignore
+        """Extract token ids from raw_gen_ai_request attributes.
+
+        - llm.hosted_vllm.prompt_token_ids: string -> List[int]
+        - llm.hosted_vllm.choices: string -> [{'token_ids': [...]}] -> take first
+        """
+        prompt_ids: List[int] = []
+        resp_ids: List[int] = []
+        logprobs: List[float] = []
+
+        # prompt
+        p = attrs.get("llm.hosted_vllm.prompt_token_ids")
+        p = self._literal_eval_maybe(p)
+        if isinstance(p, list) and all(isinstance(x, int) for x in p):  # type: ignore
+            prompt_ids = cast(List[int], p)
+
+        choices = attrs.get("llm.hosted_vllm.choices")
+        choices = self._literal_eval_maybe(choices)
+        if isinstance(choices, list) and choices:
+            cand = cast(Any, choices[0])
+            if isinstance(cand, dict):
+                tids = cast(Dict[str, Any], cand).get("token_ids")
+                if isinstance(tids, list) and all(isinstance(x, int) for x in tids):  # type: ignore
+                    resp_ids = cast(List[int], tids)
+
+                if "logprobs" in cand:
+                    logprobs_dict = cast(Dict[str, Any], cand).get("logprobs")
+                    if isinstance(logprobs_dict, dict) and "content" in logprobs_dict:
+                        content = cast(List[Dict[str, Any]], logprobs_dict["content"])
+                        logprobs = [float(item["logprob"]) for item in content if "logprob" in item]
+
+        return prompt_ids, resp_ids, logprobs
+
+    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
+        """Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.
+
+        Args:
+            source: Spans emitted by the LLM Proxy containing prompt, response, and reward data.
+
+        Returns:
+            Ordered trajectory transitions matched purely by `sequence_id`.
+        """
+        # 1) Sort deterministically by (sequence_id, start_time).
+        spans = sorted(
+            source,
+            key=lambda s: (s.sequence_id, s.start_time),
+        )
+
+        # 2) Collect LLM calls
+        llm_items: List[Dict[str, Any]] = []
+        seen_request_ids: set[str] = set()
+        for s in spans:
+            attrs = s.attributes or {}
+            prompt_ids: List[int] = []
+            resp_ids: List[int] = []
+            logprobs: List[float] = []
+
+            if s.name == "raw_gen_ai_request":
+                prompt_ids, resp_ids, logprobs = self._extract_tokens_from_raw(attrs)
+
+                if len(prompt_ids) == 0 or len(resp_ids) == 0:
+                    logger.warning(
+                        f"Span {s.span_id} is missing prompt (len={len(prompt_ids)}) or response (len={len(resp_ids)}) token ids. Ignoring this span."
+                    )
+                    continue
+                elif len(logprobs) == 0:
+                    logger.warning(f"Span {s.span_id} is missing logprobs. Ignoring logprobs for this span.")
+                    continue
+                elif len(resp_ids) != len(logprobs):
+                    logger.warning(
+                        f"Span {s.span_id} has mismatched response ids and logprobs lengths: "
+                        f"{len(resp_ids)} vs {len(logprobs)}. Ignoring this span."
+                    )
+                    continue
+
+            if prompt_ids and resp_ids and logprobs:
+                rid = self._request_id_from_attrs(attrs)
+                if rid:
+                    # Duplicated request ID. This request is already handled.
+                    if rid in seen_request_ids:
+                        continue
+                    seen_request_ids.add(rid)
+                llm_items.append(
+                    dict(
+                        span=s,
+                        seq=s.sequence_id,
+                        response_ids=resp_ids,
+                        prompt_ids=prompt_ids,
+                        request_id=rid,
+                        logprobs=logprobs,
+                    )
+                )
+
+        # Order LLM items by sequence only.
+        llm_items.sort(key=lambda x: x["seq"])
+
+        # Collect rewards by sequence only.
+        rewards: List[Tuple[int, Optional[float]]] = []
+        for s in spans:
+            val = self._maybe_reward_value(s)
+            if val is not None:
+                rewards.append((s.sequence_id, val))
+
+        # First-occurrence matching by sequence_id only:
+        # For reward at sequence R, assign to the most recent unmatched LLM with seq < R.
+        assigned: Dict[str, Optional[float]] = {}
+        for r_seq, r_val in sorted(rewards, key=lambda x: x[0]):
+            for item in reversed(llm_items):
+                sid = item["span"].span_id
+                if sid in assigned:
+                    continue
+                if item["seq"] < r_seq:
+                    assigned[sid] = r_val
+                    break
+
+        # Build triplets in LLM sequence order.
+        triplets: List[Triplet] = []
+        for item in llm_items:
+            s = item["span"]
+            triplets.append(
+                Triplet(
+                    prompt={"token_ids": item["prompt_ids"]},
+                    response={"token_ids": item["response_ids"], "logprobs": item["logprobs"]},
+                    reward=assigned.get(s.span_id, None),
+                    metadata=dict(
+                        # This is called response_id to align with the other adapters.
+                        response_id=item["request_id"],
+                        sequence_id=item["seq"],
+                        messages=self._literal_eval_maybe(s.attributes.get("llm.hosted_vllm.messages")),
+                    ),
+                )
+            )
+
+        return triplets

examples/claude_code/swebench_utils/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

examples/claude_code/swebench_utils/docker_runtime.py

@@ -0,0 +1,430 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Docker runtime management for repository setup and command execution.
+
+Provides containerized environment for repository testing with command execution,
+file operations, and state management capabilities.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import queue
+import re
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from docker.errors import DockerException, ImageNotFound
+from docker.models.containers import Container
+from swebench.harness.constants import SWEbenchInstance
+from typing_extensions import Self
+
+import docker
+
+# This will log to the console for debugging purposes.
+claude_code_logger = logging.getLogger("claude_code_agent.docker_runtime")
+
+CMD_OUTPUT_PS1_BEGIN = "\n###PS1JSON###\n"
+CMD_OUTPUT_PS1_END = "\n###PS1END###"
+CMD_OUTPUT_METADATA_PS1_REGEX = re.compile(
+    r"(?m)^\s*" + re.escape(CMD_OUTPUT_PS1_BEGIN.strip()) + r"\s*(.*?)\s*" + re.escape(CMD_OUTPUT_PS1_END.strip()),
+    re.DOTALL,
+)
+ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
+TIMEOUT_EXIT_CODE = 124
+
+MEM_LIMIT = "8g"
+CPU_CORES = 4
+
+
+VAR_PATTERNS = {
+    "exit_code": re.compile(r'"exit_code":\s*(-?\d+)\s*(?:,|\})'),
+    "username": re.compile(r'"username":\s*"([^"]*)"'),
+    "hostname": re.compile(r'"hostname":\s*"([^"]*)"'),
+    "working_dir": re.compile(r'"working_dir":\s*"([^"]*)"'),
+    "py_interpreter_path": re.compile(r'"py_interpreter_path":\s*"([^"]*)"'),
+}
+
+
+@dataclass
+class CmdOutputMetadata:
+    """
+    Additional metadata captured from PS1 shell prompt.
+
+    Provides context about command execution environment including
+    exit codes, user info, working directory, and Python interpreter.
+    """
+
+    exit_code: int = -1
+    username: str | None = None
+    hostname: str | None = None
+    working_dir: str | None = None
+    py_interpreter_path: str | None = None
+
+    @classmethod
+    def matches_ps1_metadata(cls, output: str) -> List[re.Match[str]]:
+        matches: List[re.Match[str]] = []
+        for match in CMD_OUTPUT_METADATA_PS1_REGEX.finditer(output):
+            scope = match.group(1).strip()
+            try:
+                d = json.loads(scope)  # Try to parse as JSON
+                matches.append(match)
+            except json.JSONDecodeError:
+                d = cls.best_effort_match(scope)
+                if len(d) > 0:
+                    matches.append(match)
+        return matches
+
+    @classmethod
+    def best_effort_match(cls, scope: str) -> Dict[str, Any]:
+        out: Dict[str, str] = {}
+        for field, pattern in VAR_PATTERNS.items():
+            m = pattern.search(scope)
+            if m:
+                out[field] = m.group(1)
+            else:
+                out[field] = ""
+        return out
+
+    @classmethod
+    def from_ps1_match(cls, match: re.Match[str]) -> Self:
+        """
+        Extract metadata from a PS1 prompt regex match.
+
+        Args:
+            match (re.Match[str]): Regex match containing JSON metadata
+
+        Returns:
+            Self: CmdOutputMetadata instance with parsed values
+        """
+        try:
+            metadata = json.loads(match.group(1))
+        except:
+            metadata = cls.best_effort_match(match.group(1))
+        # Create a copy of metadata to avoid modifying the original
+        processed = metadata.copy()
+        # Convert numeric fields
+        if "exit_code" in metadata:
+            try:
+                processed["exit_code"] = int(float(str(metadata["exit_code"])))
+            except (ValueError, TypeError):
+                processed["exit_code"] = -1
+        return cls(**processed)
+
+
+@dataclass
+class CommandResult:
+    """
+    Result of a command execution with output and metadata.
+
+    Attributes:
+        output (str): Command output text
+        metadata (Optional[CmdOutputMetadata]): Execution context metadata
+    """
+
+    output: str
+    metadata: Optional[CmdOutputMetadata]
+
+    def to_observation(self, strip: bool = True) -> str:
+        """
+        Convert command result to formatted observation string.
+
+        Args:
+            strip (bool): Whether to truncate long output
+
+        Returns:
+            str: Formatted observation with output and context
+        """
+        # compile regex once for efficiency
+        ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
+        output = ANSI_ESCAPE.sub("", self.output).replace("\r", "")
+
+        if len(output) > 1024 * 8 and strip:
+            output = output[: 1024 * 4] + "....stripped due to length....\n" + output[-1024 * 4 :]
+
+        if self.metadata is None:
+            return f"\n{output}\n"
+        return f"""{output}
+{self.metadata.username}@{self.metadata.hostname}:{self.metadata.working_dir} $
+
+exit code: {self.metadata.exit_code}
+"""
+
+
+class Runtime:
+    """
+    Docker container runtime for repository setup and testing.
+
+    Manages a Docker container with persistent bash session, command execution,
+    file operations, and container lifecycle management.
+    """
+
+    def __init__(self, container: Container, log_function: Callable[..., None]) -> None:
+        """
+        Initialize runtime with an existing Docker container.
+
+        Args:
+            container (Container): Docker container instance to manage
+        """
+        self.container = container
+        self.logger = log_function  # Set logger early so it's available even if init fails later
+        self.sock: Any = self.container.attach_socket(params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1})  # type: ignore
+        self.output_queue: queue.Queue[bytes] = queue.Queue()
+        self._start_output_thread()
+        self._clear_initial_prompt()
+
+        json_str = json.dumps(
+            {
+                "exit_code": "$?",
+                "username": r"\u",
+                "hostname": r"\h",
+                "working_dir": r"$(pwd)",
+                "py_interpreter_path": r'$(which python 2>/dev/null || echo "")',
+            },
+            indent=2,
+        ).replace('"', r"\"")
+        ps1 = CMD_OUTPUT_PS1_BEGIN + json_str + CMD_OUTPUT_PS1_END + "\n"
+        self.send_command(f'export PROMPT_COMMAND=\'export PS1="{ps1}"\'; export PS2=""')
+        self.send_command("apt update -qq && apt install -y -qq git")
+        self.stopped = False
+
+    def _stream_output(self):
+        while True:
+            try:
+                output = self._recv_bytes(4096)
+                if not output:
+                    break
+                self.output_queue.put(output)
+            except (OSError, ConnectionError) as e:
+                print(f"Connection error in _stream_output: {e}")
+                break
+            except Exception as e:
+                # print(f"Unexpected error in _stream_output: {e}")
+                break
+
+    def _start_output_thread(self):
+        self.output_thread = threading.Thread(target=self._stream_output, daemon=True)
+        self.output_thread.start()
+        # TODO: kill the thread if main thread is stopped
+
+    def _clear_initial_prompt(self):
+        time.sleep(0.5)
+        while not self.output_queue.empty():
+            self.output_queue.get()
+
+    def _read_raw_output(self, timeout: float = 30) -> tuple[str, Optional[CmdOutputMetadata]]:
+        accumulated_output = ""
+        start_time = time.time()
+
+        while time.time() - start_time < timeout:
+            try:
+                chunk = self.output_queue.get(timeout=0.1)
+                accumulated_output += chunk.decode("utf-8", errors="ignore")
+                # PSReadLine injects ANSI + cursor control; normalize before matching
+                accumulated_clean = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
+                ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_clean)
+                if ps1_matches:
+                    break
+            except queue.Empty:
+                continue
+        accumulated_output = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
+        ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_output)
+        metadata = CmdOutputMetadata.from_ps1_match(ps1_matches[-1]) if ps1_matches else None
+        output = self._combine_outputs_between_matches(
+            accumulated_output,
+            ps1_matches,
+        )
+        return output, metadata
+
+    def _combine_outputs_between_matches(self, pane_content: str, ps1_matches: list[re.Match[str]]) -> str:
+        if len(ps1_matches) == 1:
+            return pane_content[: ps1_matches[0].start()]
+        elif len(ps1_matches) == 0:
+            return pane_content
+        output_segments: List[str] = []
+        for i in range(len(ps1_matches) - 1):
+            output_segment = pane_content[ps1_matches[i].end() + 1 : ps1_matches[i + 1].start()]
+            output_segments.append(output_segment)
+        return "\n".join(output_segments) + "\n" if output_segments else ""
+
+    def _recv_bytes(self, n: int = 4096) -> bytes:
+        # Prefer the public API on whatever object the SDK returns
+        for m in ("recv", "read"):
+            if hasattr(self.sock, m):
+                return getattr(self.sock, m)(n)
+        # Last-resort fallback for odd wrappers that still expose ._sock
+        if hasattr(self.sock, "_sock"):
+            for m in ("recv", "read"):
+                if hasattr(self.sock._sock, m):
+                    return getattr(self.sock._sock, m)(n)
+        raise TypeError(f"Don't know how to read from {type(self.sock).__name__}")
+
+    def _send_bytes(self, data: bytes) -> None:
+        if hasattr(self.sock, "_sock"):
+            for m in ("send", "sendall", "write"):
+                if hasattr(self.sock._sock, m):
+                    getattr(self.sock._sock, m)(data)
+                    return
+        for m in ("send", "sendall", "write"):
+            if hasattr(self.sock, m):
+                getattr(self.sock, m)(data)
+                return
+
+        raise TypeError(f"Don't know how to write to {type(self.sock).__name__}")
+
+    def _log_command_result(self, result: CommandResult) -> None:
+        claude_code_logger.debug("Docker runtime command finished with metadata: %s", result.metadata)
+        if len(result.output) > 2048:
+            logged_output = result.output[:1024] + "\n(... stripped due to length ...)\n" + result.output[-1024:]
+        else:
+            logged_output = result.output
+        claude_code_logger.debug(
+            "Docker runtime command finished with output (length = %d):\n%s", len(result.output), logged_output
+        )
+        # Output to the evaluation logger simultaneously
+        self.logger(text=logged_output)
+
+    def send_command(self, command: str, timeout: float = 20 * 60) -> CommandResult:
+        # Redact sensitive API keys from the command before logging
+        redacted_command = command
+        for sensitive_var in ["ANTHROPIC_AUTH_TOKEN", "API_KEY", "SECRET_KEY"]:
+            pattern = rf"(export (.*?){re.escape(sensitive_var)}(.*?)=)[^\s]+"
+            redacted_command = re.sub(pattern, rf"\1****REDACTED****", redacted_command)
+        claude_code_logger.info("Docker runtime receiving command: %s", redacted_command)
+        # Normalize newline semantics for interactive shells
+        if not command.endswith("\n"):
+            command += "\n"
+
+        while not self.output_queue.empty():
+            self.output_queue.get()
+
+        self._send_bytes(command.encode())
+
+        output, metadata = self._read_raw_output(timeout=timeout)
+        # TODO: Check exit code of the command (claude code download fail will not be caught by this)
+        if metadata is not None:
+            result = CommandResult(output=output, metadata=metadata)
+            self._log_command_result(result)
+            return result
+
+        # handle timeout
+        self._send_bytes(b"\x03")
+
+        kill_timeout = 5.0
+        kill_output, kill_metadata = self._read_raw_output(timeout=kill_timeout)
+
+        output = output + kill_output + "\n**Exited due to timeout**\n"
+        if kill_metadata is not None:
+            kill_metadata.exit_code = TIMEOUT_EXIT_CODE
+            result = CommandResult(output=output, metadata=kill_metadata)
+            self._log_command_result(result)
+            return result
+
+        fallback_metadata = CmdOutputMetadata(
+            exit_code=TIMEOUT_EXIT_CODE,
+        )
+        result = CommandResult(output=output, metadata=fallback_metadata)
+        self._log_command_result(result)
+        return result
+
+    def cleanup(self) -> None:
+        if self.stopped:
+            return
+        try:
+            claude_code_logger.info(f"Stopping container: {self.container.id}")
+            self.container.stop()
+            claude_code_logger.info(f"Removing container: {self.container.id}")
+            self.container.remove(force=True)
+            claude_code_logger.info(f"Container removed: {self.container.id}")
+            self.stopped = True
+        except Exception as e:
+            print(f"Failed to stop container: {e}")
+
+    def __del__(self):
+        self.cleanup()
+
+    @staticmethod
+    def pull_image(image_name: str) -> bool:
+        """
+        Pull Docker image from registry.
+
+        Args:
+            image_name (str): Name of the Docker image to pull
+
+        Returns:
+            bool: True if successful, False if image not found
+        """
+        client = docker.from_env()
+        try:
+            client.images.pull(image_name)
+            return True
+        except ImageNotFound:
+            return False
+
+    @classmethod
+    def start_session(
+        cls,
+        image_name: str,
+        instance: SWEbenchInstance,
+        log_function: Callable[..., None] = lambda: None,
+    ) -> Runtime:
+        """
+        Start a Docker container session for repository testing.
+
+        Args:
+            image_name (str): Base Docker image name
+            instance (dict): SWE-bench instance data with repo info
+
+        Returns:
+            SetupRuntime: Configured runtime session ready for command execution
+
+        Raises:
+            RuntimeError: If Docker is not available
+        """
+        try:
+            docker.from_env().ping()  # type: ignore
+        except DockerException:
+            raise RuntimeError("Docker is not installed or not running.")
+
+        _ = cls.pull_image(image_name)
+        client = docker.from_env(timeout=600)
+        container_id = instance["instance_id"]
+        container_name = f"git-launch-{container_id}-{str(uuid.uuid4())[:4]}"
+        info: Dict[str, str] = client.version()  # type: ignore
+        engine_os: str = (info.get("Os") or info.get("OSType") or "").lower()  # type: ignore
+        # which operating system this code is running on, note windows can run linux containers, so engine_os != (container) platform
+        extra_hosts = {"host.docker.internal": "host-gateway"} if "linux" in engine_os else None
+
+        shell_command = "/bin/bash"
+        working_dir = "/testbed"
+
+        claude_code_logger.info(
+            f"Starting container {container_name} with image {image_name}. Shell command: {shell_command}"
+        )
+        container = client.containers.run(
+            image_name,
+            name=container_name,
+            command=shell_command,
+            stdin_open=True,
+            tty=True,
+            detach=True,
+            environment={
+                "TERM": "xterm-mono",
+            },
+            working_dir=working_dir,
+            extra_hosts=extra_hosts,
+            network_mode="host",
+            cpu_quota=int(CPU_CORES * 100000),
+            mem_limit=MEM_LIMIT,
+        )
+        claude_code_logger.info(f"Container {container_name} started with ID: {container.id}")
+
+        session = cls(container, log_function=log_function)
+
+        return session

examples/claude_code/swebench_utils/evaluation.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Evaluation module for SWE-bench instance testing and grading.
+
+This module provides core functionality for evaluating model predictions on SWE-bench
+instances. It handles containerized execution of test scripts, patch application,
+and generation of evaluation reports. The module orchestrates the complete evaluation
+process including container management, patch application, test execution, and result
+grading.
+"""
+
+import json
+import traceback
+from pathlib import Path, PurePosixPath
+from typing import Any, Dict, Optional
+
+from docker.models.containers import ExecResult
+from swebench.harness.constants import (
+    APPLY_PATCH_FAIL,
+    APPLY_PATCH_PASS,
+    DOCKER_PATCH,
+    DOCKER_USER,
+    DOCKER_WORKDIR,
+    INSTANCE_IMAGE_BUILD_DIR,
+    KEY_MODEL,
+    KEY_PREDICTION,
+    LOG_INSTANCE,
+    LOG_REPORT,
+    LOG_TEST_OUTPUT,
+    RUN_EVALUATION_LOG_DIR,
+    UTF8,
+    SWEbenchInstance,
+)
+from swebench.harness.docker_build import close_logger  # type: ignore
+from swebench.harness.docker_build import (
+    BuildImageError,
+    build_container,
+    setup_logger,
+)
+from swebench.harness.docker_utils import cleanup_container  # type: ignore
+from swebench.harness.docker_utils import exec_run_with_timeout  # type: ignore
+from swebench.harness.docker_utils import remove_image  # type: ignore
+from swebench.harness.docker_utils import should_remove  # type: ignore
+from swebench.harness.docker_utils import (
+    copy_to_container,
+)
+from swebench.harness.grading import get_eval_report
+from swebench.harness.test_spec.test_spec import TestSpec, make_test_spec
+from swebench.harness.utils import EvaluationError
+
+import docker
+
+GIT_APPLY_CMDS = [
+    "git apply --verbose",
+    "git apply --verbose --reject",
+    "patch --batch --fuzz=5 -p1 -i",
+]
+
+
+def run_instance(
+    test_spec: TestSpec,
+    pred: Dict[str, Any],
+    rm_image: bool,
+    force_rebuild: bool,
+    client: docker.DockerClient,
+    run_id: str,
+    timeout: int | None = None,
+    rewrite_reports: bool = False,
+):
+    """
+    Run a single instance with the given prediction.
+
+    Args:
+        test_spec (TestSpec): TestSpec instance
+        pred (dict): Prediction w/ model_name_or_path, model_patch, instance_id
+        rm_image (bool): Whether to remove the image after running
+        force_rebuild (bool): Whether to force rebuild the image
+        client (docker.DockerClient): Docker client
+        run_id (str): Run ID
+        timeout (int): Timeout for running tests
+        rewrite_reports (bool): True if eval run is just to reformat existing report
+    """
+    # Set up logging directory
+    instance_id = test_spec.instance_id
+    model_name_or_path = pred.get(KEY_MODEL, "None").replace("/", "__")
+    log_dir = RUN_EVALUATION_LOG_DIR / run_id / model_name_or_path / instance_id
+
+    # Set up report file
+    report_path = log_dir / LOG_REPORT
+    if rewrite_reports:
+        test_output_path = log_dir / LOG_TEST_OUTPUT
+        if not test_output_path.exists():
+            raise ValueError(f"Test output file {test_output_path} does not exist")
+        report = get_eval_report(
+            test_spec=test_spec,
+            prediction=pred,
+            test_log_path=test_output_path,
+            include_tests_status=True,
+        )
+        # Write report to report.json
+        with open(report_path, "w") as f:
+            f.write(json.dumps(report, indent=4))
+        return instance_id, report
+    if report_path.exists():
+        return instance_id, json.loads(report_path.read_text())
+
+    if not test_spec.is_remote_image:
+        # Link the image build dir in the log dir
+        build_dir = INSTANCE_IMAGE_BUILD_DIR / test_spec.instance_image_key.replace(":", "__")
+        image_build_link = log_dir / "image_build_dir"
+        if not image_build_link.exists():
+            try:
+                # link the image build dir in the log dir
+                image_build_link.symlink_to(build_dir.absolute(), target_is_directory=True)
+            except:
+                # some error, idk why
+                pass
+
+    # Set up logger
+    log_dir.mkdir(parents=True, exist_ok=True)
+    log_file = log_dir / LOG_INSTANCE
+    logger = setup_logger(instance_id, log_file)
+
+    # Run the instance
+    container = None
+    try:
+        # Build + start instance container (instance image should already be built)
+        container = build_container(test_spec, client, run_id, logger, rm_image, force_rebuild)
+        container.start()
+        logger.info(f"Container for {instance_id} started: {container.id}")
+
+        # Copy model prediction as patch file to container
+        patch_file = Path(log_dir / "patch.diff")
+        patch_file.write_text(pred[KEY_PREDICTION] or "")
+        logger.info(f"Intermediate patch for {instance_id} written to {patch_file}, now applying to container...")
+        copy_to_container(container, patch_file, PurePosixPath(DOCKER_PATCH))  # type: ignore
+
+        # Attempt to apply patch to container (TODO: FIX THIS)
+        val: Optional[ExecResult] = None
+        for git_apply_cmd in GIT_APPLY_CMDS:
+            val = container.exec_run(  # type: ignore
+                f"{git_apply_cmd} {DOCKER_PATCH}",
+                workdir=DOCKER_WORKDIR,
+                user=DOCKER_USER,
+            )
+            if val.exit_code == 0:
+                logger.info(f"{APPLY_PATCH_PASS}:\n{val.output.decode(UTF8)}")
+                break
+            else:
+                logger.info(f"Failed to apply patch to container: {git_apply_cmd}")
+        if val is not None:
+            logger.info(f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}")
+            raise EvaluationError(
+                instance_id,
+                f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}",
+                logger,
+            )
+
+        # Get git diff before running eval script
+        git_diff_output_before = (
+            container.exec_run("git -c core.fileMode=false diff", workdir=DOCKER_WORKDIR).output.decode(UTF8).strip()  # type: ignore
+        )
+        logger.info(f"Git diff before:\n{git_diff_output_before}")
+
+        eval_file = Path(log_dir / "eval.sh")
+        eval_file.write_text(test_spec.eval_script)
+        logger.info(f"Eval script for {instance_id} written to {eval_file}; copying to container...")
+        copy_to_container(container, eval_file, PurePosixPath("/eval.sh"))  # type: ignore
+
+        # Run eval script, write output to logs
+        test_output, timed_out, total_runtime = exec_run_with_timeout(container, "/bin/bash /eval.sh", timeout)
+        test_output_path = log_dir / LOG_TEST_OUTPUT
+        logger.info(f"Test runtime: {total_runtime:_.2f} seconds")
+        with open(test_output_path, "w") as f:
+            f.write(test_output)
+            logger.info(f"Test output for {instance_id} written to {test_output_path}")
+            if timed_out:
+                f.write(f"\n\nTimeout error: {timeout} seconds exceeded.")
+                raise EvaluationError(
+                    instance_id,
+                    f"Test timed out after {timeout} seconds.",
+                    logger,
+                )
+
+        # Get git diff after running eval script (ignore permission changes)
+        git_diff_output_after = (
+            container.exec_run("git -c core.fileMode=false diff", workdir=str(DOCKER_WORKDIR)).output.decode(UTF8).strip()  # type: ignore
+        )
+
+        # Check if git diff changed after running eval script
+        logger.info(f"Git diff after:\n{git_diff_output_after}")
+        if git_diff_output_after != git_diff_output_before:
+            logger.info("Git diff changed after running eval script")
+
+        # Get report from test output
+        logger.info(f"Grading answer for {instance_id}...")
+        report = get_eval_report(
+            test_spec=test_spec,
+            prediction=pred,
+            test_log_path=test_output_path,
+            include_tests_status=True,
+        )
+        logger.info(f"report: {report}\n" f"Result for {instance_id}: resolved: {report[instance_id]['resolved']}")
+
+        # Write report to report.json
+        with open(report_path, "w") as f:
+            f.write(json.dumps(report, indent=4))
+        return instance_id, report
+    except EvaluationError as e:
+        error_msg = traceback.format_exc()
+        logger.info(error_msg)
+        print(e)
+    except BuildImageError as e:
+        error_msg = traceback.format_exc()
+        logger.info(error_msg)
+        print(e)
+    except Exception as e:
+        error_msg = f"Error in evaluating model for {instance_id}: {e}\n" f"{traceback.format_exc()}"
+        logger.error(error_msg)
+    finally:
+        # Remove instance container + image, close logger
+        cleanup_container(client, container, logger)
+        if rm_image:
+            remove_image(client, test_spec.instance_image_key, logger)
+        close_logger(logger)
+    return
+
+
+def evaluate(
+    prediction: Dict[str, Any],
+    instance: SWEbenchInstance,
+    cache_level: str,
+    clean: bool,
+    force_rebuild: bool,
+    run_id: str,
+    timeout: Optional[int],
+    namespace: Optional[str],
+    instance_image_tag: str,
+    rewrite_reports: bool,
+):
+    client = docker.from_env()
+    test_spec = make_test_spec(instance, namespace=namespace, instance_image_tag=instance_image_tag)
+
+    instance_image_ids = {
+        test_spec.instance_image_key,
+    }
+    existing_images = {tag for i in client.images.list(all=True) for tag in i.tags if tag in instance_image_ids}
+
+    return run_instance(
+        test_spec,
+        prediction,
+        should_remove(test_spec.instance_image_key, cache_level, clean, existing_images),
+        force_rebuild,
+        client,
+        run_id,
+        timeout,
+        rewrite_reports,
+    )

examples/claude_code/swebench_utils/logging.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Logging utility module for SWE-bench evaluation runs.
+
+This module provides a simple logging utility function that writes evaluation
+results and logs to timestamped files organized by run ID and instance ID.
+"""
+
+import datetime
+import os
+
+
+def log_for_evaluation(run_id: str, instance_id: str, text: str) -> None:
+    """Log a message for evaluation purposes of SWE-Bench.
+
+    The format follows the SWE-Bench evaluation framework.
+
+    Args:
+        run_id: The run ID of the evaluation.
+        instance_id: The instance ID of the evaluation.
+        text: The text to log.
+    """
+    os.makedirs(f"./logs/{run_id}", exist_ok=True)
+    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    with open(f"./logs/{run_id}/{instance_id}", mode="a") as f:
+        print(f"\n\n{current_time}\n{text}\n", file=f)

examples/claude_code/templates/handle_hook.template.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# Read from stdin
+input=$(cat)
+
+# Define output file
+output_file="/tmp/hook.out"
+
+# Append input followed by two newlines
+echo -e "${input}\n\n" >> "$output_file"
+
+# Exit with status 0
+exit 0

examples/claude_code/templates/settings.template.json

@@ -0,0 +1,94 @@
+{
+  "hooks": {
+    "PreToolUse": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "PostToolUse": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "Notification": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "UserPromptSubmit": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "Stop": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "SubagentStop": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "PreCompact": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "SessionStart": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ],
+    "SessionEnd": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "/tmp/handle_hook.sh"
+          }
+        ]
+      }
+    ]
+  }
+}

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.

examples/minimal/write_traces.py

@@ -22,6 +22,7 @@ from rich.console import Console
 
 from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, setup_logging
 from agentlightning.store import InMemoryLightningStore
+from agentlightning.utils.otel import get_tracer_provider
 
 console = Console()
 
@@ -61,13 +62,13 @@ async def send_traces_via_otel(use_client: bool = False):
     assert "grpc-span-1" in span_names
     assert "grpc-span-2" in span_names
     assert "grpc-span-3" in span_names
-    assert "agentlightning.reward" in span_names
+    assert "agentlightning.annotation" in span_names
 
     last_span = traces[-1]
-    assert last_span.name == "agentlightning.reward"
-    # NOTE: Try not to rely on this attribute. It may change in the future.
+    assert last_span.name == "agentlightning.annotation"
+    # NOTE: Try not to rely on this attribute like this example do. It may change in the future.
     # Use utils from agentlightning.emitter to get the reward value.
-    assert last_span.attributes["reward"] == 1.0
+    assert last_span.attributes["agentlightning.reward.0.value"] == 1.0
 
     if use_client:
         # When using client, the resource should have rollout_id and attempt_id set
@@ -90,6 +91,9 @@ async def send_traces_via_agentops(use_client: bool = False):
     # Initialize the tracer lifespan
     # One lifespan can contain multiple traces
     with tracer.lifespan(store):
+        # Inspect current tracer provider
+        get_tracer_provider(inspect=True)
+
         # Initialize the capture of one single trace for one single rollout
         async with tracer.trace_context(
             "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id

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

@@ -98,7 +98,8 @@ torch-stable = [
   # This can work for both CPU and GPU.
   "torch>=2.8.0",
   "torchvision>=0.23.0",
-  "transformers>=4.55.0",
+  # https://github.com/huggingface/transformers/issues/42369
+  "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
@@ -192,7 +193,10 @@ sql = [
 ]
 crewai = [
   # https://github.com/crewAIInc/crewAI/issues/3959
-  "crewai[tools]==1.2.0",
+  "crewai[tools]>=1.2.0,!=1.2.1,!=1.3.0,!=1.4.0,!=1.4.1,!=1.5.0",
+]
+swebench = [
+  "swebench",
 ]
 
 # Summarize into large installable groups.
@@ -203,6 +207,7 @@ agents = [
   {include-group = "sql"},
   {include-group = "anthropic"},
   {include-group = "crewai"},
+  {include-group = "swebench"},
 ]
 
 [tool.uv]

tests/algorithm/test_apo.py

@@ -12,6 +12,7 @@ import agentlightning.algorithm.apo.apo as apo_module
 from agentlightning.adapter import TraceAdapter
 from agentlightning.adapter.messages import TraceToMessages
 from agentlightning.algorithm.apo.apo import APO, RolloutResultForAPO, VersionedPromptTemplate, batch_iter_over_dataset
+from agentlightning.semconv import AGL_ANNOTATION
 from agentlightning.types import (
     Dataset,
     NamedResources,
@@ -22,7 +23,6 @@ from agentlightning.types import (
     Rollout,
     Span,
     SpanContext,
-    SpanNames,
     TraceStatus,
 )
 
@@ -120,7 +120,7 @@ def make_reward_span(rollout_id: str, attempt_id: str, reward: float, sequence_i
         trace_id=hex_id,
         span_id=span_hex,
         parent_id=None,
-        name=SpanNames.REWARD.value,
+        name=AGL_ANNOTATION,
         status=TraceStatus(status_code="OK"),
         attributes={"reward": reward},
         events=[],
@@ -424,7 +424,7 @@ async def test_get_rollout_results_adapts_spans() -> None:
     # Verify spans were serialized
     assert len(results[0]["spans"]) == 2
     assert results[0]["spans"][0]["rollout_id"] == "r-1"
-    assert results[0]["spans"][0]["name"] == SpanNames.REWARD.value
+    assert results[0]["spans"][0]["name"] == AGL_ANNOTATION
     assert results[0]["spans"][0]["attributes"]["reward"] == 1.0
     assert results[0]["spans"][1]["attributes"]["reward"] == 2.0
 

tests/benchmark/analysis.py

@@ -0,0 +1,922 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Lightweight benchmark report for the Prometheus + Grafana stack shipped with Agent Lightning."""
+
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeGuard, cast
+from urllib import error, parse, request
+
+
+class PrometheusQueryError(RuntimeError):
+    """Raised when Prometheus returns an error payload."""
+
+
+class PrometheusClient:
+    """Tiny helper around the Prometheus HTTP API."""
+
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 10.0,
+        default_time: Optional[dt.datetime] = None,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.default_time = default_time
+
+    def query_vector(self, expr: str, eval_time: Optional[dt.datetime] = None) -> List[Mapping[str, object]]:
+        params: Dict[str, str] = {"query": expr}
+        query_time = eval_time or self.default_time
+        if query_time is not None:
+            params["time"] = query_time.isoformat()
+        payload = self._get("/api/v1/query", params)
+        status = payload.get("status")
+        if not isinstance(status, str) or status != "success":
+            error_msg = payload.get("error", "unknown error")
+            raise PrometheusQueryError(str(error_msg))
+        data_obj = payload.get("data", {})
+        if isinstance(data_obj, dict):
+            data = cast(Dict[str, Any], data_obj)
+        else:
+            data = {}
+        result_type_obj = data.get("resultType")
+        result_type = result_type_obj if isinstance(result_type_obj, str) else None
+        raw_result_obj = data.get("result", [])
+        raw_result: List[object]
+        if isinstance(raw_result_obj, list):
+            raw_result = cast(List[object], raw_result_obj)
+        else:
+            raw_result = []
+        if result_type == "scalar":
+            if len(raw_result) >= 2:
+                ts = raw_result[0]
+                value = raw_result[1]
+                return [{"metric": {}, "value": [ts, value]}]
+            return []
+        vector_result: List[Mapping[str, object]] = [
+            cast(Mapping[str, object], item) for item in raw_result if isinstance(item, Mapping)
+        ]
+        if result_type == "matrix":
+            collapsed: List[Dict[str, object]] = []
+            for series in vector_result:
+                values_obj = series.get("values")
+                if isinstance(values_obj, list) and values_obj and isinstance(values_obj[-1], Sequence):
+                    last = cast(Sequence[object], values_obj[-1])
+                else:
+                    continue
+                metric_obj = series.get("metric")
+                if isinstance(metric_obj, Mapping):
+                    metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
+                else:
+                    metric = {}
+                collapsed.append({"metric": metric, "value": list(last)})
+            return cast(List[Mapping[str, object]], collapsed)
+        if result_type == "vector":
+            return vector_result
+        return []
+
+    def query_scalar(self, expr: str, eval_time: Optional[dt.datetime] = None) -> Optional[float]:
+        samples = self.query_vector(expr, eval_time=eval_time)
+        if not samples:
+            return None
+        return _sample_value(samples[0])
+
+    def _get(self, path: str, data: Optional[Mapping[str, str]] = None) -> Dict[str, Any]:
+        encoded: Optional[bytes] = None
+        if data is not None:
+            encoded = parse.urlencode(data).encode()
+        req = request.Request(f"{self.base_url}{path}", data=encoded)
+        try:
+            with request.urlopen(req, timeout=self.timeout) as resp:
+                loaded = json.loads(resp.read().decode())
+                if isinstance(loaded, dict):
+                    return cast(Dict[str, Any], loaded)
+                return {}
+        except error.URLError as exc:  # pragma: no cover - network/infra issues
+            raise PrometheusQueryError(str(exc)) from exc
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Summarize benchmark metrics from Prometheus.")
+    parser.add_argument("--prom-url", default="http://localhost:9090", help="Base URL for the Prometheus API.")
+    parser.add_argument(
+        "--store-url",
+        default="http://localhost:4747/v1/agl",
+        help="Base URL for the Lightning Store API (without the /statistics suffix).",
+    )
+    parser.add_argument("--timeout", type=float, default=10.0, help="HTTP timeout in seconds.")
+    parser.add_argument("--start", type=str, help="ISO timestamp (e.g. 2024-05-01T12:00:00Z).")
+    parser.add_argument("--end", type=str, help="ISO timestamp (default: now).")
+    parser.add_argument(
+        "--duration",
+        type=str,
+        default="5m",
+        help="Fallback duration (e.g. 5m, 1h) used when --start is omitted.",
+    )
+    parser.add_argument("--top", type=int, default=8, help="Number of rows to show per table.")
+    return parser.parse_args(argv)
+
+
+def parse_timestamp(value: Optional[str], default: Optional[dt.datetime] = None) -> Optional[dt.datetime]:
+    if value is None:
+        return default
+    try:
+        if value.endswith("Z"):
+            value = value[:-1] + "+00:00"
+        return dt.datetime.fromisoformat(value).astimezone(dt.timezone.utc)
+    except ValueError as exc:  # pragma: no cover - invalid CLI input
+        raise SystemExit(f"Invalid timestamp '{value}': {exc}") from exc
+
+
+def parse_duration(text: str) -> dt.timedelta:
+    units = {"s": 1, "m": 60, "h": 3600}
+    if text.isdigit():
+        return dt.timedelta(seconds=int(text))
+    suffix = text[-1]
+    if suffix not in units:
+        raise SystemExit(f"Unsupported duration '{text}'. Use Ns/Nm/Nh.")
+    try:
+        value = int(text[:-1])
+    except ValueError as exc:  # pragma: no cover - invalid CLI input
+        raise SystemExit(f"Invalid duration '{text}': {exc}") from exc
+    return dt.timedelta(seconds=value * units[suffix])
+
+
+def format_window(seconds: float) -> str:
+    seconds = max(int(seconds), 1)
+    return f"{seconds}s"
+
+
+def compute_rate_window(duration_seconds: float) -> str:
+    return format_window(min(duration_seconds, 60.0))
+
+
+def compute_subquery_step(duration_seconds: float) -> str:
+    step_seconds = max(int(duration_seconds / 60), 1)
+    step_seconds = min(step_seconds, 15)
+    return f"{step_seconds}s"
+
+
+def _is_http_pair(value: Any) -> TypeGuard[Tuple[Any, Any]]:
+    if not isinstance(value, tuple):
+        return False
+    try:
+        value[0]
+        value[1]
+    except IndexError:
+        return False
+    return True
+
+
+def _sample_value(sample: Mapping[str, object]) -> Optional[float]:
+    value_obj = sample.get("value")
+    if not isinstance(value_obj, Sequence):
+        return None
+    value_seq = cast(Sequence[object], value_obj)
+    if len(value_seq) < 2:
+        return None
+    candidate = value_seq[1]
+    if isinstance(candidate, (int, float)):
+        return float(candidate)
+    if isinstance(candidate, str):
+        try:
+            return float(candidate)
+        except ValueError:
+            return None
+    return None
+
+
+def vector_to_map(
+    samples: Optional[Sequence[Mapping[str, object]]],
+    labels: Sequence[str],
+) -> Dict[Any, float]:
+    mapping: Dict[Any, float] = {}
+    if not samples:
+        return mapping
+    for sample in samples:
+        metric_obj = sample.get("metric", {})
+        if isinstance(metric_obj, Mapping):
+            metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
+        else:
+            metric = {}
+        if len(labels) == 1:
+            key: Any = str(metric.get(labels[0], ""))
+        else:
+            key = tuple(str(metric.get(label, "")) for label in labels)
+        value = _sample_value(sample)
+        if value is not None:
+            mapping[key] = value
+    return mapping
+
+
+def safe_vector(client: PrometheusClient, expr: str) -> Optional[List[Mapping[str, object]]]:
+    try:
+        return client.query_vector(expr)
+    except PrometheusQueryError as exc:
+        print(f"[warn] Prometheus query failed: {exc} (expr={expr})")
+        return None
+
+
+def safe_scalar(client: PrometheusClient, expr: str) -> Optional[float]:
+    try:
+        return client.query_scalar(expr)
+    except PrometheusQueryError as exc:
+        print(f"[warn] Prometheus query failed: {exc} (expr={expr})")
+        return None
+
+
+def fetch_store_statistics(store_url: str, timeout: float) -> Optional[Dict[str, Any]]:
+    store_url = store_url.rstrip("/")
+    stats_url = f"{store_url}/statistics"
+    req = request.Request(stats_url)
+    try:
+        with request.urlopen(req, timeout=timeout) as resp:
+            loaded = json.loads(resp.read().decode())
+            if isinstance(loaded, Mapping):
+                return dict(cast(Mapping[str, Any], loaded))
+            return None
+    except error.URLError as exc:
+        print(f"[warn] Failed to fetch store statistics: {exc} (url={stats_url})")
+        return None
+    except json.JSONDecodeError as exc:
+        print(f"[warn] Failed to decode store statistics: {exc} (url={stats_url})")
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Part 1 – high level throughput
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class CollectionThroughput:
+    name: str
+    count: Optional[float]
+    per_sec: Optional[float]
+
+
+STORE_TOTAL_FIELDS = {
+    "rollouts": "total_rollouts",
+    "spans": "total_spans",
+    "attempts": "total_attempts",
+    "resources": "total_resources",
+    "workers": "total_workers",
+}
+STORE_TOTAL_COLLECTIONS = tuple(STORE_TOTAL_FIELDS.keys())
+
+
+def _coerce_int(value: Any) -> Optional[int]:
+    if isinstance(value, bool):
+        return int(value)
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        if math.isnan(value):
+            return None
+        return int(value)
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError:
+            try:
+                return int(float(value))
+            except ValueError:
+                return None
+    return None
+
+
+def extract_store_totals(stats: Optional[Mapping[str, Any]]) -> Dict[str, Optional[int]]:
+    totals: Dict[str, Optional[int]] = {}
+    if not stats:
+        return totals
+    for display_name, field_name in STORE_TOTAL_FIELDS.items():
+        if field_name in stats:
+            totals[display_name] = _coerce_int(stats.get(field_name))
+        else:
+            totals[display_name] = None
+    return totals
+
+
+def gather_collection_throughput(
+    client: PrometheusClient, collections: Sequence[str], duration_seconds: float
+) -> List[CollectionThroughput]:
+    rows: List[CollectionThroughput] = []
+    window = format_window(duration_seconds)
+    for collection in collections:
+        # Successful insert operations reflect the number of new records.
+        expr = (
+            "sum("
+            f'increase(mongo_operation_total{{collection="{collection}", operation="insert", status="ok"}}[{window}])'
+            ")"
+        )
+        count = safe_scalar(client, expr)
+        if count is not None and count < 0:
+            count = 0.0
+        per_sec = (count / duration_seconds) if (count is not None and duration_seconds > 0) else None
+        rows.append(CollectionThroughput(collection, count, per_sec))
+    return rows
+
+
+# ---------------------------------------------------------------------------
+# Part 2 – CollectionBasedLightningStore method stats
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StoreMethodStats:
+    method: str
+    ops_mean: float
+    ops_max: Optional[float]
+    ops_min: Optional[float]
+    p50: Optional[float]
+    p95: Optional[float]
+    p99: Optional[float]
+
+
+StatsSummary = Dict[str, Optional[float]]
+
+
+@dataclass
+class RolloutOutcomeStats:
+    status: str
+    rate: Optional[float]
+    p25: Optional[float]
+    p50: Optional[float]
+    p75: Optional[float]
+    max_latency: Optional[float]
+
+
+def gather_store_methods(
+    client: PrometheusClient,
+    window: str,
+    rate_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",),
+    )
+    ops_max = vector_to_map(
+        safe_vector(client, f"max_over_time(({ops_expr})[{window}:{subquery_step}])"),
+        ("method",),
+    )
+    ops_min = vector_to_map(
+        safe_vector(client, f"min_over_time(({ops_expr})[{window}:{subquery_step}])"),
+        ("method",),
+    )
+    p50 = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.50, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        ),
+        ("method",),
+    )
+    p95 = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.95, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        ),
+        ("method",),
+    )
+    p99 = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.99, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        ),
+        ("method",),
+    )
+    rows: List[StoreMethodStats] = []
+    for method, rate in sorted(ops_mean.items(), key=lambda item: str(item[0])):
+        p50_val = p50.get(method)
+        p95_val = p95.get(method)
+        p99_val = p99.get(method)
+        rows.append(
+            StoreMethodStats(
+                str(method or "-"),
+                rate,
+                ops_max.get(method),
+                ops_min.get(method),
+                p50_val,
+                p95_val,
+                p99_val,
+            )
+        )
+
+    overall: StatsSummary = {
+        "ops_mean": safe_scalar(
+            client,
+            f"avg_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+        )
+        or 0.0,
+        "ops_max": safe_scalar(
+            client,
+            f"max_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+        ),
+        "ops_min": safe_scalar(
+            client,
+            f"min_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
+        ),
+        "p50": safe_scalar(
+            client,
+            f"histogram_quantile(0.50, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        )
+        or 0.0,
+        "p95": safe_scalar(
+            client,
+            f"histogram_quantile(0.95, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        )
+        or 0.0,
+        "p99": safe_scalar(
+            client,
+            f"histogram_quantile(0.99, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+        )
+        or 0.0,
+    }
+    return rows, overall
+
+
+def gather_rollout_outcomes(
+    client: PrometheusClient,
+    window: str,
+    rate_window: str,
+) -> List[RolloutOutcomeStats]:
+    rate_map = vector_to_map(
+        safe_vector(client, f"sum by (status)(rate(collection_store_rollout_total[{rate_window}]))"),
+        ("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}])))",
+        ),
+        ("status",),
+    )
+    p50_map = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.50, "
+            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+        ),
+        ("status",),
+    )
+    p75_map = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(0.75, "
+            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+        ),
+        ("status",),
+    )
+    max_map = vector_to_map(
+        safe_vector(
+            client,
+            f"histogram_quantile(1.00, "
+            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
+        ),
+        ("status",),
+    )
+    statuses = sorted({*rate_map.keys(), *p25_map.keys(), *p50_map.keys(), *p75_map.keys(), *max_map.keys()}, key=str)
+    stats: List[RolloutOutcomeStats] = []
+    for status in statuses:
+        stats.append(
+            RolloutOutcomeStats(
+                status=str(status or "-"),
+                rate=rate_map.get(status),
+                p25=p25_map.get(status),
+                p50=p50_map.get(status),
+                p75=p75_map.get(status),
+                max_latency=max_map.get(status),
+            )
+        )
+    return stats
+
+
+# ---------------------------------------------------------------------------
+# Part 3 – HTTP traffic
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class HttpPathStats:
+    method: str
+    path: 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,
+    subquery_step: str,
+) -> Tuple[List[HttpPathStats], StatsSummary]:
+    qps_expr = f"sum by (method, path)(rate(http_requests_total[{rate_window}]))"
+    qps_mean = vector_to_map(
+        safe_vector(client, f"avg_over_time(({qps_expr})[{window}:{subquery_step}])"),
+        ("method", "path"),
+    )
+    qps_max = vector_to_map(
+        safe_vector(client, f"max_over_time(({qps_expr})[{window}:{subquery_step}])"),
+        ("method", "path"),
+    )
+    qps_min = vector_to_map(
+        safe_vector(client, f"min_over_time(({qps_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}])))",
+        ),
+        ("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}])))",
+        ),
+        ("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}])))",
+        ),
+        ("method", "path"),
+    )
+
+    def normalize_http_key(raw_key: Any) -> Tuple[str, str]:
+        if _is_http_pair(raw_key):
+            method_raw, path_raw = raw_key
+            return (str(method_raw or "-"), str(path_raw or "-"))
+        return ("-", str(raw_key))
+
+    def normalize_dict(source: Mapping[Any, Optional[float]]) -> Dict[Tuple[str, str], Optional[float]]:
+        normalized: Dict[Tuple[str, str], Optional[float]] = {}
+        for raw_key, value in source.items():
+            normalized[normalize_http_key(raw_key)] = value
+        return normalized
+
+    qps_mean_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_mean))
+    qps_max_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_max))
+    qps_min_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_min))
+    p50_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p50))
+    p95_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p95))
+    p99_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p99))
+
+    path_stats: List[HttpPathStats] = []
+    for method_path in sorted(qps_mean_norm.keys()):
+        method_label, path_label = method_path
+        path_stats.append(
+            HttpPathStats(
+                method_label,
+                path_label,
+                qps_mean_norm.get(method_path, 0.0) or 0.0,
+                qps_max_norm.get(method_path),
+                qps_min_norm.get(method_path),
+                p50_norm.get(method_path),
+                p95_norm.get(method_path),
+                p99_norm.get(method_path),
+            )
+        )
+    overall: StatsSummary = {
+        "qps_mean": safe_scalar(
+            client,
+            f"avg_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+        )
+        or 0.0,
+        "qps_max": safe_scalar(
+            client,
+            f"max_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+        ),
+        "qps_min": safe_scalar(
+            client,
+            f"min_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
+        ),
+        "p50": safe_scalar(
+            client, f"histogram_quantile(0.50, sum by (le)(rate(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}])))"
+        )
+        or 0.0,
+        "p99": safe_scalar(
+            client, f"histogram_quantile(0.99, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
+        )
+        or 0.0,
+    }
+    return path_stats, overall
+
+
+# ---------------------------------------------------------------------------
+# Part 4 – diagnostics
+# ---------------------------------------------------------------------------
+
+
+def gather_diagnostics(client: PrometheusClient, window: str) -> Dict[str, Any]:
+    diagnostics: Dict[str, Any] = {}
+    diagnostics["mongo_ops"] = vector_to_map(
+        safe_vector(
+            client,
+            f"sum by (operation)(rate(mongo_operation_total{{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:
+        for sample in opcounters_samples:
+            metric_obj = sample.get("metric", {})
+            if isinstance(metric_obj, Mapping):
+                metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
+            else:
+                metric = {}
+            label_value = metric.get("legacy_op_type") or metric.get("type")
+            label = str(label_value) if label_value is not None else ""
+            value = _sample_value(sample)
+            if value is not None:
+                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["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)")
+    diagnostics["network_rx"] = safe_scalar(
+        client,
+        f"sum(rate(node_network_receive_bytes_total{{device!~'lo|docker.*'}}[{window}]))",
+    )
+    diagnostics["network_tx"] = safe_scalar(
+        client,
+        f"sum(rate(node_network_transmit_bytes_total{{device!~'lo|docker.*'}}[{window}]))",
+    )
+    diagnostics["disk_read_ops"] = safe_scalar(client, f"sum(rate(node_disk_reads_completed_total[{window}]))")
+    diagnostics["disk_write_ops"] = safe_scalar(client, f"sum(rate(node_disk_writes_completed_total[{window}]))")
+    diagnostics["disk_read_bytes"] = safe_scalar(client, f"sum(rate(node_disk_read_bytes_total[{window}]))")
+    diagnostics["disk_write_bytes"] = safe_scalar(client, f"sum(rate(node_disk_written_bytes_total[{window}]))")
+    return diagnostics
+
+
+# ---------------------------------------------------------------------------
+# Rendering helpers
+# ---------------------------------------------------------------------------
+
+
+def render_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> List[str]:
+    if not rows:
+        return [f"(no data for {headers})"]
+    widths = [len(h) for h in headers]
+    rendered: List[List[str]] = []
+    for row in rows:
+        rendered_row = [str(cell) for cell in row]
+        for idx, cell in enumerate(rendered_row):
+            widths[idx] = max(widths[idx], len(cell))
+        rendered.append(rendered_row)
+
+    lines = [
+        " | ".join(headers[idx].ljust(widths[idx]) for idx in range(len(headers))),
+        "-+-".join("-" * widths[idx] for idx in range(len(headers))),
+    ]
+    for row in rendered:
+        lines.append(" | ".join(row[idx].ljust(widths[idx]) for idx in range(len(headers))))
+    return lines
+
+
+def fmt_rate(value: Optional[float]) -> str:
+    if value is None or math.isnan(value):
+        return "-"
+    return f"{value:.2f}/s"
+
+
+def fmt_latency(value: Optional[float]) -> str:
+    if value is None or math.isnan(value):
+        return "-"
+    if value < 0.5:
+        return f"{value * 1e3:.2f} ms"
+    return f"{value:.2f} s"
+
+
+def fmt_bytes(value: Optional[float]) -> str:
+    if value is None or math.isnan(value):
+        return "-"
+    units = ["B", "KB", "MB", "GB", "TB", "PB"]
+    idx = 0
+    current = value
+    while current >= 1024 and idx < len(units) - 1:
+        current /= 1024
+        idx += 1
+    return f"{current:.2f} {units[idx]}"
+
+
+def fmt_percentage(value: Optional[float]) -> str:
+    if value is None or math.isnan(value):
+        return "-"
+    return f"{value * 100:4.1f}%"
+
+
+def section(title: str, body: Iterable[str]) -> List[str]:
+    lines = [f"## {title}"]
+    lines.extend(body)
+    lines.append("")
+    return lines
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    end = parse_timestamp(args.end, default=dt.datetime.now(dt.timezone.utc))
+    if end is None:
+        raise SystemExit("End timestamp could not be determined.")
+    start = parse_timestamp(args.start)
+    if start is None:
+        duration = parse_duration(args.duration)
+        start = end - duration
+    assert start is not None
+    duration_seconds = max((end - start).total_seconds(), 1.0)
+    window = format_window(duration_seconds)
+    rate_window = compute_rate_window(duration_seconds)
+    subquery_step = compute_subquery_step(duration_seconds)
+
+    client = PrometheusClient(args.prom_url, timeout=args.timeout, default_time=end)
+    store_stats = fetch_store_statistics(args.store_url, timeout=args.timeout)
+    store_totals = extract_store_totals(store_stats)
+    lines: List[str] = [
+        f"Agent Lightning benchmark report",
+        f"Range: {start.isoformat()} — {end.isoformat()} ({duration_seconds:.0f}s window)",
+        f"Prometheus: {args.prom_url}",
+        f"Store: {args.store_url}",
+        "",
+    ]
+
+    # Throughput
+    throughput_rows = gather_collection_throughput(
+        client, collections=STORE_TOTAL_COLLECTIONS, duration_seconds=duration_seconds
+    )
+    rows: List[List[str]] = []
+    for item in throughput_rows:
+        store_total = store_totals.get(item.name)
+        if store_total is not None:
+            count_value: Optional[int] = store_total
+        elif item.count is not None:
+            count_value = int(item.count)
+        else:
+            count_value = None
+        if count_value is None:
+            count_str = "-"
+        else:
+            count_str = f"{count_value:,}"
+        if count_value is not None and duration_seconds > 0:
+            per_sec_value = float(count_value) / duration_seconds
+        else:
+            per_sec_value = item.per_sec
+        rows.append([item.name, count_str, fmt_rate(per_sec_value)])
+    lines.extend(
+        section(
+            "Rollout / Attempt / Span / Resource / Worker Throughput",
+            render_table(["Collection", "Count", "Per Sec"], rows),
+        )
+    )
+
+    # Store internals
+    store_stats, store_overall = gather_store_methods(client, window, rate_window, subquery_step)
+    store_rows: List[List[str]] = [
+        [
+            stat.method,
+            fmt_rate(stat.ops_mean),
+            fmt_rate(stat.ops_max),
+            fmt_rate(stat.ops_min),
+            fmt_latency(stat.p50),
+            fmt_latency(stat.p95),
+            fmt_latency(stat.p99),
+        ]
+        for stat in store_stats
+    ]
+    store_lines = render_table(
+        ["Method", "Mean Ops/s", "Max Ops/s", "Min Ops/s", "P50", "P95", "P99"],
+        store_rows,
+    )
+    store_lines.append(
+        f"Overall: mean={fmt_rate(store_overall['ops_mean'])}, "
+        f"max={fmt_rate(store_overall['ops_max'])}, "
+        f"min={fmt_rate(store_overall['ops_min'])}, "
+        f"P50={fmt_latency(store_overall['p50'])}, "
+        f"P95={fmt_latency(store_overall['p95'])}, "
+        f"P99={fmt_latency(store_overall['p99'])}"
+    )
+    lines.extend(section("CollectionBasedLightningStore", store_lines))
+
+    rollout_outcomes = gather_rollout_outcomes(client, window, rate_window)
+    rollout_rows = [
+        [
+            stat.status,
+            fmt_rate(stat.rate),
+            fmt_latency(stat.p25),
+            fmt_latency(stat.p50),
+            fmt_latency(stat.p75),
+            fmt_latency(stat.max_latency),
+        ]
+        for stat in rollout_outcomes
+    ]
+    lines.extend(
+        section("Rollout Outcomes", render_table(["Status", "Rate", "P25", "P50", "P75", "Max"], rollout_rows))
+    )
+
+    # HTTP traffic
+    http_paths, http_overall = gather_http_paths(client, window, rate_window, subquery_step)
+    http_rows: List[List[str]] = [
+        [
+            stat.method,
+            stat.path,
+            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),
+        ]
+        for stat in http_paths
+    ]
+    http_lines = render_table(
+        ["Method", "Path", "Mean Req/s", "Max Req/s", "Min Req/s", "P50", "P95", "P99"], http_rows
+    )
+    http_lines.append(
+        f"Overall HTTP: mean={fmt_rate(http_overall['qps_mean'])}, "
+        f"max={fmt_rate(http_overall['qps_max'])}, "
+        f"min={fmt_rate(http_overall['qps_min'])}, "
+        f"P50={fmt_latency(http_overall['p50'])}, "
+        f"P95={fmt_latency(http_overall['p95'])}, "
+        f"P99={fmt_latency(http_overall['p99'])}"
+    )
+    lines.extend(section("HTTP Endpoints", http_lines))
+
+    # Diagnostics
+    diag = gather_diagnostics(client, window)
+    diagnostics_blocks: List[List[str]] = []
+
+    mongo_ops = cast(Dict[str, float], diag.get("mongo_ops", {}))
+    mongo_ops_rows = [
+        [operation or "-", fmt_rate(rate)]
+        for operation, rate in sorted(mongo_ops.items(), key=lambda item: str(item[0]))
+    ]
+    diagnostics_blocks.append(render_table(["Mongo Operation", "Ops/s"], mongo_ops_rows))
+
+    mongo_opcounters = cast(Dict[str, float], diag.get("mongo_opcounters", {}))
+    mongo_opcounters_rows = [
+        [op_type or "-", fmt_rate(rate)]
+        for op_type, rate in sorted(mongo_opcounters.items(), key=lambda item: str(item[0]))
+    ]
+    diagnostics_blocks.append(render_table(["MongoDB Opcounter", "Ops/s"], mongo_opcounters_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}"])
+    if mongo_misc_rows:
+        diagnostics_blocks.append(render_table(["Mongo Metric", "Value"], mongo_misc_rows))
+
+    node_rows: List[List[str]] = []
+    if diag.get("cpu_usage") is not None:
+        node_rows.append(["CPU usage", fmt_percentage(diag["cpu_usage"])])
+    mem_total = diag.get("memory_total")
+    mem_available = diag.get("memory_available")
+    if mem_total and mem_available:
+        used = mem_total - mem_available
+        node_rows.append(
+            ["Memory usage", f"{fmt_bytes(used)} / {fmt_bytes(mem_total)} ({fmt_percentage(used / mem_total)})"]
+        )
+    node_rows.append(["Network rx", f"{fmt_bytes(diag.get('network_rx'))}/s"])
+    node_rows.append(["Network tx", f"{fmt_bytes(diag.get('network_tx'))}/s"])
+    node_rows.append(["Disk read ops", fmt_rate(diag.get("disk_read_ops"))])
+    node_rows.append(["Disk read bytes", f"{fmt_bytes(diag.get('disk_read_bytes'))}/s"])
+    node_rows.append(["Disk write ops", fmt_rate(diag.get("disk_write_ops"))])
+    node_rows.append(["Disk write bytes", f"{fmt_bytes(diag.get('disk_write_bytes'))}/s"])
+    diagnostics_blocks.append(render_table(["Node Metric", "Value"], node_rows))
+
+    diagnostics_lines: List[str] = []
+    for idx, block in enumerate(diagnostics_blocks):
+        diagnostics_lines.extend(block)
+        if idx != len(diagnostics_blocks) - 1:
+            diagnostics_lines.append("")
+
+    lines.extend(section("Diagnostics", diagnostics_lines))
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":  # pragma: no cover - manual execution
+    main()

tests/benchmark/benchmark_store.py

@@ -2,15 +2,38 @@
 
 """Benchmarking store performance by writing and querying spans from the store."""
 
+import argparse
 import asyncio
+import os
 import random
-from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple
+import sys
+import threading
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple, cast
+
+from rich.console import Console
 
 import agentlightning as agl
-from agentlightning.emitter.utils import get_tracer
+from agentlightning.utils.otel import get_tracer
 
 from .utils import flatten_dict, random_dict
 
+console = Console()
+
+MAX_RUNTIME_SECONDS = 45 * 60
+
+
+def _abort_due_to_timeout() -> None:
+    sys.stderr.write(f"[benchmark] Exiting after exceeding the {MAX_RUNTIME_SECONDS // 60} minute timeout.\n")
+    sys.stderr.flush()
+    os._exit(1)
+
+
+def _start_timeout_guard(timeout_seconds: float) -> threading.Timer:
+    timer = threading.Timer(timeout_seconds, _abort_due_to_timeout)
+    timer.daemon = True
+    timer.start()
+    return timer
+
 
 def generate_attributes() -> Dict[str, Any]:
     return flatten_dict(
@@ -23,58 +46,51 @@ def generate_attributes() -> Dict[str, Any]:
     )
 
 
-@agl.rollout
-async def agent(task: str, llm: agl.LLM):
-    tracer = get_tracer()
-    rounds = random.randint(1, 10)
-    selected_round = random.randint(0, rounds - 1)
+def make_agent(max_rounds: int, sleep_seconds: float) -> agl.LitAgent[str]:
+    @agl.rollout
+    async def agent(task: str, llm: agl.LLM):
+        tracer = get_tracer()
+        rounds = random.randint(1, max_rounds)
+        selected_round = random.randint(0, rounds - 1)
 
-    for i in range(rounds):
-        with tracer.start_as_current_span(f"agent{i}") as span:
-            # Nested Span
-            with tracer.start_as_current_span(f"round{i}_1") as span:
-                await asyncio.sleep(random.uniform(0.0, 1.0))
-                span.set_attributes(generate_attributes())
-                if i == selected_round:
-                    span.set_attribute("task", task)
+        for i in range(rounds):
+            with tracer.start_as_current_span(f"agent{i}") as span:
+                # Nested Span
+                with tracer.start_as_current_span(f"round{i}_1") as span:
+                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+                    span.set_attributes(generate_attributes())
+                    if i == selected_round:
+                        span.set_attribute("task", task)
 
-            # Nested Span
-            with tracer.start_as_current_span(f"round{i}_2") as span:
-                await asyncio.sleep(random.uniform(0.0, 1.0))
-                span.set_attributes(generate_attributes())
+                # Nested Span
+                with tracer.start_as_current_span(f"round{i}_2") as span:
+                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+                    span.set_attributes(generate_attributes())
 
-        if random.uniform(0, 1) < 0.5:
-            agl.emit_reward(random.uniform(0.0, 1.0))
+            if random.uniform(0, 1) < 0.5:
+                agl.emit_reward(random.uniform(0.0, 1.0))
 
-    # Final Span
-    with tracer.start_as_current_span("final") as span:
-        await asyncio.sleep(random.uniform(0.0, 1.0))
-        span.set_attributes(generate_attributes())
+        # Final Span
+        with tracer.start_as_current_span("final") as span:
+            await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+            span.set_attributes(generate_attributes())
 
-    agl.emit_reward(random.uniform(1.0, 2.0))
+        agl.emit_reward(random.uniform(1.0, 2.0))
+
+    return agent
 
 
 def check_spans(spans: Sequence[agl.Span], task: str) -> None:
     """Check if the spans contain the task."""
-    found_task = False
-    last_reward_in_12 = None
-    for span in spans:
-        if span.attributes.get("task") == task:
-            found_task = True
-        if span.name == agl.SpanNames.REWARD.value:
-            if span.attributes.get("reward") is None:
-                raise ValueError("Reward is not set for a reward span")
-            rew = float(span.attributes.get("reward"))  # type: ignore
-            if rew >= 1 and rew <= 2:
-                last_reward_in_12 = True
-            else:
-                last_reward_in_12 = False
+    found_task = any(span.attributes.get("task") == task for span in spans)
+
+    final_reward = agl.find_final_reward(spans)
+    if final_reward is None:
+        raise ValueError("Final reward is not found")
+    if not (final_reward >= 1 and final_reward <= 2):
+        raise ValueError(f"Final reward {final_reward} is not in the range of 1 to 2")
     if not found_task:
         raise ValueError(f"Task {task} is not found in the spans")
-    if last_reward_in_12 is None:
-        raise ValueError("Last reward is not found")
-    elif not last_reward_in_12:
-        raise ValueError("Last reward is not in the range of 1 to 2")
 
 
 class AlgorithmBatch(agl.Algorithm):
@@ -122,6 +138,7 @@ class AlgorithmBatch(agl.Algorithm):
         submitted = 0
 
         while submitted < total_tasks:
+            print(f"Submitting batch {submitted} of {total_tasks}")
             batch_count = min(batch_size, total_tasks - submitted)
             batch_rollouts: List[Tuple[str, str]] = []
             await store.add_resources(
@@ -166,6 +183,7 @@ class AlgorithmBatch(agl.Algorithm):
         active_rollouts: Dict[str, str] = {}
 
         while completed < total_tasks:
+            console.print(f"Completed {completed} of {total_tasks} rollouts")
             if submitted < total_tasks and len(active_rollouts) < remaining_tasks:
                 batch_count = min(batch_size, total_tasks - submitted)
                 await store.add_resources(
@@ -217,6 +235,7 @@ class AlgorithmBatch(agl.Algorithm):
         async def handle_single(task_index: int) -> None:
             task_name = f"task-{task_index}"
             async with semaphore:
+                console.print(f"Submitting task {task_index} of {total_tasks}")
                 await store.add_resources(
                     {
                         "llm": agl.LLM(
@@ -241,20 +260,70 @@ class AlgorithmBatch(agl.Algorithm):
         await asyncio.gather(*all_tasks)
 
 
-def main() -> None:
-    store = agl.LightningStoreClient("http://localhost:4747")
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Benchmark LightningStore implementations with synthetic rollouts.")
+    parser.add_argument("--store-url", default="http://localhost:4747", help="Lightning Store endpoint base URL.")
+    parser.add_argument(
+        "--mode",
+        choices=("batch", "batch_partial", "single"),
+        default="batch",
+        help="Algorithm mode to exercise different submission patterns.",
+    )
+    parser.add_argument("--total-tasks", type=int, default=128 * 128, help="Total number of rollouts to submit.")
+    parser.add_argument("--batch-size", type=int, default=128, help="Batch size for batch-style modes.")
+    parser.add_argument(
+        "--remaining-tasks",
+        type=int,
+        default=512,
+        help="Target number of in-flight rollouts before submitting more (batch_partial mode).",
+    )
+    parser.add_argument("--concurrency", type=int, default=32, help="Maximum concurrent rollouts for single mode.")
+    parser.add_argument("--n-runners", type=int, default=32, help="Number of runner processes to launch.")
+    parser.add_argument("--max-rounds", type=int, default=10, help="Maximum number of rounds for each rollout.")
+    parser.add_argument("--sleep-seconds", type=float, default=1.0, help="Sleep seconds for each rollout.")
+    args = parser.parse_args(argv)
+
+    if args.total_tasks <= 0:
+        parser.error("--total-tasks must be positive")
+    if args.n_runners <= 0:
+        parser.error("--n-runners must be positive")
+    if args.mode in {"batch", "batch_partial"} and (args.batch_size is None or args.batch_size <= 0):
+        parser.error("--batch-size must be positive for batch modes")
+    if args.mode == "batch_partial" and (args.remaining_tasks is None or args.remaining_tasks <= 0):
+        parser.error("--remaining-tasks must be positive for batch_partial mode")
+    if args.mode == "single" and (args.concurrency is None or args.concurrency <= 0):
+        parser.error("--concurrency must be positive for single mode")
+    if args.max_rounds <= 0:
+        parser.error("--max-rounds must be positive")
+    if args.sleep_seconds <= 0:
+        parser.error("--sleep-seconds must be positive")
+
+    return args
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    store = agl.LightningStoreClient(args.store_url)
+    timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
     try:
         trainer = agl.Trainer(
             store=store,
-            algorithm=AlgorithmBatch(mode="batch", total_tasks=1024, batch_size=128),
-            n_runners=32,
+            algorithm=AlgorithmBatch(
+                mode=cast(Literal["batch", "batch_partial", "single"], args.mode),
+                total_tasks=args.total_tasks,
+                batch_size=args.batch_size,
+                remaining_tasks=args.remaining_tasks,
+                concurrency=args.concurrency,
+            ),
+            n_runners=args.n_runners,
             strategy={
                 "type": "cs",
                 "managed_store": False,
             },
         )
-        trainer.fit(agent)
+        trainer.fit(make_agent(max_rounds=args.max_rounds, sleep_seconds=args.sleep_seconds))
     finally:
+        timeout_guard.cancel()
         asyncio.run(store.close())
 
 

tests/emitter/test_annotation.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import opentelemetry.trace as trace_api
+import pytest
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import TraceFlags
+
+from agentlightning.emitter import annotation as annotation_module
+from agentlightning.emitter.annotation import emit_annotation
+from agentlightning.semconv import AGL_ANNOTATION
+
+
+class DummyReadableSpan(ReadableSpan):
+    def __init__(self) -> None:
+        super().__init__(
+            name="dummy",
+            context=trace_api.SpanContext(
+                trace_id=0x1,
+                span_id=0x2,
+                is_remote=False,
+                trace_flags=TraceFlags(TraceFlags.SAMPLED),
+                trace_state=trace_api.TraceState(),
+            ),
+            resource=Resource.create({}),
+        )
+
+    def __enter__(self) -> "DummyReadableSpan":
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
+        return False
+
+
+class DummyTracer:
+    def __init__(self, span: DummyReadableSpan) -> None:
+        self._span = span
+        self.last_name: str | None = None
+        self.last_attributes: Dict[str, Any] | None = None
+
+    def start_span(self, name: str, attributes: Dict[str, Any] | None = None) -> DummyReadableSpan:
+        self.last_name = name
+        self.last_attributes = attributes or {}
+        return self._span
+
+
+def test_emit_annotation_flattens_and_respects_propagation(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummyReadableSpan()
+    tracer = DummyTracer(span)
+    captured: Dict[str, Any] = {}
+
+    def fake_get_tracer(*_: Any, **kwargs: Any) -> DummyTracer:
+        captured["propagate"] = kwargs.get("use_active_span_processor")
+        return tracer
+
+    monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
+
+    result = emit_annotation({"meta": {"tag": "foo"}, "score": 1.5}, propagate=False)
+
+    assert result is span
+    assert captured["propagate"] is False
+    assert tracer.last_name == AGL_ANNOTATION
+    assert tracer.last_attributes == {"meta.tag": "foo", "score": 1.5}
+
+
+def test_emit_annotation_rejects_non_primitive_values() -> None:
+    with pytest.raises(TypeError):
+        emit_annotation({"bad": {"set": {1}}})

tests/emitter/test_emitter.py

@@ -1,105 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-import json
-from typing import Any, Dict, Optional
-
-import pytest
-from opentelemetry.semconv.attributes import exception_attributes
-
-from agentlightning.emitter import emit_exception, emit_message, emit_object
-from agentlightning.emitter import exception as exception_module
-from agentlightning.emitter import message as message_module
-from agentlightning.emitter import object as object_module
-from agentlightning.types.tracer import SpanAttributeNames, SpanNames
-
-
-class DummySpan:
-    def __init__(self) -> None:
-        self.recorded_exception: Optional[Exception] = None
-        self.status: Optional[Any] = None
-
-    def __enter__(self) -> "DummySpan":
-        return self
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
-        return False
-
-    def record_exception(self, exception: Exception) -> None:
-        self.recorded_exception = exception
-
-    def set_status(self, status: Any) -> None:
-        self.status = status
-
-
-class DummyTracer:
-    def __init__(self, span: DummySpan) -> None:
-        self._span = span
-        self.last_name: Optional[str] = None
-        self.last_attributes: Optional[Dict[str, Any]] = None
-
-    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
-        self.last_name = name
-        self.last_attributes = attributes or {}
-        return self._span
-
-
-def test_emit_message_valid(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = DummyTracer(span)
-    monkeypatch.setattr(message_module, "get_tracer", lambda: tracer)
-
-    emit_message("hello world")
-
-    assert tracer.last_name == SpanNames.MESSAGE.value
-    assert tracer.last_attributes == {SpanAttributeNames.MESSAGE.value: "hello world"}
-
-
-def test_emit_message_requires_string(caplog: pytest.LogCaptureFixture) -> None:
-    # emit_message logs an error but doesn't raise an exception for invalid input
-    emit_message(123)  # type: ignore[arg-type]
-    assert "Message must be a string" in caplog.text
-
-
-def test_emit_object_serializes_payload(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = DummyTracer(span)
-    monkeypatch.setattr(object_module, "get_tracer", lambda: tracer)
-
-    payload = {"foo": "bar", "baz": [1, 2, 3]}
-    emit_object(payload)
-
-    assert tracer.last_name == SpanNames.OBJECT.value
-    assert tracer.last_attributes is not None
-    assert json.loads(tracer.last_attributes[SpanAttributeNames.OBJECT.value]) == payload
-
-
-def test_emit_object_requires_json_serializable(caplog: pytest.LogCaptureFixture) -> None:
-    # emit_object logs an error but doesn't raise an exception for non-serializable input
-    emit_object(object())  # type: ignore[arg-type]
-    assert "Object must be JSON serializable" in caplog.text
-
-
-def test_emit_exception_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = DummyTracer(span)
-    monkeypatch.setattr(exception_module, "get_tracer", lambda: tracer)
-
-    exc: Optional[Exception] = None
-    try:
-        raise ValueError("boom")
-    except ValueError as err:
-        emit_exception(err)
-        exc = err
-
-    assert tracer.last_name == SpanNames.EXCEPTION.value
-    assert tracer.last_attributes is not None
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_ESCAPED] is True
-    assert span.recorded_exception is exc
-
-
-def test_emit_exception_requires_exception_instance(caplog: pytest.LogCaptureFixture) -> None:
-    # emit_exception logs an error but doesn't raise an exception for invalid input
-    emit_exception("boom")  # type: ignore[arg-type]
-    assert "Expected an BaseException instance" in caplog.text

tests/emitter/test_exception.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import pytest
+from opentelemetry.semconv.attributes import exception_attributes
+
+from agentlightning.emitter import emit_exception
+from agentlightning.emitter import exception as exception_module
+from agentlightning.semconv import AGL_EXCEPTION
+
+
+class DummySpan:
+    def __init__(self) -> None:
+        self.recorded_exception: Optional[Exception] = None
+
+    def __enter__(self) -> "DummySpan":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
+        return False
+
+    def record_exception(self, exception: Exception) -> None:
+        self.recorded_exception = exception
+
+
+class DummyTracer:
+    def __init__(self, span: DummySpan) -> None:
+        self._span = span
+        self.last_name: Optional[str] = None
+        self.last_attributes: Optional[Dict[str, Any]] = None
+
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
+        self.last_name = name
+        self.last_attributes = attributes or {}
+        return self._span
+
+
+def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
+    tracer = DummyTracer(span)
+
+    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
+        return tracer
+
+    monkeypatch.setattr(exception_module, "get_tracer", fake_get_tracer)
+    return tracer
+
+
+def test_emit_exception_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = _stub_tracer(monkeypatch, span)
+
+    exc: Optional[Exception] = None
+    try:
+        raise ValueError("boom")
+    except ValueError as err:
+        emit_exception(err)
+        exc = err
+
+    assert tracer.last_name == AGL_EXCEPTION
+    assert tracer.last_attributes is not None
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_ESCAPED] is True
+    assert span.recorded_exception is exc
+
+
+def test_emit_exception_requires_exception_instance() -> None:
+    with pytest.raises(TypeError):
+        emit_exception("boom")  # type: ignore[arg-type]

tests/emitter/test_message.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, cast
+
+import pytest
+
+from agentlightning.emitter import emit_message
+from agentlightning.emitter import message as message_module
+from agentlightning.emitter.message import get_message_value
+from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
+from agentlightning.types.tracer import SpanLike
+
+
+@dataclass
+class FakeSpan:
+    attributes: Optional[Dict[str, Any]]
+
+
+class DummySpan:
+    def __enter__(self) -> "DummySpan":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
+        return False
+
+
+class DummyTracer:
+    def __init__(self, span: DummySpan) -> None:
+        self._span = span
+        self.last_name: Optional[str] = None
+        self.last_attributes: Optional[Dict[str, Any]] = None
+
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
+        self.last_name = name
+        self.last_attributes = attributes or {}
+        return self._span
+
+
+def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
+    tracer = DummyTracer(span)
+
+    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
+        return tracer
+
+    monkeypatch.setattr(message_module, "get_tracer", fake_get_tracer)
+    return tracer
+
+
+def test_get_message_value_returns_string() -> None:
+    span = FakeSpan(attributes={LightningSpanAttributes.MESSAGE_BODY.value: "hello"})
+
+    assert get_message_value(cast(SpanLike, span)) == "hello"
+
+
+def test_get_message_value_returns_none_when_missing() -> None:
+    span = FakeSpan(attributes={})
+
+    assert get_message_value(cast(SpanLike, span)) is None
+
+
+def test_get_message_value_rejects_non_string() -> None:
+    span = FakeSpan(attributes={LightningSpanAttributes.MESSAGE_BODY.value: ["not", "string"]})
+
+    with pytest.raises(TypeError):
+        get_message_value(cast(SpanLike, span))
+
+
+def test_emit_message_valid(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = _stub_tracer(monkeypatch, span)
+
+    emit_message("hello world")
+
+    assert tracer.last_name == AGL_MESSAGE
+    assert tracer.last_attributes == {LightningSpanAttributes.MESSAGE_BODY.value: "hello world"}
+
+
+def test_emit_message_requires_string() -> None:
+    with pytest.raises(TypeError):
+        emit_message(123)  # type: ignore[arg-type]

tests/emitter/test_object.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import base64
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, cast
+
+import pytest
+
+from agentlightning.emitter import emit_object
+from agentlightning.emitter import object as object_module
+from agentlightning.emitter.object import encode_object, get_object_value
+from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
+from agentlightning.types.tracer import SpanLike
+
+
+@dataclass
+class FakeSpan:
+    attributes: Optional[Dict[str, Any]]
+
+
+class DummySpan:
+    def __enter__(self) -> "DummySpan":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
+        return False
+
+
+class DummyTracer:
+    def __init__(self, span: DummySpan) -> None:
+        self._span = span
+        self.last_name: Optional[str] = None
+        self.last_attributes: Optional[Dict[str, Any]] = None
+
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
+        self.last_name = name
+        self.last_attributes = attributes or {}
+        return self._span
+
+
+def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
+    tracer = DummyTracer(span)
+
+    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
+        return tracer
+
+    monkeypatch.setattr(object_module, "get_tracer", fake_get_tracer)
+    return tracer
+
+
+def test_encode_object_for_primitives() -> None:
+    encoded = encode_object("hello")
+
+    assert encoded == {
+        LightningSpanAttributes.OBJECT_TYPE.value: "str",
+        LightningSpanAttributes.OBJECT_LITERAL.value: "hello",
+    }
+
+
+def test_encode_object_for_bytes() -> None:
+    payload = b"binary"
+    encoded = encode_object(payload)
+    expected_literal = base64.b64encode(payload).decode("utf-8")
+
+    assert encoded == {
+        LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
+        LightningSpanAttributes.OBJECT_LITERAL.value: expected_literal,
+    }
+
+
+def test_encode_object_for_dict_serializes_json() -> None:
+    payload = {"key": 1, "nested": [1, 2]}
+    encoded = encode_object(payload)
+
+    assert encoded[LightningSpanAttributes.OBJECT_TYPE.value] == "dict"
+    assert json.loads(encoded[LightningSpanAttributes.OBJECT_JSON.value]) == payload
+
+
+def test_encode_object_raises_for_unserializable() -> None:
+    with pytest.raises(RuntimeError):
+        encode_object(object())
+
+
+def test_get_object_value_from_json_attribute() -> None:
+    payload = {"answer": 42}
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_JSON.value: json.dumps(payload),
+            LightningSpanAttributes.OBJECT_TYPE.value: "dict",
+        }
+    )
+
+    assert get_object_value(cast(SpanLike, span)) == payload
+
+
+def test_get_object_value_from_literal_types() -> None:
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_LITERAL.value: "123",
+            LightningSpanAttributes.OBJECT_TYPE.value: "int",
+        }
+    )
+    assert get_object_value(cast(SpanLike, span)) == 123
+
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_LITERAL.value: "3.14",
+            LightningSpanAttributes.OBJECT_TYPE.value: "float",
+        }
+    )
+    assert get_object_value(cast(SpanLike, span)) == pytest.approx(3.14)  # type: ignore
+
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_LITERAL.value: "true",
+            LightningSpanAttributes.OBJECT_TYPE.value: "bool",
+        }
+    )
+    assert get_object_value(cast(SpanLike, span)) is True
+
+    literal = base64.b64encode(b"hi").decode("utf-8")
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_LITERAL.value: literal,
+            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
+        }
+    )
+    assert get_object_value(cast(SpanLike, span)) == b"hi"
+
+
+def test_get_object_value_returns_none_when_missing() -> None:
+    span = FakeSpan(attributes={})
+
+    assert get_object_value(cast(SpanLike, span)) is None
+
+
+def test_get_object_value_raises_for_invalid_json() -> None:
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_JSON.value: "not json",
+        }
+    )
+
+    with pytest.raises(RuntimeError):
+        get_object_value(cast(SpanLike, span))
+
+
+def test_emit_object_serializes_payload(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = _stub_tracer(monkeypatch, span)
+
+    payload = {"foo": "bar", "baz": [1, 2, 3]}
+    emit_object(payload)
+
+    assert tracer.last_name == AGL_OBJECT
+    assert tracer.last_attributes is not None
+    assert json.loads(tracer.last_attributes[LightningSpanAttributes.OBJECT_JSON.value]) == payload
+
+
+def test_emit_object_requires_json_serializable() -> None:
+    with pytest.raises(RuntimeError):
+        emit_object(object())  # type: ignore[arg-type]
+
+
+def test_get_object_value_raises_for_unknown_literal_type() -> None:
+    span = FakeSpan(
+        attributes={
+            LightningSpanAttributes.OBJECT_LITERAL.value: "value",
+            LightningSpanAttributes.OBJECT_TYPE.value: "complex",
+        }
+    )
+
+    with pytest.raises(RuntimeError):
+        get_object_value(cast(SpanLike, span))

tests/emitter/test_reward.py

@@ -1,16 +1,18 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import importlib
 import json
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, cast
 
-from agentlightning.reward import (
-    find_final_reward,
-    find_reward_spans,
-    get_reward_value,
-    is_reward_span,
-)
-from agentlightning.types import SpanLike, SpanNames
+import pytest
+
+reward_module = importlib.import_module("agentlightning.emitter.reward")
+from agentlightning.emitter.reward import emit_reward, get_rewards_from_span
+from agentlightning.reward import find_final_reward, find_reward_spans, get_reward_value, is_reward_span
+from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
+from agentlightning.types import SpanLike
+from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
 
 
 @dataclass
@@ -19,10 +21,77 @@ class FakeSpan:
     attributes: Optional[Dict[str, Any]] = None
 
 
+@dataclass
+class AttributeSpan:
+    attributes: Optional[Dict[str, Any]]
+
+
 def make_span(name: str, attributes: Optional[Dict[str, Any]] = None) -> SpanLike:
     return cast(SpanLike, FakeSpan(name=name, attributes=attributes))
 
 
+def _capture_emit_annotation(monkeypatch: pytest.MonkeyPatch) -> tuple[Dict[str, Any], object]:
+    captured: Dict[str, Any] = {}
+    sentinel = object()
+
+    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> object:
+        captured["payload"] = payload
+        captured["propagate"] = propagate
+        return sentinel
+
+    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
+    return captured, sentinel
+
+
+def test_emit_reward_example_scalar(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured, sentinel = _capture_emit_annotation(monkeypatch)
+
+    result = emit_reward(1.0)
+
+    assert result is sentinel
+    assert captured["propagate"] is True
+    dimensions = captured["payload"][LightningSpanAttributes.REWARD.value]
+    assert dimensions == [{"name": "primary", "value": 1.0}]
+
+
+def test_emit_reward_example_multi_dimensional(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured, _ = _capture_emit_annotation(monkeypatch)
+
+    emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
+
+    dimensions = captured["payload"][LightningSpanAttributes.REWARD.value]
+    assert dimensions == [
+        {"name": "task_completion", "value": 1.0},
+        {"name": "efficiency", "value": 0.8},
+    ]
+
+
+def test_emit_reward_example_with_links(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured, _ = _capture_emit_annotation(monkeypatch)
+
+    link_attrs = make_link_attributes({"gen_ai.response.id": "response-123", "span_id": "abcd-efgh"})
+    emit_reward(0.5, attributes=link_attrs)
+
+    payload = captured["payload"]
+    assert payload[f"{LightningSpanAttributes.LINK.value}.0.key_match"] == "gen_ai.response.id"
+    assert payload[f"{LightningSpanAttributes.LINK.value}.0.value_match"] == "response-123"
+    reward_dimensions = payload[LightningSpanAttributes.REWARD.value]
+    assert reward_dimensions == [{"name": "primary", "value": 0.5}]
+
+
+def test_emit_reward_example_with_tags(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured, _ = _capture_emit_annotation(monkeypatch)
+
+    tag_attrs = make_tag_attributes(["fast", "reliable"])
+    emit_reward(0.7, attributes=tag_attrs)
+
+    payload = captured["payload"]
+    assert payload[f"{LightningSpanAttributes.TAG.value}.0"] == "fast"
+    assert payload[f"{LightningSpanAttributes.TAG.value}.1"] == "reliable"
+    reward_dimensions = payload[LightningSpanAttributes.REWARD.value]
+    assert reward_dimensions == [{"name": "primary", "value": 0.7}]
+
+
 def test_get_reward_value_from_agentops_dict() -> None:
     span = make_span(
         name="any",
@@ -43,7 +112,7 @@ def test_get_reward_value_from_agentops_json_string() -> None:
 
 def test_get_reward_value_from_reward_span_attributes() -> None:
     span = make_span(
-        name=SpanNames.REWARD.value,
+        name=AGL_ANNOTATION,
         attributes={"reward": 0.75},
     )
 
@@ -73,7 +142,7 @@ def test_is_reward_span_false_when_no_reward() -> None:
 
 def test_find_reward_spans_filters_correctly() -> None:
     reward_span = make_span(
-        name=SpanNames.REWARD.value,
+        name=AGL_ANNOTATION,
         attributes={"reward": 2.0},
     )
     non_reward_span = make_span(name="other", attributes={})
@@ -86,7 +155,7 @@ def test_find_reward_spans_filters_correctly() -> None:
 def test_find_final_reward_returns_last_reward_value() -> None:
     spans = [
         make_span(name="first", attributes={}),
-        make_span(name=SpanNames.REWARD.value, attributes={"reward": 1.0}),
+        make_span(name=AGL_ANNOTATION, attributes={"reward": 1.0}),
         make_span(name="agentops", attributes={"agentops.task.output": {"type": "reward", "value": 5.5}}),
     ]
 
@@ -100,3 +169,75 @@ def test_find_final_reward_returns_none_when_no_reward() -> None:
     ]
 
     assert find_final_reward(spans) is None
+
+
+def test_emit_reward_scalar_converts_to_primary_dimension(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: Dict[str, Any] = {}
+    sentinel_span = object()
+
+    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> object:
+        captured["payload"] = payload
+        captured["propagate"] = propagate
+        return sentinel_span
+
+    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
+
+    result = emit_reward(2, attributes={"extra": "value"}, propagate=False)
+
+    assert result is sentinel_span
+    assert captured["propagate"] is False
+    rewards = captured["payload"][LightningSpanAttributes.REWARD.value]
+    assert rewards == [{"name": "primary", "value": 2.0}]
+    assert captured["payload"]["extra"] == "value"
+
+
+def test_emit_reward_dict_requires_primary_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: Dict[str, Any] = {}
+
+    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> Dict[str, Any]:
+        captured["payload"] = payload
+        return payload
+
+    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
+
+    emit_reward({"score": 0.8, "other": 0.2}, primary_key="score")
+
+    rewards = captured["payload"][LightningSpanAttributes.REWARD.value]
+    assert [dim["name"] for dim in rewards] == ["score", "other"]
+
+    with pytest.raises(ValueError):
+        emit_reward({"score": 0.8}, primary_key=None)
+
+    with pytest.raises(ValueError):
+        emit_reward({"score": 0.8}, primary_key="missing")
+
+    with pytest.raises(ValueError):
+        emit_reward({"score": "bad"}, primary_key="score")
+
+
+def test_emit_reward_rejects_non_numeric() -> None:
+    with pytest.raises(TypeError):
+        emit_reward("bad")  # type: ignore[arg-type]
+
+
+def test_get_rewards_from_span_roundtrip() -> None:
+    attributes = {
+        f"{LightningSpanAttributes.REWARD.value}.0.name": "primary",
+        f"{LightningSpanAttributes.REWARD.value}.0.value": 1.0,
+        f"{LightningSpanAttributes.REWARD.value}.1.name": "aux",
+        f"{LightningSpanAttributes.REWARD.value}.1.value": 0.25,
+    }
+    span = AttributeSpan(attributes=attributes)
+
+    rewards = get_rewards_from_span(cast(SpanLike, span))
+
+    assert rewards == [
+        RewardPydanticModel(name="primary", value=1.0),
+        RewardPydanticModel(name="aux", value=0.25),
+    ]
+
+
+def test_get_rewards_from_span_returns_empty_when_missing() -> None:
+    span = AttributeSpan(attributes={})
+
+    assert get_rewards_from_span(cast(SpanLike, span)) == []

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -127,6 +127,7 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer, otlp_enabled:
         print(f">>> Span: {span.name}")
         print(f">>> Start time: {span.start_time}")
         print(f">>> End time: {span.end_time}")
+        print(f">>> Attributes: {span.attributes.keys()}")
         assert span.start_time is not None, f"Span {span.name} has no start time"
         assert span.end_time is not None, f"Span {span.name} has no end time"
 

tests/runner/test_agent_integration.py

@@ -67,8 +67,7 @@ async def test_runner_integration_basic_rollout() -> None:
     assert rollouts and rollouts[0].status == "succeeded"
     attempts = await store.query_attempts(rollouts[0].rollout_id)
     spans = await store.query_spans(rollouts[0].rollout_id, attempts[-1].attempt_id)
-    print(store.__dict__)
-    assert any(span.attributes.get("reward") == 1.0 for span in spans)
+    assert any(span.attributes.get("agentlightning.reward.0.value") == 1.0 for span in spans)
 
 
 @pytest.mark.asyncio
@@ -229,8 +228,10 @@ async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenA
 
         last_spans = [span for span in spans if span.sequence_id == max(span.sequence_id for span in spans)]
         assert len(last_spans) == 1
-        assert last_spans[0].name == "agentlightning.reward"
-        assert last_spans[0].attributes.get("reward") == 0.5
+        assert last_spans[0].name == "agentlightning.annotation"
+        assert (
+            last_spans[0].attributes.get("agentlightning.reward.0.value") == 0.5
+        ), f"Expected reward to be 0.5, found {last_spans[0].attributes}"
     finally:
         teardown_runner(runner)
         await proxy.stop()

tests/runner/test_agent_runner.py

@@ -17,10 +17,11 @@ from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.runner import LitAgentRunner
 from agentlightning.runner.base import Runner
+from agentlightning.semconv import AGL_ANNOTATION
 from agentlightning.store.base import UNSET, LightningStore, Unset
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.tracer.base import Tracer
-from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, SpanNames, Worker
+from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, Worker
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -259,7 +260,7 @@ async def test_step_emits_reward_for_float_result() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    rewards = [span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION]
     assert rewards == [0.75]
 
 
@@ -286,7 +287,7 @@ async def test_step_handles_non_llm_resource() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    rewards = [span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION]
     assert rewards == [0.1]
 
 
@@ -512,7 +513,9 @@ async def test_agent_emits_multiple_rewards() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    reward_values = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    reward_values = [
+        span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION
+    ]
     assert reward_values == [0.2, 0.6]
 
 

tests/store/dummy_store.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -101,17 +101,21 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("query_resources", args, kwargs))
         return self.return_values["query_resources"]
 
-    async def add_span(self, span: Span) -> Span:
+    async def add_span(self, span: Span) -> Optional[Span]:
         self.calls.append(("add_span", (span,), {}))
         return self.return_values["add_span"]
 
+    async def add_many_spans(self, spans: Sequence[Span]) -> List[Span]:
+        self.calls.append(("add_many_spans", (spans,), {}))
+        return self.return_values["add_many_spans"]
+
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: Optional[int] = None,
-    ) -> Span:
+    ) -> Optional[Span]:
         self.calls.append(("add_otel_span", (rollout_id, attempt_id, readable_span, sequence_id), {}))
         return self.return_values["add_otel_span"]
 
@@ -123,6 +127,10 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("get_next_span_sequence_id", (rollout_id, attempt_id), {}))
         return self.return_values["get_next_span_sequence_id"]
 
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> List[int]:
+        self.calls.append(("get_many_span_sequence_ids", (rollout_attempt_ids,), {}))
+        return self.return_values["get_many_span_sequence_ids"]
+
     async def query_spans(self, *args: Any, **kwargs: Any) -> List[Span]:
         self.calls.append(("query_spans", args, kwargs))
         return self.return_values["query_spans"]
@@ -206,10 +214,13 @@ def minimal_dummy_store() -> DummyLightningStore:
             "update_resources": None,
             "get_resources_by_id": None,
             "get_latest_resources": None,
+            "query_resources": [],
             "add_span": None,
+            "add_many_spans": [],
             "add_otel_span": None,
             "wait_for_rollouts": [],
             "get_next_span_sequence_id": 0,
+            "get_many_span_sequence_ids": [],
             "query_spans": [],
             "update_rollout": None,
             "update_attempt": None,

tests/store/test_client_server.py

@@ -141,6 +141,24 @@ async def test_server_accepts_custom_launcher_args(store_fixture: LightningStore
         await server.stop()
 
 
+@pytest.mark.asyncio
+async def test_server_client_statistics_match(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
+    """Server and client should report identical statistics."""
+    server, client = server_client
+    await client.start_rollout(input={"source": "statistics"})
+
+    server_stats = await server.statistics()
+    client_stats = await client.statistics()
+
+    assert {k: v for k, v in server_stats.items() if k != "uptime"} == {
+        k: v for k, v in client_stats.items() if k != "uptime"
+    }
+    assert server_stats["uptime"] < client_stats["uptime"]  # type: ignore
+    expected_name = server.store.__class__.__name__ if server.store is not None else server_stats["name"]  # type: ignore
+    assert server_stats["name"] == expected_name  # type: ignore
+    assert server_stats["total_rollouts"] >= 1  # type: ignore
+
+
 @pytest.mark.asyncio
 async def test_add_resources_via_server(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
     """Test that add_resources works correctly via server."""
@@ -394,7 +412,7 @@ async def test_client_server_end_to_end(
 
     client_span = _make_span(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id, 101, "client-span")
     stored_span = await client.add_span(client_span)
-    assert stored_span.name == "client-span"
+    assert stored_span is not None and stored_span.name == "client-span"
     assert await client.get_next_span_sequence_id(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id) == 102
 
     with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked:
@@ -703,6 +721,87 @@ async def test_client_query_spans_filters_and_pagination(
     assert [span.span_id for span in paged] == [spans[1].span_id]
 
 
+@pytest.mark.asyncio
+async def test_server_get_many_span_sequence_ids_and_add_many_spans_mixed_batches(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    server, _ = server_client
+
+    first = await server.start_rollout(input={"origin": "batch-server"})
+    second = await server.start_rollout(input={"origin": "batch-server-2"})
+    await server.update_rollout(first.rollout_id, status="requeuing")
+    retried = await server.start_attempt(first.rollout_id)
+
+    sequence_pairs = [
+        (first.rollout_id, first.attempt.attempt_id),
+        (second.rollout_id, second.attempt.attempt_id),
+        (first.rollout_id, retried.attempt.attempt_id),
+        (first.rollout_id, first.attempt.attempt_id),
+    ]
+    sequence_ids = await server.get_many_span_sequence_ids(sequence_pairs)
+    assert sequence_ids == [1, 1, 2, 3]
+
+    next_single = await server.get_next_span_sequence_id(first.rollout_id, first.attempt.attempt_id)
+    assert next_single == 4
+
+    batch_spans = [
+        _make_span(first.rollout_id, first.attempt.attempt_id, 10, "server-batch-1"),
+        _make_span(second.rollout_id, second.attempt.attempt_id, 11, "server-batch-2"),
+        _make_span(retried.rollout_id, retried.attempt.attempt_id, 12, "server-batch-retry"),
+    ]
+    stored_spans = await server.add_many_spans(batch_spans)
+    assert {span.name for span in stored_spans} == {
+        "server-batch-1",
+        "server-batch-2",
+        "server-batch-retry",
+    }
+
+    spans_first = await server.query_spans(first.rollout_id)
+    assert any(span.name == "server-batch-1" for span in spans_first)
+    assert any(span.name == "server-batch-retry" for span in spans_first)
+    spans_second = await server.query_spans(second.rollout_id)
+    assert any(span.name == "server-batch-2" for span in spans_second)
+
+
+@pytest.mark.asyncio
+async def test_client_handles_optional_span_results_and_batch_insert(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+    mock_readable_span: ReadableSpan,
+) -> None:
+    _, client = server_client
+
+    first = await client.start_rollout(input={"origin": "client-span"})
+    second = await client.start_rollout(input={"origin": "client-span-2"})
+    first_attempt_id = first.attempt.attempt_id
+    second_attempt_id = second.attempt.attempt_id
+
+    base_span = _make_span(first.rollout_id, first_attempt_id, 1, "client-span-1")
+    stored = await client.add_span(base_span)
+    assert stored is not None
+    assert await client.add_span(base_span) is None
+
+    batch_spans = [
+        _make_span(first.rollout_id, first_attempt_id, 2, "client-span-2"),
+        _make_span(second.rollout_id, second_attempt_id, 1, "client-span-other"),
+        base_span,
+    ]
+    inserted = await client.add_many_spans(batch_spans)
+    assert [span.name for span in inserted] == ["client-span-2", "client-span-other"]
+
+    sequence_ids = await client.get_many_span_sequence_ids(
+        [
+            (first.rollout_id, first_attempt_id),
+            (second.rollout_id, second_attempt_id),
+            (first.rollout_id, "latest"),
+        ]
+    )
+    assert sequence_ids == [3, 2, 4]
+
+    with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked_span_factory:
+        mocked_span_factory.return_value = base_span
+        assert await client.add_otel_span(first.rollout_id, first_attempt_id, mock_readable_span) is None
+
+
 @pytest.mark.asyncio
 async def test_concurrent_add_otel_span_sequence_ids_unique(
     server_client: Tuple[LightningStoreServer, LightningStoreClient], mock_readable_span: ReadableSpan
@@ -721,7 +820,7 @@ async def test_concurrent_add_otel_span_sequence_ids_unique(
         spans = await asyncio.gather(
             *[client.add_otel_span(rollout_id, attempt_id, mock_readable_span) for _ in range(20)]
         )
-    sequence_ids = [span.sequence_id for span in spans]
+    sequence_ids = [span.sequence_id for span in spans]  # type: ignore
     assert len(set(sequence_ids)) == 20
     assert set(sequence_ids) == set(range(1, 21))
 

tests/store/test_collection_memory.py

@@ -92,6 +92,20 @@ async def test_list_collection_insert_duplicate_raises(sample_collection: Collec
         await sample_collection.insert([duplicate])
 
 
+@pytest.mark.asyncio()
+async def test_list_collection_insert_rejects_duplicate_payload(sample_collection: Collection[SampleItem]) -> None:
+    """Ensure duplicate items within the same insert batch are rejected."""
+    starting_size = await sample_collection.size()
+    dup_a = SampleItem(partition="omega", index=1, name="dup-a", status="new")
+    dup_b = SampleItem(partition="omega", index=1, name="dup-b", status="new")
+
+    with pytest.raises(ValueError, match="duplicate primary key"):
+        await sample_collection.insert([dup_a, dup_b])
+
+    assert await sample_collection.size() == starting_size
+    assert await sample_collection.get({"partition": {"exact": "omega"}}) is None
+
+
 @pytest.mark.asyncio()
 async def test_list_collection_insert_wrong_type(sample_collection: Collection[SampleItem]) -> None:
     class Another(BaseModel):
@@ -770,6 +784,31 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
         assert not await span_kv.has("span-123")
 
 
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_based_collection_rejects_duplicate_payload(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection[Any](
+            client_pool,
+            temporary_mongo_database.name,
+            f"duplicate-check-{uuid4().hex}",
+            "partition-dup",
+            ["rollout_id"],
+            Rollout,
+        )
+        await collection.ensure_collection()
+        start_time = time.time()
+        first = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
+        duplicate = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
+
+        with pytest.raises(ValueError, match="duplicate primary key"):
+            await collection.insert([first, duplicate])
+
+        assert await collection.size() == 0
+
+
 @pytest.mark.mongo
 @pytest.mark.asyncio()
 async def test_mongo_ensure_collection_creates_partition_scoped_index(

tests/store/test_implementation.py

@@ -61,6 +61,34 @@ def test_paginated_result_behaves_like_sequence() -> None:
     assert repr(result2) == "<PaginatedResult (1: of 5) ['a', ...]>"
 
 
+@pytest.mark.asyncio
+async def test_statistics_updates_counts(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
+    """Statistics should reflect entity counts across the store."""
+    initial = await store_fixture.statistics()
+
+    rollout = await store_fixture.start_rollout(input={"origin": "stats"})
+    await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)
+    await store_fixture.add_resources(
+        {
+            "stat_llm": LLM(
+                resource_type="llm",
+                endpoint="http://localhost:8000/v1",
+                model="stats-model",
+            )
+        }
+    )
+    await store_fixture.update_worker("stats-worker", heartbeat_stats={"cpu": 0.5})
+
+    updated = await store_fixture.statistics()
+    assert updated["name"] == store_fixture.__class__.__name__  # type: ignore
+    assert updated["total_rollouts"] == initial.get("total_rollouts", 0) + 1  # type: ignore
+    assert updated["total_attempts"] == initial.get("total_attempts", 0) + 1  # type: ignore
+    assert updated["total_spans"] == initial.get("total_spans", 0) + 1  # type: ignore
+    assert updated["total_resources"] == initial.get("total_resources", 0) + 1  # type: ignore
+    assert updated["total_workers"] == initial.get("total_workers", 0) + 1  # type: ignore
+    assert updated["uptime"] > initial.get("uptime", 0.0)  # type: ignore
+
+
 # Core CRUD Operations Tests
 
 
@@ -1098,20 +1126,116 @@ async def test_span_sequence_generation(store_fixture: LightningStore, mock_read
     assert seq_id == 1
 
     span1 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert span1.sequence_id == 2
+    assert span1 is not None and span1.sequence_id == 2
 
     # Next span gets sequence_id 3
     seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id)
     assert seq_id == 3
 
     span2 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert span2.sequence_id == 4
+    assert span2 is not None and span2.sequence_id == 4
 
     # Different attempt reuses the same rollout_id
     seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, "attempt-does-not-exist")
     assert seq_id == 5
 
 
+@pytest.mark.asyncio
+async def test_get_many_span_sequence_ids_handles_mixed_batches(store_fixture: LightningStore) -> None:
+    """Bulk sequence allocation should work across mixed rollouts and attempts."""
+    first = await store_fixture.start_rollout(input={"origin": "bulk-seq"})
+    second = await store_fixture.start_rollout(input={"origin": "bulk-seq-2"})
+    await store_fixture.update_rollout(first.rollout_id, status="requeuing")
+    retried = await store_fixture.start_attempt(first.rollout_id)
+
+    sequence_pairs = [
+        (first.rollout_id, first.attempt.attempt_id),
+        (second.rollout_id, second.attempt.attempt_id),
+        (first.rollout_id, retried.attempt.attempt_id),
+        (first.rollout_id, first.attempt.attempt_id),
+    ]
+    sequence_ids = await store_fixture.get_many_span_sequence_ids(sequence_pairs)
+    assert sequence_ids == [1, 1, 2, 3]
+
+    next_first = await store_fixture.get_next_span_sequence_id(first.rollout_id, first.attempt.attempt_id)
+    next_second = await store_fixture.get_next_span_sequence_id(second.rollout_id, second.attempt.attempt_id)
+    assert next_first == 4
+    assert next_second == 2
+
+
+@pytest.mark.asyncio
+async def test_add_many_spans_handles_mixed_rollouts_and_attempts(store_fixture: LightningStore) -> None:
+    """Batch span insertion should handle mixed rollouts and skip duplicates."""
+    first = await store_fixture.start_rollout(input={"origin": "batch-spans"})
+    second = await store_fixture.start_rollout(input={"origin": "batch-spans-2"})
+    await store_fixture.update_rollout(first.rollout_id, status="requeuing")
+    retried = await store_fixture.start_attempt(first.rollout_id)
+
+    def _build_span(seed: int, rollout_id: str, attempt_id: str) -> Span:
+        trace_hex = f"{seed:032x}"
+        span_hex = f"{seed:016x}"
+        return Span(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=seed,
+            trace_id=trace_hex,
+            span_id=span_hex,
+            parent_id=None,
+            name=f"batch-span-{seed}",
+            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=""),
+        )
+
+    span_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)
+    span_retry = _build_span(2, retried.rollout_id, retried.attempt.attempt_id)
+    span_second = _build_span(3, second.rollout_id, second.attempt.attempt_id)
+    duplicate_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)
+
+    stored = await store_fixture.add_many_spans([span_first, span_retry, span_second, duplicate_first])
+    assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}
+
+    spans_first = await store_fixture.query_spans(first.rollout_id)
+    assert {span.span_id for span in spans_first} >= {span_first.span_id, span_retry.span_id}
+    spans_second = await store_fixture.query_spans(second.rollout_id)
+    assert {span.span_id for span in spans_second} >= {span_second.span_id}
+
+
+@pytest.mark.asyncio
+async def test_add_span_returns_none_on_duplicate(store_fixture: LightningStore) -> None:
+    """Duplicate span IDs should return None instead of raising."""
+    attempted = await store_fixture.start_rollout(input={"origin": "duplicate-span"})
+    attempt_id = attempted.attempt.attempt_id
+
+    span = Span(
+        rollout_id=attempted.rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=1,
+        trace_id="0" * 32,
+        span_id="0" * 16,
+        parent_id=None,
+        name="dup-span",
+        status=TraceStatus(status_code="OK"),
+        attributes={},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id="0" * 32, span_id="0" * 16, is_remote=False, trace_state={}),
+        parent=None,
+        resource=OtelResource(attributes={}, schema_url=""),
+    )
+
+    assert await store_fixture.add_span(span) is span
+    assert await store_fixture.add_span(span) is None
+
+
 @pytest.mark.asyncio
 async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
     """Adding a span should persist attempt heartbeat and transition status to running."""
@@ -1196,7 +1320,8 @@ async def test_duplicate_span_id_error(
 
     await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
 
-    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+    duplicate = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
+    assert duplicate is None
     assert "Duplicated span added" in caplog.text
 
 
@@ -1211,7 +1336,7 @@ async def test_span_with_explicit_sequence_id(store_fixture: LightningStore, moc
 
     # Add span with explicit sequence_id
     span = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span, sequence_id=100)
-    assert span.sequence_id == 100
+    assert span is not None and span.sequence_id == 100
 
     next_seq = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id)
     assert next_seq == 101
@@ -2153,7 +2278,7 @@ async def test_concurrent_span_additions(store_fixture: LightningStore, mock_rea
     assert rollout is not None
 
     async def add_span(index: int) -> Span:
-        return await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)
+        return await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)  # type: ignore
 
     # Add 30 spans concurrently
     tasks = [add_span(i) for i in range(30)]
@@ -2754,7 +2879,7 @@ async def test_full_lifecycle_success(store_fixture: LightningStore, mock_readab
 
     # 3. Add span (transitions to running)
     span = await store_fixture.add_otel_span(rollout.rollout_id, attempt.attempt_id, mock_readable_span)
-    assert span.sequence_id == 1
+    assert span is not None and span.sequence_id == 1
 
     # Check status transitions
     rollouts = await store_fixture.query_rollouts(status=["running"])

tests/store/test_restful.py

@@ -137,6 +137,23 @@ async def test_cors_allows_wildcard_origin() -> None:
                 assert allow_credentials == "true"
 
 
+@pytest.mark.asyncio
+async def test_statistics_endpoint_returns_counts(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    """The statistics endpoint should expose store-level metrics."""
+    server, client, session, api_endpoint = server_client
+    await client.start_rollout(input={"source": "stats-endpoint"})
+
+    async with session.get(f"{api_endpoint}/statistics") as resp:
+        assert resp.status == 200
+        payload = await resp.json()
+
+    expected_name = server.store.__class__.__name__ if server.store is not None else payload["name"]
+    assert payload["name"] == expected_name
+    assert payload["total_rollouts"] >= 1
+
+
 # Rollouts Pagination, Sorting, and Filtering Tests
 
 

tests/store/test_threading.py

@@ -177,9 +177,11 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "get_resources_by_id": resources_update,
         "get_latest_resources": resources_update,
         "add_span": span,
+        "add_many_spans": [span],
         "add_otel_span": span,
         "wait_for_rollouts": [base_rollout],
         "get_next_span_sequence_id": 42,
+        "get_many_span_sequence_ids": [101],
         "query_spans": [span],
         "update_rollout": updated_rollout,
         "update_attempt": updated_attempt,
@@ -211,9 +213,11 @@ async def test_threaded_store_delegates_all_methods() -> None:
     assert await threaded_store.get_resources_by_id("resources-1") == resources_update
     assert await threaded_store.get_latest_resources() == resources_update
     assert await threaded_store.add_span(span) == span
+    assert await threaded_store.add_many_spans([span]) == [span]
     assert await threaded_store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id=5) == span
     assert await threaded_store.wait_for_rollouts(rollout_ids=[rollout_id], timeout=1.0) == [base_rollout]
     assert await threaded_store.get_next_span_sequence_id(rollout_id, attempt_id) == 42
+    assert await threaded_store.get_many_span_sequence_ids([(rollout_id, attempt_id)]) == [101]
     assert await threaded_store.query_spans(rollout_id, attempt_id="latest") == [span]
     assert (
         await threaded_store.update_rollout(
@@ -254,9 +258,11 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "get_resources_by_id",
         "get_latest_resources",
         "add_span",
+        "add_many_spans",
         "add_otel_span",
         "wait_for_rollouts",
         "get_next_span_sequence_id",
+        "get_many_span_sequence_ids",
         "query_spans",
         "update_rollout",
         "update_attempt",
@@ -359,3 +365,28 @@ def test_threaded_store_serializes_add_resources_calls() -> None:
     assert len(updates) == num_adds
     resource_ids = {update.resources_id for update in updates}
     assert len(resource_ids) == num_adds  # All IDs should be unique
+
+
+@pytest.mark.asyncio
+async def test_threaded_store_statistics_match_underlying(inmemory_store: InMemoryLightningStore) -> None:
+    """Threaded store should proxy statistics from wrapped store."""
+    rollout = await inmemory_store.start_rollout(input={"source": "threaded-stats"})
+    await inmemory_store.add_span(make_span(rollout.rollout_id, rollout.attempt.attempt_id))
+    await inmemory_store.add_resources(
+        {
+            "threaded_llm": LLM(
+                resource_type="llm",
+                endpoint="http://localhost:9000/v1",
+                model="threaded-model",
+            )
+        }
+    )
+    await inmemory_store.update_worker("threaded-worker", heartbeat_stats={"cpu": 0.1})
+
+    threaded_store = LightningStoreThreaded(inmemory_store)
+    threaded_stats = await threaded_store.statistics()
+    inmemory_stats = await inmemory_store.statistics()
+    assert {k: v for k, v in threaded_stats.items() if k != "uptime"} == {
+        k: v for k, v in inmemory_stats.items() if k != "uptime"
+    }
+    assert threaded_stats["uptime"] < inmemory_stats["uptime"]  # type: ignore

tests/test_env_var.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import pytest
+
+from agentlightning.env_var import (
+    LightningEnvVar,
+    resolve_bool_env_var,
+    resolve_int_env_var,
+    resolve_str_env_var,
+)
+
+
+def test_resolve_bool_env_var_override_takes_precedence(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_MANAGED_STORE.value
+    monkeypatch.setenv(env_name, "0")
+
+    assert resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, override=True, fallback=False) is True
+
+
+@pytest.mark.parametrize(
+    ("raw_value", "expected"),
+    [
+        ("1", True),
+        (" YES ", True),
+        ("on", True),
+        ("0", False),
+        ("no", False),
+        ("Off", False),
+    ],
+)
+def test_resolve_bool_env_var_parses_truthy_and_falsy_values(
+    monkeypatch: pytest.MonkeyPatch, raw_value: str, expected: bool
+) -> None:
+    env_name = LightningEnvVar.AGL_MANAGED_STORE.value
+    monkeypatch.setenv(env_name, raw_value)
+
+    assert resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=not expected) is expected
+
+
+def test_resolve_bool_env_var_returns_fallback_when_unset(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_MANAGED_STORE.value
+    monkeypatch.delenv(env_name, raising=False)
+
+    assert resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=False) is False
+
+
+def test_resolve_bool_env_var_rejects_invalid_value(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_MANAGED_STORE.value
+    monkeypatch.setenv(env_name, "maybe")
+
+    with pytest.raises(ValueError):
+        resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=False)
+
+
+def test_resolve_int_env_var_reads_from_environment(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_SERVER_PORT.value
+    monkeypatch.setenv(env_name, "1234")
+
+    assert resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, fallback=4747) == 1234
+
+
+def test_resolve_int_env_var_invalid_input(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_SERVER_PORT.value
+    monkeypatch.setenv(env_name, "not-a-number")
+
+    with pytest.raises(ValueError):
+        resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, fallback=4747)
+
+
+def test_resolve_str_env_var_override_and_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    env_name = LightningEnvVar.AGL_CURRENT_ROLE.value
+    monkeypatch.setenv(env_name, "client")
+
+    assert resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, override="server", fallback="both") == "server"
+
+    monkeypatch.delenv(env_name, raising=False)
+    assert resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, fallback="both") == "both"

tests/tracer/test_agentops.py

@@ -1,15 +1,115 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import multiprocessing
 import sys
-from typing import Any, Optional, Union
+import threading
+import time
+from typing import Any, Callable, Coroutine, Iterator, List, Optional, Union
 
 import agentops
 import pytest
+import uvicorn
 from agentops.sdk.core import TraceContext
+from fastapi import FastAPI, Request
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+    ExportTraceServiceResponse,
+)
+from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.trace.status import StatusCode
+from portpicker import pick_unused_port
 
+from agentlightning.store.base import LightningStore, LightningStoreCapabilities
 from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.types import Span
+from agentlightning.utils import otlp
+
+
+class MockOTLPService:
+    """A mock OTLP server to capture trace export requests for testing purposes."""
+
+    def __init__(self) -> None:
+        self.received: List[ExportTraceServiceRequest] = []
+
+    def start_service(self) -> int:
+        app = FastAPI()
+
+        @app.post("/v1/traces")
+        async def _export_traces(request: Request):  # type: ignore
+            async def capture(message: ExportTraceServiceRequest) -> None:
+                self.received.append(message)
+
+            return await otlp.handle_otlp_export(
+                request,
+                ExportTraceServiceRequest,
+                ExportTraceServiceResponse,
+                capture,
+                signal_name="traces",
+            )
+
+        port = pick_unused_port()
+        config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")
+        self.server = uvicorn.Server(config)
+        self.thread = threading.Thread(target=self.server.run, daemon=True)
+        self.thread.start()
+        timeout = time.time() + 5
+        while not getattr(self.server, "started", False):
+            if time.time() > timeout:
+                raise RuntimeError("OTLP test server failed to start")
+            if not self.thread.is_alive():
+                raise RuntimeError("OTLP test server thread exited before startup")
+            time.sleep(0.01)
+
+        return port
+
+    def stop_service(self) -> None:
+        self.server.should_exit = True
+        self.thread.join(timeout=5)
+
+    def get_traces(self) -> List[ExportTraceServiceRequest]:
+        return self.received
+
+
+class MockLightningStore(LightningStore):
+    """A minimal stub-only LightningStore, only implements methods likely used in tests."""
+
+    def __init__(self, server_port: int = 80) -> None:
+        super().__init__()
+        self.otlp_traces = False
+        self.server_port = server_port
+
+    def enable_otlp_traces(self) -> None:
+        self.otlp_traces = True
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        if sequence_id is None:
+            sequence_id = 0
+
+        span = Span.from_opentelemetry(
+            readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id
+        )
+        return span
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            async_safe=False,
+            thread_safe=False,
+            zero_copy=False,
+            otlp_traces=self.otlp_traces,
+        )
+
+    def otlp_traces_endpoint(self) -> str:
+        return f"http://127.0.0.1:{self.server_port}/v1/traces"
 
 
 def _func_with_exception():
@@ -86,3 +186,107 @@ def _test_trace_error_status_from_instance_imp(with_exception: bool):
         agentops.end_trace = old_end_trace
         tracer.teardown_worker(0)
         tracer.teardown()
+
+
+async def _test_agentops_trace_without_store_imp():
+    tracer = AgentOpsTracer()
+    tracer.init()
+    tracer.init_worker(0)
+
+    try:
+        # Using AgentOpsTracer to trace a function without providing a store, rollout_id, or attempt_id.
+        tracer.trace_run(_func_without_exception)
+        spans = tracer.get_last_trace()
+        assert len(spans) > 0
+    finally:
+        tracer.teardown_worker(0)
+        tracer.teardown()
+
+
+async def _test_agentops_trace_with_store_disable_imp():
+    tracer = AgentOpsTracer()
+    tracer.init()
+    tracer.init_worker(0)
+
+    try:
+        # Using AgentOpsTracer to trace a function with providing a store which disabled native otlp exporter, rollout_id, and attempt_id.
+        store = MockLightningStore()
+        async with tracer.trace_context(
+            name="agentops_test", store=store, rollout_id="test_rollout_id", attempt_id="test_attempt_id"
+        ):
+            _func_without_exception()
+        spans = tracer.get_last_trace()
+        assert len(spans) > 0
+    finally:
+        tracer.teardown_worker(0)
+        tracer.teardown()
+
+
+async def _test_agentops_trace_with_store_enable_imp():
+    mock_service = MockOTLPService()
+    port = mock_service.start_service()
+
+    tracer = AgentOpsTracer()
+    tracer.init()
+    tracer.init_worker(0)
+
+    try:
+        # Using AgentOpsTracer to trace a function with providing a store which disabled native otlp exporter, rollout_id, and attempt_id.
+        store = MockLightningStore(port)
+        async with tracer.trace_context(
+            name="agentops_test", store=store, rollout_id="test_rollout_id", attempt_id="test_attempt_id"
+        ):
+            _func_without_exception()
+        spans = tracer.get_last_trace()
+        assert len(spans) > 0
+    finally:
+        tracer.teardown_worker(0)
+        tracer.teardown()
+
+        mock_service.stop_service()
+
+
+def agentops_trace_paths() -> Iterator[Callable[[], Any]]:
+    yield from [
+        _test_agentops_trace_without_store_imp,
+        _test_agentops_trace_with_store_disable_imp,
+        _test_agentops_trace_with_store_enable_imp,
+    ]
+
+
+@pytest.mark.parametrize("func_name", [f.__name__ for f in agentops_trace_paths()], ids=str)
+def test_agentops_trace_with_store_or_not(func_name: str):
+    """
+    The purpose of this test is to verify whether the following two scenarios both work correctly:
+
+    1. Using AgentOpsTracer to trace a function without providing a store, rollout_id, or attempt_id.
+    2. Using AgentOpsTracer to trace a function with providing a store which disabled native otlp exporter, rollout_id, and attempt_id.
+    3. Using AgentOpsTracer to trace a function with providing a store which enabled native otlp exporter, rollout_id, and attempt_id.
+    """
+
+    func = {f.__name__: f for f in agentops_trace_paths()}[func_name]
+
+    ctx = multiprocessing.get_context("spawn")
+    proc = ctx.Process(target=_run_async, args=(func,))
+    proc.start()
+    proc.join(30.0)  # On GPU server, the time is around 10 seconds.
+
+    if proc.is_alive():
+        proc.terminate()
+        proc.join(5)
+        if proc.is_alive():
+            proc.kill()
+
+        assert False, "Child process hung. Check test output for details."
+
+    assert proc.exitcode == 0, (
+        f"Child process for test_trace_error_status_from_instance failed with exit code {proc.exitcode}. "
+        "Check child traceback in test output."
+    )
+
+
+def _run_async(coro: Callable[[], Coroutine[Any, Any, Any]]) -> None:
+    """Small wrapper: run async function inside multiprocessing target."""
+    import asyncio
+
+    asyncio.run(coro())

tests/utils/test_otel.py

@@ -0,0 +1,533 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+from types import SimpleNamespace
+from typing import Any, Dict, List
+
+import opentelemetry.trace as trace_api
+import pytest
+from opentelemetry.sdk.trace import ReadableSpan, SynchronousMultiSpanProcessor
+from opentelemetry.trace import TraceFlags
+from pydantic import ValidationError
+
+from agentlightning.semconv import LightningSpanAttributes, LinkPydanticModel
+from agentlightning.types.tracer import Span
+from agentlightning.utils import otel
+from agentlightning.utils.otel import (
+    extract_links_from_attributes,
+    extract_tags_from_attributes,
+    filter_and_unflatten_attributes,
+    filter_attributes,
+    flatten_attributes,
+    get_tracer,
+    get_tracer_provider,
+    make_link_attributes,
+    make_tag_attributes,
+    query_linked_spans,
+    unflatten_attributes,
+)
+
+
+def _span_context(trace_id_hex: str, span_id_hex: str) -> trace_api.SpanContext:
+    return trace_api.SpanContext(
+        trace_id=int(trace_id_hex, 16),
+        span_id=int(span_id_hex, 16),
+        is_remote=False,
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        trace_state=trace_api.TraceState(),
+    )
+
+
+def test_flatten_simple_nested_dict_and_list() -> None:
+    data = {"a": {"b": 1, "c": [2, 3]}}
+    result = flatten_attributes(data)
+    assert result == {
+        "a.b": 1,
+        "a.c.0": 2,
+        "a.c.1": 3,
+    }
+
+
+def test_flatten_empty_dict() -> None:
+    data: Dict[str, Any] = {}
+    assert flatten_attributes(data) == {}
+
+
+def test_flatten_empty_list() -> None:
+    data: List[Any] = []
+    # No elements -> no keys
+    assert flatten_attributes(data) == {}
+
+
+def test_flatten_root_list_of_primitives() -> None:
+    data = [10, 20, 30]
+    result = flatten_attributes(data)
+    assert result == {
+        "0": 10,
+        "1": 20,
+        "2": 30,
+    }
+
+
+def test_flatten_nested_lists_and_dicts() -> None:
+    data: Dict[str, Any] = {
+        "users": [
+            {"name": "Alice", "tags": ["admin", "staff"]},
+            {"name": "Bob", "tags": []},
+        ]
+    }
+    result = flatten_attributes(data)
+    assert result == {
+        "users.0.name": "Alice",
+        "users.0.tags.0": "admin",
+        "users.0.tags.1": "staff",
+        "users.1.name": "Bob",
+        # Empty list yields no extra keys
+    }
+
+
+def test_flatten_mixed_types_and_none() -> None:
+    data = {
+        "a": True,
+        "b": None,
+        "c": 3.14,
+        "d": "hello",
+        "e": {"f": False},
+    }
+    result = flatten_attributes(data)
+    assert result == {
+        "a": True,
+        "b": None,
+        "c": 3.14,
+        "d": "hello",
+        "e.f": False,
+    }
+
+
+def test_flatten_non_string_key_raises_value_error() -> None:
+    data = {
+        "a": {
+            1: "bad",  # non-string key inside nested dict
+        }
+    }
+    with pytest.raises(ValueError) as excinfo:
+        flatten_attributes(data)
+
+    msg = str(excinfo.value)
+    assert "Only string keys are supported in dictionaries" in msg
+    # Ensure the offending key is mentioned
+    assert "'1'" in msg
+    assert "type <class 'int'>" in msg
+
+
+def test_flatten_root_primitive_is_allowed() -> None:
+    # Even though the type hint says Dict/List, function behavior supports primitives.
+    data = 42
+    result = flatten_attributes(data)  # type: ignore[arg-type]
+    assert result == {"": 42}
+
+
+def test_unflatten_simple_nested_dict() -> None:
+    flat = {
+        "a.b": 1,
+        "a.c": 2,
+    }
+    result = unflatten_attributes(flat)
+    assert result == {"a": {"b": 1, "c": 2}}
+
+
+def test_unflatten_consecutive_numeric_keys_to_list() -> None:
+    flat = {
+        "a.0": "x",
+        "a.1": "y",
+        "a.2": "z",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": ["x", "y", "z"],
+    }
+
+
+def test_unflatten_non_consecutive_numeric_keys_stays_dict() -> None:
+    flat = {
+        "a.0": "first",
+        "a.2": "third",
+    }
+    result = unflatten_attributes(flat)
+    # Keys are numeric but not consecutive -> remains dict
+    assert result == {
+        "a": {
+            "0": "first",
+            "2": "third",
+        }
+    }
+
+
+def test_unflatten_mixed_numeric_and_non_numeric_keys_stays_dict() -> None:
+    flat = {
+        "a.0": "zero",
+        "a.foo": "bar",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": {
+            "0": "zero",
+            "foo": "bar",
+        }
+    }
+
+
+def test_unflatten_root_list_from_numeric_keys() -> None:
+    flat = {
+        "0": "a",
+        "1": "b",
+        "2": "c",
+    }
+    result = unflatten_attributes(flat)
+    # Root dict with all numeric keys 0..n-1 becomes list
+    assert result == ["a", "b", "c"]
+
+
+def test_unflatten_empty_flat_dict_returns_empty_dict() -> None:
+    flat: Dict[str, Any] = {}
+    result = unflatten_attributes(flat)
+    assert result == {}
+
+
+def test_unflatten_nested_lists_and_dicts() -> None:
+    flat = {
+        "users.0.name": "Alice",
+        "users.0.tags.0": "admin",
+        "users.0.tags.1": "staff",
+        "users.1.name": "Bob",
+        "users.1.tags.0": "guest",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "users": [
+            {"name": "Alice", "tags": ["admin", "staff"]},
+            {"name": "Bob", "tags": ["guest"]},
+        ]
+    }
+
+
+def test_unflatten_list_of_lists() -> None:
+    flat = {
+        "a.0.0": 1,
+        "a.0.1": 2,
+        "a.1.0": 3,
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": [
+            [1, 2],
+            [3],
+        ]
+    }
+
+
+def test_unflatten_conflicting_primitive_and_nested_path_prefers_nested() -> None:
+    # "a" is first set to a primitive, then to a nested dict via "a.b"
+    flat = {
+        "a": 1,
+        "a.b": 2,
+    }
+    result = unflatten_attributes(flat)
+    # Primitive is overwritten by nested dict structure
+    assert result == {"a": {"b": 2}}
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        {"a": {"b": 1, "c": [2, 3]}},
+        {"x": [1, 2, {"y": 3}]},
+        {"root": [{"k": "v"}, {"k": "w"}]},
+        [{"name": "Alice"}, {"name": "Bob", "scores": [10, 20]}],
+    ],
+)
+def test_round_trip_flatten_then_unflatten_preserves_structure(value: Dict[str, Any] | List[Any]) -> None:
+    flat = flatten_attributes(value)  # type: ignore[arg-type]
+    reconstructed = unflatten_attributes(flat)
+    assert reconstructed == value
+
+
+@pytest.mark.parametrize(
+    "flat",
+    [
+        {"a.b": 1, "a.c": 2},
+        {"0": "x", "1": "y"},
+        {
+            "users.0.name": "Alice",
+            "users.1.name": "Bob",
+        },
+    ],
+)
+def test_round_trip_unflatten_then_flatten_preserves_flat_structure(flat: Dict[str, Any]) -> None:
+    nested = unflatten_attributes(flat)
+    re_flat = flatten_attributes(nested)
+    # Order of items in dict shouldn't matter
+    assert re_flat == flat
+
+
+def test_round_trip_with_empty_list_information_loss_is_expected() -> None:
+    """This documents the corner case: empty list flattens to {},
+    which unflattens back to {} (empty dict), losing the distinction.
+    """
+    data: List[Any] = []
+    flat = flatten_attributes(data)
+    assert flat == {}
+    reconstructed = unflatten_attributes(flat)
+    assert reconstructed == {}
+    assert reconstructed != data  # explicit documentation of the behavior
+
+
+def test_make_and_extract_link_attributes_round_trip() -> None:
+    flattened = make_link_attributes(
+        {
+            "gen_ai.response.id": "response-123",
+            "span_id": "abcd1234abcd1234",
+        }
+    )
+    assert flattened == {
+        f"{LightningSpanAttributes.LINK.value}.0.key_match": "gen_ai.response.id",
+        f"{LightningSpanAttributes.LINK.value}.0.value_match": "response-123",
+        f"{LightningSpanAttributes.LINK.value}.1.key_match": "span_id",
+        f"{LightningSpanAttributes.LINK.value}.1.value_match": "abcd1234abcd1234",
+    }
+
+    extracted = extract_links_from_attributes(flattened)
+    assert [link.model_dump() for link in extracted] == [
+        {"key_match": "gen_ai.response.id", "value_match": "response-123"},
+        {"key_match": "span_id", "value_match": "abcd1234abcd1234"},
+    ]
+
+
+def test_make_link_attributes_rejects_non_string_values() -> None:
+    with pytest.raises(ValueError) as excinfo:
+        make_link_attributes({"span_id": 123})  # type: ignore
+
+    assert "Link value must be a string" in str(excinfo.value)
+
+
+def test_make_tag_attributes_and_extract_round_trip() -> None:
+    flattened = make_tag_attributes(["fast", "reliable"])
+    assert flattened == {
+        f"{LightningSpanAttributes.TAG.value}.0": "fast",
+        f"{LightningSpanAttributes.TAG.value}.1": "reliable",
+    }
+
+    assert extract_tags_from_attributes(flattened) == ["fast", "reliable"]
+
+
+def test_extract_tags_from_attributes_rejects_non_strings() -> None:
+    attributes = {
+        f"{LightningSpanAttributes.TAG.value}.0": 1,
+    }
+
+    with pytest.raises(ValidationError):
+        extract_tags_from_attributes(attributes)
+
+
+def test_filter_attributes_keeps_exact_matches_and_children() -> None:
+    attributes = {
+        "agentlightning.link": "root",
+        "agentlightning.link.0.key_match": "trace_id",
+        "agentlightning.other": "discard",
+        "agentlightning.link_extra": "different_prefix",
+    }
+
+    filtered = filter_attributes(attributes, LightningSpanAttributes.LINK.value)
+    assert filtered == {
+        "agentlightning.link": "root",
+        "agentlightning.link.0.key_match": "trace_id",
+    }
+
+
+def test_filter_and_unflatten_attributes_strips_prefix_and_rebuilds_nested_structure() -> None:
+    attributes = {
+        f"{LightningSpanAttributes.LINK.value}.0.key_match": "trace_id",
+        f"{LightningSpanAttributes.LINK.value}.0.value_match": "aaa",
+        f"{LightningSpanAttributes.LINK.value}.1.key_match": "span_id",
+        f"{LightningSpanAttributes.LINK.value}.1.value_match": "bbb",
+    }
+
+    result = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+    assert result == [
+        {"key_match": "trace_id", "value_match": "aaa"},
+        {"key_match": "span_id", "value_match": "bbb"},
+    ]
+
+
+def test_filter_and_unflatten_attributes_rejects_exact_prefix_key() -> None:
+    attributes = {LightningSpanAttributes.LINK.value: "invalid"}
+
+    with pytest.raises(ValueError):
+        filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+
+
+def test_query_linked_spans_matches_trace_id_on_readable_span() -> None:
+    readable_span = ReadableSpan(
+        name="upstream",
+        context=_span_context("a" * 32, "b" * 16),
+        attributes={},
+    )
+
+    assert readable_span.context is not None
+    links = [
+        LinkPydanticModel(key_match="trace_id", value_match=trace_api.format_trace_id(readable_span.context.trace_id)),
+    ]
+
+    matches = query_linked_spans([readable_span], links)
+    assert matches == [readable_span]
+
+
+def test_query_linked_spans_matches_custom_span_attributes() -> None:
+    custom_span = Span.from_attributes(
+        attributes={"gen_ai.response.id": "response-123", "custom": "needle"},
+        trace_id="c" * 32,
+        span_id="d" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="gen_ai.response.id", value_match="response-123"),
+        LinkPydanticModel(key_match="custom", value_match="needle"),
+    ]
+
+    matches = query_linked_spans([custom_span], links)
+    assert matches == [custom_span]
+
+
+def test_query_linked_spans_excludes_span_with_mismatched_span_id() -> None:
+    span = Span.from_attributes(
+        attributes={"marker": "x"},
+        trace_id="e" * 32,
+        span_id="f" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="span_id", value_match="deadbeefdeadbeef"),
+        LinkPydanticModel(key_match="marker", value_match="x"),
+    ]
+
+    assert query_linked_spans([span], links) == []
+
+
+def test_query_linked_spans_requires_all_links_to_match() -> None:
+    span = Span.from_attributes(
+        attributes={"marker": "x", "other": "y"},
+        trace_id="1" * 32,
+        span_id="2" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="marker", value_match="x"),
+        LinkPydanticModel(key_match="other", value_match="z"),
+    ]
+
+    assert query_linked_spans([span], links) == []
+
+
+def test_query_linked_spans_handles_readable_span_without_context() -> None:
+    readable_span = ReadableSpan(name="orphan", context=None, attributes={"marker": "x"})
+
+    links = [LinkPydanticModel(key_match="marker", value_match="x")]
+
+    matches = query_linked_spans([readable_span], links)
+    assert matches == [readable_span]
+
+
+def test_get_tracer_provider_raises_when_tracer_uninitialized(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", None, raising=False)
+
+    with pytest.raises(RuntimeError):
+        get_tracer_provider(inspect=False)
+
+
+def test_get_tracer_provider_logs_when_provider_not_sdk(
+    caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    sentinel_provider = object()
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", sentinel_provider, raising=False)
+    monkeypatch.setattr(otel, "otel_get_tracer_provider", lambda: sentinel_provider)
+
+    caplog.set_level(logging.ERROR, logger=otel.logger.name)
+
+    returned = get_tracer_provider(inspect=False)
+
+    assert returned is sentinel_provider
+    assert any("Tracer provider is expected" in rec.getMessage() for rec in caplog.records)
+
+
+def test_get_tracer_delegates_to_active_span_processor(monkeypatch: pytest.MonkeyPatch) -> None:
+    class DummyProvider:
+        def __init__(self) -> None:
+            self.calls: List[str] = []
+
+        def get_tracer(self, name: str) -> str:
+            self.calls.append(name)
+            return f"tracer:{name}"
+
+    provider = DummyProvider()
+
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", object(), raising=False)
+    monkeypatch.setattr(otel, "get_tracer_provider", lambda inspect=True: provider)
+
+    tracer = get_tracer()
+
+    assert tracer == "tracer:agentlightning"
+    assert provider.calls == ["agentlightning"]
+
+
+def test_get_tracer_without_active_span_processor_builds_isolated_tracer(monkeypatch: pytest.MonkeyPatch) -> None:
+    provider = SimpleNamespace(
+        sampler="sampler",
+        resource="resource",
+        id_generator="id_gen",
+    )
+
+    created: Dict[str, Any] = {}
+
+    class DummyTracer:
+        def __init__(
+            self,
+            sampler: Any,
+            resource: Any,
+            span_processor: SynchronousMultiSpanProcessor,
+            id_generator: Any,
+            instrumentation_info: Any,
+            span_limits: Any,
+            instrumentation_scope: Any,
+        ) -> None:
+            created["args"] = (
+                sampler,
+                resource,
+                span_processor,
+                id_generator,
+                instrumentation_info,
+                span_limits,
+                instrumentation_scope,
+            )
+
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", object(), raising=False)
+    monkeypatch.setattr(otel, "get_tracer_provider", lambda inspect=True: provider)
+    monkeypatch.setattr(otel, "Tracer", DummyTracer)
+
+    tracer = get_tracer(use_active_span_processor=False)
+
+    assert isinstance(created.get("args"), tuple)
+    assert created["args"][0] == "sampler"
+    assert created["args"][1] == "resource"
+    assert isinstance(created["args"][2], SynchronousMultiSpanProcessor)
+    assert created["args"][3] == "id_gen"
+    assert created["args"][4].name == "agentlightning"
+    assert tracer is not None
+
+
+def test_get_tracer_raises_when_provider_not_initialized(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", None, raising=False)
+
+    with pytest.raises(RuntimeError):
+        get_tracer()

tests/utils/test_otlp.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 import gzip
 from types import SimpleNamespace
-from typing import Any, Dict, Iterable, List, Optional, cast
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, cast
 
 import pytest
 from fastapi import Request
@@ -23,8 +23,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 from starlette.types import Message, Scope
 
-from agentlightning.store import LightningStore
-from agentlightning.types.tracer import SpanNames
+from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.utils import otlp
 
 BASE_TIME_NANOS = 1_700_000_000_000_000_000
@@ -34,16 +33,18 @@ EXTRA_EVENT_TIME_OFFSET = 4_000_000_000
 EXTRA_EVENT_TIME_SECONDS = (BASE_TIME_NANOS + EXTRA_EVENT_TIME_OFFSET) / 1_000_000_000
 
 
-class _StubStore(LightningStore):
+class _StubStore:
     def __init__(self) -> None:
         self.sequence_calls: List[tuple[str, str]] = []
         self.next_value = 1
 
-    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        self.sequence_calls.append((rollout_id, attempt_id))
-        value = self.next_value
-        self.next_value += 1
-        return value
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        self.sequence_calls.extend(rollout_attempt_ids)
+        allocations: List[int] = []
+        for _ in rollout_attempt_ids:
+            allocations.append(self.next_value)
+            self.next_value += 1
+        return allocations
 
 
 def _make_request(
@@ -109,9 +110,13 @@ def _add_attribute(attrs: Iterable[KeyValue], key: str, value: object) -> None:
 def _build_span_request() -> ExportTraceServiceRequest:
     request = ExportTraceServiceRequest()
     resource_spans = request.resource_spans.add()
-    _add_attribute(resource_spans.resource.attributes, SpanNames.ROLLOUT_ID, "resource-rollout")
-    _add_attribute(resource_spans.resource.attributes, SpanNames.ATTEMPT_ID, "resource-attempt")
-    _add_attribute(resource_spans.resource.attributes, SpanNames.SPAN_SEQUENCE_ID, "5")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "resource-rollout")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "resource-attempt")
+    _add_attribute(
+        resource_spans.resource.attributes,
+        LightningResourceAttributes.SPAN_SEQUENCE_ID.value,
+        "5",
+    )
     resource_spans.schema_url = "https://example/schema"
 
     scope_spans = resource_spans.scope_spans.add()
@@ -126,9 +131,9 @@ def _build_span_request() -> ExportTraceServiceRequest:
     span.status.message = "boom"
 
     _add_attribute(span.attributes, "foo", "bar")
-    _add_attribute(span.attributes, SpanNames.ROLLOUT_ID, "span-rollout")
-    _add_attribute(span.attributes, SpanNames.ATTEMPT_ID, "span-attempt")
-    _add_attribute(span.attributes, SpanNames.SPAN_SEQUENCE_ID, "7")
+    _add_attribute(span.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "span-rollout")
+    _add_attribute(span.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "span-attempt")
+    _add_attribute(span.attributes, LightningResourceAttributes.SPAN_SEQUENCE_ID.value, "7")
 
     event = span.events.add()
     event.name = "event"
@@ -230,7 +235,7 @@ async def test_spans_from_proto_prefers_span_level_metadata() -> None:
     store = _StubStore()
     request = _build_span_request()
 
-    spans = await otlp.spans_from_proto(request, store)
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
 
     assert len(spans) == 1
     span = spans[0]
@@ -241,7 +246,7 @@ async def test_spans_from_proto_prefers_span_level_metadata() -> None:
     assert span.events[0].timestamp == pytest.approx(EVENT_TIME_SECONDS)  # type: ignore
     assert span.links[0].context.trace_id == "0404" * 8
     assert span.links[0].attributes == {"link-attr": True}
-    assert span.resource.attributes[SpanNames.ROLLOUT_ID] == "resource-rollout"
+    assert span.resource.attributes[LightningResourceAttributes.ROLLOUT_ID.value] == "resource-rollout"
     assert span.resource.schema_url == "https://example/schema"
     assert not store.sequence_calls
 
@@ -251,8 +256,8 @@ async def test_spans_from_proto_requests_sequence_ids_when_missing() -> None:
     store = _StubStore()
     request = ExportTraceServiceRequest()
     resource_spans = request.resource_spans.add()
-    _add_attribute(resource_spans.resource.attributes, SpanNames.ROLLOUT_ID, "r1")
-    _add_attribute(resource_spans.resource.attributes, SpanNames.ATTEMPT_ID, "a1")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r1")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a1")
 
     scope_span = resource_spans.scope_spans.add()
     span = scope_span.spans.add()
@@ -260,20 +265,59 @@ async def test_spans_from_proto_requests_sequence_ids_when_missing() -> None:
     span.span_id = b""
     span.name = "needs-seq"
 
-    spans = await otlp.spans_from_proto(request, store)
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
 
     assert len(spans) == 1
     assert spans[0].sequence_id == 1
     assert store.sequence_calls == [("r1", "a1")]
 
 
+@pytest.mark.asyncio
+async def test_spans_from_proto_bulk_issues_for_mixed_rollouts() -> None:
+    store = _StubStore()
+    request = ExportTraceServiceRequest()
+
+    resource_first = request.resource_spans.add()
+    _add_attribute(resource_first.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r1")
+    _add_attribute(resource_first.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a-default")
+
+    scope_first = resource_first.scope_spans.add()
+    span_missing = scope_first.spans.add()
+    span_missing.trace_id = bytes.fromhex("11" * 16)
+    span_missing.span_id = bytes.fromhex("22" * 8)
+    span_missing.name = "missing-seq"
+    _add_attribute(span_missing.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a1")
+
+    span_negative = scope_first.spans.add()
+    span_negative.trace_id = bytes.fromhex("33" * 16)
+    span_negative.span_id = bytes.fromhex("44" * 8)
+    span_negative.name = "negative-seq"
+    _add_attribute(span_negative.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a2")
+    _add_attribute(span_negative.attributes, LightningResourceAttributes.SPAN_SEQUENCE_ID.value, "-5")
+
+    resource_second = request.resource_spans.add()
+    _add_attribute(resource_second.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r2")
+    _add_attribute(resource_second.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "b1")
+
+    scope_second = resource_second.scope_spans.add()
+    span_second = scope_second.spans.add()
+    span_second.trace_id = bytes.fromhex("55" * 16)
+    span_second.span_id = bytes.fromhex("66" * 8)
+    span_second.name = "other-rollout"
+
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
+    assert [span.name for span in spans] == ["missing-seq", "negative-seq", "other-rollout"]
+    assert [span.sequence_id for span in spans] == [1, 2, 3]
+    assert store.sequence_calls == [("r1", "a1"), ("r1", "a2"), ("r2", "b1")]
+
+
 @pytest.mark.asyncio
 async def test_spans_from_proto_skips_spans_without_ids() -> None:
     store = _StubStore()
     request = ExportTraceServiceRequest()
     request.resource_spans.add()  # missing rollout and attempt
 
-    spans = await otlp.spans_from_proto(request, store)
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
 
     assert spans == []
     assert store.sequence_calls == []
@@ -399,8 +443,8 @@ def test_lightning_store_otlp_exporter_overrides_resources(monkeypatch: pytest.M
     assert result == SpanExportResult.SUCCESS
     assert captured_spans
     attributes = captured_spans[0][0]._resource.attributes  # type: ignore[attr-defined]
-    assert attributes[SpanNames.ROLLOUT_ID] == "rollout"
-    assert attributes[SpanNames.ATTEMPT_ID] == "attempt"
+    assert attributes[LightningResourceAttributes.ROLLOUT_ID.value] == "rollout"
+    assert attributes[LightningResourceAttributes.ATTEMPT_ID.value] == "attempt"
 
     exporter.disable_store_otlp()
     assert exporter._rollout_id is None