fix trainer

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: ultmaster <8463288+ultmaster@users.noreply.github.com>

.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/badge-spider.yml

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

.github/workflows/examples-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

.github/workflows/examples-rag.yml

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

.github/workflows/examples-spider.yml

@@ -33,8 +33,7 @@ jobs:
     strategy:
       matrix:
         include:
-        - python-version: '3.10'
-          setup-script: 'legacy'
+        # legacy is omitted because langchain doesn't work with legacy vllm versions
         - python-version: '3.12'
           setup-script: 'stable'
         - python-version: '3.13'
@@ -58,13 +57,13 @@ jobs:
       - name: Sync dependencies (latest)
         run: |
           uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --group torch-gpu-stable
+            --group dev --group experiment --group agents --group langchain --group torch-gpu-stable
         if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (stable & legacy)
+      - name: Sync dependencies (stable)
         run: |
           uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+            --group dev --group experiment --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
       - name: Freeze dependencies
         run: |
           set -ex

.github/workflows/tests-full.yml

@@ -56,11 +56,15 @@ jobs:
         run: uv lock --upgrade
         if: matrix.setup-script == 'latest'
       - name: Sync dependencies (latest)
-        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-stable
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-stable
         if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+      - name: Sync dependencies (stable)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
+      # Don't install langchain for legacy dependency because it has conflicts with torch.
+      - name: Sync dependencies (legacy)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-legacy
+        if: matrix.setup-script == 'legacy'
       - name: Freeze dependencies
         run: |
           set -ex
@@ -178,11 +182,15 @@ jobs:
         run: uv lock --upgrade
         if: matrix.setup-script == 'latest'
       - name: Sync dependencies (latest)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-stable
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group torch-gpu-stable
         if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+      - name: Sync dependencies (stable)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
+      # Don't install langchain for legacy dependency because it has conflicts with torch.
+      - name: Sync dependencies (legacy)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-legacy
+        if: matrix.setup-script == 'legacy'
       - name: Freeze dependencies
         run: |
           set -ex
@@ -320,3 +328,24 @@ jobs:
             echo "Waiting for llm_proxy.py to finish..."
             sleep 5
           done
+
+      - name: MultiMetrics backend example
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_metrics.py --duration 8 --prom-port 9105 --prom-host 0.0.0.0 2>&1 | tee metrics.log &
+          pid=$!
+
+          for attempt in $(seq 1 20); do
+            if curl -sSf http://localhost:9105/metrics | grep -q minimal_requests_total; then
+              echo "Metrics endpoint responding"
+              wait $pid
+              cat metrics.log
+              exit 0
+            fi
+            sleep 1
+          done
+
+          echo "Metrics endpoint did not respond"
+          exit 1

.github/workflows/tests.yml

@@ -44,6 +44,7 @@ jobs:
             --group trl \
             --group tinker \
             --group agents \
+            --group langchain \
             --no-default-groups
         if: matrix.setup == 'slow'
       # This pre-commit skips JavaScript on purpose.
@@ -139,10 +140,10 @@ jobs:
         run: uv lock --upgrade
         if: matrix.setup-script == 'latest'
       - name: Sync dependencies (latest)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-stable
         if: matrix.setup-script == 'latest'
       - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-${{ matrix.setup-script }}
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-${{ matrix.setup-script }}
         if: matrix.setup-script != 'latest'
       - name: Freeze dependencies
         run: |

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

agentlightning/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

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

agentlightning/semconv.py

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

agentlightning/store/client_server.py

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

agentlightning/store/collection/__init__.py

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

agentlightning/store/collection/base.py

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

agentlightning/store/collection/memory.py

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

agentlightning/store/memory.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import asyncio
 import logging
 import sys
-import threading
 from collections.abc import Iterable
 from collections.abc import Mapping as MappingABC
 from typing import (
@@ -17,19 +16,22 @@ from typing import (
     Literal,
     Mapping,
     Optional,
+    Sequence,
     Set,
+    Tuple,
     TypeVar,
     Union,
     cast,
 )
 
+import aiologic
 from pydantic import BaseModel
 
-from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span
+from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
 
-from .base import UNSET, LightningStoreCapabilities, 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])
 
