Bump to 0.3.1

[Release] v0.3.0 (#427 )
Fix documentation and dashboard building issues (#429 )
2025-12-24 11:37:57 +08:00 · 2025-12-24 09:46:58 +08:00 · 2025-12-23 23:55:59 +08:00 · 2025-12-23 23:32:19 +08:00 · 2025-12-23 14:57:24 +08:00 · 2025-12-22 11:12:21 +08:00
315 changed files with 65781 additions and 9492 deletions
@@ -0,0 +1,14 @@
+.venv
+**/.venv
+__pycache__
+.git
+.gitignore
+**/node_modules
+dist
+build
+.env
+docker
+.pytest_cache
+.vscode
+**/*.log
+examples/**/data
@@ -0,0 +1,29 @@
+name: Badge - Azure
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Azure
+    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-azure.yml', label: 'azure', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -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 });
@@ -0,0 +1,29 @@
+name: Badge - Claude Code
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Claude Code
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-claude-code.yml', label: 'claude-code', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Compatibility
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Backward Compatibility
+    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-compat.yml', label: 'examples-compat', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -7,6 +7,11 @@ on:
      - Examples - Spider
      - Examples - APO
      - Examples - Unsloth
+      - Examples - Tinker
+      - Examples - Azure
+      - Examples - Claude Code
+      - Examples - RAG
+      - Examples - ChartQA
    types: [completed]

  workflow_dispatch:
@@ -31,5 +36,10 @@ jobs:
              { workflow: 'examples-spider.yml', label: 'examples-spider.stable', variants: ['stable'] },
              { workflow: 'examples-apo.yml', label: 'examples-apo.stable', variants: ['stable'] },
              { workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
+              { workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
+              { workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
+              { workflow: 'examples-claude-code.yml', label: 'examples-claude-code.stable', variants: ['stable'] },
+              { 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 });
@@ -24,6 +24,6 @@ jobs:
          script: |
            const badgeAggregation = require('./scripts/badge_aggregation.js');
            const dependencies = [
-              { workflow: 'examples-spider.yml', label: 'spider', variants: ['stable', 'legacy'] },
+              { workflow: 'examples-spider.yml', label: 'spider', variants: ['stable'] },
            ];
            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Tinker
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Tinker
+    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-tinker.yml', label: 'tinker', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,31 @@
+name: Badge - Unit Test
+
+on:
+  workflow_run:
+    workflows:
+      - CPU Test
+      - GPU Test
+    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: 'tests-full.yml', label: 'tests-full', variants: ['legacy', 'stable'] },
+              { workflow: 'tests.yml', label: 'tests', variants: ['legacy', 'stable', 'Lint', 'documentation', 'JavaScript'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,441 @@
+name: Benchmark
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+  schedule:
+    # Every Monday and Thursday at 3 AM UTC+8
+    - cron: '0 19 * * 0,3'
+
+jobs:
+  benchmark:
+    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:
+        backend:
+          - id: memory
+            compose_file: compose.prometheus-memory-store.yml
+          - id: mongo
+            compose_file: compose.prometheus-mongo-store.yml
+        workload:
+          - id: scenario-minimal-scale
+            display: Minimal production scale
+            kind: scenario
+            store_workers: 4
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 45
+            args: >-
+              --mode batch
+              --total-tasks 4096
+              --batch-size 256
+              --n-runners 32
+              --max-rounds 6
+              --sleep-seconds 0.5
+          - id: scenario-medium-scale
+            display: Medium production scale
+            kind: scenario
+            store_workers: 16
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 45
+            args: >-
+              --mode batch
+              --total-tasks 10000
+              --batch-size 1000
+              --n-runners 100
+              --max-rounds 10
+              --sleep-seconds 0.1
+          - 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 300
+              --max-rounds 6
+              --sleep-seconds 0.1
+          - id: scenario-large-batch
+            display: Large batch waves
+            kind: scenario
+            store_workers: 96
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 120
+            args: >-
+              --mode batch
+              --total-tasks 50000
+              --batch-size 8192
+              --n-runners 1000
+              --max-rounds 3
+              --sleep-seconds 0.1
+          - id: scenario-long-queues
+            display: Long rollout queues
+            kind: scenario
+            store_workers: 48
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 120
+            args: >-
+              --mode batch_partial
+              --total-tasks 50000
+              --batch-size 1024
+              --n-runners 256
+              --remaining-tasks 4096
+              --max-rounds 4
+              --sleep-seconds 0.1
+          - id: scenario-high-concurrency
+            display: High-throughput concurrent requests
+            kind: scenario
+            store_workers: 96
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 120
+            args: >-
+              --mode single
+              --total-tasks 50000
+              --concurrency 2048
+              --n-runners 256
+              --max-rounds 2
+              --sleep-seconds 0.1
+          - id: scenario-heavy-traces
+            display: Heavy rollouts with deep traces
+            kind: scenario
+            store_workers: 96
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 60
+            args: >-
+              --mode batch_partial
+              --total-tasks 10000
+              --batch-size 1024
+              --remaining-tasks 256
+              --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:
+      PYTHONUNBUFFERED: "1"
+      STORE_URL: http://localhost:4747
+      STORE_API_URL: http://localhost:4747/v1/agl
+      PROM_URL: http://localhost:9090
+      GITHUB_ACTIONS_TIMEOUT_MINUTES: ${{ matrix.workload.timeout }}
+      WORKLOAD_KIND: ${{ matrix.workload.kind }}
+      WORKLOAD_ID: ${{ matrix.workload.id }}
+      BACKEND_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.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
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - name: Sync dependencies
+        run: uv sync --frozen --extra mongo --group core-stable --group dev
+
+      - name: Check disk space
+        run: df -h
+
+      - name: Reset benchmark data directories
+        run: |
+          set -euo pipefail
+          cd docker
+          rm -rf data
+          bash setup.sh
+
+      - name: Launch ${{ matrix.backend.id }} Prometheus stack
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+          docker compose -f "$COMPOSE_FILE" up -d --quiet-pull
+
+      - name: Wait for store readiness
+        run: |
+          set -euo pipefail
+          for attempt in {1..60}; do
+            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+              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
+          # 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 workload start
+        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - 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.workload.args }}
+
+      - 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: 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_FILE"
+            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_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() }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+
+      - name: Archive Prometheus metrics
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          if [ -d docker/data/prometheus ]; then
+            tar -C docker/data -czf "$ARTIFACT_DIR/${PROM_ARCHIVE_BASENAME}.tar.gz" prometheus
+          fi
+
+      - name: Upload workload artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error
+
+  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
+            needs_mongo: false
+            runner: ubuntu-latest
+          - id: mongo
+            needs_mongo: true
+            runner: ubuntu-latest
+        workload:
+          - id: high-insert
+            total_tasks: 50000
+            concurrency: 2048
+            type: insert
+          - id: medium-insert
+            total_tasks: 50000
+            concurrency: 128
+            type: insert
+          - id: low-insert
+            total_tasks: 50000
+            concurrency: 4
+            type: insert
+          - id: high-dequeue
+            total_tasks: 50000
+            concurrency: 2048
+            type: dequeue
+          - id: medium-dequeue
+            total_tasks: 50000
+            concurrency: 128
+            type: dequeue
+          - id: low-dequeue
+            total_tasks: 50000
+            concurrency: 4
+            type: dequeue
+    env:
+      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
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - name: Sync dependencies
+        run: uv sync --frozen --extra mongo --group core-stable --group dev
+
+      - name: Launch MongoDB
+        if: ${{ matrix.backend.needs_mongo }}
+        run: |
+          set -euo pipefail
+          cd docker
+          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 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 2
+          done
+          echo "MongoDB did not become ready in time" >&2
+          docker compose -f compose.mongo.yml logs mongo
+          exit 1
+
+      - name: Run collection benchmark
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          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: Show collection benchmark summary
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          if [ -f "$SUMMARY_FILE" ]; then
+            echo "Collection benchmark summary (${{ matrix.backend.id }}):"
+            cat "$SUMMARY_FILE"
+          else
+            echo "Summary file not found: $SUMMARY_FILE"
+          fi
+
+      - name: Stop MongoDB
+        if: ${{ always() && matrix.backend.needs_mongo }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f compose.mongo.yml down -v || true
+
+      - name: Upload collection artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'APO - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -0,0 +1,98 @@
+name: Examples - Azure
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: '0 20 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-azure, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Azure - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Azure - {0}', github.event_name) }}
+
+jobs:
+  azure:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-azure' ||
+      github.event.action == 'ci-all'
+    name: Azure (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
+    timeout-minutes: 400
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.12'
+          setup-script: 'stable'
+      fail-fast: false
+    steps:
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group core-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-azure-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Azure Login
+        run: |
+          az login --identity
+        shell: bash
+
+      - name: Azure OpenAI Sanity Check
+        run: |
+          source .venv/bin/activate
+          cd examples/azure
+          python capital_agent.py
+        shell: bash
+        env:
+          AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+        id: azure_openai_sanity_check
+
+      - name: Azure OpenAI Supervised Fine-tuning
+        run: |
+          source .venv/bin/activate
+          cd examples/azure
+          python train_capital_agent.py --n-iterations 2 --cleanup
+        shell: bash
+        env:
+          AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+          AZURE_OPENAI_API_VERSION: 2025-04-01-preview
+          AZURE_RESOURCE_GROUP: ${{ secrets.AZURE_RESOURCE_GROUP }}
+          AZURE_RESOURCE_NAME: ${{ secrets.AZURE_RESOURCE_NAME }}
+        id: azure_openai_finetune
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Calc-X - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -22,12 +22,12 @@ run-name: >-
      || format('Calc-X - {0}', github.event_name) }}

 jobs:
-  calc-x:
+  calc-x-perf:
    if: >
      github.event_name != 'repository_dispatch' ||
      github.event.action == 'ci-calc-x' ||
      github.event.action == 'ci-all'
-    name: Calc-X (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    name: Calc-X Performance (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
    timeout-minutes: 90
    strategy:
@@ -74,7 +74,7 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-calc-x-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-calc-x-performance-${{ matrix.python-version }}-${{ matrix.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

@@ -116,13 +116,11 @@ jobs:
      # Don't ask why. Don't touch this.
      - name: Calc-X training
        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
-          sleep 10
+          python train_calc_agent.py --val-file data/test_mini.parquet --ci
        shell: bash
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
@@ -137,14 +135,126 @@ jobs:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

-      - name: Calc-X training LLM Proxy
+  calc-x-variants:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-calc-x' ||
+      github.event.action == 'ci-all'
+    name: Calc-X Variants (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 90
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-calc-x-variants-${{ 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 Calc-X dataset
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
+          unzip calc-x-data.zip -d data
+          rm calc-x-data.zip
+
+      - name: Calc-X MCP sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run tests/test_mcp_calculator.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: Calc-X sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run legacy_calc_agent_debug.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Training with local model
        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 --llm-proxy
+          hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir data/qwen_model
+          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --model $(realpath data/qwen_model)
+          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_local_model
+
+      - name: Validate training with local model
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_local_model.outputs.project_name }} ${{ steps.calc_x_train_local_model.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with LLM Proxy
+        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 --llm-proxy
          sleep 10
        shell: bash
        env:
@@ -152,7 +262,113 @@ jobs:
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
        id: calc_x_train_llm_proxy

-      - name: Calc-X training with external store
+      - name: Validate training with LLM Proxy
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_llm_proxy.outputs.project_name }} ${{ steps.calc_x_train_llm_proxy.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Setup Docker environments
+        run: ./scripts/mongodb_docker_run.sh
+        shell: bash
+
+      - name: Training with MongoDB
+        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 --mongo-uri mongodb://localhost:27017/?replicaSet=rs0
+          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_mongo
+
+      - name: Validate training with MongoDB
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_mongo.outputs.project_name }} ${{ steps.calc_x_train_mongo.outputs.run_name }}
+        env:
+          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 trajectory level aggregation
+        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 --trajectory-level
+          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_trajectory_level
+
+      - name: Validate training with trajectory level aggregation
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_trajectory_level.outputs.project_name }} ${{ steps.calc_x_train_trajectory_level.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with Weave
+        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 --weave
+          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_weave
+
+      - name: Validate training with Weave
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_weave.outputs.project_name }} ${{ steps.calc_x_train_weave.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with external store
        run: |
          set -euo pipefail
          source .venv/bin/activate
@@ -182,7 +398,15 @@ jobs:
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
        id: calc_x_train_external_store

-      - name: Calc-X training with role-based environment variables
+      - name: Validate training with external store
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_external_store.outputs.project_name }} ${{ steps.calc_x_train_external_store.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with role-based environment variables
        run: |
          set -euo pipefail
          source .venv/bin/activate
@@ -203,3 +427,12 @@ jobs:
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_role_based_env_var
+
+      - name: Validate training with role-based environment variables
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_role_based_env_var.outputs.project_name }} ${{ steps.calc_x_train_role_based_env_var.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -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 }}
@@ -0,0 +1,151 @@
+name: Examples - Claude Code
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: "0 20 * * *"
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-claude-code, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Claude Code - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Claude Code - {0}', github.event_name) }}
+
+jobs:
+  claude-code:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-claude-code' ||
+      github.event.action == 'ci-all'
+    name: Claude Code (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+          - python-version: "3.12"
+            setup-script: "stable"
+          - python-version: "3.13"
+            setup-script: "latest"
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-gpu-stable
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Download model
+        run: |
+          source .venv/bin/activate
+          python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
+
+      - name: Launch vLLM server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+            --max-model-len 131072 \
+            --enable-auto-tool-choice \
+            --tool-call-parser qwen3_coder \
+            --port 45993 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:45993/v1/models > /dev/null 2>&1; then
+              echo "vLLM server is ready!"
+              VLLM_READY=1
+              break
+            fi
+            echo "Waiting for vLLM server to be ready... (${i})"
+            sleep 5
+          done
+          if [[ "$VLLM_READY" != "1" ]]; then
+            echo "vLLM server failed to start!"
+            exit 1
+          fi
+
+      - name: Claude Code sanity check with vLLM models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py vllm --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct --base-url http://localhost:45993/v1 --debug
+        shell: bash
+
+      - name: Upload sanity check artifacts for vLLM
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error
+
+      - name: Cleanup vLLM
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+          rm -rf examples/claude_code/data/
+          rm -rf examples/claude_code/logs/
+
+      - name: Claude Code sanity check with OpenAI models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+
+      - name: Upload sanity check artifacts for OpenAI
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Backward Compatibility - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -56,7 +56,7 @@ jobs:
            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
      - name: Override VERL (stable)
        run: |
-          uv pip install verl==0.5.0
+          uv pip install verl==0.5.0 vllm==0.10.2
        if: matrix.setup-script == 'stable'
      - name: Freeze dependencies
        run: |
@@ -0,0 +1,179 @@
+name: Examples - RAG
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: '0 22 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-rag, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'RAG - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('RAG - {0}', github.event_name) }}
+
+jobs:
+  rag:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-rag' ||
+      github.event.action == 'ci-all'
+    name: RAG (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-rag-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare RAG dataset
+        run: |
+          set -euo pipefail
+          cd examples/rag
+          mkdir -p data
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" -O data/dataset_tiny.parquet
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" -O data/chunks_candidate_tiny.pkl
+          uv run gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" -O data/index_hnsw_faiss_n32e40_tiny.index
+
+      - name: Run WIKI Retriever MCP Server
+        run: |
+          set -euo pipefail
+          cd examples/rag
+          uv run python wiki_retriever_mcp.py &
+          for i in {1..20}; do
+            sleep 5
+            if nc -z localhost 8099; then
+              echo "MCP server is up!"
+              exit 0
+            else
+              echo "Waiting for MCP server to start..."
+            fi
+          done
+          echo "MCP server failed to start within expected time."
+          exit 1
+
+      - name: Run vLLM Server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen2.5-1.5B-Instruct \
+            --enable-auto-tool-choice \
+            --tool-call-parser hermes \
+            --port 8000 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:8000/v1/models > /dev/null 2>&1; then
+              echo "vLLM server is ready!"
+              VLLM_READY=1
+              break
+            fi
+            echo "Waiting for vLLM server to be ready... (${i})"
+            sleep 5
+          done
+          if [[ "$VLLM_READY" != "1" ]]; then
+            echo "vLLM server failed to start!"
+            exit 1
+          fi
+
+      - name: Run RAG Sanity check
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/rag
+          uv run python rag_agent.py
+        shell: bash
+
+      - name: Stop vLLM Server
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+
+      - name: RAG training
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/rag
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_rag.py fast
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: rag_train
+
+      - name: Validate RAG training
+        run: |
+          set -ex
+          # Allow up to 5 rollouts to fail to produce rewards
+          uv run scripts/validate_example_wandb.py ${{ steps.rag_train.outputs.project_name }} ${{ steps.rag_train.outputs.run_name }} --reward-tolerance 5
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Spider - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -33,8 +33,7 @@ jobs:
    strategy:
      matrix:
        include:
-        - python-version: '3.10'
-          setup-script: 'legacy'
+        # legacy is omitted because langchain doesn't work with legacy vllm versions
        - python-version: '3.12'
          setup-script: 'stable'
        - python-version: '3.13'
@@ -58,13 +57,13 @@ jobs:
      - name: Sync dependencies (latest)
        run: |
          uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --group torch-gpu-stable
+            --group dev --group experiment --group agents --group langchain --group torch-gpu-stable
        if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (stable & legacy)
+      - name: Sync dependencies (stable)
        run: |
          uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+            --group dev --group experiment --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
      - name: Freeze dependencies
        run: |
          set -ex
@@ -121,7 +120,7 @@ jobs:
      - name: Validate Spider training
        run: |
          set -ex
-          uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
+          uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }} --reward-tolerance 5
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -0,0 +1,170 @@
+name: Examples - Tinker
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-tinker, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Tinker - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Tinker - {0}', github.event_name) }}
+
+jobs:
+  tinker:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-tinker' ||
+      github.event.action == 'ci-all'
+    name: Tinker (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
+    timeout-minutes: 150
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-cpu --group core-stable --group tinker
+      - name: Freeze dependencies
+        run: |
+          set -euo pipefail
+          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-tinker-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      # 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
+          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 }}
+
+      - name: Tinker Hello
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python hello.py oneclick --ci
+        shell: bash
+        env:
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Evaluate (GPT-4.1)
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          mkdir -p logs
+          python q20_evaluate.py --ci --model gpt-4.1 --output-file logs/q20_evaluate_gpt-4.1.jsonl
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Evaluate (Qwen3-30B-A3B-Instruct-2507)
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python q20_evaluate.py --ci --model Qwen/Qwen3-30B-A3B-Instruct-2507 --output-file logs/q20_evaluate_qwen3-30b-a3b.jsonl
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Training Dry Run
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python q20_train.py dryrun --model qwen4b
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Training
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          agl store --port 4747 &
+          sleep 5
+          python q20_train.py runner --n-runners 4 &
+          sleep 5
+          python q20_train.py algo --model qwen4b --ci
+          sleep 5
+
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+          pkill -f q20_train.py && echo "SIGTERM sent to q20_train.py" || echo "No q20_train.py process found"
+          while pgrep -f q20_train.py; do
+            echo "Waiting for q20_train.py to finish..."
+            sleep 5
+          done
+          echo "q20_train.py has finished."
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Unsloth - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -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!"
@@ -26,6 +26,14 @@ jobs:
      - name: Sync dependencies
        run: uv sync --frozen --no-default-groups --group dev

+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Build dashboard
+        run: cd dashboard && npm run build
+
      - name: Get current version
        id: get_version
        run: |
@@ -60,6 +60,14 @@ jobs:
      - name: Sync dependencies
        run: uv sync --frozen --no-default-groups --group dev

+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Build dashboard
+        run: cd dashboard && npm run build
+
      - name: Build package
        run: |
          uv build
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'GPU Test - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -27,7 +27,138 @@ 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: ${{ matrix.mark.runs-on }}
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        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
+          # Similar for Weave.
+          - id: weave
+            display-name: Weave
+            pytest-mark: 'weave'
+            runs-on: ubuntu-latest  # No GPU tests for Weave.
+            has-gpu: false
+          # 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 weave 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'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        if: matrix.mark.has-gpu
+        run: nvidia-smi
+      - 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.env.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        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 weave --extra mongo --group dev --group agents --group langchain --group torch-gpu-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 weave --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 weave --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 weave --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, gpu)
+        if: matrix.env.setup-script == 'legacy' && matrix.mark.has-gpu
+        run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group torch-gpu-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 weave --extra mongo --group dev --group agents --group core-legacy
+
+      - 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-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
+        run: cd dashboard && npm run build
+
+      - name: Setup Docker environments
+        run: ./scripts/mongodb_docker_run.sh
+        shell: bash
+
+      - 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 }}
+
+      # mongo, openai, gpu, all enabled by default
+      - name: Run tests
+        run: |
+          uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}${{ matrix.env.setup-script == 'legacy' && ' and not langchain' || '' }}"
+        env:
+          PYTEST_ADDOPTS: "--color=yes"
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+          AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
+
+
+  minimal-examples:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-gpu' ||
+      github.event.action == 'ci-all'
+    name: Minimal Examples with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})

    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
    timeout-minutes: 30
@@ -55,11 +186,15 @@ jobs:
        run: uv lock --upgrade
        if: matrix.setup-script == 'latest'
      - name: Sync dependencies (latest)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-stable
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group torch-gpu-stable
        if: matrix.setup-script == 'latest'
-      - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+      - name: Sync dependencies (stable)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
+      # Don't install langchain for legacy dependency because it has conflicts with torch.
+      - name: Sync dependencies (legacy)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-legacy
+        if: matrix.setup-script == 'legacy'
      - name: Freeze dependencies
        run: |
          set -ex
@@ -69,7 +204,7 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

@@ -80,10 +215,149 @@ jobs:
          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}

-      - name: Run tests
+      - name: Write Traces via Otel Tracer
        run: |
-          uv run pytest -v --durations=0 tests
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_traces.py otel
+          sleep 5
+
+      - name: Write Traces via AgentOps Tracer
        env:
-          PYTEST_ADDOPTS: "--color=yes"
          OPENAI_BASE_URL: http://localhost:12306/
          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_traces.py agentops
+          sleep 5
+
+      - name: Write Traces with Operations
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_traces.py operation
+          sleep 5
+
+      - name: Write Traces via Otel Tracer with Client
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          agl store --port 45993 --log-level DEBUG &
+          sleep 5
+          python write_traces.py otel --use-client
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+
+      - name: Write Traces via AgentOps Tracer with Client
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          agl store --port 45993 --log-level DEBUG &
+          sleep 5
+          python write_traces.py agentops --use-client
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+
+      - name: vLLM Server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python vllm_server.py Qwen/Qwen2.5-0.5B-Instruct
+
+      - name: LLM Proxy (OpenAI backend)
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+
+          python llm_proxy.py openai gpt-4.1-mini &
+
+          LLM_PROXY_READY=0
+          for attempt in $(seq 1 30); do
+            if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
+              LLM_PROXY_READY=1
+              break
+            fi
+            sleep 2
+          done
+          if [[ "$LLM_PROXY_READY" != "1" ]]; then
+            echo "LLM proxy failed to become healthy" >&2
+            exit 1
+          fi
+
+          python llm_proxy.py test gpt-4.1-mini
+
+          pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
+          while pgrep -f llm_proxy.py; do
+            echo "Waiting for llm_proxy.py to finish..."
+            sleep 5
+          done
+
+      - name: LLM Proxy (vLLM backend)
+        if: matrix.setup-script != 'legacy'  # Skip if return_token_ids is not supported
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct &
+
+          LLM_PROXY_READY=0
+          for attempt in $(seq 1 30); do
+            if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
+              LLM_PROXY_READY=1
+              break
+            fi
+            sleep 2
+          done
+          if [[ "$LLM_PROXY_READY" != "1" ]]; then
+            echo "LLM proxy failed to become healthy" >&2
+            exit 1
+          fi
+
+          python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
+
+          pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
+          while pgrep -f llm_proxy.py; do
+            echo "Waiting for llm_proxy.py to finish..."
+            sleep 5
+          done
+
+      - name: MultiMetrics backend example
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_metrics.py --duration 8 --prom-port 9105 --prom-host 0.0.0.0 2>&1 | tee metrics.log &
+          pid=$!
+
+          for attempt in $(seq 1 20); do
+            if curl -sSf http://localhost:9105/metrics | grep -q minimal_requests_total; then
+              echo "Metrics endpoint responding"
+              wait $pid
+              cat metrics.log
+              exit 0
+            fi
+            sleep 1
+          done
+
+          echo "Metrics endpoint did not respond"
+          exit 1
@@ -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,19 +33,25 @@ 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 \
            --extra apo \
+            --extra weave \
            --extra verl \
+            --extra mongo \
            --group dev \
            --group torch-cpu \
            --group torch-stable \
            --group trl \
            --group tinker \
            --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
@@ -59,7 +66,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
@@ -70,6 +77,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
@@ -102,6 +111,10 @@ jobs:
      - name: Set source commit for docs
        run: |
          echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
+      - name: Verify OpenAPI specification is up-to-date
+        run: |
+          uv run --locked --no-sync python scripts/export_openapi.py
+          git diff --exit-code docs/assets/store-openapi.json
      - name: Build documentation
        run: uv run --locked --no-sync mkdocs build --strict
      - name: Upload docs artifact
@@ -114,7 +127,32 @@ 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'
+          # Similar for Weave.
+          - id: weave
+            display-name: Weave
+            pytest-mark: 'weave'
+          # 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 weave and not llmproxy and not utils'
+        env:
          - python-version: '3.10'
            setup-script: 'legacy'
          - python-version: '3.11'
@@ -125,7 +163,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:
@@ -133,16 +171,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 core-stable
-        if: matrix.setup-script == 'latest'
+        run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-stable
+        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 core-${{ matrix.setup-script }}
-        if: matrix.setup-script != 'latest'
+        run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
+        if: matrix.env.setup-script != 'latest'
      - name: Freeze dependencies
        run: |
          set -ex
@@ -152,17 +190,28 @@ 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
+        run: cd dashboard && npm run build
+
      - name: Run tests
        run: |
-          uv run pytest -v --durations=0 tests
+          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:
@@ -172,6 +221,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
@@ -1,8 +1,10 @@
 # Agentlightning specific files
 verl_old
 meta-llama/**
-debug/*.png
+**/debug/**/*.png
+**/debug/**/*.json
 requirements-freeze*.txt
+/playground

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -207,3 +209,12 @@ cython_debug/

 # Claude
 .claude/*.local.json
+
+# Dashboard generated files
+agentlightning/dashboard/**/*.css
+agentlightning/dashboard/**/*.js
+agentlightning/dashboard/**/*.html
+agentlightning/dashboard/**/*.svg
+
+# Docker data
+docker/data/
@@ -3,13 +3,14 @@ repos:
    rev: v6.0.0
    hooks:
      - id: end-of-file-fixer
+        exclude: (.*store-openapi\.json$)
      - id: trailing-whitespace
      - id: check-yaml
        exclude: ^mkdocs\.yml$
      - id: check-toml
      - id: check-added-large-files
        args: ["--maxkb=1024"]
-        exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)
+        exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)|(.*store-openapi\.json$)
      - id: check-shebang-scripts-are-executable
      - id: detect-private-key
  - repo: https://github.com/pycqa/isort
@@ -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
@@ -4,7 +4,7 @@

 # Agent Lightning⚡

-[![Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml)
+[![Unit Tests](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
 [![Documentation](https://img.shields.io/badge/GitHub%20Pages-Documentation-blue)](https://microsoft.github.io/agent-lightning/)
 [![PyPI version](https://badge.fury.io/py/agentlightning.svg)](https://badge.fury.io/py/agentlightning)
 [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
@@ -34,12 +34,19 @@ Read more on our [documentation website](https://microsoft.github.io/agent-light
 pip install agentlightning
 ```

+For the latest nightly build (cutting-edge features), you can install from Test PyPI:
+
+```bash
+pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre agentlightning
+```
+
 Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.

 To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).

 ## ⚡ Articles

+- 12/17/2025 [Adopting the Trajectory Level Aggregation for Faster Training](https://agent-lightning.github.io/posts/trajectory_level_aggregation/) Agent-lightning blog.
 - 11/4/2025 [Tuning ANY AI agent with Tinker ✕ Agent-lightning](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-1-1d8c9a397f0e) Medium. See also [Part 2](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-2-332c5437f0dc).
 - 10/22/2025 [No More Retokenization Drift: Returning Token IDs via the OpenAI Compatible API Matters in Agent RL](https://blog.vllm.ai/2025/10/22/agent-lightning.html) vLLM blog. See also [Zhihu writeup](https://zhuanlan.zhihu.com/p/1965067274642785725).
 - 8/11/2025 [Training AI Agents to Write and Self-correct SQL with Reinforcement Learning](https://medium.com/@yugez/training-ai-agents-to-write-and-self-correct-sql-with-reinforcement-learning-571ed31281ad) Medium.
@@ -51,6 +58,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) and their blog [*Stop Wrestling with Your Agent RL: How Youtu-Agent Achieved Stable, 128-GPU Scaling Without Breaking a Sweat*](https://spotted-coconut-df8.notion.site/Stop-Wrestling-with-Your-Agent-RL-How-Youtu-Agent-Achieved-Stable-128-GPU-Scaling-Without-Breaking-2ca5e8f089ba80539a98c582b65e0233).

 ## ⚡ Architecture

@@ -69,10 +77,11 @@ No rewrites, no lock-in, just a clear path from first rollout to steady improvem
 | Workflow | Status |
 |----------|--------|
 | CPU Tests | [![tests workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
-| GPU Tests | [![tests-full workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml) |
+| Full Tests | [![tests summary 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) |
+| UI Tests | [![UI Tests](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml) |
 | Examples Integration | [![examples summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml) |
 | Latest Dependency Compatibility | [![latest summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml) |
-| Legacy Examples Compatibility | [![examples compatibility 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) |
+| Legacy Examples Compatibility | [![compat summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml) |

 ## ⚡ Citation

@@ -92,7 +101,7 @@ If you find Agent Lightning useful in your research or projects, please cite our

 ## ⚡ Contributing

-This project welcomes contributions and suggestions. Start by reading the [Contributing Guide](docs/community/contributing.md) for environment setup, branching conventions, and pull request expectations. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
+This project welcomes contributions and suggestions. Start by reading the [Contributing Guide](docs/community/contributing.md) for recommended contribution points, environment setup, branching conventions, and pull request expectations. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

 When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.

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

-__version__ = "0.2.2"
+__version__ = "0.3.1"

 from .adapter import *
 from .algorithm import *
 from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
 from .config import *
 from .emitter import *
+from .env_var import *
 from .execution import *
 from .litagent import *
 from .llm_proxy import *
-from .logging import *
+from .logging import configure_logger  # deprecated  # type: ignore
+from .logging import setup as setup_logging  # type: ignore
+from .logging import setup_module as setup_module_logging  # type: ignore
 from .runner import *
 from .server import AgentLightningServer  # deprecated  # type: ignore
 from .store import *
@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.

-from typing import Generic, List, TypeVar
+from typing import Generic, Sequence, TypeVar

 from opentelemetry.sdk.trace import ReadableSpan

@@ -66,7 +66,7 @@ class Adapter(Generic[T_from, T_to]):
        raise NotImplementedError("Adapter.adapt() is not implemented")


-class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
+class OtelTraceAdapter(Adapter[Sequence[ReadableSpan], T_to], Generic[T_to]):
    """Base class for adapters that convert OpenTelemetry trace spans into other formats.

    This specialization of [`Adapter`][agentlightning.Adapter] expects a list of
@@ -84,7 +84,7 @@ class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
    """


-class TraceAdapter(Adapter[List[Span], T_to], Generic[T_to]):
+class TraceAdapter(Adapter[Sequence[Span], T_to], Generic[T_to]):
    """Base class for adapters that convert trace spans into other formats.

    This class specializes [`Adapter`][agentlightning.Adapter] for working with
@@ -4,7 +4,7 @@ from __future__ import annotations

 import json
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, TypedDict, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, Sequence, TypedDict, Union, cast

 from pydantic import TypeAdapter

@@ -208,7 +208,7 @@ class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
        children of the associated completion span.
    """

-    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
+    def get_tool_calls(self, completion: Span, all_spans: Sequence[Span], /) -> Iterable[Dict[str, Any]]:
        """Yield tool call payloads for a completion span.

        Args:
@@ -231,7 +231,7 @@ class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
            if tool_call:
                yield tool_call

-    def adapt(self, source: List[Span], /) -> List[OpenAIMessages]:
+    def adapt(self, source: Sequence[Span], /) -> List[OpenAIMessages]:
        """Transform trace spans into OpenAI chat payloads.

        Args:
@@ -6,18 +6,62 @@ import json
 import logging
 import re
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast

 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel

-from agentlightning.types import Span, SpanNames, Triplet
+from agentlightning.emitter.reward import get_reward_value
+from agentlightning.semconv import AGL_OPERATION, AGL_REWARD, LightningSpanAttributes
+from agentlightning.types import Span, Triplet
+from agentlightning.utils.otel import filter_and_unflatten_attributes

 from .base import TraceAdapter

 logger = logging.getLogger(__name__)


+def _attributes_get_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[str]:
+    """Get a string from the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        if key in attributes:
+            if isinstance(attributes[key], str):
+                return attributes[key]
+            else:
+                logger.warning(f"Attribute {key} is found but is not a string: {attributes[key]}")
+    return None
+
+
+def _attributes_get_ids_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[List[int]]:
+    """Get a list of integers from the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        if key in attributes:
+            if (isinstance(attributes[key], list) or isinstance(attributes[key], tuple)) and all(
+                isinstance(x, int) for x in attributes[key]
+            ):
+                return list(attributes[key])
+            else:
+                logger.warning(f"Attribute {key} is found but is not a list of integers: {attributes[key]}")
+    return None
+
+
+def _attributes_unflatten_multiple(
+    attributes: Dict[str, Any], keys: List[str]
+) -> Union[Dict[str, Any], List[Any], None]:
+    """Unflatten the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        result = filter_and_unflatten_attributes(attributes, key)
+        if result:
+            return result
+    return None
+
+
 class Transition(BaseModel):
    """A single transition within a reinforcement learning trajectory.

@@ -130,7 +174,7 @@ class TraceTree:
            if not should_visit(node):
                return False
            agent_name = node.agent_name()
-            vis_name = node.id[:8] + " (" + node.span.name + ")"
+            vis_name = node.id[-8:] + " (" + node.span.name + ")"
            if agent_name is not None:
                vis_name += " [" + agent_name + "]"
            dot.node(node.id, vis_name)  # type: ignore
@@ -307,30 +351,30 @@ class TraceTree:
        if agent_name is not None:
            return agent_name

+        # Case 6: Weave
+        is_agent_type = attributes.get("type") == "agent"
+        if is_agent_type:
+            agent_name = cast(Optional[str], attributes.get("agentlightning.operation.input.name"))
+            if agent_name is not None:
+                return agent_name
+
+        # Case 7: Weave + LangChain
+        if self.span.name.startswith("langchain.Chain."):
+            attributes_lc_name = cast(Optional[str], attributes.get("lc_name"))
+            if attributes_lc_name is not None:
+                return attributes_lc_name
+
    def maybe_reward_dict(self) -> dict[str, Any]:
        """Return a reward payload if the span encodes one.

        Returns:
            Dictionary containing reward metadata, or an empty dictionary when no reward is found.
        """
-        for key in [
-            "agentops.task.output",  # newer versions of agentops
-            "agentops.entity.output",
-        ]:
-            output = self.span.attributes.get(key)  # type: ignore
-            if output:
-                if isinstance(output, dict):
-                    return output
-                elif isinstance(output, str):
-                    try:
-                        return json.loads(output)
-                    except json.JSONDecodeError:
-                        return {}
-
-        # Latest emit reward format
-        if self.span.name == SpanNames.REWARD.value and self.span.attributes:
-            return {"type": "reward", "value": self.span.attributes.get("reward", None)}
-        return {}
+        reward_value = get_reward_value(self.span)
+        if reward_value is not None:
+            return {"type": "reward", "value": reward_value}
+        else:
+            return {}

    def is_reward_span(self) -> bool:
        """Return whether the span explicitly encodes a reward.
@@ -339,7 +383,17 @@ class TraceTree:
            `True` when the span payload describes a reward, otherwise `False`.
        """
        maybe_reward = self.maybe_reward_dict()
-        return maybe_reward and maybe_reward.get("type") == "reward"  # type: ignore
+        if maybe_reward and maybe_reward.get("type") == "reward":  # type: ignore
+            return True
+
+        # Agent-lightning 0.3+
+        if (
+            self.span.name == AGL_OPERATION
+            and self.span.attributes.get(LightningSpanAttributes.OPERATION_NAME.value) == AGL_REWARD
+        ):
+            return True
+
+        return False

    def find_llm_calls(
        self,
@@ -376,7 +430,9 @@ class TraceTree:
            is_llm_call = False
        if is_llm_call:
            # Check the response id
-            response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id")  # type: ignore
+            response_id = _attributes_get_multiple(
+                self.span.attributes, ["gen_ai.response.id", "agentlightning.operation.output.id"]
+            )
            if response_id is None and within_llm_call is True:
                is_llm_call = False
            if (
@@ -388,7 +444,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:
@@ -502,12 +559,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:
@@ -517,28 +574,129 @@ 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.debug(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.

        Subclass can override this method to add more fields to the triplet,
        such as chat messages and tool calls.
        """
-        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
+        prompt_token_ids = (
+            _attributes_get_ids_multiple(
+                span.attributes,
+                [
+                    "prompt_token_ids",
+                    "agentlightning.operation.output.prompt_token_ids",  # Weave tracer
+                ],
+            )
+            or []
+        )
+        response_token_ids = (
+            _attributes_get_ids_multiple(
+                span.attributes,
+                [
+                    "response_token_ids",
+                    "agentlightning.operation.output.response_token_ids.0",  # Weave tracer
+                    "agentlightning.operation.output.choices.0.token_ids",  # Weave tracer with newer vLLM
+                    "agentlightning.operation.output.choices.0.provider_specific_fields.token_ids",  # new vLLM + new OpenAI client SDK
+                ],
+            )
+            or []
+        )

+        response_id = _attributes_get_multiple(
+            span.attributes, ["gen_ai.response.id", "agentlightning.operation.output.id"]
+        )
+        request_metadata = _attributes_unflatten_multiple(
+            span.attributes, ["gen_ai.request", "agentlightning.operation.input"]
+        )
+        response_metadata = _attributes_unflatten_multiple(
+            span.attributes, ["gen_ai.response", "agentlightning.operation.output"]
+        )
+        # Special handling for Weave tracer: messages are handled separately
+        if isinstance(request_metadata, dict):
+            request_metadata.pop("messages", None)
+        if isinstance(response_metadata, dict):
+            response_metadata.pop("choices", None)
+            response_metadata.pop("prompt_token_ids", None)
+            response_metadata.pop("response_token_ids", None)
+
+        prompt_raw_content = _attributes_unflatten_multiple(
+            span.attributes, ["gen_ai.prompt", "agentlightning.operation.input.messages"]
+        )
+        completion_raw_content = _attributes_unflatten_multiple(
+            span.attributes, ["gen_ai.completion", "agentlightning.operation.output.choices"]
+        )
+        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}
+
+        # FIXME: logprob doesn't support Weave tracer yet.
        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(
@@ -670,7 +828,7 @@ class TracerTraceToTriplet(TraceToTripletBase):
        trace_tree.visualize(filename, interested_span_match=interested_span_match)
        return trace_tree

-    def adapt(self, source: Union[List[Span], List[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
+    def adapt(self, source: Union[Sequence[Span], Sequence[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
        """Convert tracer spans into [`Triplet`][agentlightning.Triplet] trajectories.

        Args:
@@ -776,31 +934,14 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):

    def _maybe_reward_value(self, span: Span) -> Optional[float]:
        """Parse reward from typical AgentOps payloads or explicit reward spans."""
-        attrs = span.attributes or {}
-
-        # AgentOps new/old keys
-        for k in ("agentops.task.output", "agentops.entity.output"):
-            v = attrs.get(k)
-            v = self._literal_eval_maybe(v)
-            if isinstance(v, dict) and cast(Dict[str, Any], v).get("type") == "reward":
-                rv = cast(Dict[str, Any], v).get("value", None)
-                if rv is None or isinstance(rv, (int, float)):
-                    return None if rv is None else float(rv)
-
-        # Explicit reward span
-        if span.name == SpanNames.REWARD.value:
-            rv = attrs.get("reward", None)
-            if rv is None or isinstance(rv, (int, float)):
-                return None if rv is None else float(rv)
-
-        return None
+        return get_reward_value(span)

    def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
        # Prefer OpenAI-like id if present, else proxy raw id.
        rid = attrs.get("gen_ai.response.id") or attrs.get("llm.hosted_vllm.id")
        return str(rid) if isinstance(rid, str) and rid else None

-    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
+    def adapt(self, source: Sequence[Span], /) -> List[Triplet]:  # type: ignore
        """Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.

        Args:
@@ -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.
@@ -143,7 +150,7 @@ class Baseline(FastAlgorithm):
        store = self.get_store()

        for index in train_indices + val_indices:
-            queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+            queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
            if len(queuing_rollouts) <= 1:
                # Only enqueue a new rollout when there is at most 1 rollout in the queue.
                sample = dataset[index]
@@ -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:
@@ -222,7 +231,7 @@ class Baseline(FastAlgorithm):
                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
                )
                while True:
-                    queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                    queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
                    if len(queuing_rollouts) <= self.max_queue_length:
                        # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
                        sample = concatenated_dataset[index]
@@ -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,32 @@ 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.
+
+    !!! note "Trajectory aggregation (experimental)"
+
+        Trajectory-level aggregation merges an entire multi-turn rollout into a single,
+        masked training sample so GPU time is spent once per trajectory rather than N times
+        per turn. Enable it via:
+
+        ```python
+        config["agentlightning"]["trace_aggregator"] = {
+            "level": "trajectory",
+            "trajectory_max_prompt_length": 4096,
+            "trajectory_max_response_length": 34384,
+        }
+        ```
+
+        Keep conversations structured (message lists rather than manual string
+        concatenation) so prefix matching can stitch traces. `trajectory_max_prompt_length`
+        should be set to the maximum length of the prompt for the first turn, and
+        `trajectory_max_response_length` should be set to the maximum cumulative
+        length of agent responses in the full trajectory.
+        Toggle `debug=True` plus `mismatch_log_dir` when you need to inspect
+        retokenization or chat-template mismatches. See
+        [this blog post](https://agent-lightning.github.io/posts/trajectory_level_aggregation/)
+        for more details.

    Examples:
        ```python
@@ -90,7 +122,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 +139,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 +158,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 +174,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 +188,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,17 +7,25 @@ from __future__ import annotations
 import argparse
 import asyncio
 import logging
-from typing import Iterable
+from typing import Iterable, List

-from agentlightning.logging import configure_logger
+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__)


 def main(argv: Iterable[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
    parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
    parser.add_argument(
        "--cors-origin",
@@ -25,16 +33,91 @@ def main(argv: Iterable[str] | None = None) -> int:
        action="append",
        help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the store.",
+    )
+    parser.add_argument(
+        "--tracker",
+        nargs="+",
+        choices=["prometheus", "console"],
+        help="Enable metrics tracking. Repeat for multiple trackers.",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help=(
+            "Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. "
+            "Only applicable for zero-copy stores such as MongoDB backend."
+        ),
+    )
+
+    parser.add_argument(
+        "--backend",
+        choices=["memory", "mongo"],
+        default="memory",
+        help="Backend to use for the store.",
+    )
+    parser.add_argument(
+        "--mongo-uri",
+        default="mongodb://localhost:27017/?replicaSet=rs0",
+        help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
+    )
+
    args = parser.parse_args(list(argv) if argv is not None else None)

-    configure_logger()
+    setup_logging(args.log_level)

-    store = InMemoryLightningStore()
+    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(
+            thread_safe=True,  # Using thread_safe store for server
+            tracker=tracker,
+        )
+    elif args.backend == "mongo":
+        from agentlightning.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    if args.n_workers > 1:
+        logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
+        launch_mode = "mp"
+    else:
+        logger.info("Running the server using `asyncio` launch mode.")
+        launch_mode = "asyncio"
    server = LightningStoreServer(
        store,
-        host="0.0.0.0",
+        host=args.host,
        port=args.port,
        cors_allow_origins=args.cors_origins,
+        launch_mode=launch_mode,
+        tracker=tracker,
+        n_workers=args.n_workers,
    )
    try:
        asyncio.run(server.run_forever())
@@ -1,25 +1,43 @@
 # Copyright (c) Microsoft. All rights reserved.

+"""Convenient helpers for creating spans / traces.
+
+All emitters operate in two modes, switchable via the `propagate` parameter.
+The emitters first [`SpanCreationRequest`][agentlightning.SpanCreationRequest] object, then:
+
+1. When `propagate` is True, this creation request will be propagated to the active tracer
+   and a [`Span`][agentlightning.Span] instance will be created (possibly deferred).
+2. When `propagate` is False, the creation request will be returned directly. Useful for cases
+   when you don't have a tracer but you want to create a creation request for later use.
+"""
+
+from .annotation import emit_annotation, operation
 from .exception import emit_exception
-from .message import emit_message
-from .object import emit_object
+from .message import emit_message, get_message_value
+from .object import emit_object, get_object_value
 from .reward import (
    emit_reward,
    find_final_reward,
    find_reward_spans,
    get_reward_value,
+    get_rewards_from_span,
    is_reward_span,
    reward,
 )

 __all__ = [
    "reward",
+    "operation",
    "emit_reward",
    "get_reward_value",
+    "get_rewards_from_span",
    "is_reward_span",
    "find_reward_spans",
    "find_final_reward",
    "emit_message",
    "emit_object",
    "emit_exception",
+    "emit_annotation",
+    "get_message_value",
+    "get_object_value",
 ]
@@ -0,0 +1,370 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helpers for emitting annotation/operation spans."""
+
+import asyncio
+import functools
+import inspect
+import logging
+from types import TracebackType
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
+from agentlightning.tracer.base import get_active_tracer
+from agentlightning.tracer.dummy import DummyTracer
+from agentlightning.types import SpanCoreFields, SpanRecordingContext, TraceStatus
+from agentlightning.utils.otel import check_attributes_sanity, flatten_attributes, sanitize_attributes
+
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanCoreFields:
+    """Emit a new annotation span.
+
+    This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward].
+
+    Annotation spans are used to annotate a specific event or a part of rollout.
+    See [semconv][agentlightning.semconv] for conventional annotation keys in Agent-lightning.
+
+    If annotations contain nested dicts, they will be flattened before emitting.
+    Complex objects will lead to emitting failures.
+
+    Args:
+        annotation: Dictionary containing annotation key-value pairs.
+            Representatives are rewards, tags, and metadata.
+        propagate: Whether to propagate the span to tracers automatically.
+    """
+    annotation_attributes = flatten_attributes(annotation, expand_leaf_lists=False)
+    check_attributes_sanity(annotation_attributes)
+    sanitized_attributes = sanitize_attributes(annotation_attributes)
+    logger.debug("Emitting annotation span with keys %s", sanitized_attributes.keys())
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit annotation span.")
+    else:
+        tracer = DummyTracer()
+
+    return tracer.create_span(
+        name=AGL_ANNOTATION,
+        attributes=sanitized_attributes,
+        status=TraceStatus(status_code="OK"),
+    )
+
+
+class OperationContext:
+    """Context manager and decorator for tracing operations.
+
+    This class manages a tracer-backed span for a logical unit of work. It can be
+    used either:
+
+    * As a decorator, in which case inputs and outputs are inferred
+      automatically from the wrapped function's signature.
+    * As a context manager, in which case inputs and outputs can be recorded
+      explicitly via [`set_input`][agentlightning.emitter.annotation.OperationContext.set_input]
+      and [`set_output`][agentlightning.emitter.annotation.OperationContext.set_output].
+
+    Attributes:
+        name: Human-readable span name.
+        initial_attributes: Attributes applied when the span is created.
+        tracer: Tracer implementation used to create spans.
+    """
+
+    def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
+        """Initialize a new operation context.
+
+        Args:
+            name: Human-readable name of the span.
+            attributes: Initial attributes attached to the span. Values are
+                JSON-serialized where necessary.
+            propagate: Whether the span should be sent to active exporters.
+        """
+        self.name = name
+        self.initial_attributes = flatten_attributes(attributes, expand_leaf_lists=False)
+        self.propagate = propagate
+        if propagate:
+            tracer = get_active_tracer()
+            if tracer is None:
+                raise RuntimeError("No active tracer found. Cannot trace operation spans.")
+            self.tracer = tracer
+        else:
+            self.tracer = DummyTracer()
+        self._ctx_manager: Optional[ContextManager[SpanRecordingContext]] = None
+        self._recording_context: Optional[SpanRecordingContext] = None
+        self._span: Optional[SpanCoreFields] = None
+
+    def __enter__(self) -> "OperationContext":
+        """Enter the context manager and start a new span.
+
+        Returns:
+            The current :class:`OperationContext` instance with an active span.
+        """
+        sanitized_attrs = sanitize_attributes(self.initial_attributes)
+        self._ctx_manager = self.tracer.operation_context(self.name, attributes=sanitized_attrs)
+        recording_context = self._ctx_manager.__enter__()
+        self._recording_context = recording_context
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Exit the context manager and finish the span."""
+        if self._ctx_manager:
+            self._ctx_manager.__exit__(exc_type, exc_val, exc_tb)
+        if self._recording_context:
+            self._span = self._recording_context.get_recorded_span()
+        self._ctx_manager = None
+        self._recording_context = None
+
+    def span(self) -> SpanCoreFields:
+        """Get the span that was created by this context manager."""
+        if self._span is None:
+            raise RuntimeError("Span is not ready yet.")
+        return self._span
+
+    def set_input(self, *args: Any, **kwargs: Any) -> None:
+        """Record input arguments on the current span.
+
+        Positional arguments are stored under the `input.args.<index>` attributes,
+        and keyword arguments are stored under `input.<name>` attributes.
+
+        This is intended for use inside a `with operation(...) as op` block.
+
+        Args:
+            *args: Positional arguments to record.
+            **kwargs: Keyword arguments to record.
+        """
+        if not self._recording_context:
+            raise RuntimeError("No recording context found. Cannot set input.")
+
+        prefix = LightningSpanAttributes.OPERATION_INPUT.value
+        attributes: Dict[str, Any] = {}
+        if args:
+            for idx, value in enumerate(args):
+                flattened = flatten_attributes({str(idx): value})
+                for nested_key, nested_value in flattened.items():
+                    attributes[f"{prefix}.args.{nested_key}"] = nested_value
+        if kwargs:
+            for key, value in kwargs.items():
+                flattened = flatten_attributes({key: value})
+                for nested_key, nested_value in flattened.items():
+                    attributes[f"{prefix}.{nested_key}"] = nested_value
+        if attributes:
+            self._recording_context.record_attributes(sanitize_attributes(attributes))
+
+    def set_output(self, output: Any) -> None:
+        """Record the output value on the current span.
+
+        This is intended for use inside a `with operation(...) as op` block.
+
+        Args:
+            output: The output value to record.
+        """
+        if not self._recording_context:
+            raise RuntimeError("No recording context found. Cannot set output.")
+
+        flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: output})
+        self._recording_context.record_attributes(sanitize_attributes(flattened))
+
+    def __call__(self, fn: _FnType) -> _FnType:
+        """Wrap a callable so its execution is traced in a span.
+
+        When used as a decorator, a new span is created for each call to
+        the wrapped function. The bound arguments are recorded as input
+        attributes, the return value is recorded as an output attribute,
+        and any exception is recorded and marks the span as an error.
+
+        Args:
+            fn: The function or coroutine function to wrap.
+
+        Returns:
+            The wrapped callable.
+        """
+        function_name = fn.__name__
+
+        sig = inspect.signature(fn)
+
+        sanitized_init_attrs = sanitize_attributes(
+            {LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
+        )
+
+        def _record_auto_inputs(
+            recording_ctx: SpanRecordingContext, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+        ) -> None:
+            """Bind arguments to signature and log them on the span."""
+            attributes: Dict[str, Any] = {}
+            try:
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                for name, value in bound.arguments.items():
+                    parameter = sig.parameters.get(name)
+                    if parameter and parameter.kind is inspect.Parameter.VAR_POSITIONAL:
+                        attr_prefix = f"{LightningSpanAttributes.OPERATION_INPUT.value}.{name}"
+                        for idx, item in enumerate(value):
+                            flattened = flatten_attributes({str(idx): item})
+                            for nested_key, nested_value in flattened.items():
+                                attributes[f"{attr_prefix}.{nested_key}"] = nested_value
+                    else:
+                        flattened = flatten_attributes({name: value})
+                        for nested_key, nested_value in flattened.items():
+                            attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
+            except Exception:
+                if args:
+                    for idx, value in enumerate(args):
+                        flattened = flatten_attributes({str(idx): value})
+                        for nested_key, nested_value in flattened.items():
+                            attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.args.{nested_key}"] = (
+                                nested_value
+                            )
+                if kwargs:
+                    flattened = flatten_attributes({"kwargs": kwargs})
+                    for nested_key, nested_value in flattened.items():
+                        attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
+            if attributes:
+                recording_ctx.record_attributes(sanitize_attributes(attributes))
+
+        def _record_auto_outputs(recording_ctx: SpanRecordingContext, result: Any) -> None:
+            """Record the output value on the span."""
+            flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: result})
+            recording_ctx.record_attributes(sanitize_attributes(flattened))
+
+        if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Async wrapper that traces the wrapped coroutine."""
+                with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
+                    _record_auto_inputs(recording_ctx, args, kwargs)
+                    result = await fn(*args, **kwargs)
+                    _record_auto_outputs(recording_ctx, result)
+                    return result
+
+            return cast(_FnType, async_wrapper)
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Sync wrapper that traces the wrapped callable."""
+                with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
+                    _record_auto_inputs(recording_ctx, args, kwargs)
+                    result = fn(*args, **kwargs)
+                    _record_auto_outputs(recording_ctx, result)
+                    return result
+
+            return cast(_FnType, sync_wrapper)
+
+
+@overload
+def operation(
+    fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
+) -> _FnType: ...
+
+
+@overload
+def operation(
+    *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
+) -> OperationContext: ...
+
+
+@overload
+def operation(fn: _FnType, *, name: Optional[str] = None, **additional_attributes: Any) -> _FnType: ...
+
+
+@overload
+def operation(*, name: Optional[str] = None, **additional_attributes: Any) -> OperationContext: ...
+
+
+@overload
+def operation(fn: _FnType, **additional_attributes: Any) -> _FnType: ...
+
+
+@overload
+def operation(**additional_attributes: Any) -> OperationContext: ...
+
+
+def operation(
+    fn: Optional[_FnType] = None,
+    *,
+    propagate: bool = True,
+    name: Optional[str] = None,
+    **additional_attributes: Any,
+) -> Union[_FnType, OperationContext]:
+    """Entry point for tracking operations.
+
+    This helper can be used either as a decorator or as a context manager.
+    The span name is fixed to [`AGL_OPERATION`][agentlightning.semconv.AGL_OPERATION];
+    custom span names are not supported. Any keyword arguments are recorded as span attributes.
+
+    Usage as a decorator:
+
+    ```python
+    @operation
+    def func(...):
+        ...
+
+    @operation(category="compute")
+    def func(...):
+        ...
+    ```
+
+    Usage as a context manager:
+
+    ```python
+    with operation(user_id=123) as op:
+        op.set_input(data=data)
+        # ... do work ...
+        op.set_output(result)
+    ```
+
+    Args:
+        fn: When used as `@operation`, this is the wrapped function.
+            When used as `operation(**attrs)`, this should be omitted (or
+            left as `None`) and only keyword attributes are provided.
+        propagate: Whether spans should use the active span processor. When False,
+            spans will stay local and not be exported.
+        name: Optional alias that populates
+            [`LightningSpanAttributes.OPERATION_NAME`][agentlightning.semconv.LightningSpanAttributes.OPERATION_NAME]
+            when `additional_attributes` does not already define it.
+        **additional_attributes: Additional span attributes to attach at
+            creation time.
+
+    Returns:
+        Either a wrapped callable (when used as a decorator) or an
+        [`OperationContext`][agentlightning.emitter.annotation.OperationContext]
+        (when used as a context manager factory).
+    """
+
+    if name is not None:
+        if LightningSpanAttributes.OPERATION_NAME.value in additional_attributes:
+            raise ValueError("Cannot specify both `name` and `additional_attributes.operation_name`.")
+        additional_attributes[LightningSpanAttributes.OPERATION_NAME.value] = name
+
+    # Case 1: Used as @operation (bare decorator or with attributes)
+    if callable(fn):
+        # Create context with fixed name, then immediately wrap the function
+        return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn)
+
+    # Case 2: Used as operation(...) / with operation(...)
+    # Custom span names are intentionally not supported; use AGL_OPERATION.
+    if fn is not None:
+        raise ValueError("Custom span names are intentionally not supported when used as a context manager.")
+    return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)
@@ -1,46 +1,54 @@
 # Copyright (c) Microsoft. All rights reserved.

 import logging
-import traceback
+from typing import Any, Dict, Optional

-from opentelemetry.semconv.attributes import exception_attributes
-
-from agentlightning.types import SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_EXCEPTION
+from agentlightning.tracer.base import get_active_tracer
+from agentlightning.tracer.dummy import DummyTracer
+from agentlightning.types import TraceStatus
+from agentlightning.utils.otel import flatten_attributes, format_exception_attributes, sanitize_attributes

 logger = logging.getLogger(__name__)


-def emit_exception(exception: BaseException) -> None:
+def emit_exception(
+    exception: BaseException, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True
+) -> None:
    """Record an exception with OpenTelemetry metadata.

+    Classic OpenTelemetry records exceptions in a dedicated logging service.
+    We simplify the model and use trace spans to record exceptions as well.
+
    Args:
        exception: Raised exception instance to serialize into telemetry attributes.
+        attributes: Additional attributes to attach to the exception span.
+        propagate: Whether to propagate the span to exporters automatically.

    !!! note
-        The helper validates its input. Non-exception values are ignored to prevent
-        noisy telemetry and indicate programming mistakes via the logger.
+
+        The helper validates its input. If a non-exception value is provided,
+        a TypeError is raised to indicate a programming mistake.
    """
    if not isinstance(exception, BaseException):  # type: ignore
-        logger.error(f"Expected an BaseException instance, got: {type(exception)}. Skip emit_exception.")
-        return
+        raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
+    span_attributes = format_exception_attributes(exception)

-    tracer = get_tracer()
-    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
-    attributes = {
-        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
-        exception_attributes.EXCEPTION_MESSAGE: str(exception),
-        exception_attributes.EXCEPTION_ESCAPED: True,
-    }
-    if stacktrace.strip():
-        attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))

-    span = tracer.start_span(
-        SpanNames.EXCEPTION.value,
-        attributes=attributes,
-    )
    logger.debug("Emitting exception span for %s", type(exception).__name__)
-    with span:
-        span.record_exception(exception)
-        # We don't set the status of the span here. They have other semantics.
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit exception span.")
+    else:
+        tracer = DummyTracer()
+    tracer.create_span(
+        AGL_EXCEPTION,
+        attributes=span_attributes,
+        # The exception span is successful by itself.
+        status=TraceStatus(status_code="OK"),
+    )
@@ -1,33 +1,61 @@
 # Copyright (c) Microsoft. All rights reserved.

 import logging
+from typing import Any, Dict, Optional

-from agentlightning.types import SpanAttributeNames, SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
+from agentlightning.tracer.base import get_active_tracer
+from agentlightning.tracer.dummy import DummyTracer
+from agentlightning.types import Attributes, SpanLike
+from agentlightning.utils.otel import flatten_attributes, sanitize_attributes

 logger = logging.getLogger(__name__)


-def emit_message(message: str) -> None:
+def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
    """Emit a textual message as an OpenTelemetry span.

+    Commonly used for sending debugging and logging messages.
+
    Args:
        message: Human readable message to attach as a span attribute.
+        attributes: Additional attributes to attach to the message span.
+        propagate: Whether to propagate the span to exporters automatically.

    !!! note
        OpenTelemetry distinguishes between logs and spans. Emitting the message as a
        span keeps all Agent Lightning telemetry in a single data store for analysis.
    """
    if not isinstance(message, str):  # type: ignore
-        logger.error(f"Message must be a string, got: {type(message)}. Skip emit_message.")
-        return
+        raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")

-    tracer = get_tracer()
-    span = tracer.start_span(
-        SpanNames.MESSAGE.value,
-        attributes={SpanAttributeNames.MESSAGE.value: message},
-    )
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit message span.")
+    else:
+        tracer = DummyTracer()
+    span_attributes: Attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
    logger.debug("Emitting message span with message: %s", message)
-    with span:
-        pass
+    tracer.create_span(
+        AGL_MESSAGE,
+        attributes=span_attributes,
+    )
+
+
+def get_message_value(span: SpanLike) -> Optional[str]:
+    """Extract the message string from a message span.
+
+    Args:
+        span: Span-like object to extract the message from.
+    """
+    span_attributes = span.attributes or {}
+    if LightningSpanAttributes.MESSAGE_BODY.value not in span_attributes:
+        return None
+    message = span_attributes[LightningSpanAttributes.MESSAGE_BODY.value]
+    if isinstance(message, str):
+        return message
+    raise TypeError(f"Message must be a string, got: {type(message)}.")
@@ -1,37 +1,117 @@
 # Copyright (c) Microsoft. All rights reserved.

+import base64
 import json
 import logging
-from typing import Any
+from typing import Any, Dict, Optional

-from agentlightning.types import SpanAttributeNames, SpanNames
-
-from .utils import get_tracer
+from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
+from agentlightning.tracer.base import get_active_tracer
+from agentlightning.tracer.dummy import DummyTracer
+from agentlightning.types import SpanCoreFields, SpanLike, TraceStatus
+from agentlightning.utils.otel import flatten_attributes, full_qualified_name, sanitize_attributes

 logger = logging.getLogger(__name__)


-def emit_object(object: Any) -> None:
+def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> SpanCoreFields:
    """Emit an object's serialized representation as an OpenTelemetry span.

    Args:
        object: Data structure to encode as JSON and attach to the span payload.
+        attributes: Additional attributes to attach to the object span.
+        propagate: Whether to propagate the span to exporters automatically.

    !!! note
-        The payload must be JSON serializable. Non-serializable objects are ignored and
-        an error is logged to aid debugging.
+        The payload must be JSON serializable. Non-serializable objects will lead to a RuntimeError.
    """
-    try:
-        serialized = json.dumps(object)
-    except (TypeError, ValueError):
-        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
-        return
+    span_attributes = encode_object(object)
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))

-    tracer = get_tracer()
-    span = tracer.start_span(
-        SpanNames.OBJECT.value,
-        attributes={SpanAttributeNames.OBJECT.value: serialized},
+    attr_length = 0
+    if LightningSpanAttributes.OBJECT_JSON.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_JSON.value])
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_LITERAL.value])
+    logger.debug("Emitting object span with payload size %d characters", attr_length)
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit object span.")
+    else:
+        # Do not actually propagate to any store or tracer backend.
+        tracer = DummyTracer()
+
+    return tracer.create_span(
+        name=AGL_OBJECT,
+        attributes=span_attributes,
+        status=TraceStatus(status_code="OK"),
    )
-    logger.debug("Emitting object span with payload size %d characters", len(serialized))
-    with span:
-        pass
+
+
+def encode_object(object: Any) -> Dict[str, Any]:
+    """Encode an object as span attributes.
+
+    Args:
+        object: Data structure to encode as JSON.
+    """
+    span_attributes = {}
+    if isinstance(object, (str, int, float, bool)):
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: type(object).__name__,
+            LightningSpanAttributes.OBJECT_LITERAL.value: str(object),
+        }
+    elif isinstance(object, bytes):
+        b64_encoded = base64.b64encode(object).decode("utf-8")
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
+            LightningSpanAttributes.OBJECT_LITERAL.value: b64_encoded,
+        }
+    else:
+        try:
+            serialized = json.dumps(object)
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError(f"Object must be JSON serializable, got: {type(object)}.") from exc
+
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: full_qualified_name(type(object)),  # type: ignore
+            LightningSpanAttributes.OBJECT_JSON.value: serialized,
+        }
+
+    return span_attributes
+
+
+def get_object_value(span: SpanLike) -> Any:
+    """Extract the object payload from an object span.
+
+    Args:
+        span: Span object produced by Agent Lightning emitters.
+    """
+    attributes = span.attributes or {}
+    if LightningSpanAttributes.OBJECT_JSON.value in attributes:
+        serialized = attributes[LightningSpanAttributes.OBJECT_JSON.value]
+        try:
+            return json.loads(serialized)  # type: ignore
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError("Failed to deserialize object JSON from span.") from exc
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in attributes:
+        literal = attributes[LightningSpanAttributes.OBJECT_LITERAL.value]
+        obj_type = attributes.get(LightningSpanAttributes.OBJECT_TYPE.value, "str")
+        if obj_type == "str":
+            return literal
+        elif obj_type == "int":
+            # Let it raise errors if there are any
+            return int(literal)  # type: ignore
+        elif obj_type == "float":
+            return float(literal)  # type: ignore
+        elif obj_type == "bool":
+            return literal.lower() == "true"  # type: ignore
+        elif obj_type == "bytes":
+            return base64.b64decode(literal.encode("utf-8"))  # type: ignore
+        else:
+            raise RuntimeError(f"Unsupported object type for literal deserialization: {obj_type}")
+    else:
+        return None
@@ -20,13 +20,13 @@ from typing import (
    cast,
 )

-import agentops
-from agentops.sdk.decorators import operation
-from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import TypeAdapter

-from agentlightning.types import SpanLike, SpanNames
+from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
+from agentlightning.types import SpanCoreFields, SpanLike
+from agentlightning.utils.otel import filter_and_unflatten_attributes

-from .utils import get_tracer
+from .annotation import emit_annotation

 logger = logging.getLogger(__name__)

@@ -34,26 +34,36 @@ __all__ = [
    "reward",
    "emit_reward",
    "get_reward_value",
+    "get_rewards_from_span",
    "is_reward_span",
    "find_reward_spans",
    "find_final_reward",
 ]


-class RewardSpanData(TypedDict):
+class RewardDimension(TypedDict):
+    """Type representing a single dimension in a multi-dimensional reward."""
+
+    name: str
+    value: float
+
+
+class _RewardSpanData(TypedDict):
    type: Literal["reward"]
    value: Optional[float]


-FnType = TypeVar("FnType", bound=Callable[..., Any])
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])


 def _agentops_initialized() -> bool:
    """Return `True` when the AgentOps client has been configured."""
+    import agentops
+
    return agentops.get_client().initialized


-def reward(fn: FnType) -> FnType:
+def reward(fn: _FnType) -> _FnType:
    """Decorate a reward function so its outputs are tracked as spans.

    The decorator integrates with AgentOps when it is available and falls back to
@@ -70,7 +80,9 @@ def reward(fn: FnType) -> FnType:
        Wrapped callable that preserves the original signature.
    """

-    def wrap_result(result: Optional[float]) -> RewardSpanData:
+    from agentops.sdk.decorators import operation
+
+    def wrap_result(result: Optional[float]) -> _RewardSpanData:
        """Normalize the reward value into the span payload format."""
        if result is None:
            return {"type": "reward", "value": None}
@@ -94,7 +106,7 @@ def reward(fn: FnType) -> FnType:
            result: Optional[float] = None

            @operation
-            async def agentops_reward_operation() -> RewardSpanData:
+            async def agentops_reward_operation() -> _RewardSpanData:
                # The reward function we are interested in tracing
                # It takes zero inputs and return a formatted dict
                nonlocal result
@@ -118,7 +130,7 @@ def reward(fn: FnType) -> FnType:
            result: Optional[float] = None

            @operation
-            def agentops_reward_operation() -> RewardSpanData:
+            def agentops_reward_operation() -> _RewardSpanData:
                nonlocal result
                result = fn(*args, **kwargs)
                return wrap_result(result)
@@ -129,35 +141,69 @@ def reward(fn: FnType) -> FnType:
        return wrapper  # type: ignore


-def emit_reward(reward: float) -> ReadableSpan:
+def emit_reward(
+    reward: float | Dict[str, Any],
+    *,
+    primary_key: str | None = None,
+    attributes: Dict[str, Any] | None = None,
+    propagate: bool = True,
+) -> SpanCoreFields:
    """Emit a reward value as an OpenTelemetry span.

+    Examples:
+        Emit a single-dimensional reward:
+        >>> emit_reward(1.0)
+
+        Emit multi-dimensional rewards:
+        >>> emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
+
+        Emit a reward with additional attributes (for example linking to another response span):
+        >>> from agentlightning.utils.otel import make_link_attributes
+        >>> emit_reward(0.5, attributes=make_link_attributes({"gen_ai.response.id": "response-123"}))
+
+        Or adding tags onto the reward span:
+        >>> from agentlightning.utils.otel import make_tag_attributes
+        >>> emit_reward(0.7, attributes=make_tag_attributes(["fast", "reliable"]))
+
    Args:
        reward: Numeric reward to record. Integers and booleans are converted to
            floating point numbers for consistency.
+            Use a dictionary to represent a multi-dimensional reward.
+        attributes: Other optional span attributes.
+        propagate: Whether to propagate the span to exporters automatically.

    Returns:
-        Readable span capturing the recorded reward.
-
-    Raises:
-        ValueError: If the provided reward cannot be interpreted as a float or the
-            resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
+        Span core fields capturing the recorded reward.
    """
    logger.debug(f"Emitting reward: {reward}")
-    if isinstance(reward, (int, bool)):
-        reward = float(reward)
-    if not isinstance(reward, float):
-        raise ValueError(f"Reward must be a number, got: {type(reward)}")
+    reward_dimensions: List[RewardDimension] = []
+    if isinstance(reward, dict):
+        reward_dict: Dict[str, float] = {}
+        for k, v in reward.items():
+            if isinstance(v, (int, bool)):
+                reward_dict[k] = float(v)
+            elif isinstance(v, float):
+                reward_dict[k] = v
+            else:
+                raise ValueError(f"Reward value must be a number, got: {type(v)} for key {k}")
+        if primary_key is None:
+            raise ValueError("When emitting a multi-dimensional reward as a dict, primary_key must be provided.")
+        if primary_key not in reward_dict:
+            raise ValueError(f"Primary key '{primary_key}' not found in reward dict keys: {list(reward_dict.keys())}")
+        reward_dimensions.append(RewardDimension(name=primary_key, value=reward_dict[primary_key]))
+        for k, v in reward_dict.items():
+            if k != primary_key:
+                reward_dimensions.append(RewardDimension(name=k, value=v))
+    else:
+        if isinstance(reward, (int, bool)):
+            reward = float(reward)
+        elif not isinstance(reward, float):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise TypeError(f"Reward must be a number, got: {type(reward)}")
+        reward_dimensions.append(RewardDimension(name="primary", value=reward))

-    # TODO: This should use the tracer from current context by tracer
-    tracer = get_tracer()
-    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
-    # Do nothing; it's just a number
-    with span:
-        pass
-    if not isinstance(span, ReadableSpan):
-        raise ValueError(f"Span is not a ReadableSpan: {span}")
-    return span
+    return emit_annotation(
+        {LightningSpanAttributes.REWARD.value: reward_dimensions, **(attributes or {})}, propagate=propagate
+    )


 def get_reward_value(span: SpanLike) -> Optional[float]:
@@ -167,8 +213,14 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
        span: Span object produced by AgentOps or Agent Lightning emitters.

    Returns:
-        The reward encoded in the span or `None` when the span does not represent a reward.
+        The primary reward encoded in the span or `None` when the span does not represent a reward.
    """
+    # v0.3+ emit reward format
+    reward_list = get_rewards_from_span(span)
+    if reward_list:
+        # Reward list is ordered and the first element is the primary reward
+        return reward_list[0].value
+
    for key in [
        "agentops.task.output",  # newer versions of agentops
        "agentops.entity.output",
@@ -191,19 +243,45 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
                return None
            if not isinstance(reward_value, float):
                logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+            logger.warning(
+                f"Extracted reward {reward_value} from AgentOps. This format is deprecated, please migrate to using `emit_reward`."
+            )
            return cast(float, reward_value)

-    # Latest emit reward format
-    if span.name == SpanNames.REWARD.value and span.attributes:
+    # v0.2 emit reward format
+    if span.name == AGL_ANNOTATION and span.attributes:
        reward_value = span.attributes.get("reward", None)
        if reward_value is None:
            return None
        if not isinstance(reward_value, float):
            logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+        logger.warning(
+            f"Extracted reward {reward_value} from a legacy version of reward span. You might have inconsistent agent-lightning versions."
+        )
        return cast(float, reward_value)
+
    return None


+def get_rewards_from_span(span: SpanLike) -> List[RewardPydanticModel]:
+    """Extract the reward as a list from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or Agent Lightning emitters.
+
+    Returns:
+        A list of reward dimensions encoded in the span or an empty list when the span does not represent a reward.
+    """
+    if span.attributes and any(key.startswith(LightningSpanAttributes.REWARD.value) for key in span.attributes):
+        reward_attr = filter_and_unflatten_attributes(
+            cast(Any, span.attributes or {}), LightningSpanAttributes.REWARD.value
+        )
+        recovered_rewards = TypeAdapter(List[RewardPydanticModel]).validate_python(reward_attr)
+        return recovered_rewards
+    else:
+        return []
+
+
 def is_reward_span(span: SpanLike) -> bool:
    """Return ``True`` when the provided span encodes a reward value."""
    maybe_reward = get_reward_value(span)
@@ -1,22 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Utilities shared across emitter implementations."""
-
-import opentelemetry.trace as trace_api
-from opentelemetry.trace import get_tracer_provider
-
-
-def get_tracer() -> trace_api.Tracer:
-    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
-
-    Returns:
-        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
-
-    Raises:
-        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
-    """
-    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
-        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
-
-    tracer_provider = get_tracer_provider()
-    return tracer_provider.get_tracer("agentlightning")
@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Environment variable managements."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+from typing import overload
+
+__all__ = [
+    "LightningEnvVar",
+    "resolve_bool_env_var",
+    "resolve_int_env_var",
+    "resolve_str_env_var",
+]
+
+
+class LightningEnvVar(Enum):
+    """Environment variables for Agent Lightning."""
+
+    AGL_EMITTER_DEBUG = "AGL_EMITTER_DEBUG"
+    """Enable debug logging for the emitter."""
+
+    AGL_MANAGED_STORE = "AGL_MANAGED_STORE"
+    """If yes, the [`ExecutionStrategy`][agentlightning.ExecutionStrategy]
+    constructs LightningStore wrappers automatically. When `False` the provided
+    `store` is passed directly to the bundles, allowing callers to manage
+    store wrappers manually."""
+
+    AGL_CURRENT_ROLE = "AGL_CURRENT_ROLE"
+    """Which side(s) to run in this process. Used in
+    [`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]."""
+
+    AGL_SERVER_HOST = "AGL_SERVER_HOST"
+    """Interface the [`LightningStoreServer`][agentlightning.LightningStoreServer]
+    binds to when running the algorithm bundle locally."""
+
+    AGL_SERVER_PORT = "AGL_SERVER_PORT"
+    """Port the [`LightningStoreServer`][agentlightning.LightningStoreServer] listens to."""
+
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+_FALSY_VALUES = {"0", "false", "no", "off"}
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, override: bool, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, *, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None: ...
+
+
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None:
+    """Resolve a boolean environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    normalized = env_value.strip().lower()
+    if normalized in _TRUTHY_VALUES:
+        return True
+    if normalized in _FALSY_VALUES:
+        return False
+
+    raise ValueError(f"{env_var.value} must be one of {_TRUTHY_VALUES} or {_FALSY_VALUES}")
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, override: int, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, *, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None: ...
+
+
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None:
+    """Resolve an integer environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    try:
+        return int(env_value)
+    except ValueError:
+        raise ValueError(f"{env_var.value} must be an integer")
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, override: str, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, *, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None: ...
+
+
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None:
+    """Resolve a string environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    return env_value
@@ -3,7 +3,6 @@
 from __future__ import annotations

 import logging
-import os
 from typing import Protocol

 from agentlightning.store.base import LightningStore
@@ -13,47 +12,6 @@ from .events import ExecutionEvent
 logger = logging.getLogger(__name__)


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

@@ -9,10 +9,11 @@ import time
 from multiprocessing.context import BaseContext
 from typing import Callable, Iterable, Literal, cast

+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_int_env_var, resolve_str_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer

-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
 from .events import ExecutionEvent, MultiprocessingEvent

 logger = logging.getLogger(__name__)
@@ -67,10 +68,11 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        server_host: str | None = None,
        server_port: int | None = None,
        n_runners: int = 1,
-        graceful_timeout: float = 5.0,
-        terminate_timeout: float = 5.0,
+        graceful_timeout: float = 10.0,
+        terminate_timeout: float = 10.0,
        main_process: Literal["algorithm", "runner"] = "algorithm",
        managed_store: bool | None = None,
+        allowed_exit_codes: Iterable[int] = (0, -15),
    ) -> None:
        """Configure the strategy.

@@ -94,45 +96,33 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                LightningStore client/server wrappers automatically. When
                `False` the provided `store` is passed directly to the
                bundles, allowing callers to manage store wrappers manually.
+            allowed_exit_codes: Allowed exit codes for subprocesses.
+                By default, runner can exit gracefully with code 0 or terminated
+                by SIGTERM (-15).
        """
-        if role is None:
-            role_env = os.getenv("AGL_CURRENT_ROLE")
-            if role_env is None:
-                # Use both if not specified via env var or argument
-                role = "both"
-            elif role_env not in ("algorithm", "runner", "both"):
-                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
-            else:
-                role = role_env
-
-        if server_host is None:
-            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
-
-        if server_port is None:
-            server_port_env = os.getenv("AGL_SERVER_PORT")
-            if server_port_env is None:
-                server_port = 4747
-            else:
-                try:
-                    server_port = int(server_port_env)
-                except ValueError as exc:
-                    raise ValueError("AGL_SERVER_PORT must be an integer") from exc
-
-        self.role = role
+        resolved_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, override=role, fallback="both")
+        if resolved_role not in ("algorithm", "runner", "both"):
+            raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
+        self.role: Literal["algorithm", "runner", "both"] = resolved_role
        self.n_runners = n_runners
-        self.server_host = server_host
-        self.server_port = server_port
+        self.server_host = resolve_str_env_var(
+            LightningEnvVar.AGL_SERVER_HOST, override=server_host, fallback="localhost"
+        )
+        self.server_port = resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, override=server_port, fallback=4747)
        self.graceful_timeout = graceful_timeout
        self.terminate_timeout = terminate_timeout
        if main_process not in ("algorithm", "runner"):
            raise ValueError("main_process must be 'algorithm' or 'runner'")
        if main_process == "runner":
-            if role != "both":
+            if self.role != "both":
                raise ValueError("main_process='runner' is only supported when role='both'")
            if n_runners != 1:
                raise ValueError("main_process='runner' requires n_runners to be 1")
        self.main_process = main_process
-        self.managed_store = resolve_managed_store_flag(managed_store)
+        self.managed_store = resolve_bool_env_var(
+            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
+        )
+        self.allowed_exit_codes = tuple(allowed_exit_codes)

    async def _execute_algorithm(
        self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
@@ -153,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()
@@ -189,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()
@@ -220,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(
@@ -244,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,
@@ -338,10 +348,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):

    def _check_process_exitcodes(self, processes: Iterable[multiprocessing.Process]) -> None:
        """Raise an error if any managed process exited with a non-zero status."""
-        failed = [p for p in processes if p.exitcode not in (0, None)]
+        failed = [p for p in processes if p.exitcode not in self.allowed_exit_codes + (None,)]
        if failed:
            formatted = ", ".join(f"{p.name or p.pid} (exitcode={p.exitcode})" for p in failed)
-            raise RuntimeError(f"Subprocesses failed: {formatted}")
+            raise RuntimeError(f"Subprocesses failed with unexpected exit codes: {formatted}")

    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
        logger.info(
@@ -7,10 +7,11 @@ from contextlib import suppress
 from queue import SimpleQueue
 from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple

+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.threading import LightningStoreThreaded

-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
 from .events import ExecutionEvent, ThreadingEvent

 logger = logging.getLogger(__name__)
@@ -62,7 +63,9 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
        self.join_timeout = join_timeout
        self.graceful_delay = graceful_delay
        self.poll_interval = poll_interval
-        self.managed_store = resolve_managed_store_flag(managed_store)
+        self.managed_store = resolve_bool_env_var(
+            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
+        )

    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
        """Run `coro` until it finishes or a cooperative stop is requested.
@@ -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
@@ -13,7 +13,8 @@ from agentops.sdk.exporters import AuthenticatedOTLPExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.metrics.export import MetricExportResult
-from opentelemetry.sdk.trace.export import SpanExportResult
+
+from agentlightning.utils.otlp import LightningStoreOTLPExporter

 logger = logging.getLogger(__name__)

@@ -32,25 +33,27 @@ def enable_agentops_service(enabled: bool = True) -> None:
    """
    Enable or disable communication with the AgentOps service.

-    False (default): AgentOps exporters and clients will run in local mode
-    and will not attempt to communicate with the remote AgentOps service.
-    True: all exporters and clients will operate in normal mode and send data
-    to the AgentOps service as expected.
+    By default, AgentOps exporters and clients will run in local mode
+    and will NOT attempt to communicate with the remote AgentOps service.
+
+    Args:
+        enabled: If True, enable all AgentOps exporters and clients.
+            All exporters and clients will operate in normal mode and send data
+            to the [AgentOps service](https://www.agentops.ai).
    """
    global _agentops_service_enabled
    _agentops_service_enabled = enabled
-    logger.info(f"Switch set to {enabled} for exporters and clients.")
+    logger.info(f"AgentOps service enabled is set to {enabled}.")


 def _patch_exporters():
    import agentops.client.api
    import agentops.sdk.core
-    import opentelemetry.exporter.otlp.proto.http.metric_exporter
-    import opentelemetry.exporter.otlp.proto.http.trace_exporter

    agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter  # type: ignore
-    opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = BypassableOTLPMetricExporter
-    opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = BypassableOTLPSpanExporter
+    agentops.sdk.core.OTLPMetricExporter = BypassableOTLPMetricExporter
+    if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
+        agentops.sdk.core.OTLPSpanExporter = BypassableOTLPSpanExporter  # type: ignore
    agentops.client.api.V3Client = BypassableV3Client
    agentops.client.api.V4Client = BypassableV4Client

@@ -58,12 +61,11 @@ def _patch_exporters():
 def _unpatch_exporters():
    import agentops.client.api
    import agentops.sdk.core
-    import opentelemetry.exporter.otlp.proto.http.metric_exporter
-    import opentelemetry.exporter.otlp.proto.http.trace_exporter

    agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter  # type: ignore
-    opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = OTLPMetricExporter
-    opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = OTLPSpanExporter
+    agentops.sdk.core.OTLPMetricExporter = OTLPMetricExporter
+    if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
+        agentops.sdk.core.OTLPSpanExporter = OTLPSpanExporter  # type: ignore
    agentops.client.api.V3Client = V3Client
    agentops.client.api.V4Client = V4Client

@@ -243,18 +245,15 @@ def uninstrument_agentops():
        pass


-class BypassableAuthenticatedOTLPExporter(AuthenticatedOTLPExporter):
+class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, AuthenticatedOTLPExporter):
    """
    AuthenticatedOTLPExporter with switchable service control.
+
    When `_agentops_service_enabled` is False, skip export and return success.
    """

-    def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
-        if _agentops_service_enabled:
-            return super().export(*args, **kwargs)
-        else:
-            logger.debug("SwitchableAuthenticatedOTLPExporter is switched off, skipping export.")
-            return SpanExportResult.SUCCESS
+    def should_bypass(self) -> bool:
+        return not _agentops_service_enabled


 class BypassableOTLPMetricExporter(OTLPMetricExporter):
@@ -271,18 +270,16 @@ class BypassableOTLPMetricExporter(OTLPMetricExporter):
            return MetricExportResult.SUCCESS


-class BypassableOTLPSpanExporter(OTLPSpanExporter):
+class BypassableOTLPSpanExporter(LightningStoreOTLPExporter):
    """
    OTLPSpanExporter with switchable service control.
    When `_agentops_service_enabled` is False, skip export and return success.
+
+    This is used instead of BypassableAuthenticatedOTLPExporter on legacy AgentOps versions.
    """

-    def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
-        if _agentops_service_enabled:
-            return super().export(*args, **kwargs)
-        else:
-            logger.debug("SwitchableOTLPSpanExporter is switched off, skipping export.")
-            return SpanExportResult.SUCCESS
+    def should_bypass(self) -> bool:
+        return not _agentops_service_enabled


 class BypassableV3Client(V3Client):
@@ -0,0 +1,500 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import threading
+import warnings
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Iterator, List
+
+import weave.trace.weave_init
+from pydantic import validate_call
+from weave.trace_server import trace_server_interface as tsi
+from weave.trace_server.ids import generate_id
+from weave.trace_server_bindings.client_interface import TraceServerClientInterface
+from weave.trace_server_bindings.models import ServerInfoRes
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "instrument_weave",
+    "uninstrument_weave",
+    "InMemoryWeaveTraceServer",
+]
+
+
+class InMemoryWeaveTraceServer(TraceServerClientInterface):
+    """A minimal in-memory implementation of the TraceServerInterface.
+
+    It stores calls and objects in local dictionaries and returns valid Pydantic
+    responses to satisfy the Weave client and FullTraceServerInterface protocol.
+    """
+
+    def __init__(self):
+        # Minimal storage to allow basic querying in tests
+        self.calls: Dict[str, tsi.CallSchema] = {}
+        self.partial_calls: Dict[str, Dict[str, Any]] = {}
+        self.objs: Dict[str, Any] = {}
+        self.files: Dict[str, bytes] = {}
+        self.feedback: List[tsi.FeedbackCreateReq] = []
+
+        self._call_threading_lock = threading.Lock()
+
+    @classmethod
+    def from_env(cls, *args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
+        return cls()
+
+    def server_info(self) -> ServerInfoRes:
+        return ServerInfoRes(min_required_weave_python_version="0.52.22")
+
+    def ensure_project_exists(self, entity: str, project: str) -> tsi.EnsureProjectExistsRes:
+        return tsi.EnsureProjectExistsRes(project_name=project)
+
+    # --- Call API ---
+
+    @validate_call
+    def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
+        # NOTE: It's not necessary that call_end must be called after call_start.
+        request_content = req.start.model_dump(exclude_none=True)
+
+        # If id needs to be generated here, it's very likely we won't be able to find the call later.
+        # This is just to make the type checker happy.
+        call_id = request_content.get("id") or generate_id()
+        trace_id = request_content.get("trace_id") or generate_id()
+        request_content["id"] = call_id
+        request_content["trace_id"] = trace_id
+
+        with self._call_threading_lock:
+            if call_id in self.partial_calls:
+                # call_end has already been called for this call.
+                kwargs = {**request_content, **self.partial_calls[call_id]}
+                self.calls[call_id] = tsi.CallSchema(**kwargs)
+                del self.partial_calls[call_id]
+            else:
+                self.partial_calls[call_id] = request_content
+
+        return tsi.CallStartRes(id=call_id, trace_id=trace_id)
+
+    @validate_call
+    def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
+        request_content = req.end.model_dump(exclude_none=True)
+        call_id = req.end.id
+
+        with self._call_threading_lock:
+            if call_id in self.partial_calls:
+                # End request always override the start request content.
+                kwargs = {**self.partial_calls[call_id], **request_content}
+                self.calls[call_id] = tsi.CallSchema(**kwargs)
+                del self.partial_calls[call_id]
+            else:
+                self.partial_calls[call_id] = request_content
+        return tsi.CallEndRes()
+
+    @validate_call
+    def call_start_batch(self, req: tsi.CallCreateBatchReq) -> tsi.CallCreateBatchRes:
+        for item in req.batch:
+            if isinstance(item, tsi.CallStartReq):
+                self.call_start(item)
+            elif isinstance(item, tsi.CallEndReq):
+                self.call_end(item)
+        return tsi.CallCreateBatchRes(res=[])
+
+    @validate_call
+    def call_read(self, req: tsi.CallReadReq) -> tsi.CallReadRes:
+        call_data = self.calls.get(req.id)
+        return tsi.CallReadRes(call=call_data)
+
+    @validate_call
+    def calls_query(self, req: tsi.CallsQueryReq) -> tsi.CallsQueryRes:
+        return tsi.CallsQueryRes(calls=list(self.calls_query_stream(req)))
+
+    @validate_call
+    def calls_query_stream(self, req: tsi.CallsQueryReq) -> Iterator[tsi.CallSchema]:
+        yield from self.calls.values()
+
+    @validate_call
+    def calls_delete(self, req: tsi.CallsDeleteReq) -> tsi.CallsDeleteRes:
+        num_deleted = 0
+        for call_id in req.call_ids:
+            if call_id in self.calls:
+                del self.calls[call_id]
+                num_deleted += 1
+        return tsi.CallsDeleteRes(num_deleted=num_deleted)
+
+    @validate_call
+    def call_update(self, req: tsi.CallUpdateReq) -> tsi.CallUpdateRes:
+        return tsi.CallUpdateRes()
+
+    @validate_call
+    def calls_query_stats(self, req: tsi.CallsQueryStatsReq) -> tsi.CallsQueryStatsRes:
+        return tsi.CallsQueryStatsRes(count=len(self.calls))
+
+    # --- Cost API ---
+
+    @validate_call
+    def cost_create(self, req: tsi.CostCreateReq) -> tsi.CostCreateRes:
+        return tsi.CostCreateRes(ids=[(generate_id(), generate_id()) for _ in req.costs])
+
+    @validate_call
+    def cost_query(self, req: tsi.CostQueryReq) -> tsi.CostQueryRes:
+        return tsi.CostQueryRes(results=[])
+
+    @validate_call
+    def cost_purge(self, req: tsi.CostPurgeReq) -> tsi.CostPurgeRes:
+        return tsi.CostPurgeRes()
+
+    # --- Object API (Legacy V1) ---
+
+    @validate_call
+    def obj_create(self, req: tsi.ObjCreateReq) -> tsi.ObjCreateRes:
+        digest = generate_id()
+        self.objs[digest] = req.obj
+        return tsi.ObjCreateRes(digest=digest)
+
+    @validate_call
+    def obj_read(self, req: tsi.ObjReadReq) -> tsi.ObjReadRes:
+        return tsi.ObjReadRes(obj=self.objs.get(req.digest, {}))
+
+    @validate_call
+    def objs_query(self, req: tsi.ObjQueryReq) -> tsi.ObjQueryRes:
+        return tsi.ObjQueryRes(objs=[])
+
+    @validate_call
+    def obj_delete(self, req: tsi.ObjDeleteReq) -> tsi.ObjDeleteRes:
+        return tsi.ObjDeleteRes(num_deleted=0)
+
+    # --- Table API ---
+
+    @validate_call
+    def table_create(self, req: tsi.TableCreateReq) -> tsi.TableCreateRes:
+        return tsi.TableCreateRes(digest=generate_id(), row_digests=[])
+
+    @validate_call
+    def table_create_from_digests(self, req: tsi.TableCreateFromDigestsReq) -> tsi.TableCreateFromDigestsRes:
+        return tsi.TableCreateFromDigestsRes(digest=generate_id())
+
+    @validate_call
+    def table_update(self, req: tsi.TableUpdateReq) -> tsi.TableUpdateRes:
+        return tsi.TableUpdateRes(digest=generate_id(), updated_row_digests=[])
+
+    @validate_call
+    def table_query(self, req: tsi.TableQueryReq) -> tsi.TableQueryRes:
+        return tsi.TableQueryRes(rows=[])
+
+    @validate_call
+    def table_query_stream(self, req: tsi.TableQueryReq) -> Iterator[tsi.TableRowSchema]:
+        yield from []
+
+    @validate_call
+    def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
+        return tsi.TableQueryStatsRes(count=0)
+
+    @validate_call
+    def table_query_stats_batch(self, req: tsi.TableQueryStatsBatchReq) -> tsi.TableQueryStatsBatchRes:
+        return tsi.TableQueryStatsBatchRes(tables=[])
+
+    # --- Ref API ---
+
+    @validate_call
+    def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
+        return tsi.RefsReadBatchRes(vals=[])
+
+    # --- File API ---
+
+    def file_create(self, req: tsi.FileCreateReq) -> tsi.FileCreateRes:
+        self.files[req.name] = req.content
+        return tsi.FileCreateRes(digest=generate_id())
+
+    def file_content_read(self, req: tsi.FileContentReadReq) -> tsi.FileContentReadRes:
+        return tsi.FileContentReadRes(content=self.files.get(req.digest, b"dummy_content"))
+
+    def files_stats(self, req: tsi.FilesStatsReq) -> tsi.FilesStatsRes:
+        total_size = sum(len(c) for c in self.files.values())
+        return tsi.FilesStatsRes(total_size_bytes=total_size)
+
+    # --- Feedback API ---
+
+    @validate_call
+    def feedback_create(self, req: tsi.FeedbackCreateReq) -> tsi.FeedbackCreateRes:
+        req.id = req.id or generate_id()
+        self.feedback.append(req)
+        return tsi.FeedbackCreateRes(
+            id=req.id,
+            created_at=datetime.now(timezone.utc),
+            wb_user_id="dummy_user",
+            payload=req.payload,
+        )
+
+    def feedback_create_batch(self, req: tsi.FeedbackCreateBatchReq) -> tsi.FeedbackCreateBatchRes:
+        results: List[tsi.FeedbackCreateRes] = []
+        for item in req.batch:
+            res = self.feedback_create(item)
+            results.append(res)
+        return tsi.FeedbackCreateBatchRes(res=results)
+
+    @validate_call
+    def feedback_query(self, req: tsi.FeedbackQueryReq) -> tsi.FeedbackQueryRes:
+        return tsi.FeedbackQueryRes(result=[])
+
+    @validate_call
+    def feedback_purge(self, req: tsi.FeedbackPurgeReq) -> tsi.FeedbackPurgeRes:
+        self.feedback.clear()
+        return tsi.FeedbackPurgeRes()
+
+    @validate_call
+    def feedback_replace(self, req: tsi.FeedbackReplaceReq) -> tsi.FeedbackReplaceRes:
+        return tsi.FeedbackReplaceRes(
+            id=req.id or generate_id(),
+            created_at=datetime.now(timezone.utc),
+            wb_user_id="dummy",
+            payload={},
+        )
+
+    # --- Action API ---
+
+    @validate_call
+    def actions_execute_batch(self, req: tsi.ActionsExecuteBatchReq) -> tsi.ActionsExecuteBatchRes:
+        return tsi.ActionsExecuteBatchRes()
+
+    # --- Execute LLM API ---
+
+    @validate_call
+    def completions_create(self, req: tsi.CompletionsCreateReq) -> tsi.CompletionsCreateRes:
+        return tsi.CompletionsCreateRes(response={"choices": [{"text": "dummy completion"}]})
+
+    @validate_call
+    def completions_create_stream(self, req: tsi.CompletionsCreateReq) -> Iterator[dict[str, Any]]:
+        yield {"choices": [{"text": "dummy "}]}
+        yield {"choices": [{"text": "stream"}]}
+
+    # --- Execute Image Generation API ---
+
+    @validate_call
+    def image_create(self, req: tsi.ImageGenerationCreateReq) -> tsi.ImageGenerationCreateRes:
+        return tsi.ImageGenerationCreateRes(response={})
+
+    # --- Project Statistics API ---
+
+    @validate_call
+    def project_stats(self, req: tsi.ProjectStatsReq) -> tsi.ProjectStatsRes:
+        return tsi.ProjectStatsRes(
+            trace_storage_size_bytes=0,
+            objects_storage_size_bytes=0,
+            tables_storage_size_bytes=0,
+            files_storage_size_bytes=0,
+        )
+
+    # --- Thread API ---
+
+    @validate_call
+    def threads_query_stream(self, req: tsi.ThreadsQueryReq) -> Iterator[tsi.ThreadSchema]:
+        yield from []
+
+    # --- Evaluation API (V1) ---
+
+    @validate_call
+    def evaluate_model(self, req: tsi.EvaluateModelReq) -> tsi.EvaluateModelRes:
+        return tsi.EvaluateModelRes(call_id=generate_id())
+
+    @validate_call
+    def evaluation_status(self, req: tsi.EvaluationStatusReq) -> tsi.EvaluationStatusRes:
+        return tsi.EvaluationStatusRes(status=tsi.EvaluationStatusNotFound())
+
+    # --- OTEL API ---
+
+    def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
+        return tsi.OtelExportRes()
+
+    # ==========================================
+    # Object Interface (V2 APIs)
+    # ==========================================
+
+    # --- Ops ---
+    def op_create(self, req: tsi.OpCreateReq) -> tsi.OpCreateRes:
+        return tsi.OpCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
+
+    def op_read(self, req: tsi.OpReadReq) -> tsi.OpReadRes:
+        return tsi.OpReadRes(op=None)  # type: ignore
+
+    def op_list(self, req: tsi.OpListReq) -> Iterator[tsi.OpReadRes]:
+        yield from []
+
+    def op_delete(self, req: tsi.OpDeleteReq) -> tsi.OpDeleteRes:
+        return tsi.OpDeleteRes(num_deleted=0)
+
+    # --- Datasets ---
+    def dataset_create(self, req: tsi.DatasetCreateReq) -> tsi.DatasetCreateRes:
+        return tsi.DatasetCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
+
+    def dataset_read(self, req: tsi.DatasetReadReq) -> tsi.DatasetReadRes:
+        return tsi.DatasetReadRes(dataset=None)  # type: ignore
+
+    def dataset_list(self, req: tsi.DatasetListReq) -> Iterator[tsi.DatasetReadRes]:
+        yield from []
+
+    def dataset_delete(self, req: tsi.DatasetDeleteReq) -> tsi.DatasetDeleteRes:
+        return tsi.DatasetDeleteRes(num_deleted=0)
+
+    # --- Scorers ---
+    def scorer_create(self, req: tsi.ScorerCreateReq) -> tsi.ScorerCreateRes:
+        return tsi.ScorerCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0, scorer=generate_id())
+
+    def scorer_read(self, req: tsi.ScorerReadReq) -> tsi.ScorerReadRes:
+        return tsi.ScorerReadRes(scorer=None)  # type: ignore
+
+    def scorer_list(self, req: tsi.ScorerListReq) -> Iterator[tsi.ScorerReadRes]:
+        yield from []
+
+    def scorer_delete(self, req: tsi.ScorerDeleteReq) -> tsi.ScorerDeleteRes:
+        return tsi.ScorerDeleteRes(num_deleted=0)
+
+    # --- Evaluations (V2) ---
+    def evaluation_create(self, req: tsi.EvaluationCreateReq) -> tsi.EvaluationCreateRes:
+        return tsi.EvaluationCreateRes(
+            digest=generate_id(), object_id=generate_id(), version_index=0, evaluation_ref=generate_id()
+        )
+
+    def evaluation_read(self, req: tsi.EvaluationReadReq) -> tsi.EvaluationReadRes:
+        return tsi.EvaluationReadRes(evaluation=None)  # type: ignore
+
+    def evaluation_list(self, req: tsi.EvaluationListReq) -> Iterator[tsi.EvaluationReadRes]:
+        yield from []
+
+    def evaluation_delete(self, req: tsi.EvaluationDeleteReq) -> tsi.EvaluationDeleteRes:
+        return tsi.EvaluationDeleteRes(num_deleted=0)
+
+    # --- Models ---
+    def model_create(self, req: tsi.ModelCreateReq) -> tsi.ModelCreateRes:
+        return tsi.ModelCreateRes(
+            digest=generate_id(), object_id=generate_id(), version_index=0, model_ref=generate_id()
+        )
+
+    def model_read(self, req: tsi.ModelReadReq) -> tsi.ModelReadRes:
+        return tsi.ModelReadRes(model=None)  # type: ignore
+
+    def model_list(self, req: tsi.ModelListReq) -> Iterator[tsi.ModelReadRes]:
+        yield from []
+
+    def model_delete(self, req: tsi.ModelDeleteReq) -> tsi.ModelDeleteRes:
+        return tsi.ModelDeleteRes(num_deleted=0)
+
+    # --- Evaluation Runs ---
+    def evaluation_run_create(self, req: tsi.EvaluationRunCreateReq) -> tsi.EvaluationRunCreateRes:
+        return tsi.EvaluationRunCreateRes(evaluation_run_id=generate_id())
+
+    def evaluation_run_read(self, req: tsi.EvaluationRunReadReq) -> tsi.EvaluationRunReadRes:
+        return tsi.EvaluationRunReadRes(evaluation_run=None)  # type: ignore
+
+    def evaluation_run_list(self, req: tsi.EvaluationRunListReq) -> Iterator[tsi.EvaluationRunReadRes]:
+        yield from []
+
+    def evaluation_run_delete(self, req: tsi.EvaluationRunDeleteReq) -> tsi.EvaluationRunDeleteRes:
+        return tsi.EvaluationRunDeleteRes(num_deleted=0)
+
+    def evaluation_run_finish(self, req: tsi.EvaluationRunFinishReq) -> tsi.EvaluationRunFinishRes:
+        return tsi.EvaluationRunFinishRes(success=True)
+
+    # --- Predictions ---
+    def prediction_create(self, req: tsi.PredictionCreateReq) -> tsi.PredictionCreateRes:
+        return tsi.PredictionCreateRes(prediction_id=generate_id())
+
+    def prediction_read(self, req: tsi.PredictionReadReq) -> tsi.PredictionReadRes:
+        return tsi.PredictionReadRes(prediction=None)  # type: ignore
+
+    def prediction_list(self, req: tsi.PredictionListReq) -> Iterator[tsi.PredictionReadRes]:
+        yield from []
+
+    def prediction_delete(self, req: tsi.PredictionDeleteReq) -> tsi.PredictionDeleteRes:
+        return tsi.PredictionDeleteRes(num_deleted=0)
+
+    def prediction_finish(self, req: tsi.PredictionFinishReq) -> tsi.PredictionFinishRes:
+        return tsi.PredictionFinishRes(success=True)
+
+    # --- Scores ---
+    def score_create(self, req: tsi.ScoreCreateReq) -> tsi.ScoreCreateRes:
+        return tsi.ScoreCreateRes(score_id=generate_id())
+
+    def score_read(self, req: tsi.ScoreReadReq) -> tsi.ScoreReadRes:
+        return tsi.ScoreReadRes(score=None)  # type: ignore
+
+    def score_list(self, req: tsi.ScoreListReq) -> Iterator[tsi.ScoreReadRes]:
+        yield from []
+
+    def score_delete(self, req: tsi.ScoreDeleteReq) -> tsi.ScoreDeleteRes:
+        return tsi.ScoreDeleteRes(num_deleted=0)
+
+
+# Module-level storage for originals
+_original_init_weave_get_server: Callable[..., Any] | None = None
+_original_get_entity_project_from_project_name: Callable[..., Any] | None = None
+_original_get_username: Callable[..., Any] | None = None
+
+
+def init_weave_get_server_factory(server: InMemoryWeaveTraceServer) -> Callable[..., Any]:
+    # Bypass the usage of Weave remote server
+    def init_weave_get_server(*args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
+        return server
+
+    return init_weave_get_server
+
+
+def get_entity_project_from_project_name_factory(entity_name: str) -> tuple[str, str]:
+    # Bypass the usage of API
+    try:
+        assert _original_get_entity_project_from_project_name is not None
+        if _original_get_entity_project_from_project_name is not get_entity_project_from_project_name_factory:
+            return _original_get_entity_project_from_project_name(entity_name)
+        else:
+            warnings.warn("W&B integration might have been repeatedly/recursively instrumented.")
+            return "agl", "weave"
+    except weave.trace.weave_init.WeaveWandbAuthenticationException:
+        # In case API is not available.
+        return "agl", "weave"
+
+
+def get_username() -> str:
+    # Bypass the usage of API
+    try:
+        assert _original_get_username is not None
+        return _original_get_username()
+    except RuntimeError:
+        return "agl"
+    except Exception as exc:
+        warnings.warn(f"Unexpected error in get_username. Using default username. Error: {exc}")
+        return "agl"
+
+
+def instrument_weave(server: InMemoryWeaveTraceServer):
+    """Patch the Weave/W&B integration to bypass actual network calls for testing."""
+
+    global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
+    _original_init_weave_get_server = weave.trace.weave_init.init_weave_get_server
+    _original_get_entity_project_from_project_name = weave.trace.weave_init.get_entity_project_from_project_name
+    _original_get_username = weave.trace.weave_init.get_username
+    weave.trace.weave_init.init_weave_get_server = init_weave_get_server_factory(server)
+    weave.trace.weave_init.get_entity_project_from_project_name = get_entity_project_from_project_name_factory
+    weave.trace.weave_init.get_username = get_username
+
+
+def uninstrument_weave():
+    """Restore the original Weave/W&B integration methods and HTTP requests."""
+    global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
+
+    if _original_init_weave_get_server is not None:
+        weave.trace.weave_init.init_weave_get_server = _original_init_weave_get_server
+        _original_init_weave_get_server = None
+    else:
+        raise RuntimeError("Weave/W&B integration was not instrumented.")
+
+    if _original_get_entity_project_from_project_name is not None:
+        weave.trace.weave_init.get_entity_project_from_project_name = _original_get_entity_project_from_project_name
+        _original_get_entity_project_from_project_name = None
+    else:
+        raise RuntimeError("Weave/W&B integration was not instrumented.")
+
+    if _original_get_username is not None:
+        weave.trace.weave_init.get_username = _original_get_username
+        _original_get_username = None
+    else:
+        raise RuntimeError("Weave/W&B integration was not instrumented.")
@@ -198,6 +198,7 @@ class LitAgent(Generic[T]):
            * `float` representing the final reward.
            * `List[ReadableSpan]` with OpenTelemetry spans.
            * `List[Span]` with Agent Lightning spans.
+            * `List[SpanCoreFields]` with Agent Lightning spans.
        """
        raise NotImplementedError("Agents must implement the `rollout` method.")

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

+from __future__ import annotations
+
 import logging
 import os
 import platform
+import sys
+import warnings
+from logging.config import dictConfig
+from typing import Any, Dict, Optional

-__all__ = ["configure_logger"]
+from rich.console import Console
+
+__all__ = ["setup", "configure_logger", "setup_module"]


 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
@@ -15,6 +23,10 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
    not propagate to the root logger, preventing duplicate log emission when
    applications compose multiple logging configurations.

+    !!! danger
+
+        This function is deprecated in favor of [`setup_logging`][agentlightning.setup_logging].
+
    Args:
        level: Logging level applied both to the logger and the installed
            handler. Defaults to `logging.INFO`.
@@ -32,23 +44,327 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
        logger.info("agent-lightning is ready!")
        ```
    """
+    warnings.warn("This function is deprecated in favor of `setup_logging`.", DeprecationWarning, stacklevel=2)

+    return setup_module(level=level, name=name, console=True, color=True, propagate=False)
+
+
+DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s"
+DATE_FORMAT = "%H:%M:%S"
+
+
+def _to_level_value(lvl: int | str) -> int:
+    if isinstance(lvl, int):
+        return lvl
+    val = getattr(logging, str(lvl).upper(), None)
+    if val is None:
+        raise ValueError(f"Invalid log level: {lvl}")
+    return val
+
+
+def _ensure_file_handler(
+    logger: logging.Logger,
+    filename: str,
+    *,
+    level: int,
+    formatter: Optional[logging.Formatter],
+) -> None:
+    """Attach a FileHandler to `logger` for `filename` if it doesn't already exist."""
+    abspath = os.path.abspath(filename)
+
+    # Avoid duplicates
+    for h in logger.handlers:
+        if isinstance(h, logging.FileHandler) and getattr(h, "baseFilename", None) == abspath:
+            return
+
+    # Ensure directory exists
+    dirname = os.path.dirname(abspath)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+
+    fh = logging.FileHandler(abspath, encoding="utf-8")
+    fh.setLevel(level)
+    if formatter is not None:
+        fh.setFormatter(formatter)
+    else:
+        fh.setFormatter(logging.Formatter(DEFAULT_FORMAT, DATE_FORMAT))
+
+    logger.addHandler(fh)
+
+
+def setup(
+    level: int | str = "INFO",
+    *,
+    console: bool = True,
+    color: bool | Dict[str, Any] = True,
+    propagate: bool = False,
+    disable_existing_loggers: bool = False,
+    capture_warnings: bool = False,
+    submodule_levels: Optional[dict[str, int | str]] = None,
+    extra_handlers: Optional[list[logging.Handler]] = None,
+    formatter: Optional[logging.Formatter] = None,
+    apply_to: Optional[list[str]] = None,
+    files: Optional[str | dict[str, str]] = None,
+) -> None:
+    """Configures logging for the `agentlightning` logger hierarchy.
+
+    This function provides a one-stop setup utility for configuring the
+    `agentlightning` root logger and optionally its submodules or external
+    loggers. It supports console logging, colored rich output, per-submodule
+    log levels, and optional handler/formatter injection.
+
+    The setup is intentionally isolated: it does not modify the global root
+    logger or loggers belonging to other libraries unless explicitly directed
+    via `apply_to`.
+
+    Args:
+        level:
+            Logging level for the base `agentlightning` logger. Accepts either
+            an integer (e.g., `logging.DEBUG`) or a string level name
+            (e.g., `"INFO"`). Defaults to `"INFO"`.
+        console:
+            Whether to attach a console handler to the logger. Defaults to
+            `True`.
+        color:
+            Enables rich-formatted output using `RichHandler` when `True`
+            or a configuration dict. If `False`, a plain text formatter is
+            used instead. Defaults to `True`.
+        propagate:
+            Whether `agentlightning` logs should propagate to ancestor
+            loggers. Defaults to `False`.
+        disable_existing_loggers:
+            Passed to `logging.config.dictConfig`. If `True`, disables all
+            existing configured loggers before applying this configuration.
+            Defaults to `False`.
+        capture_warnings:
+            If `True`, redirects Python `warnings` emitted via the `warnings`
+            module into the logging system. Defaults to `False`.
+        submodule_levels:
+            Mapping of submodule logger names to logging levels. If a specified
+            submodule level is more verbose than the base level, a warning is emitted.
+        extra_handlers:
+            A list of user-provided handlers to attach to the `agentlightning` logger.
+            Handlers are added idempotently; duplicates are not reattached.
+        formatter:
+            A formatter to apply to any handler under `agentlightning` that does not
+            already have one assigned. Useful for customizing output without overwriting
+            formatters on custom handlers.
+        apply_to:
+            A list of additional logger names to configure identically to
+            `agentlightning` base logger. Their handlers are replaced with copies of the base
+            handlers, and propagation is disabled to avoid duplicate log emission.
+        files:
+            If a string, attach a FileHandler to the base `agentlightning` logger.
+            If a dict, for each `(logger_name, filename)` pair, attach a FileHandler
+            directly to that logger.
+            Each file handler should use the logger's effective level at creation.
+
+    Notes:
+        * On Windows, this function forces UTF-8 mode in the console to prevent
+          issues with rich output or special characters.
+        * Submodule loggers can generate records below the handler's emission
+          threshold. Whether such records appear depends on both the logger's
+          level and the handler's level.
+        * `apply_to` loggers inherit the same handlers but do not propagate
+          upward, yielding isolated, consistent behavior.
+
+    Examples:
+        Basic setup:
+
+        >>> setup()
+
+        Enabling debug mode with no color:
+
+        >>> setup(level="DEBUG", color=False)
+
+        Overriding specific submodule levels:
+
+        >>> setup(submodule_levels={"agentlightning.io": "DEBUG"})
+
+        Attaching an additional file handler:
+
+        >>> fh = logging.FileHandler("app.log")
+        >>> setup(extra_handlers=[fh])
+    """
    # Ensure UTF-8 encoding on Windows consoles
-    # Note: This change does not fully represent support for execution under the windown system.
+    # Note: This change does not fully represent support for execution under the windows system.
    # It only fixes console printing issues caused by special characters.
    # TODO: More comprehensive Windows support may be needed in the future.
    if platform.system() == "Windows":
        os.environ["PYTHONUTF8"] = "1"

-    logger = logging.getLogger(name)
-    logger.handlers.clear()  # clear existing handlers
+    base_logger = setup_module(
+        level,
+        name="agentlightning",
+        console=console,
+        color=color,
+        propagate=propagate,
+        disable_existing_loggers=disable_existing_loggers,
+    )

-    # log to stdout
-    handler = logging.StreamHandler()
-    handler.setLevel(level)
-    formatter = logging.Formatter("%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s")
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-    logger.setLevel(level)
-    logger.propagate = False  # prevent double logging
-    return logger
+    base_level_value = base_logger.level
+
+    # Apply user-provided formatter (only to handlers without one,
+    # so we don't clobber custom extra_handlers)
+    if formatter is not None:
+        for h in base_logger.handlers:
+            if h.formatter is None:
+                h.setFormatter(formatter)
+
+    # Attach user-provided handler(s) if any, idempotently
+    if extra_handlers:
+        for h in extra_handlers:
+            if h not in base_logger.handlers:
+                base_logger.addHandler(h)
+
+    # Per-submodule levels
+    if submodule_levels:
+        for name, lvl in submodule_levels.items():
+            sub_level = _to_level_value(lvl)
+
+            # Emit a warning if submodule level is lower (more verbose) than the global/base level
+            if sub_level < base_level_value:
+                base_logger.warning(
+                    "Submodule logger '%s' level %s (%s) is more verbose than base "
+                    "logger level %s (%s). Records below the base level may still be "
+                    "filtered out by handlers depending on their own levels.",
+                    name,
+                    lvl,
+                    sub_level,
+                    logging.getLevelName(base_level_value),
+                    base_level_value,
+                )
+
+            # The logger will *create* records down to the logger's level, but a handler
+            # with a higher level will still drop anything below its own threshold.
+            # Effective emission is gated by both: record.level >= logger.level AND handler.level.
+            logging.getLogger(name).setLevel(lvl)
+
+    # Attach file handlers if requested
+    if files is not None:
+        if isinstance(files, str):
+            # Single file for the entire `agentlightning` hierarchy.
+            _ensure_file_handler(
+                logger=base_logger,
+                filename=files,
+                level=base_level_value,
+                formatter=formatter,
+            )
+        else:
+            # Per-logger files
+            for logger_name, filename in files.items():
+                lg = logging.getLogger(logger_name)
+                # Use the logger's *effective* level at creation time
+                effective_level = lg.getEffectiveLevel()
+                _ensure_file_handler(
+                    logger=lg,
+                    filename=filename,
+                    level=effective_level,
+                    formatter=formatter,
+                )
+
+    # Optionally apply the same handler setup to other loggers outside this module
+    if apply_to:
+        for name in apply_to:
+            lg = logging.getLogger(name)
+            # This removes any existing handlers so we don't duplicate output
+            # and ensures these loggers share exactly the same handlers as base_logger.
+            lg.handlers.clear()
+            for h in base_logger.handlers:
+                lg.addHandler(h)
+            lg.setLevel(base_logger.level)
+            # We've attached handlers directly to these loggers; if propagate
+            # stayed True, records would bubble up to ancestor loggers and could be
+            # emitted twice (here and on the parent/root). Setting False isolates them.
+            lg.propagate = False
+
+    # Optionally capture warnings
+    if capture_warnings:
+        logging.captureWarnings(True)
+
+
+def setup_module(
+    level: int | str = "INFO",
+    *,
+    name: str = "agentlightning",
+    console: bool = True,
+    color: bool | Dict[str, Any] = True,
+    propagate: bool = False,
+    disable_existing_loggers: bool = False,
+) -> logging.Logger:
+    """Initializes and returns the base logger for `agentlightning`.
+
+    This function constructs and applies a `dictConfig` configuration for the
+    logger hierarchy rooted at `name`. It supports either rich console
+    formatting (via `RichHandler`) or plain text formatting, based on the
+    `color` argument.
+
+    Unlike [`setup_logging`][agentlightning.setup_logging], this function configures only a single logger namespace
+    and does not attach extra handlers or submodule levels. It is primarily used
+    internally by [`setup_logging`][agentlightning.setup_logging] but is also suitable for direct integration in
+    custom logging workflows.
+    """
+    root_cfg: Dict[str, Any] = {
+        "version": 1,
+        "disable_existing_loggers": disable_existing_loggers,
+        "loggers": {
+            name: {
+                "handlers": [],
+                "level": level,
+                "propagate": propagate,
+            }
+        },
+        "handlers": {},
+        "formatters": {},
+    }
+
+    # Choose formatter / handler definition
+    if color is not False and console:
+        # Console must be true to display colored outputs
+        if isinstance(color, dict):
+            rich_handler_config = color
+        else:
+            rich_handler_config: Dict[str, Any] = {
+                "rich_tracebacks": False,
+                "markup": False,
+                "show_time": True,
+                "show_path": True,
+            }
+
+            if not _has_width():
+                # e.g., in a CI environment.
+                rich_handler_config["console"] = Console(width=200)
+
+        root_cfg["handlers"]["console"] = {
+            "class": "rich.logging.RichHandler",
+            "level": level,
+            **rich_handler_config,
+        }
+        # RichHandler manages its own style; keep formatter None
+    else:
+        fmt_name = "plain"
+        root_cfg["formatters"][fmt_name] = {
+            "format": DEFAULT_FORMAT,
+            "datefmt": DATE_FORMAT,
+        }
+
+        if console:
+            root_cfg["handlers"]["console"] = {
+                "class": "logging.StreamHandler",
+                "level": level,
+                "formatter": fmt_name,
+            }
+
+    # Attach selected handlers to agentlightning
+    handler_names = list(root_cfg["handlers"].keys())
+    root_cfg["loggers"][name]["handlers"] = handler_names
+
+    # Apply dictConfig (this resets the logger handlers)
+    dictConfig(root_cfg)
+
+    return logging.getLogger(name)
+
+
+def _has_width() -> bool:
+    """Automatically determine whether the terminal has a width."""
+    return sys.stdout.isatty()
@@ -11,16 +11,30 @@ from __future__ import annotations

 import asyncio
 import logging
+import random
+import threading
 import time
-from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
+from contextlib import suppress
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypeVar,
+    cast,
+)

 from opentelemetry.sdk.trace import ReadableSpan

 from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.store.base import LightningStore
-from agentlightning.tracer.agentops import AgentOpsTracer
 from agentlightning.tracer.base import Tracer
+from agentlightning.tracer.otel import OtelTracer
 from agentlightning.types import (
    AttemptedRollout,
    Hook,
@@ -29,7 +43,9 @@ from agentlightning.types import (
    RolloutMode,
    RolloutRawResult,
    Span,
+    SpanCoreFields,
 )
+from agentlightning.utils.system_snapshot import system_snapshot

 if TYPE_CHECKING:
    from agentlightning.execution.events import ExecutionEvent
@@ -52,7 +68,16 @@ class LitAgentRunner(Runner[T_task]):
        worker_id: Identifier for the active worker process, if any.
    """

-    def __init__(self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
+    def __init__(
+        self,
+        tracer: Tracer,
+        max_rollouts: Optional[int] = None,
+        poll_interval: float = 5.0,
+        heartbeat_interval: float = 10.0,
+        interval_jitter: float = 0.5,
+        heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
+        heartbeat_include_gpu: bool = False,
+    ) -> None:
        """Initialize the agent runner.

        Args:
@@ -60,11 +85,25 @@ class LitAgentRunner(Runner[T_task]):
            max_rollouts: Optional cap on iterations processed by
                [`iter`][agentlightning.LitAgentRunner.iter].
            poll_interval: Seconds to wait between store polls when no work is available.
+            heartbeat_interval: Seconds to wait between sending heartbeats to the store.
+            interval_jitter: Jitter factor for the poll interval. The actual interval will be between
+                poll_interval - interval_jitter and poll_interval + interval_jitter.
+                This is to avoid the overload caused by the synchronization of the runners.
+            heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
+                "thread" is the default and recommended mode as it prevents blocking the event loop
+                under load. Use "asyncio" for simpler deployments with low worker counts.
+            heartbeat_include_gpu: Whether to include GPU stats in heartbeat snapshots.
+                Querying GPU stats can be slow under load, so this is disabled by default.
        """
        super().__init__()
        self._tracer = tracer
        self._max_rollouts = max_rollouts
        self._poll_interval = poll_interval
+        self._heartbeat_interval = heartbeat_interval
+        self._interval_jitter = interval_jitter
+        self._heartbeat_launch_mode = heartbeat_launch_mode
+        self._heartbeat_include_gpu = heartbeat_include_gpu
+        self._random_state = random.Random()

        # Set later
        self._agent: Optional[LitAgent[T_task]] = None
@@ -105,7 +144,7 @@ class LitAgentRunner(Runner[T_task]):
        self._store = store
        self.worker_id = worker_id

-        self._tracer.init_worker(worker_id)
+        self._tracer.init_worker(worker_id, store)

    def teardown(self, *args: Any, **kwargs: Any) -> None:
        """Teardown the runner and clean up all resources.
@@ -243,49 +282,84 @@ class LitAgentRunner(Runner[T_task]):
        """
        store = self.get_store()

-        trace_spans: list[ReadableSpan] | list[Span] = []
+        trace_spans: list[Span] = []
+        result_recognized: bool = False

        # Case 0: result is None
        if raw_result is None:
            trace_spans = self._tracer.get_last_trace()
+            result_recognized = True

        # Case 1: result is a float (final reward)
-        if isinstance(raw_result, float):
+        if isinstance(raw_result, (bool, int, float)):
+            if isinstance(raw_result, (bool, int)):
+                logger.warning(
+                    f"{self._log_prefix(rollout.rollout_id)} Reward is not a number, got: {type(raw_result)}. "
+                    "Auto converting to float."
+                )
+                raw_result = float(raw_result)
            # Preserve the existing spans before another span is emitted
            trace_spans = list(self._tracer.get_last_trace())
-            # This will emit another span to the tracer
-            reward_span = emit_reward(raw_result)
-            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
-            trace_spans.append(reward_span)
+            # This will NOT emit another span to the tracer
+            reward_span_core_fields = emit_reward(raw_result, propagate=False)
+            # We add it to the store manually
+            sequence_id = await store.get_next_span_sequence_id(rollout.rollout_id, rollout.attempt.attempt_id)
+            reward_span = Span.from_core_fields(
+                reward_span_core_fields,
+                rollout_id=rollout.rollout_id,
+                attempt_id=rollout.attempt.attempt_id,
+                sequence_id=sequence_id,
+            )
+            await store.add_span(reward_span)
+            result_recognized = True

+        # Case 2-4: result is a list
        if isinstance(raw_result, list):
-            # For rollout methods that return a list, we assume that the returned spans
-            # are the complete span set from the whole rollout
-            trace_spans = raw_result
-
            # Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
            if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
-
-                if not isinstance(
-                    self._tracer, AgentOpsTracer
-                ):  # TODO: this should be replaced with general OpenTelemetry tracer in next version
-                    for span in raw_result:
-                        await store.add_otel_span(
-                            rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
-                        )
-                else:
+                if isinstance(self._tracer, OtelTracer):
                    logger.warning(
                        f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
                        "The traces should have already been added to the store. "
-                        "No need to return anything from rollout."
+                        "Returning the traces from the rollout will result in duplicate spans."
                    )
+                for span in raw_result:
+                    added_span = await store.add_otel_span(
+                        rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
+                    )
+                    if added_span is not None:
+                        trace_spans.append(added_span)
+                    else:
+                        logger.error(
+                            f"{self._log_prefix(rollout.rollout_id)} Failed to add OpenTelemetry span to the store: {span}"
+                        )
+                result_recognized = True

            # Case 3: result is a list of Span (agentlightning spans)
            elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
                # Add the spans directly to the store
                for span in raw_result:
                    await store.add_span(cast(Span, span))
-                trace_spans = raw_result
+                trace_spans = [cast(Span, span) for span in raw_result]
+                result_recognized = True
+
+            # Case 4: result is a list of SpanCoreFields (agentlightning spans)
+            elif len(raw_result) > 0 and all(isinstance(t, SpanCoreFields) for t in raw_result):
+                # Add the spans directly to the store too, but needs to get sequence id first
+                sequence_ids = await store.get_many_span_sequence_ids(
+                    [(rollout.rollout_id, rollout.attempt.attempt_id) for _ in range(len(raw_result))]
+                )
+                trace_spans = [
+                    Span.from_core_fields(
+                        cast(SpanCoreFields, span_core_fields),
+                        rollout_id=rollout.rollout_id,
+                        attempt_id=rollout.attempt.attempt_id,
+                        sequence_id=sequence_id,
+                    )
+                    for span_core_fields, sequence_id in zip(raw_result, sequence_ids, strict=True)
+                ]
+                await store.add_many_spans(trace_spans)
+                result_recognized = True

            # Left over cases for list
            elif len(raw_result) == 0:
@@ -293,7 +367,8 @@ class LitAgentRunner(Runner[T_task]):
                    f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
                    "Please check your rollout implementation."
                )
-                trace_spans = raw_result
+                trace_spans = []
+                result_recognized = True

            else:
                types = [type(t).__name__ for t in raw_result][:10]
@@ -302,8 +377,225 @@ class LitAgentRunner(Runner[T_task]):
                    f"but got: {', '.join(types)}..."
                )

+        if not result_recognized:
+            raise TypeError(
+                f"Invalid raw result type. It's expected to be none, float, or a list of ReadableSpan or Span, "
+                f"but got: {type(raw_result).__name__}..."
+            )
+
        return trace_spans

+    async def _emit_heartbeat(self, store: LightningStore) -> None:
+        """Send a heartbeat tick to the store.
+
+        Args:
+            store: The lightning store to update.
+        """
+        logger.debug(f"{self._log_prefix()} Preparing to emit heartbeat.")
+        worker_id = self.get_worker_id()
+
+        try:
+            snapshot = await asyncio.wait_for(
+                asyncio.to_thread(system_snapshot, self._heartbeat_include_gpu),
+                timeout=self._heartbeat_interval,
+            )
+            logger.debug(f"{self._log_prefix()} Heartbeat snapshot acquired.")
+        except asyncio.TimeoutError:
+            logger.warning(
+                "%s Heartbeat snapshot acquisition timed out after %.1fs, skipping.",
+                self._log_prefix(),
+                self._heartbeat_interval,
+            )
+            return
+        except asyncio.CancelledError:
+            # bypass the exception
+            raise
+        except Exception:
+            logger.exception("%s Unable to acquire heartbeat snapshot.", self._log_prefix())
+            return
+
+        try:
+            await asyncio.wait_for(store.update_worker(worker_id, snapshot), timeout=self._heartbeat_interval)
+            logger.debug(f"{self._log_prefix()} Heartbeat updated successfully.")
+        except asyncio.CancelledError:
+            # bypass the exception
+            raise
+        except asyncio.TimeoutError:
+            logger.warning(
+                "%s update worker heartbeat timed out after %.1fs, skipping.",
+                self._log_prefix(),
+                self._heartbeat_interval,
+            )
+        except Exception:
+            logger.exception("%s Unable to update worker heartbeat.", self._log_prefix())
+
+    def _start_heartbeat_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
+        """Start a background heartbeat loop and return an async stopper."""
+
+        if self._heartbeat_interval <= 0:
+            return None
+
+        if self.worker_id is None:
+            logger.warning("%s Cannot start heartbeat loop without worker_id.", self._log_prefix())
+            return None
+
+        if self._heartbeat_launch_mode == "asyncio":
+            return self._start_heartbeat_asyncio_loop(store)
+        if self._heartbeat_launch_mode == "thread":
+            return self._start_heartbeat_thread_loop(store)
+        raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
+
+    def _start_heartbeat_asyncio_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
+        """Start a background heartbeat loop using asyncio.
+
+        Args:
+            store: The lightning store to update.
+
+        Returns:
+            An async stopper function that can be used to stop the heartbeat loop.
+        """
+
+        stop_event = asyncio.Event()
+
+        async def heartbeat_loop() -> None:
+            while not stop_event.is_set():
+                try:
+                    # Run _emit_heartbeat in thread pool to avoid blocking the event loop.
+                    # Timeout at the interval - if it takes longer, the data is stale anyway.
+                    await self._emit_heartbeat(store)
+                except Exception:
+                    logger.exception("%s Heartbeat failed.", self._log_prefix())
+                with suppress(asyncio.TimeoutError):
+                    interval = self._heartbeat_interval + self._random_state.uniform(
+                        -self._interval_jitter, self._interval_jitter
+                    )
+                    interval = max(interval, 0.01)
+                    await asyncio.wait_for(stop_event.wait(), timeout=interval)
+
+        task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
+
+        async def stop() -> None:
+            stop_event.set()
+            with suppress(asyncio.CancelledError):
+                await task
+
+        return stop
+
+    def _start_heartbeat_thread_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
+        """Start a background heartbeat loop using threading.
+
+        It uses two threads: one to produce the snapshot and one to consume it,
+        to avoid either of them blocking the event loop.
+
+        Args:
+            store: The lightning store to update.
+
+        Returns:
+            An async stopper function that can be used to stop the heartbeat loop.
+        """
+        stop_evt = threading.Event()
+        lock = threading.Lock()
+
+        latest_snapshot = None
+        latest_ts = 0.0  # time.monotonic() when snapshot was captured
+
+        # Consider snapshot stale after ~1 interval plus jitter slack.
+        stale_after = self._heartbeat_interval + self._interval_jitter + 1.0
+
+        worker_id = self.get_worker_id()
+
+        def producer() -> None:
+            nonlocal latest_snapshot, latest_ts
+            while not stop_evt.is_set():
+                try:
+                    logger.debug(f"{self._log_prefix()} Heartbeat producer: acquiring snapshot.")
+                    snap = system_snapshot(self._heartbeat_include_gpu)  # sync
+                    logger.debug(f"{self._log_prefix()} Heartbeat producer: snapshot acquired.")
+                    ts = time.monotonic()
+                    with lock:
+                        latest_snapshot = snap
+                        latest_ts = ts
+                except Exception:
+                    logger.warning("%s Heartbeat producer: system_snapshot failed.", self._log_prefix(), exc_info=True)
+
+                interval = self._heartbeat_interval + self._random_state.uniform(
+                    -self._interval_jitter, self._interval_jitter
+                )
+                stop_evt.wait(max(interval, 0.01))
+
+        def consumer() -> None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            last_warned_ts = None  # Track which snapshot we've already warned about
+            try:
+                while not stop_evt.is_set():
+                    with lock:
+                        snap = latest_snapshot
+                        ts = latest_ts
+
+                    wait_interval = max(
+                        self._heartbeat_interval
+                        + self._random_state.uniform(-self._interval_jitter, self._interval_jitter),
+                        0.01,
+                    )
+
+                    if snap is None:
+                        # probably just started
+                        logger.debug("%s Heartbeat consumer: no snapshot yet; skipping update.", self._log_prefix())
+                        stop_evt.wait(wait_interval)
+                        continue
+
+                    age = time.monotonic() - ts
+                    if age > stale_after:
+                        # Only warn once per stale snapshot (check if we haven't warned about this timestamp yet)
+                        if last_warned_ts != ts:
+                            logger.warning(
+                                "%s Heartbeat consumer: snapshot stale (age=%.2fs > %.2fs); skipping update.",
+                                self._log_prefix(),
+                                age,
+                                stale_after,
+                            )
+                            last_warned_ts = ts
+                        stop_evt.wait(wait_interval)
+                        continue
+
+                    try:
+                        logger.debug(f"{self._log_prefix()} Heartbeat consumer: updating worker.")
+                        loop.run_until_complete(
+                            asyncio.wait_for(
+                                store.update_worker(worker_id, snap),
+                                timeout=self._heartbeat_interval,
+                            )
+                        )
+                        logger.debug(f"{self._log_prefix()} Heartbeat consumer: worker updated.")
+                    except asyncio.TimeoutError:
+                        logger.warning(
+                            "%s Heartbeat consumer: update timed out after %.1fs.",
+                            self._log_prefix(),
+                            self._heartbeat_interval,
+                        )
+                    except Exception:
+                        logger.warning("%s Heartbeat consumer: update failed.", self._log_prefix(), exc_info=True)
+
+                    stop_evt.wait(wait_interval)
+            finally:
+                with suppress(Exception):
+                    loop.stop()
+                with suppress(Exception):
+                    loop.close()
+
+        t_prod = threading.Thread(target=producer, name=f"{worker_id}-heartbeat-producer", daemon=True)
+        t_cons = threading.Thread(target=consumer, name=f"{worker_id}-heartbeat-consumer", daemon=True)
+        t_prod.start()
+        t_cons.start()
+
+        async def stop() -> None:
+            stop_evt.set()
+            await asyncio.to_thread(t_prod.join)
+            await asyncio.to_thread(t_cons.join)
+
+        return stop
+
    async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
        """Sleep until the next poll interval, with optional event-based interruption.

@@ -314,11 +606,13 @@ class LitAgentRunner(Runner[T_task]):
            event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep.
                If set during the sleep period, the method returns immediately.
        """
+        interval = self._poll_interval + self._random_state.uniform(-self._interval_jitter, self._interval_jitter)
+        interval = max(interval, 0.01)
        if event is None:
-            await asyncio.sleep(self._poll_interval)
+            await asyncio.sleep(interval)
            return
        current_time = time.time()
-        next_time = current_time + self._poll_interval
+        next_time = current_time + interval
        while time.time() < next_time:
            await asyncio.sleep(0.1)
            if event.is_set():
@@ -356,6 +650,8 @@ class LitAgentRunner(Runner[T_task]):
                logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
                return rollout_id

+        logger.debug(f"{self._log_prefix(rollout_id)} Resources fetched (id={resources_update.resources_id}).")
+
        trace_spans: List[ReadableSpan] | List[Span] = []
        has_exception: bool = False

@@ -363,9 +659,11 @@ class LitAgentRunner(Runner[T_task]):
            await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)

            start_time = time.time()
+            logger.debug(f"{self._log_prefix(rollout_id)} Prepared for trace context.")
            async with self._tracer.trace_context(
-                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+                name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
            ):
+                logger.debug(f"{self._log_prefix(rollout_id)} Entered trace context.")
                await self._trigger_hooks(
                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
                )
@@ -377,21 +675,27 @@ class LitAgentRunner(Runner[T_task]):
                    rollout_method = (
                        agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
                    )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Starting async rollout method.")
                    result = await rollout_method(
                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
                    )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Async rollout method completed.")
                else:
                    rollout_method = (
                        agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
                    )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Starting sync rollout method.")
                    result = rollout_method(
                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
                    )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Sync rollout method completed.")

                await self._trigger_hooks(
                    hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
                )

+            logger.debug(f"{self._log_prefix(rollout_id)} Trace context exited.")
+
            # Possible exceptions in post_process will be caught in the overall exception handler
            trace_spans = await self._post_process_rollout_result(next_rollout, result)
            last_reward = find_final_reward(trace_spans)
@@ -450,39 +754,40 @@ class LitAgentRunner(Runner[T_task]):
        logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
        store = self.get_store()

-        while not (event is not None and event.is_set()) and (
-            self._max_rollouts is None or num_tasks_processed < self._max_rollouts
-        ):
-            # Retrieve the next rollout
-            next_rollout: Optional[Rollout] = None
-            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()
+        stop_heartbeat = self._start_heartbeat_loop(store)
+
+        try:
+            while not (event is not None and event.is_set()) and (
+                self._max_rollouts is None or num_tasks_processed < self._max_rollouts
+            ):
+                # Retrieve the next rollout
+                next_rollout: Optional[Rollout] = None
+                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."
+                        )
+                        await self._sleep_until_next_poll(event)
+                    else:
+                        break
+
                if next_rollout is None:
-                    logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
-                    await self._sleep_until_next_poll(event)
-                else:
-                    break
+                    return

-            if next_rollout is None:
-                return
+                # Execute the step
+                await self._step_impl(next_rollout)

-            try:
-                # Claim the rollout but updating the current worker id
-                await store.update_attempt(
-                    next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
-                )
-            except Exception:
-                # This exception could happen if the rollout is dequeued and the other end died for some reason
-                logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
-                continue
-
-            # Execute the step
-            await self._step_impl(next_rollout)
-
-            num_tasks_processed += 1
-            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
-                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}")
+                num_tasks_processed += 1
+                if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                    logger.info(
+                        f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}"
+                    )
+        finally:
+            if stop_heartbeat is not None:
+                await stop_heartbeat()

        logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")

@@ -525,7 +830,9 @@ class LitAgentRunner(Runner[T_task]):
        else:
            resources_id = None

-        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        attempted_rollout = await self.get_store().start_rollout(
+            input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
+        )
        rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)

        completed_rollout = await store.get_rollout_by_id(rollout_id)
@@ -12,7 +12,7 @@ from agentlightning.client import AgentLightningClient
 from agentlightning.litagent import LitAgent
 from agentlightning.litagent.litagent import is_v0_1_rollout_api
 from agentlightning.tracer.base import Tracer
-from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet

 from .base import Runner

@@ -99,7 +99,7 @@ class LegacyAgentRunner(Runner[Any]):
        trace: Any = None
        final_reward: Optional[float] = None
        triplets: Optional[List[Triplet]] = None
-        trace_spans: Optional[List[ReadableSpan]] = None
+        trace_spans: Optional[List[SpanLike]] = None

        # Handle different types of results from the agent
        # Case 1: result is a float (final reward)
@@ -108,10 +108,14 @@ class LegacyAgentRunner(Runner[Any]):
        # Case 2: result is a list of Triplets
        if isinstance(result, list) and all(isinstance(t, Triplet) for t in result):
            triplets = result  # type: ignore
-        # Case 3: result is a list of ReadableSpan (OpenTelemetry spans)
-        if isinstance(result, list) and all(isinstance(t, ReadableSpan) for t in result):
+        # Case 3.1: result is a list of ReadableSpan (OpenTelemetry spans)
+        if isinstance(result, list) and all(isinstance(t, (ReadableSpan)) for t in result):
            trace_spans = result  # type: ignore
            trace = [json.loads(readable_span.to_json()) for readable_span in trace_spans]  # type: ignore
+        # Case 3.2: result is a list of Span (Agent-lightning spans)
+        if isinstance(result, list) and all(isinstance(t, Span) for t in result):
+            trace_spans = result  # type: ignore
+            trace = [span.model_dump() for span in trace_spans]  # type: ignore
        # Case 4: result is a list of dict (trace JSON)
        if isinstance(result, list) and all(isinstance(t, dict) for t in result):
            trace = result
@@ -123,10 +127,9 @@ class LegacyAgentRunner(Runner[Any]):

        # If the agent has tracing enabled, use the tracer's last trace if not already set
        if self.tracer and (trace is None or trace_spans is None):
-            spans = self.tracer.get_last_trace()
-            if spans:
-                trace = [json.loads(readable_span.to_json()) for readable_span in spans]
-                trace_spans = spans
+            trace_spans = self.tracer.get_last_trace()  # type: ignore
+            if trace_spans:
+                trace = [cast(Span, span).model_dump() for span in trace_spans]

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

-from .base import LightningStore
+from .base import LightningStore, LightningStoreCapabilities, LightningStoreStatistics
 from .client_server import LightningStoreClient, LightningStoreServer
+from .collection_based import CollectionBasedLightningStore
 from .memory import InMemoryLightningStore
 from .threading import LightningStoreThreaded

 __all__ = [
    "LightningStore",
+    "LightningStoreCapabilities",
+    "LightningStoreStatistics",
    "LightningStoreClient",
    "LightningStoreServer",
    "InMemoryLightningStore",
+    "CollectionBasedLightningStore",
    "LightningStoreThreaded",
 ]
@@ -2,7 +2,7 @@

 from __future__ import annotations

-from typing import Any, Dict, List, Literal, Optional, Sequence
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict

 from opentelemetry.sdk.trace import ReadableSpan

@@ -10,13 +10,17 @@ from agentlightning.types import (
    Attempt,
    AttemptedRollout,
    AttemptStatus,
+    EnqueueRolloutRequest,
    NamedResources,
    ResourcesUpdate,
    Rollout,
    RolloutConfig,
+    RolloutMode,
    RolloutStatus,
    Span,
    TaskInput,
+    Worker,
+    WorkerStatus,
 )


@@ -52,6 +56,51 @@ UNSET = _UnsetType()
 Unset = _UnsetType  # Alias for convenience


+class LightningStoreCapabilities(TypedDict, total=False):
+    """Capability of a LightningStore implementation.
+
+    All keys are optional and false by default.
+    """
+
+    thread_safe: bool
+    """Whether the store is thread-safe."""
+    async_safe: bool
+    """Whether the store is async-safe."""
+    zero_copy: bool
+    """Whether the store has only one copy across all threads/processes."""
+    otlp_traces: bool
+    """Whether the store supports OTLP/HTTP traces."""
+
+
+class LightningStoreStatistics(TypedDict, total=False):
+    """Statistics of a LightningStore implementation."""
+
+    name: str
+    """Name of the store implementation."""
+    total_rollouts: int
+    """Total number of rollouts in the store."""
+    total_attempts: int
+    """Total number of attempts in the store."""
+    total_spans: int
+    """Total number of spans in the store."""
+    total_resources: int
+    """Total number of resources in the store."""
+    total_workers: int
+    """Total number of workers in the store."""
+    uptime: float
+    """Uptime of since the store has been started."""
+
+    # Memory-related statistics
+    total_span_bytes: int
+    """Total number of bytes of spans in the store."""
+    eviction_threshold_bytes: int
+    """Eviction threshold for spans in bytes."""
+    safe_threshold_bytes: int
+    """Safe threshold for spans in bytes."""
+    memory_capacity_bytes: int
+    """Memory capacity of the store in bytes."""
+
+
 class LightningStore:
    """Contract for the persistent control-plane that coordinates training rollouts.

@@ -74,13 +123,46 @@ class LightningStore:
    Unless stated otherwise, missing identifiers should result in a `ValueError`.
    """

+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            thread_safe=False,
+            async_safe=False,
+            zero_copy=False,
+            otlp_traces=False,
+        )
+
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        return {
+            "name": self.__class__.__name__,
+        }
+
+    def otlp_traces_endpoint(self) -> str:
+        """Return the OTLP/HTTP traces endpoint of the store.
+
+        The traces can have rollout ID and attempt ID (and optionally sequence ID)
+        saved in the "resource" of the spans.
+        The store, if it supports OTLP, should be able to receive the traces and save them
+        via [`add_span`][agentlightning.LightningStore.add_span] or
+        [`add_otel_span`][agentlightning.LightningStore.add_otel_span].
+
+        The endpoint should be compatible with [OTLP HTTP protocol](https://opentelemetry.io/docs/specs/otlp/).
+        It's not necessarily compatible with OTLP gRPC protocol.
+
+        The returned endpoint will usually ends with `/v1/traces`.
+        """
+        raise NotImplementedError()
+
    async def start_rollout(
        self,
        input: TaskInput,
-        mode: Literal["train", "val", "test"] | None = None,
+        mode: RolloutMode | None = None,
        resources_id: str | None = None,
        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
+        worker_id: str | None = None,
    ) -> AttemptedRollout:
        """Register a rollout and immediately create its first attempt.

@@ -103,6 +185,7 @@ class LightningStore:
            resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
            config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
            metadata: Free-form metadata persisted verbatim with the rollout.
+            worker_id: Optional worker identifier to associate the new attempt with.

        Returns:
            The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
@@ -148,7 +231,23 @@ class LightningStore:
        """
        raise NotImplementedError()

-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        """Persist multiple rollouts in `queuing` state.
+
+        The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
+        per request and preserves the input ordering. Subclasses can override to provide
+        more efficient bulk enqueue semantics.
+
+        Args:
+            rollouts: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
+                parameters. Each entry requires `input` and can optionally include other fields.
+
+        Returns:
+            Rollouts enqueued in the same order as `rollouts`.
+        """
+        raise NotImplementedError()
+
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
        """Claim the oldest queued rollout and transition it to `preparing`.

        This function do not block.
@@ -161,6 +260,11 @@ class LightningStore:
          the number of attempts already registered for the rollout plus one.
        * Return an [`AttemptedRollout`][agentlightning.AttemptedRollout] snapshot so the
          runner knows both rollout metadata and the attempt identifier.
+        * Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
+          (e.g., `last_dequeue_time`) when `worker_id` is provided.
+
+        Args:
+            worker_id: Optional worker identifier to associate the claimed attempt with.

        Returns:
            The next attempt to execute, or `None` when no eligible rollouts are queued.
@@ -170,7 +274,30 @@ class LightningStore:
        """
        raise NotImplementedError()

-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        """Claim up to `limit` queued rollouts without blocking.
+
+        The implementation can repeatedly invokes
+        [`dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] until reaching
+        the requested limit or the queue is empty. Subclasses can override it to fetch
+        multiple rollouts atomically.
+
+        Args:
+            limit: Maximum number of rollouts to claim. Non-positive values return an empty list.
+            worker_id: Optional worker identifier passed through to each dequeue call.
+
+        Returns:
+            Attempted rollouts claimed in FIFO order. May contain fewer than `limit` entries
+            when the queue is exhausted.
+        """
+        raise NotImplementedError()
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
        """Create a manual retry attempt for an existing rollout.

        This is typically invoked by runners that wish to retry outside of the
@@ -181,6 +308,7 @@ class LightningStore:

        Args:
            rollout_id: Unique identifier of the rollout receiving a new attempt.
+            worker_id: Optional worker identifier to associate the new attempt with.

        Returns:
            The rollout paired with its newly-created attempt.
@@ -191,7 +319,15 @@ class LightningStore:
        """
        raise NotImplementedError()

-    async def add_span(self, span: Span) -> Span:
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        """Persist a sequence of pre-constructed spans emitted during rollout execution.
+
+        Implementations can simply delegate to [`add_span()`][agentlightning.LightningStore.add_span] for each span.
+        However, if the store supports bulk insertion, it can implement this method to improve performance.
+        """
+        raise NotImplementedError()
+
+    async def add_span(self, span: Span) -> Optional[Span]:
        """Persist a pre-constructed span emitted during rollout execution.

        The provided [`Span`][agentlightning.Span] must already contain the `rollout_id`,
@@ -208,6 +344,7 @@ class LightningStore:

        Returns:
            The stored span record (implementations may return a copy).
+            Return `None` if the span was not added due to a duplicate.

        Raises:
            NotImplementedError: Subclasses must implement span persistence.
@@ -221,7 +358,7 @@ class LightningStore:
        attempt_id: str,
        readable_span: ReadableSpan,
        sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
        """Convert and persist an OpenTelemetry span for a particular attempt.

        Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
@@ -238,7 +375,7 @@ class LightningStore:
                automatically.

        Returns:
-            The stored span record.
+            The stored span record. Return `None` if the span was not added due to a duplicate.

        Raises:
            NotImplementedError: Subclasses must implement span persistence.
@@ -247,30 +384,77 @@ class LightningStore:
        raise NotImplementedError()

    async def query_rollouts(
-        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Rollout]:
+        self,
+        *,
+        status_in: Optional[Sequence[RolloutStatus]] = None,
+        rollout_id_in: Optional[Sequence[str]] = None,
+        rollout_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+        # Deprecated fields
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> Sequence[Rollout]:
        """Retrieve rollouts filtered by status and/or explicit identifiers.

+        This interface supports structured filtering, sorting, and pagination so
+        callers can build simple dashboards without copying data out of the
+        store. The legacy parameters `status` and `rollout_ids` remain valid and
+        are treated as aliases for `status_in` and `rollout_id_in`
+        respectively—when both the new and deprecated parameters are supplied
+        the new parameters take precedence.
+
        Args:
-            status: Optional whitelist of [`RolloutStatus`][agentlightning.RolloutStatus] values.
-            rollout_ids: Optional whitelist of rollout identifiers to include.
+            status_in: Optional whitelist of [`RolloutStatus`][agentlightning.RolloutStatus] values.
+            rollout_id_in: Optional whitelist of rollout identifiers to include.
+            rollout_id_contains: Optional substring match for rollout identifiers.
+            filter_logic: Logical operator to combine filters.
+            sort_by: Optional field to sort by. Must reference a numeric or string
+                field on [`Rollout`][agentlightning.Rollout].
+            sort_order: Direction to sort when `sort_by` is provided.
+            limit: Maximum number of rows to return. Use `-1` for "no limit".
+            offset: Number of rows to skip before returning results.
+            status: Deprecated field. Use `status_in` instead.
+            rollout_ids: Deprecated field. Use `rollout_id_in` instead.

        Returns:
-            A list of matching rollouts. Ordering is backend-defined but must be deterministic.
+            A sequence of matching rollouts (or [`AttemptedRollout`][agentlightning.AttemptedRollout]
+            when attempts exist). Ordering is deterministic when `sort_by` is set.
+            The return value is not guaranteed to be a list.

        Raises:
            NotImplementedError: Subclasses must implement the query.
        """
        raise NotImplementedError()

-    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+    async def query_attempts(
+        self,
+        rollout_id: str,
+        *,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Attempt]:
        """Return every attempt ever created for `rollout_id` in ascending sequence order.

+        The parameters allow callers to re-order or paginate the attempts so that
+        large retry histories can be streamed lazily.
+
        Args:
            rollout_id: Identifier of the rollout being inspected.
+            sort_by: Field to sort by. Must be a numeric or string field of
+                [`Attempt`][agentlightning.Attempt]. Defaults to `sequence_id` (oldest first).
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.

        Returns:
-            Attempts sorted by `sequence_id` (oldest first). Returns an empty list when none exist.
+            Sequence of Attempts. Returns an empty sequence when none exist.
+            The return value is not guaranteed to be a list.

        Raises:
            NotImplementedError: Subclasses must implement the query.
@@ -307,11 +491,35 @@ class LightningStore:
        """
        raise NotImplementedError()

-    async def query_resources(self) -> List[ResourcesUpdate]:
+    async def query_resources(
+        self,
+        *,
+        resources_id: Optional[str] = None,
+        resources_id_contains: Optional[str] = None,
+        # Filter logic is not supported here because I can't see why it's needed.
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[ResourcesUpdate]:
        """List every stored resource snapshot in insertion order.

+        Supports lightweight filtering, sorting, and pagination for embedding in
+        dashboards.
+
+        Args:
+            resources_id: Optional identifier of the resources to include.
+            resources_id_contains: Optional substring match for resources identifiers.
+            sort_by: Optional field to sort by (must be numeric or string on
+                [`ResourcesUpdate`][agentlightning.ResourcesUpdate]).
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+
        Returns:
-            A chronological list of [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
+            By default, resources are sorted in a deterministic but undefined order.
+            The return value is not guaranteed to be a list.

        Raises:
            NotImplementedError: Subclasses must implement retrieval.
@@ -368,6 +576,20 @@ class LightningStore:
        """
        raise NotImplementedError()

+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        """Bulk allocate the next strictly increasing sequence number used to order spans.
+
+        Implementations may delegate to [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
+        for each rollout and attempt.
+
+        Args:
+            rollout_attempt_ids: List of tuples of rollout and attempt identifiers.
+
+        Returns:
+            List of sequence numbers.
+        """
+        raise NotImplementedError()
+
    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
        """Block until the targeted rollouts reach a terminal status or the timeout expires.

@@ -394,19 +616,61 @@ class LightningStore:
        """
        raise NotImplementedError()

-    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+        *,
+        # Filtering
+        trace_id: Optional[str] = None,
+        trace_id_contains: Optional[str] = None,
+        span_id: Optional[str] = None,
+        span_id_contains: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        parent_id_contains: Optional[str] = None,
+        name: Optional[str] = None,
+        name_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        # Pagination
+        limit: int = -1,
+        offset: int = 0,
+        # Sorting
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+    ) -> Sequence[Span]:
        """Return the stored spans for a rollout, optionally scoped to one attempt.

-        Spans must be returned in ascending `sequence_id` order. Implementations may raise
-        a `RuntimeError` when spans were evicted or expired.
+        Supports a handful of filters that cover the most common debugging
+        scenarios (matching `trace_id`/`span_id`/`parent_id` or substring
+        matches on the span name). `attempt_id="latest"` acts as a convenience
+        that resolves the most recent attempt before evaluating filters. When
+        `attempt_id=None`, spans across every attempt are eligible. By default
+        results are sorted by `sequence_id` (oldest first). Implementations may
+        raise a `RuntimeError` when spans were evicted or expired.

        Args:
            rollout_id: Identifier of the rollout being inspected.
            attempt_id: Attempt identifier to filter by. Pass `"latest"` to retrieve only the
                most recent attempt, or `None` to return all spans across attempts.
+            trace_id: Optional trace ID to filter by.
+            trace_id_contains: Optional substring match for trace IDs.
+            span_id: Optional span ID to filter by.
+            span_id_contains: Optional substring match for span IDs.
+            parent_id: Optional parent span ID to filter by.
+            parent_id_contains: Optional substring match for parent span IDs.
+            name: Optional span name to filter by.
+            name_contains: Optional substring match for span names.
+            filter_logic: Logical operator to combine the optional filters above.
+                The `rollout_id` argument is always applied with AND semantics.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+            sort_by: Field to sort by. Must be a numeric or string field of
+                [`Span`][agentlightning.Span].
+            sort_order: Order to sort by.

        Returns:
            An ordered list of spans (possibly empty).
+            The return value is not guaranteed to be a list.

        Raises:
            NotImplementedError: Subclasses must implement the query.
@@ -502,12 +766,19 @@ class LightningStore:

        When `attempt_id` is `"latest"` the update must target the attempt with the highest
        `sequence_id`; otherwise it must target the specific attempt. Implementations should
-        propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
+        propagate status changes to the rollout (for example
+        via [`rollout_status_from_attempt()`][agentlightning.store.utils.rollout_status_from_attempt])
        once the latest attempt transitions to a terminal state.

        Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],
        parameters also default to the sentinel [`UNSET`][agentlightning.store.base.UNSET].

+        If `worker_id` is present, the worker status will be updated following the rules:
+
+        1. If attempt status is "succeeded" or "failed", the corresponding worker status will be set to "idle".
+        2. If attempt status is "unresponsive" or "timeout", the corresponding worker status will be set to "unknown".
+        3. Otherwise, the worker status will be set to "busy".
+
        Args:
            rollout_id: Identifier of the rollout whose attempt will be updated.
            attempt_id: Attempt identifier or `"latest"` as a convenience.
@@ -524,3 +795,64 @@ class LightningStore:
            ValueError: Implementations must raise when the rollout or attempt is unknown.
        """
        raise NotImplementedError()
+
+    async def query_workers(
+        self,
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
+        """Query all workers in the system.
+
+        Args:
+            status_in: Optional whitelist of [`WorkerStatus`][agentlightning.WorkerStatus] values.
+            worker_id_contains: Optional substring match for worker identifiers.
+            filter_logic: Logical operator to combine the optional filters above.
+            sort_by: Field to sort by. Must be a numeric or string field of [`Worker`][agentlightning.Worker].
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+
+        Returns:
+            Sequence of Workers. Returns an empty sequence when none exist.
+            The return value is not guaranteed to be a list.
+        """
+        raise NotImplementedError()
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        """Retrieve a single worker by identifier.
+
+        Args:
+            worker_id: Identifier of the worker.
+
+        Returns:
+            The worker record if it exists, otherwise `None`.
+
+        Raises:
+            NotImplementedError: Subclasses must implement lookup semantics.
+        """
+        raise NotImplementedError()
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        """Record a heartbeat for `worker_id` and refresh telemetry.
+
+        Implementations must treat this API as heartbeat-only: it should snapshot
+        the latest stats when provided, stamp `last_heartbeat_time` with the
+        current wall clock, and rely on other store mutations (`dequeue_rollout`,
+        `update_attempt`, etc.) to drive the worker's busy/idle status,
+        assignment, and activity timestamps.
+
+        Args:
+            worker_id: Identifier of the worker to update.
+            heartbeat_stats: Replacement worker heartbeat statistics (non-null when provided).
+        """
+        raise NotImplementedError()
@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import (
+    AtomicLabels,
+    AtomicMode,
+    Collection,
+    FilterOptions,
+    KeyValue,
+    LightningCollections,
+    PaginatedResult,
+    Queue,
+    SortOptions,
+)
+from .memory import DequeBasedQueue, DictBasedKeyValue, InMemoryLightningCollections, ListBasedCollection
+
+__all__ = [
+    "AtomicLabels",
+    "AtomicMode",
+    "Collection",
+    "Queue",
+    "KeyValue",
+    "FilterOptions",
+    "SortOptions",
+    "PaginatedResult",
+    "LightningCollections",
+    "ListBasedCollection",
+    "DequeBasedQueue",
+    "DictBasedKeyValue",
+    "InMemoryLightningCollections",
+]
@@ -0,0 +1,587 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import time
+from contextlib import asynccontextmanager
+from numbers import Real
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    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
+
+from agentlightning.types import (
+    Attempt,
+    FilterField,
+    FilterOptions,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    SortOptions,
+    Span,
+    Worker,
+)
+
+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", "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.
+"""
+
+
+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 get_current_store_methods
+
+            # Enable tracking
+            start_time = time.perf_counter()
+            status: str = "OK"
+            public_store_method, private_store_method = get_current_store_methods()
+            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."""
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the collection."""
+        raise NotImplementedError()
+
+    async def size(self) -> int:
+        """Get the number of items in the collection."""
+        raise NotImplementedError()
+
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with the given filters, sort order, and pagination.
+
+        Args:
+            filter:
+                The filters to apply to the collection. See [`FilterOptions`][agentlightning.FilterOptions].
+
+            sort:
+                The options for sorting the collection. See [`SortOptions`][agentlightning.SortOptions].
+                The field must exist in the model. If field might contain null values, in which case the behavior is undefined
+                (i.e., depending on the implementation).
+
+            limit:
+                Max number of items to return. Use -1 for "no limit".
+
+            offset:
+                Number of items to skip from the start of the *matching* items.
+
+        Returns:
+            PaginatedResult with items, limit, offset, and total matched items.
+        """
+        raise NotImplementedError()
+
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Get the first item that matches the given filters.
+
+        Args:
+            filter: The filters to apply to the collection.
+                See [`FilterOptions`][agentlightning.store.collection.FilterOptions].
+            sort: Sort options. See [`SortOptions`][agentlightning.store.collection.SortOptions].
+
+        Returns:
+            The first item that matches the given filters, or None if no item matches.
+        """
+        raise NotImplementedError()
+
+    async def insert(self, items: Sequence[T]) -> None:
+        """Add the given items to the collection.
+
+        Raises:
+            ValueError: If an item with the same primary key already exists.
+        """
+        raise NotImplementedError()
+
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+        """Update the given items in the collection.
+
+        Args:
+            items: The items to update in the collection.
+            update_fields: The fields to update. If not provided, all fields in the type will be updated.
+                Only applicable if the item type is a Pydantic BaseModel.
+
+        Raises:
+            ValueError: If an item with the primary keys does not exist.
+
+        Returns:
+            The items that were updated.
+        """
+        raise NotImplementedError()
+
+    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+        """Upsert the given items into the collection.
+
+        If the items with the same primary keys already exist, they will be updated.
+        Otherwise, they will be inserted.
+
+        The operation has three semantics configurable via `update_fields`:
+
+        - `update_or_insert` via `collection.upsert(items, update_fields=["status", "updated_at"])`.
+          If the item with the same primary keys already exists, only the specified fields will be updated.
+          Otherwise, the item will be inserted.
+        - `get_or_insert` via `collection.upsert(items, update_fields=[])`.
+          If the item with the same primary keys already exists, the item will be left unchanged.
+          Otherwise, the item will be inserted.
+        - `replace_ish` via `collection.upsert(items)`.
+          If the item with the same primary keys already exists, all fields from the item will be set.
+          Otherwise, the item will be inserted.
+
+        Returns:
+            The items that were upserted.
+        """
+        raise NotImplementedError()
+
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items from the collection.
+
+        Args:
+            items: The items to delete from the collection.
+
+        Raises:
+            ValueError: If the items with the primary keys to be deleted do not exist.
+        """
+        raise NotImplementedError()
+
+
+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:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the queue."""
+        raise NotImplementedError()
+
+    async def has(self, item: T) -> bool:
+        """Check if the given item is in the queue."""
+        raise NotImplementedError()
+
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        """Append the given items to the end of the queue.
+
+        Args:
+            items: The items to append to the end of the queue.
+
+        Returns:
+            The items that were appended to the end of the queue.
+        """
+        raise NotImplementedError()
+
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        """Pop the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to pop from the front of the queue.
+
+        Returns:
+            The items that were popped from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        """Peek the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to peek from the front of the queue.
+
+        Returns:
+            The items that were peeked from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    async def size(self) -> int:
+        """Get the number of items in the queue."""
+        raise NotImplementedError()
+
+
+class KeyValue(TrackedCollection, Generic[K, V]):
+    """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"
+
+    async def has(self, key: K) -> bool:
+        """Check if the given key is in the dictionary."""
+        raise NotImplementedError()
+
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        """Get the value for the given key, or the default value if the key is not found."""
+        raise NotImplementedError()
+
+    async def set(self, key: K, value: V) -> None:
+        """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()
+
+    async def size(self) -> int:
+        """Get the number of items in the dictionary."""
+        raise NotImplementedError()
+
+
+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."""
+        raise NotImplementedError()
+
+    @property
+    def attempts(self) -> Collection[Attempt]:
+        """Collections of attempts."""
+        raise NotImplementedError()
+
+    @property
+    def spans(self) -> Collection[Span]:
+        """Collections of spans."""
+        raise NotImplementedError()
+
+    @property
+    def resources(self) -> Collection[ResourcesUpdate]:
+        """Collections of resources."""
+        raise NotImplementedError()
+
+    @property
+    def workers(self) -> Collection[Worker]:
+        """Collections of workers."""
+        raise NotImplementedError()
+
+    @property
+    def rollout_queue(self) -> Queue[str]:
+        """Queue of rollouts (tasks)."""
+        raise NotImplementedError()
+
+    @property
+    def span_sequence_ids(self) -> KeyValue[str, int]:
+        """Dictionary (counter) of span sequence IDs."""
+        raise NotImplementedError()
+
+    def atomic(
+        self,
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
+    ) -> AsyncContextManager[Self]:
+        """Perform a atomic operation on the collections.
+
+        Subclass may use args and kwargs to support multiple levels of atomicity.
+        The arguments can be seen as tags. They only imply the behavior of the operation, not the implementation.
+
+        Args:
+            mode: The mode of atomicity. See [`AtomicMode`][agentlightning.store.collection.AtomicMode].
+            snapshot: Enable read snapshot for repeatable reads. Data consistency is guaranteed. The real behavior is implementation-dependent.
+            commit: Enable commitment for write operations. Unsuccessful operations will be rolled back depending on the implementation.
+                Recommend to use [`execute()`][agentlightning.store.collection.LightningCollections.execute] for this level to enable automatic retries.
+                Remember that the real behavior is implementation-dependent.
+            labels: Labels to add to the atomic operation (commonly used as lock names or collection names).
+            **kwargs: Keyword arguments to pass to the operation.
+        """
+        raise NotImplementedError()
+
+    async def execute(
+        self,
+        callback: Callable[[Self], Awaitable[T]],
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
+    ) -> T:
+        """Execute the given callback within an atomic operation. Retry on transient errors is implied.
+
+        See [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] for more details.
+        """
+        async with self.atomic(mode=mode, snapshot=snapshot, commit=commit, labels=labels, **kwargs) as collections:
+            return await callback(collections)
+
+
+FilterMap = Mapping[str, FilterField]
+
+
+def merge_must_filters(target: MutableMapping[str, FilterField], definition: Any) -> None:
+    """Normalize a `_must` filter group into the provided mapping.
+
+    Mainly for validation purposes.
+    """
+    if definition is None:
+        return
+
+    entries: List[Mapping[str, FilterField]] = []
+    if isinstance(definition, Mapping):
+        entries.append(cast(Mapping[str, FilterField], definition))
+    elif isinstance(definition, Sequence) and not isinstance(definition, (str, bytes)):
+        for entry in definition:  # type: ignore
+            if not isinstance(entry, Mapping):
+                raise TypeError("Each `_must` entry must be a mapping of field names to operators")
+            entries.append(cast(Mapping[str, FilterField], entry))
+    else:
+        raise TypeError("`_must` filters must be provided as a mapping or sequence of mappings")
+
+    for entry in entries:
+        for field_name, ops in entry.items():
+            existing = target.get(field_name, {})
+            merged_ops: Dict[str, Any] = dict(existing)
+            for op_name, expected in ops.items():
+                if op_name in merged_ops:
+                    raise ValueError(f"Duplicate operator '{op_name}' for field '{field_name}' in must filters")
+                merged_ops[op_name] = expected
+            target[field_name] = cast(FilterField, merged_ops)
+
+
+def normalize_filter_options(
+    filter_options: Optional[FilterOptions],
+) -> Tuple[Optional[FilterMap], Optional[FilterMap], Literal["and", "or"]]:
+    """Convert FilterOptions to the internal structure and resolve aggregate logic."""
+    if not filter_options:
+        return None, None, "and"
+
+    aggregate = cast(Literal["and", "or"], filter_options.get("_aggregate", "and"))
+    if aggregate not in ("and", "or"):
+        raise ValueError(f"Unsupported filter aggregate '{aggregate}'")
+
+    # Extract normalized filters and must filters from the filter options.
+    normalized: Dict[str, FilterField] = {}
+    must_filters: Dict[str, FilterField] = {}
+    for field_name, ops in filter_options.items():
+        if field_name == "_aggregate":
+            continue
+        if field_name == "_must":
+            merge_must_filters(must_filters, ops)
+            continue
+        normalized[field_name] = cast(FilterField, dict(ops))  # type: ignore
+
+    return (normalized or None, must_filters or None, aggregate)
+
+
+def resolve_sort_options(sort: Optional[SortOptions]) -> Tuple[Optional[str], Literal["asc", "desc"]]:
+    """Extract sort field/order from the caller-provided SortOptions."""
+    if not sort:
+        return None, "asc"
+
+    sort_name = sort.get("name")
+    if not sort_name:
+        raise ValueError("Sort options must include a 'name' field")
+
+    sort_order = sort.get("order", "asc")
+    if sort_order not in ("asc", "desc"):
+        raise ValueError(f"Unsupported sort order '{sort_order}'")
+
+    return sort_name, sort_order
@@ -0,0 +1,970 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+import weakref
+from collections import deque
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import (
+    Any,
+    Deque,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+import aiologic
+from pydantic import BaseModel
+
+from agentlightning.types import (
+    Attempt,
+    FilterField,
+    FilterOptions,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    SortOptions,
+    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
+K = TypeVar("K")
+V = TypeVar("V")
+
+logger = logging.getLogger(__name__)
+
+# Nested structure type:
+# dict[pk1] -> dict[pk2] -> ... -> item
+ListBasedCollectionItemType = Union[
+    Dict[Any, "ListBasedCollectionItemType[T]"],  # intermediate node
+    Dict[Any, T],  # leaf node dictionary
+]
+
+MutationMode = Literal["insert", "update", "upsert", "delete"]
+
+
+def _item_matches_filters(
+    item: object,
+    filters: Optional[FilterMap],
+    filter_logic: Literal["and", "or"],
+    must_filters: Optional[FilterMap] = None,
+) -> bool:
+    """Check whether an item matches the provided filter definition.
+
+    Filter format:
+
+    ```json
+    {
+        "_aggregate": "or",
+        "field_name": {
+            "exact": <value>,
+            "within": <iterable_of_allowed_values>,
+            "contains": <substring_or_element>,
+        },
+        ...
+    }
+    ```
+
+    Operators within the same field are stored in a unified pool and combined using
+    a universal logical operator.
+    """
+    if must_filters and not _item_matches_filters(item, must_filters, "and"):
+        return False
+
+    if not filters:
+        return True
+
+    all_conditions_match: List[bool] = []
+
+    for field_name, ops in filters.items():
+        item_value = getattr(item, field_name, None)
+
+        for op_name, expected in ops.items():
+            # Ignore no-op filters
+            if expected is None:
+                continue
+
+            if op_name == "exact":
+                all_conditions_match.append(item_value == expected)
+
+            elif op_name == "within":
+                try:
+                    all_conditions_match.append(item_value in expected)  # type: ignore[arg-type]
+                except TypeError:
+                    all_conditions_match.append(False)
+
+            elif op_name == "contains":
+                if item_value is None:
+                    all_conditions_match.append(False)
+                elif isinstance(item_value, str) and isinstance(expected, str):
+                    all_conditions_match.append(expected in item_value)
+                else:
+                    # Fallback: treat as generic iterable containment.
+                    try:
+                        all_conditions_match.append(expected in item_value)  # type: ignore[arg-type]
+                    except TypeError:
+                        all_conditions_match.append(False)
+            else:
+                raise ValueError(f"Unsupported filter operator '{op_name}' for field '{field_name}'")
+
+    return all(all_conditions_match) if filter_logic == "and" else any(all_conditions_match)
+
+
+def _get_sort_value(item: object, sort_by: str) -> Any:
+    """Get a sort key for the given item/field.
+
+    - If the field name ends with '_time', values are treated as comparable timestamps.
+    - For other fields we try to infer a safe default from the Pydantic model annotation.
+    """
+    value = getattr(item, sort_by, None)
+
+    if sort_by.endswith("_time"):
+        # For *_time fields, push missing values to the end.
+        return float("inf") if value is None else value
+
+    if value is None:
+        # Introspect model field type to choose a reasonable default for None.
+        model_fields = getattr(item.__class__, "model_fields", {})
+        if sort_by not in model_fields:
+            raise ValueError(
+                f"Failed to sort items by '{sort_by}': field does not exist " f"on {item.__class__.__name__}"
+            )
+
+        field_type_str = str(model_fields[sort_by].annotation)
+        if "str" in field_type_str or "Literal" in field_type_str:
+            return ""
+        if "int" in field_type_str:
+            return 0
+        if "float" in field_type_str:
+            return 0.0
+        raise ValueError(f"Failed to sort items by '{sort_by}': unsupported field type {field_type_str!r}")
+
+    return value
+
+
+class ListBasedCollection(Collection[T]):
+    """In-memory implementation of Collection using a nested dict for O(1) primary-key lookup.
+
+    The internal structure is:
+
+        {
+            pk1_value: {
+                pk2_value: {
+                    ...
+                        pkN_value: item
+                }
+            }
+        }
+
+    where the nesting depth equals the number of primary keys.
+
+    Sorting behavior:
+
+    1. If no sort_by is provided, the items are returned in the order of insertion.
+    2. If sort_by is provided, the items are sorted by the value of the sort_by field.
+    3. If the sort_by field is a timestamp, the null values are treated as infinity.
+    4. If the sort_by field is not a timestamp, the null values are treated as empty string
+       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],
+        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):
+            raise TypeError(f"Expect item to be not a dict, got {item_type.__name__}")
+        self._item_type: Type[T] = item_type
+        self._primary_keys: Tuple[str, ...] = tuple(primary_keys)
+
+        # Pre-populate the collection with the given items.
+        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
+
+    def item_type(self) -> Type[T]:
+        """Return the Pydantic model type of items stored in this collection."""
+        return self._item_type
+
+    async def size(self) -> int:
+        """Return the number of items stored in the collection."""
+        return self._size
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self._size})>"
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    def _ensure_item_type(self, item: T) -> None:
+        """Validate that the item matches the declared item_type."""
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, " f"got {type(item).__name__}")
+
+    def _extract_primary_key_values(self, item: T) -> Tuple[Any, ...]:
+        """Extract the primary key values from an item.
+
+        Raises:
+            ValueError: If any primary key is missing on the item.
+        """
+        values: List[Any] = []
+        for key in self._primary_keys:
+            if not hasattr(item, key):
+                raise ValueError(f"Item {item} does not have primary key field '{key}'")
+            values.append(getattr(item, key))
+        return tuple(values)
+
+    def _render_key_values(self, key_values: Sequence[Any]) -> str:
+        return ", ".join(f"{name}={value!r}" for name, value in zip(self._primary_keys, key_values))
+
+    def _locate_node(
+        self,
+        key_values: Sequence[Any],
+        create_missing: bool,
+    ) -> Tuple[MutableMapping[Any, Any], Any]:
+        """Locate the parent mapping and final key for an item path.
+
+        Args:
+            key_values: The sequence of primary key values.
+            create_missing: Whether to create intermediate dictionaries as needed.
+
+        Returns:
+            (parent_mapping, final_key)
+
+        Raises:
+            KeyError: If the path does not exist and create_missing is False.
+            ValueError: If the internal structure is corrupted (non-dict where dict is expected).
+        """
+        if not key_values:
+            raise ValueError("key_values must be non-empty")
+
+        current: MutableMapping[Any, Any] = self._items
+        for idx, value in enumerate(key_values):
+            is_last = idx == len(key_values) - 1
+            if is_last:
+                # At the final level, current[value] is the item (or will be).
+                return current, value  # type: ignore
+
+            # Intermediate level: current[value] must be a dict.
+            if value not in current:
+                if not create_missing:
+                    raise KeyError(f"Path does not exist for given primary keys: {self._render_key_values(key_values)}")
+                current[value] = {}
+            next_node = current[value]  # type: ignore
+            if not isinstance(next_node, dict):
+                raise ValueError(f"Internal structure corrupted: expected dict, got {type(next_node)!r}")  # type: ignore
+            current = next_node  # type: ignore
+
+        # We should always return inside the loop.
+        raise RuntimeError("Unreachable")
+
+    def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
+        """Core mutation logic shared by insert, update, upsert, and delete."""
+        self._ensure_item_type(item)
+        key_values = self._extract_primary_key_values(item)
+
+        if mode in ("insert", "upsert"):
+            parent, final_key = self._locate_node(key_values, create_missing=True)
+            exists = final_key in parent
+
+            if mode == "insert":
+                if exists:
+                    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
+                if not exists:
+                    self._size += 1
+                    parent[final_key] = item
+
+                elif update_fields is None:
+                    # update_or_insert: update all fields
+                    parent[final_key] = item
+
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+
+                    # Try to fetch the existing item
+                    existing = parent[final_key]
+                    if not isinstance(existing, self._item_type):
+                        raise ValueError(
+                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
+                        )
+
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+
+            return parent[final_key]
+
+        elif mode in ("update", "delete"):
+            # For update/delete we must not create missing paths.
+            try:
+                parent, final_key = self._locate_node(key_values, create_missing=False)
+            except KeyError:
+                raise ValueError(
+                    f"Item does not exist with primary key(s): {self._render_key_values(key_values)}"
+                ) from None
+
+            if final_key not in parent:
+                raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
+
+            if mode == "update":
+                if update_fields is None:
+                    # replace the entire item
+                    parent[final_key] = item
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+                return parent[final_key]
+            else:  # delete
+                del parent[final_key]
+                self._size -= 1
+        else:
+            raise ValueError(f"Unknown mutation mode: {mode}")
+
+    def _iter_items(
+        self,
+        root: Optional[Mapping[Any, Any]] = None,
+        filters: Optional[FilterMap] = None,
+        must_filters: Optional[FilterMap] = None,
+        filter_logic: Literal["and", "or"] = "and",
+    ) -> Iterable[T]:
+        """Iterate over all items in the nested dictionary structure, optionally applying filters."""
+        if root is None:
+            root = self._items
+        if not root:
+            return
+        stack: List[Mapping[Any, Any]] = [root]
+        while stack:
+            node = stack.pop()
+            for value in node.values():
+                # Leaf nodes contain items; intermediate nodes are dicts.
+                if isinstance(value, self._item_type):
+                    if _item_matches_filters(value, filters, filter_logic, must_filters):
+                        yield value
+                elif isinstance(value, dict):
+                    stack.append(value)  # type: ignore
+                else:
+                    raise ValueError(
+                        f"Internal structure corrupted: expected dict or {self._item_type.__name__}, "
+                        f"got {type(value)!r}"
+                    )
+
+    def _iter_matching_items(
+        self,
+        filters: Optional[FilterMap],
+        must_filters: Optional[FilterMap],
+        filter_logic: Literal["and", "or"],
+    ) -> Iterable[T]:
+        """Efficiently iterate over items matching filters, using primary-key prefix when possible."""
+        # Fast path: when optional filters can't form a prefix, fall back to scanning.
+        if filter_logic != "and" and must_filters is None:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        # Try to derive a primary-key prefix from exact filters.
+        pk_values_prefix: List[Any] = []
+        prefix_sources: List[FilterMap] = []
+        if must_filters:
+            prefix_sources.append(must_filters)
+        if filter_logic == "and" and filters:
+            prefix_sources.append(filters)
+
+        for pk in self._primary_keys:
+            # combined_ops are: [{"exact": value}, {"within": [...]}, ...]
+            combined_ops: List[FilterField] = []
+            for source in prefix_sources:
+                field_ops = source.get(pk)  # type: ignore[union-attr]
+                if field_ops:
+                    combined_ops.append(field_ops)
+            if not combined_ops:
+                break
+            # Only allow a pure {"exact": value} constraint.
+            exact_value: Any | None = None
+            allow_prefix = True
+            for ops in combined_ops:
+                if set(ops.keys()) != {"exact"}:
+                    allow_prefix = False
+                    break
+                candidate = ops.get("exact")
+                if candidate is None:
+                    allow_prefix = False
+                    break
+                if exact_value is not None and candidate != exact_value:
+                    # Contradictory exact filters mean no items can match.
+                    logger.warning(f"Contradictory exact filters for field '{pk}': {exact_value} != {candidate}")
+                    return ()
+                exact_value = candidate
+
+            if not allow_prefix:
+                break
+
+            value = exact_value
+            if value is None:
+                break
+            pk_values_prefix.append(value)
+
+        if not pk_values_prefix:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        try:
+            if len(pk_values_prefix) == len(self._primary_keys):
+                # All primary keys specified -> at most a single item.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                single_item = parent.get(final_key)
+                if isinstance(single_item, self._item_type) and _item_matches_filters(
+                    single_item,
+                    filters,
+                    filter_logic,
+                    must_filters,
+                ):
+                    return (single_item,)
+                return ()
+            else:
+                # Prefix of primary keys specified -> iterate only the subtree below that prefix.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                subtree = parent.get(final_key)
+                if isinstance(subtree, dict):
+                    return self._iter_items(
+                        subtree,  # type: ignore
+                        filters=filters,
+                        must_filters=must_filters,
+                        filter_logic=filter_logic,
+                    )
+                return ()
+        except KeyError:
+            # No items exist for this primary-key prefix.
+            return ()
+
+    @tracked("query")
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with filters, sort order, and pagination.
+
+        Args:
+            filter: Mapping of field name to operator dict along with the optional `_aggregate` logic.
+            sort: Options describing which field to sort by and in which order.
+            limit: Max number of items to return. Use -1 for "no limit".
+            offset: Number of items to skip from the start of the *matching* items.
+        """
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        # No sorting: stream through items and apply pagination on the fly.
+        if not sort_by:
+            matched_items: List[T] = []
+            total_matched = 0
+
+            for item in items_iter:
+                # Count every match for 'total'
+                total_matched += 1
+
+                # Apply offset/limit window
+                if total_matched <= offset:
+                    continue
+                if limit != -1 and len(matched_items) >= limit:
+                    # Still need to finish iteration to get accurate total_matched.
+                    continue
+
+                matched_items.append(item)
+
+            return PaginatedResult(
+                items=matched_items,
+                limit=limit,
+                offset=offset,
+                total=total_matched,
+            )
+
+        # With sorting: we must materialize all matching items to sort them.
+        all_matches: List[T] = list(items_iter)
+
+        total_matched = len(all_matches)
+        reverse = sort_order == "desc"
+        all_matches.sort(key=lambda x: _get_sort_value(x, sort_by), reverse=reverse)
+
+        if limit == -1:
+            paginated_items = all_matches[offset:]
+        else:
+            paginated_items = all_matches[offset : offset + limit]
+
+        return PaginatedResult(
+            items=paginated_items,
+            limit=limit,
+            offset=offset,
+            total=total_matched,
+        )
+
+    @tracked("get")
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Return the first (or best-sorted) item that matches the given filters, or None."""
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        if not sort_by:
+            # Just return the first matching item, if any.
+            for item in items_iter:
+                return item
+            return None
+
+        # Single-pass min/max according to sort_order.
+        best_item: Optional[T] = None
+        best_key: Any = None
+
+        for item in items_iter:
+            key = _get_sort_value(item, sort_by)
+            if best_item is None:
+                best_item = item
+                best_key = key
+                continue
+
+            if sort_order == "asc":
+                if key < best_key:
+                    best_item, best_key = item, key
+            else:
+                if key > best_key:
+                    best_item, best_key = item, key
+
+        return best_item
+
+    @tracked("insert")
+    async def insert(self, items: Sequence[T]) -> None:
+        """Insert the given items.
+
+        Raises:
+            DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
+        """
+        seen_keys: set[Tuple[Any, ...]] = set()
+        prepared: List[T] = []
+        for item in items:
+            self._ensure_item_type(item)
+            key_values = self._extract_primary_key_values(item)
+            if key_values in seen_keys:
+                raise DuplicatedPrimaryKeyError(
+                    f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
+                )
+            seen_keys.add(key_values)
+            prepared.append(item)
+
+        for item in prepared:
+            self._mutate_single(item, mode="insert")
+
+    @tracked("update")
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+        """Update the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        updated_items: List[T] = []
+        for item in items:
+            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
+            if updated is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            updated_items.append(updated)
+        return updated_items
+
+    @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] = []
+        for item in items:
+            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
+            if upserted is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            upserted_items.append(upserted)
+        return upserted_items
+
+    @tracked("delete")
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        # We use a two-phase approach to avoid partial deletion if one fails:
+        # first compute key_values to validate, then perform deletions.
+        for item in items:
+            # _mutate_single will validate existence and update size.
+            self._mutate_single(item, mode="delete")
+
+
+class DequeBasedQueue(Queue[T]):
+    """Queue implementation backed by collections.deque.
+
+    Provides O(1) amortized enqueue (append) and dequeue (popleft).
+    """
+
+    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):
+                raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+            self._items.append(item)
+        return items
+
+    @tracked("dequeue")
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        out: List[T] = []
+        for _ in range(min(limit, len(self._items))):
+            out.append(self._items.popleft())
+        return out
+
+    @tracked("peek")
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        result: List[T] = []
+        count = min(limit, len(self._items))
+        for idx, item in enumerate(self._items):
+            if idx >= count:
+                break
+            result.append(item)
+        return result
+
+    @tracked("size")
+    async def size(self) -> int:
+        return len(self._items)
+
+
+class DictBasedKeyValue(KeyValue[K, V]):
+    """KeyValue implementation backed by a plain dictionary."""
+
+    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)
+
+
+class InMemoryLightningCollections(LightningCollections):
+    """In-memory implementation of LightningCollections using Python data structures.
+
+    Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+    """
+
+    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(),
+            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+        }
+        self._rollouts = ListBasedCollection(
+            items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
+        )
+        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
+
+    @property
+    def collection_name(self) -> str:
+        return "router"
+
+    @property
+    def rollouts(self) -> ListBasedCollection[Rollout]:
+        return self._rollouts
+
+    @property
+    def attempts(self) -> ListBasedCollection[Attempt]:
+        return self._attempts
+
+    @property
+    def spans(self) -> ListBasedCollection[Span]:
+        return self._spans
+
+    @property
+    def resources(self) -> ListBasedCollection[ResourcesUpdate]:
+        return self._resources
+
+    @property
+    def workers(self) -> ListBasedCollection[Worker]:
+        return self._workers
+
+    @property
+    def rollout_queue(self) -> DequeBasedQueue[str]:
+        return self._rollout_queue
+
+    @property
+    def span_sequence_ids(self) -> DictBasedKeyValue[str, int]:
+        return self._span_sequence_ids
+
+    @asynccontextmanager
+    async def atomic(
+        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.
+
+        Skip the locking if mode is "r" and snapshot is False.
+
+        This collection implementation does NOT support rollback / commit.
+        """
+        if mode == "r" and not snapshot:
+            yield self
+            return
+        if not labels:
+            # If no labels are provided, use all locks.
+            labels = list(self._lock.keys())
+
+        # IMPORTANT: Sort the labels to ensure consistent locking order.
+        # This is necessary to avoid deadlocks when multiple threads/coroutines
+        # are trying to acquire the same locks in different orders.
+        labels = sorted(labels)
+
+        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.
+
+        Uses private API for efficiency.
+        """
+        self._spans._items.pop(rollout_id, [])  # pyright: ignore[reportPrivateUsage]
+
+
+class _LoopAwareAsyncLock:
+    """Async lock that transparently rebinds to the current event loop.
+
+    The lock intentionally remains *thread-unsafe*: callers must only use it from
+    one thread at a time. If multiple threads interact with the store, each
+    thread gets its own event loop specific lock.
+    """
+
+    def __init__(self) -> None:
+        self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+    # When serializing and deserializing, we don't need to serialize the locks.
+    # Because another process will have its own set of event loops and its own lock.
+    def __getstate__(self) -> dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self._locks = weakref.WeakKeyDictionary()
+
+    def _get_lock_for_current_loop(self) -> asyncio.Lock:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[loop] = lock
+        return lock
+
+    async def __aenter__(self) -> asyncio.Lock:
+        lock = self._get_lock_for_current_loop()
+        await lock.acquire()
+        return lock
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None or not lock.locked():
+            raise RuntimeError("Lock released without being acquired")
+        lock.release()
+
+
+class _ThreadSafeAsyncLock:
+    """A thread lock powered by aiologic that can be used in both async and sync contexts.
+
+    aiologic claims itself to be a thread-safe asyncio lock.
+    """
+
+    def __init__(self):
+        self._lock = aiologic.Lock()
+
+    async def __aenter__(self):
+        await self._lock.async_acquire()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any):
+        # .release() is non-blocking, so we can call it directly
+        self._lock.async_release()
@@ -0,0 +1,165 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+import time
+import uuid
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TypeVar, Union
+
+from agentlightning.types import Attempt, AttemptedRollout, Rollout
+from agentlightning.utils.metrics import MetricsBackend
+
+from .base import LightningStoreCapabilities, is_finished
+from .collection.mongo import MongoClientPool, MongoLightningCollections
+from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_partition_id() -> str:
+    return "pt-" + hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+
+
+class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollections]):
+    """
+    MongoDB implementation of LightningStore using MongoDB collections.
+    Data is persistent and can be shared between multiple processes.
+
+    Args:
+        mongo_uri: MongoDB connection string (defaults to local replica set).
+        mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
+        database_name: 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,
+        *,
+        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,
+        tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
+    ) -> None:
+        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'")
+
+        if partition_id is None:
+            partition_id = _generate_partition_id()
+            logger.info("No partition id provided, generated a new one: %s", partition_id)
+
+        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,
+                tracker=tracker,
+            ),
+            tracker=tracker,
+            scan_debounce_seconds=scan_debounce_seconds,
+        )
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            thread_safe=True,
+            async_safe=True,
+            zero_copy=True,
+            otlp_traces=False,
+        )
+
+    async def close(self) -> None:
+        """Close the store by closing the client pool."""
+        await self._client_pool.close()
+
+    @tracked("wait_for_rollouts")
+    @healthcheck_before
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """Wait for specified rollouts to complete with a timeout.
+
+        Concurrently wait for all rollouts to complete with a timeout.
+        """
+        start_time = time.time()
+        current_time = start_time
+        deadline = start_time + timeout if timeout is not None else None
+
+        finished_rollouts: Dict[str, Rollout] = {}
+        unfinished_rollout_ids = set(rollout_ids)
+
+        while deadline is None or current_time <= deadline:
+            async with self.collections.atomic(
+                mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
+            ) as collections:
+                # Query the rollouts that are not finished in a single query
+                rollouts = await collections.rollouts.query(
+                    filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
+                )
+            for rollout in rollouts.items:
+                if is_finished(rollout):
+                    finished_rollouts[rollout.rollout_id] = rollout
+                    unfinished_rollout_ids.remove(rollout.rollout_id)
+
+            if not unfinished_rollout_ids:
+                break
+
+            # Poll every 10 seconds by default
+            # Minus 0.1 to make sure the time is still sufficient for another call
+            rest_time = max(0.01, min(deadline - time.time() - 0.1, 10.0)) if deadline is not None else 10.0
+            await asyncio.sleep(rest_time)
+            current_time = time.time()
+
+        # 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]
+
+    @tracked("_unlocked_many_rollouts_to_attempted_rollouts")
+    async def _unlocked_many_rollouts_to_attempted_rollouts(
+        self, collections: MongoLightningCollections, rollouts: Sequence[Rollout]
+    ) -> List[Union[Rollout, AttemptedRollout]]:
+        """Query the latest attempts for the rollouts, and attach them to the rollout objects."""
+        async with collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["attempts"]) as collections:
+            attempts = await collections.attempts.query(
+                filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
+                sort={"name": "sequence_id", "order": "desc"},
+            )
+        latest_attempts: Dict[str, Attempt] = {}
+        for attempt in attempts:
+            if attempt.rollout_id not in latest_attempts:
+                latest_attempts[attempt.rollout_id] = attempt
+            # Otherwise we ignore the attempt because there's already a newer attempt
+
+        return [
+            (
+                AttemptedRollout(**rollout.model_dump(), attempt=latest_attempts[rollout.rollout_id])
+                if rollout.rollout_id in latest_attempts
+                else rollout
+            )
+            for rollout in rollouts
+        ]
@@ -3,7 +3,7 @@
 from __future__ import annotations

 import threading
-from typing import Any, Dict, List, Literal, Optional, Sequence
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple

 from opentelemetry.sdk.trace import ReadableSpan

@@ -11,6 +11,7 @@ from agentlightning.types import (
    Attempt,
    AttemptedRollout,
    AttemptStatus,
+    EnqueueRolloutRequest,
    NamedResources,
    ResourcesUpdate,
    Rollout,
@@ -18,9 +19,11 @@ from agentlightning.types import (
    RolloutStatus,
    Span,
    TaskInput,
+    Worker,
+    WorkerStatus,
 )

-from .base import UNSET, LightningStore, Unset
+from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset


 class LightningStoreThreaded(LightningStore):
@@ -35,6 +38,21 @@ class LightningStoreThreaded(LightningStore):
        self.store = store
        self._lock = threading.Lock()

+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        capabilities = self.store.capabilities
+        return {
+            **capabilities,
+            "async_safe": True,
+            "thread_safe": True,
+        }
+
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        with self._lock:
+            return await self.store.statistics()
+
    async def start_rollout(
        self,
        input: TaskInput,
@@ -42,9 +60,17 @@ class LightningStoreThreaded(LightningStore):
        resources_id: str | None = None,
        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
    ) -> AttemptedRollout:
        with self._lock:
-            return await self.store.start_rollout(input, mode, resources_id, config, metadata)
+            return await self.store.start_rollout(
+                input,
+                mode,
+                resources_id,
+                config,
+                metadata,
+                worker_id,
+            )

    async def enqueue_rollout(
        self,
@@ -57,26 +83,72 @@ class LightningStoreThreaded(LightningStore):
        with self._lock:
            return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)

-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
        with self._lock:
-            return await self.store.dequeue_rollout()
+            return await self.store.enqueue_many_rollouts(rollouts)

-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
        with self._lock:
-            return await self.store.start_attempt(rollout_id)
+            return await self.store.dequeue_rollout(worker_id=worker_id)
+
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id, worker_id)

    async def query_rollouts(
        self,
        *,
+        status_in: Optional[Sequence[RolloutStatus]] = None,
+        rollout_id_in: Optional[Sequence[str]] = None,
+        rollout_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
        status: Optional[Sequence[RolloutStatus]] = None,
        rollout_ids: Optional[Sequence[str]] = None,
-    ) -> List[Rollout]:
+    ) -> Sequence[Rollout]:
        with self._lock:
-            return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+            return await self.store.query_rollouts(
+                status_in=status_in,
+                rollout_id_in=rollout_id_in,
+                rollout_id_contains=rollout_id_contains,
+                filter_logic=filter_logic,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+                status=status,
+                rollout_ids=rollout_ids,
+            )

-    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+    async def query_attempts(
+        self,
+        rollout_id: str,
+        *,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Attempt]:
        with self._lock:
-            return await self.store.query_attempts(rollout_id)
+            return await self.store.query_attempts(
+                rollout_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )

    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
        with self._lock:
@@ -86,6 +158,26 @@ class LightningStoreThreaded(LightningStore):
        with self._lock:
            return await self.store.get_latest_attempt(rollout_id)

+    async def query_resources(
+        self,
+        *,
+        resources_id: Optional[str] = None,
+        resources_id_contains: Optional[str] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.query_resources(
+                resources_id=resources_id,
+                resources_id_contains=resources_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
        with self._lock:
            return await self.store.add_resources(resources)
@@ -102,7 +194,11 @@ class LightningStoreThreaded(LightningStore):
        with self._lock:
            return await self.store.get_latest_resources()

-    async def add_span(self, span: Span) -> Span:
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        with self._lock:
+            return await self.store.add_many_spans(spans)
+
+    async def add_span(self, span: Span) -> Optional[Span]:
        with self._lock:
            return await self.store.add_span(span)

@@ -112,7 +208,7 @@ class LightningStoreThreaded(LightningStore):
        attempt_id: str,
        readable_span: ReadableSpan,
        sequence_id: int | None = None,
-    ) -> Span:
+    ) -> Optional[Span]:
        with self._lock:
            return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)

@@ -124,13 +220,47 @@ class LightningStoreThreaded(LightningStore):
        with self._lock:
            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)

+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        with self._lock:
+            return await self.store.get_many_span_sequence_ids(rollout_attempt_ids)
+
    async def query_spans(
        self,
        rollout_id: str,
        attempt_id: str | Literal["latest"] | None = None,
-    ) -> List[Span]:
+        *,
+        trace_id: Optional[str] = None,
+        trace_id_contains: Optional[str] = None,
+        span_id: Optional[str] = None,
+        span_id_contains: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        parent_id_contains: Optional[str] = None,
+        name: Optional[str] = None,
+        name_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        limit: int = -1,
+        offset: int = 0,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+    ) -> Sequence[Span]:
        with self._lock:
-            return await self.store.query_spans(rollout_id, attempt_id)
+            return await self.store.query_spans(
+                rollout_id,
+                attempt_id,
+                trace_id=trace_id,
+                trace_id_contains=trace_id_contains,
+                span_id=span_id,
+                span_id_contains=span_id_contains,
+                parent_id=parent_id,
+                parent_id_contains=parent_id_contains,
+                name=name,
+                name_contains=name_contains,
+                filter_logic=filter_logic,
+                limit=limit,
+                offset=offset,
+                sort_by=sort_by,
+                sort_order=sort_order,
+            )

    async def update_rollout(
        self,
@@ -171,3 +301,39 @@ class LightningStoreThreaded(LightningStore):
                last_heartbeat_time=last_heartbeat_time,
                metadata=metadata,
            )
+
+    async def query_workers(
+        self,
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
+        with self._lock:
+            return await self.store.query_workers(
+                status_in=status_in,
+                worker_id_contains=worker_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        with self._lock:
+            return await self.store.get_worker_by_id(worker_id)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        with self._lock:
+            return await self.store.update_worker(
+                worker_id=worker_id,
+                heartbeat_stats=heartbeat_stats,
+            )
@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

 import time
-from typing import Awaitable, Callable, List, cast
+from typing import Awaitable, Callable, Dict, List, Tuple

 from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus

@@ -9,66 +9,102 @@ UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
 UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]


-async def propagate_status(
-    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
+LATENCY_BUCKETS = [
+    0.000001,
+    0.000002,
+    0.000005,
+    0.00001,
+    0.00002,
+    0.00005,
+    0.0001,
+    0.0002,
+    0.0005,
+    0.001,
+    0.002,
+    0.003,
+    0.005,
+    0.007,
+    0.01,
+    0.015,
+    0.02,
+    0.03,
+    0.05,
+    0.07,
+    0.1,
+    0.2,
+    0.3,
+    0.5,
+    0.7,
+    1.0,
+    2.0,
+    3.0,
+    5.0,
+    7.0,
+    10.0,
+    12.0,
+    15.0,
+    20.0,
+    25.0,
+    30.0,
+    40.0,
+    50.0,
+    60.0,
+    90.0,
+    120.0,
+    180.0,
+    240.0,
+    300.0,
+]
+
+
+async def rollout_status_from_attempt(
    attempt: Attempt,
    config: RolloutConfig,
-) -> Rollout:
+) -> RolloutStatus:
    """
    Propagate the status of an attempt to the rollout.

-    The rollout should be made sure in a state to be outdated.
-    Requeue the rollout if it should be retried.
-
-    This operation is completely unlocked. The caller is responsible for locking the store.
+    Returns:
+        The status of the rollout from the perspective of the attempt.
    """
    # Propagate the status directly to the rollout
    if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
-        return await update_rollout_status(
-            attempt.rollout_id,
-            attempt.status,
-        )
+        return attempt.status

    if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
        # Check if this status should trigger a retry
        if attempt.status in config.retry_condition:
            # If we haven't exceeded max attempts, retry
            if attempt.sequence_id < config.max_attempts:
-                return await update_rollout_status(
-                    attempt.rollout_id,
-                    "requeuing",
-                )
+                return "requeuing"

        # If we can't retry or shouldn't retry, mark as failed
-        return await update_rollout_status(
-            attempt.rollout_id,
-            "failed",
-        )
+        return "failed"

    raise ValueError(f"Invalid attempt status: {attempt.status}")


-async def healthcheck(
+async def scan_unhealthy_rollouts(
    rollouts: List[AttemptedRollout],
-    update_rollout_status: UpdateRolloutStatus,
-    update_attempt_status: UpdateAttemptStatus,
-) -> None:
+) -> Dict[Tuple[str, str], AttemptStatus]:
    """
    Perform health check on all running rollouts in the store.

    This method should be called periodically to:

-    1. Update rollout status to failed to succeeded when the attempt is done
-    2. Check for unresponsive attempts (no heartbeat or spans for a while)
-    3. Check for timed-out rollouts (running too long since start_time)
-    4. Update attempt/rollout status accordingly
+    1. Check for unresponsive attempts (no heartbeat or spans for a while)
+    2. Check for timed-out rollouts (running too long since start_time)

    This operation is completely unlocked. The caller is responsible for locking the store.

    Args:
-        store: The LightningStore instance to check rollouts from
+        rollouts: The list of running rollouts to check.
+
+    Returns:
+        A dictionary of updates to the rollouts.
    """
    current_time = time.time()
+    updates: Dict[Tuple[str, str], AttemptStatus] = {}

    for rollout in rollouts:
        config = rollout.config  # policy for retry and timeout
@@ -76,52 +112,31 @@ async def healthcheck(
        # Get the latest attempt for this rollout
        latest_attempt = rollout.attempt
        if not latest_attempt:
-            continue
-
-        # Check if the attempt has already failed or succeeded
-        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
-            await propagate_status(update_rollout_status, latest_attempt, config)
+            # This should not happen
            continue

        # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
        if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "timeout",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
            continue

        # Check for unresponsive condition (based on last heartbeat)
-        if latest_attempt.last_heartbeat_time:
-            if latest_attempt.status == "preparing":
-                # If still preparing, mark it as running
-                latest_attempt = await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "running",
-                )
+        # (1) Haven't received heartbeat for a while
+        if (
+            latest_attempt.last_heartbeat_time
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
+        ):
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue

-            # Haven't received heartbeat for a while
-            if (
-                config.unresponsive_seconds is not None
-                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
-            ):
-                await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "unresponsive",
-                )
-                continue
-
-        # Check if there's no last heartbeat (no spans) at all
+        # (2) Check if there's no last heartbeat (no spans) at all
        if (
            latest_attempt.last_heartbeat_time is None
            and config.unresponsive_seconds is not None
            and current_time - latest_attempt.start_time > config.unresponsive_seconds
        ):
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "unresponsive",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
+
+    return updates
@@ -1,7 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.

 from .agentops import AgentOpsTracer
-from .base import Tracer
+from .base import Tracer, clear_active_tracer, get_active_tracer, set_active_tracer
+from .dummy import DummyTracer
 from .otel import OtelTracer

-__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer"]
+__all__ = [
+    "AgentOpsTracer",
+    "Tracer",
+    "OtelTracer",
+    "DummyTracer",
+    "get_active_tracer",
+    "set_active_tracer",
+    "clear_active_tracer",
+]
@@ -2,25 +2,25 @@

 from __future__ import annotations

-import asyncio
 import logging
 import os
-import threading
+import warnings
 from contextlib import asynccontextmanager, contextmanager
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator, List, Optional

 import agentops
 import agentops.sdk.core
+import opentelemetry.trace as trace_api
 from agentops.sdk.core import TracingCore
-from agentops.sdk.processors import SpanProcessor
-from opentelemetry.instrumentation.utils import suppress_instrumentation
-from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
 from opentelemetry.trace.status import StatusCode

 from agentlightning.instrumentation import instrument_all, uninstrument_all
 from agentlightning.store.base import LightningStore
+from agentlightning.utils.otel import get_span_processors, get_tracer_provider

-from .base import Tracer
+from .base import with_active_tracer_context
+from .otel import LightningSpanProcessor, OtelTracer

 if TYPE_CHECKING:
    from agentops.integration.callbacks.langchain import LangchainCallbackHandler
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)


-class AgentOpsTracer(Tracer):
+class AgentOpsTracer(OtelTracer):
    """Traces agent execution using AgentOps.

    This tracer provides functionality to capture execution details using the
@@ -67,9 +67,8 @@ class AgentOpsTracer(Tracer):
    def uninstrument(self, worker_id: int):
        uninstrument_all()

-    def init_worker(self, worker_id: int):
-        super().init_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Setting up tracer...")  # worker_id included in process name
+    def _initialize_tracer_provider(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Setting up AgentOps tracer...")  # worker_id included in process name

        if self.instrument_managed:
            self.instrument(worker_id)
@@ -81,20 +80,20 @@ class AgentOpsTracer(Tracer):
                agentops.init(auto_start_session=False)  # type: ignore
                logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
            else:
-                logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
+                logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized. Skip initialization.")

-        self._lightning_span_processor = LightningSpanProcessor()
-
-        try:
-            # new versions
-            instance = agentops.sdk.core.tracer
-            # TODO: The span processor cannot be deleted once added.
-            # This might be a problem if the tracer is entered and exited multiple times.
-            instance.provider.add_span_processor(self._lightning_span_processor)  # type: ignore
-        except AttributeError:
-            # old versions
-            instance = TracingCore.get_instance()  # type: ignore
-            instance._provider.add_span_processor(self._lightning_span_processor)  # type: ignore
+        span_processors = get_span_processors(self._get_tracer_provider(), LightningSpanProcessor)
+        if len(span_processors) > 0:
+            logger.warning(
+                "LightningSpanProcessor already present in TracerProvider. You might have called init_worker() multiple times."
+                "Agent-lightning will try to reuse the existing LightningSpanProcessor."
+            )
+            if len(span_processors) > 1:
+                logger.error("More than one LightningSpanProcessors present in TracerProvider. This should not happen.")
+            self._lightning_span_processor = span_processors[0]
+        else:
+            self._lightning_span_processor = LightningSpanProcessor()
+            self._get_tracer_provider().add_span_processor(self._lightning_span_processor)  # type: ignore

    def teardown_worker(self, worker_id: int) -> None:
        super().teardown_worker(worker_id)
@@ -103,6 +102,10 @@ class AgentOpsTracer(Tracer):
            self.uninstrument(worker_id)
            logger.info(f"[Worker {worker_id}] Instrumentation removed.")

+        # NOTE: The teardown doesn't try to remove the LightningSpanProcessor from the TracerProvider.
+        # Currently there is no stable way to fully restore the AgentOps state to the initial state.
+
+    @with_active_tracer_context
    @asynccontextmanager
    async def trace_context(
        self,
@@ -111,7 +114,7 @@ class AgentOpsTracer(Tracer):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> AsyncGenerator[LightningSpanProcessor, None]:
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
        """
        Starts a new tracing context. This should be used as a context manager.

@@ -122,12 +125,18 @@ class AgentOpsTracer(Tracer):
            attempt_id: Optional attempt ID to add the spans to.

        Yields:
-            The [`LightningSpanProcessor`][agentlightning.tracer.agentops.LightningSpanProcessor] instance to collect spans.
+            The OpenTelemetry tracer instance to collect spans.
        """
-        with self._trace_context_sync(
-            name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id
-        ) as processor:
-            yield processor
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        else:
+            store = self._store
+        with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
+            yield tracer

    @contextmanager
    def _trace_context_sync(
@@ -137,47 +146,51 @@ class AgentOpsTracer(Tracer):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> Iterator[LightningSpanProcessor]:
+    ) -> Iterator[trace_api.Tracer]:
        """Implementation of `trace_context` for synchronous execution."""
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
+        tracer_provider = self._get_tracer_provider()

        kwargs: dict[str, Any] = {}
        if name is not None:
            kwargs["trace_name"] = name
        elif rollout_id is not None:
            kwargs["trace_name"] = rollout_id
+        if store is not None and rollout_id is not None and attempt_id is not None:
+            if store.capabilities.get("otlp_traces", False) is True:
+                logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
+                self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
+            else:
+                self._disable_native_otlp_exporter()
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx:
+                # AgentOps end_trace and start_trace must live inside the lightning span processor context.
+                # Otherwise some traces might not be recorded.
+                with self._agentops_trace_context(rollout_id, attempt_id, kwargs):
+                    yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
+        elif store is None and rollout_id is None and attempt_id is None:
+            self._disable_native_otlp_exporter()
+            with self._lightning_span_processor:
+                with self._agentops_trace_context(None, None, kwargs):
+                    yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+
+    @contextmanager
+    def _agentops_trace_context(self, rollout_id: Optional[str], attempt_id: Optional[str], kwargs: dict[str, Any]):
        trace = agentops.start_trace(**kwargs)
        status = StatusCode.OK  # type: ignore
        try:
-            if store is not None and rollout_id is not None and attempt_id is not None:
-                ctx = self._lightning_span_processor.with_context(
-                    store=store, rollout_id=rollout_id, attempt_id=attempt_id
-                )
-                with ctx as processor:
-                    yield processor
-            elif store is None and rollout_id is None and attempt_id is None:
-                with self._lightning_span_processor:
-                    yield self._lightning_span_processor
-            else:
-                raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+            yield
        except Exception as e:
+            # This will catch errors in user code.
            status = StatusCode.ERROR  # type: ignore
-            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
+            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}: {e}")
+            raise  # should reraise the error here so that runner can handle it
        finally:
            agentops.end_trace(trace, end_state=status)  # type: ignore

-    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.
-        """
-        if not self._lightning_span_processor:
-            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
-        return self._lightning_span_processor.spans()
-
    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
        """
        Get the Langchain callback handler for integrating with Langchain.
@@ -204,135 +217,26 @@ class AgentOpsTracer(Tracer):

    get_langchain_callback_handler = get_langchain_handler  # alias

+    def _get_tracer_provider(self) -> TracerProviderImpl:
+        try:
+            # new versions
+            instance = agentops.sdk.core.tracer
+            if instance.provider is None:
+                raise RuntimeError("AgentOps TracerProvider is not initialized.")

-class LightningSpanProcessor(SpanProcessor):
-    """Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
-    to a [`LightningStore`][agentlightning.LightningStore].
-    """
+            if get_tracer_provider() is not instance.provider:
+                logger.error(
+                    "Mismatch between global singleton TracerProvider and AgentOps TracerProvider. "
+                    "AgentOps might not work properly."
+                )

-    def __init__(self):
-        self._spans: List[ReadableSpan] = []
+            if not isinstance(instance.provider, TracerProviderImpl):  # type: ignore
+                raise RuntimeError("Unsupported TracerProvider type for AgentOps instrumentation.")

-        # Store related context and states
-        self._store: Optional[LightningStore] = None
-        self._rollout_id: Optional[str] = None
-        self._attempt_id: Optional[str] = None
-        self._lock = threading.Lock()
-
-        # private asyncio loop running in a daemon thread
-        self._loop_ready = threading.Event()
-        self._loop: Optional[asyncio.AbstractEventLoop] = None
-        self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
-        self._loop_thread.start()
-        self._loop_ready.wait()  # loop is ready
-
-    def _loop_runner(self):
-        loop = asyncio.new_event_loop()
-        self._loop = loop
-        asyncio.set_event_loop(loop)
-        self._loop_ready.set()
-        loop.run_forever()
-        loop.close()
-
-    def __enter__(self):
-        self._last_trace = None
-        self._spans = []
-        return self
-
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
-        self._store = None
-        self._rollout_id = None
-        self._attempt_id = None
-
-    def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
-        # submit to the dedicated loop and wait synchronously
-        if self._loop is None:
-            raise RuntimeError("Loop is not initialized. This should not happen.")
-
-        # If already on the exporter loop thread, schedule and return immediately.
-        # ---------------------------------------------------------------------------
-        # WHY THIS CONDITIONAL EXISTS:
-        # In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
-        # (or another finalizer) while the Python garbage collector is running on the
-        # *same thread* that owns our exporter event loop ("otel-loop").
-        #
-        # When that happens, on_end() executes on the exporter loop thread itself.
-        # If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
-        # it would deadlock immediately — because the loop cannot both wait on and run
-        # the same coroutine. The Future stays pending forever and the loop stops
-        # processing scheduled callbacks.
-        #
-        # To avoid that self-deadlock, we detect when on_end() runs on the exporter
-        # loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
-        # instead of blocking with .result().
-        #
-        # This situation can occur because Python calls __del__ in whatever thread
-        # releases the last reference, which can easily be our loop thread if the
-        # object is dereferenced during loop._run_once().
-        # ---------------------------------------------------------------------------
-        if threading.current_thread() is self._loop_thread:
-            self._loop.call_soon_threadsafe(asyncio.create_task, coro)  # type: ignore
-            return None
-
-        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)  # type: ignore
-        return fut.result(timeout=timeout)  # raises on error  # type: ignore
-
-    def shutdown(self) -> None:
-        if self._loop:
-            self._loop.call_soon_threadsafe(self._loop.stop)
-            self._loop_thread.join(timeout=5)
-            self._loop = None
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        return True
-
-    def spans(self) -> List[ReadableSpan]:
-        """
-        Get the list of spans collected by this processor.
-        This is useful for debugging and testing purposes.
-
-        Returns:
-            List of ReadableSpan objects collected during tracing.
-        """
-        return self._spans
-
-    def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
-        # simple context manager without nesting into asyncio
-        class _Ctx:
-            def __enter__(_):  # type: ignore
-                with self._lock:
-                    self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
-                    self._last_trace = None
-                    self._spans = []
-                return self
-
-            def __exit__(_, exc_type, exc, tb):  # type: ignore
-                with self._lock:
-                    self._store = self._rollout_id = self._attempt_id = None
-
-        return _Ctx()
-
-    def on_end(self, span: ReadableSpan) -> None:
-        """
-        Process a span when it ends.
-
-        Args:
-            span: The span that has ended.
-        """
-        # Skip if span is not sampled
-        if not span.context or not span.context.trace_flags.sampled:
-            return
-
-        if self._store and self._rollout_id and self._attempt_id:
-            try:
-                # Submit add_otel_span to the event loop and wait for it to complete
-                with suppress_instrumentation():
-                    self._await_in_loop(
-                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                        timeout=60.0,
-                    )
-            except Exception:
-                # log; on_end MUST NOT raise
-                logger.exception(f"Error adding span to store: {span.name}")
-
-        self._spans.append(span)
+            self._tracer_provider = instance.provider
+            return self._tracer_provider
+        except AttributeError:
+            # old versions
+            instance = TracingCore.get_instance()  # type: ignore
+            self._tracer_provider = instance._provider  # type: ignore
+            return self._tracer_provider  # type: ignore
@@ -2,13 +2,13 @@

 from __future__ import annotations

+import functools
 import logging
-from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
-
-from opentelemetry.sdk.trace import ReadableSpan
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar

 from agentlightning.store.base import LightningStore
-from agentlightning.types import ParallelWorkerBase
+from agentlightning.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus

 if TYPE_CHECKING:
    from langchain_core.callbacks.base import BaseCallbackHandler  # type: ignore
@@ -16,6 +16,14 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)


+T = TypeVar("T")
+
+
+_active_tracer: Optional[Tracer] = None
+
+T_func = Callable[..., Awaitable[Any]]
+
+
 class Tracer(ParallelWorkerBase):
    """
    An abstract base class for tracers.
@@ -51,6 +59,18 @@ class Tracer(ParallelWorkerBase):
    ```
    """

+    _store: Optional[LightningStore] = None
+
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
+        """Initialize the tracer for a worker.
+
+        Args:
+            worker_id: The ID of the worker.
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        super().init_worker(worker_id)
+        self._store = store
+
    def trace_context(
        self,
        name: Optional[str] = None,
@@ -67,11 +87,9 @@ class Tracer(ParallelWorkerBase):
        within the `with` block are collected and made available via
        [`get_last_trace`][agentlightning.Tracer.get_last_trace].

-        If a store is provided, the spans will be added to the store when tracing.
-
        Args:
            name: The name for the root span of this trace context.
-            store: The store to add the spans to.
+            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
            rollout_id: The rollout ID to add the spans to.
            attempt_id: The attempt ID to add the spans to.
        """
@@ -81,19 +99,18 @@ class Tracer(ParallelWorkerBase):
        self,
        name: Optional[str] = None,
        *,
-        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> ContextManager[Any]:
        """Internal API for CI backward compatibility."""
        raise NotImplementedError()

-    def get_last_trace(self) -> List[ReadableSpan]:
+    def get_last_trace(self) -> List[Span]:
        """
        Retrieves the raw list of captured spans from the most recent trace.

        Returns:
-            A list of OpenTelemetry `ReadableSpan` objects.
+            A list of [`Span`][agentlightning.Span] objects collected during the last trace.
        """
        raise NotImplementedError()

@@ -114,6 +131,48 @@ class Tracer(ParallelWorkerBase):
        with self._trace_context_sync(name=func.__name__):
            return func(*args, **kwargs)

+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        """Notify the tracer that a span should be created here.
+
+        It uses a fire-and-forget approach and doesn't wait for the span to be created.
+
+        Args:
+            name: The name of the span.
+            attributes: The attributes of the span.
+            timestamp: The timestamp of the span.
+            status: The status of the span.
+
+        Returns:
+            The core fields of the span.
+        """
+        raise NotImplementedError()
+
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> ContextManager[SpanRecordingContext]:
+        """Start to record an operation to a span.
+
+        Args:
+            name: The name of the operation.
+            attributes: The attributes of the operation.
+            start_time: The start time of the operation.
+            end_time: The end time of the operation.
+
+        Returns:
+            A [`SpanRecordingContext`][agentlightning.SpanRecordingContext] for recording the operation on the span.
+        """
+        raise NotImplementedError()
+
    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.
@@ -138,3 +197,91 @@ class Tracer(ParallelWorkerBase):
        """
        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
        return None
+
+    @contextmanager
+    def lifespan(self, store: Optional[LightningStore] = None):
+        """A context manager to manage the lifespan of the tracer.
+
+        This can be used to set up and tear down any necessary resources
+        for the tracer, useful for debugging purposes.
+
+        Args:
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        has_init = False
+        has_init_worker = False
+        try:
+            self.init()
+            has_init = True
+
+            self.init_worker(0, store)
+            has_init_worker = True
+
+            yield
+
+        finally:
+            if has_init_worker:
+                self.teardown_worker(0)
+            if has_init:
+                self.teardown()
+
+
+def set_active_tracer(tracer: Tracer):
+    """Set the active tracer for the current process.
+
+    Args:
+        tracer: The tracer to set as active.
+    """
+    global _active_tracer
+    if _active_tracer is not None:
+        raise ValueError("An active tracer is already set. Cannot set a new one.")
+    _active_tracer = tracer
+
+
+def clear_active_tracer():
+    """Clear the active tracer for the current process."""
+    global _active_tracer
+    _active_tracer = None
+
+
+def get_active_tracer() -> Optional[Tracer]:
+    """Get the active tracer for the current process.
+
+    Returns:
+        The active tracer, or None if no tracer is active.
+    """
+    global _active_tracer
+    return _active_tracer
+
+
+class _ActiveTracerAsyncCM(AsyncContextManager[T]):
+    def __init__(self, tracer: Tracer, inner: AsyncContextManager[T]):
+        self._tracer = tracer
+        self._inner = inner
+
+    async def __aenter__(self) -> T:
+        set_active_tracer(self._tracer)  # will raise if nested
+        try:
+            return await self._inner.__aenter__()
+        except Exception:
+            clear_active_tracer()
+            raise
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> Optional[bool]:
+        try:
+            return await self._inner.__aexit__(*args, **kwargs)
+        finally:
+            clear_active_tracer()
+
+
+def with_active_tracer_context(
+    func: Callable[..., AsyncContextManager[T]],
+) -> Callable[..., AsyncContextManager[T]]:
+    """Decorate a method returning an AsyncContextManager so tracer is active for the whole `async with`."""
+
+    @functools.wraps(func)
+    def wrapper(self: Tracer, *args: Any, **kwargs: Any) -> AsyncContextManager[T]:
+        cm = func(self, *args, **kwargs)
+        return _ActiveTracerAsyncCM(self, cm)
+
+    return wrapper
@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import (
+    Iterator,
+    Optional,
+)
+
+from agentlightning.types import (
+    Attributes,
+    SpanCoreFields,
+    SpanRecordingContext,
+    StatusCode,
+    TraceStatus,
+)
+from agentlightning.utils.otel import format_exception_attributes
+
+from .base import Tracer
+
+logger = logging.getLogger(__name__)
+
+
+class DummySpanRecordingContext(SpanRecordingContext):
+    """Context for recording operations on a dummy span, not dependent on any backend tracer."""
+
+    def __init__(self, name: str, attributes: Optional[Attributes] = None, start_time: Optional[float] = None) -> None:
+        self.name = name
+        self.attributes = attributes or {}
+        self.start_time = start_time or time.time()
+        self.end_time = None
+        self.status = TraceStatus(status_code="OK")
+
+    def record_exception(self, exception: BaseException) -> None:
+        self.record_status("ERROR", str(exception))
+        self.record_attributes(format_exception_attributes(exception))
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        self.attributes.update(attributes)
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        self.status = TraceStatus(status_code=status_code, description=description)
+
+    def finalize(self, end_time: Optional[float] = None) -> None:
+        self.end_time = end_time or time.time()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        if self.end_time is None:
+            raise ValueError("End time is not set. Call finalize() first.")
+        return SpanCoreFields(
+            name=self.name,
+            attributes=self.attributes,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            status=self.status,
+        )
+
+
+class DummyTracer(Tracer):
+    """A dummy tracer that does not trace anything, but it is compatible with the emitter API.
+
+    It doesn't rely on any backend tracer, and also doesn't use any stores.
+    """
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        if attributes is None:
+            attributes = {}
+        if timestamp is None:
+            timestamp = time.time()
+        if status is None:
+            status = TraceStatus(status_code="OK")
+        return SpanCoreFields(
+            name=name,
+            attributes=attributes,
+            start_time=timestamp,
+            end_time=timestamp,
+            status=status,
+        )
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[DummySpanRecordingContext]:
+        start_time = start_time or time.time()
+        recording_context = DummySpanRecordingContext(name, attributes, start_time)
+        try:
+            yield recording_context
+        except Exception as exc:
+            recording_context.record_exception(exc)
+            recording_context.record_status("ERROR", str(exc))
+            raise
+        finally:
+            recording_context.finalize(end_time)
@@ -1,393 +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 .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) -> None:
-        """
-        Initialize the tracer in a worker process.
-
-        Args:
-            worker_id: The ID of the worker process.
-        """
-        super().init_worker(worker_id)
-        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
@@ -2,20 +2,73 @@

 from __future__ import annotations

+import asyncio
 import logging
-from contextlib import asynccontextmanager
-from typing import AsyncGenerator, List, Optional
+import threading
+import warnings
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Awaitable, Iterator, List, Optional

 import opentelemetry.trace as trace_api
-from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+from opentelemetry.instrumentation.utils import suppress_instrumentation
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor

+from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.store.base import LightningStore
+from agentlightning.types import Attributes, Span, SpanCoreFields, SpanRecordingContext, StatusCode, TraceStatus
+from agentlightning.types.tracer import convert_timestamp
+from agentlightning.utils.otel import get_tracer_provider
+from agentlightning.utils.otlp import LightningStoreOTLPExporter

-from .agentops import LightningSpanProcessor  # FIXME: This import should be from otel to agentops
-from .base import Tracer
+from .base import Tracer, with_active_tracer_context

 logger = logging.getLogger(__name__)

+STORE_WRITE_TIMEOUT_SECONDS = 10.0
+
+
+def to_otel_status_code(status_code: StatusCode) -> trace_api.StatusCode:
+    if status_code == "UNSET":
+        return trace_api.StatusCode.UNSET
+    elif status_code == "ERROR":
+        return trace_api.StatusCode.ERROR
+    else:
+        return trace_api.StatusCode.OK
+
+
+class OtelSpanRecordingContext(SpanRecordingContext):
+    def __init__(self, span: trace_api.Span) -> None:
+        self._span = span
+
+    def record_exception(self, exception: BaseException) -> None:
+        self._span.record_exception(exception)
+        self.record_status("ERROR", str(exception))
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        self._span.set_attributes(attributes)
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        otel_status_code = to_otel_status_code(status_code)
+        self._span.set_status(otel_status_code, description)
+
+    def get_otel_span(self) -> trace_api.Span:
+        return self._span
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        if isinstance(self._span, ReadableSpan):
+            return SpanCoreFields(
+                name=self._span.name,
+                attributes=dict(self._span.attributes) if self._span.attributes else {},
+                start_time=convert_timestamp(self._span.start_time),
+                end_time=convert_timestamp(self._span.end_time),
+                status=TraceStatus.from_opentelemetry(self._span.status),
+            )
+        else:
+            raise ValueError(f"Span is not a ReadableSpan: {self._span}")
+

 class OtelTracer(Tracer):
    """Tracer that provides a basic OpenTelemetry tracer provider.
@@ -27,28 +80,47 @@ class OtelTracer(Tracer):
    def __init__(self):
        super().__init__()
        # This provider is only initialized when the worker is initialized.
-        self._tracer_provider: Optional[TracerProvider] = None
+        self._tracer_provider: Optional[trace_api.TracerProvider] = None
        self._lightning_span_processor: Optional[LightningSpanProcessor] = None
+        self._simple_span_processor: Optional[SimpleSpanProcessor] = None
+        self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
        self._initialized: bool = False

-    def init_worker(self, worker_id: int):
-        super().init_worker(worker_id)
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
+        super().init_worker(worker_id, store)
+        self._initialize_tracer_provider(worker_id)
+
+    def _initialize_tracer_provider(self, worker_id: int):
        logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")

        if self._initialized:
-            logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
+            logger.info(f"[Worker {worker_id}] Tracer provider is already initialized. Skipping initialization.")
+            return

-        tracer_provider = TracerProvider()
-        trace_api.set_tracer_provider(tracer_provider)
+        try:
+            get_tracer_provider()
+            logger.error(
+                f"[Worker {worker_id}] Tracer provider is already initialized but not by OtelTracer. OpenTelemetry may not work as expected."
+            )
+        except RuntimeError:
+            logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
+
+        self._tracer_provider = TracerProviderImpl()
+        trace_api.set_tracer_provider(self._tracer_provider)
        self._lightning_span_processor = LightningSpanProcessor()
-        tracer_provider.add_span_processor(self._lightning_span_processor)
+        self._tracer_provider.add_span_processor(self._lightning_span_processor)
+        self._otlp_span_exporter = LightningStoreOTLPExporter()
+        self._simple_span_processor = SimpleSpanProcessor(self._otlp_span_exporter)
+        self._tracer_provider.add_span_processor(self._simple_span_processor)
        self._initialized = True

+        logger.info(f"[Worker {worker_id}] OpenTelemetry tracer provider initialized.")
+
    def teardown_worker(self, worker_id: int):
        super().teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
-        self._tracer_provider = None
+        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")

+    @with_active_tracer_context
    @asynccontextmanager
    async def trace_context(
        self,
@@ -57,7 +129,7 @@ class OtelTracer(Tracer):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> AsyncGenerator[LightningSpanProcessor, None]:
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
        """
        Starts a new tracing context. This should be used as a context manager.

@@ -68,28 +140,391 @@ class OtelTracer(Tracer):
            attempt_id: Optional attempt ID to add the spans to.

        Yields:
-            The LightningSpanProcessor instance to collect spans.
+            The OpenTelemetry tracer instance to collect spans.
        """
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")

-        if store is not None and rollout_id is not None and attempt_id is not None:
-            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
-            with ctx as processor:
-                yield processor
-        elif store is None and rollout_id is None and attempt_id is None:
-            with self._lightning_span_processor:
-                yield self._lightning_span_processor
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
        else:
-            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+            store = self._store

-    def get_last_trace(self) -> List[ReadableSpan]:
+        if rollout_id is not None and attempt_id is not None:
+            if store is None:
+                raise ValueError("store is required to be initialized when rollout_id and attempt_id are provided")
+            if store.capabilities.get("otlp_traces", False) is True:
+                logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
+                self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
+            else:
+                self._disable_native_otlp_exporter()
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx:
+                yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
+        elif rollout_id is None and attempt_id is None:
+            self._disable_native_otlp_exporter()
+            with self._lightning_span_processor:
+                yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
+        else:
+            raise ValueError("rollout_id and attempt_id must be either all provided or all None")
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        # Fire the span to the current active tracer provider.
+        tracer_provider = self._get_tracer_provider()
+        tracer = tracer_provider.get_tracer(__name__)
+        span = tracer.start_span(
+            name, attributes=attributes, start_time=int(timestamp * 1_000_000_000) if timestamp else None
+        )
+        if status is not None:
+            span.set_status(to_otel_status_code(status.status_code), status.description)
+        span.end(int(timestamp * 1_000_000_000) if timestamp else None)
+
+        # The span should have been auto-created by now.
+        # Return the core fields of the span.
+        if isinstance(span, ReadableSpan):
+            return SpanCoreFields(
+                name=name,
+                attributes=dict(span.attributes) if span.attributes else {},
+                start_time=convert_timestamp(span.start_time),
+                end_time=convert_timestamp(span.end_time),
+                status=TraceStatus.from_opentelemetry(span.status),
+            )
+        else:
+            raise ValueError(f"Span is not a ReadableSpan: {span}")
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[SpanRecordingContext]:
+        if end_time is not None:
+            logger.warning("OpenTelemetry doesn't support customizing the end time of a span. End time is ignored.")
+        # Record the span to the current active tracer provider.
+        tracer_provider = self._get_tracer_provider()
+        tracer = tracer_provider.get_tracer(__name__)
+
+        # Activate the span as the current span within otel.
+        with tracer.start_as_current_span(
+            name, attributes=attributes, start_time=int(start_time * 1_000_000_000) if start_time else None
+        ) as span:
+            recording_context = OtelSpanRecordingContext(span)
+            try:
+                yield recording_context
+            except Exception as exc:
+                recording_context.record_exception(exc)
+                raise
+
+        # No need to retrieve the span here. It's already been sent to otel processor.
+
+    def get_last_trace(self) -> List[Span]:
        """
        Retrieves the raw list of captured spans from the most recent trace.

        Returns:
-            A list of OpenTelemetry `ReadableSpan` objects.
+            A list of [`Span`][agentlightning.Span] objects captured during the most recent trace.
        """
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
        return self._lightning_span_processor.spans()
+
+    def _get_tracer_provider(self) -> TracerProviderImpl:
+        if self._tracer_provider is None:
+            raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
+        if not isinstance(self._tracer_provider, TracerProviderImpl):
+            raise TypeError(f"TracerProvider is not a opentelemetry.sdk.trace.TracerProvider: {self._tracer_provider}")
+        return self._tracer_provider
+
+    def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
+        tracer_provider = self._get_tracer_provider()
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+
+        # Override the resources so that the server knows where the request comes from.
+        tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
+            Resource.create(
+                {
+                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                }
+            )
+        )
+        instrumented = False
+        candidates: List[str] = []
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # We don't need the LightningSpanProcessor any more.
+                logger.debug("LightningSpanProcessor already present in TracerProvider, disabling it.")
+                processor.disable_store_submission = True
+            elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+                # Instead, we rely on the OTLPSpanExporter to send spans to the store.
+                if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
+                    processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
+                    logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
+                    instrumented = True
+                else:
+                    candidates.append(
+                        f"{processor.__class__.__name__} with {processor.span_exporter.__class__.__name__}"
+                    )
+            else:
+                candidates.append(f"{processor.__class__.__name__}")
+
+        if not instrumented:
+            raise RuntimeError(
+                "Failed to enable native OTLP exporter: no BatchSpanProcessor or SimpleSpanProcessor with "
+                "LightningStoreOTLPExporter found in TracerProvider. Please try using a non-OTLP store."
+                "Candidates are: " + ", ".join(candidates)
+            )
+
+    def _disable_native_otlp_exporter(self):
+        tracer_provider = self._get_tracer_provider()
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+        tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
+            Resource.create(
+                {
+                    LightningResourceAttributes.ROLLOUT_ID.value: "",
+                    LightningResourceAttributes.ATTEMPT_ID.value: "",
+                }
+            )
+        )  # reset resource
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # We will be in need of the LightningSpanProcessor again.
+                logger.debug("Enabling LightningSpanProcessor in TracerProvider.")
+                processor.disable_store_submission = False
+
+
+class LightningSpanProcessor(SpanProcessor):
+    """Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
+    to a [`LightningStore`][agentlightning.LightningStore].
+
+    It serves two purposes:
+
+    1. Records all the spans in a local buffer.
+    2. Submits the spans to the event loop to be added to the store.
+    """
+
+    def __init__(self, disable_store_submission: bool = False):
+        self._disable_store_submission: bool = disable_store_submission
+        self._spans: List[Span] = []
+
+        # Store related context and states
+        self._store: Optional[LightningStore] = None
+        self._rollout_id: Optional[str] = None
+        self._attempt_id: Optional[str] = None
+        self._local_sequence_id: int = 0
+        self._lock = threading.Lock()
+
+        # private asyncio loop running in a daemon thread
+        self._loop_ready = threading.Event()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+        self._loop_init_lock = threading.Lock()
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            + f"disable_store_submission={self.disable_store_submission}, "
+            + f"store={self.store!r}, "
+            + f"rollout_id={self.rollout_id!r}, "
+            + f"attempt_id={self.attempt_id!r})"
+        )
+
+    @property
+    def store(self) -> Optional[LightningStore]:
+        """The store to submit the spans to."""
+        return self._store
+
+    @property
+    def rollout_id(self) -> Optional[str]:
+        """The rollout ID to submit the spans to."""
+        return self._rollout_id
+
+    @property
+    def attempt_id(self) -> Optional[str]:
+        """The attempt ID to submit the spans to."""
+        return self._attempt_id
+
+    @property
+    def disable_store_submission(self) -> bool:
+        """Whether to disable submitting spans to the store."""
+        return self._disable_store_submission
+
+    @disable_store_submission.setter
+    def disable_store_submission(self, value: bool) -> None:
+        self._disable_store_submission = value
+
+    def _ensure_loop(self) -> None:
+        # Fast path: loop already initialized
+        if self._loop_thread is not None and self._loop is not None:
+            return
+
+        with self._loop_init_lock:
+            # Double-check after acquiring lock
+            if self._loop_thread is not None and self._loop is not None:
+                return
+            self._loop_ready.clear()
+            self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
+            self._loop_thread.start()
+            if not self._loop_ready.wait(timeout=30.0):
+                raise RuntimeError("Timed out waiting for otel-loop thread to start")
+
+    def _loop_runner(self):
+        loop = asyncio.new_event_loop()
+        self._loop = loop
+        asyncio.set_event_loop(loop)
+        self._loop_ready.set()
+        loop.run_forever()
+        loop.close()
+
+    def __enter__(self):
+        self._last_trace = None
+        self._spans = []
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
+        self._store = None
+        self._rollout_id = None
+        self._attempt_id = None
+
+    def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
+        # submit to the dedicated loop and wait synchronously
+        self._ensure_loop()
+        if self._loop is None:
+            raise RuntimeError("Loop is not initialized. This should not happen.")
+
+        # If already on the exporter loop thread, schedule and return immediately.
+        # ---------------------------------------------------------------------------
+        # WHY THIS CONDITIONAL EXISTS:
+        # In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
+        # (or another finalizer) while the Python garbage collector is running on the
+        # *same thread* that owns our exporter event loop ("otel-loop").
+        #
+        # When that happens, on_end() executes on the exporter loop thread itself.
+        # If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
+        # it would deadlock immediately — because the loop cannot both wait on and run
+        # the same coroutine. The Future stays pending forever and the loop stops
+        # processing scheduled callbacks.
+        #
+        # To avoid that self-deadlock, we detect when on_end() runs on the exporter
+        # loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
+        # instead of blocking with .result().
+        #
+        # This situation can occur because Python calls __del__ in whatever thread
+        # releases the last reference, which can easily be our loop thread if the
+        # object is dereferenced during loop._run_once().
+        # ---------------------------------------------------------------------------
+        if threading.current_thread() is self._loop_thread:
+            self._loop.call_soon_threadsafe(asyncio.create_task, coro)  # type: ignore
+            return None
+
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)  # type: ignore
+        return fut.result(timeout=timeout)  # raises on error  # type: ignore
+
+    def shutdown(self) -> None:
+        if self._loop:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._loop = None
+        if self._loop_thread:
+            self._loop_thread.join(timeout=5)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def spans(self) -> List[Span]:
+        """
+        Get the list of spans collected by this processor.
+        This is useful for debugging and testing purposes.
+
+        Returns:
+            List of [`Span`][agentlightning.Span] objects collected during tracing.
+        """
+        return self._spans
+
+    def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
+        # simple context manager without nesting into asyncio
+        class _Ctx:
+            def __enter__(_):  # type: ignore
+                # Use _ instead of self to avoid shadowing the instance method.
+                with self._lock:
+                    self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
+                    self._last_trace = None
+                    self._spans = []
+                return self
+
+            def __exit__(_, exc_type, exc, tb):  # type: ignore
+                with self._lock:
+                    self._store = self._rollout_id = self._attempt_id = None
+
+        return _Ctx()
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """
+        Process a span when it ends.
+
+        Args:
+            span: The span that has ended.
+        """
+        # Skip if span is not sampled
+        if not span.context or not span.context.trace_flags.sampled:
+            return
+
+        if not self._disable_store_submission and self._store and self._rollout_id and self._attempt_id:
+            try:
+                # Submit add_otel_span to the event loop and wait for it to complete
+                with suppress_instrumentation():
+                    self._ensure_loop()
+                    uploaded_span = self._await_in_loop(
+                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                        timeout=STORE_WRITE_TIMEOUT_SECONDS,
+                    )
+                    if uploaded_span is not None:
+                        self._spans.append(uploaded_span)
+            except TimeoutError:
+                logger.warning(
+                    "Timed out adding span %s to store after %.1f seconds. The span will be stored locally "
+                    "but it's not guaranteed to be persisted.",
+                    span.name,
+                    STORE_WRITE_TIMEOUT_SECONDS,
+                )
+                self._spans.append(
+                    Span.from_opentelemetry(
+                        span,
+                        rollout_id=self._rollout_id,
+                        attempt_id=self._attempt_id,
+                        sequence_id=self._local_sequence_id,
+                    )
+                )
+            except Exception:
+                # log; on_end MUST NOT raise
+                logger.exception(f"Error adding span to store: {span.name}. The span will be store locally only.")
+                self._spans.append(
+                    Span.from_opentelemetry(
+                        span,
+                        rollout_id=self._rollout_id,
+                        attempt_id=self._attempt_id,
+                        sequence_id=self._local_sequence_id,
+                    )
+                )
+
+        else:
+            # Fallback path
+            created_span = Span.from_opentelemetry(
+                span,
+                rollout_id=self._rollout_id or "rollout-dummy",
+                attempt_id=self._attempt_id or "attempt-dummy",
+                sequence_id=self._local_sequence_id,
+            )
+            self._local_sequence_id += 1
+            self._spans.append(created_span)
@@ -0,0 +1,677 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures as futures
+import logging
+import os
+import re
+import weakref
+from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    cast,
+)
+
+import weave
+from opentelemetry.semconv.attributes import exception_attributes
+from weave.trace.call import Call
+from weave.trace.settings import UserSettings
+from weave.trace.weave_client import WeaveClient
+from weave.trace_server import trace_server_interface as tsi
+from weave.wandb_interface.context import set_wandb_api_context
+
+from agentlightning.instrumentation.weave import InMemoryWeaveTraceServer, instrument_weave, uninstrument_weave
+from agentlightning.semconv import LightningResourceAttributes, LightningSpanAttributes
+from agentlightning.store.base import LightningStore
+from agentlightning.types import (
+    Attributes,
+    OtelResource,
+    Span,
+    SpanContext,
+    SpanCoreFields,
+    SpanRecordingContext,
+    StatusCode,
+    TraceStatus,
+)
+from agentlightning.utils.id import generate_id
+from agentlightning.utils.otel import (
+    filter_and_unflatten_attributes,
+    flatten_attributes,
+    format_exception_attributes,
+    sanitize_attributes,
+)
+
+from .base import Tracer, with_active_tracer_context
+
+logger = logging.getLogger(__name__)
+
+
+def op_name_to_func_name(op_name: str) -> str:
+    """Convert a Weave operation name to a function name.
+
+    Weave operation names look like this: `weave:///xxx/agentlightning.tracer.weave/op/openai.chat.completions.create:019b10be-...-44d74272569c`
+    """
+    match = re.search(r"/([^/:]+):", op_name)
+    if match:
+        return match.group(1)
+    else:
+        return op_name
+
+
+def random_project_name() -> str:
+    return "agl/weave-" + generate_id(12)
+
+
+def get_timestamp_or_throw(date: Optional[datetime], field_name: str) -> float:
+    if date is None:
+        raise ValueError(f"{field_name} is required but not set")
+    return date.timestamp()
+
+
+class WeaveSpanRecordingContext(SpanRecordingContext):
+    """Universal interface for recording operations on a Weave call."""
+
+    def __init__(self, call: Call) -> None:
+        self._call = call
+
+    def record_exception(self, exception: BaseException) -> None:
+        self._call.exception = str(exception)
+        self.record_status("ERROR", str(exception))
+        self.record_attributes(format_exception_attributes(exception))
+
+    def _get_input_from_attributes(self, attributes: Attributes) -> Dict[str, Any]:
+        if LightningSpanAttributes.OPERATION_INPUT.value in attributes:
+            # This can be a very rare case. If it happens, we can just let it throw.
+            return cast(Dict[str, Any], attributes[LightningSpanAttributes.OPERATION_INPUT.value])
+        else:
+            filtered_attributes = filter_and_unflatten_attributes(
+                attributes, LightningSpanAttributes.OPERATION_INPUT.value
+            )
+            if isinstance(filtered_attributes, list):
+                return {str(i): v for i, v in enumerate(filtered_attributes)}
+            else:
+                return filtered_attributes
+
+    def _get_output_from_attributes(self, attributes: Attributes) -> Any:
+        if LightningSpanAttributes.OPERATION_OUTPUT.value in attributes:
+            return attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]
+        else:
+            return filter_and_unflatten_attributes(attributes, LightningSpanAttributes.OPERATION_OUTPUT.value)
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        input_attributes = self._get_input_from_attributes(attributes)
+        if input_attributes:
+            self._call.inputs.update(input_attributes)
+
+        output_attributes = self._get_output_from_attributes(attributes)
+        if output_attributes:
+            if self._call.output is not None:
+                logger.warning(f"Output is already set. It will be overridden: {self._call.output}")
+            self._call.output = output_attributes
+
+        if LightningSpanAttributes.OPERATION_NAME.value in attributes:
+            logger.error(
+                f"Cannot record operation name as an attribute. It will be skipped: {attributes[LightningSpanAttributes.OPERATION_NAME.value]}"
+            )
+
+        # The rest of the attributes are recorded as summary.
+        for key, value in attributes.items():
+            if (
+                not key == LightningSpanAttributes.OPERATION_INPUT.value
+                and not key.startswith(LightningSpanAttributes.OPERATION_INPUT.value + ".")
+                and not key == LightningSpanAttributes.OPERATION_OUTPUT.value
+                and not key.startswith(LightningSpanAttributes.OPERATION_OUTPUT.value + ".")
+                and not key == LightningSpanAttributes.OPERATION_NAME.value
+            ):
+                if self._call.summary is None:
+                    self._call.summary = {}
+                self._call.summary[key] = value
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        if status_code == "ERROR":
+            if not description:
+                raise ValueError("Description is required when status code is ERROR")
+            self._call.exception = description
+        elif status_code == "OK":
+            self._call.exception = None
+        # Do nothing for other status codes.
+
+    def finalize(self) -> None:
+        # Do nothing
+        pass
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        return SpanCoreFields(
+            name=self._call.op_name,
+            attributes=flatten_attributes(self._call.attributes or {}),
+            start_time=self._call.started_at.timestamp() if self._call.started_at else None,
+            end_time=self._call.ended_at.timestamp() if self._call.ended_at else None,
+            status=TraceStatus(
+                status_code="OK" if self._call.exception is None else "ERROR", description=self._call.exception
+            ),
+        )
+
+
+class WeaveTracerManagedTraceServer(InMemoryWeaveTraceServer):
+    """A managed trace server for WeaveTracer."""
+
+    def __init__(
+        self,
+        partial_call_callback: Callable[[Dict[str, Any]], None],
+        complete_call_callback: Callable[[tsi.CallSchema], None],
+    ):
+        super().__init__()
+        self.partial_call_callback = partial_call_callback
+        self.complete_call_callback = complete_call_callback
+        self._calls_already_invoked: set[str] = set()
+
+    def trigger_callbacks(self, call_id: str) -> None:
+        with self._call_threading_lock:
+            if call_id in self.calls:
+                if call_id not in self._calls_already_invoked:
+                    self._calls_already_invoked.add(call_id)
+                    self.complete_call_callback(self.calls[call_id])
+                else:
+                    logger.info(f"Call {call_id} has callback already invoked. Skipping.")
+            elif call_id in self.partial_calls:
+                self.partial_call_callback(self.partial_calls[call_id])
+            else:
+                logger.error(f"Call {call_id} not found in partial_calls or calls")
+
+    def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
+        try:
+            ret = super().call_start(req)
+            self.trigger_callbacks(ret.id)
+            return ret
+        except Exception:
+            logger.exception(f"Error calling call_start: {req}", exc_info=True)
+            raise
+
+    def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
+        try:
+            ret = super().call_end(req)
+            self.trigger_callbacks(req.end.id)
+            return ret
+        except Exception:
+            logger.exception(f"Error calling call_end: {req}", exc_info=True)
+            raise
+
+    def clear(self) -> None:
+        self._calls_already_invoked.clear()
+
+
+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.
+    """
+
+    def __init__(
+        self,
+        *,
+        project_name: str | None = None,
+        weave_user_settings: UserSettings | None = None,
+        instrument_managed: bool = True,
+    ):
+        """Initialize a WeaveTracer instance.
+
+        Args:
+            project_name: Optional project name for Weave; defaults to the current module name.
+            weave_user_settings: Optional UserSettings for Weave.
+            instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
+        """
+        super().__init__()
+        self.project_name = project_name
+        self.instrument_managed = instrument_managed
+        self.weave_user_settings = weave_user_settings or UserSettings(use_server_cache=False)
+
+        self._store: Optional[LightningStore] = None
+        self._server = WeaveTracerManagedTraceServer(
+            partial_call_callback=self.partial_call_callback, complete_call_callback=self.complete_call_callback
+        )
+
+        self._default_sequence_counter: int = 0
+        self._calls: Dict[str, tsi.CallSchema] = {}  # call_id -> call
+        self._spans: List[Span] = []  # spans in the current trace
+        self._rollout_id: Optional[str] = None
+        self._attempt_id: Optional[str] = None
+        self._partial_call_futures: Dict[str, asyncio.Future[int] | futures.Future[int]] = {}
+        self._complete_call_futures: List[asyncio.Future[None] | futures.Future[None]] = []
+        self._loop: weakref.ReferenceType[asyncio.AbstractEventLoop] | None = None
+
+    def instrument(self, worker_id: int):
+        instrument_weave(self._server)
+
+    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
+
+        # Optionally patch network calls to bypass real Weave/W&B endpoints
+        if self.instrument_managed:
+            self.instrument(worker_id)
+
+        # If WANDB_API_KEY is not set, we need to initialize Weave with a hack
+        if not os.getenv("WANDB_API_KEY"):
+            logger.info("WANDB_API_KEY is not set. Initializing Weave a mock context.")
+            set_wandb_api_context("agl", api_key=None, headers=None, cookies=None)
+        else:
+            logger.debug("WANDB_API_KEY is set. Weave will be initialized automatically.")
+
+        weave_client = weave.get_client()
+        if self.project_name is None:
+            self.project_name = random_project_name()
+
+        if weave_client is not None:
+            logger.warning("Weave client was already initialized. Reentrant calls are at your own risk.")
+            if weave_client.project == self.project_name:
+                logger.error(
+                    f"Weave client was already initialized for the same project '{self.project_name}'. It's very likely that weave won't work correctly."
+                )
+
+        # Init no matter what
+        try:
+            weave.init(project_name=self.project_name, settings=self.weave_user_settings)
+            logger.info(f"[Worker {worker_id}] Weave client initialized.")
+        except Exception as exc:
+            raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}'") from exc
+
+    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.")
+
+    @with_active_tracer_context
+    @asynccontextmanager
+    async def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[Any]:
+        """Asynchronous implementation of the tracing context.
+
+        Args:
+            name: Optional operation name.
+            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.
+        """
+
+        if rollout_id is not None and attempt_id is not None:
+            self._rollout_id = rollout_id
+            self._attempt_id = attempt_id
+        elif rollout_id is None and attempt_id is None:
+            logger.info("No rollout_id or attempt_id provided. Skipping writing to store.")
+            self._rollout_id = self._attempt_id = None
+        else:
+            raise ValueError("rollout_id and attempt_id must be either both provided or both None")
+
+        await self._init_trace_context()
+
+        weave_client = self._get_weave_client()
+
+        if weave_client.server is not self._server:
+            logger.error(
+                "Weave client is not using the correct trace server. You might have multiple WeaveTracer instances running in the same process. "
+                f"Expected {self._server}, got {weave_client.server}"
+            )
+
+        arg_op = name or weave_client.project
+        arg_inputs: dict[str, str] = {}
+        if rollout_id is not None:
+            arg_inputs[LightningResourceAttributes.ROLLOUT_ID.value] = rollout_id
+        if attempt_id is not None:
+            arg_inputs[LightningResourceAttributes.ATTEMPT_ID.value] = attempt_id
+
+        try:
+            # Create a new trace call object in Weave
+            trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+                op=arg_op, inputs=arg_inputs
+            )
+
+            try:
+                yield trace_call
+                # Finish trace even if no exception
+                weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+            except Exception as exc:
+                # Finish trace and log any exception
+                weave_client.finish_call(trace_call, exception=exc)  # pyright: ignore[reportUnknownMemberType]
+                logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={exc}")
+                raise
+
+        finally:
+            try:
+                weave_client.flush()
+                # It's possible that the call end futures are from a dedicated Weave thread pool,
+                await asyncio.gather(*[asyncio.wrap_future(future) for future in self._complete_call_futures])
+
+            finally:
+                # Mandatory cleanup
+                self._rollout_id = None
+                self._attempt_id = None
+                self._server.clear()
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        if timestamp is not None:
+            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
+        weave_client = self._get_weave_client()
+        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+            op=name,
+            attributes=attributes,
+            inputs={},
+        )
+        # Immediately finish the call
+        weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+        # We don't wait for the call to be propagated to the server.
+        start_time = trace_call.started_at.timestamp() if trace_call.started_at else None
+        end_time = trace_call.ended_at.timestamp() if trace_call.ended_at else None
+        trace_status = (
+            TraceStatus(status_code="OK")
+            if trace_call.exception is None
+            else TraceStatus(status_code="ERROR", description=trace_call.exception)
+        )
+        return SpanCoreFields(
+            name=name,
+            attributes=flatten_attributes(trace_call.attributes or {}),
+            start_time=start_time,
+            end_time=end_time,
+            status=trace_status,
+        )
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[SpanRecordingContext]:
+        if start_time is not None:
+            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
+        if end_time is not None:
+            logger.warning("Weave doesn't support customizing the end time of a call. Timestamp is ignored.")
+        weave_client = self._get_weave_client()
+        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+            op=name,
+            attributes=attributes,
+            inputs={},
+        )
+        recording_context = WeaveSpanRecordingContext(trace_call)
+        try:
+            yield recording_context
+        except Exception as exc:
+            recording_context.record_exception(exc)
+            raise
+        finally:
+            weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+
+    async def _init_trace_context(self) -> None:
+        """Initialize the trace context."""
+        self._spans.clear()
+        self._calls.clear()
+        self._partial_call_futures.clear()
+        self._complete_call_futures.clear()
+        self._loop = weakref.ref(asyncio.get_running_loop())
+
+    def _get_weave_client(self) -> WeaveClient:
+        """Get the Weave client."""
+        weave_client = weave.get_client()
+        if not weave_client:
+            raise RuntimeError("Weave client is not initialized. Call init_worker() first.")
+        return weave_client
+
+    def _ensure_loop(self) -> tuple[asyncio.AbstractEventLoop, bool]:
+        """Returns a usable event loop and a boolean indicating whether it's the current running loop.
+
+        Prefer using the main loop if it's possible. Otherwise, use the current running loop.
+        """
+        # Get the current running loop
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        # Get the main loop, which can be a different loop
+        if self._loop is not None:
+            main_loop = self._loop()
+        else:
+            main_loop = None
+
+        if main_loop is not None:
+            return main_loop, id(main_loop) == id(running_loop)
+        elif running_loop is not None:
+            return running_loop, True
+        else:
+            raise RuntimeError("No running event loop found. This should not happen.")
+
+    def get_last_trace(self) -> List[Span]:
+        return self._spans
+
+    def partial_call_callback(self, request_content: Dict[str, Any]) -> None:
+        call_id = request_content.get("id")
+        if call_id is None:
+            raise ValueError("Call ID is required even for partial calls")
+
+        if call_id in self._partial_call_futures:
+            raise ValueError(f"Call {call_id} already has a start future")
+
+        # The callback must possibly be called from a dedicated Weave thread pool,
+        # but it should be executed on the main event loop.
+        try:
+            loop, is_current_loop = self._ensure_loop()
+            if is_current_loop:
+                task = loop.create_task(self.partial_call_handler(request_content))
+            else:
+                # Schedule the task on the dedicated loop
+                task = asyncio.run_coroutine_threadsafe(self.partial_call_handler(request_content), loop)
+            self._partial_call_futures[call_id] = task
+        except Exception as exc:
+            logger.exception(f"Error creating call start task: {exc}", exc_info=True)
+
+    def complete_call_callback(self, call: tsi.CallSchema) -> None:
+        try:
+            loop, is_current_loop = self._ensure_loop()
+            if is_current_loop:
+                task = loop.create_task(self.complete_call_handler(call))
+            else:
+                # Schedule the task on the dedicated loop
+                task = asyncio.run_coroutine_threadsafe(self.complete_call_handler(call), loop)
+            self._complete_call_futures.append(task)
+        except Exception as exc:
+            logger.exception(f"Error creating call finish task: {exc}", exc_info=True)
+
+    async def _get_next_sequence_id(self) -> int:
+        """Get the next sequence ID for a span.
+
+        Use store to get the next sequence ID if available, otherwise use a default counter.
+        """
+        if self._rollout_id and self._attempt_id and self._store:
+            return await self._store.get_next_span_sequence_id(self._rollout_id, self._attempt_id)
+        else:
+            self._default_sequence_counter += 1
+            return self._default_sequence_counter
+
+    async def partial_call_handler(self, request_content: Dict[str, Any]) -> int:
+        """Handler called when a Weave Call starts.
+
+        Args:
+            request_content: The partial Weave Call object.
+
+        Returns:
+            The sequence ID for the call.
+        """
+        sequence_id = await self._get_next_sequence_id()
+        return sequence_id
+
+    async def complete_call_handler(self, call: tsi.CallSchema) -> None:
+        """Handler called when a Weave Call finishes.
+
+        Converts the call (including nested children) into spans and stores them in LightningStore.
+        """
+        # Make sure the corresponding call_start_future is complete
+        if call.id in self._partial_call_futures:
+            sequence_id = await asyncio.wrap_future(self._partial_call_futures[call.id])
+            del self._partial_call_futures[call.id]
+        else:
+            # Fetch a new sequence ID as the call_start is somehow missing
+            if call.id in self._calls:
+                logger.warning(
+                    f"Call {call.id} is already in calls. The call is already completed. Overwriting the call."
+                )
+            else:
+                logger.warning(f"Call {call.id} has no start future. Fetching a new sequence ID.")
+            sequence_id = await self._get_next_sequence_id()
+
+        self._calls[call.id] = call
+
+        span = await self.convert_call_to_span(call, self._rollout_id, self._attempt_id, sequence_id)
+        self._spans.append(span)
+        if self._store and self._rollout_id and self._attempt_id:
+            try:
+                await self._store.add_span(span)
+            except Exception as exc:
+                logger.exception(f"Error adding span to store: {exc}")
+
+    async def convert_call_to_span(
+        self,
+        call: tsi.CallSchema,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        """Convert a Weave Call (with nested children) into a Agent-lightning Span.
+
+        `rollout_id` and `attempt_id` are required to attach the spans to the store.
+
+        Args:
+            call: The Weave Call object.
+            rollout_id: Optional rollout ID to attach to spans.
+            attempt_id: Optional attempt ID to attach to spans.
+            sequence_id: Optional sequence ID to attach to spans.
+
+        Returns:
+            List of converted spans.
+        """
+        rollout_id = rollout_id or "rollout-dummy"
+        attempt_id = attempt_id or "attempt-dummy"
+        sequence_id = sequence_id or 0
+
+        start_ts: float = call.started_at.timestamp()
+        end_ts: Optional[float] = call.ended_at.timestamp() if call.ended_at else None
+
+        if call.exception:
+            status = TraceStatus(status_code="ERROR", description=call.exception)
+        else:
+            status = TraceStatus(status_code="OK")
+
+        attributes: Dict[str, Any] = {
+            LightningSpanAttributes.OPERATION_NAME.value: call.op_name,
+            # op_name can be possibly overridden by the attributes.
+            **call.attributes,
+        }
+        if call.inputs:
+            attributes[LightningSpanAttributes.OPERATION_INPUT.value] = call.inputs
+        if call.output:
+            attributes[LightningSpanAttributes.OPERATION_OUTPUT.value] = call.output
+        if call.summary:
+            # attributes can be possibly overridden by the summary.
+            attributes.update(call.summary)
+        if call.exception:
+            attributes[exception_attributes.EXCEPTION_MESSAGE] = call.exception
+
+        sanitized_attributes = sanitize_attributes(flatten_attributes(attributes, expand_leaf_lists=False))
+
+        context = SpanContext(
+            trace_id=call.trace_id,
+            span_id=call.id,
+            is_remote=False,
+            trace_state={},
+        )
+
+        # Get context for parent
+        if call.parent_id:
+            parent_call = self._calls.get(call.parent_id)
+            if parent_call:
+                parent_context = SpanContext(
+                    trace_id=parent_call.trace_id,
+                    span_id=parent_call.id,
+                    is_remote=False,
+                    trace_state={},
+                )
+            else:
+                parent_context = None
+        else:
+            parent_context = None
+
+        # Build the Span object
+        return Span(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            trace_id=call.trace_id,
+            span_id=call.id,
+            parent_id=call.parent_id,
+            name=op_name_to_func_name(call.op_name),
+            status=status,
+            attributes=sanitized_attributes,
+            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={
+                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                    LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id,
+                    LightningResourceAttributes.TRACER_NAME.value: "weave",
+                },
+                schema_url="",
+            ),
+        )
@@ -152,6 +152,13 @@ class Trainer(TrainerLegacy):
        # super().__init__() will call TrainerLegacy's initialization, which is not intended.
        self.worker_id: Optional[int] = None

+        if dev:
+            warnings.warn(
+                "Trainer(dev=True) is deprecated and will be removed in future versions. "
+                "Please use Trainer.dev(...) instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
        self._dev = dev
        self.daemon = daemon
        self._client: AgentLightningClient | None = None  # Will be initialized in fit or fit_v0
@@ -213,10 +220,6 @@ class Trainer(TrainerLegacy):
        # We might be able to support a list of resources in future.
        self.initial_resources = initial_resources

-        # The active store for the current execution context
-        self.store = self._make_store(store)
-        self.runner = self._make_runner(runner)
-
        self.port = port

        self.strategy = self._make_strategy(
@@ -224,6 +227,11 @@ class Trainer(TrainerLegacy):
            n_runners=self.n_runners,
            port=port,
        )
+
+        # The active store for the current execution context
+        self.store = self._make_store(store, self.strategy)
+        self.runner = self._make_runner(runner)
+
        if hasattr(self.strategy, "n_runners"):
            strategy_runners = getattr(self.strategy, "n_runners")
            if isinstance(strategy_runners, int) and strategy_runners > 0:
@@ -282,13 +290,19 @@ class Trainer(TrainerLegacy):
            type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
        )

-    def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
-        """Resolve the store implementation backing rollouts, attempts, spans, and resources."""
+    def _make_store(self, store: ComponentSpec[LightningStore], strategy: ExecutionStrategy) -> LightningStore:
+        """Resolve the store implementation backing rollouts, attempts, spans, and resources.
+
+        By default, it's always a in-memory store. If using a client/server execution strategy,
+        the in-memory store will be initialized in a thread-safe manner.
+        """
+        is_client_server = isinstance(strategy, ClientServerExecutionStrategy)
+        default_store_factory = lambda: InMemoryLightningStore(thread_safe=is_client_server)
        return build_component(
            store,
            expected_type=LightningStore,
            spec_name="store",
-            default_factory=InMemoryLightningStore,
+            default_factory=default_store_factory,
            invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
            type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
        )
@@ -10,20 +10,25 @@ from typing import (
    Callable,
    Dict,
    Generic,
+    Iterator,
    List,
    Literal,
+    Mapping,
    Optional,
    Protocol,
+    Sequence,
    SupportsIndex,
+    TypedDict,
    TypeVar,
    Union,
    cast,
+    overload,
 )

 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field, model_validator

-from .tracer import Span
+from .tracer import Span, SpanCoreFields

 if TYPE_CHECKING:
    from agentlightning.litagent import LitAgent
@@ -48,7 +53,14 @@ __all__ = [
    "Rollout",
    "Attempt",
    "AttemptedRollout",
+    "EnqueueRolloutRequest",
    "Hook",
+    "Worker",
+    "WorkerStatus",
+    "PaginatedResult",
+    "FilterOptions",
+    "SortOptions",
+    "FilterField",
 ]

 T_co = TypeVar("T_co", covariant=True)
@@ -200,6 +212,50 @@ class AttemptedRollout(Rollout):
        return self


+class EnqueueRolloutRequest(BaseModel):
+    """Payload describing a rollout to be queued via [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout].
+
+    A subset of fields from [`Rollout`][agentlightning.Rollout] used for queuing new rollouts.
+    """
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+    config: Optional[RolloutConfig] = None
+    """Retry and timeout configuration associated with the rollout."""
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
+WorkerStatus = Literal["idle", "busy", "unknown"]
+
+
+class Worker(BaseModel):
+    """Worker information. This is actually the same as Runner info."""
+
+    worker_id: str
+    """The ID of the worker."""
+    status: WorkerStatus = "unknown"
+    """The status of the worker."""
+    heartbeat_stats: Optional[Dict[str, Any]] = None
+    """Statistics about the worker's heartbeat."""
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported the stats."""
+    last_dequeue_time: Optional[float] = None
+    """The last time when the worker has tried to dequeue a rollout."""
+    last_busy_time: Optional[float] = None
+    """The last time when the worker has started an attempt and became busy."""
+    last_idle_time: Optional[float] = None
+    """The last time when the worker has triggered the end of an attempt and became idle."""
+    current_rollout_id: Optional[str] = None
+    """The ID of the current rollout that the worker is processing."""
+    current_attempt_id: Optional[str] = None
+    """The ID of the current attempt that the worker is processing."""
+
+
 TaskInput = Any
 """Task input type. Accepts arbitrary payloads."""

@@ -251,6 +307,7 @@ RolloutRawResult = Union[
    float,  # only final reward
    List[ReadableSpan],  # constructed OTEL spans by user
    List[Span],  # constructed Span objects by user
+    List[SpanCoreFields],  # constructed SpanCoreFields objects by user
 ]
 """Rollout result type.

@@ -393,3 +450,104 @@ class Hook(ParallelWorkerBase):
        Subclasses can override this method for cleanup or additional
        logging. By default, this is a no-op.
        """
+
+
+class FilterField(TypedDict, total=False):
+    """An operator dict for a single field."""
+
+    exact: Any
+    within: Sequence[Any]
+    contains: str
+
+
+FilterOptions = Mapping[
+    Union[str, Literal["_aggregate", "_must"]],
+    Union[FilterField, Literal["and", "or"], Mapping[str, FilterField]],
+]
+"""A mapping of field name -> operator dict.
+
+Each operator dict can contain:
+
+- "exact": value for exact equality.
+- "within": iterable of allowed values.
+- "contains": substring to search for in string fields.
+
+The filter can also have a special field called "_aggregate" that can be used to specify the logic
+to combine the results of the filters:
+
+- "and": all conditions must match. This is the default value if not specified.
+- "or": at least one condition must match.
+
+All conditions within a field and between different fields are
+stored in a unified pool and combined using `_aggregate`.
+
+The filter can also have a special group called "_must", which is a mapping of filters that must all match,
+no matter whether the aggregate logic is "and" or "or".
+
+Example:
+
+```json
+{
+    "_aggregate": "or",
+    "_must": {
+        "city": {"exact": "New York"},
+        "timezone": {"within": ["America/New_York", "America/Los_Angeles"]},
+    },
+    "status": {"exact": "active"},
+    "id": {"within": [1, 2, 3]},
+    "name": {"contains": "foo"},
+}
+```
+"""
+
+
+class SortOptions(TypedDict):
+    """Options for sorting the collection."""
+
+    name: str
+    """The name of the field to sort by."""
+    order: Literal["asc", "desc"]
+    """The order to sort by."""
+
+
+T_item = TypeVar("T_item")
+
+
+class PaginatedResult(BaseModel, Sequence[T_item]):
+    """Result of a paginated query.
+
+    Behaves like a sequence, but also carries pagination metadata (limit, offset, total).
+    """
+
+    items: Sequence[T_item]
+    """Items in the result."""
+    limit: int
+    """Limit of the result."""
+    offset: int
+    """Offset of the result."""
+    total: int
+    """Total number of items in the collection."""
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    @overload
+    def __getitem__(self, index: int) -> T_item: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[T_item]: ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[T_item, Sequence[T_item]]:
+        return self.items[index]
+
+    # Overriding __iter__ enables list(paginated_result) to work as expected,
+    # but changes Pydantic's default dict iteration behavior (which would otherwise
+    # iterate over field names).
+    def __iter__(self) -> Iterator[T_item]:  # type: ignore
+        return iter(self.items)
+
+    def __repr__(self) -> str:
+        first_item_repr = repr(self.items[0]) if self.items else "empty"
+        items_repr = f"[{first_item_repr}, ...]" if len(self.items) > 1 else first_item_repr
+        slice_repr = f"{self.offset}:" if self.limit == -1 else f"{self.offset}:{self.offset + self.limit}"
+        return f"<PaginatedResult ({slice_repr} of {self.total}) {items_repr}>"
@@ -2,11 +2,13 @@

 from __future__ import annotations

+import time
+
 """Data models that mirror OpenTelemetry spans for Agent Lightning."""

 import json
 from enum import Enum
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, Union

 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.resources import Resource
@@ -16,6 +18,8 @@ from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace.status import Status as OtelStatus
 from pydantic import BaseModel, ConfigDict

+from agentlightning.semconv import AGL_VIRTUAL
+
 __all__ = [
    "AttributeValue",
    "Attributes",
@@ -29,6 +33,9 @@ __all__ = [
    "SpanNames",
    "SpanAttributeNames",
    "SpanLike",
+    "StatusCode",
+    "SpanCoreFields",
+    "SpanRecordingContext",
 ]


@@ -81,6 +88,8 @@ Attributes = Dict[str, AttributeValue]
 """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
 TraceState = Dict[str, str]
 """Mapping from trace state key to its value. Same as OpenTelemetry `TraceState` type."""
+StatusCode = Literal["UNSET", "OK", "ERROR"]
+"""The status code of the span."""


 class SpanContext(BaseModel):
@@ -113,7 +122,7 @@ class SpanContext(BaseModel):
 class TraceStatus(BaseModel):
    """Serializable variant of `opentelemetry.trace.Status`."""

-    status_code: str
+    status_code: StatusCode
    """The status code of the span. Same as OpenTelemetry `Status.status_code` type."""
    description: Optional[str] = None
    """The description of the span. Same as OpenTelemetry `Status.description` type."""
@@ -201,6 +210,44 @@ class OtelResource(BaseModel):
        )


+class SpanCoreFields(BaseModel):
+    """Core fields of a span. Used by span creators who don't care about the full span model.
+
+    If the spans are managed by some OTel tracer provider, it's not advised to create spans via this path.
+    """
+
+    name: str
+    """The name of the span."""
+    status: TraceStatus
+    """The status of the span."""
+    attributes: Attributes
+    """The attributes of the span."""
+    start_time: Optional[float]
+    """The start time of the span."""
+    end_time: Optional[float]
+    """The end time of the span."""
+
+
+class SpanRecordingContext(Protocol):
+    """Context for recording operations on a span. It doesn't have to finalize the span; the caller will do it."""
+
+    def record_exception(self, exception: BaseException) -> None:
+        """Record an exception on the span."""
+        raise NotImplementedError()
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        """Record attributes on the span."""
+        raise NotImplementedError()
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        """Record the status of the span."""
+        raise NotImplementedError()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        """Get the recording of the span."""
+        raise NotImplementedError()
+
+
 class Span(BaseModel):
    """Agent Lightning's canonical span model used for persistence and analytics.

@@ -338,6 +385,7 @@ class Span(BaseModel):
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        resource: Optional[OtelResource] = None,
+        status: Optional[TraceStatus] = None,
    ) -> "Span":
        """Build a synthetic span from raw attributes.
        Different from the [`from_opentelemetry`][agentlightning.Span.from_opentelemetry] method,
@@ -355,6 +403,7 @@ class Span(BaseModel):
            start_time: Span start timestamp in seconds.
            end_time: Span end timestamp in seconds.
            resource: Explicit resource information to attach to the span.
+            status: Optional status of the span.

        Returns:
            [`Span`][agentlightning.Span] populated with the provided attributes.
@@ -379,10 +428,10 @@ class Span(BaseModel):
                is_remote=False,
                trace_state={},
            ),
-            name=name or SpanNames.VIRTUAL.value,
+            name=name or AGL_VIRTUAL,
            resource=resource or OtelResource(attributes={}, schema_url=""),
            attributes=attributes,
-            status=TraceStatus(status_code="OK"),
+            status=status or TraceStatus(status_code="OK"),
            events=[],
            links=[],
            parent=(
@@ -397,9 +446,40 @@ class Span(BaseModel):
            ),
        )

+    @classmethod
+    def from_core_fields(
+        cls,
+        core: SpanCoreFields,
+        *,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        """Build a span from a core span.
+
+        Args:
+            core: Core span to build from.
+            rollout_id: Optional rollout identifier associated with the span.
+            attempt_id: Optional attempt identifier associated with the span.
+            sequence_id: Optional sequence number to preserve ordering.
+
+        Returns:
+            [`Span`][agentlightning.Span] populated with the provided attributes.
+        """
+        return cls.from_attributes(
+            attributes=core.attributes,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            name=core.name,
+            start_time=core.start_time or time.time(),
+            end_time=core.end_time,
+            status=core.status,
+        )
+

 class SpanNames(str, Enum):
-    """Enumerated span names recognised by Agent-lightning."""
+    """Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""

    REWARD = "agentlightning.reward"
    """The name of the reward span."""
@@ -411,10 +491,16 @@ class SpanNames(str, Enum):
    """The name of the exception span."""
    VIRTUAL = "agentlightning.virtual"
    """The name of the virtual span. It represents derived spans without concrete operations."""
+    ROLLOUT_ID = "agentlightning.rollout_id"
+    """The name of the rollout ID."""
+    ATTEMPT_ID = "agentlightning.attempt_id"
+    """The name of the attempt ID."""
+    SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
+    """The name of the span sequence ID."""


 class SpanAttributeNames(str, Enum):
-    """Canonical attribute names written by Agent Lightning emitters."""
+    """Canonical attribute names written by Agent Lightning emitters. Deprecated in favor of [semconv][agentlightning.semconv]."""

    MESSAGE = "message"
    """The name of the message attribute."""
@@ -0,0 +1,18 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import hashlib
+import uuid
+
+__all__ = ["generate_id"]
+
+
+def generate_id(length: int) -> str:
+    """Generate a random ID of the given length.
+
+    Args:
+        length: The length of the ID to generate.
+
+    Returns:
+        A random ID of the given length.
+    """
+    return hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:length]
@@ -0,0 +1,541 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities shared for OpenTelemetry span (attributes) support."""
+
+import json
+import logging
+import traceback
+from typing import Any, Dict, List, Sequence, Type, TypeVar, Union, cast
+from warnings import filterwarnings
+
+import opentelemetry.trace as trace_api
+from agentops.sdk.exporters import OTLPSpanExporter
+from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SpanProcessor, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
+from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
+from opentelemetry.semconv.attributes import exception_attributes
+from opentelemetry.trace import get_tracer_provider as otel_get_tracer_provider
+from pydantic import TypeAdapter
+
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
+from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
+from agentlightning.types import Attributes, AttributeValue, SpanLike
+from agentlightning.utils.otlp import LightningStoreOTLPExporter
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "full_qualified_name",
+    "get_tracer_provider",
+    "get_tracer",
+    "make_tag_attributes",
+    "extract_tags_from_attributes",
+    "make_link_attributes",
+    "query_linked_spans",
+    "extract_links_from_attributes",
+    "filter_attributes",
+    "filter_and_unflatten_attributes",
+    "flatten_attributes",
+    "unflatten_attributes",
+    "sanitize_attribute_value",
+    "sanitize_attributes",
+    "sanitize_list_attribute_sanity",
+    "check_attributes_sanity",
+    "format_exception_attributes",
+]
+
+T_SpanLike = TypeVar("T_SpanLike", bound=SpanLike)
+T_SpanProcessor = TypeVar("T_SpanProcessor", bound=SpanProcessor)
+
+
+def full_qualified_name(obj: type) -> str:
+    if str(obj.__module__) == "builtins":
+        return obj.__qualname__
+    return f"{obj.__module__}.{obj.__qualname__}"
+
+
+def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
+    """Get the OpenTelemetry tracer provider configured for Agent Lightning.
+
+    Args:
+        inspect: Whether to inspect the tracer provider and log its configuration.
+            When it's on, make sure you also set the logger level to DEBUG to see the logs.
+    """
+    from agentlightning.tracer.otel import LightningSpanProcessor
+
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+    tracer_provider = otel_get_tracer_provider()
+    if not isinstance(tracer_provider, TracerProviderImpl):
+        logger.error(
+            "Tracer provider is expected to be an instance of opentelemetry.sdk.trace.TracerProvider, found: %s",
+            full_qualified_name(type(tracer_provider)),
+        )
+        return cast(TracerProviderImpl, tracer_provider)
+
+    if not inspect:
+        return tracer_provider
+
+    emitter_debug = resolve_bool_env_var(LightningEnvVar.AGL_EMITTER_DEBUG, fallback=None)
+    logger_effective_level = logger.getEffectiveLevel()
+    if emitter_debug is True and logger_effective_level > logging.DEBUG:
+        logger.warning(
+            "Emitter debug logging is enabled but logging level is not set to DEBUG. Nothing will be logged."
+        )
+
+    if emitter_debug is None:
+        # Set to true by default if the logging level is lower than DEBUG
+        emitter_debug = logging.DEBUG >= logger_effective_level
+
+    if emitter_debug:
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+        processors: List[str] = []
+        active_span_processor_cls = active_span_processor.__class__.__name__
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # The legacy case for tracers without OTLP support.
+                processors.append(f"{active_span_processor_cls} - {processor!r}")
+            elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+                processor_cls = processor.__class__.__name__
+                if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
+                    # This should be the main path now.
+                    processors.append(f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter!r}")
+                elif isinstance(processor.span_exporter, OTLPSpanExporter):
+                    # You need to be careful if the code goes into this path.
+                    endpoint = processor.span_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
+                    processors.append(
+                        f"{active_span_processor_cls} - {processor_cls} - "
+                        f"{processor.span_exporter.__class__.__name__}(endpoint={endpoint!r})"
+                    )
+                else:
+                    # Other cases like Console Span Exporter.
+                    processors.append(
+                        f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter.__class__.__name__}"
+                    )
+            else:
+                processors.append(f"{active_span_processor_cls} - {processor.__class__.__name__}")
+
+        logger.debug(f"Tracer provider: {tracer_provider!r}. Active span processors:")
+        for processor in processors:
+            logger.debug("  * " + processor)
+
+    return tracer_provider
+
+
+def get_span_processors(
+    tracer_provider: TracerProviderImpl, expected_type: Type[T_SpanProcessor]
+) -> List[T_SpanProcessor]:
+    """Get the span processors from the tracer provider.
+
+    Args:
+        tracer_provider: The tracer provider to get the span processors from.
+        expected_type: The type of the span processors to get.
+
+    Returns:
+        A list of span processors of the expected type.
+    """
+    processors: List[T_SpanProcessor] = []
+    for processor in tracer_provider._active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+        if isinstance(processor, expected_type):
+            processors.append(processor)
+    return processors
+
+
+def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
+    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
+
+    Args:
+        use_active_span_processor: Whether to use the active span processor.
+
+    Returns:
+        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
+
+    Raises:
+        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
+    """
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = get_tracer_provider(inspect=True)  # inspection is on by default
+
+    if use_active_span_processor:
+        return tracer_provider.get_tracer("agentlightning")
+
+    else:
+        filterwarnings(
+            "ignore",
+            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
+            category=DeprecationWarning,
+            module="opentelemetry.sdk.trace",
+        )
+
+        return Tracer(
+            tracer_provider.sampler,
+            tracer_provider.resource,
+            # We use an empty span processor to avoid emitting spans to the tracer
+            SynchronousMultiSpanProcessor(),
+            tracer_provider.id_generator,
+            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
+            SpanLimits(),
+            InstrumentationScope(
+                "agentlightning",
+                "",
+                "",
+                {},
+            ),
+        )
+
+
+def make_tag_attributes(tags: List[str]) -> Dict[str, Any]:
+    """Convert a list of tags into flattened attributes for span tagging.
+
+    There is no syntax enforced for tags, they are just strings. For example:
+
+    ```python
+    ["gen_ai.model:gpt-4", "reward.extrinsic"]
+    ```
+    """
+    return flatten_attributes({LightningSpanAttributes.TAG.value: tags}, expand_leaf_lists=True)
+
+
+def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
+    """Extract tag attributes from flattened span attributes.
+
+    Args:
+        attributes: A dictionary of flattened span attributes.
+    """
+    maybe_tag_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.TAG.value)
+    return TypeAdapter(List[str]).validate_python(maybe_tag_list)
+
+
+def make_link_attributes(links: Dict[str, str]) -> Dict[str, Any]:
+    """Convert a dictionary of links into flattened attributes for span linking.
+
+    Links example:
+
+    ```python
+    {
+        "gen_ai.response.id": "response-123",
+        "span_id": "abcd-efgh-ijkl",
+    }
+    ```
+    """
+    link_list: List[Dict[str, str]] = []
+    for key, value in links.items():
+        if not isinstance(value, str):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise ValueError(f"Link value must be a string, got {type(value)} for key '{key}'")
+        link_list.append({LinkAttributes.KEY_MATCH.value: key, LinkAttributes.VALUE_MATCH.value: value})
+    return flatten_attributes({LightningSpanAttributes.LINK.value: link_list}, expand_leaf_lists=True)
+
+
+def query_linked_spans(spans: Sequence[T_SpanLike], links: List[LinkPydanticModel]) -> List[T_SpanLike]:
+    """Query spans that are linked by the given link attributes.
+
+    Args:
+        spans: A sequence of spans to search.
+        links: A list of link attributes to match.
+
+    Returns:
+        A list of spans that match the given link attributes.
+    """
+    matched_spans: List[T_SpanLike] = []
+
+    for span in spans:
+        span_attributes = span.attributes or {}
+        is_match = True
+        for link in links:
+            # trace_id and span_id must be full match.
+            if link.key_match == "trace_id":
+                if isinstance(span, ReadableSpan):
+                    trace_id = trace_api.format_trace_id(span.context.trace_id) if span.context else None
+                else:
+                    trace_id = span.trace_id
+                if trace_id != link.value_match:
+                    is_match = False
+                    break
+
+            elif link.key_match == "span_id":
+                if isinstance(span, ReadableSpan):
+                    span_id = trace_api.format_span_id(span.context.span_id) if span.context else None
+                else:
+                    span_id = span.span_id
+                if span_id != link.value_match:
+                    is_match = False
+                    break
+
+            else:
+                attribute = span_attributes.get(link.key_match)
+                # attributes must also be a full match currently.
+                if attribute != link.value_match:
+                    is_match = False
+                    break
+
+        if is_match:
+            matched_spans.append(span)
+
+    return matched_spans
+
+
+def extract_links_from_attributes(attributes: Dict[str, Any]) -> List[LinkPydanticModel]:
+    """Extract link attributes from flattened span attributes.
+
+    Args:
+        attributes: A dictionary of flattened span attributes.
+    """
+    maybe_link_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+    return TypeAdapter(List[LinkPydanticModel]).validate_python(maybe_link_list)
+
+
+def filter_attributes(attributes: Dict[str, Any], prefix: str) -> Dict[str, Any]:
+    """Filter attributes that start with the given prefix.
+
+    The attribute must start with `prefix.` or be exactly `prefix` to be included.
+
+    Args:
+        attributes: A dictionary of span attributes.
+        prefix: The prefix to filter by.
+
+    Returns:
+        A dictionary of attributes that start with the given prefix.
+    """
+    return {k: v for k, v in attributes.items() if k.startswith(prefix + ".") or k == prefix}
+
+
+def filter_and_unflatten_attributes(attributes: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
+    """Filter attributes that start with the given prefix and unflatten them.
+    The prefix will be removed during unflattening.
+
+    Args:
+        attributes: A dictionary of span attributes.
+        prefix: The prefix to filter by.
+
+    Returns:
+        A nested dictionary or list of attributes that start with the given prefix.
+    """
+    filtered_attributes = filter_attributes(attributes, prefix)
+    stripped_attributes: Dict[str, Any] = {}
+    for k, v in filtered_attributes.items():
+        if k == prefix:
+            raise ValueError(f"Cannot unflatten attribute with key exactly equal to prefix: {prefix}")
+        else:
+            stripped_key = k[len(prefix) + 1 :]  # +1 to remove the dot
+        stripped_attributes[stripped_key] = v
+    return unflatten_attributes(stripped_attributes)
+
+
+def flatten_attributes(
+    nested_data: Union[Dict[str, Any], List[Any]], *, expand_leaf_lists: bool = False
+) -> Dict[str, Any]:
+    """Flatten a nested dictionary or list into a flat dictionary with dotted keys.
+
+    This function recursively traverses dictionaries and lists, producing a flat
+    key-value mapping where nested paths are represented via dot-separated keys.
+    Lists are indexed numerically.
+
+    Example:
+
+        >>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}}, expand_leaf_lists=True)
+        {"a.b": 1, "a.c.0": 2, "a.c.1": 3}
+
+    Args:
+        nested_data: A nested structure composed of dictionaries, lists, or primitive values.
+        expand_leaf_lists: Whether to expand lists composed only of primitive values.
+            When `False` (the default), lists of str/int/float/bool are treated as
+            leaf values and stored without enumerating their indices.
+
+    Returns:
+        A flat dictionary mapping dotted-string paths to primitive values.
+    """
+
+    flat: Dict[str, Any] = {}
+
+    def _primitive_type(value: Any) -> Union[type[str], type[int], type[float], type[bool]]:
+        if isinstance(value, bool):
+            return bool
+        if isinstance(value, int):
+            return int
+        if isinstance(value, float):
+            return float
+        return str
+
+    def _walk(value: Any, prefix: str = "") -> None:
+        if isinstance(value, dict):
+            for k, v in cast(Dict[Any, Any], value).items():
+                if not isinstance(k, str):
+                    raise ValueError(
+                        f"Only string keys are supported in dictionaries, got '{k}' of type {type(k)} in {prefix}"
+                    )
+                new_prefix = f"{prefix}.{k}" if prefix else k
+                _walk(v, new_prefix)
+        elif isinstance(value, list):
+            maybe_list = cast(List[Any], value)
+            is_leaf_candidate = bool(maybe_list) and all(
+                isinstance(item, (str, int, float, bool)) for item in maybe_list
+            )
+            if not expand_leaf_lists and is_leaf_candidate and prefix:
+                primitive_types = {_primitive_type(item) for item in maybe_list}
+                if len(primitive_types) == 1:
+                    flat[prefix] = maybe_list
+                    return
+                logger.warning(
+                    "List attribute '%s' contains mixed primitive types %s; expanding indexed keys instead.",
+                    prefix,
+                    primitive_types,
+                )
+
+            for idx, item in enumerate(maybe_list):
+                new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
+                _walk(item, new_prefix)
+        else:
+            flat[prefix] = value
+
+    _walk(nested_data)
+    return flat
+
+
+def unflatten_attributes(flat_data: Dict[str, Any]) -> Union[Dict[str, Any], List[Any]]:
+    """Reconstruct a nested dictionary/list structure from a flat dictionary.
+
+    Keys are dot-separated paths. Segments that are digit strings will only
+    become list indices if *all* keys in that dict form a consecutive
+    0..n-1 range. Otherwise they remain dict keys.
+
+    Example:
+
+        >>> unflatten_attributes({"a.b": 1, "a.c.0": 2, "a.c.1": 3})
+        {"a": {"b": 1, "c": [2, 3]}}
+
+    Args:
+        flat_data: A dictionary whose keys are dot-separated paths and whose
+            values are primitive data elements.
+
+    Returns:
+        A nested dictionary (and lists where appropriate) corresponding to
+        the flattened structure.
+    """
+    # 1) Build a pure dict tree first (no lists yet)
+    root: Dict[str, Any] = {}
+
+    for flat_key, value in flat_data.items():
+        parts = flat_key.split(".")
+        curr: Dict[str, Any] = root
+
+        for part in parts[:-1]:
+            # Ensure intermediate node is a dict
+            if part not in curr or not isinstance(curr[part], dict):
+                curr[part] = {}
+            curr = curr[part]  # type: ignore[assignment]
+
+        curr[parts[-1]] = value
+
+    # 2) Recursively convert dicts-with-consecutive-numeric-keys into lists
+    def convert(node: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
+        if isinstance(node, dict):
+            # First convert children
+            for k, v in list(node.items()):
+                node[k] = convert(v)
+
+            if not node:
+                # empty dict stays dict
+                return node
+
+            # Check if keys are all numeric strings
+            keys = list(node.keys())
+            if all(isinstance(k, str) and k.isdigit() for k in keys):  # pyright: ignore[reportUnnecessaryIsInstance]
+                indices = sorted(int(k) for k in keys)
+                # Must be exactly 0..n-1
+                if indices == list(range(len(indices))):
+                    return [node[str(i)] for i in range(len(indices))]
+
+            return node
+
+        if isinstance(node, list):  # pyright: ignore[reportUnnecessaryIsInstance]
+            return [convert(v) for v in node]
+
+        # Keep as is
+        return node
+
+    return convert(root)
+
+
+def sanitize_attribute_value(object: Any, force: bool = True) -> AttributeValue:
+    """Sanitize an attribute value to be a valid OpenTelemetry attribute value."""
+    if isinstance(object, (str, int, float, bool)):
+        return object
+
+    if isinstance(object, list):
+        try:
+            return sanitize_list_attribute_sanity(cast(List[Any], object))
+        except ValueError as exc:
+            logger.warning(f"Failed to sanitize list attribute. Fallback to JSON serialization: {exc}")
+
+    try:
+        # This include null, dict, etc.
+        serialized = json.dumps(object, default=str if force else None)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Object must be JSON serializable, got: {type(cast(Any, object))}.") from exc
+    return serialized
+
+
+def sanitize_attributes(attributes: Dict[str, Any], force: bool = True) -> Attributes:
+    """Sanitize a dictionary of attributes to be a valid OpenTelemetry attributes.
+
+    Args:
+        attributes: A dictionary of attributes to sanitize.
+        force: Whether to force sanitization even when the value is not JSON serializable.
+    """
+    result: Attributes = {}
+    for k, v in attributes.items():
+        try:
+            result[k] = sanitize_attribute_value(v, force=force)
+        except ValueError as exc:
+            raise ValueError(f"Failed to sanitize attribute '{k}': {exc}") from exc
+    return result
+
+
+def sanitize_list_attribute_sanity(maybe_list: List[Any]) -> AttributeValue:
+    """Try to sanitize a list of attributes to be a valid OpenTelemetry attribute value.
+
+    Raise error if the list contains multiple types of primitive values.
+    """
+    if all(isinstance(item, str) for item in maybe_list):
+        return list[str](maybe_list)
+    if all(isinstance(item, bool) for item in maybe_list):
+        return list[bool](maybe_list)
+    if all(isinstance(item, (int, bool)) for item in maybe_list):
+        return [int(item) for item in maybe_list]
+    if all(isinstance(item, (float, int, bool)) for item in maybe_list):
+        return [float(item) for item in maybe_list]
+
+    list_types: List[Any] = [type(item) for item in maybe_list]
+    raise ValueError(f"List must contain only one type of primitive values, got: {set(list_types)}.")
+
+
+def check_attributes_sanity(attributes: Dict[Any, Any]) -> None:
+    """Check if a dictionary of attributes is a valid OpenTelemetry attributes."""
+    for k, v in attributes.items():
+        if not isinstance(k, str):
+            raise ValueError(f"Attribute key must be a string, got {type(k)} for key '{k}'")
+        if isinstance(v, list):
+            try:
+                sanitize_list_attribute_sanity(cast(List[Any], v))
+            except ValueError as exc:
+                raise ValueError(f"Failed to sanitize list attribute '{k}': {exc}") from exc
+        elif not isinstance(v, (str, int, float, bool)):
+            raise ValueError(
+                f"Attribute value must be a string, int, float, bool, or list of these, got {type(v)} for value '{v}'"
+            )
+
+
+def format_exception_attributes(exception: BaseException) -> Attributes:
+    """Format an exception into a dictionary of attributes."""
+    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
+    span_attributes: Attributes = {
+        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
+        exception_attributes.EXCEPTION_MESSAGE: str(exception),
+        exception_attributes.EXCEPTION_ESCAPED: True,
+    }
+    if stacktrace.strip():
+        span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+    return span_attributes
@@ -0,0 +1,475 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import gzip
+import logging
+from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar
+
+from fastapi import Request, Response
+from google.protobuf import json_format
+from google.rpc.status_pb2 import Status
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import (
+    ExportLogsServiceRequest,
+    ExportLogsServiceResponse,
+)
+from opentelemetry.proto.collector.metrics.v1.metrics_service_pb2 import (
+    ExportMetricsServiceRequest,
+    ExportMetricsServiceResponse,
+)
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+    ExportTraceServiceResponse,
+)
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
+from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
+from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
+from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.util.types import AttributeValue
+
+from agentlightning.semconv import LightningResourceAttributes
+from agentlightning.types.tracer import (
+    Attributes,
+    Event,
+    Link,
+    OtelResource,
+    Span,
+    SpanContext,
+    StatusCode,
+    TraceStatus,
+    convert_timestamp,
+)
+
+PROTOBUF_CT = "application/x-protobuf"
+
+logger = logging.getLogger(__name__)
+
+
+T_request = TypeVar("T_request", ExportLogsServiceRequest, ExportMetricsServiceRequest, ExportTraceServiceRequest)
+T_response = TypeVar("T_response", ExportLogsServiceResponse, ExportMetricsServiceResponse, ExportTraceServiceResponse)
+
+
+async def handle_otlp_export(
+    request: Request,
+    request_message_cls: Type[T_request],
+    response_message_cls: Type[T_response],
+    message_callback: Optional[Callable[[T_request], Awaitable[None]]],
+    signal_name: str,
+) -> Response:
+    """
+    Generic handler for /v1/traces, /v1/metrics, /v1/logs.
+
+    Convert the OTLP Protobuf request to a JSON-like object.
+    """
+    content_type = request.headers.get("Content-Type", "").split(";")[0].strip()
+
+    if content_type != PROTOBUF_CT:
+        # For brevity we only support binary protobuf here.
+        return _bad_request_response(
+            request,
+            f"Unsupported Content-Type '{content_type}', expected '{PROTOBUF_CT}'",
+            content_type=PROTOBUF_CT,
+        )
+
+    raw_body = await request.body()
+    body = _read_body_maybe_gzip(request, raw_body)
+
+    # Empty request is allowed and should still succeed.
+    if not body:
+        req_msg = request_message_cls()
+    else:
+        req_msg = request_message_cls()
+        try:
+            req_msg.ParseFromString(body)
+        except Exception as exc:
+            return _bad_request_response(request, f"Unable to parse OTLP {signal_name} payload: {exc}")
+
+    if message_callback is not None:
+        await message_callback(req_msg)
+
+    # Build success response. Partial success field is left unset.
+    resp_msg = response_message_cls()
+
+    # Encode response in the same Content-Type as request.
+    if content_type == PROTOBUF_CT:
+        resp_bytes = resp_msg.SerializeToString()
+    else:
+        resp_bytes = json_format.MessageToJson(resp_msg).encode("utf-8")
+
+    resp_bytes, headers = _maybe_gzip_response(request, resp_bytes)
+
+    return Response(
+        content=resp_bytes,
+        media_type=content_type,
+        status_code=200,
+        headers=headers,
+    )
+
+
+async def spans_from_proto(
+    request: ExportTraceServiceRequest,
+    sequence_id_bulk_issuer: Callable[[Sequence[Tuple[str, str]]], Awaitable[Sequence[int]]],
+) -> List[Span]:
+    """Parse an OTLP proto payload into List[Span].
+
+    A store is needed here for generating a sequence ID for each span.
+    """
+    output_spans: List[Span] = []
+
+    for resource_spans in request.resource_spans:
+        # Resource-level attributes & IDs
+        resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
+        # rollout_id, attempt_id from resource attributes when present.
+        rollout_id_resource = resource_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
+        attempt_id_resource = resource_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
+        # If sequence id is provided, all the spans will share the same sequence ID.
+        # unless otherwise overridden by span-level attributes.
+        sequence_id_resource = resource_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
+
+        otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
+
+        # Each ScopeSpans contains multiple spans
+        for scope_spans in resource_spans.scope_spans:
+            for proto_span in scope_spans.spans:
+                trace_id_hex = _bytes_to_trace_id_hex(proto_span.trace_id)
+                span_id_hex = _bytes_to_span_id_hex(proto_span.span_id)
+                parent_id_hex = _bytes_to_span_id_hex(proto_span.parent_span_id) if proto_span.parent_span_id else None
+
+                # Status
+                status_code_str = _STATUS_CODE_MAP.get(proto_span.status.code, "UNSET")
+                status = TraceStatus(
+                    status_code=status_code_str,
+                    description=proto_span.status.message or None,
+                )
+
+                # Attributes
+                span_attrs = _kv_list_to_dict(proto_span.attributes)
+
+                # Context
+                context = SpanContext(
+                    trace_id=trace_id_hex,
+                    span_id=span_id_hex,
+                    is_remote=False,
+                    trace_state={},
+                )
+
+                # Try to get if span attributes contain something like rollout_id or attempt_id
+                # Override the resource-level attributes with the span-level attributes if present.
+                rollout_id_span = span_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
+                attempt_id_span = span_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
+                sequence_id_span = span_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
+
+                # Normalize to regular strings and ints
+                rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
+                attempt_id_raw = attempt_id_span if attempt_id_span is not None else attempt_id_resource
+                sequence_id_raw = sequence_id_span if sequence_id_span is not None else sequence_id_resource
+
+                rollout_id, attempt_id = _normalize_rollout_attempt_id(rollout_id_raw, attempt_id_raw)
+                sequence_id = _normalize_sequence_id(sequence_id_raw)
+
+                if rollout_id is None or attempt_id is None:
+                    logger.warning(
+                        "Both rollout_id and attempt_id must be present in resource attributes. "
+                        "Spans will not be able to log to the store because of missing IDs: rollout_id=%s, attempt_id=%s, sequence_id=%s",
+                        rollout_id,
+                        attempt_id,
+                        sequence_id,
+                    )
+                    continue
+
+                # Generate a new sequence ID if not provided
+                if sequence_id is None:
+                    current_sequence_id = -1
+                elif sequence_id < 0:
+                    logger.error(
+                        "Invalid sequence_id value in resource attributes: %r. Must be a positive integer. Regenerating one.",
+                        sequence_id,
+                    )
+                    current_sequence_id = -1
+                else:
+                    current_sequence_id = sequence_id
+
+                # Build Span
+                span = Span(
+                    rollout_id=rollout_id,
+                    attempt_id=attempt_id,
+                    sequence_id=current_sequence_id,
+                    trace_id=trace_id_hex,
+                    span_id=span_id_hex,
+                    parent_id=parent_id_hex,
+                    name=proto_span.name,
+                    status=status,
+                    attributes=span_attrs,
+                    events=_events_from_proto(proto_span),
+                    links=_links_from_proto(proto_span),
+                    start_time=convert_timestamp(proto_span.start_time_unix_nano),
+                    end_time=convert_timestamp(proto_span.end_time_unix_nano),
+                    context=context,
+                    parent=None,  # OTLP only has parent_span_id; we don't have full SpanContext
+                    resource=otel_resource,
+                )
+
+                output_spans.append(span)
+
+    # Finalize the sequence IDs
+    bulk_issue_requests = [(span.rollout_id, span.attempt_id) for span in output_spans if span.sequence_id < 0]
+    bulk_sequence_ids = await sequence_id_bulk_issuer(bulk_issue_requests)
+    for span, sequence_id in zip(
+        [span for span in output_spans if span.sequence_id < 0], bulk_sequence_ids, strict=True
+    ):
+        span.sequence_id = sequence_id
+
+    return output_spans
+
+
+class LightningStoreOTLPExporter(OTLPSpanExporter):
+    """OTLP Exporter that write to a LightningStore-compatible backend.
+
+    The backend requires two special attributes on each span:
+
+    - `agentlightning.rollout_id`: The rollout ID to associate the span with.
+    - `agentlightning.attempt_id`: The attempt ID to associate the span with.
+
+    It can optionally use the following attribute to sequence spans:
+
+    - `agentlightning.span_sequence_id`: A decimal string representing the sequence ID of the span.
+    """
+
+    _default_endpoint: Optional[str] = None
+    _rollout_id: Optional[str] = None
+    _attempt_id: Optional[str] = None
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            + f"endpoint={self.endpoint!r}, "
+            + f"rollout_id={self.rollout_id!r}, "
+            + f"attempt_id={self.attempt_id!r}, "
+            + f"should_bypass={self.should_bypass()!r})"
+        )
+
+    @property
+    def endpoint(self) -> Optional[str]:
+        """The endpoint to submit the spans to."""
+        if hasattr(self, "_endpoint"):
+            return self._endpoint
+        return None
+
+    @property
+    def rollout_id(self) -> Optional[str]:
+        """The rollout ID to submit the spans to."""
+        if hasattr(self, "_rollout_id"):
+            return self._rollout_id
+        return None
+
+    @property
+    def attempt_id(self) -> Optional[str]:
+        """The attempt ID to submit the spans to."""
+        if hasattr(self, "_attempt_id"):
+            return self._attempt_id
+        return None
+
+    def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
+        """Enable storing OTLP data to a specific LightningStore rollout/attempt."""
+        self._rollout_id = rollout_id
+        self._attempt_id = attempt_id
+
+        self._default_endpoint = self._endpoint
+        self._endpoint = endpoint
+
+    def disable_store_otlp(self) -> None:
+        """Disable storing OTLP data to LightningStore."""
+        self._rollout_id = None
+        self._attempt_id = None
+        if self._default_endpoint is not None:
+            self._endpoint = self._default_endpoint
+
+    def should_bypass(self) -> bool:
+        """Check if the exporter should bypass the default export if rollout_id and attempt_id are not set."""
+        return True
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        if self._rollout_id is not None and self._attempt_id is not None:
+            # rollout_id and attempt_id are present in resource attributes
+            # It means that the server supports OTLP endpoint.
+            for span in spans:
+                # Override the resources so that the server knows where the request comes from.
+                span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                    Resource.create(
+                        {
+                            LightningResourceAttributes.ROLLOUT_ID.value: self._rollout_id,
+                            LightningResourceAttributes.ATTEMPT_ID.value: self._attempt_id,
+                        }
+                    )
+                )
+            return super().export(spans)
+        elif not self.should_bypass():
+            logger.debug("Rollout ID and Attempt ID not set; using default OTLP exporter behavior.")
+            return super().export(spans)
+        else:
+            logger.debug("Rollout ID and Attempt ID not set; bypassing export.")
+            return SpanExportResult.SUCCESS
+
+
+def _read_body_maybe_gzip(request: Request, raw_body: bytes) -> bytes:
+    """
+    Decompress body if Content-Encoding: gzip; otherwise return as is.
+    """
+    encoding = request.headers.get("Content-Encoding", "").lower()
+    if encoding == "gzip":
+        return gzip.decompress(raw_body)
+    return raw_body
+
+
+def _maybe_gzip_response(request: Request, payload: bytes) -> Tuple[bytes, Dict[str, str]]:
+    """
+    If Accept-Encoding includes gzip, gzip the payload and set Content-Encoding header.
+    """
+    ae = request.headers.get("Accept-Encoding", "")
+    tokens = [token.split(";")[0].strip().lower() for token in ae.split(",") if token.strip()]
+    headers: Dict[str, str] = {}
+    if "gzip" in tokens:
+        payload = gzip.compress(payload)
+        headers["Content-Encoding"] = "gzip"
+    return payload, headers
+
+
+def _bad_request_response(request: Request, message: str, content_type: str = PROTOBUF_CT) -> Response:
+    """
+    Build a 400 response whose body is a protobuf Status message, encoded
+    in the same Content-Type as the request (OTLP/HTTP requirement).
+    """
+    status_msg = Status(message=message)
+
+    if content_type == PROTOBUF_CT:
+        body = status_msg.SerializeToString()
+    else:
+        # Fallback: JSON representation of Status.
+        body = json_format.MessageToJson(status_msg).encode("utf-8")
+
+    body, headers = _maybe_gzip_response(request, body)
+
+    return Response(
+        content=body,
+        status_code=400,
+        media_type=content_type,
+        headers=headers,
+    )
+
+
+def _normalize_rollout_attempt_id(
+    rollout_id: Optional[AttributeValue], attempt_id: Optional[AttributeValue]
+) -> Tuple[Optional[str], Optional[str]]:
+    """Normalize a rollout or attempt ID to a string."""
+    rollout_id_str = str(rollout_id) if rollout_id is not None else None
+    attempt_id_str = str(attempt_id) if attempt_id is not None else None
+    return rollout_id_str, attempt_id_str
+
+
+def _normalize_sequence_id(sequence_id: Optional[AttributeValue]) -> Optional[int]:
+    """Normalize a sequence ID to an integer."""
+    if sequence_id is None:
+        return None
+    try:
+        sequence_id_int = int(str(sequence_id))
+    except (ValueError, TypeError):
+        logger.warning(
+            "Invalid sequence_id value in resource attributes: %r. Must be an integer or string representing an integer. Assuming None.",
+            sequence_id,
+        )
+        sequence_id_int = None
+    return sequence_id_int
+
+
+def _any_value_to_python(value: AnyValue) -> Any:
+    """Convert OTLP AnyValue -> plain Python value."""
+    kind = value.WhichOneof("value")
+    if kind is None:
+        return None
+    if kind == "string_value":
+        return value.string_value
+    if kind == "bool_value":
+        return value.bool_value
+    if kind == "int_value":
+        return int(value.int_value)
+    if kind == "double_value":
+        return float(value.double_value)
+    if kind == "array_value":
+        return [_any_value_to_python(v) for v in value.array_value.values]
+    if kind == "kvlist_value":
+        # Map<string, AnyValue> -> dict
+        return {kv.key: _any_value_to_python(kv.value) for kv in value.kvlist_value.values}
+    if kind == "bytes_value":
+        # Serialize bytes as hex string to stay JSON-friendly
+        return value.bytes_value.hex()
+    return None
+
+
+def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
+    """Convert repeated KeyValue -> Attributes dict."""
+    return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
+
+
+_STATUS_CODE_MAP: Mapping[ProtoStatus.StatusCode.ValueType, StatusCode] = {
+    ProtoStatus.STATUS_CODE_UNSET: "UNSET",
+    ProtoStatus.STATUS_CODE_OK: "OK",
+    ProtoStatus.STATUS_CODE_ERROR: "ERROR",
+}
+
+
+def _bytes_to_trace_id_hex(b: bytes) -> str:
+    # OTLP uses 16-byte trace IDs; format as 32-char hex
+    if not b:
+        return "0" * 32
+    return b.hex().rjust(32, "0")
+
+
+def _bytes_to_span_id_hex(b: bytes) -> str:
+    # OTLP uses 8-byte span IDs; format as 16-char hex
+    if not b:
+        return "0" * 16
+    return b.hex().rjust(16, "0")
+
+
+def _events_from_proto(span: ProtoSpan) -> List[Event]:
+    """Event converter from OTLP ProtoSpan to List[Event]."""
+    return [
+        Event(
+            name=e.name,
+            attributes=_kv_list_to_dict(e.attributes),
+            timestamp=convert_timestamp(e.time_unix_nano),
+        )
+        for e in span.events
+    ]
+
+
+def _links_from_proto(span: ProtoSpan) -> List[Link]:
+    """Link converter from OTLP ProtoSpan to List[Link]."""
+    links: List[Link] = []
+    for link in span.links:
+        trace_id_hex = _bytes_to_trace_id_hex(link.trace_id)
+        span_id_hex = _bytes_to_span_id_hex(link.span_id)
+        ctx = SpanContext(
+            trace_id=trace_id_hex,
+            span_id=span_id_hex,
+            is_remote=False,
+            trace_state={},  # OTLP trace_state is currently a string; you can parse if needed
+        )
+        links.append(
+            Link(
+                context=ctx,
+                attributes=_kv_list_to_dict(link.attributes) or None,
+            )
+        )
+    return links
+
+
+def _resource_from_proto(resource: ProtoResource, schema_url: str = "") -> OtelResource:
+    return OtelResource(
+        attributes=_kv_list_to_dict(resource.attributes),
+        schema_url=schema_url or "",
+    )
@@ -6,15 +6,17 @@ import asyncio
 import inspect
 import logging
 import multiprocessing
+import os
 import queue
 import signal
+import socket
 import threading
 import time
 import traceback
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass
 from multiprocessing.process import BaseProcess
-from typing import Any, AsyncContextManager, AsyncIterator, Dict, Literal, Optional
+from typing import Any, AsyncContextManager, AsyncIterator, Dict, Literal, Optional, cast

 import aiohttp
 import requests
@@ -24,18 +26,21 @@ from gunicorn.app.base import BaseApplication
 from gunicorn.arbiter import Arbiter
 from portpicker import pick_unused_port

-__all__ = ["PythonServerLauncher", "PythonServerLauncherArgs"]
+__all__ = ["PythonServerLauncher", "PythonServerLauncherArgs", "LaunchMode"]


 LaunchMode = Literal["asyncio", "thread", "mp"]
+"""The launch mode for the server."""


@dataclass
 class PythonServerLauncherArgs:
    port: Optional[int] = None
    """The TCP port to listen on. If not provided, the server will use a random available port."""
-    host: str = "127.0.0.1"
+    host: Optional[str] = None
    """The hostname or IP address to bind the server to."""
+    access_host: Optional[str] = None
+    """The hostname or IP address to advertise to the client. If not provided, the server will use the default outbound IPv4 address for this machine."""
    launch_mode: LaunchMode = "asyncio"
    """The launch mode. `asyncio` is the default mode to runs the server in the current thread.
    `thread` runs the server in a separate thread. `mp` runs the server in a separate process."""
@@ -49,6 +54,8 @@ class PythonServerLauncherArgs:
    """
    log_level: int = logging.INFO
    """The log level to use."""
+    access_log: bool = False
+    """Whether to turn on access logs."""
    startup_timeout: float = 60.0
    """The timeout to wait for the server to start up."""
    kill_unhealthy_server: bool = True
@@ -59,6 +66,8 @@ class PythonServerLauncherArgs:
    """The timeout to wait for the thread to join."""
    process_join_timeout: float = 10.0
    """The timeout to wait for the process to join."""
+    timeout_keep_alive: int = 30
+    """The timeout to keep the connection alive."""


@dataclass
@@ -152,21 +161,30 @@ async def run_uvicorn_asyncio(

        if not uvicorn_server.started:
            # Normally, the program will not reach this point, as the server will throw the exception itself earlier.
-            raise RuntimeError(f"Server did not start up within {timeout:.2f} seconds.") from server_start_exception
+            raise RuntimeError(
+                f"Server did not start up within {time.time() - start_time:.2f} seconds."
+            ) from server_start_exception

-        logger.debug(f"Server started up in {time.time() - start_time:.2f} seconds.")
+        logger.info(f"Server started up in {time.time() - start_time:.2f} seconds.")

        # Check for health endpoint status if provided
        if health_url is not None:
+            logger.info(f"Probing health endpoint {health_url}...")
            async with aiohttp.ClientSession() as session:
                while time.time() < deadline:
-                    with suppress(Exception):
+                    try:
                        async with session.get(health_url) as resp:
                            if resp.status == 200:
-                                logger.debug(
+                                logger.info(
                                    f"Server is healthy at {health_url} in {time.time() - start_time:.2f} seconds."
                                )
                                return
+                            else:
+                                logger.debug(
+                                    f"Server is NOT healthy at {health_url} in {time.time() - start_time:.2f} seconds. Got status {resp.status}."
+                                )
+                    except Exception as e:
+                        logger.debug(f"Error probing health endpoint {health_url}: {str(e)}")
                    await asyncio.sleep(0.1)

            # If the server is not healthy, kill it if requested.
@@ -187,7 +205,7 @@ async def run_uvicorn_asyncio(
                )

        else:
-            logger.debug("Server does not provide a health check endpoint. Skipping health check.")
+            logger.info("Server does not provide a health check endpoint. Skipping health check.")

    async def _serve_server() -> None:
        nonlocal server_start_exception
@@ -555,6 +573,27 @@ def run_gunicorn(
        watchdog_thread.join(timeout=5.0)


+def _get_default_ipv4_address() -> str:
+    """Determine the default outbound IPv4 address for this machine.
+
+    Implementation:
+        Opens a UDP socket and "connects" to a public address to force route
+        selection, then inspects the socket's local address. No packets are sent.
+
+    Returns:
+        str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except Exception:
+        return "127.0.0.1"
+    finally:
+        s.close()
+
+
 class PythonServerLauncher:
    """Unified launcher for FastAPI, using uvicorn or gunicorn per mode/worker count.

@@ -573,7 +612,16 @@ class PythonServerLauncher:
        self.app = app
        self.args = args
        self.serve_context = serve_context
+        self._host: Optional[str] = self.args.host
        self._port: Optional[int] = self.args.port
+        self._access_host: Optional[str] = self.args.access_host
+        self.initialize()
+
+    def initialize(self):
+        # ensure the host/port/access_host are set
+        self._ensure_host()
+        self._ensure_port()
+        self._ensure_access_host()

        # uvicorn (in-proc asyncio)
        self._uvicorn_server: Optional[uvicorn.Server] = None
@@ -592,17 +640,35 @@ class PythonServerLauncher:
        # is_running flag
        self._is_running: bool = False

+    def __getstate__(self):
+        """Control pickling to prevent server state from being sent to subprocesses."""
+        return {
+            "app": self.app,
+            "args": self.args,
+            "serve_context": self.serve_context,
+            "_host": self._host,
+            "_port": self._port,
+            "_access_host": self._access_host,
+        }
+
+    def __setstate__(self, state: Dict[str, Any]):
+        self.app = state["app"]
+        self.args = cast(PythonServerLauncherArgs, state["args"])
+        self.serve_context = state["serve_context"]
+        self._host = state["_host"]
+        self._port = state["_port"]
+        self._access_host = state["_access_host"]
+        self.initialize()
+
    @property
    def endpoint(self) -> str:
        """Return the externally advertised host:port pair regardless of accessibility."""
-        return f"http://{self.args.host}:{self._ensure_port()}"
+        return f"http://{self._ensure_host()}:{self._ensure_port()}"

    @property
-    def access_url(self) -> str:
+    def access_endpoint(self) -> str:
        """Return a loopback-friendly URL so health checks succeed even when binding to 0.0.0.0."""
-        # Probe host normalization for 0.0.0.0
-        host_for_probe = "127.0.0.1" if self.args.host in ("0.0.0.0", "::") else self.args.host
-        return f"http://{host_for_probe}:{self._ensure_port()}"
+        return f"http://{self._ensure_access_host()}:{self._ensure_port()}"

    @property
    def health_url(self) -> Optional[str]:
@@ -612,7 +678,7 @@ class PythonServerLauncher:
        path = self.args.healthcheck_url
        if not path.startswith("/"):
            path = "/" + path
-        return f"{self.access_url}{path}"
+        return f"{self.access_endpoint}{path}"

    async def start(self):
        """Starts the server according to launch_mode and n_workers."""
@@ -699,19 +765,41 @@ class PythonServerLauncher:
            return f"{module}:app"
        return "unknown:app"

+    def _ensure_host(self) -> str:
+        if self._host is None:
+            logger.warning("No host provided, using 0.0.0.0.")
+            self._host = "0.0.0.0"
+        return self._host
+
    def _ensure_port(self) -> int:
        if self._port is None:
            logger.warning("No port provided, using pick_unused_port to pick a random unused port.")
            self._port = pick_unused_port()
        return self._port

+    def _ensure_access_host(self) -> str:
+        if self._access_host is None:
+            if self.args.access_host is None:
+                if self._ensure_host() in ("0.0.0.0", "::"):
+                    # Probe host normalization for 0.0.0.0
+                    logger.warning("No access host provided, using default outbound IPv4 address for this machine.")
+                    self._access_host = _get_default_ipv4_address()
+                else:
+                    logger.warning("No access host provided, using the host provided.")
+                    self._access_host = self._ensure_host()
+            else:
+                self._access_host = self.args.access_host
+        return self._access_host  # type: ignore
+
    def _create_uvicorn_server(self) -> uvicorn.Server:
        config = uvicorn.Config(
            app=self.app,
-            host=self.args.host,
+            host=self._ensure_host(),
            port=self._ensure_port(),
            log_level=self.args.log_level,
+            access_log=self.args.access_log,
            loop="asyncio",
+            timeout_keep_alive=self.args.timeout_keep_alive,
        )
        return uvicorn.Server(config)

@@ -782,15 +870,20 @@ class PythonServerLauncher:
        try:
            evt: ChildEvent = await asyncio.to_thread(self._thread_event_queue.get, True, timeout)
        except queue.Empty:
-            logger.error("Threaded server failed to start and sends no event. This should not happen.")
+            if not self._thread.is_alive():
+                raise RuntimeError("Threaded server failed to start and is not alive. No error event was received.")
+            logger.error(
+                "Threaded server failed to start and sends no event. This should not happen. Shutting down server."
+            )
            await self._stop_uvicorn_thread()
-            return
+            raise RuntimeError("Threaded server failed to start and sends no event. This should not happen.")

        if evt.kind == "error":
            logger.error("Threaded server failed to start (%s): %s\n%s", evt.exc_type, evt.message, evt.traceback)
            await asyncio.to_thread(self._thread.join, self.args.thread_join_timeout)
            if self._thread.is_alive():
-                raise RuntimeError(evt.message or "Threaded server failed to start and refused to shut down.")
+                logger.error("Threaded server failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
        else:
            logger.info("Threaded server started successfully.")
            self._is_running = True
@@ -819,6 +912,7 @@ class PythonServerLauncher:
        if self.is_running():
            raise RuntimeError("Server process is already running. Stopping it first.")

+        host = self._ensure_host()
        port = self._ensure_port()

        try:
@@ -834,17 +928,22 @@ class PythonServerLauncher:
        if self.args.n_workers > 1:
            logger.info(f"Starting Gunicorn server...")
            options = {
-                "bind": f"{self.args.host}:{port}",
+                "bind": f"{host}:{port}",
                "workers": int(self.args.n_workers),
                "worker_class": "uvicorn_worker.UvicornWorker",
                "loglevel": logging.getLevelName(self.args.log_level).lower(),
-                "accesslog": None,
+                "accesslog": "-" if self.args.access_log else None,
                "errorlog": "-",
                "preload_app": True,
                "graceful_timeout": int(
                    self.args.process_join_timeout / 2
                ),  # Allow half the timeout for graceful shutdown
            }
+            if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+                from agentlightning.utils.metrics import shutdown_metrics
+
+                options["child_exit"] = shutdown_metrics  # type: ignore
+
            self._gunicorn_app = GunicornApp(self.app, options)

            self._proc = ctx.Process(
@@ -883,9 +982,13 @@ class PythonServerLauncher:
        try:
            evt: ChildEvent = await asyncio.to_thread(self._mp_event_queue.get, True, timeout)
        except queue.Empty:
-            logger.error("Server process failed to start and sends no event. This should not happen.")
+            if not self._proc.is_alive():
+                raise RuntimeError("Server process failed to start and is not alive. No error event was received.")
+            logger.error(
+                "Server process failed to start and sends no event. This should not happen. Shutting down server."
+            )
            await self._stop_serving_process()
-            return
+            raise RuntimeError("Server process failed to start and sends no event. This should not happen.")

        if evt.kind == "error":
            logger.error(
@@ -897,7 +1000,8 @@ class PythonServerLauncher:
            )
            await asyncio.to_thread(self._proc.join, self.args.process_join_timeout)
            if self._proc.is_alive():
-                raise RuntimeError(evt.message or "Server process failed to start and refused to shut down.")
+                logger.error("Server process failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
        else:
            logger.info("Subprocess server started successfully.")
            self._is_running = True
@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import platform
+import socket
+from contextlib import suppress
+from datetime import datetime
+from typing import Any, Dict, List, cast
+
+import psutil
+from gpustat import GPUStat, GPUStatCollection
+
+
+def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
+    """Capture a snapshot of the system's hardware and software information.
+
+    Args:
+        include_gpu: Whether to include GPU information.
+
+    Returns:
+        A dictionary containing the system's hardware and software information.
+    """
+    # CPU
+    cpu = {
+        "cpu_name": platform.processor(),
+        "cpu_cores": psutil.cpu_count(logical=False),
+        "cpu_threads": psutil.cpu_count(logical=True),
+        "cpu_usage_pct": psutil.cpu_percent(0.0),
+    }
+
+    # Memory
+    vm = psutil.virtual_memory()
+    mem = {
+        "mem_used_gb": round(vm.used / (2**30), 2),
+        "mem_total_gb": round(vm.total / (2**30), 2),
+        "mem_pct": vm.percent,
+    }
+
+    # Disk
+    du = psutil.disk_usage("/")
+    disk = {
+        "disk_used_gb": round(du.used / (2**30), 2),
+        "disk_total_gb": round(du.total / (2**30), 2),
+        "disk_pct": du.percent,
+    }
+
+    # GPU (only query if explicitly requested)
+    gpus: List[Dict[str, Any]] = []
+    if include_gpu:
+        with suppress(Exception):
+            for g in GPUStatCollection.new_query().gpus:  # type: ignore
+                g = cast(GPUStat, g)
+                gpus.append(
+                    {
+                        "gpu": g.name,  # type: ignore
+                        "util_pct": g.utilization,
+                        "mem_used_mb": g.memory_used,
+                        "mem_total_mb": g.memory_total,
+                        "temp_c": g.temperature,
+                    }
+                )
+
+    # Network
+    net = psutil.net_io_counters()
+    netinfo = {
+        "bytes_sent_mb": round(net.bytes_sent / (2**20), 2),
+        "bytes_recv_mb": round(net.bytes_recv / (2**20), 2),
+    }
+
+    # OS / meta
+    return {
+        "timestamp": datetime.now().isoformat(timespec="seconds"),
+        "host": socket.gethostname(),
+        "os": platform.platform(),
+        **cpu,
+        **mem,
+        **disk,
+        **netinfo,
+        **({"gpus": gpus} if include_gpu else {}),
+    }
@@ -8,6 +8,12 @@ defaults:

 agentlightning:
  port: 9999
+  trace_aggregator:
+    level: transition  # transition or trajectory, docs refer to https://agent-lightning.github.io/posts/trajectory_level_aggregation/
+    trajectory_max_prompt_length: 2048  # supported in trajectory level aggregation, suggest to set as maximum length for the prompt in first turn
+    trajectory_max_response_length: 8192  # supported in trajectory level aggregation, suggest to set as maximum length for the cumulative agent responses in the full trajectory, i.e., n_turns * (max_response_length + max_prompt_length)
+    debug: False  # supported in trajectory level aggregation, enable to diagnose trace merging failures
+    mismatch_log_dir: ./mismatch_cases  # supported in trajectory level aggregation with debug=True, directory to store logs of mismatch cases

 data:
  filter_overlong_prompts: false
@@ -2,6 +2,7 @@

 import asyncio
 import json
+import os
 import random
 import socket
 import threading
@@ -9,7 +10,7 @@ import time
 import uuid
 from collections import defaultdict
 from collections.abc import Mapping
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple, cast

 import numpy as np
 import requests
@@ -18,13 +19,11 @@ from flask import Flask, Response, abort, request
 from tensordict import TensorDict
 from verl import DataProto

-from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, configure_logger
+from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy
 from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store.base import LightningStore
-from agentlightning.types import Rollout, RolloutConfig, Task
-
-configure_logger()
+from agentlightning.types import EnqueueRolloutRequest, Rollout, RolloutConfig, Task

 __all__ = [
    "AgentModeDaemon",
@@ -33,6 +32,85 @@ __all__ = [
 ]


+def ids_startswith(
+    full_ids: List[int], prefix_ids: List[int], tokenizer: Any, debug: bool = False
+) -> Tuple[bool, Tuple[bool, bool, bool]]:
+    is_prefix: bool
+    template_mismatch, retoken_mismatch, others_mismatch = False, False, False
+    if full_ids[: len(prefix_ids)] == prefix_ids:
+        is_prefix = True
+        return True, (template_mismatch, retoken_mismatch, others_mismatch)
+    else:
+        is_prefix = False
+
+    if not debug:
+        return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
+
+    def _special_token_sequence(ids: List[int]) -> List[int]:
+        return [id for id in ids if id in tokenizer.all_special_ids]
+
+    def _none_special_token_sequence(ids: List[int]) -> List[int]:
+        return [id for id in ids if id not in tokenizer.all_special_ids]
+
+    # First, handle special tokens
+    full_special_ids = _special_token_sequence(full_ids)
+    prefix_special_ids = _special_token_sequence(prefix_ids)
+    if sum(1 for a, b in zip(full_special_ids, prefix_special_ids) if a != b) > 0:
+        template_mismatch = True
+
+    # Next, handle string content
+    full_content_ids = _none_special_token_sequence(full_ids)
+    prefix_content_ids = _none_special_token_sequence(prefix_ids)
+    full_string = tokenizer.decode(full_ids, skip_special_tokens=True)
+    prefix_string = tokenizer.decode(prefix_ids, skip_special_tokens=True)
+    if full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and full_string.startswith(prefix_string):
+        retoken_mismatch = True
+    elif full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and not full_string.startswith(
+        prefix_string
+    ):
+        others_mismatch = True
+    return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
+
+
+def log_mismatch_detail(
+    diagnostic: Tuple[bool, bool, bool],
+    full_ids: List[int],
+    prefix_ids: List[int],
+    global_steps: int,
+    rollout_id: str,
+    turn_id: int,
+    log_dir: str | None = None,
+):
+    if log_dir is None:
+        return
+    os.makedirs(log_dir, exist_ok=True)
+    template_mismatch, retoken_mismatch, others_mismatch = diagnostic
+    if template_mismatch:
+        with open(os.path.join(log_dir, "template_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+    if retoken_mismatch:
+        with open(os.path.join(log_dir, "retoken_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+    if others_mismatch:
+        with open(os.path.join(log_dir, "others_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+
+
 def get_left_padded_ids_and_attention_mask(
    ids: List[int], max_length: int, pad_token_id: int
 ) -> Tuple[List[int], List[int]]:
@@ -146,6 +224,9 @@ class AgentModeDaemon:
        llm_proxy: LLMProxy | None = None,
        store: LightningStore | None = None,
        adapter: TraceToTripletBase | None = None,
+        processor: Any = None,
+        image_base_dir: Optional[str] = None,
+        trace_aggregator: Dict[str, Any] = {"level": "transition"},
    ):
        self.mode = mode
        self.llm_timeout_seconds = llm_timeout_seconds
@@ -185,7 +266,13 @@ 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
+        self.trace_aggregator = trace_aggregator
+
+        # 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] = []
@@ -204,6 +291,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.
@@ -294,7 +450,7 @@ class AgentModeDaemon:
        self._proxy_thread.start()
        print(f"Proxy server running on port {self.proxy_port}")

-    def _update_proxy_server_v1(self):
+    async def _update_proxy_server_v1(self):
        model_name = self.train_information.get("model")
        if not model_name:
            raise ValueError("Model name is not set.")
@@ -313,12 +469,7 @@ class AgentModeDaemon:
            ],
        )

-        if self.llm_proxy.is_running():
-            # FIXME: Need to switch to a different port right now
-            # because the forked processes carried the old fd
-            self.llm_proxy.restart(_port=_find_available_port())
-        else:
-            self.llm_proxy.start()
+        await self.llm_proxy.restart()

    def start(self):
        """Starts the main AgentLightningServer and the proxy server."""
@@ -352,7 +503,7 @@ class AgentModeDaemon:
        if server_addresses != self.backend_llm_server_addresses:
            self.backend_llm_server_addresses = server_addresses
            if self.mode == "v1" and not self.llm_proxy.is_running():
-                self._update_proxy_server_v1()
+                await self._update_proxy_server_v1()
        self.is_train = is_train

        # 1. Update resources on the server for clients to use
@@ -384,42 +535,57 @@ class AgentModeDaemon:
        num_samples = len(data[keys[0]])
        rollouts_per_sample = self.train_rollout_n if is_train else 1

+        enqueue_rollout_requests: List[EnqueueRolloutRequest] = []
+        data_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
+
        for i in range(num_samples):
            data_id = str(uuid.uuid4())
            original_sample = {key: data[key][i] for key in keys}
            original_sample["data_id"] = data_id
+            data_id_to_original_sample[data_id] = original_sample

            # For training, each sample is rolled out multiple times
+            # Data ID is different from Rollout ID, as one data can have multiple rollouts.
            for _ in range(rollouts_per_sample):
                task_metadata = {"data_id": data_id, "is_train": is_train}
-
-                # Data ID is different from Rollout ID, as one data can have multiple rollouts.
                if self.mode == "v0":
+                    # Queue immediately
                    rollout_id = await self.server.queue_task(
                        sample=_to_native(original_sample),
                        mode="train" if is_train else "val",
                        resources_id=resources_id,
                        metadata=task_metadata,
                    )
-                else:
-                    rollout = await self.store.enqueue_rollout(
-                        input=_to_native(original_sample),
-                        mode="train" if is_train else "val",
-                        resources_id=resources_id,
-                        metadata=task_metadata,
-                    )
-                    await self.store.update_rollout(
-                        rollout_id=rollout.rollout_id,
-                        config=RolloutConfig(
-                            unresponsive_seconds=self.llm_timeout_seconds,
-                            timeout_seconds=self.llm_timeout_seconds,
-                        ),
-                    )
-                    rollout_id = rollout.rollout_id

-                # Store original sample data to reconstruct batch information later
-                self._task_id_to_original_sample[rollout_id] = original_sample
-                self._total_tasks_queued += 1
+                    # Store original sample data to reconstruct batch information later
+                    self._task_id_to_original_sample[rollout_id] = original_sample
+                    self._total_tasks_queued += 1
+                else:
+                    # Collect tasks to enqueue in batch and queue them later
+                    enqueue_rollout_requests.append(
+                        EnqueueRolloutRequest(
+                            input=_to_native(original_sample),
+                            mode="train" if is_train else "val",
+                            resources_id=resources_id,
+                            config=RolloutConfig(
+                                unresponsive_seconds=self.llm_timeout_seconds,
+                                timeout_seconds=self.llm_timeout_seconds,
+                            ),
+                            metadata=task_metadata,
+                        )
+                    )
+
+        if self.mode == "v1":
+            # Enqueue all the tasks in a single batch
+            rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
+            self._task_id_to_original_sample.update(
+                {
+                    # Recover the original data and store it for later use.
+                    rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
+                    for rollout in rollouts
+                }
+            )
+            self._total_tasks_queued += len(rollouts)

    def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
        """Synchronous wrapper for setting up data and server resources."""
@@ -436,7 +602,7 @@ class AgentModeDaemon:
                raise RuntimeError("Internal loop is not running.")
            future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
        try:
-            future.result(timeout=60)  # Wait for completion with a timeout
+            future.result(timeout=300)  # Wait for completion with a timeout
        except Exception as e:
            print(f"Failed to set up data on server: {e}")
            raise
@@ -564,14 +730,17 @@ class AgentModeDaemon:
        )  # FIXME: Evaluate whether grouping stats by source is actually needed.

        for rollout_id, rollout in self._completed_rollouts_v0.items():
+            final_reward_raw: Optional[float] = rollout.final_reward
            final_reward = self._fillna_reward(rollout)
            if not rollout.triplets:
                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
-                sample_stat_list.append({"reward": final_reward})
+                sample_stat_list.append({"reward": final_reward, "has_reward": final_reward_raw is not None})
                continue
            response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
+
            if "data_source" in self._task_id_to_original_sample[rollout_id]:
                # When a test sample includes a 'data_source' field, record per-source statistics for test results.
+                # TODO: This is a flawed design. We should have a better way to handle this.
                data_source = self._task_id_to_original_sample[rollout_id]["data_source"]
                sample_stat_list_by_source[data_source].append(
                    {
@@ -579,6 +748,7 @@ class AgentModeDaemon:
                        "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
                        "turn_count": len(rollout.triplets),
                        "reward": final_reward,
+                        "has_reward": final_reward_raw is not None,
                    }
                )
            sample_stat_list.append(
@@ -587,6 +757,7 @@ class AgentModeDaemon:
                    "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
                    "turn_count": len(rollout.triplets),
                    "reward": final_reward,
+                    "has_reward": final_reward_raw is not None,
                }
            )
        metric_dict: Dict[str, Any] = {}
@@ -601,6 +772,9 @@ class AgentModeDaemon:
                {
                    f"val/{data_source}/n_rollouts": len(sample_stats),
                    f"val/{data_source}/n_rollouts_w_trace": len(stats_w_trace_by_source[data_source]),
+                    f"val/{data_source}/n_rollouts_w_reward": len(
+                        [stat for stat in sample_stats if stat["has_reward"]]
+                    ),
                    f"val/{data_source}/reward": np.mean(
                        [stat["reward"] for stat in sample_stats]
                    ),  # each rollout must have a reward (fillna if missing)
@@ -619,6 +793,7 @@ class AgentModeDaemon:
            {
                "val/n_rollouts": len(sample_stat_list),
                "val/n_rollouts_w_trace": len(stats_w_trace),
+                "val/n_rollouts_w_reward": len([stat for stat in sample_stat_list if stat["has_reward"]]),
                "val/reward": np.mean(
                    [stat["reward"] for stat in sample_stat_list]
                ),  # each rollout must have a reward (fillna if missing)
@@ -629,7 +804,9 @@ class AgentModeDaemon:
        )
        return metric_dict

-    def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
+    def get_train_data_batch(
+        self, max_prompt_length: int, max_response_length: int, device: torch.device, global_steps: int
+    ):
        """
        Processes completed rollouts to generate a training data batch.

@@ -643,9 +820,10 @@ class AgentModeDaemon:
        # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
        finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {}
        finished_id_to_final_reward: Dict[str, float] = {}
+        sample_with_reward_count = 0
        for rollout_id, rollout in self._completed_rollouts_v0.items():
            original_sample = self._task_id_to_original_sample[rollout_id]
-
+            sample_with_reward_count += int(rollout.final_reward is not None)
            final_reward = self._fillna_reward(rollout)

            if not rollout.triplets:
@@ -654,10 +832,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 = {
@@ -687,60 +869,204 @@ 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():
-            for turn_index, trace in enumerate(sample_info["trace_list"]):
+        if self.trace_aggregator.get("level", "transition") == "transition":
+            for rollout_id, sample_info in finished_id_to_sample_info.items():
+                for turn_index, trace in enumerate(sample_info["trace_list"]):

-                reward_list.append(sample_info["reward"])
-                prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
+                    reward_list.append(sample_info["reward"])
+                    prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]

-                # Mark samples with prompts exceeding max_prompt_length to be dropped later
-                if len(prompt_ids) > max_prompt_length:
-                    prompt_ids = prompt_ids[:max_prompt_length]
-                    is_drop_list.append(True)
-                else:
-                    is_drop_list.append(False)
+                    # Mark samples with prompts exceeding max_prompt_length to be dropped later
+                    if len(prompt_ids) > max_prompt_length:
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        is_drop_list.append(True)
+                    else:
+                        is_drop_list.append(False)

-                # Truncate responses that exceed max_response_length
-                if len(response_ids) > max_response_length:
-                    response_ids = response_ids[:max_response_length]
-                    n_trunc_sample_because_of_response += 1
+                    # Truncate responses that exceed max_response_length
+                    if len(response_ids) > max_response_length:
+                        response_ids = response_ids[:max_response_length]
+                        n_trunc_sample_because_of_response += 1

-                # Pad prompts to the left and responses to the right
-                one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
-                    prompt_ids, max_prompt_length, self.pad_token_id
-                )
-                one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
-                    response_ids, max_response_length, self.pad_token_id
-                )
+                    # Pad prompts to the left and responses to the right
+                    one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
+                        prompt_ids, max_prompt_length, self.pad_token_id
+                    )
+                    one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
+                        response_ids, max_response_length, self.pad_token_id
+                    )

-                input_ids_list.append(one_input_ids)
-                input_attention_mask_list.append(one_input_attention_mask)
-                response_ids_list.append(one_response_ids)
-                response_attention_mask_list.append(one_response_attention_mask)
-                data_id_list.append(sample_info["data_id"])
-                rollout_id_list.append(rollout_id)
-                turn_index_list.append(turn_index)
+                    input_ids_list.append(one_input_ids)
+                    input_attention_mask_list.append(one_input_attention_mask)
+                    response_ids_list.append(one_response_ids)
+                    response_attention_mask_list.append(one_response_attention_mask)
+                    data_id_list.append(sample_info["data_id"])
+                    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))
+
+        elif self.trace_aggregator.get("level", "transition") == "trajectory":
+            assert not self._use_mrope, "M-RoPE is not supported in trajectory level yet."
+
+            response_mask_list: List[List[int]] = []
+            unmerged_count: int = 0
+            template_mismatch_count, retoken_mismatch_count, others_mismatch_count = 0, 0, 0
+            response_per_turn_list: List[int] = []
+
+            for rollout_id, sample_info in finished_id_to_sample_info.items():
+                merged_trace_idx: List[List[int]] = []
+
+                # Identify which turns can be merged based on token ids prefix matching
+                current_merged_trace_idx: List[int] = []
+                current_context: List[int] = []
+                for turn_index, trace in enumerate(sample_info["trace_list"]):
+                    response_per_turn_list.append(len(trace["response_ids"]))
+                    is_prefix, diagnostic = ids_startswith(
+                        trace["prompt_ids"] + trace["response_ids"],
+                        current_context,
+                        self.tokenizer,
+                        self.trace_aggregator.get("debug", False),
+                    )
+                    if not is_prefix and self.trace_aggregator.get("debug", False) == True:
+                        template_mismatch_count += diagnostic[0]
+                        retoken_mismatch_count += diagnostic[1]
+                        others_mismatch_count += diagnostic[2]
+                        log_mismatch_detail(
+                            diagnostic,
+                            trace["prompt_ids"] + trace["response_ids"],
+                            current_context,
+                            global_steps,
+                            rollout_id,
+                            turn_index,
+                            self.trace_aggregator.get("mismatch_log_dir", None),
+                        )
+
+                    if is_prefix:
+                        current_context = trace["prompt_ids"] + trace["response_ids"]
+                        current_merged_trace_idx.append(turn_index)
+                    else:
+                        merged_trace_idx.append(current_merged_trace_idx)
+                        current_merged_trace_idx = [turn_index]
+                        current_context = trace["prompt_ids"] + trace["response_ids"]
+
+                if current_merged_trace_idx not in merged_trace_idx:
+                    merged_trace_idx.append(current_merged_trace_idx)
+
+                if len(merged_trace_idx) > 1:
+                    unmerged_count += 1
+
+                # Merge all trace segments in merged_trace_idx into training samples
+                for current_merged_trace_idx in merged_trace_idx:
+                    prompt_ids = sample_info["trace_list"][current_merged_trace_idx[0]]["prompt_ids"]
+
+                    # if the merged_trace_idx doesn't start with the beginning of the prompt_ids, we need to adjust it
+                    if current_merged_trace_idx[0] > 0 and len(prompt_ids) > max_prompt_length:
+                        response_ids = prompt_ids[max_prompt_length:]
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        response_mask = [1] * len(response_ids)
+                    else:
+                        response_ids = []
+                        response_mask = []
+
+                    prompt_length = len(prompt_ids)
+                    response_ids += sample_info["trace_list"][current_merged_trace_idx[0]]["response_ids"]
+                    response_mask += [1] * len(response_ids)
+                    for turn_index in current_merged_trace_idx[1:]:
+                        trace = sample_info["trace_list"][turn_index]
+                        new_prompt_length = len(trace["prompt_ids"]) - len(response_ids) - prompt_length
+                        response_ids += trace["prompt_ids"][-new_prompt_length:]
+                        response_ids += trace["response_ids"]
+                        response_mask += [0] * new_prompt_length
+                        response_mask += [1] * len(trace["response_ids"])
+
+                    reward_list.append(sample_info["reward"])
+
+                    # Mark samples with prompts exceeding max_prompt_length to be dropped later
+                    if len(prompt_ids) > max_prompt_length:
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        is_drop_list.append(True)
+                    else:
+                        is_drop_list.append(False)
+
+                    # Truncate responses that exceed max_response_length
+                    if len(response_ids) > max_response_length:
+                        response_ids = response_ids[:max_response_length]
+                        response_mask = response_mask[:max_response_length]
+                        n_trunc_sample_because_of_response += 1
+
+                    # Pad prompts to the left and responses to the right
+                    one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
+                        prompt_ids, max_prompt_length, self.pad_token_id
+                    )
+                    one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
+                        response_ids, max_response_length, self.pad_token_id
+                    )
+                    one_response_mask, _ = get_right_padded_ids_and_attention_mask(
+                        response_mask, max_response_length, 0
+                    )
+
+                    input_ids_list.append(one_input_ids)
+                    input_attention_mask_list.append(one_input_attention_mask)
+                    response_ids_list.append(one_response_ids)
+                    response_attention_mask_list.append(one_response_attention_mask)
+                    response_mask_list.append(one_response_mask)
+                    data_id_list.append(sample_info["data_id"])
+                    rollout_id_list.append(rollout_id)
+                    # turn_index_list.append(current_merged_trace_idx)
+        else:
+            raise ValueError(f"Unknown trace_aggregator level: {self.trace_aggregator.get('level')}")

        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)
        batch_response_ids = torch.LongTensor(response_ids_list).to(device)
        response_attention_mask = torch.LongTensor(response_attention_mask_list).to(device)
+        response_mask = (
+            torch.LongTensor(response_mask_list).to(device) if self.trace_aggregator.get("level", "transition") == "trajectory" else None  # type: ignore
+        )

        # 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:]
@@ -755,7 +1081,12 @@ class AgentModeDaemon:
                "position_ids": position_ids,
                "is_drop_mask": is_drop_mask,
                "token_level_scores": token_level_scores.contiguous(),
-            },
+                **(
+                    {"response_mask": response_mask}
+                    if self.trace_aggregator.get("level", "transition") == "trajectory"
+                    else {}
+                ),
+            },  # type: ignore
            batch_size=n_transition,
        )
        data_proto = DataProto(batch=batch)
@@ -764,14 +1095,41 @@ class AgentModeDaemon:
            "training/reward": np.mean(list(finished_id_to_final_reward.values())),
            "training/n_rollouts": len(finished_id_to_final_reward),
            "training/n_rollouts_w_trace": len(finished_id_to_sample_info),
+            "training/n_rollouts_w_reward": sample_with_reward_count,
            "training/n_truncated_triplets": n_trunc_sample_because_of_response,
            "training/n_triplets": n_transition,
+            # log data, only for debug testing
+            **(
+                {
+                    "training/n_unmerged_rollouts": unmerged_count,  # type: ignore
+                    "training/n_triplets_by_turn": len(response_per_turn_list),  # type: ignore
+                    "training/avg_response_length_by_turn": np.mean(response_per_turn_list),  # type: ignore
+                    "training/max_response_length_by_turn": np.max(response_per_turn_list),  # type: ignore
+                    "training/min_response_length_by_turn": np.min(response_per_turn_list),  # type: ignore
+                }
+                if self.trace_aggregator.get("level", "transition") == "trajectory"
+                else {}
+            ),
+            **(
+                {
+                    "training/template_mismatch_triplets": template_mismatch_count,  # type: ignore
+                    "training/retoken_mismatch_triplets": retoken_mismatch_count,  # type: ignore
+                    "training/others_mismatch_triplets": others_mismatch_count,  # type: ignore
+                    "training/template_mismatch_ratio": template_mismatch_count / len(response_per_turn_list),  # type: ignore
+                    "training/retoken_mismatch_ratio": retoken_mismatch_count / len(response_per_turn_list),  # type: ignore
+                    "training/others_mismatch_ratio": others_mismatch_count / len(response_per_turn_list),  # type: ignore
+                }
+                if self.trace_aggregator.get("level", "transition") == "trajectory"
+                and self.trace_aggregator.get("debug", False)
+                else {}
+            ),
        }

        # Add non-tensor data for advantage calculation and logging
        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)  # type: ignore
        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)  # type: ignore
-        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
+        if self.trace_aggregator.get("level", "transition") == "transition":
+            data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore

        return data_proto, data_metrics

@@ -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,10 +8,11 @@ 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
+import verl
 from codetiming import Timer
 from omegaconf import OmegaConf
 from tqdm import tqdm
@@ -173,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."
@@ -198,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)
@@ -217,9 +255,18 @@ class AgentLightningTrainer(RayPPOTrainer):
                )
                self.agent_mode_daemon.run_until_all_finished()
                batch, agent_metrics = self.agent_mode_daemon.get_train_data_batch(
-                    max_prompt_length=self.config.data.max_prompt_length,
-                    max_response_length=self.config.data.max_response_length,
+                    max_prompt_length=(
+                        self.config.agentlightning.trace_aggregator.trajectory_max_prompt_length
+                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
+                        else self.config.data.max_prompt_length
+                    ),
+                    max_response_length=(
+                        self.config.agentlightning.trace_aggregator.trajectory_max_response_length
+                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
+                        else self.config.data.max_response_length
+                    ),
                    device=gen_batch.batch["fake_ids"].device,
+                    global_steps=self.global_steps,
                )
                metrics.update(agent_metrics)
                self.agent_mode_daemon.clear_data_and_server()
@@ -244,7 +291,8 @@ class AgentLightningTrainer(RayPPOTrainer):
            # uid is used for algorithm like GRPO, should be aligned to data id
            batch.non_tensor_batch["uid"] = batch.non_tensor_batch["data_id_list"]

-            batch.batch["response_mask"] = compute_response_mask(batch)
+            if "response_mask" not in batch.batch:
+                batch.batch["response_mask"] = compute_response_mask(batch)

            # compute global_valid tokens
            batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
@@ -275,7 +323,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
@@ -403,14 +451,20 @@ class AgentLightningTrainer(RayPPOTrainer):
        assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
        if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
            raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
-        self.agent_mode_daemon = AgentModeDaemon(
+        verl_version = verl.__version__
+        if verl_version == "0.5.0":
+            # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
+            # However, it is possible that verl updates the naming and causes incompatibility.
+            # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
+            model = "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:])
+        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 = self.daemon_cls(
            self.config.agentlightning.port,
            self.config.actor_rollout_ref.rollout.n,
            train_information={
-                # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
-                # However, it is possible that verl updates the naming and causes incompatibility.
-                # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
-                "model": "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:]),
+                "model": model,
                "temperature": self.config.actor_rollout_ref.rollout.temperature,
            },
            tokenizer=self.tokenizer,
@@ -420,6 +474,9 @@ 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),
+            trace_aggregator=self.config.agentlightning.trace_aggregator,
        )
        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.2.x**.

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

@@ -14,7 +14,7 @@ The example is designed to run on a single node with 8 GPUs, each having at leas
 | `retrieval_launch.sh` | Launches the retrieval service backed by the processed corpus |
 | `retrieval_server.py` | FastAPI server that powers document retrieval during training |
 | `search_r1_agent.py` | Agent-Lightning rollout script implementing the Search-R1 workflow |
-| `train.sh` | Starts the RL training server that coordinates GRPO optimization |
+| `train_search_r1_agent.py` | RL training script that coordinates GRPO optimization |
 | `qa_em.py` | Exact-match evaluation utilities for validating model predictions |

 ---
@@ -54,7 +54,7 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv

 ---

-## Run RL Training (GRPO) with Llama-3.2-3b-base
+## Run RL Training (GRPO) with Llama-3.2-3B-Instruct

 1. **Start Ray**

@@ -65,26 +65,28 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv
   > If you plan to use WandB for experiment tracking, set the environment variable
   > `WANDB_API_KEY` before starting Ray.

-2. **Launch the Agent**
-
-   ```bash
-   python search_r1_agent.py
-   ```
-
-   This script automatically launches **128 agent workers** by default. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.
-
-
-3. **Start the Training Server**
+2. **Start the Training Server**
   In another terminal, run:

   ```bash
-   bash train.sh
+   python train_search_r1_agent.py llama
   ```

-   This script starts the RL training server.
+   This script starts the RL training. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.

 ---

-## Evaluation
+## Benchmark Results

-Evaluation scripts and benchmark results will be released soon.
+We evaluated Search-R1 across seven diverse question-answering benchmarks, covering both General QA (NQ, TriviaQA, PopQA) and complex multi-hop reasoning tasks (HotpotQA, 2WikiMultiHopQA, Musique, and Bamboogle).
+
+The following tables compare the performance of the original Search-R1 implementation and the Agent-Lightning version across various base models.
+
+| Model | Source | NQ | TriviaQA | PopQA | HotpotQA | 2Wiki | Musique | Bamboogle |
+| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| **Qwen2.5-3B-Instruct** | **Search-R1 (Original)** | 34.1 | 54.5 | 37.8 | 32.4 | 31.9 | 10.3 | 26.4 |
+| | **Agent-Lightning** | **45.3** | **61.7** | **43.8** | **42.6** | **36.4** | **17.1** | **37.6** |
+| **Qwen2.5-7B-Instruct** | **Search-R1 (Original)** | 39.3 | 61.0 | 39.7 | 37.0 | 41.4 | 14.6 | 36.8 |
+| | **Agent-Lightning** | **46.5** | **65.9** | **46.8** | **43.7** | **46.2** | **20.3** | **47.2** |
+| **Llama-3.2-3B** | **Search-R1 (Reproduced)** | 26.3 | 49.0 | 23.0 | 21.6 | 27.3 | 4.5 | 9.7 |
+| | **Agent-Lightning** | **29.6** | **51.9** | **25.7** | **23.2** | **28.3** | **5.8** | 9.6 |
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yuge Zhang	a1b3150bdc	Bump to 0.3.1	2025-12-24 11:37:57 +08:00
Yuge Zhang	3b5d733861	[Release] v0.3.0 (#427 ) Deploy Documentation / deploy (push) Has been cancelled Details PyPI Release / check-version (push) Has been cancelled Details PyPI Release / publish-pypi (push) Has been cancelled Details	2025-12-24 09:46:58 +08:00
Yuge Zhang	158f5df28e	Fix documentation and dashboard building issues (#429 )	2025-12-23 23:55:59 +08:00
Yuge Zhang	40dc59205b	Scale out benchmark parameters (#428 )	2025-12-23 23:32:19 +08:00
Yuge Zhang	c1a43b6c3a	Update parallelization guides (#426 )	2025-12-23 14:57:24 +08:00
Jiahang Xu	4235731a0d	Feat: Add `trace_aggregator` to support both transition and trajectory aggregation (#134 )	2025-12-22 11:12:21 +08:00
Yuge Zhang	22b80b38bf	v0.3 Documentation Update (#422 )	2025-12-18 01:12:50 +08:00
Yuge Zhang	9f178accaf	Minor optimizations to store benchmark (#421 )	2025-12-18 00:10:22 +08:00
Yuge Zhang	68a47d5087	Make weave import optional (#423 )	2025-12-17 20:31:48 +08:00
Yuge Zhang	4b36b25aad	Fix Weave get username (#420 )	2025-12-17 13:00:24 +08:00
Wang Zilong	a13e09fc6c	add youtu agent blog link in community projects (#416 )	2025-12-17 09:00:04 +08:00
Yuge Zhang	e63c340ebd	Benchmark minor improvements (#418 )	2025-12-17 02:14:32 +08:00
Yuge Zhang	e62b7ca252	Support Weave tracer in TracerTraceToTriplet (#415 )	2025-12-17 00:15:31 +08:00
Jiahang Xu	f66d87745f	Adapt the Search R1 Example to AGL v0.2 (#412 ) Co-authored-by: SiyunZhao <siyunzhao@microsoft.com>	2025-12-16 23:32:20 +08:00
Jiahang Xu	52090e9dd5	Update benchmark results to Search-R1 v0.1 (#417 )	2025-12-16 23:28:10 +08:00
Yuge Zhang	fdaf3f1777	Fix unsloth config issue (#414 )	2025-12-16 10:41:42 +08:00
Yuge Zhang	087c7d350a	Reimplement Weave tracer and unify emitter interface (#411 )	2025-12-15 15:20:59 +08:00
Yuge Zhang	2203070ef0	Misc CI fixes and Move Search-R1 to contrib (#410 )	2025-12-13 15:12:23 +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
Copilot	4adf4e3ea4	Fix sequence ID sorting in traces table (#371 ) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ultmaster <8463288+ultmaster@users.noreply.github.com>	2025-12-06 12:31:42 +08:00
Yuge Zhang	0294eb5d32	GitHub Actions for RAG example (#357 )	2025-12-06 12:04:42 +08:00
Leonardo Pinheiro	f9fe772e10	Update langchain to 1.x (#364 )	2025-12-05 21:35:39 +08:00
Yuge Zhang	9f8a25ffdc	Store Benchmark - Part 4 (#356 )	2025-12-05 12:00:11 +08:00
Yuge Zhang	3082ac0ee0	Centralized metrics helper (#368 )	2025-12-05 08:47:10 +08:00
Yuge Zhang	56e5c7ce62	Operation emitter (#359 )	2025-12-04 15:16:29 +08:00
Yuge Zhang	21892cc6d3	Skip vllm 0.12.0 (#361 )	2025-12-04 14:21:23 +08:00
Wang Zilong	34811cb454	Update RAG example to v0.2.x (#349 )	2025-12-03 15:46:57 +08:00
Yuge Zhang	003b8c6f83	Store Benchmark - Part 3 (#344 )	2025-12-03 01:10:38 +08:00
Yuge Zhang	63b6d42669	Claude Code Example README update (#348 )	2025-12-02 01:01:30 +08:00
Yuge Zhang	8c219175f5	Add CI for Claude Code (#346 )	2025-12-01 23:48:51 +08:00
Yuge Zhang	931ddcfdcc	Store Benchmark - Part 2 (#342 )	2025-11-29 07:32:09 +08:00
Yuge Zhang	ce80b09a4a	Patch LiteLLM root span (#341 )	2025-11-28 11:34:03 +08:00
Yuge Zhang	f0546ca6c5	Semantic Convention (#340 )	2025-11-28 01:22:42 +08:00
Ni Hao	3a3bfeef31	add test code to agentops's tracer (#324 )	2025-11-27 21:25:47 +08:00
Geng Zhang	a733950b74	Support Claude Code as LitAgent (#332 )	2025-11-27 18:39:26 +08:00
Yuge Zhang	662fd90784	Upgrade transformers and CrewAI versions (#336 )	2025-11-26 09:25:31 +08:00
Yuge Zhang	475c2adb91	Add Examples Catalog and Refine Contribution Guide (#331 )	2025-11-23 16:17:16 +00:00
Yuge Zhang	bffc7013f9	Store Benchmark - Part 1 (#328 )	2025-11-22 23:35:29 +08:00
Yuge Zhang	4cf8fb94e7	Github Actions Workflow for Tinker and Azure (#327 )	2025-11-22 01:47:53 +08:00
Yuge Zhang	ab185a5c5a	MongoDB-based Lightning Store (#323 )	2025-11-21 11:49:54 +08:00
Yuge Zhang	d581cbcd63	Upgrade VM image (#325 )	2025-11-20 17:49:16 +08:00
Yuge Zhang	3459caa1de	Fix OpenAI Agents 0.6 compatibility and pin vLLM < 0.11.1 (#322 )	2025-11-20 07:13:15 +08:00
Yuge Zhang	f3fd58e72a	Put store init in the right place of tracer (#321 )	2025-11-19 20:35:27 +08:00
Yuge Zhang	b3cb5e1337	Minor improvements to make RL workflow more robust (#319 )	2025-11-18 15:40:51 +08:00
Yuge Zhang	3761c0f54c	Support native advanced queries in LightningStore (#318 )	2025-11-18 10:54:55 +08:00
Yuge Zhang	d4334182be	Adding check traces with reward for VERL (#317 )	2025-11-17 21:18:15 +08:00
Yuge Zhang	57c3c0525e	Collection-based Lightning Store (#315 )	2025-11-17 18:51:51 +08:00
Yuge Zhang	e356593f73	Bump to 0.3.0 (#316 )	2025-11-17 17:32:42 +08:00
Yuge Zhang	0e033831d5	Support OTLP in LightningStore (#313 )	2025-11-15 16:34:09 +08:00
xiaochulaoban	0d721228d5	Added the README and script files for training sql_agent on NPU (#272 ) Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>	2025-11-15 01:27:07 +08:00
Yuge Zhang	e49b75b7d8	Check all matching jobs per variant (#310 )	2025-11-13 17:10:50 +00:00
Yuge Zhang	eab691b1a1	Refactor logging (#306 )	2025-11-13 22:48:52 +08:00
Yuge Zhang	fd6494873d	Make health timeout configurable (#305 )	2025-11-13 19:46:02 +08:00
Yuge Zhang	6cbfc1fee0	Fix CI Badge and make Calc-X pipeline faster (#304 )	2025-11-13 18:06:18 +08:00
Yuge Zhang	b986ae132a	Use PythonServerLauncher in LightningStoreServer (#303 )	2025-11-13 14:22:54 +08:00
Yuge Zhang	f24a47969e	Increase graceful timeout on CI (#302 )	2025-11-13 10:15:14 +08:00
Yuge Zhang	a0bc1827d9	[Release] v0.2.2 (#298 )	2025-11-12 23:54:35 +08:00
Yuge Zhang	f2869cea30	Fix local model support in VERL (#299 )	2025-11-12 22:56:10 +08:00
Geng Zhang	77cf447717	fix stream response for anthropic and openai api (#293 ) Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>	2025-11-12 21:29:02 +08:00
Yuge Zhang	790ed3efb3	View worker status on Dashboard (#296 )	2025-11-12 21:27:31 +08:00
Yuge Zhang	5ae7933d41	Use unified server launcher for LiteLLM Proxy (#292 )	2025-11-12 02:45:29 +08:00
Yuge Zhang	2ab977ed18	Dashboard - build into Python package (#291 )	2025-11-11 16:29:02 +08:00
				`@@ -0,0 +1 @@`
				`# Put contrib-related gitignore files here.`