minor fix

move out search r1
Move search-r1 example
2025-12-13 12:47:53 +08:00 · 2025-12-13 12:44:47 +08:00 · 2025-12-13 12:41:44 +08:00 · 2025-12-13 12:37:20 +08:00 · 2025-12-13 12:32:12 +08:00 · 2025-12-13 00:25:54 +08:00
114 changed files with 10702 additions and 5897 deletions
@@ -0,0 +1,29 @@
+name: Badge - ChartQA
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - ChartQA
+    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-chartqa.yml', label: 'chartqa', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -10,6 +10,8 @@ on:
      - Examples - Tinker
      - Examples - Azure
      - Examples - Claude Code
+      - Examples - RAG
+      - Examples - ChartQA
    types: [completed]

  workflow_dispatch:
@@ -37,5 +39,7 @@ jobs:
              { 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'] },
+              { workflow: 'examples-rag.yml', label: 'examples-rag.stable', variants: ['stable'] },
+              { workflow: 'examples-chartqa.yml', label: 'examples-chartqa.stable', variants: ['stable'] },
            ];
            await badgeAggregation({ github, context, core, dependencies });
@@ -7,6 +7,8 @@ on:
      - Examples - Spider
      - Examples - APO
      - Examples - Unsloth
+      - Examples - RAG
+      - Examples - Claude Code
      - GPU Test
    types: [completed]

@@ -32,6 +34,8 @@ jobs:
              { workflow: 'examples-spider.yml', label: 'spider.latest', variants: ['latest'] },
              { workflow: 'examples-apo.yml', label: 'apo.latest', variants: ['latest'] },
              { workflow: 'examples-unsloth.yml', label: 'unsloth.latest', variants: ['latest'] },
+              { workflow: 'examples-claude-code.yml', label: 'claude-code.latest', variants: ['latest'] },
+              { workflow: 'examples-rag.yml', label: 'rag.latest', variants: ['latest'] },
              { workflow: 'tests-full.yml', label: 'tests-full.latest', variants: ['latest'] },
            ];
            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - RAG
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - RAG
+    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-rag.yml', label: 'rag', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -6,9 +6,9 @@ on:

 jobs:
  benchmark:
-    name: Benchmark (${{ matrix.backend.id }}, ${{ matrix.scenario.display }})
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
-    timeout-minutes: 60
+    name: ${{ matrix.workload.kind }} (${{ matrix.backend.id }}, ${{ matrix.workload.display }})
+    runs-on: ${{ matrix.workload.runner }}
+    timeout-minutes: ${{ matrix.workload.timeout }}
    strategy:
      fail-fast: false
      matrix:
@@ -17,10 +17,15 @@ jobs:
            compose_file: compose.prometheus-memory-store.yml
          - id: mongo
            compose_file: compose.prometheus-mongo-store.yml
-        scenario:
-          - id: minimal-production
+        workload:
+          - id: scenario-minimal-scale
            display: Minimal production scale
+            kind: scenario
            store_workers: 4
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode batch
              --total-tasks 4096
@@ -28,9 +33,14 @@ jobs:
              --n-runners 32
              --max-rounds 6
              --sleep-seconds 0.5
-          - id: medium-production
+          - id: scenario-medium-scale
            display: Medium production scale
+            kind: scenario
            store_workers: 16
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode batch
              --total-tasks 10000
@@ -38,9 +48,29 @@ jobs:
              --n-runners 100
              --max-rounds 10
              --sleep-seconds 0.1
-          - id: large-batch
+          - id: scenario-midhigh-scale
+            display: Mid-high production scale
+            kind: scenario
+            store_workers: 24
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
+            args: >-
+              --mode batch
+              --total-tasks 20000
+              --batch-size 2048
+              --n-runners 256
+              --max-rounds 8
+              --sleep-seconds 0.1
+          - id: scenario-large-batch
            display: Large batch waves
+            kind: scenario
            store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode batch
              --total-tasks 100000
@@ -48,9 +78,14 @@ jobs:
              --n-runners 256
              --max-rounds 6
              --sleep-seconds 0.1
-          - id: long-queues
+          - id: scenario-long-queues
            display: Long rollout queues
+            kind: scenario
            store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode batch_partial
              --total-tasks 100000
@@ -59,9 +94,14 @@ jobs:
              --remaining-tasks 4096
              --max-rounds 4
              --sleep-seconds 0.1
-          - id: high-concurrency
+          - id: scenario-high-concurrency
            display: High-throughput concurrent requests
+            kind: scenario
            store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode single
              --total-tasks 100000
@@ -69,9 +109,14 @@ jobs:
              --n-runners 256
              --max-rounds 2
              --sleep-seconds 0.1
-          - id: heavy-traces
+          - id: scenario-heavy-traces
            display: Heavy rollouts with deep traces
+            kind: scenario
            store_workers: 64
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
            args: >-
              --mode batch_partial
              --total-tasks 10000
@@ -80,15 +125,63 @@ jobs:
              --n-runners 512
              --max-rounds 20
              --sleep-seconds 1.0
+
+          - id: micro-worker
+            display: Update worker
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: worker
+          - id: micro-dequeue-empty
+            display: Dequeue empty
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-empty
+          - id: micro-rollout
+            display: Rollout + span
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: rollout
+          - id: micro-dequeue-update-attempt
+            display: Dequeue + update attempt
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-update-attempt
+          - id: micro-dequeue-only
+            display: Dequeue only
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-only
+          - id: micro-metrics
+            display: Multi-metric fan-out
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 15
+            cli: metrics
    env:
      STORE_URL: http://localhost:4747
      STORE_API_URL: http://localhost:4747/v1/agl
      PROM_URL: http://localhost:9090
-      SCENARIO_ID: ${{ matrix.scenario.id }}
+      WORKLOAD_KIND: ${{ matrix.workload.kind }}
+      WORKLOAD_ID: ${{ matrix.workload.id }}
      BACKEND_ID: ${{ matrix.backend.id }}
-      ARTIFACT_DIR: artifacts/${{ matrix.scenario.id }}-${{ matrix.backend.id }}
+      ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.workload.id, matrix.backend.id) }}
      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
-      AGL_STORE_N_WORKERS: ${{ matrix.scenario.store_workers }}
+      AGL_STORE_N_WORKERS: ${{ matrix.workload.store_workers }}
+      ANALYSIS_FILE: ${{ format('analysis-{0}.log', matrix.workload.id) }}
+      SUMMARY_FILE: ${{ format('summary-{0}.log', matrix.workload.id) }}
+      PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
+      ARTIFACT_NAME: ${{ format('{0}-{1}', matrix.workload.id, matrix.backend.id) }}
    steps:
      - uses: actions/checkout@v4

@@ -122,38 +215,65 @@ jobs:
          set -euo pipefail
          for attempt in {1..60}; do
            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+              sleep 1
+              curl -fsS "$STORE_API_URL/rollouts"  # Warm up the scraper
+              sleep 15  # Allow some time for the baseline metrics to be established
              exit 0
            fi
            sleep 1
          done
          echo "Store did not become ready in time" >&2
-          docker compose -f "$COMPOSE_FILE" logs app
+          # show logs for debugging
+          cd docker && docker compose -f "$COMPOSE_FILE" logs app
          exit 1

      - name: Prepare artifact directory
        run: mkdir -p "$ARTIFACT_DIR"

-      - name: Record benchmark start
+      - name: Record workload start
        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"

-      - name: Run ${{ matrix.scenario.display }} workload
+      - name: (Scenario) Run ${{ matrix.workload.display }} workload
+        if: ${{ matrix.workload.kind == 'scenario' }}
        run: |
          set -euo pipefail
          uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
            --store-url "$STORE_URL" \
-            ${{ matrix.scenario.args }}
+            ${{ matrix.workload.args }}

-      - name: Record benchmark end
+      - name: (Micro) Run ${{ matrix.workload.display }}
+        if: ${{ matrix.workload.kind == 'micro' }}
+        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_FILE" \
+            "${{ matrix.workload.cli }}" | tee "$ARTIFACT_DIR/${{ matrix.workload.id }}.txt"
+
+      - name: Record workload end
        if: ${{ always() }}
        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"

-      - name: Run benchmark analysis
+      - name: Show micro benchmark summary
+        if: ${{ always() && matrix.workload.kind == 'micro' }}
+        run: |
+          set -euo pipefail
+          summary_file="$ARTIFACT_DIR/$SUMMARY_FILE"
+          if [ -f "$summary_file" ]; then
+            echo "Micro benchmark summary ($WORKLOAD_ID/$BACKEND_ID):"
+            cat "$summary_file"
+          else
+            echo "Summary file not found: $summary_file"
+          fi
+
+      - name: Run workload 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"
+            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/$ANALYSIS_FILE"
            exit 1
          fi
          uv run --locked --no-sync python -m tests.benchmark.analysis \
@@ -161,7 +281,22 @@ jobs:
            --store-url "$STORE_API_URL" \
            --start "$BENCHMARK_START" \
            --end "$BENCHMARK_END" \
-            | tee "$ARTIFACT_DIR/analysis.txt"
+            | tee "$ARTIFACT_DIR/$ANALYSIS_FILE"
+
+      - name: Collect docker logs
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          cd docker
+          readarray -t services < <(docker compose -f "$COMPOSE_FILE" config --services)
+          if [ "${#services[@]}" -eq 0 ]; then
+            echo "No services defined in compose file."
+            exit 0
+          fi
+          for service in "${services[@]}"; do
+            docker compose -f "$COMPOSE_FILE" logs "$service" > "../$ARTIFACT_DIR/docker-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
+          done

      - name: Stop ${{ matrix.backend.id }} Prometheus stack
        if: ${{ always() }}
@@ -176,51 +311,61 @@ jobs:
          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
+            tar -C docker/data -czf "$ARTIFACT_DIR/${PROM_ARCHIVE_BASENAME}.tar.gz" prometheus
          fi

-      - name: Upload benchmark artifacts
+      - name: Upload workload artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v4
        with:
-          name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
+          name: ${{ env.ARTIFACT_NAME }}
          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
+  collection-benchmarks:
+    name: collection (${{ matrix.backend.id }}, ${{ matrix.workload.id }})
+    runs-on: ${{ matrix.backend.runner }}
+    timeout-minutes: 15
    strategy:
      fail-fast: false
      matrix:
        backend:
          - id: memory
-            compose_file: compose.prometheus-memory-store.yml
+            needs_mongo: false
+            runner: ubuntu-latest
          - 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
+            needs_mongo: true
+            runner: ubuntu-latest
+        workload:
+          - id: high-insert
+            total_tasks: 100000
+            concurrency: 2048
+            type: insert
+          - id: medium-insert
+            total_tasks: 100000
+            concurrency: 128
+            type: insert
+          - id: low-insert
+            total_tasks: 100000
+            concurrency: 4
+            type: insert
+          - id: high-dequeue
+            total_tasks: 100000
+            concurrency: 2048
+            type: dequeue
+          - id: medium-dequeue
+            total_tasks: 100000
+            concurrency: 128
+            type: dequeue
+          - id: low-dequeue
+            total_tasks: 100000
+            concurrency: 4
+            type: dequeue
    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
+      ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.backend.id, matrix.workload.id) }}
+      SUMMARY_FILE: ${{ format('artifacts/{0}-{1}/summary-{0}-{1}.jsonl', matrix.backend.id, matrix.workload.id) }}
+      ARTIFACT_NAME: ${{ format('collections-{0}-{1}', matrix.backend.id, matrix.workload.id) }}
+      MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
    steps:
      - uses: actions/checkout@v4

@@ -232,103 +377,60 @@ jobs:
      - name: Sync dependencies
        run: uv sync --frozen --extra mongo --group core-stable --group dev

-      - name: Reset benchmark data directories
+      - name: Launch MongoDB
+        if: ${{ matrix.backend.needs_mongo }}
        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
+          docker compose -f compose.mongo.yml down -v || true
+          docker compose -f compose.mongo.yml up -d --quiet-pull
          for attempt in {1..60}; do
-            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+            if docker compose -f compose.mongo.yml exec -T mongo mongosh --quiet --eval 'db.runCommand({ping:1})' >/dev/null 2>&1; then
              exit 0
            fi
-            sleep 1
+            sleep 2
          done
-          echo "Store did not become ready in time" >&2
-          cd docker && docker compose -f "$COMPOSE_FILE" logs app
+          echo "MongoDB did not become ready in time" >&2
+          docker compose -f compose.mongo.yml logs mongo
          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 }}
+      - name: Run collection benchmark
        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"
+          echo "Running collection benchmark (backend=${{ matrix.backend.id }}, workload=${{ matrix.workload.id }})"
+          uv run --locked --no-sync python -m tests.benchmark.collection_benchmark \
+            "${{ matrix.workload.type }}" \
+            --backend "${{ matrix.backend.id }}" \
+            --total-tasks "${{ matrix.workload.total_tasks }}" \
+            --concurrency "${{ matrix.workload.concurrency }}" \
+            --task-prefix "${{ matrix.backend.id }}-${{ matrix.workload.id }}" \
+            --summary-file "$SUMMARY_FILE" \
+            --mongo-uri "$MONGO_URI" \
+            --mongo-database agentlightning_collection_bench

-      - name: Record micro benchmark end
-        if: ${{ always() }}
-        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
-
-      - name: Run micro benchmark analysis
+      - name: Show collection benchmark summary
        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"
+          if [ -f "$SUMMARY_FILE" ]; then
+            echo "Collection benchmark summary (${{ matrix.backend.id }}):"
+            cat "$SUMMARY_FILE"
          else
-            echo "Summary file not found: $summary_file"
+            echo "Summary file not found: $SUMMARY_FILE"
          fi

-      - name: Stop ${{ matrix.backend.id }} Prometheus stack
-        if: ${{ always() }}
+      - name: Stop MongoDB
+        if: ${{ always() && matrix.backend.needs_mongo }}
        run: |
          set -euo pipefail
          cd docker
-          docker compose -f "$COMPOSE_FILE" down -v || true
+          docker compose -f compose.mongo.yml 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
+      - name: Upload collection artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v4
        with:
-          name: micro-benchmark-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+          name: ${{ env.ARTIFACT_NAME }}
          path: ${{ env.ARTIFACT_DIR }}
          if-no-files-found: error
@@ -270,6 +270,31 @@ jobs:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

+      - name: Training with LoRA
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --lora
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_lora
+        if: matrix.setup-script != 'legacy'
+
+      - name: Validate training with LoRA
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_lora.outputs.project_name }} ${{ steps.calc_x_train_lora.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        if: matrix.setup-script != 'legacy'
+
      - name: Training with external store
        run: |
          set -euo pipefail
@@ -0,0 +1,168 @@
+name: Examples - ChartQA
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: "0 22 * * *"
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-chartqa, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'ChartQA - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('ChartQA - {0}', github.event_name) }}
+
+jobs:
+  chartqa:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-chartqa' ||
+      github.event.action == 'ci-all'
+    name: ChartQA (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'
+      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 --extra verl \
+            --group dev --group experiment --group image --group langchain --group vllm-0-10-2 --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-chartqa-${{ 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 ChartQA dataset
+        run: |
+          set -euo pipefail
+          cd examples/chartqa
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1fWRt9hehg8_uV7BDWSCwKTycM60JcmGN/view?usp=sharing" -O chartqa-data.zip
+          unzip chartqa-data.zip
+          rm chartqa-data.zip
+        shell: bash
+
+      - name: ChartQA sanity check with GPT
+        run: |
+          set -euo pipefail
+          cd examples/chartqa
+          uv run python debug_chartqa_agent.py
+        shell: bash
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Run vLLM Server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/chartqa
+          uv run --no-sync vllm serve Qwen/Qwen2-VL-2B-Instruct \
+            --gpu-memory-utilization 0.9 \
+            --max-model-len 4096 \
+            --allowed-local-media-path "$(pwd)/data" \
+            --enable-prefix-caching \
+            --port 8088 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:8088/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: ChartQA sanity check with vLLM
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/chartqa
+          uv run python debug_chartqa_agent.py
+        shell: bash
+        env:
+          USE_LLM_PROXY: "1"
+          OPENAI_API_BASE: http://localhost:8088/v1
+          OPENAI_MODEL: Qwen/Qwen2-VL-2B-Instruct
+
+      - 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: ChartQA training
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/chartqa
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_chartqa_agent.py ci
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: chartqa_train
+
+      - name: Validate ChartQA training
+        run: |
+          set -euo pipefail
+          uv run scripts/validate_example_wandb.py ${{ steps.chartqa_train.outputs.project_name }} ${{ steps.chartqa_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -74,7 +74,7 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-rag-${{ matrix.python-version }}-${{ matrix.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

@@ -68,13 +68,23 @@ jobs:
          path: requirements-freeze.txt
          compression-level: 0

-      - name: Tinker LLM sanity check
+      # TODO: Currently only test the client tracer implementation.
+      - name: Tinker LLM sanity check (tracer text)
        run: |
          set -euo pipefail
          source .venv/bin/activate
          cd examples/tinker
-          # TODO: Currently only test the client tracer implementation.
-          python -m tests.test_tinker_llm
+          python -m tests.test_tinker_llm tracer-text
+        shell: bash
+        env:
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker LLM sanity check (tracer tool)
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python -m tests.test_tinker_llm tracer-tool
        shell: bash
        env:
          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
@@ -0,0 +1,18 @@
+# Pre-defined workflow with workflow_dispatch trigger,
+# convenient for testing and debugging.
+
+name: Playground
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+
+jobs:
+  playground:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Run script
+        run: |
+          echo "Hello, world!"
@@ -27,13 +27,37 @@ jobs:
      github.event_name != 'repository_dispatch' ||
      github.event.action == 'ci-gpu' ||
      github.event.action == 'ci-all'
-    name: GPU Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
+    name: Full Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})

-    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    runs-on: ${{ matrix.mark.runs-on }}
    timeout-minutes: 30
    strategy:
      matrix:
-        include:
+        mark:
+          - id: store
+            display-name: Store
+            pytest-mark: 'store'  # store tests should not require gpu
+            runs-on: ubuntu-latest
+            has-gpu: false
+          # AgentOps needs to be separated because it injects tricky global state.
+          - id: agentops
+            display-name: AgentOps
+            pytest-mark: 'agentops'  # including agentops+litellm tests here
+            runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+            has-gpu: true
+          # Other tests that require GPU
+          - id: gpu
+            display-name: GPU required
+            pytest-mark: '(gpu or llmproxy) and not agentops'
+            runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+            has-gpu: true
+          # Other uncovered tests
+          - id: others
+            display-name: Others
+            pytest-mark: 'not store and not agentops and not gpu and not llmproxy'
+            runs-on: ubuntu-latest
+            has-gpu: false
+        env:
        - python-version: '3.10'
          setup-script: 'legacy'
        - python-version: '3.12'
@@ -43,6 +67,7 @@ jobs:
      fail-fast: false
    steps:
      - name: Check GPU status
+        if: matrix.mark.has-gpu
        run: nvidia-smi
      - uses: actions/checkout@v4
        with:
@@ -51,20 +76,32 @@ jobs:
      - uses: astral-sh/setup-uv@v7
        with:
          enable-cache: true
-          python-version: ${{ matrix.python-version }}
+          python-version: ${{ matrix.env.python-version }}
      - name: Upgrade dependencies (latest)
        run: uv lock --upgrade
-        if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (latest)
+        if: matrix.env.setup-script == 'latest'
+
+      - name: Sync dependencies (latest, gpu)
+        if: matrix.env.setup-script == 'latest' && matrix.mark.has-gpu
        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)
-        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 vllm/pytorch on CPU counterparts
+      - name: Sync dependencies (latest, cpu)
+        if: matrix.env.setup-script == 'latest' && !matrix.mark.has-gpu
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group core-stable
+      - name: Sync dependencies (stable, gpu)
+        if: matrix.env.setup-script == 'stable' && matrix.mark.has-gpu
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.env.setup-script }}
+      - name: Sync dependencies (stable, cpu)
+        if: matrix.env.setup-script == 'stable' && !matrix.mark.has-gpu
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group core-stable
      # Don't install langchain for legacy dependency because it has conflicts with torch.
-      - name: Sync dependencies (legacy)
+      - name: Sync dependencies (legacy, gpu)
+        if: matrix.env.setup-script == 'legacy' && matrix.mark.has-gpu
        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: Sync dependencies (legacy, cpu)
+        if: matrix.env.setup-script == 'legacy' && !matrix.mark.has-gpu
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group core-legacy
+
      - name: Freeze dependencies
        run: |
          set -ex
@@ -74,13 +111,15 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-tests-full-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-tests-full-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

      - uses: actions/setup-node@v6
        with:
          node-version: '22'
+          cache: 'npm'
+          cache-dependency-path: dashboard/package-lock.json
      - name: Install JavaScript dependencies
        run: cd dashboard && npm ci
      - name: Build dashboard
@@ -139,9 +178,10 @@ jobs:
          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}

+      # mongo, openai, gpu, all enabled by default
      - name: Run tests
        run: |
-          uv run pytest -v --durations=0 tests
+          uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}"
        env:
          PYTEST_ADDOPTS: "--color=yes"
          OPENAI_BASE_URL: http://localhost:12306/
@@ -19,7 +19,8 @@ jobs:
  lint:
    strategy:
      matrix:
-        setup: [fast, slow]
+        setup: [fast, slow, next]
+      fail-fast: false
    name: Lint - ${{ matrix.setup }}
    runs-on: ubuntu-latest
    timeout-minutes: 10
@@ -32,6 +33,9 @@ jobs:
      - name: Sync dependencies (fast)
        run: uv sync --frozen --group dev --no-default-groups
        if: matrix.setup == 'fast'
+      - name: Upgrade dependencies (next)
+        run: uv lock --upgrade
+        if: matrix.setup == 'next'
      - name: Sync dependencies (slow)
        run: |
          uv sync --frozen \
@@ -46,7 +50,7 @@ jobs:
            --group agents \
            --group langchain \
            --no-default-groups
-        if: matrix.setup == 'slow'
+        if: matrix.setup != 'fast'
      # This pre-commit skips JavaScript on purpose.
      - name: Run pre-commit
        uses: pre-commit/action@v3.0.1
@@ -61,7 +65,7 @@ jobs:
        if: matrix.setup == 'fast'
      - name: Run pyright (slow)
        run: uv run --locked --no-sync pyright -p pyrightconfig.json
-        if: matrix.setup == 'slow'
+        if: matrix.setup != 'fast'

  lint-js:
    name: Lint - JavaScript
@@ -72,6 +76,8 @@ jobs:
      - uses: actions/setup-node@v6
        with:
          node-version: '22'
+          cache: 'npm'
+          cache-dependency-path: dashboard/package-lock.json
      - name: Install dependencies
        run: cd dashboard && npm ci
      - name: Run ESLint
@@ -116,7 +122,28 @@ jobs:
  test:
    strategy:
      matrix:
-        include:
+        mark:
+          # store has many tests and is a good isolated group.
+          - id: store
+            display-name: Store
+            pytest-mark: 'store'
+          # AgentOps needs to be separated because it injects tricky global state.
+          - id: agentops
+            display-name: AgentOps
+            pytest-mark: 'agentops'
+          # litellm proxy tests are slow
+          - id: llmproxy
+            display-name: LLM proxy
+            pytest-mark: 'llmproxy'
+          # Robustness of utilities is important. There are many tests.
+          - id: utils
+            display-name: Utilities
+            pytest-mark: 'utils'
+          # unmarked tests: adapter, execution engine, etc.
+          - id: others
+            display-name: Others
+            pytest-mark: 'not store and not agentops and not llmproxy and not utils'
+        env:
          - python-version: '3.10'
            setup-script: 'legacy'
          - python-version: '3.11'
@@ -127,7 +154,7 @@ jobs:
            setup-script: 'latest'
      fail-fast: false

-    name: Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
+    name: Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})
    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
@@ -135,16 +162,16 @@ jobs:
      - uses: astral-sh/setup-uv@v7
        with:
          enable-cache: true
-          python-version: ${{ matrix.python-version }}
+          python-version: ${{ matrix.env.python-version }}
      - name: Upgrade dependencies (latest)
        run: uv lock --upgrade
-        if: matrix.setup-script == 'latest'
+        if: matrix.env.setup-script == 'latest'
      - name: Sync dependencies (latest)
        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-stable
-        if: matrix.setup-script == 'latest'
+        if: matrix.env.setup-script == 'latest'
      - name: Sync dependencies (stable & legacy)
-        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'
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
+        if: matrix.env.setup-script != 'latest'
      - name: Freeze dependencies
        run: |
          set -ex
@@ -154,13 +181,15 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

      - uses: actions/setup-node@v6
        with:
          node-version: '22'
+          cache: 'npm'
+          cache-dependency-path: dashboard/package-lock.json
      - name: Install JavaScript dependencies
        run: cd dashboard && npm ci
      - name: Build dashboard
@@ -168,12 +197,12 @@ jobs:

      - name: Run tests
        run: |
-          uv run pytest -v --durations=0 tests -m "not mongo"
+          uv run pytest -v --durations=0 tests -m "not mongo and not openai and not gpu and (${{ matrix.mark.pytest-mark }})"
        env:
          PYTEST_ADDOPTS: "--color=yes"

  test-js:
-    name: Test - JavaScript
+    name: Test (JavaScript)
    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
@@ -183,6 +212,8 @@ jobs:
      - uses: actions/setup-node@v6
        with:
          node-version: '22'
+          cache: 'npm'
+          cache-dependency-path: dashboard/package-lock.json
      - uses: astral-sh/setup-uv@v7
        with:
          enable-cache: true
@@ -3,6 +3,7 @@ verl_old
 meta-llama/**
 debug/*.png
 requirements-freeze*.txt
+/playground

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -0,0 +1,41 @@
+# Repository Guidelines
+
+## Architecture Overview
+Agent Lightning runs through a continuous loop: runners and tracers emit spans, `LightningStore` (`agentlightning/store/`) keeps them synchronized, and algorithms in `agentlightning/algorithm/` consume those traces to improve behavior.
+
+## Project Structure & Module Organization
+- `agentlightning/`: adapters, execution stack, training loop, tracer, reward logic, and the `agl` CLI.
+- `docs/` & `examples/`: narrative and procedural docs (assets in `docs/assets/`, navigation in `mkdocs.yml`) plus runnable workflows whose READMEs point to their companion how-to guides. `docs/how-to` covers task-focused instructions, while `docs/tutorials` explains concepts and subsystems.
+- `dashboard/`, `scripts/`, `tests/`: UI bundles, release/dataset/CI automation, and mirrored coverage of the runtime tree. Record download steps rather than committing binaries.
+
+## Build, Test, and Development Commands
+- `uv sync --group dev` — provision tooling once per environment.
+- `uv run --no-sync pytest -v` — execute the full suite; add a path or `-k expr` to narrow the run.
+- `uv run --no-sync pyright` — enforce static typing parity with CI.
+- `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` and `uv run --no-sync mkdocs build --strict` — keep formatting tidy and documentation valid.
+Always commit the refreshed `uv.lock` when dependencies shift, and mention optional groups (VERL, APO, GPU) in PR notes.
+
+## Common Issues & Fixes
+- When `uv run` errors with `Permission denied` under `~/.cache`, override both cache locations inline: ``UV_CACHE="$(pwd)/.cache_uv" XDG_CACHE_HOME="$(pwd)/.cache_xdg" uv run --no-sync <command>``.
+
+## Coding Style & Naming Conventions
+- Target `requires-python >= 3.10`, four-space indentation, 120-character lines (though docstrings may run longer), and formatter-owned diffs (Black + isort, `black` profile). Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
+- Maintain exhaustive type hints (pyright enforces them) and prefer shared dataclasses or Pydantic models from `agentlightning.types`.
+- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points. Use mkdocs styles: `[][]` syntax for cross-references and single backticks for inline code blocks.
+- Writing logs is encouraged, especially for long functions with multiple steps and try-except blocks that catch all exceptions. Use `logging.getLogger(__name__)` to get loggers. Distinguish between DEBUG, INFO, WARNING, and ERROR logs.
+
+## Testing Guidelines
+- Mirror runtime directories under `tests/` and match filenames for quick traceability.
+- Parametrize pytest cases and apply markers (`openai`, `gpu`, `agentops`, `mongo`, `llmproxy`) so optional suites can be skipped via selectors like `-m "not mongo"` yet still exercised in CI.
+- Lean on fixtures, favor real stores/spans/agents over mocks, and drive coverage across the majority of branches.
+- If an imported module is missing from the environment, check whether `uv sync` has been run with the right groups. Do not make stubs for external dependencies unless necessary.
+
+## Example Contributions
+- Ship each example with a README that includes smoke-test instructions so maintainers can validate quickly. The README must contain an "Included Files" section summarizing every file and its role.
+- Keep runnable example modules self-contained with a module-level docstring describing CLI usage. Document important or educational classes/functions with targeted docstrings and inline comments where clarity matters.
+- Add a CI workflow per example named `examples-<name>.yml` in `.github/workflows/`. Register it in `badge-<name>.yml`, `badge-examples.yml`, and `badge-latest.yml` when applicable so badges stay accurate.
+
+## Commit & Pull Request Guidelines
+- Branch from a fresh `main` using `feature/<slug>`, `fix/<slug>`, `docs/<slug>`, or `chore/<slug>`.
+- Write imperative, scoped commits, reference issues with `Fixes #123`, and rerun pre-commit plus the relevant pytest/doc builds before pushing.
+- Use PR descriptions to summarize intent, list verification commands, call out dependency or docs-navigation updates, and link new docs/examples via `mkdocs.yml` or `examples/README.md`. Include logs for dashboard changes.
@@ -0,0 +1 @@
+AGENTS.md
@@ -57,6 +57,7 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.

 - [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
 - [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
+- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl).

 ## ⚡ Architecture

@@ -13,6 +13,7 @@ from pydantic import BaseModel

 from agentlightning.emitter.reward import get_reward_value
 from agentlightning.types import Span, Triplet
+from agentlightning.utils.otel import filter_and_unflatten_attributes

 from .base import TraceAdapter

@@ -376,7 +377,8 @@ class TraceTree:

            if is_llm_call:
                llm_calls.append((self, within_matching_subtree))  # type: ignore
-                existing_llm_call_response_ids = existing_llm_call_response_ids or set()
+                if existing_llm_call_response_ids is None:
+                    existing_llm_call_response_ids = set()
                if response_id is not None:
                    existing_llm_call_response_ids.add(response_id)
                if within_llm_call is not None:
@@ -490,12 +492,12 @@ class TraceTree:
                assign_to: List[Tuple[str, int]] = []
                for child in item.children:
                    if child.id in llm_call_ids:
-                        assign_to.append(child.id)  # type: ignore
+                        assign_to.append((child.id, child.end_time))  # type: ignore

-                    agentops_output = item.maybe_reward_dict()
+                    agentops_output = child.maybe_reward_dict()
                    if agentops_output and agentops_output.get("type") == "reward":
                        for assign_to_id, assign_to_end_time in reversed(assign_to):
-                            if assign_to_end_time > item.start_time:  # type: ignore
+                            if assign_to_end_time > child.start_time:  # type: ignore
                                # This reward happens before the end of the LLM call.
                                continue
                            if assign_to_id in rewards:
@@ -505,6 +507,60 @@ class TraceTree:

        return rewards

+    def extract_prompt_image_urls(self, prompt_raw_content: Any) -> List[str]:
+        """Extract image URLs from the span attributes, in order of appearance.
+
+        Args:
+            prompt_raw_content: The raw content of the prompt, which can be in one of several formats:
+
+                - List[dict]: A list of message entries, each being a dict with at least a "content" key.
+                - Dict[str, Any]: A dictionary, often with numeric string keys (e.g., `{"0": {...}, "1": {...}}`), where each value is a message entry.
+                  If the dict does not have numeric keys, it is treated as a single message entry.
+        """
+        message_entries: List[Any] = []
+        if isinstance(prompt_raw_content, list):
+            message_entries = cast(List[Any], prompt_raw_content)
+        elif isinstance(prompt_raw_content, dict):
+            # Common when the attributes expand to {"0": {...}, "prompt_filter_results": ...}
+            numeric_keys = [
+                key
+                for key in cast(Dict[str, Any], prompt_raw_content).keys()
+                if isinstance(key, str) and key.isdigit()  # pyright: ignore[reportUnnecessaryIsInstance]
+            ]
+            if numeric_keys:
+                for key in sorted(numeric_keys, key=int):
+                    message_entries.append(prompt_raw_content[key])
+            else:
+                message_entries = [prompt_raw_content]
+        else:
+            return []
+
+        image_urls: List[str] = []
+        for message in cast(List[Dict[str, Any]], message_entries):
+            if (
+                not isinstance(message, dict)  # pyright: ignore[reportUnnecessaryIsInstance]
+                or "content" not in message
+            ):
+                continue
+            content = message["content"]
+            if isinstance(content, str):
+                try:
+                    content = json.loads(content)  # This content should now be a list
+                except json.JSONDecodeError:
+                    logger.warning(f"Failed to parse message content as JSON: {content}")
+                    continue
+            if isinstance(content, list):
+                for content_part in cast(List[Dict[str, Any]], content):
+                    if not isinstance(content_part, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+                        continue
+                    if content_part.get("type") == "image_url":
+                        image_url_dict = cast(Dict[str, Any], content_part.get("image_url"))
+                        if not isinstance(image_url_dict, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+                            continue
+                        if "url" in image_url_dict:
+                            image_urls.append(image_url_dict["url"])
+        return image_urls
+
    def span_to_triplet(self, span: Span, agent_name: str) -> Triplet:
        """Convert a span to a triplet.

@@ -514,19 +570,27 @@ class TraceTree:
        prompt_token_ids = span.attributes.get("prompt_token_ids", [])  # type: ignore
        response_token_ids = span.attributes.get("response_token_ids", [])  # type: ignore
        response_id = span.attributes.get("gen_ai.response.id", None)  # type: ignore
+        request_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.request")
+        response_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.response")
+        prompt_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.prompt")
+        completion_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.completion")
+        image_urls = self.extract_prompt_image_urls(prompt_raw_content)
+
+        prompt_payload = {"token_ids": prompt_token_ids, "raw_content": prompt_raw_content, "image_urls": image_urls}
+        response_payload = {"token_ids": response_token_ids, "raw_content": completion_raw_content}

        logprobs_content = span.attributes.get("logprobs.content", None)  # type: ignore
        if isinstance(logprobs_content, str):
            logprobs_content = json.loads(logprobs_content)
-            response: Dict[str, Any] = {"token_ids": response_token_ids, "logprobs": logprobs_content}
-        else:
-            response = {"token_ids": response_token_ids}
+            response_payload["logprobs"] = logprobs_content

        return Triplet(
-            prompt={"token_ids": prompt_token_ids},
-            response=response,
+            prompt=prompt_payload,
+            response=response_payload,
            reward=None,
-            metadata=dict(response_id=response_id, agent_name=agent_name),
+            metadata=dict(
+                request=request_metadata, response=response_metadata, response_id=response_id, agent_name=agent_name
+            ),
        )

    def to_trajectory(
@@ -7,23 +7,44 @@ APO with textual gradients that read rollout spans and outputs to modify the pro
 - rollout: same pattern as your example, but task is a dict (T_task)
 """

+from __future__ import annotations
+
 import asyncio
 import logging
 import random
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Counter,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TypedDict,
+    TypeVar,
+    cast,
+)

 import poml
 from openai import AsyncOpenAI

 from agentlightning.adapter.messages import TraceToMessages
 from agentlightning.algorithm.base import Algorithm
-from agentlightning.algorithm.utils import batch_iter_over_dataset
+from agentlightning.algorithm.utils import batch_iter_over_dataset, with_llm_proxy, with_store
 from agentlightning.reward import find_final_reward
 from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus

+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
 logger = logging.getLogger(__name__)

 T_task = TypeVar("T_task")
@@ -360,8 +381,10 @@ class APO(Algorithm, Generic[T_task]):
            )
        return new_prompt

+    @with_store
    async def get_rollout_results(
        self,
+        store: LightningStore,
        rollout: List[Rollout],
        *,
        prefix: Optional[str] = None,
@@ -379,7 +402,6 @@ class APO(Algorithm, Generic[T_task]):
            List of rollout results formatted for APO processing.
        """
        rollout_results: List[RolloutResultForAPO] = []
-        store = self.get_store()
        adapter = self.get_adapter()
        for r in rollout:
            spans = await store.query_spans(r.rollout_id)
@@ -776,8 +798,12 @@ class APO(Algorithm, Generic[T_task]):
                prefix=prefix,
            )

+    @with_llm_proxy()
+    @with_store
    async def run(
        self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
        train_dataset: Optional[Dataset[T_task]] = None,
        val_dataset: Optional[Dataset[T_task]] = None,
    ) -> None:
@@ -5,11 +5,16 @@ from __future__ import annotations
 import asyncio
 import logging
 from datetime import datetime
-from typing import Any, List, Literal, Optional
+from typing import TYPE_CHECKING, Any, List, Literal, Optional

 from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span

 from .base import Algorithm
+from .utils import with_llm_proxy, with_store
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore

 logger = logging.getLogger(__name__)

@@ -36,6 +41,8 @@ class Baseline(FastAlgorithm):
    finish, and logs every collected span and reward. It is primarily useful as
    a smoke test for the platform plumbing rather than a performant trainer.

+    The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started.
+
    Args:
        n_epochs: Number of dataset passes to execute for both the train and val
            splits during developer experiments.
@@ -180,8 +187,12 @@ class Baseline(FastAlgorithm):

            await asyncio.sleep(self.polling_interval)

+    @with_llm_proxy()
+    @with_store
    async def run(
        self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> None:
@@ -202,8 +213,6 @@ class Baseline(FastAlgorithm):
        logger.debug(f"Train indices: {train_indices}")
        logger.debug(f"Val indices: {val_indices}")

-        store = self.get_store()
-
        # Currently we only supports a single resource update at the start.
        initial_resources = self.get_initial_resources()
        if initial_resources is not None:
@@ -1,11 +1,42 @@
 # Copyright (c) Microsoft. All rights reserved.

+from __future__ import annotations
+
+import functools
+import logging
 import random
-from typing import Iterator, List, Sequence, TypeVar
+from collections.abc import Coroutine
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Concatenate,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    ParamSpec,
+    Sequence,
+    TypeVar,
+    overload,
+)

 from agentlightning.types import Dataset

+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
+    from .base import Algorithm
+
 T_task = TypeVar("T_task")
+T_algo = TypeVar("T_algo", bound="Algorithm")
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+logger = logging.getLogger(__name__)


 def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
@@ -41,3 +72,106 @@ def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterat
                if len(current_batch) == batch_size:
                    yield [dataset[index] for index in current_batch]
                    current_batch = []
+
+
+def with_store(
+    func: Callable[Concatenate[T_algo, LightningStore, P], Coroutine[Any, Any, R]],
+) -> Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]]:
+    """Inject the algorithm's `LightningStore` into coroutine methods.
+
+    The decorator calls `Algorithm.get_store()` once per invocation and passes the
+    resulting store as an explicit argument to the wrapped coroutine. Decorated
+    methods therefore receive the resolved store even when invoked by helper
+    utilities rather than directly by the algorithm.
+
+    Args:
+        func: The coroutine that expects `(self, store, *args, **kwargs)`.
+
+    Returns:
+        A coroutine wrapper that automatically retrieves the store and forwards it
+        to `func`.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: T_algo, *args: P.args, **kwargs: P.kwargs) -> R:
+        store = self.get_store()
+        return await func(self, store, *args, **kwargs)
+
+    return wrapper
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[False] = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, Optional[LLMProxy], P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[True],
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, LLMProxy, P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+def with_llm_proxy(
+    required: bool = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[..., Coroutine[Any, Any, Any]]],
+    Callable[..., Coroutine[Any, Any, Any]],
+]:
+    """Resolve and optionally lifecycle-manage the configured LLM proxy.
+
+    Args:
+        required: When True, raises `ValueError` if the algorithm does not have an
+            [`LLMProxy`][agentlightning.LLMProxy] set. When False, the wrapped coroutine receives
+            `None` if no proxy is available.
+        auto_start: When True, [`LLMProxy.start()`][agentlightning.LLMProxy.start] is invoked if the proxy is not
+            already running before calling `func` and [`LLMProxy.stop()`][agentlightning.LLMProxy.stop] is
+            called afterwards.
+
+    Returns:
+        A decorator that injects the [`LLMProxy`][agentlightning.LLMProxy] (or `None`) as the first
+        argument after `self` and manages automatic startup/shutdown when requested.
+    """
+
+    def decorator(
+        func: Callable[..., Coroutine[Any, Any, Any]],
+    ) -> Callable[..., Coroutine[Any, Any, Any]]:
+        @functools.wraps(func)
+        async def wrapper(self: Algorithm, *args: Any, **kwargs: Any) -> Any:
+            llm_proxy = self.get_llm_proxy()
+
+            if required and llm_proxy is None:
+                raise ValueError(
+                    "LLM proxy is required but not configured. Call set_llm_proxy() before using this method."
+                )
+
+            auto_started = False
+            if auto_start and llm_proxy is not None:
+                if llm_proxy.is_running():
+                    logger.info("Proxy is already running, skipping start")
+                else:
+                    logger.info("Starting proxy, managed by the algorithm")
+                    await llm_proxy.start()
+                    auto_started = True
+
+            try:
+                # At type level, overloads guarantee that if `required=True`
+                # then `func` expects a non-optional LLMProxy.
+                return await func(self, llm_proxy, *args, **kwargs)
+            finally:
+                if auto_started and llm_proxy is not None:
+                    logger.info("Stopping proxy, managed by the algorithm")
+                    await llm_proxy.stop()
+
+        return wrapper
+
+    return decorator
@@ -1,6 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.

-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional, Type

 from hydra import compose, initialize
 from omegaconf import OmegaConf
@@ -10,6 +12,10 @@ from agentlightning.client import AgentLightningClient
 from agentlightning.types import Dataset
 from agentlightning.verl.entrypoint import run_ppo  # type: ignore

+if TYPE_CHECKING:
+    from agentlightning.verl.daemon import AgentModeDaemon
+    from agentlightning.verl.trainer import AgentLightningTrainer
+

 class VERL(Algorithm):
    """VERL-powered algorithm that delegates training to the VERL PPO runner.
@@ -23,6 +29,8 @@ class VERL(Algorithm):
        config: Dictionary mirroring the overrides passed to the VERL CLI. The
            overrides are merged with VERL's packaged defaults via Hydra before
            launching training.
+        trainer_cls: Optional override for the trainer class. Experimental.
+        daemon_cls: Optional override for the daemon class. Experimental.

    Examples:
        ```python
@@ -90,7 +98,12 @@ class VERL(Algorithm):
        ```
    """

-    def __init__(self, config: dict[str, Any]):
+    def __init__(
+        self,
+        config: dict[str, Any],
+        trainer_cls: Optional[Type[AgentLightningTrainer]] = None,
+        daemon_cls: Optional[Type[AgentModeDaemon]] = None,
+    ):
        super().__init__()

        # Compose the base config exactly like your decorator:
@@ -102,6 +115,8 @@ class VERL(Algorithm):
        # Allow adding new fields
        OmegaConf.set_struct(base_cfg, False)
        self.config = OmegaConf.merge(base_cfg, override_conf)
+        self.trainer_cls = trainer_cls
+        self.daemon_cls = daemon_cls

    def run(
        self,
@@ -119,6 +134,11 @@ class VERL(Algorithm):
                adapter have been garbage-collected when using the V1 execution
                mode.
        """
+        from agentlightning.verl.daemon import AgentModeDaemon
+        from agentlightning.verl.trainer import AgentLightningTrainer
+
+        trainer_cls = self.trainer_cls or AgentLightningTrainer
+        daemon_cls = self.daemon_cls or AgentModeDaemon
        try:
            store = self.get_store()
        except Exception:
@@ -130,6 +150,8 @@ class VERL(Algorithm):
                store=None,
                llm_proxy=None,
                adapter=None,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
            )
        else:
            print("Store is set. Assuming v1 execution mode.")
@@ -142,6 +164,8 @@ class VERL(Algorithm):
                store=store,
                llm_proxy=llm_proxy,
                adapter=adapter,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
            )

    def get_client(self) -> AgentLightningClient:
@@ -12,6 +12,7 @@ from typing import Dict, Iterable, Tuple
 _SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
    "vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
    "store": ("agentlightning.cli.store", "Run a LightningStore server."),
+    "prometheus": ("agentlightning.cli.prometheus", "Serve Prometheus metrics from the multiprocess registry."),
    "agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
 }

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Serve Prometheus metrics from the Agent Lightning multiprocess registry."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+from pathlib import Path
+from typing import Iterable
+
+from fastapi import FastAPI
+from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]
+
+from agentlightning.logging import setup as setup_logging
+from agentlightning.utils.metrics import get_prometheus_registry
+from agentlightning.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_prometheus_dir() -> str:
+    """Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
+
+    directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
+    if directory is None:
+        raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
+
+    Path(directory).mkdir(parents=True, exist_ok=True)
+    logger.info("Serving Prometheus multiprocess metrics from %s", directory)
+    return directory
+
+
+def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
+    """Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
+
+    Args:
+        metrics_path: URL path to expose the Prometheus metrics endpoint on.
+
+    Returns:
+        A FastAPI application ready to serve metrics.
+    """
+
+    if not metrics_path.startswith("/"):
+        raise ValueError("metrics_path must start with '/'.")
+
+    normalized_path = metrics_path.rstrip("/")
+    if normalized_path in ("", "/"):
+        raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
+
+    app = FastAPI(title="Agent Lightning Prometheus exporter", docs_url=None, redoc_url=None)
+    metrics_app = make_asgi_app(registry=get_prometheus_registry())  # pyright: ignore[reportUnknownVariableType]
+    app.mount(normalized_path, metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+
+    @app.get("/health")
+    async def healthcheck() -> dict[str, str]:  # pyright: ignore[reportUnusedFunction]
+        return {"status": "ok"}
+
+    return app
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
+    parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
+    parser.add_argument(
+        "--metrics-path",
+        default="/v1/prometheus",
+        help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the metrics server.",
+    )
+    parser.add_argument(
+        "--access-log",
+        action="store_true",
+        help="Enable uvicorn access logs. Disabled by default to reduce noise.",
+    )
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+    ensure_prometheus_dir()
+
+    try:
+        app = create_prometheus_app(args.metrics_path)
+    except ValueError as exc:
+        logger.error("Failed to configure prometheus app: %s", exc)
+        return 1
+
+    launcher_args = PythonServerLauncherArgs(
+        host=args.host,
+        port=args.port,
+        log_level=getattr(logging, args.log_level),
+        access_log=args.access_log,
+        healthcheck_url="/health",
+    )
+    launcher = PythonServerLauncher(app, launcher_args)
+
+    try:
+        asyncio.run(launcher.run_forever())
+    except KeyboardInterrupt:
+        logger.info("Received shutdown signal. Stopping Prometheus server.")
+    except RuntimeError as exc:
+        logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -7,11 +7,18 @@ from __future__ import annotations
 import argparse
 import asyncio
 import logging
-from typing import Iterable
+from typing import Iterable, List

 from agentlightning import setup_logging
 from agentlightning.store.client_server import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+    setup_multiprocess_prometheus,
+)

 logger = logging.getLogger(__name__)

@@ -33,9 +40,10 @@ def main(argv: Iterable[str] | None = None) -> int:
        help="Configure the logging level for the store.",
    )
    parser.add_argument(
-        "--prometheus",
-        action="store_true",
-        help="Enable Prometheus metrics.",
+        "--tracker",
+        nargs="+",
+        choices=["prometheus", "console"],
+        help="Enable metrics tracking. Repeat for multiple trackers.",
    )
    parser.add_argument(
        "--n-workers",
@@ -63,14 +71,36 @@ def main(argv: Iterable[str] | None = None) -> int:

    setup_logging(args.log_level)

+    trackers: List[MetricsBackend] = []
+    if args.tracker:
+        if "prometheus" in args.tracker:
+            logger.info("Enabling Prometheus metrics tracking.")
+            if args.n_workers > 1:
+                # This has to be done before prometheus_client is imported
+                setup_multiprocess_prometheus()
+                logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
+            trackers.append(PrometheusMetricsBackend())
+
+        if "console" in args.tracker:
+            logger.info("Enabling console metrics tracking.")
+            trackers.append(ConsoleMetricsBackend())
+
+    if len(trackers) == 0:
+        tracker: MetricsBackend | None = None
+    elif len(trackers) == 1:
+        tracker = trackers[0]
+    else:
+        tracker = MultiMetricsBackend(trackers)
+
    if args.backend == "memory":
        store = InMemoryLightningStore(
-            prometheus=args.prometheus, thread_safe=True
-        )  # Using thread_safe store for server
+            thread_safe=True,  # Using thread_safe store for server
+            tracker=tracker,
+        )
    elif args.backend == "mongo":
        from agentlightning.store.mongo import MongoLightningStore

-        store = MongoLightningStore(client=args.mongo_uri, prometheus=args.prometheus)
+        store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
    else:
        raise ValueError(f"Invalid backend: {args.backend}")

@@ -86,7 +116,7 @@ def main(argv: Iterable[str] | None = None) -> int:
        port=args.port,
        cors_allow_origins=args.cors_origins,
        launch_mode=launch_mode,
-        prometheus=args.prometheus,
+        tracker=tracker,
        n_workers=args.n_workers,
    )
    try:
@@ -143,6 +143,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                logger.debug("Algorithm bundle starting against endpoint %s", wrapper_store.endpoint)
            await algorithm(wrapper_store, stop_evt)
            logger.debug("Algorithm bundle completed successfully")
+        except asyncio.CancelledError:
+            logger.info("Algorithm received CancelledError; signaling stop event")
+            stop_evt.set()
+            raise
        except KeyboardInterrupt:
            logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
            stop_evt.set()
@@ -179,6 +183,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                logger.debug("Runner %s executing with provided store", worker_id)
            await runner(client_store, worker_id, stop_evt)
            logger.debug("Runner %s completed successfully", worker_id)
+        except asyncio.CancelledError:
+            logger.debug("Runner %s received CancelledError; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
        except KeyboardInterrupt:
            logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
            stop_evt.set()
@@ -210,7 +218,13 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        def _runner_sync(runner: RunnerBundle, worker_id: int, store: LightningStore, stop_evt: ExecutionEvent) -> None:
            # Runners are executed in child processes; each process owns its own
            # event loop to keep the asyncio scheduler isolated.
-            asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
+            try:
+                asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
+            except KeyboardInterrupt:
+                logger.warning("Runner (asyncio) %s received KeyboardInterrupt; exiting gracefully", worker_id)
+            except BaseException as exc:
+                logger.exception("Runner (asyncio) %s crashed by %s; signaling stop event", worker_id, exc)
+                raise

        for i in range(self.n_runners):
            process = cast(
@@ -234,7 +248,13 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        """Used when `main_process == "runner"`."""

        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
-            asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+            try:
+                asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+            except KeyboardInterrupt:
+                logger.warning("Algorithm (asyncio.run) received KeyboardInterrupt; exiting gracefully")
+            except BaseException as exc:
+                logger.exception("Algorithm (asyncio.run) crashed by %s; signaling stop event", exc)
+                raise

        process = cast(
            multiprocessing.Process,
@@ -6,6 +6,7 @@ AGENTOPS_INSTALLED: bool = False
 AGENTOPS_LANGCHAIN_INSTALLED: bool = False
 LITELLM_INSTALLED: bool = False
 VLLM_INSTALLED: bool = False
+WEAVE_INSTALLED: bool = False

 try:
    from . import agentops  # type: ignore
@@ -38,6 +39,13 @@ try:
 except ImportError:
    pass

+try:
+    from . import weave  # type: ignore
+
+    WEAVE_INSTALLED = True  # type: ignore
+except ImportError:
+    pass
+

 def instrument_all():
    """Instrument all the instrumentation libraries."""
@@ -111,3 +119,20 @@ def uninstrument_all():
            warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
    else:
        warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")
+
+
+def instrument_weave():
+    if WEAVE_INSTALLED:
+        from .weave import instrument_weave
+
+        instrument_weave()
+
+
+def uninstrument_weave():
+    if WEAVE_INSTALLED:
+        try:
+            from .weave import uninstrument_weave
+
+            uninstrument_weave()
+        except ImportError:
+            warnings.warn("weave is installed but uninstrument_weave could not be imported.")
@@ -0,0 +1,139 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import os
+from typing import Any, Callable, Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "instrument_weave",
+    "uninstrument_weave",
+]
+
+# Module-level storage for originals
+_original_default_entity_name_getter: Callable[..., Any] | None = None
+_original_upsert_project_getter: Callable[..., Any] | None = None
+_original_weave_get = False
+_original_weave_post = False
+
+
+def instrument_weave():
+    """
+    Patch the Weave/W&B integration to bypass actual network calls for testing.
+
+    - Mocks HTTP POST/GET requests
+    - Patches wandb.Api methods
+    - Silences Weave logging
+    - Sets dummy WANDB_API_KEY if not provided
+    """
+    try:
+        import weave
+        from weave.compat import wandb  # type: ignore
+    except ImportError:
+        logger.warning("Weave or wandb not installed; cannot uninstrument.")
+        return
+
+    _weave_tracer_entity_name = "weave_tracer_entity"
+
+    def default_entity_name_getter(_self) -> str:  # type: ignore
+        return _weave_tracer_entity_name
+
+    def upsert_project_getter(
+        _self, project: str, description: Optional[str] = None, entity: Optional[str] = None  # type: ignore
+    ) -> dict[str, Any]:
+        return {
+            "upsertModel": {
+                "model": {
+                    "name": project,
+                    "description": description or "",
+                    "entity": entity or _weave_tracer_entity_name,
+                }
+            },
+            "project": "weave_tracer_project",
+        }
+
+    # Mock network requests to avoid real HTTP calls
+    def post(url: str, *args: Any, **kwargs: Any) -> requests.Response:
+        response = requests.Response()
+        response.status_code = 200
+        response._content = b'{"digest": "mocked_digest"}'
+        return response
+
+    def get(url: str, *args: Any, **kwargs: Any) -> requests.Response:
+        response = requests.Response()
+        response.status_code = 200
+        response._content = b'{"min_required_weave_python_version": "0.52.14"}'
+        return response
+
+    # Patch API methods and HTTP requests
+    global _original_default_entity_name_getter
+    global _original_upsert_project_getter
+    global _original_weave_post
+    global _original_weave_get
+    _original_default_entity_name_getter = wandb.Api.default_entity_name  # type: ignore
+    _original_upsert_project_getter = wandb.Api.upsert_project  # type: ignore
+    _original_weave_post = weave.utils.http_requests.session.post  # type: ignore
+    _original_weave_get = weave.utils.http_requests.session.get  # type: ignore
+
+    # Patch API methods and HTTP requests
+    wandb.Api.default_entity_name = default_entity_name_getter  # type: ignore
+    wandb.Api.upsert_project = upsert_project_getter  # type: ignore
+    weave.utils.http_requests.session.post = post  # type: ignore
+    weave.utils.http_requests.session.get = get  # type: ignore
+
+    # Silence Weave logging
+    for name in logging.root.manager.loggerDict:
+        if name.startswith("weave"):
+            logging.getLogger(name).disabled = True
+
+    # Set dummy API key if missing
+    if not os.environ.get("WANDB_API_KEY"):
+        os.environ["WANDB_API_KEY"] = "dumped_api_key_for_weave_tracer"
+
+    # if needed in future tests, enable this and replace WF_TRACE_SERVER_URL to local server
+    # full_url = f"http://127.0.0.1:{_port}"
+    # os.environ["WF_TRACE_SERVER_URL"] = full_url
+
+
+def uninstrument_weave():
+    """
+    Restore the original Weave/W&B integration methods and HTTP requests.
+    """
+    try:
+        import weave
+        from weave.compat import wandb  # type: ignore
+    except ImportError:
+        logger.warning("Weave or wandb not installed; cannot uninstrument.")
+        return
+
+    global _original_default_entity_name_getter
+    if _original_default_entity_name_getter is not None:
+        wandb.Api.default_entity_name = _original_default_entity_name_getter  # type: ignore
+        _original_default_entity_name_getter = None
+        logger.info("restored wandb.Api.default_entity_name")
+
+    global _original_upsert_project_getter
+    if _original_upsert_project_getter is not None:
+        wandb.Api.upsert_project = _original_upsert_project_getter  # type: ignore
+        _original_upsert_project_getter = None
+        logger.info("restored wandb.Api.upsert_project")
+
+    global _original_weave_post
+    if _original_weave_post is not None:
+        weave.utils.http_requests.session.post = _original_weave_post  # type: ignore
+        _original_weave_post = None
+        logger.info("restored weave.utils.http_requests.session.post")
+
+    global _original_weave_get
+    if _original_weave_get is not None:
+        weave.utils.http_requests.session.get = _original_weave_get  # type: ignore
+        _original_weave_get = None
+        logger.info("restored weave.utils.http_requests.session.get")
+
+    # Restore Weave logging
+    for name in logging.root.manager.loggerDict:
+        if name.startswith("weave"):
+            logging.getLogger(name).disabled = False
@@ -582,6 +582,7 @@ class LitAgentRunner(Runner[T_task]):
                while not (event is not None and event.is_set()):
                    logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
                    next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
+                    logger.debug(f"{self._log_prefix()} Next rollout retrieved: {next_rollout}")
                    if next_rollout is None:
                        logger.debug(
                            f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."
@@ -59,10 +59,12 @@ from agentlightning.types import (
    Worker,
    WorkerStatus,
 )
+from agentlightning.utils.metrics import MetricsBackend, get_prometheus_registry
 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, LightningStoreStatistics, Unset
+from .collection.base import resolve_error_type
 from .utils import LATENCY_BUCKETS

 server_logger = logging.getLogger("agentlightning.store.server")
@@ -238,7 +240,7 @@ class LightningStoreServer(LightningStore):
        launcher_args: The arguments to use for the server launcher.
            It's not allowed to set `host`, `port`, `launch_mode` together with `launcher_args`.
        n_workers: The number of workers to run in the server. Only applicable for `mp` launch mode.
-        prometheus: Whether to enable Prometheus metrics.
+        tracker: The metrics tracker to use for the server.
    """

    def __init__(
@@ -250,7 +252,7 @@ class LightningStoreServer(LightningStore):
        launch_mode: LaunchMode = "thread",
        launcher_args: PythonServerLauncherArgs | None = None,
        n_workers: int = 1,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
    ):
        super().__init__()
        self.store = store
@@ -268,6 +270,7 @@ class LightningStoreServer(LightningStore):
                port=port,
                launch_mode=launch_mode,
                healthcheck_url=API_V1_AGL_PREFIX + "/health",
+                n_workers=n_workers,
            )

        store_capabilities = self.store.capabilities
@@ -287,7 +290,7 @@ class LightningStoreServer(LightningStore):
            app=self.app,
            args=self.launcher_args,
        )
-        self._prometheus = prometheus
+        self._tracker = tracker

        self._lock: threading.Lock = threading.Lock()
        self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
@@ -332,7 +335,6 @@ class LightningStoreServer(LightningStore):
        return {
            "launcher_args": self.launcher_args,
            "server_launcher": self.server_launcher,
-            "_prometheus": self._prometheus,
            "_owner_pid": self._owner_pid,
        }

@@ -350,11 +352,12 @@ class LightningStoreServer(LightningStore):
        self.store = None
        self.launcher_args = state["launcher_args"]
        self.server_launcher = state["server_launcher"]
-        self._prometheus = state["_prometheus"]
+        self._tracker = None
        self._owner_pid = state["_owner_pid"]
        self._cors_allow_origins = state.get("_cors_allow_origins")
        self._client = None
        self._lock = threading.Lock()
+        self._prometheus_registry = None
        # Do NOT reconstruct app, _uvicorn_config, _uvicorn_server
        # to avoid transferring server state to subprocess

@@ -435,8 +438,8 @@ class LightningStoreServer(LightningStore):
        api = APIRouter(prefix=API_V1_PREFIX)

        # The outermost-layer of monitoring
-        if self._prometheus:
-            self._setup_prometheus(api=api, app=self.app)
+        if self._tracker is not None:
+            self._setup_metrics(api=api, app=self.app)

        # TODO: This should only be enabled in development mode.
        @self.app.middleware("http")
@@ -844,41 +847,21 @@ class LightningStoreServer(LightningStore):
        # Finally, mount the dashboard assets
        self._setup_dashboard()

-    def _setup_prometheus(self, api: APIRouter, app: FastAPI):
+    def _setup_metrics(self, api: APIRouter, app: FastAPI):
        """Setup Prometheus metrics endpoints."""
-        try:
-            from prometheus_client import make_asgi_app  # type: ignore
-            from prometheus_client import (
-                REGISTRY,
-                CollectorRegistry,
-                Counter,
-                Histogram,
-                multiprocess,
-            )
-        except ImportError:
-            raise ImportError(
-                "Prometheus client is not installed. Please either install it or set prometheus to False."
-            )
+        if self._tracker is None:
+            return

-        # 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"],
+        self._tracker.register_counter(
+            "agl.http.total",
+            ["path", "method", "status"],
+            group_level=2,
        )
-
-        HTTP_LATENCY = Histogram(
-            "http_request_duration_seconds",
-            "Latency of HTTP requests",
-            ["method", "path", "status_code"],
+        self._tracker.register_histogram(
+            "agl.http.latency",
+            ["path", "method", "status"],
            buckets=LATENCY_BUCKETS,
+            group_level=2,
        )

        def get_template_path(path: str) -> str:
@@ -904,9 +887,12 @@ class LightningStoreServer(LightningStore):
            return path

        @app.middleware("http")
-        async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
+        async def tracking_middleware(  # pyright: ignore[reportUnusedFunction]
            request: Request, call_next: Callable[[Request], Awaitable[Response]]
        ) -> Response:
+            if self._tracker is None:
+                return await call_next(request)
+
            start = time.perf_counter()
            status = 520  # Default to 520 if things crash hard

@@ -918,9 +904,8 @@ class LightningStoreServer(LightningStore):
                # 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
+            except Exception as exc:
+                status = resolve_error_type(exc)
                raise
            finally:
                # This block executes NO MATTER WHAT happens above
@@ -930,13 +915,25 @@ class LightningStoreServer(LightningStore):
                path = get_template_path(request.url.path)
                method = request.method

-                HTTP_REQUESTS.labels(method, path, status).inc()
-                HTTP_LATENCY.labels(method, path, status).observe(elapsed)
+                await self._tracker.inc_counter(
+                    "agl.http.total",
+                    labels={"method": method, "path": path, "status": str(status)},
+                )
+                await self._tracker.observe_histogram(
+                    "agl.http.latency",
+                    value=elapsed,
+                    labels={"method": method, "path": path, "status": str(status)},
+                )

-        metrics_app = make_asgi_app(registry=registry)  # type: ignore
+        if self._tracker.has_prometheus():
+            from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]

-        # This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
-        app.mount(api.prefix + "/prometheus", metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+            metrics_app = make_asgi_app(  # pyright: ignore[reportUnknownVariableType]
+                registry=get_prometheus_registry()
+            )
+
+            # 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."""
@@ -2,6 +2,10 @@

 from __future__ import annotations

+import functools
+import time
+from contextlib import asynccontextmanager
+from numbers import Real
 from typing import (
    TYPE_CHECKING,
    Any,
@@ -18,10 +22,14 @@ from typing import (
    Sequence,
    Tuple,
    Type,
+    TypeGuard,
    TypeVar,
    cast,
 )

+from agentlightning.store.utils import LATENCY_BUCKETS
+from agentlightning.utils.metrics import MetricsBackend
+
 if TYPE_CHECKING:
    from typing import Self

@@ -40,19 +48,146 @@ from agentlightning.types import (
 T = TypeVar("T")  # Recommended to be a BaseModel
 K = TypeVar("K")
 V = TypeVar("V")
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])

 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"]
+AtomicLabels = Literal[
+    "rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids", "generic"
+]
 """Labels for atomic operations.

 These labels are used to identify the collections that are affected by the atomic operation.
+
+The `generic` label is used to identify atomic operations that are not associated with any specific collection.
 """


-class Collection(Generic[T]):
-    """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
+def resolve_error_type(exc: BaseException | None) -> str:
+    if exc is None:
+        return "N/A"
+
+    try:
+        from .mongo import resolve_mongo_error_type
+
+        error_type = resolve_mongo_error_type(exc)
+        if error_type is not None:
+            return error_type
+    except ImportError:
+        # If the mongo backend is not available, fall back to using the exception's class name.
+        pass
+
+    return exc.__class__.__name__
+
+
+def tracked(operation: str):
+    """Decorator to track the execution of the decorated method."""
+
+    def decorator(func: T_callable) -> T_callable:
+
+        @functools.wraps(func)
+        async def wrapper(self: TrackedCollection, *args: Any, **kwargs: Any) -> Any:
+            async with self.tracking_context(operation, self.collection_name):
+                return await func(self, *args, **kwargs)
+
+        return cast(T_callable, wrapper)
+
+    return decorator
+
+
+def ensure_numeric(value: Any, *, description: str) -> TypeGuard[Real]:
+    """Validate that *value* behaves like a real number.
+
+    Returns true or crashes.
+    """
+
+    if isinstance(value, bool):
+        raise TypeError(f"{description} must be numeric; got bool")
+    if not isinstance(value, Real):
+        raise TypeError(f"{description} must be numeric; got {type(value).__name__}")
+    return True
+
+
+class DuplicatedPrimaryKeyError(ValueError):
+    """Error raised when a duplicate key is encountered."""
+
+    pass
+
+
+class TrackedCollection:
+    """An object that can be tracked by the metrics backend."""
+
+    def __init__(self, tracker: MetricsBackend | None = None):
+        self._tracker = tracker
+
+    @property
+    def tracker(self) -> MetricsBackend | None:
+        return self._tracker
+
+    @property
+    def collection_name(self) -> str:
+        """The identifier of the collection."""
+        raise NotImplementedError()
+
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, Any]:
+        """Extra labels to add to the tracking context."""
+        return {}
+
+    @asynccontextmanager
+    async def tracking_context(self, operation: str, collection: str):
+        """Context manager to track the execution of the decorated method.
+
+        Args:
+            operation: The operation to track.
+            collection: The collection to track.
+        """
+        if self._tracker is None:
+            # no-op context manager
+            yield
+
+        else:
+            from agentlightning.store.collection_based import nearest_lightning_store_method_from_stack
+
+            # Enable tracking
+            start_time = time.perf_counter()
+            status: str = "OK"
+            public_store_method, private_store_method = nearest_lightning_store_method_from_stack()
+            try:
+                yield
+            except BaseException as exc:
+                status = resolve_error_type(exc)
+                raise
+            finally:
+                elapsed = time.perf_counter() - start_time
+                await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
+                    "agl.collections.total",
+                    labels={
+                        "store_pubmeth": public_store_method,
+                        "store_privmeth": private_store_method,
+                        "operation": operation,
+                        "collection": collection,
+                        "status": status,
+                        **self.extra_tracking_labels,
+                    },
+                )
+                await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
+                    "agl.collections.latency",
+                    value=elapsed,
+                    labels={
+                        "store_pubmeth": public_store_method,
+                        "store_privmeth": private_store_method,
+                        "operation": operation,
+                        "collection": collection,
+                        "status": status,
+                        **self.extra_tracking_labels,
+                    },
+                )
+
+
+class Collection(TrackedCollection, Generic[T]):
+    """Standard collection interface. Behaves like a list of items. Supporting addition, updating, and deletion of items."""

    def primary_keys(self) -> Sequence[str]:
        """Get the primary keys of the collection."""
@@ -174,7 +309,7 @@ class Collection(Generic[T]):
        raise NotImplementedError()


-class Queue(Generic[T]):
+class Queue(TrackedCollection, Generic[T]):
    """Behaves like a deque. Supporting appending items to the end and popping items from the front."""

    def __repr__(self) -> str:
@@ -228,7 +363,7 @@ class Queue(Generic[T]):
        raise NotImplementedError()


-class KeyValue(Generic[K, V]):
+class KeyValue(TrackedCollection, Generic[K, V]):
    """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""

    def __repr__(self) -> str:
@@ -246,6 +381,22 @@ class KeyValue(Generic[K, V]):
        """Set the value for the given key."""
        raise NotImplementedError()

+    async def inc(self, key: K, amount: V) -> V:
+        """Increase the numeric value for the given key by `amount` and return the new value.
+
+        Raises:
+            TypeError: If the existing value or `amount` is not numeric.
+        """
+        raise NotImplementedError()
+
+    async def chmax(self, key: K, value: V) -> V:
+        """Set the value for the given key to the maximum of the current and new value.
+
+        Raises:
+            TypeError: If the existing value or `value` is not numeric.
+        """
+        raise NotImplementedError()
+
    async def pop(self, key: K, default: V | None = None) -> V | None:
        """Pop the value for the given key, or the default value if the key is not found."""
        raise NotImplementedError()
@@ -255,13 +406,35 @@ class KeyValue(Generic[K, V]):
        raise NotImplementedError()


-class LightningCollections:
+class LightningCollections(TrackedCollection):
    """Collections of rollouts, attempts, spans, resources, and workers.

    [LightningStore][agentlightning.LightningStore] implementations can use this as a storage base
    to implement the store API.
    """

+    def __init__(self, tracker: MetricsBackend | None = None, extra_labels: Optional[Sequence[str]] = None):
+        super().__init__(tracker=tracker)
+        self.register_collection_metrics(extra_labels)
+
+    def register_collection_metrics(self, extra_labels: Optional[Sequence[str]] = None) -> None:
+        if self._tracker is None:
+            return
+        labels = ["store_pubmeth", "operation", "collection", "store_privmeth", "status"]
+        if extra_labels is not None:
+            labels.extend(extra_labels)
+        self._tracker.register_histogram(
+            "agl.collections.latency",
+            labels,
+            buckets=LATENCY_BUCKETS,
+            group_level=2,
+        )
+        self._tracker.register_counter("agl.collections.total", labels, group_level=2)
+
+    @property
+    def tracker(self) -> MetricsBackend | None:
+        return self._tracker
+
    @property
    def rollouts(self) -> Collection[Rollout]:
        """Collections of rollouts."""
@@ -4,7 +4,7 @@ from __future__ import annotations

 import asyncio
 import logging
-import time
+import uuid
 import weakref
 from collections import deque
 from contextlib import AsyncExitStack, asynccontextmanager
@@ -23,12 +23,12 @@ from typing import (
    Type,
    TypeVar,
    Union,
+    cast,
 )

 import aiologic
 from pydantic import BaseModel

-from agentlightning.store.utils import LATENCY_BUCKETS
 from agentlightning.types import (
    Attempt,
    FilterField,
@@ -40,16 +40,21 @@ from agentlightning.types import (
    Span,
    Worker,
 )
+from agentlightning.utils.metrics import MetricsBackend

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

 T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
@@ -192,10 +197,19 @@ class ListBasedCollection(Collection[T]):
       if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
    """

-    def __init__(self, items: List[T], item_type: Type[T], primary_keys: Sequence[str]):
+    def __init__(
+        self,
+        items: List[T],
+        item_type: Type[T],
+        primary_keys: Sequence[str],
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
        if not primary_keys:
            raise ValueError("primary_keys must be non-empty")

+        self._id = id if id is not None else str(uuid.uuid4())
        self._items: Dict[Any, Any] = {}
        self._size: int = 0
        if issubclass(item_type, dict):
@@ -207,6 +221,10 @@ class ListBasedCollection(Collection[T]):
        for item in items or []:
            self._mutate_single(item, mode="insert")

+    @property
+    def collection_name(self) -> str:
+        return self._id
+
    def primary_keys(self) -> Sequence[str]:
        """Return the primary key field names for this collection."""
        return self._primary_keys
@@ -299,7 +317,9 @@ class ListBasedCollection(Collection[T]):

            if mode == "insert":
                if exists:
-                    raise ValueError(f"Item already exists with primary key(s): {self._render_key_values(key_values)}")
+                    raise DuplicatedPrimaryKeyError(
+                        f"Item already exists with primary key(s): {self._render_key_values(key_values)}"
+                    )
                parent[final_key] = item
                self._size += 1
            else:  # upsert
@@ -483,6 +503,7 @@ class ListBasedCollection(Collection[T]):
            # No items exist for this primary-key prefix.
            return ()

+    @tracked("query")
    async def query(
        self,
        filter: Optional[FilterOptions] = None,
@@ -546,6 +567,7 @@ class ListBasedCollection(Collection[T]):
            total=total_matched,
        )

+    @tracked("get")
    async def get(
        self,
        filter: Optional[FilterOptions] = None,
@@ -582,11 +604,12 @@ class ListBasedCollection(Collection[T]):

        return best_item

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

        Raises:
-            ValueError: If any item with the same primary keys already exists.
+            DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
        """
        seen_keys: set[Tuple[Any, ...]] = set()
        prepared: List[T] = []
@@ -594,8 +617,8 @@ class ListBasedCollection(Collection[T]):
            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)}"
+                raise DuplicatedPrimaryKeyError(
+                    f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
                )
            seen_keys.add(key_values)
            prepared.append(item)
@@ -603,6 +626,7 @@ class ListBasedCollection(Collection[T]):
        for item in prepared:
            self._mutate_single(item, mode="insert")

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

@@ -617,6 +641,7 @@ class ListBasedCollection(Collection[T]):
            updated_items.append(updated)
        return updated_items

+    @tracked("upsert")
    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
        """Upsert the given items (insert if missing, otherwise update)."""
        upserted_items: List[T] = []
@@ -627,6 +652,7 @@ class ListBasedCollection(Collection[T]):
            upserted_items.append(upserted)
        return upserted_items

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

@@ -646,23 +672,37 @@ class DequeBasedQueue(Queue[T]):
    Provides O(1) amortized enqueue (append) and dequeue (popleft).
    """

-    def __init__(self, item_type: Type[T], items: Optional[Sequence[T]] = None):
+    def __init__(
+        self,
+        item_type: Type[T],
+        items: Optional[Sequence[T]] = None,
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
        self._items: Deque[T] = deque()
        self._item_type: Type[T] = item_type
+        self._id = id if id is not None else str(uuid.uuid4())
        if items:
            self._items.extend(items)

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

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

+    @tracked("has")
    async def has(self, item: T) -> bool:
        if not isinstance(item, self._item_type):
            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
        return item in self._items

+    @tracked("enqueue")
    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
        for item in items:
            if not isinstance(item, self._item_type):
@@ -670,6 +710,7 @@ class DequeBasedQueue(Queue[T]):
            self._items.append(item)
        return items

+    @tracked("dequeue")
    async def dequeue(self, limit: int = 1) -> Sequence[T]:
        if limit <= 0:
            return []
@@ -678,6 +719,7 @@ class DequeBasedQueue(Queue[T]):
            out.append(self._items.popleft())
        return out

+    @tracked("peek")
    async def peek(self, limit: int = 1) -> Sequence[T]:
        if limit <= 0:
            return []
@@ -689,6 +731,7 @@ class DequeBasedQueue(Queue[T]):
            result.append(item)
        return result

+    @tracked("size")
    async def size(self) -> int:
        return len(self._items)

@@ -696,21 +739,61 @@ class DequeBasedQueue(Queue[T]):
 class DictBasedKeyValue(KeyValue[K, V]):
    """KeyValue implementation backed by a plain dictionary."""

-    def __init__(self, data: Optional[Mapping[K, V]] = None):
+    def __init__(
+        self, data: Optional[Mapping[K, V]] = None, id: Optional[str] = None, tracker: Optional[MetricsBackend] = None
+    ):
+        super().__init__(tracker=tracker)
        self._values: Dict[K, V] = dict(data) if data else {}
+        self._id = id if id is not None else str(uuid.uuid4())

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

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

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

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

+    @tracked("size")
    async def size(self) -> int:
        return len(self._values)

@@ -721,8 +804,9 @@ class InMemoryLightningCollections(LightningCollections):
    Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
    """

-    def __init__(self, lock_type: Literal["thread", "asyncio"], prometheus: bool = False):
-        self._lock = {
+    def __init__(self, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
+        super().__init__(tracker=tracker)
+        self._lock: Mapping[AtomicLabels, _LoopAwareAsyncLock | _ThreadSafeAsyncLock] = {
            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
@@ -730,32 +814,31 @@ class InMemoryLightningCollections(LightningCollections):
            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
        }
-        self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
-        self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
-        self._spans = ListBasedCollection(
-            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"]
+        self._rollouts = ListBasedCollection(
+            items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
        )
-        self._resources = ListBasedCollection(items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"])
-        self._workers = ListBasedCollection(items=[], item_type=Worker, primary_keys=["worker_id"])
-        self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
-        self._span_sequence_ids = DictBasedKeyValue[str, int](data={})  # rollout_id -> sequence_id
+        self._attempts = ListBasedCollection(
+            items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"], id="attempts", tracker=tracker
+        )
+        self._spans = ListBasedCollection(
+            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"], id="spans", tracker=tracker
+        )
+        self._resources = ListBasedCollection(
+            items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"], id="resources", tracker=tracker
+        )
+        self._workers = ListBasedCollection(
+            items=[], item_type=Worker, primary_keys=["worker_id"], id="workers", tracker=tracker
+        )
+        self._rollout_queue = DequeBasedQueue(items=[], item_type=str, id="rollout_queue", tracker=tracker)
+        self._span_sequence_ids = DictBasedKeyValue[str, int](
+            data={}, id="span_sequence_ids", tracker=tracker
+        )  # rollout_id -> sequence_id

-        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 collection_name(self) -> str:
+        return "router"

    @property
    def rollouts(self) -> ListBasedCollection[Rollout]:
@@ -787,7 +870,12 @@ class InMemoryLightningCollections(LightningCollections):

    @asynccontextmanager
    async def atomic(
-        self, *, mode: AtomicMode = "rw", snapshot: bool = False, labels: Optional[Sequence[str]] = None, **kwargs: Any
+        self,
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
    ):
        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.

@@ -807,17 +895,15 @@ class InMemoryLightningCollections(LightningCollections):
        # 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 with self.tracking_context(operation="atomic", collection=self.collection_name):
+            managers = [(label, self._lock[label]) for label in labels]
+            async with AsyncExitStack() as stack:
+                for label, manager in managers:
+                    async with self.tracking_context(operation="lock", collection=label):
+                        await stack.enter_async_context(manager)
+                yield self

+    @tracked("evict_spans_for_rollout")
    async def evict_spans_for_rollout(self, rollout_id: str) -> None:
        """Evict all spans for a given rollout ID.

@@ -16,6 +16,7 @@ from __future__ import annotations
 import asyncio
 import functools
 import hashlib
+import inspect
 import logging
 import time
 import uuid
@@ -60,6 +61,7 @@ from agentlightning.types import (
    Worker,
    WorkerStatus,
 )
+from agentlightning.utils.metrics import MetricsBackend

 from .base import (
    UNSET,
@@ -71,7 +73,7 @@ from .base import (
    is_queuing,
 )
 from .collection import FilterOptions, LightningCollections
-from .collection.base import AtomicLabels
+from .collection.base import AtomicLabels, DuplicatedPrimaryKeyError
 from .utils import LATENCY_BUCKETS, rollout_status_from_attempt, scan_unhealthy_rollouts

 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -119,34 +121,42 @@ def _with_collections_execute(labels: Sequence[AtomicLabels]):
 def tracked(name: str):
    """Decorator to track the execution of the decorated method with Prometheus."""

-    _public_methods = frozenset([name for name in LightningStore.__dict__ if not name.startswith("_")])
-
    def decorator(func: T_callable) -> T_callable:

        @functools.wraps(func)
        async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
-            # For backtracking in Mongo-collection methods.
-            # Only track the public methods (+healthcheck)
-            if name in _public_methods or name == "_healthcheck":
-                method_name = name  # pyright: ignore[reportUnusedVariable]
-            else:
-                method_name = None  # pyright: ignore[reportUnusedVariable]
+            # Backtracking where this method comes from
+            public_meth_in_stack, _ = nearest_lightning_store_method_from_stack()

-            if not self._prometheus:  # pyright: ignore[reportPrivateUsage]
+            # For backtracking in collection methods.
+            # Only track the public methods (+healthcheck)
+            if name in COLLECTION_STORE_PUBLIC_METHODS:
+                public_method_name = name  # pyright: ignore[reportUnusedVariable]
+                public_meth_in_stack = name  # We are in a public method already.
+            if name in COLLECTION_STORE_ALL_METHODS:
+                private_method_name = name  # pyright: ignore[reportUnusedVariable]
+
+            if self._tracker is None:  # pyright: ignore[reportPrivateUsage]
                # Skip the tracking because tracking is not configured
                return await func(self, *args, **kwargs)

            start_time = time.perf_counter()
+            status: str = "OK"
            try:
-                ret = await func(self, *args, **kwargs)
-                self._total_metric.labels(name, "OK").inc()  # pyright: ignore[reportPrivateUsage]
-                return ret
-            except Exception as exc:
-                self._total_metric.labels(name, exc.__class__.__name__).inc()  # pyright: ignore[reportPrivateUsage]
+                return await func(self, *args, **kwargs)
+            except BaseException as exc:
+                status = exc.__class__.__name__
                raise
            finally:
                elapsed = time.perf_counter() - start_time
-                self._latency_metric.labels(name).observe(elapsed)  # pyright: ignore[reportPrivateUsage]
+                await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
+                    "agl.store.total", labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status}
+                )
+                await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
+                    "agl.store.latency",
+                    value=elapsed,
+                    labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
+                )

        return cast(T_callable, wrapper)

@@ -215,40 +225,56 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        read_snapshot: Make sure read operations are atomic. If set to true,
            all read operations like `query_rollouts` will have better consistency.
            It may use an isolated snapshot that supports repeatable reads.
-        prometheus: Enable Prometheus tracking.
+        tracker: Enable metrics tracking.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing. The debounce is a non-perfect traffic control.
+            It's isolated for each store instance if there are multiple worker replicas.
    """

-    def __init__(self, collections: T_collections, *, read_snapshot: bool = False, prometheus: bool = False):
+    def __init__(
+        self,
+        collections: T_collections,
+        *,
+        read_snapshot: bool = False,
+        tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
+    ) -> None:
        # rollouts and spans' storage
        self.collections = collections
        self._read_snapshot = read_snapshot
-        self._prometheus = prometheus
+        self._tracker = tracker
        self._launch_time = time.time()

-        if prometheus:
-            from prometheus_client import Counter, Histogram
+        # Control scan debounce to avoid overloading the store.
+        self._scan_debounce_seconds = scan_debounce_seconds
+        last_scan_time = self._launch_time
+        if self._scan_debounce_seconds > 0:
+            # Allow the first scan immediately after instantiation
+            last_scan_time -= self._scan_debounce_seconds
+        self._last_scan_entrance_time = last_scan_time

-            self._latency_metric = Histogram(
-                "collection_store_latency_seconds",
-                "Latency of CollectionBasedLightningStore methods",
-                ["method"],
+        if self._tracker is not None:
+            self._tracker.register_histogram(
+                "agl.store.latency",
+                ["method", "store_pubmeth", "status"],
                buckets=LATENCY_BUCKETS,
+                group_level=1,
            )
-            self._total_metric = Counter(
-                "collection_store_total",
-                "Total MongoDB operations",
-                ["method", "error_type"],
+            self._tracker.register_counter(
+                "agl.store.total",
+                ["method", "store_pubmeth", "status"],
+                group_level=1,
            )
-            self._rollout_counter = Counter(
-                "collection_store_rollout_total",
-                "Total rollouts",
+            self._tracker.register_counter(
+                "agl.rollouts.total",
                ["status", "mode"],
+                group_level=1,
            )
-            self._rollout_duration_metric = Histogram(
-                "collection_store_rollout_duration_seconds",
-                "Duration of rollouts",
+            self._tracker.register_histogram(
+                "agl.rollouts.duration",
                ["status", "mode"],
                buckets=LATENCY_BUCKETS,
+                group_level=1,
            )

    async def statistics(self) -> LightningStoreStatistics:
@@ -610,6 +636,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
            if not dequeued:
                break
            rollout_id = dequeued[0]
+            logger.debug("Rollout ID %s has been dequeued by Worker ID %s", rollout_id, worker_id)

            post_dequeue_result = await self._post_dequeue_rollouts([rollout_id], worker_id)
            if post_dequeue_result:
@@ -617,6 +644,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                attempted_rollout, _ = post_dequeue_result[0]
                if worker_id is not None:
                    await self._sync_workers_with_attempts([attempted_rollout.attempt], dequeue=True)
+                logger.debug("Rollout has been prepared for Worker ID %s: %s", worker_id, attempted_rollout)
                return attempted_rollout

            # else continue the loop
@@ -967,37 +995,36 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        return await self._get_latest_resources()

    @tracked("_issue_many_span_sequence_ids")
-    @_with_collections_execute(labels=["span_sequence_ids"])
-    async def _issue_many_span_sequence_ids(self, collections: T_collections, rollout_ids: List[str]) -> List[int]:
+    async def _issue_many_span_sequence_ids(self, rollout_ids: List[str]) -> List[int]:
        """Issue a new span sequence ID for a given rollout."""
-        # Cache the next sequence IDs for the rollouts (for both RW)
-        next_sequence_ids_cache: Dict[str, int] = {}
+        if not rollout_ids:
+            return []
+
+        request_counts: Dict[str, int] = defaultdict(int)
+        for rollout_id in rollout_ids:
+            request_counts[rollout_id] += 1
+
+        latest_values: Dict[str, int] = {}
+        for rollout_id, count in request_counts.items():
+            async with self.collections.atomic(mode="rw", snapshot=False, labels=["span_sequence_ids"]) as collections:
+                latest_values[rollout_id] = await collections.span_sequence_ids.inc(rollout_id, count)
+
+        next_value_tracker: Dict[str, int] = {
+            rollout_id: latest_values[rollout_id] - request_counts[rollout_id] for rollout_id in request_counts
+        }
+
        result: List[int] = []
        for rollout_id in rollout_ids:
-            if rollout_id not in next_sequence_ids_cache:
-                retrieved_id = await collections.span_sequence_ids.get(rollout_id)
-                if retrieved_id is None:
-                    retrieved_id = 0
-                next_sequence_ids_cache[rollout_id] = retrieved_id
-
-            # Increment the sequence ID for the rollout
-            next_sequence_ids_cache[rollout_id] += 1
-            result.append(next_sequence_ids_cache[rollout_id])
-
-        # Propagate the cache to storage
-        for rollout_id, sequence_id in next_sequence_ids_cache.items():
-            await collections.span_sequence_ids.set(rollout_id, sequence_id)
+            next_value_tracker[rollout_id] += 1
+            result.append(next_value_tracker[rollout_id])

        return result

    @tracked("_sync_span_sequence_id")
-    @_with_collections_execute(labels=["span_sequence_ids"])
-    async def _sync_span_sequence_id(self, collections: T_collections, rollout_id: str, sequence_id: int) -> None:
+    async def _sync_span_sequence_id(self, rollout_id: str, sequence_id: int) -> None:
        """Sync the span sequence ID for a given rollout from the input span sequence ID."""
-        existing_sequence_id = await collections.span_sequence_ids.get(rollout_id)
-        if existing_sequence_id is None:
-            existing_sequence_id = 0
-        await collections.span_sequence_ids.set(rollout_id, max(existing_sequence_id, sequence_id))
+        async with self.collections.atomic(mode="rw", snapshot=False, labels=["span_sequence_ids"]) as collections:
+            await collections.span_sequence_ids.chmax(rollout_id, sequence_id)

    @tracked("get_next_span_sequence_id")
    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
@@ -1079,13 +1106,11 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
            try:
                await collections.spans.insert([span])
                return True
-            except ValueError as e:
-                if "already exists" in str(e) or "contains duplicate" in str(e):
-                    logger.error(
-                        f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
-                    )
-                    return False
-                raise
+            except DuplicatedPrimaryKeyError:
+                logger.error(
+                    f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
+                )
+                return False

        successful_spans: List[Span] = []
        try:
@@ -1093,22 +1118,20 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
            async with self.collections.atomic(
                mode="w", snapshot=self._read_snapshot, commit=False, labels=["spans"]
            ) as collections:
+                # FIXME: Part of the insertion might complete though the full operation fails.
+                # In that case, the "insert spans" return values might not be accurate.
                await collections.spans.insert(spans)
            successful_spans.extend(spans)
-        except ValueError as e:
-            if "already exists" in str(e) or "contains duplicate" in str(e):
-                # There is a duplicate span, we warn it
-                # We fallback to adding the spans one by one
-                async def _add_many_spans_fallback(collections: T_collections):
-                    for span in spans:
-                        if await _add_span_fallback(collections, span):
-                            successful_spans.append(span)
-
-                await self.collections.execute(
-                    _add_many_spans_fallback, mode="w", snapshot=self._read_snapshot, commit=True, labels=["spans"]
-                )
-            else:
-                raise
+        except DuplicatedPrimaryKeyError:
+            # There is a duplicate span, we warn it
+            # We fallback to adding the spans one by one
+            for span in spans:
+                async with self.collections.atomic(
+                    mode="w", snapshot=self._read_snapshot, labels=["spans"]
+                ) as collections:
+                    # No need to commit here, it will be simple atomic write operations
+                    if await _add_span_fallback(collections, span):
+                        successful_spans.append(span)

        return successful_spans

@@ -1140,6 +1163,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        if successful_spans:
            await self._post_add_spans(successful_spans, rollout_id, attempt_id)

+        logger.debug("Added %d spans for rollout %s, attempt %s", len(successful_spans), rollout_id, attempt_id)
+
        return successful_spans

    @tracked("_post_add_spans")
@@ -1158,48 +1183,46 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        if not spans:
            return

-        async def _update_rollout_attempt(collections: T_collections) -> Optional[Tuple[Rollout, Sequence[str]]]:
-            attempt = await collections.attempts.get(
-                {"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
-            )
-            if attempt is None:
-                return None
-            rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
-            if rollout is None:
-                return None
-
-            # Update attempt heartbeat and ensure persistence
-            attempt.last_heartbeat_time = time.time()
-            if attempt.status in ["preparing", "unresponsive"]:
-                attempt.status = "running"
-            await collections.attempts.update([attempt], update_fields=["last_heartbeat_time", "status"])
-
-            # If the status has already timed out or failed, do not change it (but heartbeat is still recorded)
-
-            # Update rollout status if it's the latest attempt
-            rollout_updated: bool = False
-            updated_fields: List[str] = []
-            latest_attempt = await self._unlocked_get_latest_attempt(collections, rollout.rollout_id)
-            if latest_attempt is not None and attempt.attempt_id == latest_attempt.attempt_id:
-                if rollout.status in ["preparing", "queueing", "requeuing"]:
-                    # If rollout is currently preparing or queuing, set it to running
-                    rollout.status = "running"
-                    await collections.rollouts.update([rollout], update_fields=["status"])
-                    rollout_updated = True
-                    updated_fields = ["status"]
-                # Otherwise, the rollout has succeeded or failed, do nothing
-            return (rollout, updated_fields) if rollout_updated else None
-
-        rollout_update = await self.collections.execute(
-            _update_rollout_attempt,
-            mode="rw",
-            snapshot=self._read_snapshot,
-            commit=True,
-            labels=["rollouts", "attempts"],
-        )
+        rollout_update = await self._on_attempt_heartbeat(rollout_id=rollout_id, attempt_id=attempt_id)
        if rollout_update is not None:
            await self._post_update_rollout([rollout_update])

+    @tracked("_on_attempt_heartbeat")
+    @_with_collections_execute(labels=["rollouts", "attempts"])
+    async def _on_attempt_heartbeat(
+        self, collections: T_collections, rollout_id: str, attempt_id: str
+    ) -> Optional[Tuple[Rollout, Sequence[str]]]:
+        attempt = await collections.attempts.get(
+            {"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
+        )
+        if attempt is None:
+            return None
+        rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+        if rollout is None:
+            return None
+
+        # Update attempt heartbeat and ensure persistence
+        attempt.last_heartbeat_time = time.time()
+        if attempt.status in ["preparing", "unresponsive"]:
+            attempt.status = "running"
+        await collections.attempts.update([attempt], update_fields=["last_heartbeat_time", "status"])
+
+        # If the status has already timed out or failed, do not change it (but heartbeat is still recorded)
+
+        # Update rollout status if it's the latest attempt
+        rollout_updated: bool = False
+        updated_fields: List[str] = []
+        latest_attempt = await self._unlocked_get_latest_attempt(collections, rollout.rollout_id)
+        if latest_attempt is not None and attempt.attempt_id == latest_attempt.attempt_id:
+            if rollout.status in ["preparing", "queueing", "requeuing"]:
+                # If rollout is currently preparing or queuing, set it to running
+                rollout.status = "running"
+                await collections.rollouts.update([rollout], update_fields=["status"])
+                rollout_updated = True
+                updated_fields = ["status"]
+            # Otherwise, the rollout has succeeded or failed, do nothing
+        return (rollout, updated_fields) if rollout_updated else None
+
    @tracked("wait_for_rollouts")
    @healthcheck_before
    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
@@ -1220,7 +1243,17 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                logger.error(f"Error waiting for rollout {rollout_id}: {rollout}")

        # Filter out the exceptions
-        return [rollout for rollout in rollouts if isinstance(rollout, Rollout)]
+        ret = [rollout for rollout in rollouts if isinstance(rollout, Rollout)]
+        finished_rollout_ids = set([rollout.rollout_id for rollout in ret])
+        unfinished_rollout_ids = set(rollout_ids) - finished_rollout_ids
+        logger.debug(
+            "Waiting for rollouts. Number of finished rollouts: %d; number of unfinished rollouts: %d",
+            len(finished_rollout_ids),
+            len(unfinished_rollout_ids),
+        )
+        if len(unfinished_rollout_ids) < 30:
+            logger.debug("Unfinished rollouts: %s", unfinished_rollout_ids)
+        return ret

    @tracked("wait_for_rollout")
    async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
@@ -1467,10 +1500,17 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        for rollout, updated_fields in rollouts:
            # Sometimes "end_time" is set but it's not really updated.
            if "end_time" in updated_fields and is_finished(rollout):
-                if self._prometheus:
-                    self._rollout_counter.labels(rollout.status, rollout.mode).inc()
-                    self._rollout_duration_metric.labels(rollout.status, rollout.mode).observe(
-                        cast(float, rollout.end_time) - rollout.start_time
+                if self._tracker is not None:
+                    labels = {
+                        "status": rollout.status,
+                        "mode": rollout.mode if rollout.mode is not None else "unknown",
+                    }
+                    duration = cast(float, rollout.end_time) - rollout.start_time
+                    await self._tracker.inc_counter("agl.rollouts.total", labels=labels)
+                    await self._tracker.observe_histogram(
+                        "agl.rollouts.duration",
+                        value=duration,
+                        labels=labels,
                    )

        if not skip_enqueue:
@@ -1644,6 +1684,9 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
    @tracked("_scan_for_unhealthy_rollouts")
    async def _scan_for_unhealthy_rollouts(self) -> None:
        """Perform healthcheck against all running rollouts in the store."""
+        if not await self._should_scan_for_unhealthy_rollouts():
+            return
+
        rollouts, attempts_sync_required = await self._find_and_update_unhealthy_rollouts()

        if rollouts:
@@ -1653,6 +1696,25 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
        if attempts_sync_required:
            await self._sync_workers_with_attempts(attempts_sync_required)

+    @tracked("_should_scan_for_unhealthy_rollouts")
+    async def _should_scan_for_unhealthy_rollouts(self) -> bool:
+        """Check if the scan for unhealthy rollouts should be performed."""
+        if self._scan_debounce_seconds <= 0:
+            return True
+
+        now = time.time()
+        should_scan = now - self._last_scan_entrance_time >= self._scan_debounce_seconds
+        if not should_scan:
+            return False
+
+        # Someone else may be racing for the same scan. Double-check inside the lock.
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["generic"]):
+            now = time.time()
+            if now - self._last_scan_entrance_time < self._scan_debounce_seconds:
+                return False
+            self._last_scan_entrance_time = now
+            return True
+
    @tracked("_find_and_update_unhealthy_rollouts")
    @_with_collections_execute(labels=["rollouts", "attempts"])
    async def _find_and_update_unhealthy_rollouts(
@@ -1681,3 +1743,50 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
            if worker_sync_required:
                attempts.append(attempt)
        return rollouts, attempts
+
+
+# _scan_for_unhealthy_rollouts is somehow standalone and automatically invoked.
+COLLECTION_STORE_PUBLIC_METHODS = frozenset(
+    [name for name in LightningStore.__dict__ if not name.startswith("_")] + ["_scan_for_unhealthy_rollouts"]
+)
+
+COLLECTION_STORE_ALL_METHODS = frozenset([name for name in CollectionBasedLightningStore.__dict__])
+
+_UNKNOWN_STORE_METHOD = "unknown"
+
+
+def nearest_lightning_store_method_from_stack() -> Tuple[str, str]:
+    """Stack introspection so that we capture the nearest public API method from the
+    call stack whenever metrics are recorded.
+
+    Returns:
+        A tuple of public method name and nearest private method name.
+    """
+    frame = inspect.currentframe()
+    final_public_method_name = final_private_method_name = _UNKNOWN_STORE_METHOD
+    try:
+        if frame is not None:
+            frame = frame.f_back
+            while frame is not None:
+                self_obj = frame.f_locals.get("self")
+                public_method_name = frame.f_locals.get("public_method_name")
+                private_method_name = frame.f_locals.get("private_method_name")
+                if (
+                    final_public_method_name == _UNKNOWN_STORE_METHOD
+                    and public_method_name in COLLECTION_STORE_PUBLIC_METHODS
+                    and isinstance(self_obj, LightningStore)
+                ):
+                    final_public_method_name = public_method_name
+                if (
+                    final_private_method_name == _UNKNOWN_STORE_METHOD
+                    and private_method_name in COLLECTION_STORE_ALL_METHODS
+                    and isinstance(self_obj, LightningStore)
+                ):
+                    final_private_method_name = private_method_name
+                frame = frame.f_back
+    except Exception as exc:
+        logger.debug("Error during stack introspection for LightningStore method: %s", exc)
+    finally:
+        del frame
+
+    return final_public_method_name, final_private_method_name
@@ -28,6 +28,7 @@ import aiologic
 from pydantic import BaseModel

 from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
+from agentlightning.utils.metrics import MetricsBackend

 from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
@@ -72,12 +73,16 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
    Thread-safe and async-compatible but data is not persistent.

    Args:
+        thread_safe: Whether the store is thread-safe.
        eviction_memory_threshold: The threshold for evicting spans in bytes.
            By default, it's 70% of the total VRAM available.
        safe_memory_threshold: The threshold for safe memory usage in bytes.
            By default, it's 80% of the eviction threshold.
        span_size_estimator: A function to estimate the size of a span in bytes.
            By default, it's a simple size estimator that uses sys.getsizeof.
+        tracker: The metrics tracker to use.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing.
    """

    def __init__(
@@ -87,13 +92,13 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
        eviction_memory_threshold: float | int | None = None,
        safe_memory_threshold: float | int | None = None,
        span_size_estimator: Callable[[Span], int] | None = None,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
    ):
        super().__init__(
-            collections=InMemoryLightningCollections(
-                lock_type="thread" if thread_safe else "asyncio", prometheus=prometheus
-            ),
-            prometheus=prometheus,
+            collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio", tracker=tracker),
+            tracker=tracker,
+            scan_debounce_seconds=scan_debounce_seconds,
        )

        self._thread_safe = thread_safe
@@ -7,24 +7,13 @@ 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 typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TypeVar, Union

 from agentlightning.types import Attempt, AttemptedRollout, Rollout
+from agentlightning.utils.metrics import MetricsBackend

 from .base import LightningStoreCapabilities, is_finished
-from .collection.mongo import MongoClientPool, MongoLightningCollections, MongoOperationPrometheusTracker
+from .collection.mongo import MongoClientPool, MongoLightningCollections
 from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked

 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -42,27 +31,28 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
    Data is persistent and can be shared between multiple processes.

    Args:
-        client: The MongoDB client. Could be a string URI or an instance of AsyncMongoClient.
-        database: The MongoDB database. Could be a string name or an instance of AsyncDatabase.
-            You must provide at least one of client or database.
+        mongo_uri: MongoDB connection string (defaults to local replica set).
+        mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
+        database: The MongoDB database name. Defaults to ``agentlightning``.
        partition_id: The partition id. Useful when sharing the database among multiple Agent-lightning trainers.
+        tracker: The metrics tracker to use.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing.
    """

    def __init__(
        self,
        *,
-        client: AsyncMongoClient[Mapping[str, Any]] | str,
+        mongo_uri: str = "mongodb://localhost:27017/?replicaSet=rs0",
+        mongo_client_kwargs: Mapping[str, Any] | None = None,
        database_name: str | None = None,
        partition_id: str | None = None,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
    ) -> None:
-        self._enable_prometheus = prometheus
-        self._auto_created_client = False
-        if isinstance(client, str):
-            self._client = AsyncMongoClient[Mapping[str, Any]](client)
-            self._auto_created_client = True
-        else:
-            self._client = client
+        self._mongo_uri = mongo_uri
+        self._mongo_client_kwargs = dict(mongo_client_kwargs or {})
+
        if database_name is None:
            database_name = "agentlightning"
            logger.info("No database name provided, using default 'agentlightning'")
@@ -71,16 +61,20 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
            partition_id = _generate_partition_id()
            logger.info("No partition id provided, generated a new one: %s", partition_id)

-        self._client_pool = MongoClientPool(self._client)
+        self._client_pool = MongoClientPool[Mapping[str, Any]](
+            mongo_uri=self._mongo_uri,
+            mongo_client_kwargs=self._mongo_client_kwargs,
+        )

        super().__init__(
            collections=MongoLightningCollections(
                self._client_pool,
                database_name,
                partition_id,
-                prometheus_tracker=MongoOperationPrometheusTracker(enabled=self._enable_prometheus),
+                tracker=tracker,
            ),
-            prometheus=self._enable_prometheus,
+            tracker=tracker,
+            scan_debounce_seconds=scan_debounce_seconds,
        )

    @property
@@ -96,9 +90,6 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
    async def close(self) -> None:
        """Close the store by closing the client pool."""
        await self._client_pool.close()
-        # 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
@@ -136,6 +127,15 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
            await asyncio.sleep(rest_time)
            current_time = time.time()

+        # Logging will help debugging when there are stuck rollouts.
+        logger.debug(
+            "Waiting for rollouts. Number of finished rollouts: %d; number of unfinished rollouts: %d",
+            len(finished_rollouts),
+            len(unfinished_rollout_ids),
+        )
+        if len(unfinished_rollout_ids) < 30:
+            logger.debug("Unfinished rollouts: %s", unfinished_rollout_ids)
+
        # 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]

@@ -3,5 +3,6 @@
 from .agentops import AgentOpsTracer
 from .base import Tracer
 from .otel import OtelTracer
+from .weave import WeaveTracer

-__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer"]
+__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer", "WeaveTracer"]
@@ -1,396 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-import asyncio
-import logging
-import multiprocessing
-import queue
-import uuid
-from contextlib import asynccontextmanager, contextmanager
-from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
-from urllib.parse import urlparse
-
-from httpdbg.hooks.all import httprecord
-from httpdbg.records import HTTPRecords
-from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import SpanKind, Status, StatusCode
-from opentelemetry.trace.span import (
-    SpanContext,
-    TraceFlags,
-    TraceState,
-)
-
-from agentlightning.store import LightningStore
-
-from .base import Tracer
-
-logger = logging.getLogger(__name__)
-
-
-class HttpTracer(Tracer):
-    """
-    A tracer implementation that captures HTTP requests using httpdbg.
-
-    This tracer hooks into the Python HTTP libraries and captures all
-    HTTP requests and responses made during the traced code execution.
-    The captured requests are converted to OpenTelemetry spans for
-    compatibility with the rest of the tracing ecosystem.
-
-    Caution: The current implementation of HttpTracer is very fragile,
-    and we do not recommend using it in production.
-    It is primarily for demonstration and testing purposes.
-
-    Deprecated: This tracer is deprecated and will be removed in a future version.
-    Please use LLMProxy as an alternative.
-
-    Attributes:
-        include_headers: Whether to include HTTP headers in the spans.
-            Headers may contain sensitive information. Use with caution.
-        include_body: Whether to include HTTP request and response bodies in the spans.
-            Bodies may be large and contain sensitive information. Use with caution.
-        include_agentlightning_requests: Whether to include requests initiated by AgentLightning itself.
-        subprocess_mode: Whether to run trace_run and trace_run_async in subprocesses for isolation.
-        subprocess_timeout: Timeout for subprocess execution in seconds.
-    """
-
-    AGENTLIGHTNING_HEADERS = {"x-agentlightning-client"}
-
-    def __init__(
-        self,
-        include_headers: bool = False,
-        include_body: bool = False,
-        include_agentlightning_requests: bool = False,
-        subprocess_mode: bool = True,
-        subprocess_timeout: float = 3600.0,
-    ):
-        super().__init__()
-        self._last_records: Optional[HTTPRecords] = None
-        self.include_headers = include_headers
-        self.include_body = include_body
-        self.include_agentlightning_requests = include_agentlightning_requests
-        self.subprocess_mode = subprocess_mode
-        self.subprocess_timeout = subprocess_timeout
-
-    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
-        """
-        Initialize the tracer in a worker process.
-
-        Args:
-            worker_id: The ID of the worker process.
-            store: The store to add the spans to.
-        """
-        super().init_worker(worker_id, store)
-        logger.info(f"[Worker {worker_id}] HttpTracer initialized.")
-
-    @asynccontextmanager
-    async def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> AsyncGenerator[HTTPRecords, None]:
-        """
-        Starts a new HTTP tracing context. This should be used as a context manager.
-
-        Args:
-            name: Optional name for the tracing context.
-        """
-        with self._trace_context_sync(name=name, **kwargs) as records:
-            yield records
-
-    @contextmanager
-    def _trace_context_sync(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
-        """
-        Starts a new HTTP tracing context. This should be used as a context manager.
-
-        Args:
-            name: Optional name for the tracing context.
-
-        Yields:
-            The HTTPRecords instance containing traced HTTP activities.
-        """
-        records = HTTPRecords()
-        with httprecord(records):
-            self._last_records = records
-            yield records
-
-    def get_last_trace(self) -> List[ReadableSpan]:
-        """
-        Retrieves the raw list of captured spans from the most recent trace.
-
-        Returns:
-            A list of OpenTelemetry `ReadableSpan` objects converted from HTTP records.
-        """
-        if self._last_records is None:
-            return []
-
-        return self._convert_to_spans(self._last_records)
-
-    def _convert_to_spans(self, records: HTTPRecords) -> List[ReadableSpan]:
-        """
-        Convert HTTPRecords to OpenTelemetry spans.
-
-        Args:
-            records: The HTTPRecords instance containing HTTP traces.
-
-        Returns:
-            A list of ReadableSpan objects representing the HTTP activities.
-        """
-        spans: List[ReadableSpan] = []
-
-        # Create a trace ID that will be shared by all spans in this trace
-        trace_id = int(uuid.uuid4().hex[:16], 16)
-
-        for record in records.requests.values():
-            # Skip AgentLightning requests if include_agentlightning_requests is False
-            should_skip = False
-            if not self.include_agentlightning_requests and record.request and record.request.headers:
-                for header in record.request.headers:
-                    if header.name.lower() in self.AGENTLIGHTNING_HEADERS and header.value.lower() == "true":
-                        should_skip = True
-                        break
-
-            if should_skip:
-                continue
-
-            # Create a span ID for this specific HTTP request
-            span_id = int(uuid.uuid4().hex[:8], 16)
-
-            # Create a span context
-            span_context = SpanContext(
-                trace_id=trace_id,
-                span_id=span_id,
-                is_remote=False,
-                trace_flags=TraceFlags(TraceFlags.SAMPLED),
-                trace_state=TraceState(),
-            )
-
-            # Extract important information from the HTTP record
-            method = record.method
-            url = record.url
-            parsed_url = urlparse(url)
-            status_code = record.status_code
-
-            # Create attributes dictionary
-            attributes: Dict[str, Any] = {
-                "http.method": method,
-                "http.url": url,
-                "http.target": parsed_url.path,
-                "http.host": parsed_url.netloc,
-            }
-
-            if status_code is not None and status_code > 0:  # type: ignore
-                attributes["http.status_code"] = status_code
-
-            # Calculate duration - from begin time to last update
-            duration = None
-            if hasattr(record, "last_update") and record.last_update and record.tbegin:
-                duration = (record.last_update - record.tbegin).total_seconds()
-                attributes["http.duration_ms"] = duration * 1000  # Convert to ms
-
-            # Optionally include headers
-            if self.include_headers and record.request and record.request.headers:
-                for header in record.request.headers:
-                    header_name = header.name.lower()
-                    attributes[f"http.request.header.{header_name}"] = header.value
-
-            if self.include_headers and record.response and record.response.headers:
-                for header in record.response.headers:
-                    header_name = header.name.lower()
-                    attributes[f"http.response.header.{header_name}"] = header.value
-
-            # Optionally include body - preserve complete content for analysis
-            if self.include_body and record.request:
-                body_content = record.request.content
-                if body_content:
-                    # Store raw body content for later parsing/analysis
-                    attributes["http.request.body"] = body_content
-
-            if self.include_body and record.response:
-                body_content = record.response.content
-                if body_content:
-                    # Store raw body content for later parsing/analysis
-                    attributes["http.response.body"] = body_content
-
-            # Determine span status
-            span_status = StatusCode.OK
-            if status_code and status_code >= 400 or record.exception:
-                span_status = StatusCode.ERROR
-
-            # Create start and end timestamps in nanoseconds
-            # If we have duration, use it, otherwise default to current time - 1ms
-            start_time_ns = int(record.tbegin.timestamp() * 1e9)
-            if duration:
-                end_time_ns = int((record.tbegin.timestamp() + duration) * 1e9)
-            else:
-                end_time_ns = int(record.last_update.timestamp() * 1e9)
-
-            span = ReadableSpan(
-                name=f"HTTP {method} {url}",
-                context=span_context,
-                parent=None,
-                kind=SpanKind.CLIENT,
-                status=Status(span_status),
-                start_time=start_time_ns,
-                end_time=end_time_ns,
-                attributes=attributes,
-                events=[],
-                links=[],
-                resource=None,
-            )
-
-            spans.append(span)
-
-        return spans
-
-    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
-        """
-        A convenience wrapper to trace the execution of a single synchronous function.
-
-        If subprocess_mode is enabled, the function will be executed in an isolated subprocess
-        to prevent HTTP hooks from affecting the parent process.
-
-        Args:
-            func: The synchronous function to execute and trace.
-            *args: Positional arguments to pass to the function.
-            **kwargs: Keyword arguments to pass to the function.
-
-        Returns:
-            The return value of the function.
-        """
-        if self.subprocess_mode:
-            return self._trace_run_subprocess(func, args, kwargs)
-        else:
-            return super().trace_run(func, *args, **kwargs)
-
-    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
-        """
-        A convenience wrapper to trace the execution of a single asynchronous function.
-
-        If subprocess_mode is enabled, the function will be executed in an isolated subprocess
-        to prevent HTTP hooks from affecting the parent process.
-
-        Args:
-            func: The asynchronous function to execute and trace.
-            *args: Positional arguments to pass to the function.
-            **kwargs: Keyword arguments to pass to the function.
-
-        Returns:
-            The return value of the function.
-        """
-        if self.subprocess_mode:
-            loop = asyncio.get_event_loop()
-            return await loop.run_in_executor(
-                None, self._trace_run_subprocess, func, args, kwargs, True  # True for async
-            )
-        else:
-            return await super().trace_run_async(func, *args, **kwargs)
-
-    def _trace_run_subprocess(
-        self,
-        func: Callable[..., Any],
-        args: Optional[Tuple[Any, ...]] = None,
-        kwargs: Optional[Dict[str, Any]] = None,
-        is_async: bool = False,
-    ) -> Any:
-        """
-        Execute a function in a subprocess with HTTP tracing.
-
-        Args:
-            func: The function to execute.
-            args: Positional arguments to pass to the function.
-            kwargs: Keyword arguments to pass to the function.
-            is_async: Whether the function is asynchronous.
-
-        Returns:
-            The return value of the function.
-        """
-        if args is None:
-            args = ()
-        if kwargs is None:
-            kwargs = {}
-
-        # Create a queue to receive results from the subprocess
-        result_queue = multiprocessing.Queue()  # type: ignore
-
-        # Create and start the subprocess
-        process = multiprocessing.Process(
-            target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async)  # type: ignore
-        )
-        process.start()
-
-        try:
-            # Wait for the process to complete and get the result
-            process.join(timeout=self.subprocess_timeout)
-            result = result_queue.get_nowait()  # type: ignore
-
-            if result["success"]:
-                # Store the captured records for get_last_trace()
-                self._last_records = result["records"]
-                return result["return_value"]  # type: ignore
-            else:
-                if "records" in result:
-                    self._last_records = result["records"]
-                # Re-raise the exception that occurred in the subprocess
-                raise result["exception"]
-
-        except multiprocessing.TimeoutError:
-            process.terminate()
-            process.join()
-            raise TimeoutError(f"Subprocess execution timed out after {self.subprocess_timeout} seconds.")
-        except queue.Empty:
-            logger.error("Traced result is empty. This may indicate a timeout or an issue with the subprocess.")
-        finally:
-            if process.is_alive():
-                process.terminate()
-                process.join()
-
-    def _subprocess_worker(
-        self,
-        func: Callable[..., Any],
-        args: Tuple[Any, ...],
-        kwargs: Dict[str, Any],
-        result_queue: multiprocessing.Queue,  # type: ignore
-        is_async: bool,
-    ) -> None:
-        """
-        Worker function that runs in the subprocess to execute the traced function.
-
-        Args:
-            func: The function to execute.
-            args: Positional arguments.
-            kwargs: Keyword arguments.
-            result_queue: Queue to send results back to parent process.
-            is_async: Whether the function is asynchronous.
-        """
-        # Create a new tracer instance in the subprocess (without subprocess mode to avoid recursion)
-        subprocess_tracer = HttpTracer(
-            include_headers=self.include_headers,
-            include_body=self.include_body,
-            include_agentlightning_requests=self.include_agentlightning_requests,
-            subprocess_mode=False,  # Disable subprocess mode in the worker
-        )
-
-        try:
-            if is_async:
-                # Run async function in new event loop
-                import asyncio
-
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-                try:
-                    return_value = loop.run_until_complete(subprocess_tracer.trace_run_async(func, *args, **kwargs))
-                finally:
-                    loop.close()
-            else:
-                # Run sync function
-                return_value = subprocess_tracer.trace_run(func, *args, **kwargs)
-
-            # Get the captured records
-            records = subprocess_tracer._last_records
-
-            # Send success result back to parent
-            result_queue.put({"success": True, "return_value": return_value, "records": records})  # type: ignore
-
-        except Exception as e:
-            # Log the exception
-            logger.exception(f"Error in subprocess worker in http tracer: {e}")
-
-            # Get the captured records even when there's an exception
-            records = subprocess_tracer._last_records
-            # Send error result back to parent
-            result_queue.put({"success": False, "exception": e, "records": records})  # type: ignore
@@ -0,0 +1,307 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import os
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Tuple, Union
+
+from agentlightning.instrumentation import instrument_weave, uninstrument_weave
+from agentlightning.store.base import LightningStore
+from agentlightning.types.tracer import OtelResource, Span, SpanContext, TraceStatus
+
+from .base import Tracer
+
+if TYPE_CHECKING:
+    from weave.trace.call import Call  # type: ignore
+
+JSONPrimitive = Union[str, int, float, bool, None]
+
+logger = logging.getLogger(__name__)
+
+
+class WeaveTracer(Tracer):
+    """
+    Tracer implementation using Weave for telemetry and trace logging.
+
+    This replaces AgentOpsTracer with a Weave-based manual trace context. It tracks:
+    - Function/method calls
+    - Input/Output data
+    - Exceptions
+    and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
+
+    Attributes:
+        project_name: Name of the Weave project. Used to initialize the Weave client.
+        _store: Optional LightningStore instance for storing collected spans.
+        instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
+    """
+
+    def __init__(
+        self, *, project_name: str | None = None, wandb_api_key: str | None = None, instrument_managed: bool = True
+    ):
+        """
+        Initialize a WeaveTracer instance.
+
+        Args:
+            project_name: Optional project name for Weave; defaults to the current module name.
+            wandb_api_key: Optional W&B API key; sets environment variable if provided.
+            instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
+        """
+        super().__init__()
+        self.project_name = project_name or __name__
+        self.sequence_id = 0
+        self._store: Optional[LightningStore] = None
+        self.instrument_managed = instrument_managed
+
+        if wandb_api_key:
+            os.environ["WANDB_API_KEY"] = wandb_api_key
+
+    def instrument(self, worker_id: int):
+        instrument_weave()
+
+    def uninstrument(self, worker_id: int):
+        uninstrument_weave()
+
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
+        """
+        Initialize the tracer for a worker thread/process.
+
+        Args:
+            worker_id: Identifier of the worker.
+            store: Optional LightningStore for storing spans.
+        """
+        super().init_worker(worker_id, store)
+        logger.info(f"[Worker {worker_id}] Setting up Weave tracer...")
+        self._store = store
+
+        try:
+            import weave
+        except ImportError:
+            raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
+
+        # Optionally patch network calls to bypass real Weave/W&B endpoints
+        if self.instrument_managed:
+            self.instrument(worker_id)
+
+        # Initialize the Weave client if not already initialized
+        if weave.get_client() is None:  # type: ignore
+            try:
+                weave.init(project_name=self.project_name)  # type: ignore
+                logger.info(f"[Worker {worker_id}] Weave client initialized.")
+            except Exception as e:
+                raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}': {e}")
+
+    def teardown_worker(self, worker_id: int):
+        """
+        Clean up tracer resources for the worker.
+
+        Args:
+            worker_id: Identifier of the worker.
+        """
+        super().teardown_worker(worker_id)
+
+        if self.instrument_managed:
+            self.uninstrument(worker_id)
+            logger.info(f"[Worker {worker_id}] Instrumentation removed.")
+
+    @asynccontextmanager
+    async def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncIterator[Any]:
+        """
+        Synchronous implementation of the tracing context.
+
+        Args:
+            name: Optional operation name.
+            store: Optional LightningStore instance.
+            rollout_id: Optional rollout ID.
+            attempt_id: Optional attempt ID.
+
+        Raises:
+            ValueError: If store, rollout_id, and attempt_id are inconsistently provided.
+            RuntimeError: If Weave is not installed or client is uninitialized.
+        """
+        arg_op = name or self.project_name
+        arg_inputs: dict[str, str] | None = {"rollout_id": rollout_id or "", "attempt_id": attempt_id or ""}
+
+        if store is not None and rollout_id is not None and attempt_id is not None:
+            self._rollout_id = rollout_id
+            self._attempt_id = attempt_id
+            self._store = store
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided")
+
+        try:
+            import datetime
+
+            import weave
+        except ImportError:
+            raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
+
+        weave_client = weave.get_client()  # type: ignore
+        if not weave_client:
+            raise RuntimeError("Weave client is not initialized. Call init_worker() first.")
+
+        # Create a new trace call object in Weave
+        trace_call = weave_client.create_call(op=arg_op, inputs=arg_inputs)  # type: ignore
+        trace_call.started_at = datetime.datetime.now(tz=datetime.timezone.utc)
+
+        try:
+            yield trace_call
+        except Exception as e:
+            # Finish trace and log any exception
+            weave_client.finish_call(trace_call, exception=e)  # type: ignore
+            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
+        finally:
+            # Finish trace even if no exception
+            weave_client.finish_call(trace_call)  # type: ignore
+            await self._on_finish_handler(trace_call)  # type: ignore
+
+    async def _on_finish_handler(self, call: "Call", *args: Any, **kwargs: Any) -> None:  # type: ignore
+        """
+        Handler called when a Weave Call finishes.
+
+        Converts the call (including nested children) into spans and stores them in LightningStore.
+        """
+        spans, self.sequence_id = self.convert_call_to_spans(call, self._rollout_id, self._attempt_id, self.sequence_id)  # type: ignore
+
+        if self._store and self._rollout_id and self._attempt_id:
+            try:
+                await self._store.add_many_spans(spans)
+            except Exception as e:
+                logger.exception(f"Error adding span to store: {e}")
+
+    def convert_call_to_spans(
+        self,
+        call: "Call",  # type: ignore
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        seq_start: int = 0,
+    ) -> tuple[List[Span], int]:
+        """
+        Recursively convert a Weave Call (with nested children) into a flat list of Agent Lightning Spans.
+
+        Args:
+            call: The Weave Call object.
+            rollout_id: Optional rollout ID to attach to spans.
+            attempt_id: Optional attempt ID to attach to spans.
+            seq_start: Sequence number to start from.
+
+        Returns:
+            Tuple of (list_of_spans, next_sequence_id).
+        """
+        spans: List[Span] = []
+        sequence_id = seq_start
+
+        rollout_id = rollout_id or ""  # type: ignore
+        attempt_id = attempt_id or ""  # type: ignore
+
+        start_dt = getattr(call, "started_at", None)  # type: ignore
+        start_ts: Optional[float] = start_dt.timestamp() if start_dt else None
+
+        end_dt = getattr(call, "ended_at", None)  # type: ignore
+        end_ts: Optional[float] = end_dt.timestamp() if end_dt else None
+
+        trace_id = str(getattr(call, "trace_id", None))  # type: ignore
+        span_id = str(getattr(call, "id", None))  # type: ignore
+        parent_id = str(getattr(call, "parent_id", None)) if getattr(call, "parent_id", None) else None  # type: ignore
+
+        exception = getattr(call, "exception", None)  # type: ignore
+        status_code = "ERROR" if exception else "OK"
+
+        def sanitize(
+            inputs: Dict[str, Any],
+            output: Dict[str, Any],
+        ) -> Dict[str, str | JSONPrimitive]:
+            stack: List[Tuple[Any, str]] = [
+                (inputs or {}, "input"),
+                (output or {}, "output"),
+            ]
+
+            attributes: Dict[str, str | JSONPrimitive] = {}
+
+            while stack:
+                value, key = stack.pop()
+
+                if isinstance(value, dict):
+                    for k, v in value.items():  # type: ignore
+                        stack.append((v, f"{key}.{k}"))  # type: ignore
+                elif isinstance(value, (list, tuple)):
+                    for i, v in enumerate(value):  # type: ignore
+                        stack.append((v, f"{key}.{i}"))  # type: ignore
+                else:
+                    if value is None:
+                        attributes[key] = "None"
+                    elif isinstance(value, (str, int, float, bool)):
+                        attributes[key] = value
+                    else:
+                        try:
+                            attributes[key] = str(value)
+                        except Exception:
+                            attributes[key] = "None"
+
+            return attributes
+
+        inputs = getattr(call, "inputs", {})  # type: ignore
+        output = getattr(call, "output", {})  # type: ignore
+        attributes = sanitize(inputs, output)
+
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_state={},
+        )
+
+        parent_context = (
+            SpanContext(
+                trace_id=trace_id,
+                span_id=parent_id,
+                is_remote=False,
+                trace_state={},
+            )
+            if parent_id
+            else None
+        )
+
+        # Build the Span object
+        span = Span(
+            rollout_id=rollout_id or "",
+            attempt_id=attempt_id or "",
+            sequence_id=sequence_id,
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_id=parent_id,
+            name=getattr(call, "func_name", "unknown"),  # type: ignore
+            status=TraceStatus(status_code=status_code),
+            attributes=attributes,  # type: ignore
+            events=[],  # Weave calls do not generate events
+            links=[],  # Weave calls do not generate links
+            start_time=start_ts,
+            end_time=end_ts,
+            context=context,
+            parent=parent_context,
+            resource=OtelResource(attributes={}, schema_url=""),
+        )
+
+        spans.append(span)
+        sequence_id += 1
+
+        children: List["Call"] = getattr(call, "_children", [])  # type: ignore
+        # Recursively process child calls
+        for child in children:  # type: ignore
+            child_spans, sequence_id = self.convert_call_to_spans(  # type: ignore
+                child,  # type: ignore
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                seq_start=sequence_id,
+            )
+            spans.extend(child_spans)
+
+        return spans, sequence_id
@@ -153,7 +153,7 @@ class Trainer(TrainerLegacy):
        self.worker_id: Optional[int] = None

        if dev:
-            logger.warning(
+            warnings.warn(
                "Trainer(dev=True) is deprecated and will be removed in future versions. "
                "Please use Trainer.dev(...) instead.",
                DeprecationWarning,
@@ -16,16 +16,18 @@ 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

+import aiologic
+
 if TYPE_CHECKING:
    from prometheus_client import CollectorRegistry

 LabelDict = Dict[str, str]
-LabelKey = Tuple[Tuple[str, str], ...]  # normalized, sorted (key, value) pairs
+# Label metadata
+LabelKey = Tuple[Tuple[str, str], ...]  # normalized (key, value) pairs in registration order

 logger = logging.getLogger(__name__)

@@ -45,7 +47,7 @@ def _validate_labels(
        expected_names: Expected label names as a tuple.

    Returns:
-        A tuple of (key, value) pairs sorted by registered label order.
+        A tuple of (key, value) pairs honoring the registered label order.

    Raises:
        ValueError: If label keys do not match expected_names.
@@ -67,11 +69,17 @@ def _normalize_label_names(label_names: Optional[Sequence[str]]) -> Tuple[str, .
        label_names: Iterable of label names or None.

    Returns:
-        A tuple of label names sorted alphabetically.
+        A tuple of label names preserving their original order.
    """
    if not label_names:
        return ()
-    return tuple(sorted(label_names))
+    return tuple(label_names)
+
+
+def _normalize_prometheus_metric_name(metric_name: str) -> str:
+    """Normalizes Prometheus metric names by replacing unsupported characters."""
+
+    return metric_name.replace(".", "_")


@dataclass(frozen=True)
@@ -80,6 +88,7 @@ class _CounterDef:

    name: str
    label_names: Tuple[str, ...]
+    group_level: Optional[int] = None


@dataclass(frozen=True)
@@ -89,6 +98,7 @@ class _HistogramDef:
    name: str
    label_names: Tuple[str, ...]
    buckets: Tuple[float, ...]
+    group_level: Optional[int] = None


@dataclass
@@ -110,16 +120,25 @@ class _HistogramState:
 class MetricsBackend:
    """Abstract base class for metrics backends."""

+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return False
+
    def register_counter(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a counter metric.

        Args:
            name: Metric name.
-            label_names: List of label names. Order is not important.
+            label_names: List of label names. Order determines the truncation
+                priority for group-level logging.
+            group_level: Optional per-metric grouping depth for backends that
+                support label grouping (Console). Global backend settings take
+                precedence when provided.

        Raises:
            ValueError: If the metric is already registered with a different
@@ -132,14 +151,19 @@ class MetricsBackend:
        name: str,
        label_names: Optional[Sequence[str]] = None,
        buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a histogram metric.

        Args:
            name: Metric name.
-            label_names: List of label names. Order is not important.
+            label_names: List of label names. Order determines the truncation
+                priority for group-level logging.
            buckets: Bucket boundaries (exclusive upper bounds). If None, the
                backend may choose defaults.
+            group_level: Optional per-metric grouping depth for backends that
+                support label grouping (Console). Global backend settings take
+                precedence when provided.

        Raises:
            ValueError: If the metric is already registered with a different
@@ -147,7 +171,7 @@ class MetricsBackend:
        """
        raise NotImplementedError()

-    def inc_counter(
+    async def inc_counter(
        self,
        name: str,
        amount: float = 1.0,
@@ -166,7 +190,7 @@ class MetricsBackend:
        """
        raise NotImplementedError()

-    def observe_histogram(
+    async def observe_histogram(
        self,
        name: str,
        value: float,
@@ -199,22 +223,32 @@ class ConsoleMetricsBackend(MetricsBackend):

    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:
+    Label grouping: When logging, label dictionaries are truncated to the first
+    `group_level` label pairs (following the registered label order) and metrics
+    with identical truncated labels are aggregated together. For example:

-        labels = {"method": "GET", "path": "/", "status": "200"}
-        group_level = 2 -> logged labels {"method": "GET", "path": "/"}
+    ```python
+    labels = {"method": "GET", "path": "/", "status": "200"}
+    group_level = 2  # aggregated labels {"method": "GET", "path": "/"}
+    ```

-    If `group_level` is None or < 1, all labels are logged.
+    If `group_level` is None or < 1, all label combinations for a metric are
+    merged into a single log entry (equivalent to grouping by zero labels).
+    Individual counters or histograms can set their own `group_level` during
+    registration; those values apply only when the backend-level `group_level`
+    is unset, allowing selective overrides.

-    Thread-safety: A single lock protects shared state mutation, pruning, and snapshotting.
-    Percentile computation, formatting, and printing are done after releasing the lock.
+    Thread-safety: Runtime updates and snapshotting use two aiologic locks: one for mutating
+    shared state and another that serializes the global logging decision/snapshot capture so
+    other tasks can continue writing. Metric registration happens during initialization,
+    so it is intentionally left lock-free; this assumption is documented here to avoid
+    blocking writes unnecessarily.
    """

    def __init__(
        self,
        window_seconds: Optional[float] = 60.0,
-        log_interval_seconds: float = 5.0,
+        log_interval_seconds: float = 10.0,
        group_level: Optional[int] = None,
    ) -> None:
        """Initializes ConsoleMetricsBackend.
@@ -226,8 +260,9 @@ class ConsoleMetricsBackend(MetricsBackend):
                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.
+                `group_level` labels (following registered order) are retained and metric
+                events sharing those labels are aggregated. If None or < 1,
+                all label combinations collapse into a single group per metric.
        """
        self.window_seconds = window_seconds
        self.log_interval_seconds = log_interval_seconds
@@ -243,40 +278,42 @@ class ConsoleMetricsBackend(MetricsBackend):
        # Global last log time (for all metrics)
        self._last_log_time: Optional[float] = None

-        self._lock = threading.Lock()
+        self._write_lock = aiologic.Lock()
+        self._snapshot_lock = aiologic.Lock()

    def register_counter(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = 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)
+        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_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
+        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)
+        self._counters[name] = _CounterDef(name=name, label_names=label_tuple, group_level=group_level)

    def register_histogram(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
        buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a histogram metric.

@@ -288,29 +325,29 @@ class ConsoleMetricsBackend(MetricsBackend):
        else:
            bucket_tuple = tuple(buckets)

-        with self._lock:
-            existing_counter = self._counters.get(name)
-            existing_hist = self._histograms.get(name)
+        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_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
+        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,
-            )
+        self._histograms[name] = _HistogramDef(
+            name=name,
+            label_names=label_tuple,
+            buckets=bucket_tuple,
+            group_level=group_level,
+        )

-    def inc_counter(
+    async def inc_counter(
        self,
        name: str,
        amount: float = 1.0,
@@ -330,7 +367,7 @@ class ConsoleMetricsBackend(MetricsBackend):
        label_key = _validate_labels("counter", name, labels, definition.label_names)
        state_key = (name, label_key)

-        with self._lock:
+        async with self._write_lock:
            state = self._counter_state.get(state_key)
            if state is None:
                state = _CounterState(timestamps=[], amounts=[])
@@ -340,18 +377,19 @@ class ConsoleMetricsBackend(MetricsBackend):
            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
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        should_log = False
+        snapshot_time = now

-        if should_log and (counter_snaps or hist_snaps):
+        async with self._snapshot_lock:
+            should_log = self._should_log_locked(now)
+        if should_log:
+            async with self._write_lock:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
            self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)

-    def observe_histogram(
+    async def observe_histogram(
        self,
        name: str,
        value: float,
@@ -371,7 +409,7 @@ class ConsoleMetricsBackend(MetricsBackend):
        label_key = _validate_labels("histogram", name, labels, definition.label_names)
        state_key = (name, label_key)

-        with self._lock:
+        async with self._write_lock:
            state = self._hist_state.get(state_key)
            if state is None:
                state = _HistogramState(timestamps=[], values=[])
@@ -381,15 +419,17 @@ class ConsoleMetricsBackend(MetricsBackend):
            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
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        should_log = False
+        snapshot_time = now

-        if should_log and (counter_snaps or hist_snaps):
+        async with self._snapshot_lock:
+            should_log = self._should_log_locked(now)
+
+        if should_log:
+            async with self._write_lock:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
            self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)

    def _prune_events(
@@ -490,21 +530,22 @@ class ConsoleMetricsBackend(MetricsBackend):

        return counter_snaps, hist_snaps

-    def _truncate_labels_for_logging(self, labels: LabelDict) -> LabelDict:
+    def _truncate_labels_for_logging(self, labels: LabelDict, group_level: Optional[int]) -> LabelDict:
        """Returns a label dict truncated to the configured group depth.

        Args:
            labels: Original label dictionary.
+            group_level: Effective grouping depth for this metric.

        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.
+            chosen by registered label order. If group_level is None or < 1,
+            returns an empty dict so that all label combinations collapse together.
        """
-        if self.group_level is None or self.group_level < 1:
-            return dict(labels)
-        items = sorted(labels.items())
-        return dict(items[: self.group_level])
+        if group_level is None or group_level < 1:
+            return {}
+        items = list(labels.items())
+        return dict(items[:group_level])

    def _log(self, message: str) -> None:
        """Logs a message via the module logger."""
@@ -523,21 +564,101 @@ class ConsoleMetricsBackend(MetricsBackend):
            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)
+        for name, labels, timestamps, amounts in self._group_counter_snapshots(counter_snaps):
+            line = self._log_counter(name, 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)
+        for name, labels, values, buckets in self._group_histogram_snapshots(hist_snaps):
+            line = self._log_histogram(name, labels, values, buckets, snapshot_time)
            if line:
                entries.append(line)

        if entries:
+            entries.sort()
            self._log("  ".join(entries))

+    def _effective_group_level(self, metric_name: str, *, is_histogram: bool) -> Optional[int]:
+        """Returns the active group level for a metric, honoring per-metric overrides."""
+        if self.group_level is not None:
+            return self.group_level
+        if is_histogram:
+            definition = self._histograms.get(metric_name)
+        else:
+            definition = self._counters.get(metric_name)
+        if definition is None:
+            return None
+        return definition.group_level
+
+    def _group_counter_snapshots(
+        self,
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]],
+    ) -> List[Tuple[str, LabelDict, List[float], List[float]]]:
+        grouped: Dict[Tuple[str, Tuple[Tuple[str, str], ...]], Dict[str, Any]] = {}
+        for name, labels, timestamps, amounts in counter_snaps:
+            group_level = self._effective_group_level(name, is_histogram=False)
+            truncated_labels = self._truncate_labels_for_logging(labels, group_level)
+            key = (name, tuple(truncated_labels.items()))
+            entry = grouped.setdefault(
+                key,
+                {"name": name, "labels": truncated_labels, "timestamps": [], "amounts": []},
+            )
+            entry["timestamps"].extend(timestamps)
+            entry["amounts"].extend(amounts)
+
+        grouped_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        for entry in grouped.values():
+            timestamps = entry["timestamps"]
+            amounts = entry["amounts"]
+            if not timestamps:
+                continue
+            combined = sorted(zip(timestamps, amounts), key=lambda item: item[0])
+            ordered_timestamps = [ts for ts, _ in combined]
+            ordered_amounts = [amt for _, amt in combined]
+            grouped_snaps.append(
+                (
+                    entry["name"],
+                    entry["labels"],
+                    ordered_timestamps,
+                    ordered_amounts,
+                )
+            )
+
+        return grouped_snaps
+
+    def _group_histogram_snapshots(
+        self,
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]],
+    ) -> List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]]:
+        grouped: Dict[Tuple[str, Tuple[Tuple[str, str], ...]], Dict[str, Any]] = {}
+        for name, labels, values, buckets in hist_snaps:
+            group_level = self._effective_group_level(name, is_histogram=True)
+            truncated_labels = self._truncate_labels_for_logging(labels, group_level)
+            key = (name, tuple(truncated_labels.items()))
+            entry = grouped.setdefault(
+                key,
+                {"name": name, "labels": truncated_labels, "values": [], "buckets": buckets},
+            )
+            if entry["buckets"] != buckets:
+                raise ValueError(f"Histogram buckets mismatch for metric '{name}'.")
+            entry["values"].extend(values)
+
+        grouped_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        for entry in grouped.values():
+            values = entry["values"]
+            if not values:
+                continue
+            grouped_snaps.append(
+                (
+                    entry["name"],
+                    entry["labels"],
+                    list(values),
+                    entry["buckets"],
+                )
+            )
+
+        return grouped_snaps
+
    def _log_counter(
        self,
        name: str,
@@ -598,8 +719,8 @@ class ConsoleMetricsBackend(MetricsBackend):

 def _format_label_string(labels: LabelDict) -> str:
    if not labels:
-        return "{}"
-    ordered = ",".join(f"{key}={value}" for key, value in sorted(labels.items()))
+        return ""
+    ordered = ",".join(f"{key}={value}" for key, value in labels.items())
    return f"{{{ordered}}}"


@@ -622,6 +743,9 @@ class PrometheusMetricsBackend(MetricsBackend):

    Thread-safety: Registration is protected by a lock. Metric updates assume metrics
    are registered during initialization and then remain stable.
+
+    Due to the nature of Prometheus, this backend is only suitable for recording high-volume metrics.
+    Low-volume metrics might be lost if the event has only appeared once.
    """

    def __init__(self) -> None:
@@ -641,46 +765,50 @@ class PrometheusMetricsBackend(MetricsBackend):
        self._histograms: Dict[str, _HistogramDef] = {}
        self._prom_counters: Dict[str, Any] = {}
        self._prom_histograms: Dict[str, Any] = {}
+        self._prom_metric_names: Dict[str, str] = {}

-        self._lock = threading.Lock()
+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return True

    def register_counter(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = 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.")
+        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
+        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_name = self._register_prometheus_metric_name(name)
+        self._counters[name] = _CounterDef(name=name, label_names=label_tuple, group_level=group_level)

-            prom_counter = PromCounter(
-                name,
-                f"Counter {name}",
-                labelnames=label_tuple,
-            )
-            self._prom_counters[name] = prom_counter
+        prom_counter = PromCounter(
+            prom_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,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a Prometheus histogram metric."""
        from prometheus_client import Histogram as PromHistogram
@@ -688,43 +816,44 @@ class PrometheusMetricsBackend(MetricsBackend):
        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.")
+        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
+        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,
+        prom_name = self._register_prometheus_metric_name(name)
+        self._histograms[name] = _HistogramDef(
+            name=name,
+            label_names=label_tuple,
+            buckets=bucket_tuple,
+            group_level=group_level,
+        )
+
+        if bucket_tuple:
+            prom_hist = PromHistogram(
+                prom_name,
+                f"Histogram {name}",
+                labelnames=label_tuple,
                buckets=bucket_tuple,
            )
+        else:
+            prom_hist = PromHistogram(
+                prom_name,
+                f"Histogram {name}",
+                labelnames=label_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

-            self._prom_histograms[name] = prom_hist
-
-    def inc_counter(
+    async def inc_counter(
        self,
        name: str,
        amount: float = 1.0,
@@ -743,7 +872,7 @@ class PrometheusMetricsBackend(MetricsBackend):
        else:
            prom_counter.inc(amount)

-    def observe_histogram(
+    async def observe_histogram(
        self,
        name: str,
        value: float,
@@ -762,6 +891,19 @@ class PrometheusMetricsBackend(MetricsBackend):
        else:
            prom_hist.observe(value)

+    def _register_prometheus_metric_name(self, name: str) -> str:
+        """Registers the normalized Prometheus metric name and ensures uniqueness."""
+
+        normalized = _normalize_prometheus_metric_name(name)
+        existing = self._prom_metric_names.get(normalized)
+        if existing is not None and existing != name:
+            raise ValueError(
+                f"Prometheus metric name conflict: '{name}' normalizes to '{normalized}', "
+                f"which is already used by '{existing}'. Consider renaming one of the metrics."
+            )
+        self._prom_metric_names.setdefault(normalized, name)
+        return normalized
+

 class MultiMetricsBackend(MetricsBackend):
    """Metrics backend that forwards calls to multiple underlying backends."""
@@ -779,20 +921,26 @@ class MultiMetricsBackend(MetricsBackend):
            raise ValueError("MultiMetricsBackend requires at least one backend.")
        self._backends = list(backends)

+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return any(backend.has_prometheus() for backend in self._backends)
+
    def register_counter(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a counter metric in all underlying backends."""
        for backend in self._backends:
-            backend.register_counter(name, label_names=label_names)
+            backend.register_counter(name, label_names=label_names, group_level=group_level)

    def register_histogram(
        self,
        name: str,
        label_names: Optional[Sequence[str]] = None,
        buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
    ) -> None:
        """Registers a histogram metric in all underlying backends."""
        for backend in self._backends:
@@ -800,9 +948,10 @@ class MultiMetricsBackend(MetricsBackend):
                name,
                label_names=label_names,
                buckets=buckets,
+                group_level=group_level,
            )

-    def inc_counter(
+    async def inc_counter(
        self,
        name: str,
        amount: float = 1.0,
@@ -810,9 +959,9 @@ class MultiMetricsBackend(MetricsBackend):
    ) -> None:
        """Increments a counter metric in all underlying backends."""
        for backend in self._backends:
-            backend.inc_counter(name, amount=amount, labels=labels)
+            await backend.inc_counter(name, amount=amount, labels=labels)

-    def observe_histogram(
+    async def observe_histogram(
        self,
        name: str,
        value: float,
@@ -820,9 +969,10 @@ class MultiMetricsBackend(MetricsBackend):
    ) -> None:
        """Records a histogram observation in all underlying backends."""
        for backend in self._backends:
-            backend.observe_histogram(name, value=value, labels=labels)
+            await backend.observe_histogram(name, value=value, labels=labels)


+# This variable should be carried into forked processes
 _prometheus_multiproc_dir: tempfile.TemporaryDirectory[str] | None = None


@@ -840,7 +990,7 @@ def setup_multiprocess_prometheus():
        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."
+            "Found PROMETHEUS_MULTIPROC_DIR was set by user. This directory must be wiped between multiple runs."
        )


@@ -849,7 +999,7 @@ def get_prometheus_registry() -> CollectorRegistry:
    from prometheus_client import REGISTRY, CollectorRegistry, multiprocess

    if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None:
-        logger.debug("Using multiprocess registry for prometheus metrics")
+        logger.info("Using multiprocess registry for prometheus metrics: %s", os.getenv("PROMETHEUS_MULTIPROC_DIR"))
        registry = CollectorRegistry()
        multiprocess.MultiProcessCollector(registry)
        return registry
@@ -857,17 +1007,19 @@ def get_prometheus_registry() -> CollectorRegistry:
    return REGISTRY


-def shutdown_metrics():
+def shutdown_metrics(server: Any = None, worker: Any = None, *args: Any, **kwargs: Any) -> None:
    """Shutdown prometheus metrics."""

-    from prometheus_client import multiprocess
+    if _prometheus_multiproc_dir is not None:
+        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))
+        path = _prometheus_multiproc_dir
+        try:
+            if hasattr(worker, "pid"):
+                pid = worker.pid
+            else:
+                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))
@@ -940,9 +940,9 @@ class PythonServerLauncher:
                ),  # Allow half the timeout for graceful shutdown
            }
            if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
-                from prometheus_client import multiprocess
+                from agentlightning.utils.metrics import shutdown_metrics

-                options["child_exit"] = lambda server, worker: multiprocess.mark_process_dead(worker.pid)  # type: ignore
+                options["child_exit"] = shutdown_metrics  # type: ignore

            self._gunicorn_app = GunicornApp(self.app, options)

@@ -144,6 +144,8 @@ class AgentModeDaemon:
        llm_proxy: LLMProxy | None = None,
        store: LightningStore | None = None,
        adapter: TraceToTripletBase | None = None,
+        processor: Any = None,
+        image_base_dir: Optional[str] = None,
    ):
        self.mode = mode
        self.llm_timeout_seconds = llm_timeout_seconds
@@ -183,7 +185,12 @@ class AgentModeDaemon:
        self.mini_batch_size = mini_batch_size
        self.pad_token_id = pad_token_id
        self.tokenizer = tokenizer
+        self.processor = processor
        self.reward_fillna_value = reward_fillna_value
+        self.image_base_dir = image_base_dir
+
+        # Check if model requires multimodal position_ids (e.g., Qwen2-VL)
+        self._use_mrope = self._is_mrope_model()

        # Internal State
        self.backend_llm_server_addresses: List[str] = []
@@ -202,6 +209,75 @@ class AgentModeDaemon:
        loop.run_forever()
        loop.close()

+    # Multimodal utilities for M-RoPE position embeddings
+
+    def _is_mrope_model(self) -> bool:
+        """Check if processor requires M-RoPE position embeddings."""
+        if self.processor is None or not hasattr(self.processor, "image_processor"):
+            return False
+        name = self.processor.image_processor.__class__.__name__
+        return "Qwen2VLImageProcessor" in name or "Qwen3VLImageProcessor" in name
+
+    def _resolve_image_path(self, path: str) -> str:
+        """Resolve relative image path with base directory."""
+        import os
+
+        if os.path.isabs(path):
+            return path
+        if self.image_base_dir is None:
+            raise ValueError(f"Relative path '{path}' requires 'image_base_dir' to be set.")
+        return os.path.join(self.image_base_dir, path)
+
+    def _get_image_grid_thw(self, image_urls: List[str]) -> Optional[torch.Tensor]:
+        """Compute image_grid_thw from image URLs for M-RoPE computation.
+
+        Args:
+            image_urls: List of image URLs extracted from triplet prompt payload.
+                URLs can be http(s):// URLs or file:// URIs, or data: URIs.
+        """
+        from PIL import Image
+        from verl.utils.dataset.vision_utils import process_image  # pyright: ignore[reportUnknownVariableType]
+
+        if self.processor is None or not image_urls:
+            return None
+
+        def to_image_uri(url: str) -> str:
+            # Already a proper URI (http, https, file, data)
+            if url.startswith(("http://", "https://", "file://", "data:")):
+                return url
+            # Treat as a file path that needs resolution
+            resolved = self._resolve_image_path(url)
+            return f"file://{resolved}"
+
+        images: List[Image.Image] = [process_image({"image": to_image_uri(url)}) for url in image_urls]
+        model_inputs = self.processor(text=["dummy"], images=images, return_tensors="pt")
+        return model_inputs.get("image_grid_thw")
+
+    def _compute_mrope_position_ids(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        image_grid_thw: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute 4D position_ids for M-RoPE models."""
+        from typing import Callable
+
+        get_rope_index: Callable[..., torch.Tensor]
+        if "Qwen3VL" in self.processor.__class__.__name__:
+            from verl.models.transformers.qwen3_vl import get_rope_index  # pyright: ignore[reportUnknownVariableType]
+        else:
+            from verl.models.transformers.qwen2_vl import get_rope_index  # pyright: ignore[reportUnknownVariableType]
+
+        vision_pos = get_rope_index(
+            self.processor, input_ids=input_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask
+        )
+
+        valid_mask = attention_mask.bool()
+        text_pos = torch.zeros((1, len(input_ids)), dtype=torch.long, device=input_ids.device)
+        text_pos[0, valid_mask] = torch.arange(valid_mask.sum().item(), device=input_ids.device)
+
+        return torch.cat([text_pos, vision_pos], dim=0)
+
    def _start_proxy_server_v0(self):
        """
        Initializes and runs a Flask-based proxy server in a separate thread.
@@ -672,10 +748,14 @@ class AgentModeDaemon:
                continue

            # The client should report triplets that contain prompt_ids and response_ids.
-            # Example triplet.prompt: {"token_ids": [...]}
+            # Example triplet.prompt: {"token_ids": [...], "image_urls": [...]}
            # Example triplet.response: {"token_ids": [...]}
            trace_list = [
-                {"prompt_ids": t.prompt.get("token_ids", []), "response_ids": t.response.get("token_ids", [])}
+                {
+                    "prompt_ids": t.prompt.get("token_ids", []),
+                    "response_ids": t.response.get("token_ids", []),
+                    "image_urls": t.prompt.get("image_urls", []),
+                }
                for t in rollout.triplets
            ]
            info = {
@@ -705,6 +785,7 @@ class AgentModeDaemon:
        rollout_id_list: List[str] = []
        turn_index_list: List[int] = []
        is_drop_list: List[bool] = []
+        image_grid_thw_list: List[Optional[torch.Tensor]] = []  # For Qwen2-VL mrope
        n_trunc_sample_because_of_response = 0

        for rollout_id, sample_info in finished_id_to_sample_info.items():
@@ -741,6 +822,11 @@ class AgentModeDaemon:
                rollout_id_list.append(rollout_id)
                turn_index_list.append(turn_index)

+                # Compute image_grid_thw for this triplet using image_urls from prompt
+                if self._use_mrope:
+                    image_urls = trace.get("image_urls", [])
+                    image_grid_thw_list.append(self._get_image_grid_thw(image_urls))
+
        n_transition = len(input_ids_list)
        batch_input_ids = torch.LongTensor(input_ids_list).to(device)
        input_attention_mask = torch.LongTensor(input_attention_mask_list).to(device)
@@ -750,15 +836,38 @@ class AgentModeDaemon:
        # Concatenate prompts and responses to form the full sequence
        batch_seq = torch.cat([batch_input_ids, batch_response_ids], dim=-1)
        attention_mask = torch.cat([input_attention_mask, response_attention_mask], dim=-1)
-        position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
+
+        # Compute position_ids - use mrope for Qwen2-VL, standard 2D otherwise
+        if self._use_mrope:
+            # For Qwen2-VL: compute 4D position_ids (batch_size, 4, seq_length)
+            position_ids_list: list[torch.Tensor] = []
+            for i in range(n_transition):
+                pos_ids = self._compute_mrope_position_ids(
+                    input_ids=batch_seq[i],
+                    attention_mask=attention_mask[i],
+                    image_grid_thw=image_grid_thw_list[i] if image_grid_thw_list else None,
+                )  # (4, seq_length)
+                position_ids_list.append(pos_ids)
+            # Stack to (batch_size, 4, seq_length)
+            position_ids = torch.stack(position_ids_list, dim=0)
+        else:
+            # Standard 2D position_ids (batch_size, seq_length)
+            position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
+
        is_drop_mask = torch.BoolTensor(is_drop_list).to(device)
        scores = torch.tensor(reward_list, dtype=torch.bfloat16).to(device)

        # Create token-level scores by placing the final reward at the last token position
        token_level_scores = torch.zeros_like(attention_mask, dtype=scores.dtype)
+        # For mrope (3D position_ids), use the first dimension (text position_ids) for eos calculation
+        if self._use_mrope:
+            # position_ids is (batch_size, 4, seq_length), use first dim for text positions
+            text_position_ids = position_ids[:, 0, :]  # (batch_size, seq_length)
+            eos_mask_idx = torch.argmax(text_position_ids * attention_mask, dim=-1)  # (bsz,)
+        else:
+            eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1)  # (bsz,)
        # At the eos_mask_idx position of each sample, fill in the corresponding scores.
        # torch.arange(n_transition) generates [0,1,2,...,bsz-1] as indices for the batch dimension.
-        eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1)  # (bsz,)
        token_level_scores[torch.arange(n_transition), eos_mask_idx] = scores
        # Only take the last response_length part of the sequence to get the token-level scores for the model's response part.
        token_level_scores = token_level_scores[:, -max_response_length:]
@@ -1,13 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.

-# type: ignore
+# pyright: reportUnknownVariableType=false
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownArgumentType=false

-from importlib.metadata import version
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Type

 import hydra
 import ray
-from packaging import version as packaging_version
+from ray.actor import ActorClass
 from verl.trainer.main_ppo import create_rl_sampler
 from verl.trainer.ppo.reward import load_reward_manager

@@ -17,7 +20,10 @@ from agentlightning.store.base import LightningStore
 from agentlightning.types import Dataset

 from .dataset import AgentDataset, LoadedDataset
-from .trainer import AgentLightningTrainer
+
+if TYPE_CHECKING:
+    from .daemon import AgentModeDaemon
+    from .trainer import AgentLightningTrainer

 __all__ = [
    "main",
@@ -27,8 +33,20 @@ __all__ = [


@hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
-def main(config):
-    run_ppo(config, train_dataset=None, val_dataset=None, store=None, llm_proxy=None, adapter=None)
+def main(config: Any):
+    from .daemon import AgentModeDaemon
+    from .trainer import AgentLightningTrainer
+
+    run_ppo(
+        config,
+        train_dataset=None,
+        val_dataset=None,
+        store=None,
+        llm_proxy=None,
+        adapter=None,
+        trainer_cls=AgentLightningTrainer,
+        daemon_cls=AgentModeDaemon,
+    )


 def run_ppo(
@@ -38,6 +56,8 @@ def run_ppo(
    store: LightningStore | None,
    llm_proxy: LLMProxy | None,
    adapter: TraceAdapter[Any] | None,
+    trainer_cls: Type[AgentLightningTrainer],
+    daemon_cls: Type[AgentModeDaemon],
 ) -> None:
    if not ray.is_initialized():
        # this is for local ray cluster
@@ -56,13 +76,15 @@ def run_ppo(

    runner = TaskRunner.remote()
    ray.get(
-        runner.run.remote(
+        runner.run.remote(  # type: ignore
            config=config,
            train_dataset=train_dataset,
            val_dataset=val_dataset,
            store=store,
            llm_proxy=llm_proxy,
            adapter=adapter,
+            trainer_cls=trainer_cls,
+            daemon_cls=daemon_cls,
        )
    )

@@ -72,11 +94,13 @@ class TaskRunner:
    def run(
        self,
        config: Any,
-        train_dataset: Dataset | None,
-        val_dataset: Dataset | None,
+        train_dataset: Dataset[Any] | None,
+        val_dataset: Dataset[Any] | None,
        store: LightningStore | None,
        llm_proxy: LLMProxy | None,
-        adapter: TraceAdapter | None,
+        adapter: TraceAdapter[Any] | None,
+        trainer_cls: Type[AgentLightningTrainer],
+        daemon_cls: Type[AgentModeDaemon],
    ):
        # print initial config
        from pprint import pprint
@@ -91,7 +115,7 @@ class TaskRunner:
        local_path = copy_to_local(config.actor_rollout_ref.model.path)

        # instantiate tokenizer
-        from verl.utils import hf_processor, hf_tokenizer
+        from verl.utils.tokenizer import hf_processor, hf_tokenizer

        trust_remote_code = config.data.get("trust_remote_code", False)
        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
@@ -112,7 +136,8 @@ class TaskRunner:

        elif config.actor_rollout_ref.actor.strategy == "megatron":
            assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
-            from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup
+            # FIXME: This import is outdated
+            from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup  # type: ignore
            from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker

            actor_rollout_cls = ActorRolloutRefWorker
@@ -121,9 +146,16 @@ class TaskRunner:
        else:
            raise NotImplementedError

-        from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
+        from verl.trainer.ppo.ray_trainer import ResourcePoolManager

-        role_worker_mapping = {
+        try:
+            # verl >= 0.6.0
+            from verl.trainer.ppo.utils import Role
+        except ImportError:
+            # Fallback for verl <= 0.5.0
+            from verl.trainer.ppo.ray_trainer import Role  # type: ignore
+
+        role_worker_mapping: dict[Role, ActorClass[Any]] = {
            Role.ActorRollout: ray.remote(actor_rollout_cls),
            Role.Critic: ray.remote(CriticWorker),
        }
@@ -190,7 +222,7 @@ class TaskRunner:
            val_dataset = LoadedDataset(val_dataset)

        train_sampler = create_rl_sampler(config.data, train_dataset)
-        trainer = AgentLightningTrainer(
+        trainer = trainer_cls(
            config=config,
            tokenizer=tokenizer,
            processor=processor,
@@ -206,6 +238,7 @@ class TaskRunner:
            store=store,
            llm_proxy=llm_proxy,
            adapter=adapter,
+            daemon_cls=daemon_cls,
        )
        trainer.init_workers()
        trainer.fit()
@@ -8,7 +8,7 @@ import random
 from contextlib import contextmanager
 from copy import deepcopy
 from pprint import pprint
-from typing import Dict, Tuple
+from typing import Dict, Tuple, Type

 import numpy as np
 import torch
@@ -174,12 +174,18 @@ class AgentLightningTrainer(RayPPOTrainer):
    """

    def __init__(
-        self, store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter | None, **kwargs
+        self,
+        store: LightningStore | None,
+        llm_proxy: LLMProxy | None,
+        adapter: TraceAdapter | None,
+        daemon_cls: Type[AgentModeDaemon],
+        **kwargs,
    ):
        super().__init__(**kwargs)
        self.store = store
        self.llm_proxy = llm_proxy
        self.adapter = adapter
+        self.daemon_cls = daemon_cls

    def _validate(self):
        assert len(self.val_dataloader) == 1, "Please set val_batch_size to None for better throughput."
@@ -199,6 +205,37 @@ class AgentLightningTrainer(RayPPOTrainer):
        self.async_rollout_manager.sleep()
        return test_metrics

+    def _compute_reference_log_prob(self, batch: DataProto) -> DataProto:
+        """Compute reference log probability using the correct worker based on LoRA configuration.
+
+        In verl 0.6.0+, when LoRA is detected (indicated by ref_in_actor=True),
+        the reference policy is computed by the actor rollout worker instead of a separate
+        ref policy worker. This method handles both scenarios by checking the ref_in_actor flag.
+        Note: verl sets ref_in_actor=True when it detects LoRA configuration (e.g., lora_rank > 0 or lora_adapter_path is set).
+
+        Args:
+            batch: The data batch to compute reference log probabilities for.
+
+        Returns:
+            DataProto with reference log probabilities added.
+
+        Raises:
+            RuntimeError: If the required worker is not available.
+        """
+        if getattr(self, "ref_in_actor", False):
+            actor_worker = getattr(self, "actor_rollout_wg", None)
+            if actor_worker is None:
+                raise RuntimeError("actor_rollout_wg is required when ref_in_actor is True.")
+            return actor_worker.compute_ref_log_prob(batch)
+
+        ref_worker = getattr(self, "ref_policy_wg", None)
+        if ref_worker is None:
+            raise RuntimeError(
+                "Reference policy worker was not initialized. "
+                "Ensure `use_reference_policy` is enabled and the VERL config exposes the ref worker."
+            )
+        return ref_worker.compute_ref_log_prob(batch)
+
    def _train_step(self, batch_dict: dict) -> dict:
        # Isolate in a separate method to automatically recycle the variables before validation.
        batch: DataProto = DataProto.from_single_dict(batch_dict)
@@ -276,7 +313,7 @@ class AgentLightningTrainer(RayPPOTrainer):
            if self.use_reference_policy:
                # compute reference log_prob
                with _timer("ref", timing_raw):
-                    ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
+                    ref_log_prob = self._compute_reference_log_prob(batch)
                    batch = batch.union(ref_log_prob)

            # compute values
@@ -413,7 +450,7 @@ class AgentLightningTrainer(RayPPOTrainer):
        else:
            # For other versions (e.g., 0.6.0), we use the full path to the model.
            model = self.config.actor_rollout_ref.model.path
-        self.agent_mode_daemon = AgentModeDaemon(
+        self.agent_mode_daemon = self.daemon_cls(
            self.config.agentlightning.port,
            self.config.actor_rollout_ref.rollout.n,
            train_information={
@@ -427,6 +464,8 @@ class AgentLightningTrainer(RayPPOTrainer):
            store=self.store,
            llm_proxy=self.llm_proxy,
            adapter=self.adapter,
+            processor=self.processor,  # For Qwen2-VL mrope position_ids
+            image_base_dir=getattr(self.config.data, "image_base_dir", None),
        )
        self.agent_mode_daemon.start()

@@ -0,0 +1 @@
+# Put contrib-related gitignore files here.
@@ -0,0 +1,4 @@
+# Put code owner definitions here.
+
+# Recipes
+recipes/search_r1 @SiyunZhao @JiahangXu
@@ -0,0 +1,21 @@
+# Contrib Area
+
+This tree hosts experimental integrations, third-party recipes, and curated recipes that are not ready for the main `agentlightning/`, `examples/`, or `docs/` trees. Treat it as an incubator: keep contributions self-contained, clearly owned, and reproducible so downstream users can vendor them without guesswork.
+
+## When to add something here
+
+- You are iterating on a runtime extension that would bloat the primary `agentlightning/` namespace.
+- You want to share a recipe that assembles existing components for a focused agent training or optimization workflow and needs more context than the main examples directory allows.
+- You need automation scripts or download helpers that will help the community but should not live under `scripts/` at the repo root.
+
+If a contribution starts depending on core release cadence, tight CI guarantees, or repo-wide infrastructure, talk to maintainers about graduating it out of `contrib/`.
+
+## Directory map
+
+- `agentlightning/` — Namespace packages, utilities, and adapters that extend the published wheel. Place new code under `agentlightning/contrib/<feature>/` so `import agentlightning.contrib.<feature>` works for downstream users.
+- `recipes/` — Task-focused example bundles that solve a specific problem and derive certain results. Each recipe belongs in its own directory with a README that documents usage, result reports, and ownership.
+- `scripts/` — Shared automation, dataset download steps, or reproducibility helpers that support the contrib modules above.
+
+When adding folders, document the intent in a local README, link to companion docs or examples, and update `CODEOWNERS` so future fixes reach the right reviewers quickly.
+
+Questions or proposals for new subtrees can be discussed in Discord, GitHub issues, or GitHub Discussions before opening a PR. For the canonical requirements and review checklist, see the “Agent-lightning Contrib” section of [`docs/community/contributing.md`](../docs/community/contributing.md).
@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.
@@ -2,7 +2,7 @@

 ## Overview

-This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.1.x**.
+This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.1.2**.

 The example is designed to run on a single node with 8 GPUs, each having at least 40 GB of memory.

@@ -3,6 +3,8 @@
 import { useMemo, useState } from 'react';
 import type { Meta, StoryObj } from '@storybook/react';
 import { IconSearch } from '@tabler/icons-react';
+import { waitFor, within } from '@testing-library/dom';
+import userEvent from '@testing-library/user-event';
 import { Box, Stack, TextInput, Title } from '@mantine/core';
 import type { Span } from '@/types';
 import { compareRecords } from '@/utils/table';
@@ -570,4 +572,19 @@ const sequenceSortTestSpans: Span[] = [

 export const SequenceIdSortTest: Story = {
  render: () => <TracesTableStoryWrapper maxWidth={1200} spans={sequenceSortTestSpans} />,
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    const seqHeader = await canvas.findByRole('button', { name: /seq\./i });
+
+    await userEvent.click(seqHeader);
+
+    await waitFor(() => {
+      const rows = canvas.getAllByRole('row');
+      const firstRow = rows[1];
+      if (!firstRow) {
+        throw new Error('Expected at least one data row after sorting by sequence ID');
+      }
+      within(firstRow).getByText('task_sequence_14');
+    });
+  },
 };
@@ -5,7 +5,7 @@ services:
      - "--config.file=/etc/prometheus/prometheus.yml"
      - "--storage.tsdb.path=/prometheus"
    volumes:
-      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
      - ${AGL_MONITORING_DATA_PATH:?Set AGL_MONITORING_DATA_PATH to the metrics directory}/prometheus:/prometheus
    ports:
      - "9090:9090"
@@ -4,7 +4,27 @@ services:
      file: compose.store.yml
      service: app

-    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend memory
+    depends_on:
+      - app-exporter  # Wait for the exporter to be ready first
+
+    command: agl store --host 0.0.0.0 --port 4747 --tracker console prometheus --backend memory
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
+
+  app-exporter:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    command: agl prometheus --host 0.0.0.0 --port 4748
+    ports:
+      - "4748:4748"
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc

  node-exporter:
    image: prom/node-exporter:latest
@@ -22,10 +42,11 @@ services:
      - "--storage.tsdb.path=/prometheus"
      - "--storage.tsdb.retention.time=1h"
    volumes:
-      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
      - ./data/prometheus:/prometheus
    depends_on:
      - app
+      - app-exporter
      - node-exporter
    ports:
      - "9090:9090"
@@ -51,3 +72,10 @@ services:
      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
      - GF_AUTH_DISABLE_LOGIN_FORM=true
      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
+
+volumes:
+  prometheus_multiproc:
+    driver: local
+    driver_opts:
+      type: tmpfs
+      device: tmpfs
@@ -21,18 +21,33 @@ services:

    depends_on:
      - mongo
+      - app-exporter  # Wait for the exporter to be ready first

    command:
      - /bin/bash
      - -c
      - |
-          mkdir -p /tmp/prometheus &&
          agl store --host 0.0.0.0 --port 4747 \
-            --prometheus --backend mongo \
+            --tracker console prometheus --backend mongo \
            --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
            --n-workers ${AGL_STORE_N_WORKERS:-32}
    environment:
-      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
+
+  app-exporter:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    command: agl prometheus --host 0.0.0.0 --port 4748
+    ports:
+      - "4748:4748"
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc

  mongodb-exporter:
    image: percona/mongodb_exporter:0.47.1
@@ -61,10 +76,11 @@ services:
      - "--storage.tsdb.path=/prometheus"
      - "--storage.tsdb.retention.time=1h"
    volumes:
-      - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.mongo.yml:/etc/prometheus/prometheus.yml:ro
      - ./data/prometheus:/prometheus
    depends_on:
      - app
+      - app-exporter
      - mongodb-exporter
      - node-exporter
    ports:
@@ -91,3 +107,10 @@ services:
      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
      - GF_AUTH_DISABLE_LOGIN_FORM=true
      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
+
+volumes:
+  prometheus_multiproc:
+    driver: local
+    driver_opts:
+      type: tmpfs
+      device: tmpfs
@@ -1,11 +1,11 @@
 global:
-  scrape_interval: 2s
-  evaluation_interval: 2s
+  scrape_interval: 5s
+  evaluation_interval: 5s

 scrape_configs:
  - job_name: app
    static_configs:
-      - targets: ["app:4747"]
+      - targets: ["app-exporter:4748"]
    metrics_path: /v1/prometheus/

  - job_name: node
@@ -1,11 +1,11 @@
 global:
-  scrape_interval: 2s
-  evaluation_interval: 2s
+  scrape_interval: 5s
+  evaluation_interval: 5s

 scrape_configs:
  - job_name: app
    static_configs:
-      - targets: ["app:4747"]
+      - targets: ["app-exporter:4748"]
    metrics_path: /v1/prometheus/

  - job_name: node
@@ -46,6 +46,7 @@ Bonus points for examples that:
 - Ship CI or self-test coverage so we know they still work as the core evolves.  **Otherwise, we would have to mark the example as unmaintained because we won't be able to test the examples manually before each release.**
 - Include a [`docs/how-to/`]({{ src("docs/how-to/") }}) guide (or a detailed README if no how-to exists) without duplicating content in multiple places.
 - Favor simple, dependency-light code over heavy abstractions.
+- Ship a README that documents smoke-test instructions and includes an "Included Files" section summarizing every file and its role; keep the runnable module self-contained with a module-level docstring explaining CLI usage, plus targeted docstrings or inline comments for educational functions/classes.

 !!! warning "Please discuss first"

@@ -74,6 +75,27 @@ Most brand-new algorithms ultimately land as “new examples,” so read that se

 Have a project that builds on Agent-lightning but does not belong in the main repo? Fork it or depend on it externally, then let us know. We can showcase notable projects in [Community Projects](../index.md) and the main [README]({{ src("README.md") }}).

+### Agent-lightning Contrib
+
+[`contrib/`]({{ src("contrib") }}) is where work-in-progress or third-party integrations, and curated recipes live before they are hardened enough for the core runtime tree. Think of it as an incubator: additions should remain easy to consume, clearly owned, and scoped so downstream users can vendor them with minimal risk.
+
+The following types of contributions are welcome in the contrib area:
+
+- **Recipes** that assemble multiple Agent Lightning components for a narrow task (`contrib/recipes/<topic>/`). Each recipe must be self-contained, include running instructions and result reports.
+- **Runtime extensions** that would bloat the primary `agentlightning/` namespace (`contrib/agentlightning/contrib/<feature>/`). These should mirror the published wheel layout so that `import agentlightning.contrib.<feature>` works out of the box.
+- **Supporting scripts and assets** (`contrib/scripts/`) that automate dataset downloads, environment preparation, or benchmarks required by contrib modules.
+
+If you are unsure where a contribution should live, start a thread in Discord or open an issue before writing code. The [contrib README]({{ src("contrib/README.md") }}) also lists the directory expectations.
+
+A quick checklist for contributions to be accepted:
+
+1. **Document everything.** Include configuration steps, environment variables, and sample commands so contributors can reproduce the results without guesswork. Pin to a specific version of Agent-lightning and other dependencies to avoid unexpected changes if you don't want to update the recipe frequently.
+2. **Keep quality predictable.** Match the repo’s style guide, apply exhaustive type hints, and run `uv run --no-sync pyright` plus targeted `pytest` suites for any Python module you touch.
+3. **Ship reproducibility artifacts.** Store only scripts or instructions for downloading datasets, weights, or binaries. Never upload large artifacts or credentials directly.
+4. **Update ownership.** Add `CODEOWNERS` entries when new directories appear so maintainers know who can review follow-up fixes.
+
+Contrib entries do not need the same maturity level as core code, but they must still meet the baseline above. Submissions that lack documentation, hide ownership, or depend on untracked assets are typically rejected until those gaps are resolved.
+
 ### Other Contribution Ideas

 - **Tests.** Add or improve cases in [`tests/`]({{ src("tests") }}) (unit, integration, or end-to-end).
@@ -126,13 +148,13 @@ After `uv sync`, run commands via `uv run ...` (add `--no-sync` once the environ
 Formatting and linting are enforced through [pre-commit](https://pre-commit.com/). Install once, then run before each push:

 ```bash
-uv run pre-commit install
-uv run pre-commit run --all-files --show-diff-on-failure --color=always
+uv run --no-sync pre-commit install
+uv run --no-sync pre-commit run --all-files --show-diff-on-failure --color=always
 ```

 Once installed, the hooks run automatically on every `git commit`. Running the pre-commit hooks locally keeps CI green and diffs manageable.

-### 3. Branch From a Fresh `main`
+### 3. Branch from Fresh `main` and Code

 Start all work from the latest upstream state:

@@ -165,20 +187,28 @@ Use lowercase with hyphens, e.g., `feature/async-runner-hooks`.

    Remember to register new docs in [`mkdocs.yml`]({{ src("mkdocs.yml") }}), add examples to [examples/README]({{ src("examples/README.md") }}), and update the [Examples Catalog](../how-to/examples-catalog.md).

+Before you start coding, bring the shared coding conventions with you:
+
+- Target `requires-python >= 3.10`, four-space indentation, ~120-character lines (docstrings may run longer), and formatter-owned diffs (Black + isort with the `black` profile).
+- Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
+- Maintain exhaustive type hints (pyright enforces them), write succinct Google-style docstrings (with `[][]` cross-references).
+- Prefer dataclasses or Pydantic models from `agentlightning.types`.
+- Log via `logging.getLogger(__name__)` with targeted DEBUG/INFO/WARNING/ERROR calls—especially for long multi-step functions or broad `try/except` blocks.
+
 ### 4. Test and Validate

-Most contributions require automated checks. Prefix commands with `uv run` so they use the project environment.
+Most contributions require automated checks. Once `uv sync` locks dependencies, prefix commands with `uv run --no-sync ...` so they share the same environment as CI.

 **Full test suite**

 ```bash
-uv run pytest -v
+uv run --no-sync pytest -v
 ```

 **Targeted tests**

 ```bash
-uv run pytest tests/path/to/test_file.py -k test_name
+uv run --no-sync pytest tests/path/to/test_file.py -k test_name
 ```

 **Optional/gated tests:** GPU-specific suites or API-dependent tests run automatically when the required hardware or environment variables (such as `OPENAI_API_KEY`) are present.
@@ -186,7 +216,7 @@ uv run pytest tests/path/to/test_file.py -k test_name
 **Static analysis:**

 ```bash
-uv run pyright
+uv run --no-sync pyright
 ```

 If you have touched code under `examples/`, you should run the example-specific smoke tests. Each directory includes a README with example-specific smoke tests—run those too.
@@ -196,8 +226,8 @@ If you have touched code under `examples/`, you should run the example-specific
    Keep API references under [docs/reference]({{ src("docs/reference/") }}) up to date. Doc-only changes should still build cleanly:

    ```bash
-    uv run mkdocs serve --strict   # live reload
-    uv run mkdocs build --strict   # CI-equivalent
+    uv run --no-sync mkdocs serve --strict   # live reload
+    uv run --no-sync mkdocs build --strict   # CI-equivalent
    ```

    `--strict` elevates warnings to errors so you catch issues before CI.
@@ -205,7 +235,7 @@ If you have touched code under `examples/`, you should run the example-specific
 Before opening a PR, double-check the basics:

 - Run `uv lock` if you changed dependencies.
- Run `uv run pre-commit run --all-files` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
+- Run `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
 - Execute the relevant commands from the test list above.
 - Validate each affected example via its README instructions.

@@ -30,6 +30,14 @@

    [:octicons-repo-24: Browse source]({{ src("examples/calc_x") }})

+-   :material-chart-box:{ .lg .middle } __ChartQA vision-language RL__
+
+    ---
+
+    LangGraph-powered workflow for answering chart questions end to end: rollout the multi-modality agent with GPT or vLLM, and train with VERL/GRPO plus self-refinement loops.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/chartqa") }})
+
 -   :material-code-braces:{ .lg .middle } __Claude Code SWE-bench__

    ---
@@ -35,6 +35,8 @@ This documentation is organized into the following parts:

 - [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
 - [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
+- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl).
+

 ## Citation

@@ -14,17 +14,18 @@
 ## agl

 ```text
-usage: agl [-h] {vllm,store,agentops}
+usage: agl [-h] {vllm,store,prometheus,agentops}

 Agent Lightning CLI entry point.

 Available subcommands:
-  vllm      Run the vLLM CLI with Agent Lightning instrumentation.
-  store     Run a LightningStore server.
-  agentops  Start the AgentOps server manager.
+  vllm        Run the vLLM CLI with Agent Lightning instrumentation.
+  store       Run a LightningStore server.
+  prometheus  Serve Prometheus metrics from the multiprocess registry.
+  agentops    Start the AgentOps server manager.

 positional arguments:
-  {vllm,store,agentops}
+  {vllm,store,prometheus,agentops}
                        Subcommand to run.

 options:
@@ -73,6 +74,26 @@ options:
  --port PORT  Port to run the server on
 ```

+## agl prometheus
+
+Expose the Prometheus multiprocess registry on a dedicated FastAPI server. This is useful when the main LightningStore service is under heavy load; exporters can scrape this auxiliary endpoint instead.
+
+```text
+usage: agl prometheus [-h] [--host HOST] [--port PORT] [--metrics-path METRICS_PATH] [--log-level {DEBUG,INFO,WARNING,ERROR}] [--access-log]
+
+Serve Prometheus metrics outside the LightningStore server.
+
+options:
+  -h, --help            show this help message and exit
+  --host HOST           Host to bind the metrics server to.
+  --port PORT           Port to expose the Prometheus metrics on.
+  --metrics-path METRICS_PATH
+                        HTTP path used to expose metrics. Must start with '/' and not be the root path.
+  --log-level {DEBUG,INFO,WARNING,ERROR}
+                        Configure the logging level for the metrics server.
+  --access-log          Enable uvicorn access logs. Disabled by default to reduce noise.
+```
+
 ## agl agentops

 Start a mock AgentOps server to bypass the online service of AgentOps.
@@ -56,6 +56,20 @@

 ## Utilities

+::: agentlightning.utils.metrics.MetricsBackend
+
+::: agentlightning.utils.metrics.ConsoleMetricsBackend
+
+::: agentlightning.utils.metrics.PrometheusMetricsBackend
+
+::: agentlightning.utils.metrics.MultiMetricsBackend
+
+::: agentlightning.utils.metrics.setup_multiprocess_prometheus
+
+::: agentlightning.utils.metrics.get_prometheus_registry
+
+::: agentlightning.utils.metrics.shutdown_metrics
+
 ::: agentlightning.utils.server_launcher.PythonServerLauncher

 ::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
@@ -2,17 +2,34 @@

 This catalog highlights the examples shipped with Agent-lightning.

+Community-contributed examples and recipes are available in the [contrib](../contrib) directory.
+
 | Example | Description | CI Maintenance |
 |---------|-------------|----------------|
 | [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) |
+| [chartqa](./chartqa) | Vision-language ChartQA agent that reasons over charts with LangGraph and VERL plus multi-step self-refinement. | [![chartqa workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-chartqa.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.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 |
+| [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | [![rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml) |
 | [spider](./spider) | Text-to-SQL reinforcement learning training on the Spider dataset using LangGraph. | [![spider workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-spider.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml) |
 | [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | [![tinker workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-tinker.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml) |
 | [unsloth](./unsloth) | Supervised fine-tuning example powered by Unsloth with 4-bit quantization and LoRA. | [![unsloth workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unsloth.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-unsloth.yml) |

-*NOTE: CI status avoids taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
+## `examples-*` workflow status
+
+CI status above avoids taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. The following table displays the raw `examples-*` workflow status whenever the project is maintained by CI.
+
+| Workflow | Status |
+|----------|--------|
+| `examples-apo.yml` | [![examples-apo workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml) |
+| `examples-azure.yml` | [![examples-azure workflow 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) |
+| `examples-calc-x.yml` | [![examples-calc-x workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-calc-x.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-calc-x.yml) |
+| `examples-chartqa.yml` | [![examples-chartqa workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml) |
+| `examples-claude-code.yml` | [![examples-claude-code workflow 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) |
+| `examples-compat.yml` | [![examples-compat workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml) |
+| `examples-rag.yml` | [![examples-rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml) |
+| `examples-spider.yml` | [![examples-spider workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml) |
+| `examples-tinker.yml` | [![examples-tinker workflow 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) |
+| `examples-unsloth.yml` | [![examples-unsloth workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-unsloth.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-unsloth.yml) |
@@ -117,6 +117,9 @@ def train(
    ci_fast: bool,
    n_runners: int,
    external_store_address: str,
+    lora: bool,
+    lora_rank: int,
+    lora_adapter_path: Optional[str],
 ):
    """The training entrypoint function for Calc-X agent with VERL algorithm.

@@ -129,6 +132,9 @@ def train(
        n_runners: The number of runners for the Trainer.
        ci_fast: Whether to cap the training loop at a single step (implies CI toggles).
        external_store_address: Connects to an external store instead of creating a new one in memory.
+        lora: Whether to enable LoRA training.
+        lora_rank: LoRA rank to use when LoRA is enabled.
+        lora_adapter_path: Optional path to a pre-trained LoRA adapter to load.
    """
    # Load datasets (respect CLI file paths)
    train_dataset = cast(agl.Dataset[MathProblem], HuggingFaceDataset.from_parquet(train_file).to_list())  # type: ignore
@@ -144,6 +150,15 @@ def train(
    if model:
        config["actor_rollout_ref"]["model"]["path"] = model

+    # Enable LoRA configuration if requested
+    if lora:
+        config["actor_rollout_ref"]["model"]["lora_rank"] = lora_rank
+        print(f"LoRA enabled: lora_rank={lora_rank}")
+        if lora_adapter_path:
+            config["actor_rollout_ref"]["model"]["lora_adapter_path"] = lora_adapter_path
+            print(f"Loading LoRA adapter from: {lora_adapter_path}")
+        print("LoRA configuration will trigger verl to set ref_in_actor=True (LoRA mode)")
+
    # CI toggle keeps everything else the same but you can tweak the lightweight bits here if desired
    if ci or ci_fast:
        # Config the experiment name and project name so that they are available to CI
@@ -218,6 +233,23 @@ def main():
        help="Connect to an external store instead of creating a new one in memory",
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
+    parser.add_argument(
+        "--lora",
+        action="store_true",
+        help="Enable LoRA training. When enabled, the reference policy is computed by the actor rollout worker.",
+    )
+    parser.add_argument(
+        "--lora-rank",
+        type=int,
+        default=32,
+        help="LoRA rank to use when --lora is enabled (default: 32)",
+    )
+    parser.add_argument(
+        "--lora-adapter-path",
+        type=str,
+        default=None,
+        help="Optional path to a pre-trained LoRA adapter to load when --lora is enabled",
+    )

    args = parser.parse_args()

@@ -243,6 +275,9 @@ def main():
        ci_fast=args.ci_fast,
        n_runners=args.n_runners,
        external_store_address=args.external_store_address,
+        lora=args.lora,
+        lora_rank=args.lora_rank,
+        lora_adapter_path=args.lora_adapter_path,
    )


@@ -0,0 +1,113 @@
+# ChartQA Example
+
+[![chartqa workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-chartqa.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml)
+
+This example demonstrates training a visual reasoning agent on the ChartQA dataset using Agent-Lightning with the VERL algorithm and LangGraph framework. The agent answers questions about charts through a multi-step workflow with self-refinement. It's compatible with Agent-lightning v0.3.0 or later.
+
+## Requirements
+
+This example requires a single node with at least one 40GB GPU. Install dependencies with:
+
+```bash
+uv sync --frozen \
+    --group dev \
+    --group experiment \
+    --group image \
+    --group langchain \
+    --group vllm-0-10-2 \
+    --group torch-gpu-stable
+```
+
+**Currently vLLM 0.10.2 is the only tested version. You might see issues like `cu_seqlens_q must be on CUDA` or flash-attn installation failures if you use other versions.** (See https://github.com/vllm-project/vllm/issues/27340)
+
+## Dataset
+
+Download the ChartQA dataset and prepare it for training:
+
+```bash
+cd examples/chartqa
+python prepare_data.py
+```
+
+This downloads the ChartQA dataset from HuggingFace (`HuggingFaceM4/ChartQA`), saves images locally, and creates parquet files for training/testing. No HuggingFace token is required (public dataset).
+
+**Dataset Statistics:**
+
+- Training: ~18,000 chart question-answer pairs
+- Test: ~2,500 pairs
+- Chart types: Bar, line, pie, scatter, etc.
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `chartqa_agent.py` | Chart reasoning agent using LangGraph with multi-step workflow (observe → extract → calculate → check → refine) |
+| `train_chartqa_agent.py` | Training script using VERL algorithm with configurable hyperparameters (debug, qwen) |
+| `debug_chartqa_agent.py` | Debugging script to test the agent with cloud APIs or a local vLLM proxy |
+| `prepare_data.py` | Script to download ChartQA dataset from HuggingFace and prepare parquet files |
+| `prompts.py` | Prompt templates for the agent workflow |
+| `multimodal_utils.py` | Utility functions for encoding images to base64 |
+| `env_var.py` | Environment variables and configurations |
+| `data/` | Directory containing images and parquet files after download |
+
+## Running Examples
+
+### Debugging with Cloud API (Default)
+
+For quick testing with OpenAI or other cloud APIs (no local GPU required):
+
+```bash
+export OPENAI_API_KEY=<your-api-key>
+python debug_chartqa_agent.py
+```
+
+For other providers (Azure, etc.), set `OPENAI_API_BASE`:
+
+```bash
+export OPENAI_API_BASE=https://your-resource.openai.azure.com/v1
+export OPENAI_MODEL=gpt-4o
+python debug_chartqa_agent.py
+```
+
+### Debugging with Local Model (LLMProxy)
+
+To test the agent with a local vLLM server and LLMProxy:
+
+```bash
+# Start a vLLM server (specify image path for VLM)
+export CHARTQA_DATA_DIR=<path to chartqa data>
+vllm serve Qwen/Qwen2-VL-2B-Instruct \
+    --gpu-memory-utilization 0.6 \
+    --max-model-len 4096 \
+    --allowed-local-media-path $CHARTQA_DATA_DIR \
+    --enable-prefix-caching \
+    --port 8088
+
+# Run the agent with LLMProxy
+USE_LLM_PROXY=1 \
+    OPENAI_API_BASE=http://localhost:8088/v1 \
+    OPENAI_MODEL=Qwen/Qwen2-VL-2B-Instruct \
+    python debug_chartqa_agent.py
+```
+
+### Training with Local Model
+
+```bash
+python train_chartqa_agent.py debug --n-runners 2
+```
+
+You can also use an external store server (recommended for distributed setups), first start the store:
+
+```bash
+agl store --port 4747
+```
+
+Then run the training script with the external store address:
+
+```bash
+AGL_MANAGED_STORE=0 python train_chartqa_agent.py qwen --external-store-address http://localhost:4747
+```
+
+If you want to track experiments with Weights & Biases, set the `WANDB_API_KEY` environment variable before training.
+
+The script automatically launches agent workers and the training server. The agent workers execute chart reasoning rollouts using the vision-language model, while the training server applies the VERL algorithm (GRPO) to improve the model based on answer accuracy rewards.
@@ -0,0 +1,410 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""ChartQA agent demonstrating LangGraph-based visual reasoning with refinement.
+
+This module defines `ChartQAAgent` plus the supporting prompt utilities used by
+`debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+
+1. `analyze_chart` observes and summarizes the chart.
+2. `extract_data` calls a text-only LLM to extract the requested values.
+3. `calculate_answer` runs calculations grounded in prior steps.
+4. `check_answer` verifies reasoning quality.
+5. `refine_answer` conditionally patches mistakes before responding.
+
+Example usage can be found in `debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from typing import Any, Dict, Literal, cast
+
+import env_var as chartqa_env_var
+import termcolor
+from langchain.chat_models import BaseChatModel, init_chat_model
+from langchain_core.messages import AnyMessage, BaseMessage, HumanMessage
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from multimodal_utils import encode_image_to_base64
+from prompts import (
+    ANALYZE_CHART_PROMPT,
+    CALCULATE_ANSWER_PROMPT,
+    CHECK_ANSWER_PROMPT,
+    EXTRACT_DATA_PROMPT,
+    REFINE_ANSWER_PROMPT,
+)
+
+import agentlightning as agl
+
+logger = logging.getLogger("chartqa_agent")
+
+
+class ChartState(MessagesState):
+    question: str
+    image_path: str
+    observation: str
+    extracted_data: str
+    calculation: str
+    answer: str
+    feedback: str
+    num_turns: int
+    messages: list[AnyMessage]
+
+
+class ChartQAAgent(agl.LitAgent[Dict[str, Any]]):
+    """LangGraph-powered ChartQA agent with multi-step reasoning and refinement.
+
+    The implementation shares the same [`agl.LitAgent`][agentlightning.LitAgent] interface as
+    the Calc-X sample agent but augments it with image handling and LangGraph state tracking.
+    """
+
+    def __init__(
+        self,
+        model_name: str | None = None,
+        max_turns: int = 3,
+        debug: bool = False,
+        endpoint: str | None = None,
+        temperature: float = 0.0,
+        use_base64_images: bool = False,
+    ):
+        self.debug = debug
+        self.max_turns = max_turns
+        self.use_base64_images = use_base64_images
+        self.model_name = model_name
+        self.endpoint = endpoint
+        self.temperature = temperature
+
+        self._llm: BaseChatModel | None = None
+        self._graph: CompiledStateGraph[ChartState] | None = None
+
+    def _create_llm(self) -> BaseChatModel:
+        if self.model_name is None:
+            raise ValueError("model_name is required for creating LLM")
+        return init_chat_model(
+            self.model_name,
+            model_provider="openai",
+            openai_api_base=self.endpoint,
+            openai_api_key=chartqa_env_var.OPENAI_API_KEY,
+            temperature=self.temperature,
+            max_retries=2,
+            max_tokens=1024,
+            timeout=300,
+        )
+
+    def update_llm_config(self, model_name: str, endpoint: str | None, temperature: float | None) -> None:
+        """Update the LLM configuration. Re-create the LLM if the configuration is changed."""
+        updated: bool = False
+        if model_name != self.model_name:
+            self.model_name = model_name
+            updated = True
+        if endpoint != self.endpoint:
+            self.endpoint = endpoint
+            updated = True
+        if temperature != self.temperature:
+            self.temperature = temperature
+            updated = True
+        if updated:
+            self._llm = self._create_llm()
+
+    def _ensure_llm(self) -> BaseChatModel:
+        """Ensure the LLM is created and cached."""
+        if self._llm is None:
+            self._llm = self._create_llm()
+        return self._llm
+
+    def invoke_prompt(self, prompt: Any) -> AnyMessage:
+        """Invoke LLM with prompt."""
+        if self.debug:
+            for message in prompt.messages:
+                termcolor.cprint(message.pretty_repr(), "blue")
+
+        try:
+            result = self._ensure_llm().invoke(prompt)
+        except Exception as e:
+            logger.error(f"Failed to invoke prompt: {e}")
+            result = self._ensure_llm().invoke([HumanMessage(content="Please provide a reasonable answer.")])
+
+        if self.debug:
+            termcolor.cprint(result.pretty_repr(), "green")
+
+        return result  # type: ignore
+
+    def invoke_prompt_with_image(self, prompt_text: str, image_path: str) -> str:
+        """Invoke vision-language model with image.
+
+        Handles both local vLLM (file:// URLs) and cloud APIs (base64 encoding).
+        Cloud APIs (OpenAI, Anthropic, Google, Azure, etc.) require base64 encoding.
+        """
+        # Determine image URL format based on endpoint
+        if self.use_base64_images:
+            # Cloud APIs require base64 encoding for local files
+            image_url = encode_image_to_base64(image_path)
+        else:
+            # Local vLLM supports file:// URLs
+            if not image_path.startswith("file://"):
+                image_path = f"file://{os.path.realpath(image_path)}"
+            image_url = image_path
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt_text},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ]
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Call] {prompt_text[:100]}...", "blue")
+
+        try:
+            result = self._ensure_llm().invoke(messages)
+            response = result.content if hasattr(result, "content") else str(result)  # type: ignore
+        except Exception as e:
+            logger.error(f"Failed to invoke VLM: {e}")
+            response = "<observe>Unable to analyze chart</observe>"
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Response] {response[:200]}...", "green")
+
+        return response  # type: ignore
+
+    def extract_content(self, text: str, tag: str) -> str:
+        """Extract content between XML-style tags."""
+        match = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
+        return match.group(1).strip() if match else ""
+
+    def analyze_chart(self, state: ChartState) -> ChartState:
+        """Step 1: Observe and describe the chart."""
+        prompt: Any = ANALYZE_CHART_PROMPT.invoke({"question": state["question"]})  # type: ignore
+        prompt_text = prompt.messages[1].content
+
+        result_text = self.invoke_prompt_with_image(prompt_text, state["image_path"])
+
+        observation = self.extract_content(result_text, "observe")
+        if not observation:
+            observation = result_text
+
+        return {  # type: ignore
+            **state,
+            "observation": observation,
+            "num_turns": 1,
+            "messages": [HumanMessage(content=result_text)],
+        }
+
+    def extract_data(self, state: ChartState) -> ChartState:
+        """Step 2: Extract specific data values."""
+        prompt: Any = EXTRACT_DATA_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        extracted_data = self.extract_content(result.content, "extract")  # type: ignore
+        if not extracted_data:
+            extracted_data = result.content  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,  # type: ignore
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def calculate_answer(self, state: ChartState) -> ChartState:
+        """Step 3: Calculate and provide answer."""
+        prompt: Any = CALCULATE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        calculation = self.extract_content(result.content, "calculate")  # type: ignore
+        answer = self.extract_content(result.content, "answer")  # type: ignore
+        if not answer:
+            answer = cast(str, result.content)  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "calculation": calculation,
+            "answer": answer,
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def check_answer(self, state: ChartState) -> ChartState:
+        """Step 4: Verify answer quality."""
+        prompt: Any = CHECK_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", "No calculation shown"),
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        if self.debug:
+            termcolor.cprint(f"[Check] {result.content}", "yellow")  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "feedback": result.content,  # type: ignore
+            "messages": [*state.get("messages", []), *prompt.messages, result],
+        }
+
+    def refine_answer(self, state: ChartState) -> ChartState:
+        """Step 5: Refine answer based on feedback."""
+        prompt: Any = REFINE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", ""),
+                "feedback": state["feedback"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+        content: str = result.content  # type: ignore
+
+        new_extracted = self.extract_content(content, "extract")
+        extracted_data = new_extracted if new_extracted else state["extracted_data"]
+
+        new_calculation = self.extract_content(content, "calculate")
+
+        new_answer = self.extract_content(content, "answer")
+        if not new_answer:
+            new_answer = content
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,
+            "calculation": new_calculation,
+            "answer": new_answer,
+            "num_turns": state.get("num_turns", 0) + 1,
+            "messages": [*prompt.messages, result],
+        }
+
+    def should_continue(self, state: ChartState) -> Literal[END, "refine_answer"]:  # type: ignore
+        """Determine if refinement is needed."""
+        if state["messages"] and isinstance(
+            state["messages"][-1], BaseMessage
+        ):  # pyright: ignore[reportUnnecessaryIsInstance]
+            last_message = state["messages"][-1]
+            if "THE ANSWER IS CORRECT" in last_message.content:  # type: ignore
+                if "THE ANSWER IS INCORRECT" in last_message.content:  # type: ignore
+                    correct_index = last_message.content.rfind("THE ANSWER IS CORRECT")  # type: ignore
+                    incorrect_index = last_message.content.rfind("THE ANSWER IS INCORRECT")  # type: ignore
+                    if correct_index > incorrect_index:
+                        return END
+                else:
+                    return END
+
+        if state.get("num_turns", 0) >= self.max_turns:
+            return END
+
+        return "refine_answer"
+
+    def graph(self) -> CompiledStateGraph[ChartState]:
+        """Build the workflow graph with refinement loop."""
+        # Check if the graph is already built
+        if self._graph is not None:
+            return self._graph
+
+        builder = StateGraph(ChartState)
+        builder.add_node(self.analyze_chart)  # type: ignore
+        builder.add_node(self.extract_data)  # type: ignore
+        builder.add_node(self.calculate_answer)  # type: ignore
+        builder.add_node(self.check_answer)  # type: ignore
+        builder.add_node(self.refine_answer)  # type: ignore
+
+        builder.add_edge(START, "analyze_chart")
+        builder.add_edge("analyze_chart", "extract_data")
+        builder.add_edge("extract_data", "calculate_answer")
+        builder.add_edge("calculate_answer", "check_answer")
+        builder.add_conditional_edges(
+            "check_answer",
+            self.should_continue,  # type: ignore
+        )
+        builder.add_edge("refine_answer", "extract_data")
+
+        self._graph = builder.compile()  # type: ignore
+        return self._graph
+
+    def rollout(self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout) -> float | None:
+        """AgentLightning wrapper for ChartQA agent."""
+
+        question = task["question"]
+
+        rollout = cast(agl.AttemptedRollout, rollout)
+        llm = cast(agl.LLM, resources["main_llm"])
+
+        image_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, task["image_path"])
+        ground_truth = task["answer"]
+
+        if not os.path.exists(image_path):
+            logger.error(f"Image {image_path} does not exist. Skipping.")
+            return None
+
+        # The new rollout could have a different endpoint or temperature.
+        # Update the LLM if necessary.
+        self.update_llm_config(
+            model_name=llm.model,
+            endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            temperature=llm.sampling_parameters.get("temperature", 0.0),
+        )
+
+        try:
+            handler = self.tracer.get_langchain_handler()
+            result = self.graph().invoke(  # type: ignore
+                {"question": question, "image_path": image_path},  # type: ignore
+                {"callbacks": [handler] if handler else [], "recursion_limit": 100},
+            )
+        except Exception as e:
+            error_msg = f"[Rollout {rollout.rollout_id}] Error during agent invocation: {e}"
+            logger.error(error_msg, exc_info=True)
+            # Return 0.0 as reward to indicate failure
+            return 0.0
+
+        predicted_answer = result["answer"]
+        reward = evaluate_answer(predicted_answer, ground_truth, raise_on_error=False)
+
+        return reward
+
+
+def evaluate_answer(predicted: str, ground_truth: str, raise_on_error: bool = False) -> float:
+    """Evaluate answer accuracy."""
+    try:
+        pred = predicted.lower().strip()
+        gt = ground_truth.lower().strip()
+
+        # Exact match
+        if pred == gt:
+            return 1.0
+
+        # Try numeric comparison
+        try:
+            pred_num = float(pred.replace(",", ""))
+            gt_num = float(gt.replace(",", ""))
+            if abs(pred_num - gt_num) / max(abs(gt_num), 1e-9) < 0.02:
+                return 1.0
+        except (ValueError, AttributeError):
+            pass
+
+        # Partial credit for substring match
+        if pred in gt or gt in pred:
+            return 0.5
+
+        return 0.0
+    except Exception as e:
+        if raise_on_error:
+            raise
+        logger.exception(f"Error evaluating answer: {e}")
+        return 0.0
@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Debugging helpers for the ChartQA agent.
+
+Example usage for OpenAI API:
+
+```bash
+python debug_chartqa_agent.py
+```
+
+Example usage for self-hosted model.
+
+```
+vllm serve Qwen/Qwen2-VL-2B-Instruct \
+    --gpu-memory-utilization 0.6 \
+    --max-model-len 4096 \
+    --allowed-local-media-path $CHARTQA_DATA_DIR \
+    --enable-prefix-caching \
+    --port 8088
+USE_LLM_PROXY=1 OPENAI_API_BASE=http://localhost:8088/v1 OPENAI_MODEL=Qwen/Qwen2-VL-2B-Instruct python debug_chartqa_agent.py
+```
+
+Ensure `CHARTQA_DATA_DIR` points to a directory with the prepared parquet file by running `python prepare_data.py` beforehand.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, cast
+
+import env_var as chartqa_env_var
+import pandas as pd
+from chartqa_agent import ChartQAAgent
+
+import agentlightning as agl
+
+logger = logging.getLogger("chartqa_agent")
+
+
+def create_llm_proxy_for_chartqa(vllm_endpoint: str, port: int = 8081) -> agl.LLMProxy:
+    """Create an LLMProxy configured for ChartQA with token ID capture.
+
+    Args:
+        vllm_endpoint: Base URL for the hosted vLLM server.
+        port: Local port where the proxy should listen.
+
+    Returns:
+        An [`LLMProxy`][agentlightning.LLMProxy] instance launched in a thread.
+    """
+    store = agl.LightningStoreThreaded(agl.InMemoryLightningStore())
+
+    llm_proxy = agl.LLMProxy(
+        port=port,
+        store=store,
+        model_list=[
+            {
+                "model_name": "Qwen/Qwen2-VL-2B-Instruct",
+                "litellm_params": {
+                    "model": "hosted_vllm/Qwen/Qwen2-VL-2B-Instruct",
+                    "api_base": vllm_endpoint,
+                },
+            }
+        ],
+        callbacks=["return_token_ids"],
+        launch_mode="thread",
+    )
+
+    return llm_proxy
+
+
+def debug_chartqa_agent(use_llm_proxy: bool = False) -> None:
+    """Debug the ChartQA agent against cloud APIs or a local vLLM proxy.
+
+    Args:
+        use_llm_proxy: When `True`, spin up an LLMProxy that points to a local vLLM endpoint.
+
+    Raises:
+        FileNotFoundError: If the prepared ChartQA parquet file is missing.
+    """
+    test_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "test_chartqa.parquet")
+
+    if not os.path.exists(test_data_path):
+        raise FileNotFoundError(f"Test data file {test_data_path} does not exist. Please run prepare_data.py first.")
+
+    df = pd.read_parquet(test_data_path).head(10)  # type: ignore
+    test_data = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+
+    model = chartqa_env_var.OPENAI_MODEL
+    endpoint = chartqa_env_var.OPENAI_API_BASE
+    logger.info(
+        "Debug data: %s samples, model: %s, endpoint: %s, llm_proxy=%s",
+        len(test_data),
+        model,
+        endpoint,
+        use_llm_proxy,
+    )
+
+    llm_endpoint = endpoint
+    trainer_kwargs: Dict[str, Any] = {}
+
+    if use_llm_proxy:
+        proxy_port = 8089
+        llm_proxy = create_llm_proxy_for_chartqa(endpoint, port=proxy_port)
+        trainer_kwargs["llm_proxy"] = llm_proxy
+        trainer_kwargs["n_workers"] = 2
+        llm_endpoint = f"http://localhost:{proxy_port}/v1"
+        agent = ChartQAAgent()
+    else:
+        trainer_kwargs["n_workers"] = 1
+        agent = ChartQAAgent(use_base64_images=True)
+
+    trainer = agl.Trainer(
+        initial_resources={
+            "main_llm": agl.LLM(
+                endpoint=llm_endpoint,
+                model=model,
+                sampling_parameters={"temperature": 0.0},
+            )
+        },
+        **trainer_kwargs,
+    )
+
+    trainer.dev(agent, test_data)
+
+
+if __name__ == "__main__":
+    agl.setup_logging(apply_to=["chartqa_agent"])
+    debug_chartqa_agent(use_llm_proxy=chartqa_env_var.USE_LLM_PROXY)
@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+
+__all__ = [
+    "CHARTQA_ROOT_DIR",
+    "CHARTQA_DATA_DIR",
+    "CHARTQA_IMAGES_DIR",
+    "USE_BASE64_IMAGES",
+    "USE_LLM_PROXY",
+    "OPENAI_API_BASE",
+    "OPENAI_API_KEY",
+    "OPENAI_MODEL",
+]
+
+CHARTQA_ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+CHARTQA_DATA_DIR = os.getenv("CHARTQA_DATA_DIR", os.path.realpath(os.path.join(CHARTQA_ROOT_DIR, "data")))
+
+CHARTQA_IMAGES_DIR = os.getenv("CHARTQA_IMAGES_DIR", os.path.realpath(os.path.join(CHARTQA_ROOT_DIR, "data", "images")))
+
+USE_BASE64_IMAGES = os.getenv("USE_BASE64_IMAGES", "false").lower() in ("1", "true", "yes")
+
+USE_LLM_PROXY = os.getenv("USE_LLM_PROXY", "false").lower() in ("1", "true", "yes")
+
+OPENAI_API_BASE = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "token-abc123")
+
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")
@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Multimodal support utilities for Agent Lightning.
+
+This module provides helper functions for working with multimodal agents,
+particularly for vision-language tasks.
+"""
+
+from __future__ import annotations
+
+import base64
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Union
+
+import requests
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+__all__ = [
+    "encode_image_to_base64",
+    "create_image_message",
+]
+
+
+def encode_image_to_base64(image: Union[str, Path, PILImage], max_size: int = 2048) -> str:
+    """
+    Encode an image to base64 string for multimodal LLM APIs.
+
+    Args:
+        image: Image source (file path, URL, or PIL Image object)
+        max_size: Maximum dimension for resizing
+
+    Returns:
+        Base64 encoded image string with data URI prefix
+
+    Raises:
+        ImportError: If PIL (Pillow) is not installed
+        TypeError: If image type is not supported
+
+    Examples:
+        >>> encoded = encode_image_to_base64("photo.jpg")
+        >>> encoded[:30]
+        'data:image/jpeg;base64,/9j/4A...'
+
+        >>> from PIL import Image
+        >>> img = Image.open("photo.jpg")
+        >>> encoded = encode_image_to_base64(img)
+    """
+    # Load image
+    if isinstance(image, (str, Path)):
+        image_str = str(image)
+        if image_str.startswith(("http://", "https://")):
+            response = requests.get(image_str, timeout=30)
+            response.raise_for_status()
+            img = Image.open(BytesIO(response.content))
+        else:
+            img = Image.open(image_str)
+    elif hasattr(image, "mode"):
+        # PIL Image object
+        img = image
+    else:
+        raise TypeError(f"Unsupported image type: {type(image)}")
+
+    # Convert to RGB
+    if img.mode == "RGBA":
+        background = Image.new("RGB", img.size, (255, 255, 255))
+        background.paste(img, mask=img.split()[3])
+        img = background
+    elif img.mode != "RGB":
+        img = img.convert("RGB")
+
+    # Resize if needed
+    if max(img.size) > max_size:
+        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+
+    # Encode
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG", quality=85)
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def create_image_message(text: str, image: Union[str, Path, PILImage], use_base64: bool = True) -> dict[str, Any]:
+    """
+    Create an OpenAI-compatible multimodal message.
+
+    Args:
+        text: The text prompt/question
+        image: Image source (path, URL, or PIL Image)
+        use_base64: If True, encode as base64; if False, use URL directly
+
+    Returns:
+        Message dict with role="user" and multimodal content
+
+    Examples:
+        >>> msg = create_image_message("What's in the image?", "photo.jpg")
+        >>> msg["role"]
+        'user'
+        >>> len(msg["content"])
+        2
+    """
+    content: list[dict[str, Any]] = [{"type": "text", "text": text}]
+
+    if isinstance(image, str) and image.startswith(("http://", "https://")) and not use_base64:
+        content.append({"type": "image_url", "image_url": {"url": image}})
+    else:
+        encoded = encode_image_to_base64(image)
+        content.append({"type": "image_url", "image_url": {"url": encoded}})
+
+    return {"role": "user", "content": content}
@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Prepare ChartQA dataset from HuggingFace for training."""
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pandas as pd
+from datasets import load_dataset  # pyright: ignore[reportUnknownVariableType]
+
+
+def prepare_chartqa():
+    """Download ChartQA and convert to parquet format."""
+    data_dir = Path("data")
+    images_dir = data_dir / "images"
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    dataset = load_dataset("HuggingFaceM4/ChartQA")
+
+    for split in ["train", "test"]:
+        tasks: List[Dict[str, Any]] = []
+        dataset_length = len(dataset[split])  # type: ignore
+        for idx, item in enumerate(dataset[split]):  # pyright: ignore[reportUnknownArgumentType]
+            if idx % 1000 == 0:
+                print(f"Processing {split} item {idx} (out of {dataset_length})")
+            image_filename = f"{split}_{idx:06d}.png"
+            image_path = images_dir / image_filename
+            if not image_path.exists():
+                item["image"].save(image_path)
+
+            tasks.append(
+                {
+                    "id": f"{split}_{idx}",
+                    "image_path": f"images/{image_filename}",
+                    "question": item["query"],
+                    "answer": str(item["label"]),
+                }
+            )
+
+        pd.DataFrame(tasks).to_parquet(data_dir / f"{split}_chartqa.parquet", index=False)  # type: ignore
+
+
+if __name__ == "__main__":
+    prepare_chartqa()
@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Prompts for ChartQA agent workflow."""
+
+from langchain_core.prompts import ChatPromptTemplate
+
+ANALYZE_CHART_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a visual reasoning expert analyzing charts and graphs.
+Given a chart image and a question, first carefully observe and describe the chart.
+
+Instructions:
+- Identify the chart type (bar chart, line chart, pie chart, scatter plot, etc.)
+- Note the axes labels and units (if applicable)
+- Describe the data series or categories shown
+- Observe key patterns, trends, or noteworthy values
+- Pay attention to legends, titles, and annotations
+
+## Output Format ##
+
+Provide your observation inside <observe> and </observe> tags.
+
+Example:
+<observe>
+Bar chart showing GDP of 5 countries. X-axis shows country names, Y-axis shows GDP in trillions of USD.
+Data values: USA appears highest at around 25, China second at around 20, followed by India, UK, and France.
+</observe>
+""".strip(),
+        ),
+        ("user", "Question: {question}"),
+    ]
+)
+
+EXTRACT_DATA_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+Based on your observation of the chart, extract the specific data values needed to answer the question.
+
+Instructions:
+- Extract only the data relevant to the question
+- Be precise with values (read carefully from the chart)
+- Include labels/categories with each value
+- Use appropriate units
+
+## Output Format ##
+
+Provide extracted data inside <extract> and </extract> tags.
+Format: Label1: Value1, Label2: Value2, ...
+
+Example:
+<extract>
+USA: 25, China: 20, India: 15, UK: 10, France: 8
+</extract>
+""".strip(),
+        ),
+        (
+            "user",
+            """Observation: {observation}
+
+Question: {question}
+
+Please extract the relevant data values.""",
+        ),
+    ]
+)
+
+CALCULATE_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+Using the extracted data, perform any necessary calculations to answer the question.
+
+Instructions:
+- Show your calculation steps clearly
+- Use correct mathematical operations
+- Pay attention to the question (average, sum, difference, maximum, etc.)
+- Provide a precise numerical answer if applicable
+- Keep the answer concise (typically 1-10 words)
+
+## Output Format ##
+
+Show calculation inside <calculate> and </calculate> tags (if needed).
+Provide final answer inside <answer> and </answer> tags.
+
+Example:
+<calculate>
+Average = (25 + 20 + 15 + 10 + 8) / 5 = 78 / 5 = 15.6
+</calculate>
+<answer>
+15.6
+</answer>
+""".strip(),
+        ),
+        (
+            "user",
+            """Extracted Data: {extracted_data}
+
+Question: {question}
+
+Please calculate and provide the answer.""",
+        ),
+    ]
+)
+
+CHECK_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a chart analysis expert with strong attention to detail.
+Review the answer for potential mistakes.
+
+Common mistakes to check:
+- Incorrect data extraction from chart (misread values)
+- Arithmetic errors in calculations
+- Misunderstanding the question type (average vs. sum vs. difference)
+- Wrong number of data points counted
+- Incorrect units or scale interpretation
+- Off-by-one errors
+
+## Chart Information ##
+
+Observation: {observation}
+Extracted Data: {extracted_data}
+
+## Output Format ##
+
+If any mistakes are found, list each error clearly.
+After listing mistakes (if any), conclude with **ONE** of the following exact phrases in all caps:
+- If mistakes are found: `THE ANSWER IS INCORRECT.`
+- If no mistakes are found: `THE ANSWER IS CORRECT.`
+
+DO NOT write the corrected answer in this response. You only need to report mistakes.
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {question}
+
+Current Answer: {answer}
+
+Calculation shown:
+{calculation}
+
+Please review this answer for correctness.""",
+        ),
+    ]
+)
+
+REFINE_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a chart analysis agent.
+The previous answer had errors. Based on the feedback, provide a corrected answer.
+
+Instructions:
+- Re-examine the chart observation carefully
+- Correct any data extraction errors by re-extracting if needed
+- Fix calculation mistakes
+- Address all points mentioned in the feedback
+
+## Chart Observation ##
+
+{observation}
+
+## Output Format ##
+
+If you need to re-extract data, provide it inside <extract> and </extract> tags.
+Show corrected calculation inside <calculate> and </calculate> tags.
+Provide corrected answer inside <answer> and </answer> tags.
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {question}
+
+## Previous Attempt ##
+
+Extracted Data: {extracted_data}
+Calculation: {calculation}
+Answer: {answer}
+
+## Feedback ##
+
+{feedback}
+
+Please provide the corrected answer.""",
+        ),
+    ]
+)
@@ -0,0 +1,215 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Training helper for ChartQA modeled VERL workflow.
+
+Example usage:
+
+```bash
+python train_chartqa_agent.py debug --n-runners 2
+```
+
+or:
+
+```bash
+AGL_MANAGED_STORE=0 python train_chartqa_agent.py qwen --external-store-address http://localhost:9999
+```
+
+Make sure to run `python prepare_data.py` so the parquet files referenced here exist.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional, cast
+
+import env_var as chartqa_env_var
+import pandas as pd
+from chartqa_agent import ChartQAAgent
+
+import agentlightning as agl
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
+
+RL_CONFIG: Dict[str, Any] = {
+    "algorithm": {"adv_estimator": "grpo", "use_kl_in_reward": False},
+    "data": {
+        "image_base_dir": chartqa_env_var.CHARTQA_IMAGES_DIR,
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 1024,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 1,
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+            "enable_prefix_caching": True,
+            "engine_kwargs": {"vllm": {"allowed_local_media_path": chartqa_env_var.CHARTQA_IMAGES_DIR}},
+        },
+        "actor": {
+            "ppo_mini_batch_size": 32,
+            "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, "optimizer_offload": True},
+        },
+        "ref": {"log_prob_micro_batch_size_per_gpu": 1, "fsdp_config": {"param_offload": True}},
+        "model": {
+            "path": "Qwen/Qwen2-VL-2B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": False,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "chartqa",
+        "nnodes": 1,
+    },
+}
+
+
+def config_ci() -> Dict[str, Any]:
+    """Return a CI-friendly RL config for ChartQA."""
+    # For CI testing, we need to set the experiment name and project name so that
+    # they are available to subsequent steps.
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    random_suffix = uuid.uuid4().hex[:8]
+    EXPERIMENT_NAME = f"chartqa_ci_{timestamp}_{random_suffix}"
+    PROJECT_NAME = "AgentLightningCI"
+    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")
+
+    config = deepcopy(RL_CONFIG)
+    config["data"]["train_batch_size"] = 16
+    config["trainer"]["n_gpus_per_node"] = 1
+    config["trainer"]["total_training_steps"] = 4
+    config["trainer"]["val_before_train"] = True
+    config["trainer"]["test_freq"] = 2
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    return config
+
+
+def config_debug() -> Dict[str, Any]:
+    """Return a short debugging config for smoke testing ChartQA training."""
+    config = deepcopy(RL_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.5
+    config["trainer"]["total_training_steps"] = 10
+    config["trainer"]["test_freq"] = 2
+    return config
+
+
+def config_qwen() -> Dict[str, Any]:
+    """Return a Qwen-focused config with validation before each epoch."""
+    config = deepcopy(RL_CONFIG)
+    config["trainer"]["val_before_train"] = True
+    config["trainer"]["n_gpus_per_node"] = 2
+    config["trainer"]["total_epochs"] = 2
+    config["trainer"]["test_freq"] = 32
+    return config
+
+
+def train(
+    config: Dict[str, Any],
+    train_data: agl.Dataset[Any],
+    val_data: agl.Dataset[Any],
+    external_store_address: str,
+    n_runners: int,
+    debug: bool,
+) -> None:
+    """Run VERL training for ChartQA.
+
+    Args:
+        config: VERL configuration produced by one of the helpers above.
+        train_data: Training dataset of ChartQA samples.
+        val_data: Validation dataset for periodic evaluation.
+        external_store_address: Optional address of an existing LightningStore to reuse.
+        n_runners: Number of runners passed to [`Trainer.fit`][agentlightning.Trainer.fit].
+        debug: Enables verbose logging tied to `--debug`.
+    """
+    agl.setup_logging(level="DEBUG" if debug else "INFO", apply_to=["agentlightning", __name__])
+    agent = ChartQAAgent()
+    algorithm = agl.VERL(config)
+
+    if external_store_address:
+        store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address)
+    else:
+        store = None
+
+    trainer = agl.Trainer(
+        n_runners=n_runners,
+        algorithm=algorithm,
+        store=store,
+    )
+
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main():
+    """Parse CLI arguments and kick off ChartQA training."""
+    agl.setup_logging(apply_to=["chartqa_agent"])
+    parser = argparse.ArgumentParser(description="Train ChartQA agent")
+    parser.add_argument("config", choices=["debug", "qwen", "ci"], help="Training configuration")
+    parser.add_argument("--n-runners", type=int, default=10, help="Number of runners for Trainer")
+    parser.add_argument(
+        "--external-store-address",
+        type=str,
+        default=None,
+        help="Connect to an external store instead of creating a new one in memory (e.g., http://localhost:4747)",
+    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
+    args = parser.parse_args()
+
+    if args.external_store_address:
+        print(f"Connecting to external store at: {args.external_store_address}")
+        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!"
+            )
+
+    CONFIGS = {
+        "debug": config_debug,
+        "qwen": config_qwen,
+        "ci": config_ci,
+    }
+
+    train_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "train_chartqa.parquet")
+    val_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "test_chartqa.parquet")
+
+    train_data = pd.read_parquet(train_data_path).to_dict(orient="records")  # type: ignore
+
+    if args.config in ["debug", "ci"]:
+        val_data = pd.read_parquet(val_data_path).sample(n=100, random_state=42).to_dict(orient="records")  # type: ignore
+    else:
+        val_data = pd.read_parquet(val_data_path).to_dict(orient="records")  # type: ignore
+
+    train(
+        config=CONFIGS[args.config](),
+        train_data=cast(agl.Dataset[Any], train_data),
+        val_data=cast(agl.Dataset[Any], val_data),
+        external_store_address=args.external_store_address,
+        n_runners=args.n_runners,
+        debug=args.debug,
+    )
+
+
+if __name__ == "__main__":
+    main()
@@ -17,6 +17,7 @@ endpoint binds to `0.0.0.0:9105`.
 from __future__ import annotations

 import argparse
+import asyncio
 import random
 import signal
 import sys
@@ -43,7 +44,7 @@ def _register_metrics(backend: MetricsBackend) -> None:
    )


-def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
+async def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
    statuses = ["200", "404", "500"]
    end_time = time.time() + duration
    random.seed(1337)
@@ -51,9 +52,9 @@ def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence
        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)
+        await backend.inc_counter("minimal_requests_total", labels={"operation": operation, "status": status})
+        await backend.observe_histogram("minimal_latency_seconds", value=latency, labels={"operation": operation})
+        await asyncio.sleep(0.25)


 def main() -> None:
@@ -85,7 +86,7 @@ def main() -> None:

    original_handler = signal.signal(signal.SIGINT, _handle_interrupt)
    try:
-        _emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"])
+        asyncio.run(_emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"]))
    finally:
        signal.signal(signal.SIGINT, original_handler)

@@ -1,5 +1,7 @@
 # RAG Agent Example

+[![rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml)
+
 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
@@ -20,6 +20,7 @@ from litellm.types.utils import ChoiceLogprobs as LitellmChoiceLogprobs
 from litellm.types.utils import Choices
 from litellm.types.utils import Message as LitellmMessage
 from litellm.types.utils import ModelResponse
+from litellm.types.utils import TopLogprob as LitellmTopLogprob
 from litellm.utils import custom_llm_setup
 from pydantic import TypeAdapter
 from tinker.types import ModelInput, SampleResponse, SamplingParams
@@ -104,10 +105,9 @@ class TinkerLLM(CustomLLM):
        """
        self.sampling_client = sampling_client

-    def _validate_messages(self, messages: Any) -> TypeGuard[List[TinkerMessage]]:
-        TypeAdapter(List[TinkerMessage]).validate_python(messages)
+    def _canonicalize_messages(self, messages: Any) -> List[TinkerMessage]:
+        return TypeAdapter(List[TinkerMessage]).validate_python(messages)
        # Exception will be raised if validation fails
-        return True

    def _validate_role(self, role: str) -> TypeGuard[Literal["assistant", "user", "system", "tool", "function"]]:
        if role not in ["assistant", "user", "system", "tool", "function"]:
@@ -115,13 +115,11 @@ class TinkerLLM(CustomLLM):
        return True

    def _parse_tool_call(self, tool_call: TinkerToolCall) -> ChatCompletionMessageToolCall:
-        if set(tool_call.keys()) != {"name", "args"}:
-            logger.warning(f"Found unexpected tool call keys: {tool_call.keys()}")
        return ChatCompletionMessageToolCall(
-            id=generate_id("tinker-tool-call-"),
+            id=tool_call.id or generate_id("tinker-tool-call-"),
            function={
-                "name": tool_call["name"],
-                "arguments": tool_call["args"],
+                "name": tool_call.function.name,
+                "arguments": tool_call.function.arguments,
            },
            type="function",
        )
@@ -150,10 +148,8 @@ class TinkerLLM(CustomLLM):
    def _prepare_model_input(self, **kwargs: Any) -> ModelInput:
        """LiteLLM messages -> Tinker ModelInput."""
        messages = kwargs.pop("messages", None)
-        if self._validate_messages(messages):
-            return self.renderer.build_generation_prompt(messages)
-        else:
-            assert False, "This should never happen"
+        canonical_messages = self._canonicalize_messages(messages)
+        return self.renderer.build_generation_prompt(canonical_messages)

    def _parse_response(self, model_input: ModelInput, response: SampleResponse) -> ModelResponse:
        """Tinker Response -> LiteLLM Response.
@@ -173,7 +169,8 @@ class TinkerLLM(CustomLLM):
                            token=token,
                            bytes=bytes,
                            logprob=logprob,
-                            top_logprobs=[],
+                            # NOTE: This top logprob is not the real top logprob. It's just used to fool the LiteLLM type validator.
+                            top_logprobs=[LitellmTopLogprob(token=token, bytes=bytes, logprob=logprob)],
                        )
                        for token, bytes, logprob in zip(token_strings, bytes_list, seq.logprobs)
                    ]
@@ -186,6 +183,7 @@ class TinkerLLM(CustomLLM):
                role = parsed_response["role"]
                if not self._validate_role(role):
                    assert False, "This should never happen"
+                # FIXME: The content should not be still there if tool call has been parsed.
                content = parsed_response["content"]
                # NOTE(yuge): I thought about adding this to make it more robust to empty responses,
                # but later I found it's a configuration error in my renderer. So I think it's better
@@ -6,15 +6,17 @@ with LiteLLM and Agent-lightning.
 It should be included in CI in future if we decided to maintain this example.
 """

+import argparse
 import asyncio
-from typing import cast
+import json
+from typing import Any, Awaitable, Callable, Dict, cast

 import openai
 import tinker
 from agl_tinker.llm import TinkerLLM
 from agl_tinker.rollout import reconstruct_transitions
 from rich.console import Console
-from tinker_cookbook.renderers import Qwen3Renderer
+from tinker_cookbook.renderers import Qwen3InstructRenderer
 from transformers import AutoTokenizer, PreTrainedTokenizer

 from agentlightning import (
@@ -30,13 +32,97 @@ from agentlightning.store import LightningStoreThreaded

 setup_logging(apply_to=["agl_tinker"])

+_tool_call_system_prompt = """
+You must call the provided tool once before responding to the user.

-async def test_tracer():
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{"name": "echo_text", "description": "Echo back any provided text.", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "Text to repeat back."}}, "required": ["text"]}}
+</tools>
+
+For each function call, return a json object with function name and args within <tool_call></tool_call> XML tags:
+<tool_call>
+{"name": <function-name>, "args": <args-json-object>}
+</tool_call>
+"""
+
+
+def _run_tool_call_roundtrip(client: openai.OpenAI, *, model_name: str) -> None:
+    """Force a tool call, parse the args, and feed back the tool result."""
+    prompt_messages: list[Dict[str, str]] = [
+        # FIXME: Currently the tool call definition needs to be hard-coded into the system prompt.
+        {"role": "system", "content": _tool_call_system_prompt},
+        {"role": "user", "content": "Use the tool to echo 'Agent Lightning loves tool calls'."},
+    ]
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=cast(Any, prompt_messages),
+        max_tokens=256,
+        temperature=0.0,
+        # tools=cast(Any, tools),
+        # tool_choice="auto",
+    )
+    print("First response:", response)
+    tool_calls = response.choices[0].message.tool_calls or []
+    if not tool_calls:
+        raise AssertionError("Model did not emit a tool call when forced to do so.")
+    tool_call = tool_calls[0]
+    if tool_call.type != "function" or tool_call.function is None:  # pyright: ignore[reportUnnecessaryComparison]
+        raise AssertionError("Unexpected tool call payload from model.")
+    arguments = tool_call.function.arguments or "{}"
+    tool_args = json.loads(arguments)
+    tool_result = tool_args.get("text", "")
+    followup_messages: list[Dict[str, Any]] = [
+        *prompt_messages,
+        {
+            "role": "assistant",
+            "content": "",  # FIXME: Content must be here to make validation happy.
+            "tool_calls": [
+                {
+                    "id": tool_call.id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_call.function.name,
+                        "arguments": tool_call.function.arguments,
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "name": tool_call.function.name,
+            "content": f"Echoed text: {tool_result}",
+        },
+    ]
+    followup_response = client.chat.completions.create(
+        model=model_name,
+        messages=cast(Any, followup_messages),
+        max_tokens=64,
+        temperature=0.5,
+    )
+    print("Followup response:", followup_response)
+
+
+def _run_text_completion(client: openai.OpenAI, *, model_name: str) -> None:
+    """Simple text-only completion to contrast the tool-call scenario."""
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=[{"role": "user", "content": "Hello world!"}],
+        max_tokens=20,
+        temperature=0.5,
+        top_p=0.9,
+        seed=11,
+    )
+    print(response)
+
+
+async def _run_tracer_test(*, use_tool_call: bool) -> None:
    console = Console()
    model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"

    tokenizer = cast(PreTrainedTokenizer, AutoTokenizer.from_pretrained(model_name))  # type: ignore
-    renderer = Qwen3Renderer(tokenizer)  # type: ignore
+    renderer = Qwen3InstructRenderer(tokenizer)  # type: ignore
    service_client = tinker.ServiceClient()
    sampling_client = service_client.create_sampling_client(base_model=model_name)
    tinker_llm = TinkerLLM(
@@ -54,6 +140,9 @@ async def test_tracer():
        launch_mode="thread",
    )

+    scenario = "tool-call" if use_tool_call else "text-only"
+    console.print(f"Running tracer test scenario: {scenario}")
+
    try:
        tracer = AgentOpsTracer()
        tracer.init()
@@ -64,24 +153,15 @@ async def test_tracer():
        await llm_proxy.start()
        console.print("LLM proxy started")

-        # client = openai.OpenAI(
-        #     base_url=f"http://localhost:4000/rollout/{rollout.rollout_id}/attempt/{rollout.attempt.attempt_id}",
-        #     api_key="dummy",
-        # )
        client = openai.OpenAI(base_url="http://localhost:4000/v1", api_key="dummy")

        async with tracer.trace_context(
-            name="test_llm", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+            name=f"test_llm_{scenario}", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
        ):
-            response = client.chat.completions.create(
-                model=model_name,
-                messages=[{"role": "user", "content": "Hello world!"}],
-                max_tokens=10,
-                temperature=0.5,
-                top_p=0.9,
-                seed=43,
-            )
-            print(response)
+            if use_tool_call:
+                _run_tool_call_roundtrip(client, model_name=model_name)
+            else:
+                _run_text_completion(client, model_name=model_name)
            emit_reward(8.0)

        print(f"Found {len(tracer.get_last_trace())} spans in the tracer")
@@ -105,13 +185,25 @@ async def test_tracer():
        console.print("LLM proxy stopped")


+async def test_tracer_text_only():
+    await _run_tracer_test(use_tool_call=False)
+
+
+async def test_tracer_tool_call():
+    await _run_tracer_test(use_tool_call=True)
+
+
+async def test_tracer():
+    await test_tracer_tool_call()
+
+
 async def test_llm_proxy():
    # FIXME: The llm proxy adapter needs some fixes to make this test work
    console = Console()
    model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"

    tokenizer = cast(PreTrainedTokenizer, AutoTokenizer.from_pretrained(model_name))  # type: ignore
-    renderer = Qwen3Renderer(tokenizer)  # type: ignore
+    renderer = Qwen3InstructRenderer(tokenizer)  # type: ignore
    service_client = tinker.ServiceClient()
    sampling_client = service_client.create_sampling_client(base_model=model_name)
    tinker_llm = TinkerLLM(
@@ -164,5 +256,19 @@ async def test_llm_proxy():
        console.print("LLM proxy stopped")


+CLI_VARIANTS: Dict[str, Callable[[], Awaitable[None]]] = {
+    "tracer-tool": test_tracer_tool_call,
+    "tracer-text": test_tracer_text_only,
+    "llm-proxy": test_llm_proxy,
+}
+
+
 if __name__ == "__main__":
-    asyncio.run(test_tracer())
+    parser = argparse.ArgumentParser(description="Manually run the async Tinker LLM integration tests.")
+    parser.add_argument(
+        "variant",
+        choices=sorted(CLI_VARIANTS.keys()),
+        help="Which async test to run.",
+    )
+    args = parser.parse_args()
+    asyncio.run(CLI_VARIANTS[args.variant]())  # type: ignore
@@ -11,7 +11,6 @@ dependencies = [
  "setproctitle",
  "flask",
  "agentops>=0.4.13",
-  "httpdbg",
  "uvicorn",
  "fastapi",
  "aiohttp",
@@ -39,6 +38,10 @@ verl = [
  "vllm>=0.8.4", # Due to interface change of ExternalZeroMQDistributedExecutor
 ]

+weave = [
+  "weave",
+]
+
 # Store-related dependencies.
 mongo = [
  "pymongo",
@@ -115,11 +118,21 @@ torch-legacy = [
  "torchvision==0.22.0",
  "transformers==4.53.3",
  "tokenizers>=0.21,<0.22",
-  "flash-attn==2.8.1",
  "vllm==0.9.2",
  "litellm[proxy]==1.74.15",
 ]

+# Specific vllm versions for compatibility testing or workarounds.
+# Use these when you need a pinned vllm version separate from torch groups.
+vllm-0-10-2 = [
+  {include-group = "torch-stable"},
+  "vllm==0.10.2",
+]
+vllm-0-11-0 = [
+  {include-group = "torch-stable"},
+  "vllm==0.11.0",
+]
+
 # Flash-attention must build with CUDA toolkit.
 # Use this instead of --group torch-stable --group torch-gpu
 torch-gpu-stable = [
@@ -157,10 +170,21 @@ trl = [
 tinker = [
  {include-group = "torch-stable"},
  "tinker>=0.2.2",
-  "tinker_cookbook",
+  "tinker_cookbook>=0.1.0",
  "wandb",
 ]

+# For Multi-modality supports.
+image = [
+  # NOTE: It's tied to vLLM 0.10.2 but it will blow the uv.lock file if we include it.
+  # {include-group = "vllm-0-10-2"},
+  "datasets",
+  "Pillow",
+  "pandas",
+  "pyarrow",
+  "qwen-vl-utils",
+]
+
 # Agent-related dependencies.
 autogen = [
  "autogen-agentchat",
@@ -232,6 +256,10 @@ conflicts = [
    { group = "torch-stable" },
    { group = "torch-legacy" },
  ],
+  [
+    { group = "vllm-0-10-2" },
+    { group = "vllm-0-11-0" },
+  ],
  # langchain >= 1.0 requires openai >= 1.109.1 (via langchain-openai),
  # but torch-legacy uses vllm==0.9.2 which requires openai<=1.90.0
  [
@@ -267,7 +295,6 @@ torch = [
  { index = "pytorch-cu128", group = "torch-cu128" },
  { index = "pytorch-cpu", group = "torch-cpu" },
 ]
-tinker_cookbook = { git = "https://github.com/thinking-machines-lab/tinker-cookbook" }

 [[tool.uv.index]]
 name = "pypi"
@@ -322,6 +349,9 @@ markers = [
  "agentops: tests that require AgentOps",
  "llmproxy: tests that require LiteLLM",
  "mongo: tests that require MongoDB",
+  "store: tests for agentlightning.store module",
+  "prometheus: tests that require Prometheus",
+  "utils: tests for utility functions",
 ]

 [tool.black]
@@ -9,7 +9,9 @@
    "agentlightning/algorithm/verl",
    "agentlightning/cli/vllm.py",
    "agentlightning/store/collection/mongo.py",
-    "agentlightning/store/mongo.py"
+    "agentlightning/store/mongo.py",
+    "agentlightning/tracer/weave.py",
+    "contrib/**"
  ],

  "pythonVersion": "3.12",
@@ -1,7 +1,7 @@
 {
  // Paths to check
  "include": ["agentlightning", "tests", "examples"],
-  "exclude": ["**/data", "**/assets"],
+  "exclude": ["**/data", "**/assets", "contrib/**"],

  // Lock Python version for consistent semantics
  "pythonVersion": "3.12",
@@ -0,0 +1,698 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import itertools
+import json
+from typing import Any, Dict, List, Optional
+
+import pytest
+
+from agentlightning.adapter.triplet import RewardMatchPolicy, TracerTraceToTriplet, TraceTree
+from agentlightning.types import Span
+from agentlightning.types.tracer import SpanNames
+from agentlightning.utils.otel import filter_and_unflatten_attributes
+
+_SEQ = itertools.count()
+
+
+def qwen_multimodal_attrs(response_id: str) -> Dict[str, Any]:
+    """Simplified attributes derived from the provided Qwen trace dump."""
+    prompt_content = json.dumps(
+        [
+            {"type": "text", "text": "Question: How many food items are shown in the bar graph?"},
+            {"type": "image_url", "image_url": {"url": "file:///root/test.png"}},
+        ]
+    )
+    return {
+        "gen_ai.request.type": "chat",
+        "gen_ai.system": "OpenAI",
+        "gen_ai.request.model": "Qwen/Qwen2-VL-2B-Instruct",
+        "gen_ai.request.temperature": 0.0,
+        "gen_ai.request.streaming": False,
+        "gen_ai.request.headers": "{'X-Stainless-Raw-Response': 'true'}",
+        "gen_ai.prompt.0.role": "user",
+        "gen_ai.prompt.0.content": prompt_content,
+        "gen_ai.response.id": response_id,
+        "gen_ai.response.model": "Qwen/Qwen2-VL-2B-Instruct",
+        "gen_ai.usage.total_tokens": 12,
+        "gen_ai.usage.prompt_tokens": 10,
+        "gen_ai.usage.completion_tokens": 2,
+        "gen_ai.completion.0.content": "The bar graph shows 10 food items.",
+        "gen_ai.completion.0.finish_reason": "stop",
+        "gen_ai.completion.0.role": "assistant",
+        # Shortened token arrays to keep the fixture readable.
+        "prompt_token_ids": (151644, 8948, 198, 2610),
+        "response_token_ids": (785, 3619, 4771),
+    }
+
+
+def gpt_multimodal_attrs(response_id: str) -> Dict[str, Any]:
+    """Simplified attributes derived from the provided GPT-4o trace dump."""
+    prompt_content = json.dumps(
+        [
+            {"type": "text", "text": "Question: How many food items are shown in the bar graph?"},
+            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,AAA..."}},
+        ]
+    )
+    prompt_filter_results = json.dumps(
+        [
+            {
+                "prompt_index": 1,
+                "content_filter_result": {"sexual": {"filtered": False, "severity": "safe"}},
+            },
+            {"prompt_index": 0, "content_filter_result": {}},
+        ]
+    )
+    completion_filter_results = json.dumps(
+        {
+            "hate": {"filtered": False, "severity": "safe"},
+            "violence": {"filtered": False, "severity": "safe"},
+        }
+    )
+    return {
+        "gen_ai.request.type": "chat",
+        "gen_ai.system": "OpenAI",
+        "gen_ai.request.model": "gpt-4.1-mini",
+        "gen_ai.request.temperature": 0.0,
+        "gen_ai.request.streaming": False,
+        "gen_ai.request.headers": "{'X-Stainless-Raw-Response': 'true'}",
+        "gen_ai.openai.system_fingerprint": "fp_3dcd5944f5",
+        "gen_ai.prompt.0.role": "user",
+        "gen_ai.prompt.0.content": prompt_content,
+        "gen_ai.prompt.prompt_filter_results": prompt_filter_results,
+        "gen_ai.response.id": response_id,
+        "gen_ai.response.model": "gpt-4.1-mini-2025-04-14",
+        "gen_ai.usage.total_tokens": 9,
+        "gen_ai.usage.prompt_tokens": 7,
+        "gen_ai.usage.completion_tokens": 2,
+        "gen_ai.completion.0.finish_reason": "stop",
+        "gen_ai.completion.0.role": "assistant",
+        "gen_ai.completion.0.content": "The bar graph shows 13 food items.",
+        "gen_ai.completion.0.content_filter_results": completion_filter_results,
+    }
+
+
+def make_span(
+    span_id: str,
+    name: str,
+    *,
+    parent_id: Optional[str],
+    start_time: float,
+    end_time: float,
+    attributes: Optional[Dict[str, Any]] = None,
+) -> Span:
+    return Span.from_attributes(
+        rollout_id="rollout-1",
+        attempt_id="attempt-1",
+        sequence_id=next(_SEQ),
+        trace_id="trace-1",
+        span_id=span_id,
+        parent_id=parent_id,
+        name=name,
+        attributes=attributes or {},
+        start_time=start_time,
+        end_time=end_time,
+    )
+
+
+def make_llm_span(
+    span_id: str,
+    *,
+    parent_id: str,
+    start: float,
+    end: float,
+    prompt_ids: Optional[List[int]] = None,
+    response_ids: Optional[List[int]] = None,
+    response_id: Optional[str] = None,
+    extra_attrs: Optional[Dict[str, Any]] = None,
+) -> Span:
+    attrs: Dict[str, Any] = {
+        "prompt_token_ids": prompt_ids or [],
+        "response_token_ids": response_ids or [],
+    }
+    if response_id is not None:
+        attrs["gen_ai.response.id"] = response_id
+    if extra_attrs:
+        attrs.update(extra_attrs)
+    return make_span(
+        span_id,
+        "openai.chat.completion",
+        parent_id=parent_id,
+        start_time=start,
+        end_time=end,
+        attributes=attrs,
+    )
+
+
+def reward_attributes(value: float) -> Dict[str, Any]:
+    return {
+        "agentops.task.output": json.dumps({"type": "reward", "value": value}),
+    }
+
+
+def make_trace_tree_root() -> TraceTree:
+    """Create a minimal trace tree root for helper-only tests."""
+    root_span = make_span(
+        span_id="trace-root",
+        name="agent.session",
+        parent_id=None,
+        start_time=0.0,
+        end_time=1.0,
+    )
+    return TraceTree(root_span.span_id, root_span)
+
+
+def test_trace_tree_from_spans_orders_children_and_agent_names():
+    root = make_span(
+        "root",
+        "agent.session",
+        parent_id=None,
+        start_time=0.0,
+        end_time=10.0,
+        attributes={"agent.name": "primary-agent"},
+    )
+    llm = make_llm_span(
+        "llm",
+        parent_id="root",
+        start=1.0,
+        end=2.0,
+        prompt_ids=[1, 2],
+        response_ids=[3, 4],
+        response_id="resp-1",
+    )
+
+    tree = TraceTree.from_spans([llm, root])
+
+    assert tree.id == "root"
+    assert [child.id for child in tree.children] == ["llm"]
+    assert tree.find_id("llm") is tree.children[0]
+    assert tree.names_tuple() == ("agent.session [primary-agent]", [("openai.chat.completion", [])])
+    as_json = tree.to_json()
+    assert as_json["children"][0]["span"]["name"] == "openai.chat.completion"
+
+
+def test_trace_tree_virtual_root_for_multiple_roots():
+    first_root = make_span("root-a", "agent.first", parent_id=None, start_time=0.0, end_time=5.0)
+    second_root = make_span("root-b", "agent.second", parent_id=None, start_time=5.0, end_time=9.0)
+
+    tree = TraceTree.from_spans([first_root, second_root])
+
+    assert tree.id == "virtual-root"
+    assert tree.span.name == "virtual-root"
+    assert tree.start_time == first_root.start_time
+    assert tree.end_time == second_root.end_time
+    assert {child.id for child in tree.children} == {"root-a", "root-b"}
+
+
+def test_trace_tree_handles_missing_parent_and_empty_input():
+    with pytest.raises(ValueError):
+        TraceTree.from_spans([])
+
+    orphan_child = make_span(
+        "child",
+        "agent.child",
+        parent_id="ghost-parent",
+        start_time=2.0,
+        end_time=4.0,
+        attributes={"agent.name": "nested"},
+    )
+    llm = make_llm_span(
+        "grandchild",
+        parent_id="child",
+        start=3.0,
+        end=3.5,
+        prompt_ids=[1],
+        response_ids=[2],
+        response_id="resp-nested",
+    )
+
+    tree = TraceTree.from_spans([llm, orphan_child])
+
+    assert tree.id == "ghost-parent"
+    assert tree.span.name == SpanNames.VIRTUAL.value
+    assert tree.span.rollout_id == orphan_child.rollout_id
+    assert [child.id for child in tree.children] == ["child"]
+    assert tree.children[0].children[0].id == "grandchild"
+
+
+def test_trace_tree_repair_hierarchy_moves_llm_span_under_agent():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "planner"},
+    )
+    llm = make_llm_span(
+        "llm",
+        parent_id="root",
+        start=2.0,
+        end=3.0,
+        prompt_ids=[42],
+        response_ids=[7],
+        response_id="resp-planner",
+    )
+
+    tree = TraceTree.from_spans([root, agent, llm])
+    assert any(child.id == "llm" for child in tree.children)
+
+    tree.repair_hierarchy()
+
+    assert not any(child.id == "llm" for child in tree.children)
+    agent_node = tree.find_id("agent")
+    assert agent_node is not None
+    assert [child.id for child in agent_node.children] == ["llm"]
+
+
+def test_trace_tree_to_trajectory_skips_empty_and_dedupes_llm_calls():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "primary-agent"},
+    )
+    first = make_llm_span(
+        "llm-1",
+        parent_id="agent",
+        start=2.0,
+        end=3.0,
+        prompt_ids=[1, 2],
+        response_ids=[3, 4],
+        response_id="resp-1",
+    )
+    duplicate = make_llm_span(
+        "llm-2",
+        parent_id="agent",
+        start=3.2,
+        end=3.8,
+        prompt_ids=[9],
+        response_ids=[8],
+        response_id="resp-1",
+    )
+    empty_tokens = make_llm_span(
+        "llm-3",
+        parent_id="agent",
+        start=4.0,
+        end=5.0,
+        prompt_ids=[],
+        response_ids=[],
+        response_id="resp-2",
+    )
+    reward = make_span(
+        "reward",
+        "agent.reward",
+        parent_id="agent",
+        start_time=6.0,
+        end_time=6.1,
+        attributes=reward_attributes(0.5),
+    )
+
+    tree = TraceTree.from_spans([root, agent, first, duplicate, empty_tokens, reward])
+
+    trajectory = tree.to_trajectory(
+        agent_match="primary-agent",
+        dedup_llm_call=True,
+        _skip_empty_token_spans=True,
+    )
+    assert len(trajectory) == 1
+    triplet = trajectory[0]
+    assert triplet.prompt["token_ids"] == [1, 2]
+    assert triplet.response["token_ids"] == [3, 4]
+    assert triplet.metadata["response_id"] == "resp-1"
+    assert triplet.metadata["agent_name"] == "primary-agent"
+    assert triplet.reward == 0.5
+
+    with_final_reward = tree.to_trajectory(
+        agent_match="primary-agent",
+        dedup_llm_call=True,
+        _skip_empty_token_spans=True,
+        final_reward=1.0,
+    )
+    assert len(with_final_reward) == 1
+    assert with_final_reward[0].reward == 1.0
+
+
+def test_tracer_trace_to_triplet_repair_required_for_agent_filter():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "planner"},
+    )
+    llm_outside_agent = make_llm_span(
+        "llm",
+        parent_id="root",
+        start=2.0,
+        end=3.0,
+        prompt_ids=[7],
+        response_ids=[8],
+        response_id="resp-planner",
+    )
+    reward = make_span(
+        "reward",
+        "agent.reward",
+        parent_id="agent",
+        start_time=4.0,
+        end_time=4.5,
+        attributes=reward_attributes(0.3),
+    )
+    spans = [root, agent, llm_outside_agent, reward]
+
+    adapter = TracerTraceToTriplet(agent_match="planner")
+    triplets = adapter.adapt(spans)
+    assert len(triplets) == 1
+    assert triplets[0].metadata["agent_name"] == "planner"
+    assert triplets[0].reward == 0.3
+
+    adapter_without_repair = TracerTraceToTriplet(repair_hierarchy=False, agent_match="planner")
+    assert adapter_without_repair.adapt(spans) == []
+
+
+def test_tracer_trace_to_triplet_dedup_and_skip_empty_token_spans():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "collector"},
+    )
+    kept_llm = make_llm_span(
+        "llm-1",
+        parent_id="agent",
+        start=2.0,
+        end=3.0,
+        prompt_ids=[10],
+        response_ids=[20],
+        response_id="resp-shared",
+    )
+    duplicate_llm = make_llm_span(
+        "llm-2",
+        parent_id="agent",
+        start=3.5,
+        end=4.2,
+        prompt_ids=[99],
+        response_ids=[98],
+        response_id="resp-shared",
+    )
+    missing_tokens = make_llm_span(
+        "llm-3",
+        parent_id="agent",
+        start=5.0,
+        end=5.5,
+        prompt_ids=[],
+        response_ids=[],
+        response_id="resp-3",
+    )
+    reward = make_span(
+        "reward",
+        "agent.reward",
+        parent_id="agent",
+        start_time=6.0,
+        end_time=6.5,
+        attributes=reward_attributes(0.25),
+    )
+    spans = [root, agent, kept_llm, duplicate_llm, missing_tokens, reward]
+
+    adapter = TracerTraceToTriplet(_skip_empty_token_spans=True)
+    triplets = adapter.adapt(spans)
+
+    assert len(triplets) == 1
+    assert triplets[0].prompt["token_ids"] == [10]
+    assert triplets[0].response["token_ids"] == [20]
+    assert triplets[0].metadata["response_id"] == "resp-shared"
+    assert triplets[0].reward == 0.25
+
+
+def test_trace_tree_find_llm_calls_dedupes_across_agents():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent_a = make_span(
+        "agent-a",
+        "agent.node",
+        parent_id="root",
+        start_time=0.5,
+        end_time=5.0,
+        attributes={"agent.name": "vision-a"},
+    )
+    agent_b = make_span(
+        "agent-b",
+        "agent.node",
+        parent_id="root",
+        start_time=5.1,
+        end_time=9.5,
+        attributes={"agent.name": "vision-b"},
+    )
+    shared_response_id = "chatcmpl-shared"
+    llm_a = make_span(
+        "llm-a",
+        "openai.chat.completion",
+        parent_id="agent-a",
+        start_time=1.0,
+        end_time=2.0,
+        attributes=qwen_multimodal_attrs(shared_response_id),
+    )
+    llm_b = make_span(
+        "llm-b",
+        "openai.chat.completion",
+        parent_id="agent-b",
+        start_time=6.0,
+        end_time=7.0,
+        attributes=gpt_multimodal_attrs(shared_response_id),
+    )
+
+    tree = TraceTree.from_spans([root, agent_a, agent_b, llm_a, llm_b])
+    matches = tree.find_llm_calls(
+        llm_call_match=r"openai\.chat\.completion",
+        agent_match=None,
+        within_matching_subtree="*",
+        within_reward=False,
+        within_llm_call=False,
+        existing_llm_call_response_ids=set(),
+    )
+
+    assert len(matches) == 1
+    assert matches[0][0].id == "llm-a"
+
+
+def test_tracer_trace_to_triplet_handles_multimodal_payloads():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=15.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=0.5,
+        end_time=14.5,
+        attributes={"agent.name": "vision-agent"},
+    )
+    llm_first = make_span(
+        "llm-qwen",
+        "openai.chat.completion",
+        parent_id="agent",
+        start_time=1.0,
+        end_time=2.0,
+        attributes=qwen_multimodal_attrs("chatcmpl-qwen"),
+    )
+    llm_second = make_span(
+        "llm-gpt",
+        "openai.chat.completion",
+        parent_id="agent",
+        start_time=10.0,
+        end_time=11.0,
+        attributes=gpt_multimodal_attrs("chatcmpl-gpt"),
+    )
+    reward = make_span(
+        "reward",
+        "agent.reward",
+        parent_id="agent",
+        start_time=12.0,
+        end_time=12.5,
+        attributes=reward_attributes(0.7),
+    )
+
+    assert llm_first.attributes["gen_ai.request.headers"] == "{'X-Stainless-Raw-Response': 'true'}"
+    assert llm_second.attributes["gen_ai.openai.system_fingerprint"] == "fp_3dcd5944f5"
+    assert "gen_ai.prompt.prompt_filter_results" in llm_second.attributes
+    qwen_prompt = json.loads(llm_first.attributes["gen_ai.prompt.0.content"])  # type: ignore
+    assert qwen_prompt[0]["type"] == "text"
+    assert qwen_prompt[1]["image_url"]["url"].startswith("file://")
+    gpt_prompt = json.loads(llm_second.attributes["gen_ai.prompt.0.content"])  # type: ignore
+    assert gpt_prompt[1]["image_url"]["url"].startswith("data:image/jpeg")
+    assert llm_first.attributes["gen_ai.completion.0.content"] == "The bar graph shows 10 food items."
+    assert llm_second.attributes["gen_ai.completion.0.content"] == "The bar graph shows 13 food items."
+
+    adapter = TracerTraceToTriplet(agent_match="vision-agent")
+    triplets = adapter.adapt([root, agent, llm_first, llm_second, reward])
+
+    assert len(triplets) == 2
+    first, second = triplets
+    assert list(first.prompt["token_ids"]) == [151644, 8948, 198, 2610]
+    assert list(first.response["token_ids"]) == [785, 3619, 4771]
+    assert first.metadata["response_id"] == "chatcmpl-qwen"
+    assert first.metadata["agent_name"] == "vision-agent"
+    assert first.reward is None
+
+    assert second.prompt["token_ids"] == []
+    assert second.response["token_ids"] == []
+    assert second.metadata["response_id"] == "chatcmpl-gpt"
+    assert triplets[0].metadata["agent_name"] == "vision-agent"
+    assert triplets[1].metadata["agent_name"] == "vision-agent"
+    qwen_prompt_raw = triplets[0].prompt["raw_content"]
+    assert qwen_prompt_raw == filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.prompt")
+    assert triplets[0].prompt["image_urls"] == ["file:///root/test.png"]
+    qwen_content = json.loads(qwen_prompt_raw[0]["content"])
+    assert qwen_content[1]["image_url"]["url"] == "file:///root/test.png"
+    qwen_request = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.request")
+    qwen_response = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.response")
+    assert triplets[0].metadata["request"] == qwen_request
+    assert triplets[0].metadata["response"] == qwen_response
+    qwen_completion = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.completion")
+    assert triplets[0].response["raw_content"] == qwen_completion
+
+    gpt_prompt_raw = triplets[1].prompt["raw_content"]
+    assert gpt_prompt_raw == filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.prompt")
+    gpt_content = json.loads(gpt_prompt_raw["0"]["content"])
+    assert gpt_content[1]["image_url"]["url"].startswith("data:image/jpeg")
+    assert triplets[1].prompt["image_urls"] == ["data:image/jpeg;base64,AAA..."]
+    gpt_request = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.request")
+    gpt_response = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.response")
+    assert triplets[1].metadata["request"] == gpt_request
+    assert triplets[1].metadata["response"] == gpt_response
+    gpt_completion = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.completion")
+    assert triplets[1].response["raw_content"] == gpt_completion
+    assert second.reward == 0.7
+
+
+def test_extract_prompt_image_urls_from_list_payload():
+    tree = make_trace_tree_root()
+    prompt_raw_content = [
+        {
+            "role": "user",
+            "content": json.dumps(
+                [
+                    {"type": "text", "text": "describe the image"},
+                    {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
+                ]
+            ),
+        },
+        {
+            "role": "assistant",
+            "content": json.dumps(
+                [
+                    {"type": "text", "text": "another prompt"},
+                    {"type": "image_url", "image_url": {"url": "https://example.com/b.png"}},
+                ]
+            ),
+        },
+    ]
+
+    image_urls = tree.extract_prompt_image_urls(prompt_raw_content)
+
+    assert image_urls == ["https://example.com/a.png", "https://example.com/b.png"]
+
+
+def test_tracer_trace_to_triplet_reward_match_first_sibling():
+    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
+    agent = make_span(
+        "agent",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "sibling-agent"},
+    )
+    other_agent = make_span(
+        "agent-2",
+        "agent.node",
+        parent_id="root",
+        start_time=1.0,
+        end_time=9.0,
+        attributes={"agent.name": "sibling-agent"},
+    )
+    llm_1 = make_llm_span(
+        "llm-1",
+        parent_id="agent",
+        start=2.0,
+        end=3.0,
+        prompt_ids=[1],
+        response_ids=[2],
+        response_id="resp-1",
+    )
+    reward = make_span(
+        "reward",
+        "agent.reward",
+        parent_id="agent",
+        start_time=3.5,
+        end_time=3.6,
+        attributes=reward_attributes(0.8),
+    )
+    llm_2 = make_llm_span(
+        "llm-2",
+        parent_id="agent-2",
+        start=3.1,
+        end=3.2,
+        prompt_ids=[3],
+        response_ids=[4],
+        response_id="resp-2",
+    )
+
+    spans = [root, agent, other_agent, llm_1, reward, llm_2]
+
+    adapter = TracerTraceToTriplet(
+        agent_match="sibling-agent",
+        reward_match=RewardMatchPolicy.FIRST_SIBLING,
+        _skip_empty_token_spans=True,
+    )
+    triplets = adapter.adapt(spans)
+
+    assert len(triplets) == 2
+    t1, t2 = triplets
+
+    assert t1.metadata["response_id"] == "resp-1"
+    assert t1.reward == 0.8
+
+    assert t2.metadata["response_id"] == "resp-2"
+    assert t2.reward is None
+
+
+def test_extract_prompt_image_urls_handles_numeric_dict_keys():
+    tree = make_trace_tree_root()
+    prompt_raw_content = {
+        "1": {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "second"},
+                {"type": "image_url", "image_url": {"url": "https://example.com/second.png"}},
+            ],
+        },
+        "0": {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": "https://example.com/first.png"}},
+                {"type": "text", "text": "first"},
+            ],
+        },
+    }
+
+    image_urls = tree.extract_prompt_image_urls(prompt_raw_content)
+
+    assert image_urls == ["https://example.com/first.png", "https://example.com/second.png"]
+
+
+def test_extract_prompt_image_urls_gracefully_handles_invalid_payloads():
+    tree = make_trace_tree_root()
+    invalid_prompt_content = [
+        {"role": "user", "content": "not-a-json"},
+        {"role": "assistant", "content": json.dumps([{"type": "text", "text": "no images here"}])},
+        {"role": "system"},
+    ]
+
+    assert tree.extract_prompt_image_urls(invalid_prompt_content) == []
+    assert tree.extract_prompt_image_urls("unexpected-string") == []
@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, cast
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.algorithm.utils import with_llm_proxy, with_store
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+
+
+class _BaseAlgorithm(Algorithm):
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Satisfy the abstract interface without invoking training logic."""
+        return None
+
+
+class _StubLLMProxy:
+    """Test double that tracks lifecycle calls."""
+
+    def __init__(self) -> None:
+        self.start_calls = 0
+        self.stop_calls = 0
+        self.running = False
+
+    def is_running(self) -> bool:
+        return self.running
+
+    async def start(self) -> None:
+        self.start_calls += 1
+        self.running = True
+
+    async def stop(self) -> None:
+        self.stop_calls += 1
+        self.running = False
+
+
+@pytest.mark.asyncio
+async def test_with_store_injects_store_argument():
+    class StoreAlgorithm(_BaseAlgorithm):
+        @with_store
+        async def record_store(self, store: LightningStore, payload: str) -> None:
+            self.seen_store = store  # type: ignore[attr-defined]
+            self.seen_payload = payload  # type: ignore[attr-defined]
+
+    algorithm = StoreAlgorithm()
+    fake_store = MagicMock(spec=LightningStore)
+    algorithm.set_store(fake_store)
+
+    await algorithm.record_store("batch-1")
+
+    assert algorithm.seen_store is fake_store  # type: ignore[attr-defined]
+    assert algorithm.seen_payload == "batch-1"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_allows_optional_injection():
+    class OptionalProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def record_proxy(self, llm_proxy: LLMProxy | None, marker: str) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+            self.marker = marker  # type: ignore[attr-defined]
+
+    algorithm = OptionalProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    await algorithm.record_proxy("optional")
+
+    assert algorithm.seen_proxy is None  # type: ignore[attr-defined]
+    assert algorithm.marker == "optional"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_required_raises_when_missing():
+    class RequiredProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy(required=True)
+        async def record_proxy(self, llm_proxy: LLMProxy) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = RequiredProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    with pytest.raises(ValueError):
+        await algorithm.record_proxy()
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_auto_start_and_stop():
+    class AutoProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def use_proxy(self, llm_proxy: LLMProxy | None) -> None:
+            if llm_proxy is None:
+                raise AssertionError("LLM proxy should be injected")
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = AutoProxyAlgorithm()
+    proxy = _StubLLMProxy()
+    algorithm.set_llm_proxy(cast(LLMProxy, proxy))
+
+    await algorithm.use_proxy()
+
+    assert algorithm.seen_proxy is proxy  # type: ignore[attr-defined]
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1
+
+    # When already running, no extra start/stop should be requested.
+    proxy.running = True
+    await algorithm.use_proxy()
+
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1
@@ -19,7 +19,7 @@ from .utils import flatten_dict, random_dict

 console = Console()

-MAX_RUNTIME_SECONDS = 45 * 60
+MAX_RUNTIME_SECONDS = 30 * 60


 def _abort_due_to_timeout() -> None:
@@ -157,11 +157,13 @@ class AlgorithmBatch(agl.Algorithm):

            pending = {rollout_id: task_name for rollout_id, task_name in batch_rollouts}
            completed_ids: Set[str] = set()
+            completed_ids_last_updated: int = 0
            while len(completed_ids) < len(batch_rollouts):
                finished_rollouts = await store.wait_for_rollouts(
                    rollout_ids=[rollout_id for rollout_id, _ in batch_rollouts],
                    timeout=0.0,
                )
+                complete_ids_updated: bool = False
                for rollout in finished_rollouts:
                    rollout_id = rollout.rollout_id
                    if rollout_id in completed_ids:
@@ -171,6 +173,18 @@ class AlgorithmBatch(agl.Algorithm):
                    spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
                    check_spans(spans, pending[rollout_id])
                    completed_ids.add(rollout_id)
+                    complete_ids_updated = True
+
+                # Check and warn for stale rollouts
+                if complete_ids_updated:
+                    completed_ids_last_updated = 0
+                else:
+                    completed_ids_last_updated += 1
+                    if completed_ids_last_updated >= 10:
+                        unfinished_ids = set(rollout_id for rollout_id, _ in batch_rollouts) - completed_ids
+                        print(f"Stale rollouts: {unfinished_ids}")
+                        completed_ids_last_updated = 0
+
                await asyncio.sleep(5.0)

    async def algorithm_batch_with_completion_threshold(self, total_tasks: int, batch_size: int, remaining_tasks: int):
@@ -303,6 +317,7 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:

 def main(argv: Optional[Sequence[str]] = None) -> None:
    args = parse_args(argv)
+    agl.setup_logging()
    store = agl.LightningStoreClient(args.store_url)
    timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
    try:
@@ -0,0 +1,507 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Collection-level contention benchmarks for Agent Lightning."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import multiprocessing as mp
+import random
+import threading
+import time
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import asdict, dataclass
+from multiprocessing.process import BaseProcess
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Any, AsyncContextManager, Callable, Dict, List, Mapping, Sequence
+
+from pymongo import AsyncMongoClient
+from rich.console import Console
+from rich.table import Table
+
+from agentlightning.store.collection.base import LightningCollections
+from agentlightning.store.collection.memory import InMemoryLightningCollections
+from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections
+from agentlightning.types import Rollout, RolloutConfig
+
+console = Console()
+
+DEFAULT_TOTAL_TASKS = 100_000
+DEFAULT_CONCURRENCY = 1_024
+DEFAULT_TASK_PREFIX = "collection-bench"
+MONGO_DEFAULT_DB = "agentlightning_collection_bench"
+
+
+@dataclass
+class WorkerResult:
+    durations: List[float]
+    failures: int
+
+
+@dataclass
+class BenchmarkResult:
+    backend: str
+    name: str
+    total_tasks: int
+    concurrency: int
+    successes: int
+    failures: int
+    duration: float
+    throughput: float
+    avg_latency: float
+    p50_latency: float
+    p95_latency: float
+    p99_latency: float
+    min_latency: float
+    max_latency: float
+    success_rate: float
+    ops_per_worker: float
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Benchmark LightningStore collections without the store server.")
+    parser.add_argument("benchmark", choices=("insert", "dequeue"), help="Benchmarks to run.")
+    parser.add_argument("--backend", choices=("memory", "mongo"), default="memory", help="Collection backend to test.")
+    parser.add_argument("--total-tasks", type=int, default=DEFAULT_TOTAL_TASKS, help="Total operations to run.")
+    parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY, help="Number of concurrent workers.")
+    parser.add_argument("--task-prefix", default=DEFAULT_TASK_PREFIX, help="Base prefix for generated workload IDs.")
+    parser.add_argument("--summary-file", help="Optional newline-delimited JSON summary output.")
+    parser.add_argument(
+        "--mongo-uri", default="mongodb://localhost:27017/?replicaSet=rs0", help="Mongo connection URI."
+    )
+    parser.add_argument("--mongo-database", default=MONGO_DEFAULT_DB, help="Mongo database for benchmark artifacts.")
+    return parser.parse_args(argv)
+
+
+def _percentile(values: Sequence[float], percentile: float) -> float:
+    if not values:
+        return 0.0
+    if len(values) == 1:
+        return values[0]
+    rank = (len(values) - 1) * percentile
+    lower = math.floor(rank)
+    upper = math.ceil(rank)
+    if lower == upper:
+        return values[int(rank)]
+    return values[lower] * (upper - rank) + values[upper] * (rank - lower)
+
+
+def _aggregate_results(
+    *,
+    backend: str,
+    name: str,
+    results: Sequence[WorkerResult],
+    concurrency: int,
+    total_tasks: int,
+    duration: float,
+) -> BenchmarkResult:
+    successes = sum(len(result.durations) for result in results)
+    failures = sum(result.failures for result in results)
+    latencies = [lat for result in results for lat in result.durations]
+    throughput = successes / duration if duration > 0 else 0.0
+    avg_latency = (sum(latencies) / len(latencies)) if latencies else 0.0
+    sorted_latencies = sorted(latencies)
+    return BenchmarkResult(
+        backend=backend,
+        name=name,
+        total_tasks=total_tasks,
+        concurrency=concurrency,
+        successes=successes,
+        failures=failures,
+        duration=duration,
+        throughput=throughput,
+        avg_latency=avg_latency,
+        p50_latency=_percentile(sorted_latencies, 0.50),
+        p95_latency=_percentile(sorted_latencies, 0.95),
+        p99_latency=_percentile(sorted_latencies, 0.99),
+        min_latency=sorted_latencies[0] if sorted_latencies else 0.0,
+        max_latency=sorted_latencies[-1] if sorted_latencies else 0.0,
+        success_rate=(successes / (successes + failures)) if (successes + failures) else 0.0,
+        ops_per_worker=(successes / concurrency) if concurrency else 0.0,
+    )
+
+
+def _render_results(results: Sequence[BenchmarkResult]) -> None:
+    if not results:
+        console.print("[yellow]No benchmark results to display.[/yellow]")
+        return
+    table = Table(title="Collection Benchmarks", show_lines=False)
+    table.add_column("Backend")
+    table.add_column("Benchmark")
+    table.add_column("Successes", justify="right")
+    table.add_column("Failures", justify="right")
+    table.add_column("Throughput (req/s)", justify="right")
+    table.add_column("Avg Latency (ms)", justify="right")
+    table.add_column("P95 (ms)", justify="right")
+    table.add_column("P99 (ms)", justify="right")
+    table.add_column("Success Rate", justify="right")
+    for result in results:
+        table.add_row(
+            result.backend,
+            result.name,
+            f"{result.successes:,}",
+            f"{result.failures:,}",
+            f"{result.throughput:,.2f}",
+            f"{result.avg_latency * 1e3:,.2f}",
+            f"{result.p95_latency * 1e3:,.2f}",
+            f"{result.p99_latency * 1e3:,.2f}",
+            f"{result.success_rate * 100:,.2f}%",
+        )
+    console.print(table)
+
+
+def _write_summary(results: Sequence[BenchmarkResult], file_path: Path) -> None:
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    with file_path.open("a", encoding="utf-8") as handle:
+        for result in results:
+            handle.write(json.dumps(result.to_dict()) + "\n")
+
+
+def _make_rollout(worker_index: int, sequence: int, task_prefix: str) -> Rollout:
+    rollout_id = f"{task_prefix}-ro-{worker_index}-{sequence}-{uuid.uuid4().hex}"
+    current_time = time.time()
+    return Rollout(
+        rollout_id=rollout_id,
+        input={"task": rollout_id},
+        start_time=current_time,
+        end_time=None,
+        mode="train",
+        resources_id=None,
+        status="queuing",
+        config=RolloutConfig(),
+        metadata={},
+    )
+
+
+async def _preload_queue(collections: LightningCollections, total_tasks: int, task_prefix: str) -> None:
+    batch: List[str] = []
+    for idx in range(total_tasks):
+        batch.append(f"{task_prefix}-queue-{idx}")
+        if len(batch) >= 512:
+            async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+                await collections_atomic.rollout_queue.enqueue(batch)
+            batch.clear()
+    if batch:
+        async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+            await collections_atomic.rollout_queue.enqueue(batch)
+
+
+async def _reset_mongo_database(uri: str, database: str) -> None:
+    client = AsyncMongoClient[Mapping[str, Any]](uri)
+    try:
+        await client.drop_database(database)
+    finally:
+        await client.close()
+
+
+class BaseBenchmark:
+    """Shared control flow for collection benchmarks across backends."""
+
+    def __init__(
+        self, *, backend: str, total_tasks: int, concurrency: int, task_prefix: str, name: str, kind: str
+    ) -> None:
+        self.backend = backend
+        self.total_tasks = total_tasks
+        self.concurrency = concurrency
+        self.task_prefix = task_prefix
+        self.name = name
+        self.kind = kind
+
+    def run(self) -> BenchmarkResult:
+        asyncio.run(self.setup())
+        start = time.perf_counter()
+
+        results = self.spawn_workers(worker_fn=self.worker_entrypoint)
+        duration = time.perf_counter() - start
+        return _aggregate_results(
+            backend=self.backend,
+            name=self.name,
+            results=results,
+            concurrency=self.concurrency,
+            total_tasks=self.total_tasks,
+            duration=duration,
+        )
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        raise NotImplementedError()
+
+    def worker_entrypoint(self, worker_index: int, task_queue: Any, start_barrier: Any) -> WorkerResult:
+        start_barrier.wait()
+        console.print(f"Worker {worker_index} starting")
+
+        async def _runner() -> WorkerResult:
+            async with self.worker_context() as collections:
+                if self.kind == "insert":
+                    return await insert_worker_async(
+                        collections,
+                        worker_index=worker_index,
+                        task_queue=task_queue,
+                        task_prefix=self.task_prefix,
+                    )
+                if self.kind == "dequeue":
+                    return await dequeue_worker_async(
+                        collections,
+                        worker_index=worker_index,
+                        task_queue=task_queue,
+                    )
+                raise ValueError(f"Unknown benchmark kind: {self.kind}")
+
+        return asyncio.run(_runner())
+
+    def worker_context(self, *args: Any, **kwargs: Any) -> AsyncContextManager[LightningCollections]:
+        """Provide the execution context for the benchmark workers."""
+        raise NotImplementedError()
+
+    async def setup(self) -> None:
+        """Prepare backend-specific state before running workers."""
+        if self.kind == "dequeue":
+            async with self.worker_context() as collections:
+                await _preload_queue(collections, self.total_tasks, self.task_prefix)
+
+
+class MemoryBenchmark(BaseBenchmark):
+
+    def __init__(
+        self,
+        *,
+        total_tasks: int,
+        concurrency: int,
+        task_prefix: str,
+        kind: str,
+    ) -> None:
+        super().__init__(
+            total_tasks=total_tasks,
+            concurrency=concurrency,
+            task_prefix=task_prefix,
+            name=f"collection-{kind}",
+            backend="memory",
+            kind=kind,
+        )
+        self.collections = InMemoryLightningCollections(lock_type="thread")
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        task_queue: Queue[int] = Queue()
+        for task_id in range(self.total_tasks):
+            task_queue.put(task_id)
+        start_barrier = threading.Barrier(self.concurrency)
+        results: List[WorkerResult | None] = [None] * self.concurrency
+
+        def _thread_target(worker_index: int) -> None:
+            results[worker_index] = worker_fn(worker_index, task_queue, start_barrier)
+
+        threads: List[threading.Thread] = []
+        for worker_index in range(self.concurrency):
+            thread = threading.Thread(target=_thread_target, args=(worker_index,))
+            thread.start()
+            threads.append(thread)
+        for thread in threads:
+            thread.join()
+
+        return [result for result in results if result is not None]
+
+    @asynccontextmanager
+    async def worker_context(self, *args: Any, **kwargs: Any):
+        yield self.collections
+
+
+class MongoBenchmark(BaseBenchmark):
+    def __init__(
+        self,
+        *,
+        total_tasks: int,
+        concurrency: int,
+        task_prefix: str,
+        kind: str,
+        mongo_uri: str,
+        mongo_database: str,
+    ) -> None:
+        super().__init__(
+            total_tasks=total_tasks,
+            concurrency=concurrency,
+            task_prefix=task_prefix,
+            name=f"collection-{kind}",
+            backend="mongo",
+            kind=kind,
+        )
+        self.mongo_uri = mongo_uri
+        self.mongo_database = mongo_database
+        self.partition_id = f"partition-{uuid.uuid4().hex}"
+
+    async def setup(self) -> None:
+        await _reset_mongo_database(self.mongo_uri, self.mongo_database)
+        return await super().setup()
+
+    @asynccontextmanager
+    async def worker_context(self):
+        pool = MongoClientPool[Mapping[str, Any]](mongo_uri=self.mongo_uri)
+        collections = MongoLightningCollections(
+            client_pool=pool,
+            database_name=self.mongo_database,
+            partition_id=self.partition_id,
+            tracker=None,
+        )
+
+        try:
+            yield collections
+        finally:
+            await pool.close()
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        ctx = mp.get_context("fork")
+        task_queue = ctx.Queue()
+        for task_id in range(self.total_tasks):
+            task_queue.put(task_id)
+        start_barrier = ctx.Barrier(self.concurrency)
+        result_queue = ctx.Queue()
+
+        processes: List[BaseProcess] = []
+        for worker_index in range(self.concurrency):
+            process = ctx.Process(
+                target=_process_worker_target,
+                args=(self, worker_index, task_queue, start_barrier, result_queue),
+            )
+            process.start()
+            processes.append(process)
+
+        collected: List[WorkerResult] = []
+        errors: List[Exception] = []
+        for _ in range(self.concurrency):
+            item = result_queue.get()
+            if isinstance(item, Exception):
+                errors.append(item)
+            else:
+                collected.append(item)
+
+        for process in processes:
+            process.join()
+
+        if errors:
+            raise RuntimeError("One or more worker processes failed") from errors[0]
+
+        return collected
+
+
+def _process_worker_target(
+    benchmark: BaseBenchmark,
+    worker_index: int,
+    task_queue: Any,
+    start_barrier: Any,
+    result_queue: Any,
+) -> None:
+    try:
+        result = benchmark.worker_entrypoint(worker_index, task_queue, start_barrier)
+    except Exception as exc:
+        result_queue.put(exc)
+        raise
+    else:
+        result_queue.put(result)
+
+
+async def insert_worker_async(
+    collections: LightningCollections,
+    *,
+    worker_index: int,
+    task_queue: Any,
+    task_prefix: str,
+) -> WorkerResult:
+    durations: List[float] = []
+    failures = 0
+    while True:
+        try:
+            sequence = task_queue.get_nowait()
+        except Empty:
+            break
+        rollout = _make_rollout(worker_index, sequence, task_prefix)
+        req_start = time.perf_counter()
+        try:
+            async with collections.atomic(mode="rw", labels=["rollouts"]) as collections_atomic:
+                if random.uniform(0, 1) < 0.01:
+                    console.print("Inserting rollout:", rollout.rollout_id)
+                await collections_atomic.rollouts.insert([rollout])
+            durations.append(time.perf_counter() - req_start)
+        except Exception:
+            failures += 1
+    return WorkerResult(durations=durations, failures=failures)
+
+
+async def dequeue_worker_async(
+    collections: LightningCollections,
+    *,
+    worker_index: int,
+    task_queue: Any,
+) -> WorkerResult:
+    del worker_index  # unused but kept for symmetry
+    durations: List[float] = []
+    failures = 0
+    while True:
+        try:
+            task_queue.get_nowait()
+        except Empty:
+            break
+        req_start = time.perf_counter()
+        try:
+            async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+                items = await collections_atomic.rollout_queue.dequeue(limit=1)
+                if items and random.uniform(0, 1) < 0.01:
+                    console.print("Dequeued items:", items[0])
+        except Exception:
+            failures += 1
+            continue
+        if not items:
+            break
+        durations.append(time.perf_counter() - req_start)
+    return WorkerResult(durations=durations, failures=failures)
+
+
+def run_benchmark(args: argparse.Namespace, benchmark_kind: str) -> BenchmarkResult:
+    params = {
+        "total_tasks": args.total_tasks,
+        "concurrency": args.concurrency,
+        "task_prefix": args.task_prefix,
+    }
+    if args.backend == "memory":
+        return MemoryBenchmark(kind=benchmark_kind, **params).run()
+
+    mongo_params = {
+        **params,
+        "mongo_uri": args.mongo_uri,
+        "mongo_database": args.mongo_database,
+    }
+    return MongoBenchmark(kind=benchmark_kind, **mongo_params).run()
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.total_tasks <= 0:
+        raise ValueError("total-tasks must be positive")
+    if args.concurrency <= 0:
+        raise ValueError("concurrency must be positive")
+
+    results: List[BenchmarkResult] = []
+    results.append(run_benchmark(args, args.benchmark))
+
+    _render_results(results)
+
+    if args.summary_file:
+        _write_summary(results, Path(args.summary_file))
+
+
+if __name__ == "__main__":  # pragma: no cover - manual execution
+    main()
@@ -10,17 +10,32 @@ import multiprocessing
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional, Sequence
+from typing import List, Optional, Sequence

 from rich.console import Console

 import agentlightning as agl
-from agentlightning.types.tracer import OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.types import EnqueueRolloutRequest, OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.utils.metrics import ConsoleMetricsBackend, MultiMetricsBackend
 from agentlightning.utils.system_snapshot import system_snapshot

+from .utils import flatten_dict, random_dict
+
 console = Console()


+async def _enqueue_rollouts_for_benchmark(store_url: str, *, total_rollouts: int, task_prefix: str) -> None:
+    """Utility that enqueues a fixed number of rollouts for a benchmark."""
+    store = agl.LightningStoreClient(store_url)
+    console.print(f"Enqueuing {total_rollouts} rollouts for {task_prefix} benchmark")
+    try:
+        await store.enqueue_many_rollouts(
+            [EnqueueRolloutRequest(input={"task": f"{task_prefix}-Task-{i}"}) for i in range(total_rollouts)]
+        )
+    finally:
+        await store.close()
+
+
 def _close_store_client(store: agl.LightningStoreClient) -> None:
    try:
        asyncio.run(store.close())
@@ -28,7 +43,7 @@ def _close_store_client(store: agl.LightningStoreClient) -> None:
        pass


-def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
+def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str, attribute_size: int) -> Span:
    trace_hex = f"{sequence_id:032x}"
    span_hex = f"{sequence_id:016x}"
    return Span(
@@ -40,7 +55,14 @@ def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) ->
        parent_id=None,
        name=name,
        status=TraceStatus(status_code="OK"),
-        attributes={},
+        attributes=flatten_dict(
+            random_dict(
+                depth=1,
+                breadth=attribute_size,
+                key_length=(3, 20),
+                value_length=(5, 300),
+            )
+        ),
        events=[],
        links=[],
        start_time=None,
@@ -77,8 +99,8 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser.add_argument("--summary-file", help="File to append final benchmark summary.")
    parser.add_argument(
        "mode",
-        choices=("worker", "dequeue-empty", "rollout"),
-        help="Mode to exercise different operations.",
+        choices=("worker", "dequeue-empty", "dequeue-only", "rollout", "dequeue-update-attempt", "metrics"),
+        help="Mode to exercise different operations (metrics targets MultiMetricsBackend fan-out).",
    )
    args = parser.parse_args(argv)
    return args
@@ -173,6 +195,7 @@ def _rollout_flow_task(args: tuple[str, int, int]) -> bool:
                attempt_id,
                task_id * spans_per_attempt + seq,
                f"micro-span-{seq}",
+                attribute_size=1,
            )
            await store.add_span(span)
        console.print(f"Updating attempt {attempt_id} for task {task_id} with {spans_per_attempt} spans")
@@ -207,6 +230,194 @@ def simulate_rollout_with_spans(store_url: str, spans_per_attempt: int = 4) -> B
    return BenchmarkSummary(mode="rollout", total_tasks=len(task_ids), successes=successes, duration=duration)


+def _dequeue_only_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"[Dequeue-Only Task {task_id}] Dequeueing rollout for worker {worker_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        attempted = await store.dequeue_rollout()  # no worker_id
+        if attempted is None:
+            console.print(f"[Dequeue-Only Task {task_id}] No rollout available to dequeue")
+            return False
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing only worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_rollouts(store_url: str) -> BenchmarkSummary:
+    """Benchmark simple dequeues without any additional mutations."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="DequeueOnly"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_only_task, [(store_url, worker_id, task_id) for worker_id, task_id in worker_jobs]
+        )
+
+    async def _query_remaining_rollouts() -> List[str]:
+        store = agl.LightningStoreClient(store_url)
+        try:
+            remaining_rollouts = await store.query_rollouts(status_in=["queuing"])
+            return [item.rollout_id for item in remaining_rollouts]
+        finally:
+            await store.close()
+
+    end_time = time.time()
+    remaining_rollouts = asyncio.run(_query_remaining_rollouts())
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Remaining rollouts: {remaining_rollouts}")
+    console.print(f"Remaining rollouts count: {len(remaining_rollouts)}")
+    console.print(f"Dequeue-only success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(mode="dequeue-only", total_tasks=len(worker_jobs), successes=successes, duration=duration)
+
+
+def _dequeue_and_update_attempt_task(args: tuple[str, str, str, int]) -> bool:
+    store_url, worker_id, task_id, spans_per_attempt = args
+    console.print(f"Dequeueing and update attempt with worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        console.print(f"[Task {task_id}] Dequeueing rollout")
+        attempted = await store.dequeue_rollout(worker_id=worker_id)
+        if attempted is None:
+            console.print(f"[Task {task_id}] No rollout available to dequeue")
+            return False
+        console.print(f"[Task {task_id}] Retrieving span sequence IDs")
+        sequence_ids = await store.get_many_span_sequence_ids(
+            [(attempted.rollout_id, attempted.attempt.attempt_id) for _ in range(spans_per_attempt)]
+        )
+        if len(sequence_ids) != spans_per_attempt:
+            console.print(
+                f"[Task {task_id}] Unable to retrieve enough span sequence IDs: "
+                f"expected={spans_per_attempt} got={len(sequence_ids)}"
+            )
+            return False
+        console.print(f"[Task {task_id}] Adding {spans_per_attempt} spans")
+        spans = [
+            _make_span(
+                attempted.rollout_id,
+                attempted.attempt.attempt_id,
+                sequence_id,
+                f"micro-span-{sequence_id}",
+                attribute_size=32,
+            )
+            for sequence_id in sequence_ids
+        ]
+        stored_spans = await store.add_many_spans(spans)
+        if len(stored_spans) != len(spans):
+            console.print(
+                f"[Task {task_id}] Only stored {len(stored_spans)}/{len(spans)} spans for "
+                f"rollout_id={attempted.rollout_id} attempt_id={attempted.attempt.attempt_id}"
+            )
+            return False
+        console.print(
+            f"[Task {task_id}] Updating attempt to succeeded: rollout_id={attempted.rollout_id} "
+            f"attempt_id={attempted.attempt.attempt_id}"
+        )
+        await store.update_attempt(attempted.rollout_id, attempted.attempt.attempt_id, status="succeeded")
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing and updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_and_update_attempts(store_url: str, spans_per_attempt: int = 4) -> BenchmarkSummary:
+    """Simulate dequeueing rollouts and updating attempts with spans."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="Dequeue"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_and_update_attempt_task,
+            [(store_url, worker_id, task_id, spans_per_attempt) for worker_id, task_id in worker_jobs],
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Dequeue and update attempt success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(
+        mode="dequeue-update-attempt", total_tasks=len(worker_jobs), successes=successes, duration=duration
+    )
+
+
+def benchmark_multi_metrics_backend(iterations: int = 10_000_000) -> BenchmarkSummary:
+    """Benchmark MultiMetricsBackend fan-out cost."""
+
+    console.print(f"Benchmarking MultiMetricsBackend for {iterations} iterations (2 metric ops per iteration)")
+
+    agl.setup_logging()
+
+    console_backend = ConsoleMetricsBackend(window_seconds=0.5, log_interval_seconds=0.1, group_level=None)
+    console_backend_secondary = ConsoleMetricsBackend(
+        window_seconds=None, log_interval_seconds=1_000_000.0, group_level=None
+    )
+    backend = MultiMetricsBackend([console_backend, console_backend_secondary])
+
+    backend.register_counter("benchmark.metrics.counter", label_names=["worker"])
+    backend.register_histogram(
+        "benchmark.metrics.latency",
+        label_names=["worker"],
+        buckets=(0.001, 0.005, 0.05, 0.5, 1.0),
+    )
+    labels = {"worker": "benchmark"}
+
+    async def _exercise_metrics() -> None:
+        for i in range(iterations):
+            await backend.inc_counter("benchmark.metrics.counter", labels=labels)
+            await backend.observe_histogram(
+                "benchmark.metrics.latency",
+                value=(i % 100) / 100.0,
+                labels=labels,
+            )
+
+    start_time = time.time()
+    asyncio.run(_exercise_metrics())
+    duration = time.time() - start_time
+    total_ops = iterations * 2
+    throughput = total_ops / duration if duration > 0 else 0.0
+
+    console.print(f"Executed {total_ops} metric updates in {duration:.3f}s ({throughput:.1f} ops/s)")
+    return BenchmarkSummary(mode="metrics", total_tasks=total_ops, successes=total_ops, duration=duration)
+
+
 def record_summary(summary: BenchmarkSummary, summary_file: Optional[str]) -> None:
    message = (
        f"[summary] mode={summary.mode} success_rate={summary.success_rate:.3f} "
@@ -227,11 +438,19 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
        summary = simulate_many_update_workers(args.store_url)
    elif args.mode == "dequeue-empty":
        summary = simulate_dequeue_empty_and_update_workers(args.store_url)
+    elif args.mode == "dequeue-only":
+        summary = dequeue_rollouts(args.store_url)
    elif args.mode == "rollout":
        summary = simulate_rollout_with_spans(args.store_url)
+    elif args.mode == "dequeue-update-attempt":
+        summary = dequeue_and_update_attempts(args.store_url)
+    elif args.mode == "metrics":
+        summary = benchmark_multi_metrics_backend()
    else:
        raise ValueError(f"Invalid mode: {args.mode}")
    record_summary(summary, args.summary_file)
+    if summary.success_rate < 1.0:
+        raise ValueError(f"Benchmark failed with success rate {summary.success_rate:.3f}")


 if __name__ == "__main__":
@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import types
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+
+class _CounterChild:
+    def __init__(self) -> None:
+        self.value = 0.0
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.value += amount
+
+
+class _HistogramChild:
+    def __init__(self) -> None:
+        self.values: List[float] = []
+
+    def observe(self, value: float) -> None:
+        self.values.append(value)
+
+
+class PrometheusStub(types.ModuleType):
+    """Minimal prometheus_client replacement for unit tests."""
+
+    def __init__(self, real_client: Any) -> None:
+        super().__init__("prometheus_client")
+        self.counter_instances: List[_PromCounter] = []
+        self.histogram_instances: List[_PromHistogram] = []
+        self.real_client = real_client
+
+        class CollectorRegistry:
+            pass
+
+        class _Multiprocess:
+            def __init__(self) -> None:
+                self.registry: Optional[CollectorRegistry] = None
+
+            def MultiProcessCollector(self, registry: CollectorRegistry) -> None:
+                self.registry = registry
+
+        self.CollectorRegistry = CollectorRegistry
+        self.REGISTRY = CollectorRegistry()
+        self.multiprocess = _Multiprocess()
+
+        self.Counter = _PromCounterFactory(self)
+        self.Histogram = _PromHistogramFactory(self)
+        self.make_asgi_app = self._make_asgi_app
+
+    def _make_asgi_app(self, *args: Any, **kwargs: Any) -> Any:
+        return self.real_client.make_asgi_app(*args, **kwargs)
+
+
+class _PromCounterFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(self, name: str, doc: str, labelnames: Sequence[str]) -> _PromCounter:
+        counter = _PromCounter(name, doc, labelnames)
+        counter._register(self._owner.counter_instances)  # pyright: ignore[reportPrivateUsage]
+        return counter
+
+
+class _PromHistogramFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(
+        self,
+        name: str,
+        doc: str,
+        labelnames: Sequence[str],
+        buckets: Sequence[float] | None = None,
+    ) -> _PromHistogram:
+        histogram = _PromHistogram(name, doc, labelnames, buckets or ())
+        histogram._register(self._owner.histogram_instances)  # pyright: ignore[reportPrivateUsage]
+        return histogram
+
+
+class _PromCounter:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.default = _CounterChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _CounterChild] = {}
+
+    def _register(self, sink: List["_PromCounter"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _CounterChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _CounterChild())
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.default.inc(amount)
+
+
+class _PromHistogram:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str], buckets: Sequence[float]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.buckets = tuple(buckets)
+        self.default = _HistogramChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _HistogramChild] = {}
+
+    def _register(self, sink: List["_PromHistogram"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _HistogramChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _HistogramChild())
+
+    def observe(self, value: float) -> None:
+        self.default.observe(value)
+
+
+def make_prometheus_stub() -> PrometheusStub:
+    """Factory helper for tests."""
+    import prometheus_client
+
+    return PrometheusStub(prometheus_client)
@@ -394,6 +394,11 @@ async def _kbint_in_runner(store: LightningStore, worker_id: int, event: Executi
    raise KeyboardInterrupt()


+async def _cancel_in_runner(store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
+    _ = (store, worker_id, event)
+    raise asyncio.CancelledError()
+
+
 async def _timeout_error_in_runner(store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
    # Provoke client's validation (pre-request), then raise TimeoutError.
    with pytest.raises(ValueError):
@@ -1236,3 +1241,64 @@ def test_execute_main_runner_store_state_isolated_in_subprocess(store: DummyLigh
    assert (
        len(store.calls) == initial_call_count
    ), "Store state should not be modified in main process when main_process='runner'"
+
+
+def test_spawn_runners_handles_keyboard_interrupt_gracefully(store: LightningStore) -> None:
+    """
+    Test that KeyboardInterrupt (Ctrl+C) is caught by _runner_sync
+    and results in a graceful exit (exitcode 0).
+    """
+    strat = ClientServerExecutionStrategy(
+        role="runner",
+        n_runners=1,
+        server_host="127.0.0.1",
+        server_port=_free_port(),
+    )
+    ctx = get_context()
+    stop_evt: ExecutionEvent = MpEvent()
+
+    processes = strat._spawn_runners(_kbint_in_runner, store, stop_evt, ctx=ctx)  # pyright: ignore[reportPrivateUsage]
+
+    try:
+        for p in processes:
+            p.join(timeout=2.0)
+
+        for p in processes:
+            assert not p.is_alive()
+            assert p.exitcode == 0, f"Runner {p.name} should exit gracefully on KeyboardInterrupt"
+
+    finally:
+        for p in processes:
+            if p.is_alive():
+                p.terminate()
+                p.join()
+
+
+def test_spawn_runners_treats_cancelled_error_as_crash(store: LightningStore) -> None:
+    """
+    Test that asyncio.CancelledError in __spawn_runners causes a crash (exitcode != 0).
+    """
+    strat = ClientServerExecutionStrategy(
+        role="runner",
+        n_runners=1,
+        server_host="127.0.0.1",
+        server_port=_free_port(),
+    )
+    ctx = get_context()
+    stop_evt: ExecutionEvent = MpEvent()
+
+    processes = strat._spawn_runners(_cancel_in_runner, store, stop_evt, ctx=ctx)  # pyright: ignore[reportPrivateUsage]
+
+    try:
+        for p in processes:
+            p.join(timeout=2.0)
+
+        for p in processes:
+            assert not p.is_alive()
+            assert p.exitcode != 0, f"Runner {p.name} should crash on CancelledError (exitcode={p.exitcode})"
+
+    finally:
+        for p in processes:
+            if p.is_alive():
+                p.terminate()
+                p.join()
@@ -1,7 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.

+# FIXME: This file will have side-effects on other tests if the tests failed and agentops service is not disabled.
+
 from unittest.mock import MagicMock, patch

+import pytest
 from opentelemetry.sdk.metrics.export import MetricExportResult
 from opentelemetry.sdk.trace.export import SpanExportResult

@@ -13,6 +16,7 @@ from agentlightning.instrumentation.agentops import (
 )


+@pytest.mark.agentops
 def test_switchable_authenticated_exporter():
    switchable_authenticated_exporter = BypassableAuthenticatedOTLPExporter(endpoint="http://dummy", jwt="dummy")

@@ -30,6 +34,7 @@ def test_switchable_authenticated_exporter():
        assert mock_export.call_count == 1


+@pytest.mark.agentops
 def test_switchable_otlp_metric_exporter():

    switchable_otlp_metric_exporter = BypassableOTLPMetricExporter()
@@ -47,6 +52,7 @@ def test_switchable_otlp_metric_exporter():
        assert mock_export.call_count == 1


+@pytest.mark.agentops
 def test_switchable_otlp_span_exporter():

    switchable_otlp_span_exporter = BypassableOTLPSpanExporter()
@@ -27,6 +27,8 @@ from agentlightning.utils.server_launcher import PythonServerLauncherArgs
 from ..common.network import get_free_port
 from ..common.tracer import clear_tracer_provider

+pytestmark = pytest.mark.llmproxy
+

 class _FakeSpanContext:
    def __init__(self, span_id: int):
@@ -32,15 +32,7 @@ from agentlightning.types import LLM, Span
 from ..common.tracer import clear_tracer_provider
 from ..common.vllm import VLLM_VERSION, RemoteOpenAIServer

-try:
-    import torch  # type: ignore
-
-    GPU_AVAILABLE = torch.cuda.is_available()
-except Exception:
-    GPU_AVAILABLE = False  # type: ignore
-
-if not GPU_AVAILABLE:
-    pytest.skip(reason="GPU not available", allow_module_level=True)
+pytestmark = [pytest.mark.gpu, pytest.mark.llmproxy]


@pytest.fixture(scope="module")
@@ -2,8 +2,9 @@

 import os
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Dict, Optional, cast
+from typing import Any, AsyncGenerator, Awaitable, Dict, Optional, cast

+import litellm
 import openai
 import pytest

@@ -21,19 +22,38 @@ from ..common.tracer import clear_tracer_provider
 from ..common.vllm import VLLM_AVAILABLE, RemoteOpenAIServer


-async def init_runner(
-    agent: LitAgent[Any],
-    *,
-    resources: Optional[Dict[str, LLM]] = None,
-) -> tuple[LitAgentRunner[Any], InMemoryLightningStore]:
-    store = InMemoryLightningStore()
-    llm_resource: NamedResources = resources or {"llm": LLM(endpoint="http://localhost", model="dummy")}  # type: ignore[assignment]
-    await store.update_resources("default", llm_resource)
+class InitRunnerFunction:

-    runner = LitAgentRunner[Any](tracer=AgentOpsTracer(), poll_interval=0.01)
-    runner.init(agent)
-    runner.init_worker(worker_id=0, store=store)
-    return runner, store
+    def __call__(
+        self,
+        agent: LitAgent[Any],
+        *,
+        resources: Optional[Dict[str, LLM]] = None,
+    ) -> Awaitable[tuple[LitAgentRunner[Any], InMemoryLightningStore]]: ...
+
+
+@pytest.fixture(
+    params=[
+        pytest.param("agentops", marks=pytest.mark.agentops),
+    ]
+)
+def init_runner(request: pytest.FixtureRequest) -> InitRunnerFunction:
+    async def init_runner_fn(
+        agent: LitAgent[Any],
+        *,
+        resources: Optional[Dict[str, LLM]] = None,
+    ) -> tuple[LitAgentRunner[Any], InMemoryLightningStore]:
+        store = InMemoryLightningStore()
+        llm_resource: NamedResources = resources or {"llm": LLM(endpoint="http://localhost", model="dummy")}  # type: ignore[assignment]
+        await store.update_resources("default", llm_resource)
+
+        # This is always AgentOpsTracer for now
+        runner = LitAgentRunner[Any](tracer=AgentOpsTracer(), poll_interval=0.01)
+        runner.init(agent)
+        runner.init_worker(worker_id=0, store=store)
+        return runner, store
+
+    return init_runner_fn  # type: ignore


 def teardown_runner(runner: LitAgentRunner[Any]) -> None:
@@ -51,7 +71,7 @@ def setup_module():


@pytest.mark.asyncio
-async def test_runner_integration_basic_rollout() -> None:
+async def test_runner_integration_basic_rollout(init_runner: InitRunnerFunction) -> None:
    class EchoAgent(LitAgent[str]):
        async def validation_rollout_async(self, task: str, resources: Dict[str, Any], rollout: Any) -> None:
            emit_reward(1.0)
@@ -71,11 +91,8 @@ async def test_runner_integration_basic_rollout() -> None:


@pytest.mark.asyncio
-@pytest.mark.skipif(
-    not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")),
-    reason="OpenAI endpoint or key not configured",
-)
-async def test_runner_integration_with_openai() -> None:
+@pytest.mark.openai
+async def test_runner_integration_with_openai(init_runner: InitRunnerFunction) -> None:
    class OpenAIAgent(LitAgent[str]):
        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
            llm = cast(LLM, resources["llm"])
@@ -87,6 +104,9 @@ async def test_runner_integration_with_openai() -> None:
            assert response.choices, "OpenAI response should contain choices"
            return 0.0

+    if not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")):
+        raise RuntimeError("OpenAI endpoint or key not configured")
+
    base_url = os.environ["OPENAI_BASE_URL"]
    api_key = os.environ["OPENAI_API_KEY"]
    model = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
@@ -104,23 +124,25 @@ async def test_runner_integration_with_openai() -> None:


@pytest.mark.asyncio
+@pytest.mark.openai
@pytest.mark.skipif(
    not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")),
    reason="OpenAI endpoint or key not configured",
 )
-async def test_runner_integration_with_litellm_proxy() -> None:
-    litellm = pytest.importorskip("litellm")
-
+async def test_runner_integration_with_litellm_proxy(init_runner: InitRunnerFunction) -> None:
    class LiteLLMAgent(LitAgent[str]):
        def validation_rollout(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
            llm = cast(LLM, resources["llm"])
-            response = litellm.completion(
+            response = litellm.completion(  # type: ignore
                model=llm.model,
                messages=[{"role": "user", "content": task}],
            )
-            assert response.get("choices"), "litellm proxy should return choices"
+            assert response.get("choices"), "litellm proxy should return choices"  # type: ignore
            return 0.0

+    if not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")):
+        raise RuntimeError("OpenAI endpoint or key not configured")
+
    agent = LiteLLMAgent()
    resources = {"llm": LLM(endpoint="http://dummy", model="openai/gpt-4o-mini")}
    runner, store = await init_runner(agent, resources=resources)
@@ -154,7 +176,11 @@ def server():


 class LLMProxyWithClearedTracerProvider(LLMProxy):
-    """LLMProxy that clears the tracer provider before serving."""
+    """LLMProxy that clears the tracer provider before serving.
+
+    It will be run in a separate process, so the tracer provider initialized there does not
+    interfere with the main process's tracer provider.
+    """

    @asynccontextmanager
    async def _serve_context(self) -> AsyncGenerator[None, None]:
@@ -165,11 +191,10 @@ class LLMProxyWithClearedTracerProvider(LLMProxy):


@pytest.mark.asyncio
-async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenAIServer) -> None:
-    torch = pytest.importorskip("torch")
-    if not torch.cuda.is_available():
-        pytest.skip("GPU not available")
-
+@pytest.mark.gpu
+async def test_runner_integration_with_spawned_litellm_proxy(
+    init_runner: InitRunnerFunction, server: RemoteOpenAIServer
+) -> None:
    class ProxyAgent(LitAgent[str]):
        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: Rollout) -> float:
            attempted_rollout = cast(AttemptedRollout, rollout)
@@ -5,7 +5,7 @@ from __future__ import annotations
 import os
 import time
 from itertools import count
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Sequence
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Mapping, Sequence
 from unittest.mock import Mock
 from uuid import uuid4

@@ -15,6 +15,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field
 from pytest import FixtureRequest

+from agentlightning.store import collection_based
 from agentlightning.store.base import LightningStore
 from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection, Queue
 from agentlightning.store.collection.base import Collection
@@ -26,6 +27,10 @@ if TYPE_CHECKING:

 __all__ = [
    "inmemory_store",
+    "inmemory_debounced_store",
+    "mongo_debounced_store",
+    "debounced_store",
+    "fake_time",
    "mock_readable_span",
    "sample_items",
    "sample_collection",
@@ -35,16 +40,25 @@ __all__ = [
    "dict_key_value",
    "dict_key_value_data",
    "temporary_mongo_database",
+    "mongo_uri",
+    "mongo_client_kwargs",
 ]


 mongo_uri = os.getenv("AGL_TEST_MONGO_URI", "mongodb://localhost:27017/?replicaSet=rs0")
+mongo_client_kwargs: Dict[str, Any] = {"serverSelectionTimeoutMS": 5000}


@pytest.fixture
 def inmemory_store() -> InMemoryLightningStore:
    """Create a fresh InMemoryLightningStore instance."""
-    return InMemoryLightningStore()
+    return InMemoryLightningStore(scan_debounce_seconds=0)
+
+
+@pytest.fixture
+def inmemory_debounced_store(fake_time: _FakeTime) -> InMemoryLightningStore:
+    """Create an InMemoryLightningStore configured with scan debouncing."""
+    return InMemoryLightningStore(scan_debounce_seconds=5.0)


@pytest_asyncio.fixture
@@ -52,7 +66,29 @@ async def mongo_store(temporary_mongo_database: AsyncDatabase[Any]):
    """Fixture for MongoDB store implementation."""
    from agentlightning.store.mongo import MongoLightningStore

-    db = MongoLightningStore(client=temporary_mongo_database.client, database_name=temporary_mongo_database.name)
+    db = MongoLightningStore(
+        mongo_uri=mongo_uri,
+        mongo_client_kwargs=mongo_client_kwargs,
+        database_name=temporary_mongo_database.name,
+        scan_debounce_seconds=0,
+    )
+    try:
+        yield db
+    finally:
+        await db.close()
+
+
+@pytest_asyncio.fixture
+async def mongo_debounced_store(fake_time: _FakeTime, temporary_mongo_database: AsyncDatabase[Any]):
+    """Fixture for MongoDB store implementation with scan debouncing."""
+    from agentlightning.store.mongo import MongoLightningStore
+
+    db = MongoLightningStore(
+        mongo_uri=mongo_uri,
+        mongo_client_kwargs=mongo_client_kwargs,
+        database_name=temporary_mongo_database.name,
+        scan_debounce_seconds=5.0,
+    )
    try:
        yield db
    finally:
@@ -70,6 +106,17 @@ def store_fixture(request: FixtureRequest) -> AsyncGenerator[LightningStore, Non
    return request.getfixturevalue(request.param)


+@pytest.fixture(
+    params=[
+        "inmemory_debounced_store",
+        pytest.param("mongo_debounced_store", marks=pytest.mark.mongo),
+    ]
+)
+def debounced_store(request: FixtureRequest) -> LightningStore:
+    """Parameterized fixture for debounced store implementations."""
+    return request.getfixturevalue(request.param)
+
+
@pytest.fixture
 def mock_readable_span() -> ReadableSpan:
    """Create a mock ReadableSpan for testing."""
@@ -123,11 +170,35 @@ class QueueItem(BaseModel):
    idx: int


+class _FakeTime:
+    """Simple controllable clock for scan debouncing tests."""
+
+    def __init__(self, start: float = 0.0) -> None:
+        self._value = start
+
+    def time(self) -> float:
+        return self._value
+
+    def set(self, value: float) -> None:
+        self._value = value
+
+    def advance(self, delta: float) -> None:
+        self._value += delta
+
+
+@pytest.fixture
+def fake_time(monkeypatch: pytest.MonkeyPatch) -> _FakeTime:
+    """Patch collection_based.time.time with a controllable clock."""
+    controller = _FakeTime()
+    monkeypatch.setattr(collection_based.time, "time", controller.time)
+    return controller
+
+
@pytest_asyncio.fixture
 async def mongo_client():
    from pymongo import AsyncMongoClient

-    client = AsyncMongoClient[Any](mongo_uri, serverSelectionTimeoutMS=5000)
+    client = AsyncMongoClient[Any](mongo_uri, **mongo_client_kwargs)
    try:
        await client.admin.command("ping")
    except Exception as exc:  # depends on external service
@@ -270,7 +341,9 @@ def sample_collection_memory(sample_items: Sequence[SampleItem]) -> ListBasedCol
 async def sample_collection_mongo(temporary_mongo_database: AsyncDatabase[Any], sample_items: Sequence[SampleItem]):
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        collection = MongoBasedCollection(
            client_pool,
            temporary_mongo_database.name,
@@ -307,7 +380,9 @@ def deque_queue_memory() -> DequeBasedQueue[QueueItem]:
 async def deque_queue_mongo(temporary_mongo_database: AsyncDatabase[Any]):
    from agentlightning.store.collection.mongo import MongoBasedQueue, MongoClientPool

-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        queue = MongoBasedQueue[QueueItem](
            client_pool,
            temporary_mongo_database.name,
@@ -347,7 +422,9 @@ def dict_key_value_memory(dict_key_value_data: Dict[str, int]) -> DictBasedKeyVa
 async def dict_key_value_mongo(temporary_mongo_database: AsyncDatabase[Any], dict_key_value_data: Dict[str, int]):
    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool

-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        key_value = MongoBasedKeyValue[str, int](
            client_pool,
            temporary_mongo_database.name,
@@ -28,8 +28,18 @@ from agentlightning.types import (
    Span,
    TraceStatus,
 )
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+)
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncherArgs

+from ..common.prometheus_stub import make_prometheus_stub
+
+pytestmark = [pytest.mark.store]
+

 def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
    return Span(
@@ -80,6 +90,66 @@ async def server_client(
        await server.stop()


+async def _exercise_server_metrics_backend(tracker: MetricsBackend) -> None:
+    port = pick_unused_port()
+    store = InMemoryLightningStore(tracker=tracker)
+    server = LightningStoreServer(store, "127.0.0.1", port, tracker=tracker)
+    await server.start()
+    client = LightningStoreClient(server.endpoint)
+    try:
+        await _run_server_side_operations(server)
+        await _run_client_side_operations(client)
+    finally:
+        await client.close()
+        await server.stop()
+
+
+async def _run_server_side_operations(server: LightningStoreServer) -> None:
+    await server.update_resources("metrics-server", {})
+    await server.get_latest_resources()
+    started = await server.start_rollout(input={"origin": "server"}, config=RolloutConfig(timeout_seconds=1.0))
+    queued = await server.enqueue_rollout(input={"origin": "server-queue"})
+    dequeued = await server.dequeue_rollout(worker_id="metrics-server-worker")
+    assert dequeued is not None
+
+    await server.add_span(_make_span(dequeued.rollout_id, dequeued.attempt.attempt_id, 0, "server-span"))
+    await server.update_attempt(queued.rollout_id, dequeued.attempt.attempt_id, status="running")
+    await server.update_attempt(queued.rollout_id, dequeued.attempt.attempt_id, status="succeeded")
+    await server.update_rollout(queued.rollout_id, status="succeeded")
+    await server.wait_for_rollouts(rollout_ids=[queued.rollout_id], timeout=0.1)
+    assert started is not None
+
+
+async def _run_client_side_operations(client: LightningStoreClient) -> None:
+    await client.update_resources("metrics-client", {})
+    await client.get_latest_resources()
+
+    await client.start_rollout(input={"origin": "client"}, mode="train", config=RolloutConfig(timeout_seconds=2.0))
+    queued = await client.enqueue_rollout(
+        input={"origin": "client-queue"}, config=RolloutConfig(unresponsive_seconds=5.0)
+    )
+    dequeued = await client.dequeue_rollout(worker_id="metrics-client-worker")
+    assert dequeued is not None
+
+    span = _make_span(dequeued.rollout_id, dequeued.attempt.attempt_id, 1, "client-span")
+    await client.add_span(span)
+
+    await client.update_attempt(
+        dequeued.rollout_id,
+        dequeued.attempt.attempt_id,
+        status="running",
+        worker_id="metrics-client-worker",
+    )
+    await client.update_attempt(dequeued.rollout_id, dequeued.attempt.attempt_id, status="succeeded")
+    await client.update_rollout(dequeued.rollout_id, status="succeeded")
+
+    await client.wait_for_rollouts(rollout_ids=[dequeued.rollout_id], timeout=0.1)
+    await client.query_rollouts()
+    await client.query_attempts(dequeued.rollout_id)
+    await client.get_worker_by_id("metrics-client-worker")
+    assert queued.rollout_id == dequeued.rollout_id
+
+
@pytest.mark.asyncio
 async def test_mp_server_does_not_work_with_inmemory_store() -> None:
    store = InMemoryLightningStore()
@@ -200,6 +270,51 @@ async def test_client_start_attempt_propagates_worker_id(
    assert worker.current_attempt_id == retry.attempt.attempt_id


+@pytest.mark.asyncio
+async def test_console_metrics_backend_tracks_http_and_store_metrics() -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=3600.0, group_level=3)
+    await _exercise_server_metrics_backend(backend)
+
+    counter_metrics = {name for name, _ in backend._counter_state.keys()}  # pyright: ignore[reportPrivateUsage]
+    hist_metrics = {name for name, _ in backend._hist_state.keys()}  # pyright: ignore[reportPrivateUsage]
+    assert "agl.http.total" in counter_metrics
+    assert "agl.store.total" in counter_metrics
+    assert "agl.http.latency" in hist_metrics
+
+
+@pytest.mark.asyncio
+@pytest.mark.prometheus
+async def test_prometheus_metrics_backend_tracks_http_metrics(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+    backend = PrometheusMetricsBackend()
+    await _exercise_server_metrics_backend(backend)
+
+    http_counter = next(inst for inst in stub.counter_instances if inst.name == "agl_http_total")
+    http_histogram = next(inst for inst in stub.histogram_instances if inst.name == "agl_http_latency")
+    assert any(child.value > 0 for child in http_counter.children.values())
+    assert any(child.values for child in http_histogram.children.values())
+
+
+@pytest.mark.asyncio
+@pytest.mark.prometheus
+async def test_multi_metrics_backend_updates_all_children(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+    console_backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=3600.0, group_level=3)
+    prom_backend = PrometheusMetricsBackend()
+    backend = MultiMetricsBackend([console_backend, prom_backend])
+    await _exercise_server_metrics_backend(backend)
+
+    console_counters = {
+        name for name, _ in console_backend._counter_state.keys()  # pyright: ignore[reportPrivateUsage]
+    }
+    assert "agl.http.total" in console_counters
+
+    prom_counter = next(inst for inst in stub.counter_instances if inst.name == "agl_http_total")
+    assert any(child.value > 0 for child in prom_counter.children.values())
+
+
@pytest.mark.asyncio
 async def test_client_enqueue_many_rollouts_uses_batch_payload(monkeypatch: MonkeyPatch) -> None:
    client = LightningStoreClient("http://localhost:9000")
@@ -18,26 +18,30 @@ from typing import (
    Sequence,
    Tuple,
    Union,
+    cast,
 )
 from uuid import uuid4

+import pydantic
 import pytest
 from pydantic import BaseModel

 import agentlightning.store.collection.memory as memory_module
-from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
+from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection
 from agentlightning.store.collection.base import Collection
 from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
 from agentlightning.store.collection.memory import _LoopAwareAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.store.collection.memory import _ThreadSafeAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.types import Rollout
-from tests.store.conftest import QueueItem, SampleItem
+from tests.store.conftest import QueueItem, SampleItem, mongo_client_kwargs, mongo_uri

 if TYPE_CHECKING:
    from pymongo.asynchronous.database import AsyncDatabase

    from agentlightning.store.collection.mongo import MongoLightningCollections

+pytestmark = [pytest.mark.store]
+

 def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
    return ListBasedCollection(list(items), SampleItem, ("partition", "index"))
@@ -117,11 +121,13 @@ async def test_list_collection_insert_rejects_duplicate_payload(sample_collectio
    dup_a = SampleItem(partition="omega", index=1, name="dup-a", status="new")
    dup_b = SampleItem(partition="omega", index=1, name="dup-b", status="new")

-    with pytest.raises(ValueError, match="duplicate primary key"):
+    with pytest.raises(ValueError, match=r"(duplicated|Duplicated) primary key"):
        await sample_collection.insert([dup_a, dup_b])

-    assert await sample_collection.size() == starting_size
-    assert await sample_collection.get({"partition": {"exact": "omega"}}) is None
+    if isinstance(sample_collection, ListBasedCollection):
+        # Only ListBasedCollection supports this rejecting duplicate items within the same insert batch.
+        assert await sample_collection.size() == starting_size
+        assert await sample_collection.get({"partition": {"exact": "omega"}}) is None


@pytest.mark.asyncio()
@@ -865,6 +871,150 @@ async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyVa
    assert await dict_key_value.size() == 1


+@pytest.mark.asyncio()
+async def test_key_value_inc_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.inc("alpha", 2)
+    assert new_value == 3
+    assert await dict_key_value.get("alpha") == 3
+
+
+@pytest.mark.asyncio()
+async def test_key_value_inc_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.inc("gamma", 4)
+    assert new_value == 4
+    assert await dict_key_value.has("gamma")
+
+
+@pytest.mark.asyncio()
+async def test_key_value_inc_rejects_non_numeric_amount(dict_key_value: KeyValue[str, int]) -> None:
+    with pytest.raises(TypeError):
+        await dict_key_value.inc("alpha", cast(Any, "invalid"))
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.chmax("alpha", 10)
+    assert new_value == 10
+    assert await dict_key_value.get("alpha") == 10
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_ignores_smaller(dict_key_value: KeyValue[str, int]) -> None:
+    initial = await dict_key_value.get("alpha")
+    result = await dict_key_value.chmax("alpha", 0)
+    assert result == initial
+    assert await dict_key_value.get("alpha") == initial
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
+    result = await dict_key_value.chmax("gamma", 7)
+    assert result == 7
+    assert await dict_key_value.get("gamma") == 7
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_rejects_non_numeric_value(dict_key_value: KeyValue[str, int]) -> None:
+    with pytest.raises(TypeError):
+        await dict_key_value.chmax("alpha", cast(Any, "wrong"))
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_inc_rejects_non_numeric_value(dict_key_value_memory: DictBasedKeyValue[str, Any]) -> None:
+    await dict_key_value_memory.set("alpha", cast(Any, "na"))
+    with pytest.raises(TypeError):
+        await dict_key_value_memory.inc("alpha", 1)
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_chmax_rejects_non_numeric_value(
+    dict_key_value_memory: DictBasedKeyValue[str, Any],
+) -> None:
+    await dict_key_value_memory.set("alpha", cast(Any, "na"))
+    with pytest.raises(TypeError):
+        await dict_key_value_memory.chmax("alpha", 1)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_inc_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-inc-{uuid4().hex}",
+            "partition-inc",
+            str,
+            int,
+        )
+        collection = await key_value.ensure_collection()
+        await collection.insert_one(
+            {
+                "partition_id": "partition-inc",
+                "key": "alpha",
+                "value": "oops",
+            }
+        )
+
+        with pytest.raises(TypeError):
+            await key_value.inc("alpha", 1)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_chmax_behaves_like_max(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-chmax-{uuid4().hex}",
+            "partition-chmax",
+            str,
+            int,
+        )
+        assert await key_value.chmax("alpha", 5) == 5
+        assert await key_value.chmax("alpha", 3) == 5
+        assert await key_value.chmax("alpha", 9) == 9
+        assert await key_value.get("alpha") == 9
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_chmax_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-chmax-bad-{uuid4().hex}",
+            "partition-chmax-bad",
+            str,
+            int,
+        )
+        collection = await key_value.ensure_collection()
+        await collection.insert_one(
+            {
+                "partition_id": "partition-chmax-bad",
+                "key": "alpha",
+                "value": "oops",
+            }
+        )
+
+        with pytest.raises((TypeError, pydantic.ValidationError)):
+            await key_value.chmax("alpha", 1)
+
+
 def test_thread_safe_async_lock_blocks_threads() -> None:
    lock = _ThreadSafeAsyncLock()
    allow_second = threading.Event()
@@ -1071,7 +1221,9 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
        MongoClientPool,
    )

-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        collection = MongoBasedCollection[Any](
            client_pool, temporary_mongo_database.name, "test", "test-123", ["rollout_id"], Rollout
        )
@@ -1110,31 +1262,6 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
        assert not await span_kv.has("span-123")


-@pytest.mark.mongo
-@pytest.mark.asyncio()
-async def test_mongo_based_collection_rejects_duplicate_payload(temporary_mongo_database: AsyncDatabase[Any]) -> None:
-    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
-
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
-        collection = MongoBasedCollection[Any](
-            client_pool,
-            temporary_mongo_database.name,
-            f"duplicate-check-{uuid4().hex}",
-            "partition-dup",
-            ["rollout_id"],
-            Rollout,
-        )
-        await collection.ensure_collection()
-        start_time = time.time()
-        first = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-        duplicate = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-
-        with pytest.raises(ValueError, match="duplicate primary key"):
-            await collection.insert([first, duplicate])
-
-        assert await collection.size() == 0
-
-
@pytest.mark.mongo
@pytest.mark.asyncio()
 async def test_mongo_ensure_collection_creates_partition_scoped_index(
@@ -1143,7 +1270,9 @@ async def test_mongo_ensure_collection_creates_partition_scoped_index(
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

    collection_name = f"ensure-{uuid4().hex}"
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        collection = MongoBasedCollection[Any](
            client_pool,
            temporary_mongo_database.name,
@@ -1173,7 +1302,9 @@ async def test_mongo_ensure_collection_survives_concurrent_calls(temporary_mongo
    collection_name = f"ensure-{uuid4().hex}"

    async def ensure_once() -> None:
-        async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        async with MongoClientPool[Mapping[str, Any]](
+            mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+        ) as client_pool:
            collection = MongoBasedCollection(
                client_pool,
                temporary_mongo_database.name,
@@ -1204,7 +1335,9 @@ async def test_mongo_ensure_collection_repeats_without_altering_indexes(
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

    collection_name = f"ensure-{uuid4().hex}"
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        collection = MongoBasedCollection(
            client_pool, temporary_mongo_database.name, collection_name, "partition-repeat", ["index"], SampleItem
        )
@@ -1225,7 +1358,9 @@ async def _with_mongo_collections(
 ) -> Any:
    from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections

-    async with MongoClientPool(db.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
        collections = MongoLightningCollections(
            client_pool=client_pool,
            database_name=db.name,
@@ -19,12 +19,13 @@ import asyncio
 import logging
 import sys
 import time
-from typing import List, Optional, Sequence, cast
+from typing import Any, List, Optional, Protocol, Sequence, cast
 from unittest.mock import Mock

 import pytest
 from pydantic import BaseModel

+from agentlightning.store import CollectionBasedLightningStore
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
@@ -45,7 +46,13 @@ from agentlightning.types import (
    TraceStatus,
 )

-# Typing tests
+pytestmark = [pytest.mark.store]
+
+
+class FakeTimeController(Protocol):
+    def set(self, value: float) -> None: ...
+
+    def advance(self, delta: float) -> None: ...


 def test_paginated_result_behaves_like_sequence() -> None:
@@ -1308,7 +1315,8 @@ async def test_add_many_spans_handles_mixed_rollouts_and_attempts(store_fixture:
    duplicate_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)

    stored = await store_fixture.add_many_spans([span_first, span_retry, span_second, duplicate_first])
-    assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}
+    if isinstance(store_fixture, InMemoryLightningStore):
+        assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}

    spans_first = await store_fixture.query_spans(first.rollout_id)
    assert {span.span_id for span in spans_first} >= {span_first.span_id, span_retry.span_id}
@@ -3035,7 +3043,7 @@ async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore,
    """Test that healthcheck detects and handles unresponsive conditions."""
    # Create rollout with short unresponsive timeout but no retry for unresponsive
    config = RolloutConfig(
-        unresponsive_seconds=0.1,  # Very short unresponsive timeout
+        unresponsive_seconds=0.2,  # Very short unresponsive timeout
        max_attempts=3,
        retry_condition=["timeout"],  # Note: "unresponsive" not in retry_condition
    )
@@ -3054,7 +3062,7 @@ async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore,
    assert running_attempts[0].last_heartbeat_time is not None

    # Wait for unresponsive timeout
-    await asyncio.sleep(0.15)  # Wait longer than unresponsive_seconds
+    await asyncio.sleep(0.25)  # Wait longer than unresponsive_seconds

    # Verify attempt was marked as unresponsive
    attempts_after = await store_fixture.query_attempts(rollout.rollout_id)
@@ -3383,3 +3391,56 @@ async def test_query_resources_returns_all_fields(store_fixture: LightningStore)
        assert res.update_time > 0
        assert res.version >= 1
        assert res.resources is not None
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_allows_initial_scan(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """The first watchdog scan should run immediately even with debouncing enabled."""
+    fake_time.set(100.0)
+
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_blocks_until_interval(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """Subsequent scans wait until the debounce window elapses."""
+    fake_time.set(200.0)
+
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+    assert not await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+    fake_time.advance(debounced_store._scan_debounce_seconds - 1)  # pyright: ignore[reportPrivateUsage]
+    assert not await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+    fake_time.advance(1.0)
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_allows_single_concurrent_scan(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """When multiple coroutines race, only one should trigger the scan."""
+    fake_time.set(300.0)
+
+    results = await asyncio.gather(
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+    )
+
+    assert results.count(True) == 1
+    assert results.count(False) == 2
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_disabled_when_zero(store_fixture: CollectionBasedLightningStore[Any]) -> None:
+    """Setting debounce to zero should run the scan every time."""
+    assert store_fixture._scan_debounce_seconds == 0  # pyright: ignore[reportPrivateUsage]
+
+    assert await store_fixture._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+    assert await store_fixture._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yuge Zhang	bb6f1594a8	minor fix	2025-12-13 12:47:53 +08:00
Yuge Zhang	af28550701	move out search r1	2025-12-13 12:44:47 +08:00
Yuge Zhang	79df678ac1	Move search-r1 example	2025-12-13 12:41:44 +08:00
Yuge Zhang	05cc455cb1	fix tests	2025-12-13 12:37:20 +08:00
Yuge Zhang	93f2177639	Fix script invocation in legacy	2025-12-13 12:32:12 +08:00
jinghuan-Chen	a6078caa6c	fix TraceTree/match_rewards assign_to elements. (#403 )	2025-12-13 00:25:54 +08:00
Yuge Zhang	f1a8072546	Add ChartQA to catalog (#409 )	2025-12-12 23:39:41 +08:00
Totoluo	60f9955606	Multi-modal example: ChartQA (#379 ) Co-authored-by: Totoluo <52833580+Yingluo-momo@users.noreply.github.com> Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>	2025-12-12 17:22:35 +08:00
Yuge Zhang	1d199b21c7	Support customizing trainer and daemon in VERL (#407 )	2025-12-12 13:31:30 +08:00
Yuge Zhang	5f62ecb6f4	Split vllm 0.10.2 from 0.11.0 (#394 )	2025-12-12 12:10:53 +08:00
Yuge Zhang	1948c2ba6d	Sunset HTTP tracer and Refactor tests (#402 )	2025-12-11 23:57:23 +08:00
Yuge Zhang	1e36e660b1	Support `with_llm_proxy` and `with_store` in algorithms (#398 )	2025-12-11 16:50:42 +08:00
Yuge Zhang	267b9936bb	Support image urls export in TracerTraceToTriplets (#400 )	2025-12-11 16:47:28 +08:00
Wang Zilong	14714ded2b	add youtu-agent in community projects (#399 )	2025-12-11 15:40:49 +08:00
Yuge Zhang	c3f5cc7a39	Initialize contribution area (#396 )	2025-12-10 23:46:06 +08:00
Yuge Zhang	ee0fffd3a2	Upgrade tinker dependency (#393 )	2025-12-10 23:23:43 +08:00
Yuge Zhang	94d1cd780e	Store Benchmark - Part 6 (#388 )	2025-12-10 21:11:29 +08:00
Vasu	bbd5c2a30a	fix: handle ref_in_actor flag for LoRA compatibility with verl 0.6.0 (#386 )	2025-12-10 18:30:38 +08:00
Yuge Zhang	c6f4e6c283	Add playground workflow (#391 )	2025-12-10 18:27:21 +08:00
Ni Hao	8c504518bb	add weave tracer (#277 )	2025-12-09 18:59:11 +08:00
etsplz	337cce7fdc	Fix redundant cancel tracebacks on ctrl+c (issue #343 ) (#370 )	2025-12-08 17:18:07 +08:00
Yuge Zhang	42c63d7a01	Store Benchmark - Part 5 (#380 )	2025-12-08 12:43:18 +08:00
Yuge Zhang	5ecd23792d	Fix trainer dev warning (#378 )	2025-12-08 00:43:56 +08:00
Yuge Zhang	ad89e173e1	Fix TracesTable story (#375 )	2025-12-06 14:25:09 +08:00
Yuge Zhang	feebaec24c	Add AGENTS.md (#374 )	2025-12-06 13:16:26 +08:00
				`@@ -0,0 +1 @@`
				`# Put contrib-related gitignore files here.`