@@ -81,12 +83,20 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def __init__(
         self,
         *,
+        thread_safe: bool = False,
         eviction_memory_threshold: float | int | None = None,
         safe_memory_threshold: float | int | None = None,
         span_size_estimator: Callable[[Span], int] | None = None,
+        prometheus: bool = False,
     ):
-        super().__init__(collections=InMemoryLightningCollections())
+        super().__init__(
+            collections=InMemoryLightningCollections(
+                lock_type="thread" if thread_safe else "asyncio", prometheus=prometheus
+            ),
+            prometheus=prometheus,
+        )
 
+        self._thread_safe = thread_safe
         self._start_time_by_rollout: Dict[str, float] = {}
         self._span_bytes_by_rollout: Dict[str, int] = Counter()
         self._total_span_bytes: int = 0
@@ -120,7 +130,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         self._custom_span_size_estimator = span_size_estimator
 
         # Completion tracking for wait_for_rollouts (cross-loop safe)
-        self._completion_events: Dict[str, threading.Event] = {}
+        self._completion_events: Dict[str, aiologic.Event] = {}
 
         # Running rollouts cache, including preparing and running rollouts
         self._running_rollout_ids: Set[str] = set()
@@ -132,15 +142,26 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
-            thread_safe=False,
+            thread_safe=self._thread_safe,
             async_safe=True,
             zero_copy=False,
             otlp_traces=False,
         )
 
+    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:
+        async with self.collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["rollouts"]) as collections:
             rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
             if rollout and is_finished(rollout):
                 return rollout
@@ -168,47 +189,85 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
             # If event was set (not timeout), check if rollout is finished
             if result:
-                async with self.collections.atomic() as collections:
+                async with self.collections.atomic(
+                    mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
+                ) as collections:
                     rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
                     if rollout and is_finished(rollout):
                         return rollout
 
         return None
 
-    async def on_rollout_update(self, rollout: Rollout) -> None:
+    @tracked("add_resources_inmemory")
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        ret = await super().add_resources(resources)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
+            self._latest_resources_id = ret.resources_id
+        return ret
+
+    @tracked("update_resources_inmemory")
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        ret = await super().update_resources(resources_id, resources)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
+            self._latest_resources_id = ret.resources_id
+        return ret
+
+    @tracked("_post_update_rollout_inmemory")
+    async def _post_update_rollout(
+        self, rollouts: Sequence[Tuple[Rollout, Sequence[str]]], skip_enqueue: bool = False
+    ) -> None:
         """Update the running rollout ids set when the rollout updates."""
-        if is_running(rollout):
-            self._running_rollout_ids.add(rollout.rollout_id)
-        else:
-            self._running_rollout_ids.discard(rollout.rollout_id)
+        await super()._post_update_rollout(rollouts, skip_enqueue=skip_enqueue)
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["rollouts"]):
+            for rollout, _ in rollouts:
+                if is_running(rollout):
+                    self._running_rollout_ids.add(rollout.rollout_id)
+                else:
+                    self._running_rollout_ids.discard(rollout.rollout_id)
 
-        if is_finished(rollout):
-            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-            self._completion_events[rollout.rollout_id].set()
-        else:
-            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-        # Rollout status can never transition from finished to running (unlike attempt)
-        # so we don't need to clear the completion event even in case of retrying.
+                if is_finished(rollout):
+                    self._completion_events.setdefault(rollout.rollout_id, aiologic.Event())
+                    self._completion_events[rollout.rollout_id].set()
+                else:
+                    self._completion_events.setdefault(rollout.rollout_id, aiologic.Event())
+                # Rollout status can never transition from finished to running (unlike attempt)
+                # so we don't need to clear the completion event even in case of retrying.
 
-        if rollout.rollout_id not in self._start_time_by_rollout:
-            self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
+                if rollout.rollout_id not in self._start_time_by_rollout:
+                    self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
 
-    async def get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
-        """Accelerated version of `get_running_rollouts` for in-memory store. Used for healthcheck."""
-        rollouts = await collections.rollouts.query(filter={"rollout_id": {"within": list(self._running_rollout_ids)}})
-        running_rollouts: List[AttemptedRollout] = []
-        for rollout in rollouts.items:
-            latest_attempt = await collections.attempts.get(
-                filter={"rollout_id": {"exact": rollout.rollout_id}},
-                sort={"name": "sequence_id", "order": "desc"},
-            )
-            if not latest_attempt:
-                # The rollout is running but has no attempts, this should not happen
-                logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
-                continue
-            running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
+    @tracked("_unlocked_query_rollouts_by_rollout_ids")
+    async def _unlocked_query_rollouts_by_rollout_ids(
+        self, collections: InMemoryLightningCollections, rollout_ids: Sequence[str]
+    ) -> List[Rollout]:
+        """Always use exact. This is faster than within filter for in-memory store."""
+        if len(rollout_ids) == 0:
+            return []
+
+        rollouts = [await collections.rollouts.get({"rollout_id": {"exact": rollout_id}}) for rollout_id in rollout_ids]
+        return [rollout for rollout in rollouts if rollout is not None]
+
+    @tracked("_unlocked_get_running_rollouts")
+    async def _unlocked_get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
+        """Accelerated version of `_unlocked_get_running_rollouts` for in-memory store. Used for healthcheck."""
+        async with self.collections.atomic(
+            mode="r", snapshot=self._read_snapshot, labels=["rollouts", "attempts"]
+        ) as collections:
+            rollouts = await self._unlocked_query_rollouts_by_rollout_ids(collections, list(self._running_rollout_ids))
+            running_rollouts: List[AttemptedRollout] = []
+            for rollout in rollouts:
+                latest_attempt = await collections.attempts.get(
+                    filter={"rollout_id": {"exact": rollout.rollout_id}},
+                    sort={"name": "sequence_id", "order": "desc"},
+                )
+                if not latest_attempt:
+                    # The rollout is running but has no attempts, this should not happen
+                    logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
+                    continue
+                running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
         return running_rollouts
 
+    @tracked("query_spans_inmemory")  # Since this method calls super, we need to track it separately
     async def query_spans(
         self,
         rollout_id: str,
@@ -219,23 +278,28 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
         return await super().query_spans(rollout_id, attempt_id, **kwargs)
 
-    async def _add_span_unlocked(self, collections: InMemoryLightningCollections, span: Span) -> Span:
+    @tracked("_post_add_spans")
+    async def _post_add_spans(self, spans: Sequence[Span], rollout_id: str, attempt_id: str) -> None:
         """In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""
 
-        await super()._add_span_unlocked(collections, span)
-        self._account_span_size(span)
-        await self._maybe_evict_spans(collections)
+        await super()._post_add_spans(spans, rollout_id, attempt_id)
+        async with self.collections.atomic(
+            mode="rw", snapshot=self._read_snapshot, labels=["rollouts", "spans"]
+        ) as collections:
+            for span in spans:
+                await self._account_span_size(span)
+            await self._maybe_evict_spans(collections)
 
-        return span
-
-    async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
+    @tracked("_get_latest_resources_inmemory")
+    async def _get_latest_resources(self) -> Optional[ResourcesUpdate]:
         if isinstance(self._latest_resources_id, Unset):
-            latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
-            if latest_resources:
-                self._latest_resources_id = latest_resources.resources_id
-            else:
-                self._latest_resources_id = None
-        return self._latest_resources_id
+            return await super()._get_latest_resources()
+        if self._latest_resources_id is not None:
+            async with self.collections.atomic(
+                mode="r", snapshot=self._read_snapshot, labels=["resources"]
+            ) as collections:
+                return await collections.resources.get(filter={"resources_id": {"exact": self._latest_resources_id}})
+        return None
 
     @staticmethod
     def _resolve_memory_threshold(
@@ -267,7 +331,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 +342,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 +365,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,67 @@ 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:
+            async with self.collections.atomic(
+                mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
+            ) as collections:
+                # Query the rollouts that are not finished in a single query
+                rollouts = await collections.rollouts.query(
+                    filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
+                )
+            for rollout in rollouts.items:
+                if is_finished(rollout):
+                    finished_rollouts[rollout.rollout_id] = rollout
+                    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("_unlocked_many_rollouts_to_attempted_rollouts")
+    async def _unlocked_many_rollouts_to_attempted_rollouts(
+        self, collections: MongoLightningCollections, rollouts: Sequence[Rollout]
+    ) -> List[Union[Rollout, AttemptedRollout]]:
+        """Query the latest attempts for the rollouts, and attach them to the rollout objects."""
+        async with collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["attempts"]) as collections:
+            attempts = await collections.attempts.query(
+                filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
+                sort={"name": "sequence_id", "order": "desc"},
+            )
+        latest_attempts: Dict[str, Attempt] = {}
+        for attempt in attempts:
+            if attempt.rollout_id not in latest_attempts:
+                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
 
@@ -11,6 +11,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -22,7 +23,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 +48,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,
@@ -54,9 +60,17 @@ class LightningStoreThreaded(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         with self._lock:
-            return await self.store.start_rollout(input, mode, resources_id, config, metadata)
+            return await self.store.start_rollout(
+                input,
+                mode,
+                resources_id,
+                config,
+                metadata,
+                worker_id,
+            )
 
     async def enqueue_rollout(
         self,
@@ -69,13 +83,26 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
 
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        with self._lock:
+            return await self.store.enqueue_many_rollouts(rollouts)
+
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         with self._lock:
             return await self.store.dequeue_rollout(worker_id=worker_id)
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
         with self._lock:
-            return await self.store.start_attempt(rollout_id)
+            return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id, worker_id)
 
     async def query_rollouts(
         self,
@@ -167,7 +194,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 +208,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 +220,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

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

agentlightning/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/trainer/trainer.py

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

agentlightning/types/core.py

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

agentlightning/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/metrics.py

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

agentlightning/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

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

dashboard/.gitignore

@@ -130,3 +130,6 @@ dist
 .pnp.*
 
 .DS_Store
+
+# Storybook build output
+storybook-static

dashboard/src/components/TracesTable.story.tsx

@@ -365,3 +365,209 @@ export const NestedSpans: Story = {
     <TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.traceId === 'trace-nested456')} />
   ),
 };
+
+// Test data with sequence IDs that would sort incorrectly if treated as strings
+const sequenceSortTestSpans: Span[] = [
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 2,
+    traceId: 'trace-seq-002',
+    spanId: 'span-seq-002',
+    parentId: null,
+    name: 'task_sequence_2',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 200,
+    endTime: now - 190,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 10,
+    traceId: 'trace-seq-010',
+    spanId: 'span-seq-010',
+    parentId: null,
+    name: 'task_sequence_10',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 180,
+    endTime: now - 170,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 3,
+    traceId: 'trace-seq-003',
+    spanId: 'span-seq-003',
+    parentId: null,
+    name: 'task_sequence_3',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 160,
+    endTime: now - 150,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 11,
+    traceId: 'trace-seq-011',
+    spanId: 'span-seq-011',
+    parentId: null,
+    name: 'task_sequence_11',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 140,
+    endTime: now - 130,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 9,
+    traceId: 'trace-seq-009',
+    spanId: 'span-seq-009',
+    parentId: null,
+    name: 'task_sequence_9',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 120,
+    endTime: now - 110,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 12,
+    traceId: 'trace-seq-012',
+    spanId: 'span-seq-012',
+    parentId: null,
+    name: 'task_sequence_12',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 100,
+    endTime: now - 90,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 4,
+    traceId: 'trace-seq-004',
+    spanId: 'span-seq-004',
+    parentId: null,
+    name: 'task_sequence_4',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 80,
+    endTime: now - 70,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 6,
+    traceId: 'trace-seq-006',
+    spanId: 'span-seq-006',
+    parentId: null,
+    name: 'task_sequence_6',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 60,
+    endTime: now - 50,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 7,
+    traceId: 'trace-seq-007',
+    spanId: 'span-seq-007',
+    parentId: null,
+    name: 'task_sequence_7',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 40,
+    endTime: now - 30,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 13,
+    traceId: 'trace-seq-013',
+    spanId: 'span-seq-013',
+    parentId: null,
+    name: 'task_sequence_13',
+    status: { status_code: 'OK', description: null },
+    attributes: {},
+    startTime: now - 20,
+    endTime: now - 10,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+  {
+    rolloutId: 'ro-seq-test',
+    attemptId: 'at-seq-test',
+    sequenceId: 14,
+    traceId: 'trace-seq-014',
+    spanId: 'span-seq-014',
+    parentId: null,
+    name: 'task_sequence_14',
+    status: { status_code: 'UNSET', description: null },
+    attributes: {},
+    startTime: now - 5,
+    endTime: now,
+    events: [],
+    links: [],
+    context: {},
+    parent: null,
+    resource: {},
+  },
+];
+
+export const SequenceIdSortTest: Story = {
+  render: () => <TracesTableStoryWrapper maxWidth={1200} spans={sequenceSortTestSpans} />,
+};

dashboard/src/features/traces/selectors.ts

@@ -23,6 +23,7 @@ export const selectTracesViewMode = (state: RootState) => selectTracesState(stat
 
 const TRACES_SORT_FIELD_MAP: Record<string, string> = {
   name: 'name',
+  sequenceId: 'sequence_id',
   traceId: 'trace_id',
   spanId: 'span_id',
   parentId: 'parent_id',

docker/compose.grafana.yml

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

docker/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/community/contributing.md

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

docs/deep-dive/store.md

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

docs/how-to/examples-catalog.md

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

docs/reference/agent.md

@@ -22,6 +22,10 @@
 
 ## Emitter
 
+::: agentlightning.operation
+
+::: agentlightning.emit_annotation
+
 ::: agentlightning.emit_reward
 
 ::: agentlightning.emit_message
@@ -30,7 +34,11 @@
 
 ::: agentlightning.emit_exception
 
-## Reward Helpers
+## Emitter Helpers
+
+::: agentlightning.get_message_value
+
+::: agentlightning.get_object_value
 
 ::: agentlightning.find_final_reward
 
@@ -38,8 +46,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,18 @@
 
 ::: agentlightning.algorithm.decorator.FunctionalAlgorithm
 
+## LitAgent
+
 ::: agentlightning.litagent.decorator.FunctionalLitAgent
 
 ::: agentlightning.litagent.decorator.llm_rollout
 
 ::: agentlightning.litagent.decorator.prompt_rollout
 
+::: agentlightning.emitter.annotation.OperationContext
+
+## LLM Proxy
+
 ::: agentlightning.llm_proxy.ModelConfig
 
 ::: agentlightning.llm_proxy.LightningSpanExporter
@@ -34,24 +42,58 @@
 
 ::: agentlightning.llm_proxy.RolloutAttemptMiddleware
 
+## Store
+
 ::: agentlightning.store.base.UNSET
 
-::: agentlightning.store.utils.propagate_status
+::: agentlightning.store.utils.rollout_status_from_attempt
+
+::: agentlightning.store.utils.scan_unhealthy_rollouts
+
+## Tracing and OpenTelemetry
 
 ::: 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/store.md

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

docs/reference/types.md

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

examples/.gitignore

@@ -4,6 +4,7 @@ outputs/
 checkpoints/
 calc-x-data.zip
 spider-data.zip
+claude_code/logs/
 agentops.log
 unsloth/models/
 unsloth/unsloth_compiled_cache/
@@ -11,3 +12,6 @@ unsloth/unsloth_training_checkpoints/
 apo/pomltrace/
 tinker/logs/
 tinker/crewai_*.html
+rag/dataset_tiny.parquet
+rag/chunks_candidate_tiny.pkl
+rag/index_hnsw_faiss_n32e40_tiny.index

examples/README.md

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

examples/azure/README.md

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

examples/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.
@@ -9,6 +11,7 @@ Each module have been documented with its own CLI usage in the module-level docs
 | Component | Demonstrated In | Highlights |
 | --- | --- | --- |
 | LightningStore + OTLP ingestion | `write_traces.py` | Shows how `OtelTracer` and `AgentOpsTracer` open rollouts, emit spans, and optionally forward them to a remote store client. |
+| MultiMetrics backend | `write_metrics.py` | Emits counters/histograms through `ConsoleMetricsBackend` and `PrometheusMetricsBackend` simultaneously, exposing `/metrics` for scraping. |
 | LLM proxying | `llm_proxy.py` | Guards either OpenAI or a local vLLM deployment with `LLMProxy`, proving how requests are routed through `/rollout/<id>/attempt/<id>` namespaces and captured in the store. |
 | vLLM lifecycle | `vllm_server.py` | Minimal context manager that shells out to `vllm serve`, monitors readiness, and tears down the process safely. |
 

examples/minimal/write_metrics.py

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

examples/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/rag/README.md

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

examples/rag/metric_utils.py

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

examples/rag/rag_agent.py

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

examples/rag/train.sh

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

examples/rag/train_rag.py

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

examples/rag/wiki_retriever_mcp.py

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

examples/rag/wiki_retriever_mcp/wiki_retriever_install.sh

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

examples/tinker/README.md

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