.

Co-authored-by: SiyunZhao <siyunzhao@microsoft.com>

.dockerignore

@@ -0,0 +1,14 @@
+.venv
+**/.venv
+__pycache__
+.git
+.gitignore
+**/node_modules
+dist
+build
+.env
+docker
+.pytest_cache
+.vscode
+**/*.log
+examples/**/data

.github/workflows/backport.yml

@@ -0,0 +1,32 @@
+name: Backport Merged Pull Request
+on:
+  pull_request_target:
+    types: [closed]
+permissions:
+  contents: write
+  issues: write
+  pull-requests: write
+
+# NOTE:
+# Microsoft requires rotating BOT_PAT every 3 months.
+# Log onto agent-lightning-bot account and rotate the PAT if needed.
+
+jobs:
+  backport:
+    name: Backport pull request
+    runs-on: ubuntu-latest
+    # Don't run on closed unmerged pull requests
+    if: github.event.pull_request.merged
+    steps:
+      - uses: actions/checkout@v4
+      - name: Create backport pull requests
+        uses: korthout/backport-action@v3
+        with:
+          branch_name: 'backport/${pull_number}/${target_branch}'
+          label_pattern: ^(stable/[^ ]+)$
+          github_token: ${{ secrets.BOT_PAT }}
+          add_labels: backport
+          add_author_as_assignee: true
+          git_committer_name: agent-lightning-bot
+          # This email address is not monitored.
+          git_committer_email: agl.msft@outlook.com

.github/workflows/badge-apo.yml

@@ -0,0 +1,29 @@
+name: Badge - APO
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - APO
+    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-apo.yml', label: 'apo', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-azure.yml

@@ -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 });

.github/workflows/badge-calc-x.yml

@@ -0,0 +1,29 @@
+name: Badge - Calc-X
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+    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-calc-x.yml', label: 'calc-x', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-chartqa.yml

@@ -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 });

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

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

.github/workflows/badge-compat.yml

@@ -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 });

.github/workflows/badge-examples.yml

@@ -0,0 +1,45 @@
+name: Badge - Examples
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+      - Examples - Spider
+      - Examples - APO
+      - Examples - Unsloth
+      - Examples - Tinker
+      - Examples - Azure
+      - Examples - Claude Code
+      - Examples - RAG
+      - 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-calc-x.yml', label: 'examples-calc-x.stable', variants: ['stable'] },
+              { 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 });

.github/workflows/badge-latest.yml

@@ -0,0 +1,41 @@
+name: Badge - Latest
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+      - Examples - Spider
+      - Examples - APO
+      - Examples - Unsloth
+      - Examples - RAG
+      - Examples - Claude Code
+      - 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: 'examples-calc-x.yml', label: 'calc-x.latest', variants: ['latest'] },
+              { 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 });

.github/workflows/badge-rag.yml

@@ -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 });

.github/workflows/badge-spider.yml

@@ -0,0 +1,29 @@
+name: Badge - Spider
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Spider
+    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-spider.yml', label: 'spider', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-tinker.yml

@@ -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 });

.github/workflows/badge-unit.yml

@@ -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 });

.github/workflows/badge-unsloth.yml

@@ -0,0 +1,29 @@
+name: Badge - Unsloth
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Unsloth
+    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-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/benchmark.yml

@@ -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 256
+              --max-rounds 8
+              --sleep-seconds 0.1
+          - id: scenario-large-batch
+            display: Large batch waves
+            kind: scenario
+            store_workers: 64
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 120
+            args: >-
+              --mode batch
+              --total-tasks 50000
+              --batch-size 8192
+              --n-runners 256
+              --max-rounds 6
+              --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: 64
+            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

.github/workflows/dashboard.yml

@@ -0,0 +1,33 @@
+name: Dashboard
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+  push:
+    branches: [ main, stable/**/* ]
+
+jobs:
+  dashboard:
+    name: Chromatic
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Run Chromatic
+        uses: chromaui/action@v13
+        with:
+          projectToken: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
+          workingDir: dashboard
+          exitZeroOnChanges: false

.github/workflows/docs.yml

@@ -8,6 +8,10 @@ on:
       - 'v*'
   workflow_dispatch:
 
+concurrency:
+  group: docs-deploy
+  cancel-in-progress: false
+
 permissions:
   contents: write
   pages: write
@@ -20,15 +24,14 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
         with:
           python-version: '3.12'
-
-      - name: Install dependencies
-        run: |
-          ./scripts/setup_stable.sh
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev
 
       - name: Configure Git
         run: |
@@ -51,11 +54,11 @@ jobs:
       - name: Deploy versioned docs
         if: startsWith(github.ref, 'refs/tags/')
         run: |
-          mike deploy --push --update-aliases ${{ steps.version.outputs.version }} stable
+          uv run --locked --no-sync mike deploy --push --update-aliases ${{ steps.version.outputs.version }} stable
 
       - name: Deploy dev docs
         if: github.ref == 'refs/heads/main'
         run: |
-          mike deploy --push latest
+          uv run --locked --no-sync mike deploy --push latest
           # Always set stable to default
-          mike set-default --push stable
+          uv run --locked --no-sync mike set-default --push stable

.github/workflows/examples-apo.yml

@@ -0,0 +1,116 @@
+name: Examples - APO
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-apo, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'APO - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('APO - {0}', github.event_name) }}
+
+jobs:
+  apo:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-apo' ||
+      github.event.action == 'ci-all'
+    name: APO (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    # This job is run on GitHub hosted runners rather than self-hosted runners because it needs no GPU.
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    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:
+      - 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 apo \
+            --group dev --group experiment --group agents --group core-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra apo \
+            --group dev --group experiment --group agents --group core-${{ 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-apo-${{ 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: APO custom algorithm
+        run: |
+          set -ex
+          cd examples/apo
+          uv run apo_custom_algorithm_trainer.py | tee _ci_apo.log
+          # Check whether the log contains "Best prompt found:"
+          grep "Best prompt found:" _ci_apo.log
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: APO custom algorithm debugger
+        run: |
+          set -ex
+          cd examples/apo
+          uv run apo_debug.py --mode runner
+          uv run apo_debug.py --mode hook
+          uv run apo_debug.py --mode trainer
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: APO built-in algorithm
+        run: |
+          set -ex
+          cd examples/apo
+          uv run room_selector_apo.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: matrix.setup-script != 'legacy'

.github/workflows/examples-azure.yml

@@ -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

.github/workflows/examples-calc-x.yml

@@ -0,0 +1,438 @@
+name: Examples - Calc-X
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-calc-x, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Calc-X - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Calc-X - {0}', github.event_name) }}
+
+jobs:
+  calc-x-perf:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-calc-x' ||
+      github.event.action == 'ci-all'
+    name: Calc-X Performance (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 --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 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-performance-${{ 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
+
+      # Calc-X training suddenly works after running the sanity check.
+      # And it has to be run before Spider training.
+      # The client side used to hang in many of my attempts.
+      # Don't ask why. Don't touch this.
+      - name: Calc-X training
+        run: |
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          python train_calc_agent.py --val-file data/test_mini.parquet --ci
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train
+
+      - name: Validate Calc-X training
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+  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
+          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:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_llm_proxy
+
+      - 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
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+
+          agl store --port 4747 &
+          sleep 5
+          AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=runner python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast &
+          sleep 5
+          AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast
+
+          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 train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
+          while pgrep -f train_calc_agent.py; do
+            echo "Waiting for train_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "train_calc_agent.py has finished."
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_external_store
+
+      - 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
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+
+          PYTHONUNBUFFERED=1 AGL_SERVER_HOST=127.0.0.1 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=runner python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast &
+          sleep 5
+          PYTHONUNBUFFERED=1 AGL_SERVER_HOST=0.0.0.0 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast
+
+          pkill -f train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
+          while pgrep -f train_calc_agent.py; do
+            echo "Waiting for train_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "train_calc_agent.py has finished."
+        shell: bash
+        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 }}

.github/workflows/examples-chartqa.yml

@@ -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 }}

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

@@ -0,0 +1,151 @@
+name: Examples - Claude Code
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: "0 20 * * *"
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-claude-code, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Claude Code - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Claude Code - {0}', github.event_name) }}
+
+jobs:
+  claude-code:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-claude-code' ||
+      github.event.action == 'ci-all'
+    name: Claude Code (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+          - python-version: "3.12"
+            setup-script: "stable"
+          - python-version: "3.13"
+            setup-script: "latest"
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-gpu-stable
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Download model
+        run: |
+          source .venv/bin/activate
+          python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
+
+      - name: Launch vLLM server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+            --max-model-len 131072 \
+            --enable-auto-tool-choice \
+            --tool-call-parser qwen3_coder \
+            --port 45993 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:45993/v1/models > /dev/null 2>&1; then
+              echo "vLLM server is ready!"
+              VLLM_READY=1
+              break
+            fi
+            echo "Waiting for vLLM server to be ready... (${i})"
+            sleep 5
+          done
+          if [[ "$VLLM_READY" != "1" ]]; then
+            echo "vLLM server failed to start!"
+            exit 1
+          fi
+
+      - name: Claude Code sanity check with vLLM models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py vllm --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct --base-url http://localhost:45993/v1 --debug
+        shell: bash
+
+      - name: Upload sanity check artifacts for vLLM
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error
+
+      - name: Cleanup vLLM
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+          rm -rf examples/claude_code/data/
+          rm -rf examples/claude_code/logs/
+
+      - name: Claude Code sanity check with OpenAI models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+
+      - name: Upload sanity check artifacts for OpenAI
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error

.github/workflows/examples-compat.yml

@@ -0,0 +1,151 @@
+name: Examples - Backward Compatibility
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: '0 22 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-compat, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Backward Compatibility - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Backward Compatibility - {0}', github.event_name) }}
+
+jobs:
+  backward-compatibility:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-compat' ||
+      github.event.action == 'ci-all'
+    name: Backward Compatibility (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - 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: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups --extra apo --extra verl \
+            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
+      - name: Override VERL (stable)
+        run: |
+          uv pip install verl==0.5.0 vllm==0.10.2
+        if: matrix.setup-script == '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-backward-compatibility-${{ 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: APO example (legacy client-server style)
+        run: |
+          set -ex
+          cd examples/apo
+          uv run legacy_apo_client.py &
+          sleep 3  # Wait for the client to be up
+          uv run legacy_apo_server.py
+          pkill -f legacy_apo_client.py && echo "SIGTERM sent to legacy_apo_client.py" || echo "No legacy_apo_client.py process found"
+          while pgrep -f legacy_apo_client.py; do
+            echo "Waiting for legacy_apo_client.py to finish..."
+            sleep 5
+          done
+          echo "legacy_apo_client.py has finished."
+          sleep 10
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - 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: Calc-X training (legacy client-server style)
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python legacy_calc_agent.py &
+          bash legacy_train.sh
+          pkill -f legacy_calc_agent.py && echo "SIGTERM sent to legacy_calc_agent.py" || echo "No legacy_calc_agent.py process found"
+          while pgrep -f legacy_calc_agent.py; do
+            echo "Waiting for legacy_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "legacy_calc_agent.py has finished."
+          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
+
+      - name: Validate Calc-X training
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

.github/workflows/examples-rag.yml

@@ -0,0 +1,179 @@
+name: Examples - RAG
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: '0 22 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-rag, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'RAG - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('RAG - {0}', github.event_name) }}
+
+jobs:
+  rag:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-rag' ||
+      github.event.action == 'ci-all'
+    name: RAG (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group rag --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-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 }}

.github/workflows/examples-spider.yml

@@ -0,0 +1,126 @@
+name: Examples - Spider
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: '0 20 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-spider, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Spider - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Spider - {0}', github.event_name) }}
+
+jobs:
+  spider:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-spider' ||
+      github.event.action == 'ci-all'
+    name: Spider (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+        # legacy is omitted because langchain doesn't work with legacy vllm versions
+        - 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 langchain --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --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
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare Spider dataset
+        run: |
+          set -ex
+          cd examples/spider
+          uv run gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
+          unzip -q spider-data.zip -d data
+          rm spider-data.zip
+
+      - name: Spider sanity check
+        run: |
+          set -ex
+          cd examples/spider
+          uv run sql_agent.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: success() || failure()
+
+      - name: Spider training
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/spider
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_sql_agent.py fast
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: spider_train
+
+      - 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 }} --reward-tolerance 5
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

.github/workflows/examples-tinker.yml

@@ -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 }}

.github/workflows/examples-unsloth.yml

@@ -0,0 +1,129 @@
+name: Examples - Unsloth
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-unsloth, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Unsloth - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Unsloth - {0}', github.event_name) }}
+
+jobs:
+  unsloth:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-unsloth' ||
+      github.event.action == 'ci-all'
+    name: Unsloth (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        # Legacy versions are not supported for Unsloth examples.
+        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 --extra verl \
+            --group dev --group experiment --group trl --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-unsloth-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Prepare Unsloth model
+        run: |
+          set -ex
+          cd examples/unsloth
+          rm -rf models
+          uv run hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
+
+      - name: Unsloth SFT example
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/unsloth
+
+          agl store --port 4747 &
+          sleep 5
+          python sft_rollout_runners.py &
+          sleep 5
+          python sft_algorithm.py
+
+          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 sft_rollout_runners.py && echo "SIGTERM sent to sft_rollout_runners.py" || echo "No sft_rollout_runners.py process found"
+          while pgrep -f sft_rollout_runners.py; do
+            echo "Waiting for sft_rollout_runners.py to finish..."
+            sleep 5
+          done
+          echo "sft_rollout_runners.py has finished."
+          sleep 10
+
+          # Check models/version_2 must exist
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Unsloth SFT example all-in-one
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/unsloth
+          rm -rf models/version_1 models/version_2
+
+          python sft_allinone.py
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

.github/workflows/examples.yml

@@ -1,330 +0,0 @@
-name: Examples Test
-permissions:
-  contents: read
-on:
-  schedule:
-    # Every day at 3 AM UTC+8
-    - cron: '0 19 * * *'
-
-  workflow_dispatch:
-
-jobs:
-  examples:
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
-    timeout-minutes: 90
-    strategy:
-      matrix:
-        setup: [stable, latest]
-      fail-fast: false
-    steps:
-      - name: Check GPU status
-        run: nvidia-smi
-      - name: Check disk space
-        run: df -h
-      - uses: actions/checkout@v4
-      - name: Create a virtual environment
-        run: python3 -m venv .venv
-      - name: Install dependencies (${{ matrix.setup }})
-        run: |
-          . .venv/bin/activate
-          ./scripts/setup_${{ matrix.setup }}_gpu.sh
-      - name: Freeze dependencies
-        run: |
-          . .venv/bin/activate
-          which python
-          which pip
-          which uvx
-          pip list | tee requirements-freeze.txt
-      - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: dependencies-${{ matrix.setup }}
-          path: requirements-freeze.txt
-          compression-level: 0
-
-      - name: Launch LiteLLM Proxy
-        run: |
-          set -ex
-          . .venv/bin/activate
-          litellm --config scripts/litellm_ci.yaml --port 12306 &
-          sleep 10  # Wait for the proxy to be up
-        env:
-          AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
-          AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
-
-      - name: Verify LiteLLM Proxy
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/litellm_sanity_check.py
-        env:
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-
-      - name: Prepare Unsloth model
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/unsloth
-          rm -rf models
-          hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
-
-      - name: Prepare Spider dataset
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/spider
-          gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
-          unzip -q spider-data.zip -d data
-          rm spider-data.zip
-      - name: Prepare Calc-X dataset
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
-          unzip calc-x-data.zip -d data
-          rm calc-x-data.zip
-
-      # APO Examples test
-      - name: APO example (legacy)
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/apo
-          python legacy_apo_client.py &
-          sleep 3  # Wait for the client to be up
-          python legacy_apo_server.py
-          pkill -f legacy_apo_client.py && echo "SIGTERM sent to legacy_apo_client.py" || echo "No legacy_apo_client.py process found"
-          while pgrep -f legacy_apo_client.py; do
-            echo "Waiting for legacy_apo_client.py to finish..."
-            sleep 5
-          done
-          echo "legacy_apo_client.py has finished."
-          sleep 10
-        env:
-          OPENAI_API_BASE: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-      - name: APO example
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/apo
-          python apo.py | tee _ci_apo.log
-          # Check whether the log contains "Best prompt found:"
-          grep "Best prompt found:" _ci_apo.log
-        env:
-          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-      - name: APO example debug sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/apo
-          python apo_debug.py --mode runner
-          python apo_debug.py --mode trainer
-        env:
-          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-
-      - name: APO built-in algorithm
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/apo
-          python room_selector_apo.py
-        env:
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-        if: success() || failure()
-
-      - name: Spider sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/spider
-          python sql_agent.py --trainer.n-workers 1 --trainer.dev true --trainer.max-tasks 2
-        env:
-          VERL_API_BASE: http://localhost:9999/
-          OPENAI_API_BASE: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-        if: success() || failure()
-      - name: Calc-X MCP sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          python tests/test_mcp_calculator.py
-        env:
-          OPENAI_API_BASE: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-      - name: Calc-X sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          python calc_agent_dev.py
-        env:
-          OPENAI_API_BASE: http://localhost:12306/
-          OPENAI_API_KEY: dummy
-
-      # Calc-X training suddenly works after running the sanity check.
-      # And it has to be run before Spider training.
-      # The client side used to hang in many of my attempts.
-      # Don't ask why. Don't touch this.
-      - name: Calc-X training v0.1
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/calc_x
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python calc_agent.py &
-          bash train_ci.sh
-          pkill -f calc_agent.py && echo "SIGTERM sent to calc_agent.py" || echo "No calc_agent.py process found"
-          while pgrep -f calc_agent.py; do
-            echo "Waiting for calc_agent.py to finish..."
-            sleep 5
-          done
-          echo "calc_agent.py has finished."
-          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
-        if: success() || failure()
-
-      - name: Validate Calc-X training
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
-        env:
-          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
-          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
-
-      - name: Calc-X training v0.2
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/calc_x
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python calc_agent_v0_2.py
-          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_v0_2
-        if: success() || failure()
-
-      - name: Calc-X training v0.2 LLM Proxy
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/calc_x
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python calc_agent_v0_2_llm_proxy.py
-          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_v0_2_llm_proxy
-        if: success() || failure()
-
-      - name: Spider training
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/spider
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python sql_agent.py --trainer.n-workers 10 &
-          bash train_ci.sh
-          pkill -f sql_agent.py && echo "SIGTERM sent to sql_agent.py" || echo "No sql_agent.py process found"
-          while pgrep -f sql_agent.py; do
-            echo "Waiting for sql_agent.py to finish..."
-            sleep 5
-          done
-          echo "sql_agent.py has finished."
-          sleep 10
-        shell: bash
-        env:
-          VERL_API_BASE: http://localhost:9991/
-          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
-          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
-        id: spider_train
-        if: success() || failure()
-
-      - name: Validate Spider training
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
-        env:
-          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
-          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
-
-      # Unsloth Examples test
-      - name: Unsloth SFT example
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/unsloth
-
-          agl store --port 4747 &
-          sleep 5
-          python sft_rollout_runners.py &
-          sleep 5
-          python sft_algorithm.py
-
-          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 sft_rollout_runners.py && echo "SIGTERM sent to sft_rollout_runners.py" || echo "No sft_rollout_runners.py process found"
-          while pgrep -f sft_rollout_runners.py; do
-            echo "Waiting for sft_rollout_runners.py to finish..."
-            sleep 5
-          done
-          echo "sft_rollout_runners.py has finished."
-          sleep 10
-
-          # Check models/version_2 must exist
-          if [ ! -d "models/version_2" ]; then
-            echo "models/version_2 does not exist"
-            exit 1
-          fi
-        env:
-          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
-          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
-        if: ${{ (success() || failure()) && matrix.setup == 'latest' }}
-
-      - name: Unsloth SFT example all-in-one
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/unsloth
-          rm -rf models/version_1 models/version_2
-
-          python sft_allinone.py
-          if [ ! -d "models/version_2" ]; then
-            echo "models/version_2 does not exist"
-            exit 1
-          fi
-        env:
-          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
-          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
-        if: matrix.setup == 'latest'
-
-      # Cleanup
-      - name: Cleanup
-        run: ./scripts/cleanup.sh
-        if: success() || failure()

.github/workflows/issue-comment.yml

@@ -0,0 +1,309 @@
+name: Issue Comment
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  pull-requests: write
+  issues: write
+  contents: write
+  actions: read
+
+jobs:
+  dispatch:
+    # Only run for comments on pull requests AND when the comment starts with "/ci"
+    if: >
+      github.event.issue.pull_request != null &&
+      startsWith(github.event.comment.body, '/ci')
+    runs-on: ubuntu-latest
+    outputs:
+      dispatched: ${{ steps.dispatch.outputs.dispatched }}
+      event_types: ${{ steps.dispatch.outputs.event_types }}
+      correlation_id: ${{ steps.dispatch.outputs.correlation_id }}
+      trigger_comment_id: ${{ steps.dispatch.outputs.trigger_comment_id }}
+      ack_comment_id: ${{ steps.ack.outputs.comment_id }}
+    steps:
+      - name: Guardrail — allow only members/collaborators
+        id: guard
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const allowed = ['MEMBER','OWNER','COLLABORATOR'];
+            const assoc = context.payload.comment.author_association;
+            if (!allowed.includes(assoc)) {
+              core.notice(`Ignoring /ci from ${context.payload.comment.user.login} (author_association=${assoc}).`);
+              core.setOutput('skip', 'true');
+            }
+
+      - name: Trigger repository dispatch
+        id: dispatch
+        if: steps.guard.outputs.skip != 'true'
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const owner = context.repo.owner;
+            const repo = context.repo.repo;
+            const pull_number = context.payload.issue.number;
+            const comment = context.payload.comment;
+
+            // Fetch current PR state
+            const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number });
+
+            // Add reaction so folks know we saw it
+            try {
+              await github.rest.reactions.createForIssueComment({
+                owner,
+                repo,
+                comment_id: comment.id,
+                content: 'rocket'
+              });
+            } catch (e) {
+              core.info('Could not add reaction (likely due to permissions). Continuing.');
+            }
+
+            const labels = (pr.labels ?? []).map(label => label.name);
+            const directCiLabels = labels.filter(label => label.startsWith('ci-'));
+            const hasCiAll = directCiLabels.includes('ci-all');
+            const dedupe = new Set(
+              directCiLabels.filter(label => label !== 'ci-all')
+            );
+
+            if (!hasCiAll && dedupe.size === 0) {
+              core.notice('No ci-* labels found on the pull request; nothing to dispatch.');
+              core.setOutput('dispatched', 'false');
+              core.setOutput('event_types', '');
+              return;
+            }
+
+            const correlation_id = `id-${comment.id}-${Date.now().toString(36)}`;
+
+            const clientPayload = {
+              correlation_id,
+              pull_number,
+              pr_ref: `refs/pull/${pull_number}/merge`,
+              pr_head_ref: pr.head.ref,
+              pr_head_sha: pr.head.sha,
+              pr_base_ref: pr.base.ref,
+              pr_base_sha: pr.base.sha,
+              trigger_comment_id: comment.id,
+              trigger_comment_user: comment.user.login,
+            };
+
+            const eventTypes = hasCiAll
+              ? ['ci-all']
+              : Array.from(dedupe);
+            for (const eventType of eventTypes) {
+              await github.rest.repos.createDispatchEvent({
+                owner,
+                repo,
+                event_type: eventType,
+                client_payload: { ...clientPayload, ci_label: eventType }
+              });
+              core.notice(`Dispatched '${eventType}' event for PR #${pull_number}.`);
+            }
+
+            core.setOutput('dispatched', 'true');
+            core.setOutput('event_types', eventTypes.join(','));
+            core.setOutput('correlation_id', correlation_id);
+            core.setOutput('trigger_comment_id', String(comment.id));
+
+      - name: Acknowledge in thread (optional)
+        if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched == 'true'
+        id: ack
+        uses: actions/github-script@v8
+        env:
+          EVENT_TYPES: ${{ steps.dispatch.outputs.event_types }}
+          CORRELATION_ID: ${{ steps.dispatch.outputs.correlation_id }}
+        with:
+          script: |
+            const eventTypes = (process.env.EVENT_TYPES || '')
+              .split(',')
+              .map(label => label.trim())
+              .filter(Boolean);
+            const formatted = eventTypes.map(label => `\`repository_dispatch:${label}\``).join(', ');
+            const { owner, repo } = context.repo;
+            const issue_number = context.payload.issue.number;
+            const body = [
+              `✅ CI trigger requested by @${context.payload.comment.user.login}.`,
+              `Fired ${formatted}.`,
+              '',
+              `_Collecting run links for correlation \`${process.env.CORRELATION_ID}\`…_`
+            ].join('\n');
+            const { data: comment } = await github.rest.issues.createComment({
+              owner, repo, issue_number,
+              body
+            });
+            core.setOutput('comment_id', String(comment.id));
+
+      - name: Notify missing ci label
+        if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched != 'true'
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const { owner, repo } = context.repo;
+            const issue_number = context.payload.issue.number;
+            await github.rest.issues.createComment({
+              owner,
+              repo,
+              issue_number,
+              body: `⚠️ CI trigger ignored because the pull request has no \`ci-*\` labels (e.g. \`ci-apo\`, \`ci-calc-x\`). Add the desired labels and try \`/ci\` again.`
+            });
+
+  watch:
+    needs: dispatch
+    if: needs.dispatch.outputs.dispatched == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 180
+    steps:
+      - name: Track dispatched runs and update comment
+        uses: actions/github-script@v8
+        env:
+          CORRELATION_ID: ${{ needs.dispatch.outputs.correlation_id }}
+          ACK_COMMENT_ID: ${{ needs.dispatch.outputs.ack_comment_id }}
+          TRIGGER_COMMENT_ID: ${{ needs.dispatch.outputs.trigger_comment_id }}
+        with:
+          script: |
+            const owner = context.repo.owner;
+            const repo = context.repo.repo;
+            const correlationId = process.env.CORRELATION_ID;
+            if (!correlationId) {
+              core.warning('No correlation id supplied; nothing to watch.');
+              return;
+            }
+
+            const ackCommentId = Number(process.env.ACK_COMMENT_ID || 0);
+            if (!ackCommentId) {
+              core.warning('No comment id available for updates; skipping watch.');
+              return;
+            }
+            const triggerCommentId = Number(process.env.TRIGGER_COMMENT_ID || 0);
+            if (!triggerCommentId) {
+              core.warning('No trigger comment id available; skipping watch.');
+              return;
+            }
+
+            const prefix = `🚀 CI Watcher for correlation ${correlationId} triggered by comment ${triggerCommentId}`;
+            core.notice(`Watching workflow runs for correlation '${correlationId}' using comment ${ackCommentId}.`);
+
+            function fmt(run) {
+              const status = run.status;
+              const conclusion = run.conclusion;
+              const badge = status === 'completed'
+                ? (conclusion === 'success' ? '🟢' : conclusion === 'failure' ? '🔴' : '🟡')
+                : (status === 'in_progress' ? '🟣' : '⚪️');
+              const title = run.display_title || run.name || `run ${run.id}`;
+              const statusText = status === 'completed' ? `${status}/${conclusion}` : status;
+              return `- ${badge} [${title}](${run.html_url}) — \`${statusText}\``;
+            }
+
+            const signatureOf = runs =>
+              runs
+                .map(run => `${run.id}:${run.status}/${run.conclusion || ''}`)
+                .sort()
+                .join('|');
+
+            const deadlineMs = Date.now() + 175 * 60 * 1000; // 175 minutes
+            let found = [];
+
+            async function searchOnce() {
+              const runs = await github.paginate(
+                github.rest.actions.listWorkflowRunsForRepo,
+                { owner, repo, event: 'repository_dispatch', per_page: 100 }
+              );
+              const cutoff = new Date(Date.now() - 60 * 60 * 1000); // last hour
+              return runs.filter(run => {
+                const createdAt = new Date(run.created_at);
+                const title = String(run.display_title || run.name || '');
+                return createdAt >= cutoff && title.includes(correlationId);
+              });
+            }
+
+            while (Date.now() < deadlineMs) {
+              found = await searchOnce();
+              if (found.length > 0) {
+                core.notice(`Discovered ${found.length} workflow run(s) for correlation '${correlationId}'.`);
+                break;
+              }
+              core.notice(`No runs found yet for correlation '${correlationId}'; retrying shortly.`);
+              await new Promise(res => setTimeout(res, 10000));
+            }
+
+            if (found.length === 0) {
+              core.notice(`Watcher timed out with no runs for correlation '${correlationId}'; notifying thread.`);
+              await github.rest.issues.updateComment({
+                owner,
+                repo,
+                comment_id: ackCommentId,
+                body: [
+                  prefix,
+                  `⚠️ I couldn't find any workflow runs for correlation \`${correlationId}\`.`,
+                  `They may be delayed or misconfigured.`
+                ].join('\n')
+              });
+              return;
+            }
+
+            const runIds = new Set(found.map(run => run.id));
+            let lastSignature = '';
+
+            async function refreshRuns() {
+              const ids = Array.from(runIds);
+              const refreshed = [];
+              for (const id of ids) {
+                const { data } = await github.rest.actions.getWorkflowRun({
+                  owner,
+                  repo,
+                  run_id: id
+                });
+                refreshed.push(data);
+              }
+              return refreshed;
+            }
+
+            async function updateCommentIfChanged(runs, allDone) {
+              const signature = signatureOf(runs);
+              if (signature === lastSignature) {
+                // Run statuses unchanged; skipping comment update.
+                return;
+              }
+              lastSignature = signature;
+              core.notice(`Updating comment ${ackCommentId} with ${runs.length} run status entries (allDone=${allDone}).`);
+              await github.rest.issues.updateComment({
+                owner,
+                repo,
+                comment_id: ackCommentId,
+                body: [
+                  prefix,
+                  `🏃‍♀️ Tracking ${runs.length} workflow run(s):`,
+                  '',
+                  ...runs.map(fmt),
+                  '',
+                  allDone ? '✅ All runs completed.' : '_Still running…_'
+                ].join('\n')
+              });
+            }
+
+            await updateCommentIfChanged(found, found.every(run => run.status === 'completed'));
+
+            while (Date.now() < deadlineMs) {
+              const latest = await searchOnce();
+              for (const run of latest) {
+                if (!runIds.has(run.id)) {
+                  runIds.add(run.id);
+                  core.notice(`Detected additional run ${run.id} (${run.name || run.display_title || 'unnamed'}) for correlation '${correlationId}'.`);
+                }
+              }
+              const current = await refreshRuns();
+              const allDone = current.every(run => run.status === 'completed');
+              await updateCommentIfChanged(current, allDone);
+              if (allDone) {
+                core.notice(`All runs for correlation '${correlationId}' completed; stopping watcher.`);
+                break;
+              }
+              await new Promise(res => setTimeout(res, 60000));
+            }
+
+            if (Date.now() >= deadlineMs) {
+              core.warning(`Watcher hit the deadline while monitoring correlation '${correlationId}'.`);
+            }

.github/workflows/playground.yml

@@ -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!"

.github/workflows/pypi-nightly.yml

@@ -2,8 +2,8 @@ name: PyPI Nightly Build
 
 on:
   schedule:
-    # Run daily at 6:00 AM UTC
-    - cron: '0 6 * * *'
+    # Run daily at 6:00 AM UTC+8
+    - cron: '0 22 * * *'
   workflow_dispatch:  # Allow manual trigger
 
 jobs:
@@ -14,18 +14,25 @@ jobs:
       contents: read
 
     steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v6
         with:
           python-version: '3.12'
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev
 
-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[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
@@ -44,16 +51,9 @@ jobs:
 
       - name: Build package
         run: |
-          hatch build
+          uv build
 
       - name: Publish to Test PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://test.pypi.org/legacy/
-
-      - name: Test installation from Test PyPI
-        run: |
-          # Wait a bit for the package to be available
-          sleep 30
-          pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
-          python -c "import agentlightning; print('Package installed successfully')"

.github/workflows/pypi-release.yml

@@ -48,34 +48,34 @@ jobs:
       contents: read
 
     steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v6
         with:
           python-version: '3.12'
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev
 
-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[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: |
-          hatch build
+          uv build
 
       - name: Verify package contents
         run: |
-          python -m tarfile -l dist/*.tar.gz
-          python -m zipfile -l dist/*.whl
+          uv run --locked --no-sync python -m tarfile -l dist/*.tar.gz
+          uv run --locked --no-sync python -m zipfile -l dist/*.whl
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
-
-      - name: Test installation from PyPI
-        run: |
-          # Wait a bit for the package to be available
-          sleep 30
-          pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
-          python -c "import agentlightning; print('Package installed successfully')"

.github/workflows/tests-full.yml

@@ -8,62 +8,356 @@ on:
 
   workflow_dispatch:
 
+  repository_dispatch:
+    types: [ci-gpu, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'GPU Test - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('GPU Test - {0}', github.event_name) }}
+
 jobs:
   tests-full:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-gpu' ||
+      github.event.action == 'ci-all'
+    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
     strategy:
       matrix:
-        setup: [stable, latest]
+        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
       - uses: actions/checkout@v4
-      - name: Create a virtual environment
-        run: python3 -m venv .venv
-      - name: Install dependencies (${{ matrix.setup }})
-        run: |
-          . .venv/bin/activate
-          ./scripts/setup_${{ matrix.setup }}_gpu.sh
+        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 apo --group dev --group agents --group langchain --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script == 'stable'
+      # Don't install 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: |
-          . .venv/bin/activate
-          which python
-          which pip
-          which uvx
-          pip list | tee requirements-freeze.txt
+          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-${{ matrix.setup }}
+          name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
           compression-level: 0
 
       - name: Launch LiteLLM Proxy
         run: |
-          set -ex
-          . .venv/bin/activate
-          litellm --config scripts/litellm_ci.yaml --port 12306 &
-          sleep 10  # Wait for the proxy to be up
+          ./scripts/litellm_run.sh
         env:
-          AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
-          AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
-      - name: Verify LiteLLM Proxy
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/litellm_sanity_check.py
-        env:
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
+          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: |
-          set -ex
-          . .venv/bin/activate
-          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

.github/workflows/tests.yml

@@ -5,9 +5,9 @@ permissions:
 
 on:
   push:
-    branches: [ main ]
+    branches: [ main, stable/**/* ]
   pull_request:
-    branches: [ main ]
+    branches: [ main, stable/**/* ]
   workflow_dispatch:
 
   schedule:
@@ -16,70 +16,107 @@ on:
 
 jobs:
 
-  lint-fast:
-    name: Lint - Fast
+  lint:
+    strategy:
+      matrix:
+        setup: [fast, slow, next]
+      fail-fast: false
+    name: Lint - ${{ matrix.setup }}
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v7
         with:
+          enable-cache: true
           python-version: '3.12'
-      - name: Install dependencies
+      - 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: |
-          python -m pip install --upgrade pip
-          pip install -e .[dev]
+          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 != 'fast'
+      # This pre-commit skips JavaScript on purpose.
       - name: Run pre-commit
         uses: pre-commit/action@v3.0.1
       - name: Check Python headers
-        run: |
-          python scripts/check_python_headers.py
+        run: uv run --locked --no-sync scripts/check_headers.py
       - name: Run Black
-        run: black --check .
+        run: uv run --locked --no-sync black --check .
       - name: Run isort
-        run: isort --check-only .
-      - name: Run pyright
-        run: pyright -p pyrightconfig.fast.json
+        run: uv run --locked --no-sync isort --check-only .
+      - name: Run pyright (fast)
+        run: uv run --locked --no-sync pyright -p pyrightconfig.fast.json
+        if: matrix.setup == 'fast'
+      - name: Run pyright (slow)
+        run: uv run --locked --no-sync pyright -p pyrightconfig.json
+        if: matrix.setup != 'fast'
 
-  lint-slow:
-    name: Lint - Slow
+  lint-js:
+    name: Lint - JavaScript
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v6
         with:
-          python-version: '3.12'
+          node-version: '22'
+          cache: 'npm'
+          cache-dependency-path: dashboard/package-lock.json
       - name: Install dependencies
-        run: |
-          ./scripts/setup_type_checking.sh
-      - name: Run Black
-        run: black --check .
-      - name: Run isort
-        run: isort --check-only .
-      - name: Run pyright
-        run: pyright -p pyrightconfig.json
+        run: cd dashboard && npm ci
+      - name: Run ESLint
+        run: cd dashboard && npm run eslint
+      - name: Run Prettier
+        run: cd dashboard && npm run prettier
+      - name: Run Stylelint
+        run: cd dashboard && npm run stylelint
+      - name: Run Typecheck
+        run: cd dashboard && npm run typecheck
+      - name: Verify build
+        run: cd dashboard && npm run build
 
   docs:
     name: Build documentation
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v6
         with:
           python-version: '3.12'
-      - name: Install documentation dependencies
-        run: |
-          ./scripts/setup_stable.sh
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev
       - name: Set source commit for docs
         run: |
           echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
-      - name: Build documentation
+      - name: Verify OpenAPI specification is up-to-date
         run: |
-          mkdocs build --strict
+          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
         uses: actions/upload-artifact@v4
         with:
@@ -90,37 +127,109 @@ 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'
             setup-script: 'stable'
           - python-version: '3.12'
+            setup-script: 'stable'
+          - python-version: '3.13'
             setup-script: 'latest'
-          - python-version: '3.12'
-            setup-script: 'stable'
       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:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v7
         with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        run: |
-          ./scripts/setup_${{ matrix.setup-script }}.sh
+          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)
+        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 --extra weave --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
+        if: matrix.env.setup-script != 'latest'
       - name: Freeze dependencies
         run: |
-          pip list | tee requirements-freeze-${{ matrix.python-version }}-${{ matrix.setup-script }}.txt
+          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-python-${{ matrix.python-version }}-${{ matrix.setup-script }}
-          path: requirements-freeze-${{ matrix.python-version }}-${{ matrix.setup-script }}.txt
+          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: |
-          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)
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - 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
+          python-version: '3.12'
+      - name: Sync Python dependencies
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Run vitest
+        run: cd dashboard && npm run vitest

.gitignore

@@ -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__/
@@ -189,6 +191,9 @@ cython_debug/
 #  you could uncomment the following to ignore the enitre vscode folder
 .vscode/
 
+# Emacs backup files
+*~
+
 # Ruff stuff:
 .ruff_cache/
 
@@ -204,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/

.pre-commit-config.yaml

@@ -3,11 +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$)|(.*store-openapi\.json$)
       - id: check-shebang-scripts-are-executable
       - id: detect-private-key
   - repo: https://github.com/pycqa/isort
@@ -22,3 +25,53 @@ repos:
         pass_filenames: false
         always_run: true
         args: ["."]
+
+  - repo: local
+    hooks:
+    - id: prettier
+      name: prettier (dashboard)
+      language: system
+      pass_filenames: false
+      always_run: true
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx prettier --cache --write "**/*.{ts,tsx,mjs,cjs}"
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '
+
+    - id: eslint
+      name: eslint (dashboard)
+      language: system
+      pass_filenames: false
+      always_run: true
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx eslint --cache --fix .
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '
+
+    - id: stylelint
+      name: stylelint (dashboard)
+      language: system
+      pass_filenames: false
+      always_run: true
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx stylelint --cache --fix "**/*.css"
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '

.python-version

@@ -0,0 +1 @@
+3.12

AGENTS.md

@@ -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.

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

README.md

@@ -1,13 +1,14 @@
-<div style="text-align:center; margin-bottom:20px;">
-  <img src="docs/assets/readme-banner.png" alt="Agent-lightning-banner" style="max-width:600px"/>
-</div>
+<p align="center">
+  <img src="docs/assets/readme-banner.svg" alt="Agent-lightning-banner" style="width:600px"/>
+</p>
 
 # Agent Lightning⚡
 
-[![CPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml)
-[![GPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples.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)
+[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/microsoft/agent-lightning)
 [![Discord](https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&logoColor=white)](https://discord.gg/RYk7CdvDR7)
 
 **The absolute trainer to light up AI agents.**
@@ -17,14 +18,36 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 ## ⚡ Core Features
 
 - Turn your agent into an optimizable beast with **ZERO CODE CHANGE** (almost)! 💤
-- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, ...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
+- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, Microsoft Agent Framework...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
 - **Selectively** optimize one or more agents in a multi-agent system. 🎯
-- Embraces Reinforcement Learning, Automatic Prompt Optimization and more **algorithms**. 🤗
+- Embraces **Algorithms** like Reinforcement Learning, Automatic Prompt Optimization, Supervised Fine-tuning and more. 🤗
 
-![Agent-Lightning-code-diff](docs/assets/readme-diff.png)
+Read more on our [documentation website](https://microsoft.github.io/agent-lightning/).
 
-## ⚡ Resources
+<p align="center">
+  <img src="docs/assets/readme-diff.svg" alt="Agent-Lightning Core Quickstart" style="width:100%"/>
+</p>
 
+## ⚡ Installation
+
+```bash
+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/ 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
+
+- 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.
 - 8/5/2025 [Agent Lightning: Train ANY AI Agents with Reinforcement Learning](https://arxiv.org/abs/2508.03680) arXiv paper.
 - 7/26/2025 [We discovered an approach to train any AI agent with RL, with (almost) zero code changes.](https://www.reddit.com/r/LocalLLaMA/comments/1m9m670/we_discovered_an_approach_to_train_any_ai_agent/) Reddit.
@@ -34,115 +57,30 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 
 - [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.
-
-## ⚡ Installation
-
-First, let's get your environment set up. We'll be using `/path/to/agentlightning` to refer to the directory containing this README file.
-
-### 1. Set Up Your Environment
-
-We strongly recommend creating a new virtual environment to avoid conflicts with other packages. You can use either `conda` or `venv`. **Python 3.10 or later** is recommended.
-
-### 2. Install Core Training Dependencies (Optional)
-
-If you are running RL with Agent-Lightning, the next step is to install the essential packages: `PyTorch`, `FlashAttention`, `vLLM` and `VERL`. The following versions and installation order have been tested and are confirmed to work.
-
-```bash
-pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
-pip install flash-attn --no-build-isolation
-pip install vllm==0.9.2
-pip install verl==0.5.0
-```
-
-See `scripts/setup_stable_gpu.sh` for a full installation script.
-
-### 3. Install Agent Lightning
-
-Now, you're ready to install Agent Lightning itself.
-
-```bash
-pip install agentlightning
-```
-
-### 4. Install Agent Frameworks (Optional)
-
-If you plan to use other agent frameworks, you can install them with the following commands. If you don't need these, feel free to skip this step.
-We recommend doing this as the final step to avoid dependency versions being overwritten by mistake.
-
-```bash
-# AutoGen (Recommended to install first)
-pip install "autogen-agentchat" "autogen-ext[openai]"
-
-# LiteLLM
-pip install "litellm[proxy]"
-
-# MCP
-pip install mcp
-
-# UV
-pip install uv
-
-# OpenAI Agents
-pip install openai-agents
-
-# LangChain
-pip install langgraph "langchain[openai]" langchain-community langchain-text-splitters
-
-# SQL-related dependencies
-pip install sqlparse nltk
-```
-
-Don't worry if dependency conflicts arise during this step. Follow the installation order above and the conflicts generally do not matter.
-
-## ⚡ Examples
-
-For more detailed examples, please see the `examples` folder:
-
-1. [calc_x](examples/calc_x): An agent built with AutoGen with calculator tool use, trained on Calc-X dataset with Reinforcement Learning.
-2. [spider](examples/spider): A write-check-rewrite looped agent with LangGraph with SQL execution; selectively optimize write and rewrite on Spider dataset with Reinforcement Learning.
-3. [apo](examples/apo): An example to customize an optimization algorithm: Automatic Prompt Optimization.
-
-## ⚡ Important Caveats
-
-1.  **AgentOps Integration**: Agent Lightning uses [AgentOps](https://github.com/AgentOps-AI/agentops) for agent tracking by default. If you're already using AgentOps in your own code, you'll need to disable our managed AgentOps client by modifying the `tracer` parameter of trainer.
-2.  **Debugging Traces**: If you encounter issues with tracing, you can visualize the trace tree using `tracer.last_trace().visualize("tree_graph")`. Please note that this API is experimental and may change in future releases.
-3.  **Launching the Server and Agents**: Currently, the training server and agent clients must be launched in separate processes. You can open two terminal windows or run one of them in the background. The launching order generally doesn't matter.
-4.  **Environment Variables**: The environment variables and working directory at the time of `ray init` are important. If you run into "file not found" errors, try restarting Ray from your current working directory.
-5.  **Handling Timeouts**: The training server may hang if samples fail or time out on the agent side. To prevent this, we recommend setting limits on the prompt and response lengths, as this is the most common cause of failures.
-6.  **VERL Failures**: Save checkpoints frequently, as VERL with vLLM may sometimes experience out-of-memory issues. If you encounter a VERL failure, you can resume training from the last checkpoint.
+- [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
 
-Currently, Agent Lightning is built around a **training server** and one or multiple **agents**.
+Agent Lightning keeps the moving parts to a minimum so you can focus on your idea, not the plumbing. Your agent continues to run as usual; you can still use any agent framework you like; you drop in the lightweight `agl.emit_xxx()` helper, or let the tracer collect every prompt, tool call, and reward. Those events become structured spans that flow into the LightningStore, a central hub that keeps tasks, resources, and traces in sync.
 
-* The **server** manages the training data, prepares samples for the agents, and provides the LLM endpoint.
-* **Agents** retrieve samples from the server, process them (which may involve interacting with the LLM), and send the results back. These results, or "trajectories," are lists of prompts and responses from the LLM.
-* The **server** then collects these trajectories and computes the losses to optimize the language models.
+On the other side of the store sits the algorithm you choose, or write yourself. The algorithm reads spans, learns from them, and posts updated resources such as refined prompt templates or new policy weights. The Trainer ties it all together: it streams datasets to runners, ferries resources between the store and the algorithm, and updates the inference engine when improvements land. You can either stop there, or simply let the same loop keep turning.
 
-![Agent-Lightning-architecture](docs/assets/readme-architecture.png)
+No rewrites, no lock-in, just a clear path from first rollout to steady improvement.
 
-## ⚡ Development Instructions
+<p align="center">
+  <img src="docs/assets/readme-architecture.svg" alt="Agent-lightning Architecture" style="width:100%"/>
+</p>
 
-Install with development dependencies:
+## ⚡ CI Status
 
-```
-git clone https://github.com/microsoft/agent-lightning
-cd agent-lightning
-pip install -e .[dev]
-```
-
-Please run pre-commit hooks before checking in code:
-
-```
-pre-commit install
-pre-commit run --all-files --show-diff-on-failure --color=always
-```
-
-Serve documentation locally:
-
-```bash
-mkdocs serve
-```
+| 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) |
+| 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 | [![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
 
@@ -162,7 +100,7 @@ If you find Agent Lightning useful in your research or projects, please cite our
 
 ## ⚡ Contributing
 
-This project welcomes contributions and suggestions. 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.
 

agentlightning/__init__.py

@@ -1,16 +1,19 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 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 *

agentlightning/adapter/__init__.py

@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import Adapter, TraceAdapter
+from .base import Adapter, OtelTraceAdapter, TraceAdapter
 from .messages import TraceToMessages
 from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
 
 __all__ = [
     "TraceAdapter",
+    "OtelTraceAdapter",
     "Adapter",
     "TraceToTripletBase",
     "TracerTraceToTriplet",

agentlightning/adapter/base.py

@@ -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
 
@@ -13,18 +13,20 @@ T_to = TypeVar("T_to")
 class Adapter(Generic[T_from, T_to]):
     """Base class for synchronous adapters that convert data from one format to another.
 
-    This class defines a simple protocol for transformation:
+    The class defines a minimal protocol so that adapters can be treated like callables while
+    still allowing subclasses to supply the concrete transformation logic.
 
-    - The `__call__` method makes adapters callable, so they can be used like functions.
-    - Subclasses must implement the `adapt` method to define the actual conversion logic.
+    !!! note
+        Subclasses must override [`adapt()`][agentlightning.Adapter.adapt] to provide
+        the actual conversion.
 
-    Type parameters:
+    Type Variables:
 
-    - T_from: The source data type (input).
-    - T_to: The target data type (output).
+        T_from: Source data type supplied to the adapter.
 
-    Example:
+        T_to: Target data type produced by the adapter.
 
+    Examples:
         >>> class IntToStrAdapter(Adapter[int, str]):
         ...     def adapt(self, source: int) -> str:
         ...         return str(source)
@@ -37,8 +39,9 @@ class Adapter(Generic[T_from, T_to]):
     def __call__(self, source: T_from, /) -> T_to:
         """Convert the data to the target format.
 
-        This method delegates to `adapt` and allows the adapter
-        to be invoked as a function.
+        This method delegates to [`adapt()`][agentlightning.Adapter.adapt] so that an
+        instance of [`Adapter`][agentlightning.Adapter] can be used like a standard
+        function.
 
         Args:
             source: Input data in the source format.
@@ -51,8 +54,8 @@ class Adapter(Generic[T_from, T_to]):
     def adapt(self, source: T_from, /) -> T_to:
         """Convert the data to the target format.
 
-        Subclasses should override this method with the concrete
-        transformation logic.
+        Subclasses must override this method with the concrete transformation logic. The base
+        implementation raises `NotImplementedError` to make the requirement explicit.
 
         Args:
             source: Input data in the source format.
@@ -63,20 +66,15 @@ 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 class specializes `Adapter` for working with OpenTelemetry `ReadableSpan`
-    objects. It expects a list of spans as input and produces a custom target format
-    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+    This specialization of [`Adapter`][agentlightning.Adapter] expects a list of
+    `opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
+    reinforcement learning trajectories, structured logs, or analytics-ready payloads.
 
-    Subclasses should override `adapt` to define the desired conversion.
-
-    Type parameters:
-        T_to: The target data type that spans should be converted into.
-
-    Example:
-        >>> class TraceToDictAdapter(TraceAdapter[dict]):
+    Examples:
+        >>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
         ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
         ...         return {"count": len(spans)}
         ...
@@ -86,10 +84,11 @@ 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` for working with trace spans. It expects a list of
-    Agent-lightning spans as input and produces a custom target format
-    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+    This class specializes [`Adapter`][agentlightning.Adapter] for working with
+    [`Span`][agentlightning.Span] instances emitted by Agent Lightning instrumentation.
+    Subclasses receive entire trace slices and return a format suited for the downstream consumer,
+    for example reinforcement learning training data or observability metrics.
     """

agentlightning/adapter/messages.py

@@ -1,28 +1,48 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import json
 from collections import defaultdict
-from typing import 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 openai.types.chat import (
-    ChatCompletionAssistantMessageParam,
-    ChatCompletionFunctionToolParam,
-    ChatCompletionMessageFunctionToolCallParam,
-    ChatCompletionMessageParam,
-)
 from pydantic import TypeAdapter
 
 from agentlightning.types import Span
 
 from .base import TraceAdapter
 
+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
 
 class OpenAIMessages(TypedDict):
+    """OpenAI-style chat messages with optional tool definitions.
+
+    Attributes:
+        messages: Ordered chat messages that describe the conversation.
+        tools: Tool specifications available to the assistant, if any.
+    """
+
     messages: List[ChatCompletionMessageParam]
     tools: Optional[List[ChatCompletionFunctionToolParam]]
 
 
 class _RawSpanInfo(TypedDict):
+    """Intermediate representation parsed from a span.
+
+    Attributes:
+        prompt: Prompt messages reconstructed from span attributes.
+        completion: Assistant completions following tool invocations.
+        request: Request payload recorded in the trace.
+        response: Response payload recorded in the trace.
+        tools: Tool call metadata extracted from child spans.
+    """
+
     prompt: List[Dict[str, Any]]
     completion: List[Dict[str, Any]]
     request: Dict[str, Any]
@@ -31,16 +51,20 @@ class _RawSpanInfo(TypedDict):
 
 
 def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
-    """
-    Convert a flat dict with keys like 'gen_ai.prompt.0.role'
-    into structured nested dicts or lists under the given prefix.
+    """Convert flattened trace attributes into nested structures.
+
+    Attributes emitted by the tracing pipeline often arrive as dotted paths (for example
+    `gen_ai.prompt.0.role`). This helper groups those keys into nested dictionaries or lists so that
+    downstream processing can operate on structured data.
 
     Args:
-        data: Flat dictionary (keys are dotted paths).
-        prefix: Top-level key to extract (e.g., 'gen_ai.prompt').
+        data: Flat dictionary whose keys are dotted paths.
+        prefix: Top-level key (for example `gen_ai.prompt`) that determines which attributes are
+            grouped.
 
     Returns:
-        A nested dict (if no index detected) or list (if indexed).
+        A nested dictionary (no numeric index detected) or list (numeric indices detected) containing
+        the grouped values.
     """
     result: Union[Dict[str, Any], List[Any]] = {}
 
@@ -80,12 +104,28 @@ def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any],
 
 
 def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
-    """
-    Convert raw tool call traces + prompt/completion list
-    into OpenAI fine-tuning JSONL format (tool calling style).
+    """Convert raw trace payloads into OpenAI-style chat messages.
 
-    https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning-functions
+    The function consumes an iterable produced by
+    [`TraceToMessages.adapt()`][agentlightning.TraceToMessages.adapt] and yields
+    structures that match the OpenAI fine-tuning JSONL schema, including tool definitions.
+
+    Args:
+        prompt_completion_list: Raw prompt/completion/tool payloads extracted from a trace.
+
+    Returns:
+        A generator that yields [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages]
+        entries compatible with the OpenAI Functions fine-tuning format.
     """
+
+    # Import locally to avoid legacy OpenAI version type import errors
+    from openai.types.chat import (
+        ChatCompletionAssistantMessageParam,
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
     for pc_entry in prompt_completion_list:
         messages: List[ChatCompletionMessageParam] = []
 
@@ -157,25 +197,29 @@ def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Ge
 
 
 class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
-    """
-    Adapter that converts OpenTelemetry trace spans into OpenAI-compatible message format.
+    """Convert trace spans into OpenAI-compatible conversation messages.
 
-    This adapter processes trace spans containing LLM conversation data and transforms them
-    into structured OpenAI message format suitable for fine-tuning or analysis. It extracts
-    prompts, completions, tool calls, and function definitions from trace attributes and
-    reconstructs the conversation flow.
+    The adapter reconstructs prompts, completions, tool calls, and function definitions from
+    `gen_ai.*` span attributes. The resulting objects match the JSONL structure expected by the
+    OpenAI fine-tuning pipeline.
 
-    The adapter handles:
-    - Converting flat trace attributes into structured message objects
-    - Extracting and matching tool calls with their corresponding requests
-    - Building proper OpenAI ChatCompletionMessage objects with roles, content, and tool calls
-    - Generating function definitions for tools used in conversations
+    !!! warning
+        The adapter assumes all spans share a common trace and that tool call spans are direct
+        children of the associated completion span.
     """
 
-    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
-        """Find tool calls in the trace. Returns a dict with the tool call id, name, and arguments.
+    def get_tool_calls(self, completion: Span, all_spans: Sequence[Span], /) -> Iterable[Dict[str, Any]]:
+        """Yield tool call payloads for a completion span.
 
-        The spans that are direct children of the completion span are the tool calls.
+        Args:
+            completion: The completion span whose descendants should be inspected.
+            all_spans: The complete span list belonging to the trace.
+
+        Yields:
+            Dictionaries describing tool calls with identifiers, names, and arguments.
+
+        Raises:
+            ValueError: If a candidate tool span cannot be converted into a dictionary.
         """
         # Get all the spans that are children of the completion span
         children = [span for span in all_spans if span.parent_id == completion.span_id]
@@ -187,7 +231,16 @@ 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:
+            source: Spans containing `gen_ai.*` attributes emitted by the tracing pipeline.
+
+        Returns:
+            A list of [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages] entries that
+            capture prompts, completions, tools, and metadata.
+        """
         raw_prompt_completions: List[_RawSpanInfo] = []
 
         for span in source:

agentlightning/adapter/triplet.py

@@ -3,23 +3,74 @@
 from __future__ import annotations
 
 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 SpanNames, Triplet
-from agentlightning.types.tracer import Span
+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):
-    """
-    Transition class representing one transition in a trajectory.
-    State and action are a list of token IDs.
+    """A single transition within a reinforcement learning trajectory.
+
+    Attributes:
+        state: Token identifiers describing the model input state.
+        action: Token identifiers representing the model output.
+        response_id: Identifier of the LLM response used to deduplicate spans.
+        agent_name: Human-readable agent name captured from the trace.
+        reward: Scalar reward associated with the transition, if available.
     """
 
     state: List[int]
@@ -31,22 +82,27 @@ class Transition(BaseModel):
 
 
 class RewardMatchPolicy(str, Enum):
-    """How to find the reward for each transition from the trace.
-    In all cases, the reward must have data `{"type": "reward", "value": <float>|None}`,
-    as defined in `reward.py`.
+    """Strategies for matching rewards to LLM call spans.
+
+    !!! note
+        Each reward span must expose a payload shaped like `{"type": "reward", "value": <float>|None}`
+        as described in `reward.py`.
     """
 
     FIRST_SIBLING = "first_sibling"
-    """Use the first sibling in the current trace subtree as the reward, except another LLM call match is found."""
+    """Use the first sibling in the current trace subtree as the reward unless another LLM call match is found."""
 
     FIRST_OCCURRENCE = "first_occurrence"
-    """Use the first occurrence of the reward (in start time order) that occur after the current LLM call match.
-    """
+    """Use the first reward encountered in chronological order after the current LLM call match."""
 
 
 class TraceTree:
-    """
-    A trace item, along with its span and children.
+    """Tree representation of a trace span and its descendants.
+
+    Attributes:
+        id: Unique identifier for the span node.
+        span: [`Span`][agentlightning.Span] backing this node.
+        children: Child nodes connected to the current span.
     """
 
     def __init__(
@@ -80,10 +136,16 @@ class TraceTree:
         self.children.append(child)
 
     def visualize(self, filename: str, interested_span_match: str | None = None) -> None:
-        """
-        Visualize the trace tree using graphviz.
-        For debugging purposes only.
-        Use `interested_span_match` to filter the spans (and its ancesters) to be visualized.
+        """Render the trace tree with Graphviz for debugging purposes.
+
+        Args:
+            filename: Base filename for the generated `.png` diagram.
+            interested_span_match: Optional regular expression used to keep only matching spans
+                (and their ancestors) in the output.
+
+        !!! note
+            The method requires the optional `graphviz` dependency to be available in the runtime
+            environment.
         """
         import graphviz
 
@@ -112,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
@@ -125,9 +187,11 @@ class TraceTree:
         dot.render(filename, format="png", cleanup=True)  # type: ignore
 
     def names_tuple(self) -> Tuple[str, List[Any]]:
-        """Return the span name, and a list of children.
-        Each child is also a tuple of span name and a list of children.
-        Useful for debugging and testing.
+        """Return the span name alongside nested child names.
+
+        Returns:
+            A tuple of the current span name and a list of tuples for each child containing the
+            child name and its descendants.
         """
         name = self.span.name
         agent_name = self.agent_name()
@@ -140,15 +204,14 @@ class TraceTree:
         return name, children_names
 
     def traverse(self) -> List["TraceTree"]:
-        """
-        Traverse the trace tree and return a list of all spans.
-        """
+        """Traverse the tree depth first and return every node."""
         spans: List["TraceTree"] = [self]
         for child in self.children:
             spans.extend(child.traverse())
         return spans
 
     def to_json(self) -> dict[str, Any]:
+        """Convert the tree node into a JSON-serialisable structure."""
         if isinstance(self.span, ReadableSpan):
             span_data = json.loads(self.span.to_json())
         else:
@@ -161,10 +224,17 @@ class TraceTree:
 
     @classmethod
     def from_spans(cls, spans: List[Span]) -> "TraceTree":
-        """
-        Create a TraceTree from a list of spans.
-        All spans without parents found will be considered as candidate root spans.
-        If multiple root spans are found, a virtual root span will be created as the parent of all root spans.
+        """Construct a tree from a flat list of spans.
+
+        Args:
+            spans: Spans that collectively form a single trace segment.
+
+        Returns:
+            A [`TraceTree`][agentlightning.adapter.triplet.TraceTree] rooted at either the
+            discovered root span or a synthetic root when multiple roots are present.
+
+        Raises:
+            ValueError: If the span list is empty or no root span can be inferred.
         """
 
         if not spans:
@@ -245,8 +315,11 @@ class TraceTree:
             return root_span
 
     def agent_name(self) -> Optional[str]:
-        """Return the name of agent span. Return the agent or None (not an agent at all).
-        Extend this function to support more agent frameworks."""
+        """Return the agent name associated with the span, if any.
+
+        Returns:
+            Agent name extracted from known attributes, otherwise `None`.
+        """
         attributes = self.span.attributes
         if attributes is None:  # type: ignore
             return None
@@ -278,29 +351,49 @@ class TraceTree:
         if agent_name is not None:
             return agent_name
 
-    def maybe_reward_dict(self) -> dict[str, Any]:
-        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 {}
+        # 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
 
-        # 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 {}
+        # 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.
+        """
+        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.
+
+        Returns:
+            `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,
@@ -312,12 +405,19 @@ class TraceTree:
         within_llm_call: Optional[bool] = None,
         existing_llm_call_response_ids: Optional[set[str]] = None,
     ) -> List[Tuple["TraceTree", str]]:
-        """Find all LLM calls in the trace tree.
+        """Find LLM call spans matching the supplied filters.
 
-        The LLM call is defined as a span with type = request and name matching `llm_call_match`.
-        If `agent_match` is not None, it must also reside in an agent span (type = agent) with name matched.
+        Args:
+            llm_call_match: Regular expression used to match span names that qualify as LLM calls.
+            agent_match: Optional regular expression that must match the enclosing agent span name.
+            within_matching_subtree: Marker propagated through recursive calls to record matching agents.
+            within_reward: When `True`, suppresses LLM matches under reward spans.
+            within_llm_call: When `True`, prevents duplicate matches for nested LLM calls.
+            existing_llm_call_response_ids: Known response identifiers used to deduplicate spans.
 
-        Return a list of traces and the agent names (why it's selected).
+        Returns:
+            A list of tuples pairing the matching node with the agent subtree label that triggered the
+            match.
         """
         llm_calls: List[Tuple[TraceTree, str]] = []
 
@@ -330,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 (
@@ -342,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:
@@ -373,19 +476,26 @@ class TraceTree:
         return llm_calls
 
     def repair_hierarchy(self) -> None:
-        """
-        We find that sometimes the hierarchy is not correct, due to the way the spans are created.
-        The spans within the agent frameworks (e.g., OpenAI Agent SDK) and spans within the LLM frameworks
-        (e.g., Anthropic) are created in two systems.
-        So the inner LLM completion span does not necessarily have an agent span as a parent.
-        Rather they sometimes directly become children of the root span.
-        This becomes a problem when we want to select the LLM completion span with agent as filter.
-        To repair the hierarchy, for each children of the root span, we find a span over the whole tree,
-        with duration covering the current span and being closest to the current span.
+        """Repair missing parent-child relationships introduced by mixed tracing systems.
 
-        This function modifies the tree in place.
+        Some agent frameworks emit spans via multiple subsystems, which can cause LLM completion
+        spans to float directly under the root span instead of being nested under the correct agent.
+        The method re-parents those spans to the closest ancestor that fully envelopes the child in
+        time.
+
+        If we don't, when we want to select the LLM completion span with agent as filter.
+        We will never get the correct span underneath.
         """
+        # If the current node has only one child, recursively repair its hierarchy directly.
+        # This special-case handling is needed because when a trace is manually ended
+        # (via agentops.end_trace), the AgentOps provider automatically wraps all spans
+        # under an extra synthetic root node (e.g., "run_one.session").
+        if len(self.children) == 1:
+            self.children[0].repair_hierarchy()
+            return
+
         nodes_to_repair = list(self.children)
+
         for repair_node in nodes_to_repair:
             if len(self.children) == 1:
                 # If there is only one child, we don't need to repair the hierarchy.
@@ -410,7 +520,16 @@ class TraceTree:
                 closest_parent.children.append(repair_node)
 
     def match_rewards(self, reward_match: str, llm_calls: List["TraceTree"]) -> dict[str, Optional[float]]:
-        """Match the rewards to the LLM calls."""
+        """Assign rewards to previously matched LLM calls.
+
+        Args:
+            reward_match: Strategy identifier from
+                [`RewardMatchPolicy`][agentlightning.adapter.triplet.RewardMatchPolicy].
+            llm_calls: Trace nodes representing LLM call spans.
+
+        Returns:
+            Mapping from span identifier to reward value or `None` when no reward is available.
+        """
         llm_call_ids = set([llm_call.id for llm_call in llm_calls])
         rewards: dict[str, Optional[float]] = {}
 
@@ -440,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:
@@ -455,6 +574,131 @@ 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 = (
+            _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_payload["logprobs"] = logprobs_content
+
+        return Triplet(
+            prompt=prompt_payload,
+            response=response_payload,
+            reward=None,
+            metadata=dict(
+                request=request_metadata, response=response_metadata, response_id=response_id, agent_name=agent_name
+            ),
+        )
+
     def to_trajectory(
         self,
         llm_call_match: str = r"openai\.chat\.completion",
@@ -463,20 +707,21 @@ class TraceTree:
         dedup_llm_call: bool = True,
         reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
         final_reward: Optional[float] = None,
+        _skip_empty_token_spans: bool = False,
     ) -> List[Triplet]:
-        """Convert the trace tree to a trajectory.
+        """Convert the trace tree into a trajectory of [`Triplet`][agentlightning.Triplet] items.
 
-        First, we find all the LLM calls (span type = request, `llm_call_match` matching the span name).
-        If the agent match is set, we check, for each LLM call,
-        if it resides in an agent (span type = agent, `agent_match` matching the span name).
-        The above sets the basis for the trajectory, as we use the prompt token IDs and response token IDs for each LLM call,
-        as the state and action of each transition.
+        Args:
+            llm_call_match: Regular expression for LLM call span names.
+            agent_match: Optional regular expression for agent span names.
+            exclude_llm_call_in_reward: When `True`, prevents searching for rewards under the LLM
+                call subtree.
+            dedup_llm_call: When `True`, deduplicates spans using the LLM response identifier.
+            reward_match: Reward matching policy used to associate reward spans with LLM calls.
+            final_reward: Optional reward appended to the final transition when provided.
 
-        Then, we find the reward for each transition.
-        The reward is searched on the trace tree, after the LLM call,
-        until the next LLM call or the end of the tree depending on the policy.
-        It can be enforced to a sibling or the first occurrence in the time order, depending on the policy.
-        If a reward is never found for a transition, it is set to None.
+        Returns:
+            A list of [`Triplet`][agentlightning.Triplet] objects ordered by call sequence.
         """
         # Find all LLM calls
         llm_calls = self.find_llm_calls(
@@ -487,25 +732,23 @@ class TraceTree:
             within_llm_call=False if dedup_llm_call else None,
             existing_llm_call_response_ids=set(),
         )
-        id_transitions = [
-            (
-                llm_call.id,
-                Triplet(
-                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},  # type: ignore
-                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},  # type: ignore
-                    reward=None,
-                    metadata=dict(
-                        response_id=llm_call.span.attributes.get(  # type: ignore
-                            "gen_ai.response.id", None
-                        ),  # it works at least for OpenAI
-                        agent_name=agent_name,
-                    ),
-                ),
-            )
-            for llm_call, agent_name in llm_calls
-        ]
 
-        rewards = self.match_rewards(reward_match, [call for call, _ in llm_calls])
+        id_transitions: List[Tuple[str, Triplet]] = []
+        # We need to filter out the LLM calls with unrecorded token IDs
+        filtered_llm_calls: List[Tuple[TraceTree, str]] = []
+        for llm_call, agent_name in llm_calls:
+            triplet = self.span_to_triplet(llm_call.span, agent_name)
+            # This is a hot-fix for Tinker+CrewAI, which has some anonymous requests outside the trained agent.
+            # TODO: We might need to reconsider this.
+            if _skip_empty_token_spans and (
+                not triplet.prompt.get("token_ids") or not triplet.response.get("token_ids")
+            ):
+                logger.warning(f"Skipping LLM call with unrecorded token IDs: {triplet}")
+                continue
+            filtered_llm_calls.append((llm_call, agent_name))
+            id_transitions.append((llm_call.id, triplet))
+
+        rewards = self.match_rewards(reward_match, [call for call, _ in filtered_llm_calls])
         transitions = [
             transition.model_copy(update={"reward": rewards.get(id, None)}) for id, transition in id_transitions
         ]
@@ -522,22 +765,22 @@ class TraceTree:
 
 
 class TraceToTripletBase(TraceAdapter[List[Triplet]]):
-    """
-    Base class for trace triplet adapters.
-    """
+    """Base class for adapters that emit [`Triplet`][agentlightning.Triplet] trajectories."""
 
 
 class TracerTraceToTriplet(TraceToTripletBase):
-    """
-    An adapter to convert OpenTelemetry spans to triplet data.
+    """Convert tracer-emitted spans into triplet trajectories.
 
     Attributes:
-        repair_hierarchy: When `repair_hierarchy` is set to True, the trace will be repaired with the time information.
-            See `TraceTree.repair_hierarchy` for more details.
-        llm_call_match: Regular expression pattern to match LLM call span names.
-        agent_match: Optional regular expression pattern to match agent span names. If None, all agents are matched.
-        exclude_llm_call_in_reward: Whether to exclude LLM calls that occur within reward spans.
-        reward_match: Policy for matching rewards to LLM calls.
+        repair_hierarchy: When `True`, repair the span tree using
+            [`TraceTree.repair_hierarchy()`][agentlightning.adapter.triplet.TraceTree.repair_hierarchy]
+            before matching calls and rewards.
+        llm_call_match: Regular expression pattern that selects LLM call span names.
+        agent_match: Optional regular expression pattern for agent span names. When omitted, spans
+            from any agent are considered.
+        exclude_llm_call_in_reward: When `True`, ignore matches under reward spans while searching
+            for rewards.
+        reward_match: Strategy used to associate rewards with LLM calls.
     """
 
     def __init__(
@@ -547,12 +790,14 @@ class TracerTraceToTriplet(TraceToTripletBase):
         agent_match: Optional[str] = None,
         exclude_llm_call_in_reward: bool = True,
         reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
+        _skip_empty_token_spans: bool = False,
     ):
         self.repair_hierarchy = repair_hierarchy
         self.llm_call_match = llm_call_match
         self.agent_match = agent_match
         self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
         self.reward_match = reward_match
+        self._skip_empty_token_spans = _skip_empty_token_spans
 
     def visualize(
         self,
@@ -561,16 +806,17 @@ class TracerTraceToTriplet(TraceToTripletBase):
         filename: str = "trace_tree",
         interested_span_match: str | None = None,
     ) -> TraceTree:
-        """
-        Visualize the trace tree.
+        """Visualize the trace tree built from the supplied spans.
 
         Args:
-            source (List[Span]): The list of OpenTelemetry spans to visualize.
-            filename (str): The base filename for the output visualization (default: "trace_tree").
-            interested_span_match (str | None): Optional regular expression pattern to highlight or focus on specific spans in the visualization.
+            source: Collection of Agent Lightning [`Span`][agentlightning.Span] objects
+                or raw `opentelemetry.sdk.trace.ReadableSpan` instances.
+            filename: Base filename for the generated image; `.png` is appended automatically.
+            interested_span_match: Optional regular expression used to highlight a subset of spans.
 
         Returns:
-            TraceTree: The constructed trace tree object.
+            The [`TraceTree`][agentlightning.adapter.triplet.TraceTree] built from the provided
+            spans.
         """
         source_normalized = [
             Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
@@ -582,8 +828,15 @@ 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
-        """Convert OpenTelemetry spans to a list of Triplet objects."""
+    def adapt(self, source: Union[Sequence[Span], Sequence[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
+        """Convert tracer spans into [`Triplet`][agentlightning.Triplet] trajectories.
+
+        Args:
+            source: Agent Lightning spans or raw OpenTelemetry spans that form a trace.
+
+        Returns:
+            Ordered list of trajectory transitions with prompt, response, and reward information.
+        """
         source_normalized = [
             Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
             for span in source
@@ -596,30 +849,29 @@ class TracerTraceToTriplet(TraceToTripletBase):
             agent_match=self.agent_match,
             exclude_llm_call_in_reward=self.exclude_llm_call_in_reward,
             reward_match=self.reward_match,
+            _skip_empty_token_spans=self._skip_empty_token_spans,
         )
         return trajectory
 
 
 class LlmProxyTraceToTriplet(TraceToTripletBase):
-    """
-    Converting telemetry data emitted by the LLM Proxy to triplet data.
-    This adapter is very experimental. Should only be used when the TracerTraceToTriplet does not work at all.
+    """Convert telemetry emitted by the LLM Proxy into triplet trajectories.
 
-    IMPORTANT: Do NOT rely on timestamps here. Proxy spans can be emitted from different
-    machines with unsynchronized clocks. We therefore treat `sequence_id` as the only
-    reliable ordering primitive and perform "first occurrence" reward matching using
-    sequence order only.
+    !!! warning
+        This adapter is experimental and might be merged with
+        [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] in the future.
+
+    !!! danger
+        Do not rely on timestamps when using this adapter. Proxy spans can originate on different
+        machines with unsynchronised clocks, so `sequence_id` is treated as the sole source of
+        ordering.
 
     Strategy:
 
-    1) Sort spans by (sequence_id, start_time).
-    2) Extract LLM calls that expose prompt/response token IDs from either:
-       - litellm_request         (sometimes only metadata, ignore if no token ids)
-       - raw_gen_ai_request      (llm.hosted_vllm.* stringified fields)
-    3) Extract rewards from spans whose attributes contain an AgentOps-style
-       reward payload or explicit REWARD span.
-    4) For each reward with sequence R, assign it to the most recent *unmatched* LLM call
-       with sequence < R. Ignore timestamps completely.
+    1. Sort spans by `(sequence_id, start_time)` for deterministic processing.
+    2. Extract token identifiers from `litellm_request` or `raw_gen_ai_request` spans.
+    3. Extract rewards from spans exposing AgentOps-style payloads or explicit reward spans.
+    4. Match each reward to the most recent unmatched LLM call whose sequence is smaller.
     """
 
     def _literal_eval_maybe(self, v: Any) -> Any:
@@ -681,34 +933,23 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
         return cast(List[int], prompt_ids), cast(List[int], resp_ids)
 
     def _maybe_reward_value(self, span: Span) -> Optional[float]:
-        """
-        Parse reward from typical AgentOps payload or explicit REWARD span.
-        """
-        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
+        """Parse reward from typical AgentOps payloads or explicit reward spans."""
+        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:
+            source: Spans emitted by the LLM Proxy containing prompt, response, and reward data.
+
+        Returns:
+            Ordered trajectory transitions matched purely by `sequence_id`.
+        """
         # 1) Sort deterministically by (sequence_id, start_time).
         spans = sorted(
             source,

agentlightning/algorithm/__init__.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from .base import BaseAlgorithm
+from .base import Algorithm
 from .decorator import algo
 from .fast import Baseline, FastAlgorithm
 
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
     from .apo import APO as APOType
     from .verl import VERL as VERLType
 
-__all__ = ["BaseAlgorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
+__all__ = ["Algorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
 
 # Shortcuts for usages like algo.APO(...)
 

agentlightning/algorithm/apo/apo.py

@@ -7,22 +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 BaseAlgorithm
+from agentlightning.algorithm.base import Algorithm
+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")
@@ -56,42 +78,7 @@ APPLY_EDIT_PROMPT_FILES = [
 ]
 
 
-def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
-    """
-    Create an infinite iterator that yields batches from the dataset.
-
-    When batch_size >= dataset size, yields the entire shuffled dataset repeatedly.
-    When batch_size < dataset size, yields batches of the specified size, reshuffling
-    after each complete pass through the dataset.
-
-    Args:
-        dataset: The dataset to iterate over.
-        batch_size: The desired batch size.
-
-    Yields:
-        Sequences of tasks from the dataset. Each task appears at most once per epoch.
-    """
-    if batch_size >= len(dataset):
-        while True:
-            dataset_copy = [dataset[i] for i in range(len(dataset))]
-            random.shuffle(dataset_copy)
-            yield dataset_copy
-
-    else:
-        current_batch: List[int] = []
-        while True:
-            indices = list(range(len(dataset)))
-            random.shuffle(indices)
-            for index in indices:
-                if index in current_batch:
-                    continue
-                current_batch.append(index)
-                if len(current_batch) == batch_size:
-                    yield [dataset[index] for index in current_batch]
-                    current_batch = []
-
-
-class APO(BaseAlgorithm, Generic[T_task]):
+class APO(Algorithm, Generic[T_task]):
     """Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
 
     APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
@@ -99,14 +86,16 @@ class APO(BaseAlgorithm, Generic[T_task]):
     computes critiques based on the results, and applies edits to generate improved prompts.
 
     The algorithm operates in rounds, where each round:
+
     1. Samples parent prompts from the current beam
     2. Generates new prompts by computing textual gradients and applying edits
     3. Evaluates all candidates on a validation set
     4. Selects the top-k prompts for the next round
 
     Based on the ideas from:
-    - ProTeGi: https://aclanthology.org/2023.emnlp-main.494.pdf
-    - TextGrad: https://github.com/zou-group/textgrad
+
+    - [ProTeGi](https://aclanthology.org/2023.emnlp-main.494.pdf)
+    - [TextGrad](https://github.com/zou-group/textgrad)
     """
 
     def __init__(
@@ -337,6 +326,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
         Generate an improved prompt by computing a textual gradient and applying an edit.
 
         This is the main optimization step that:
+
         1. Computes a critique (textual gradient) based on rollout performance
         2. Uses another LLM to apply the critique and generate an improved prompt
 
@@ -391,8 +381,10 @@ class APO(BaseAlgorithm, Generic[T_task]):
             )
         return new_prompt
 
+    @with_store
     async def get_rollout_results(
         self,
+        store: LightningStore,
         rollout: List[Rollout],
         *,
         prefix: Optional[str] = None,
@@ -410,7 +402,6 @@ class APO(BaseAlgorithm, 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)
@@ -443,6 +434,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
         Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.
 
         This method:
+
         1. Adds the prompt as a named resource to the store
         2. Enqueues rollouts for each task in the dataset
         3. Waits for rollouts to complete (with timeout)
@@ -587,6 +579,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
         Generate new candidate prompts from parents using textual gradients.
 
         For each parent prompt, generates branch_factor new candidates by:
+
         1. Evaluating the parent on a training batch
         2. Computing textual gradient
         3. Applying edit to generate improved prompt
@@ -805,8 +798,12 @@ class APO(BaseAlgorithm, 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:
@@ -814,6 +811,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
         Execute the APO algorithm to optimize prompts through beam search with textual gradients.
 
         The algorithm performs iterative prompt optimization over multiple rounds:
+
         - Each round: samples parent prompts, generates new candidates via textual gradients,
           evaluates all candidates on validation data, and keeps the top performers
         - Tracks the historically best prompt across all rounds

agentlightning/algorithm/base.py

@@ -22,7 +22,7 @@ if TYPE_CHECKING:
     from agentlightning.trainer import Trainer
 
 
-class BaseAlgorithm:
+class Algorithm:
     """Algorithm is the strategy, or tuner to train the agent."""
 
     _trainer_ref: weakref.ReferenceType[Trainer] | None = None

agentlightning/algorithm/decorator.py

@@ -26,7 +26,7 @@ from agentlightning.types import Dataset, NamedResources
 if TYPE_CHECKING:
     from agentlightning.llm_proxy import LLMProxy
 
-from .base import BaseAlgorithm
+from .base import Algorithm
 
 # Algorithm function signature types
 # We've missed a lot of combinations here.
@@ -100,12 +100,13 @@ AsyncFlag = Literal[True, False]
 AF = TypeVar("AF", bound=AsyncFlag)
 
 
-class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
-    """A BaseAlgorithm that wraps a function-based algorithm implementation.
+class FunctionalAlgorithm(Algorithm, Generic[AF]):
+    """An algorithm wrapper built from a callable implementation.
 
-    This class allows users to define algorithm behavior using a simple function
-    that takes train_dataset and val_dataset parameters, rather than implementing
-    a full BaseAlgorithm subclass.
+    Functional algorithms let you provide an ordinary function instead of
+    subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
+    the callable signature to supply optional dependencies
+    such as the store, adapter, and LLM proxy.
     """
 
     @overload
@@ -115,13 +116,12 @@ class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
     def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
 
     def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
-        """
-        Initialize the FunctionalAlgorithm with an algorithm function.
+        """Wrap a function that implements algorithm behaviour.
 
         Args:
-            algorithm_func: A function that defines the algorithm's behavior.
-                           Can be sync or async with signature:
-                           (train_dataset, val_dataset) -> None
+            algorithm_func: Sync or async callable implementing the algorithm
+                contract. Arguments are detected automatically based on the
+                function signature.
         """
         super().__init__()
         self._algorithm_func = algorithm_func
@@ -156,14 +156,20 @@ class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
         train_dataset: Optional[Dataset[Any]] = None,
         val_dataset: Optional[Dataset[Any]] = None,
     ) -> Union[None, Awaitable[None]]:
-        """Execute the algorithm using the wrapped function.
+        """Execute the wrapped function with injected dependencies.
 
         Args:
-            train_dataset: The dataset to train on.
-            val_dataset: The dataset to validate on.
+            train_dataset: Optional training dataset passed through when the
+                callable declares a `train_dataset` parameter.
+            val_dataset: Optional validation dataset passed through when the
+                callable declares a `val_dataset` parameter.
 
         Returns:
-            None or Awaitable[None] if the function is async.
+            None for sync callables or an awaitable when the callable is async.
+
+        Raises:
+            TypeError: If a dataset is provided but the function signature does
+                not accept the corresponding argument.
         """
         kwargs: Dict[str, Any] = {}
         if "store" in self._sig.parameters:
@@ -217,40 +223,42 @@ def algo(
         AlgorithmFuncAsyncFallback,
     ],
 ) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
-    """Create a BaseAlgorithm from a function.
+    """Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].
 
-    This decorator allows you to define an algorithm using a simple function
-    instead of creating a full BaseAlgorithm subclass. The returned FunctionalAlgorithm
-    instance is callable, preserving the original function's behavior.
+    The decorator inspects the callable signature to decide which dependencies
+    to inject at runtime, enabling concise algorithm definitions that still
+    leverage the full training runtime.
 
     Args:
-        func: A function that defines the algorithm's behavior with signature:
-              (train_dataset, val_dataset) -> None
-              Can be sync or async.
+        func: Function implementing the algorithm logic. May be synchronous or
+            asynchronous. The function can expect all of, or a subset of the following parameters:
+
+            - `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
+            - `train_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `val_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
+            - `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
+            - `initial_resources`: [`NamedResources`][agentlightning.NamedResources],
+
+            If the function does not expect a parameter, the wrapper will not inject it into the call.
+            Using `*args` and `**kwargs` will not work and no parameters will be injected.
 
     Returns:
-        A callable FunctionalAlgorithm instance that preserves the original function's
-        type hints and behavior while providing all algorithm functionality.
+        FunctionalAlgorithm that proxies the callable while exposing the
+        `Algorithm` interface.
 
-    Example:
-        @algo
-        def my_algorithm(train_dataset, val_dataset):
-            # Algorithm logic here
-            for task in train_dataset:
-                # Process training tasks
-                pass
+    Examples:
+        ```python
+        from agentlightning.algorithm.decorator import algo
 
         @algo
-        async def my_async_algorithm(train_dataset, val_dataset):
-            # Async algorithm logic here
-            async for task in train_dataset:
-                # Process training tasks asynchronously
-                pass
+        def batching_algorithm(*, store, train_dataset, val_dataset):
+            for sample in train_dataset:
+                store.enqueue_rollout(input=sample, mode="train")
 
-        # Function is still callable with original behavior
-        my_algorithm(train_data, val_data)
-
-        # Algorithm methods are also available
-        my_algorithm.run(train_data, val_data)
+        @algo
+        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
+            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
+        ```
     """
     return FunctionalAlgorithm(func)

agentlightning/algorithm/fast.py

@@ -5,23 +5,28 @@ from __future__ import annotations
 import asyncio
 import logging
 from datetime import datetime
-from typing import Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Literal, Optional
 
-from agentlightning.llm_proxy import ModelConfig
-from agentlightning.types import Dataset, Rollout, RolloutStatus
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
 
-from .base import BaseAlgorithm
+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__)
 
 __all__ = ["FastAlgorithm", "Baseline"]
 
 
-class FastAlgorithm(BaseAlgorithm):
-    """Algorithm that can run fast and qualify for dev mode.
+class FastAlgorithm(Algorithm):
+    """Base class for lightweight algorithms optimised for developer workflows.
 
-    Fast algorithms enable agent developers to quickly iterate on agent development
-    without waiting for a long training to complete.
+    Fast algorithms prioritise short feedback loops so an agent developer can run
+    small-scale experiments without waiting for long-running training jobs to
+    finish.
     """
 
 
@@ -30,40 +35,79 @@ def _timestamp_to_iso_str(timestamp: float) -> str:
 
 
 class Baseline(FastAlgorithm):
-    """A dummy implementation of algorithm interface that puts all dataset into the queue, and waits for all rollouts to complete.
+    """Reference implementation that streams the full dataset through the rollout queue.
 
-    Logs all collected spans and rewards.
+    The baseline algorithm batches task submissions, waits for each rollout to
+    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:
-        model_list: Optional list of models to load into the llm proxy.
-            If both model_list and llm_proxy is provided, llm_proxy will be launched.
-            Not implemented yet.
-        n_epochs: Number of epochs to run through the dev dataset.
-        train_split: Fraction of dev dataset to use for training vs validation. Must be between 0 and 1.
-        polling_interval: Time interval (in seconds) to poll the store for queue length and for completed rollouts.
-        max_queue_length: Maximum number of rollouts to keep in the queue at any time.
+        n_epochs: Number of dataset passes to execute for both the train and val
+            splits during developer experiments.
+        train_split: Fraction of the concatenated dataset to treat as training
+            data. Must be strictly between 0 and 1.
+        polling_interval: Interval, in seconds, to poll the store for queue
+            depth and rollout completion.
+        max_queue_length: Number of rollouts allowed to wait in the queue before
+            throttling additional submissions.
+        span_verbosity: Level of detail to include when logging span metadata.
+
+    Raises:
+        ValueError: If `train_split` falls outside the `(0, 1)` interval.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.fast import Baseline
+
+        algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values")
+        trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val)
+        ```
     """
 
     def __init__(
         self,
         *,
-        model_list: Optional[List[ModelConfig]] = None,
         n_epochs: int = 1,
         train_split: float = 0.5,
         polling_interval: float = 5.0,
         max_queue_length: int = 4,
+        span_verbosity: Literal["keys", "key_values", "none"] = "keys",
     ) -> None:
         super().__init__()
         self.n_epochs = n_epochs
         self.train_split = train_split
         self.polling_interval = polling_interval
         self.max_queue_length = max_queue_length
+        self.span_verbosity = span_verbosity
         if not (0.0 < self.train_split < 1.0):
             raise ValueError("train_split must be between 0 and 1.")
 
         self._finished_rollout_count = 0
 
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        """Format a span for logging based on the configured verbosity."""
+        if self.span_verbosity == "none":
+            return ""
+
+        prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
+        elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
+
+        msg = (
+            prefix_msg
+            + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
+            + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
+            + f"{elapsed} seconds. "
+        )
+        if self.span_verbosity == "key_values":
+            msg += f"Attributes: {span.attributes}"
+        else:
+            msg += f"Attribute keys: {list(span.attributes.keys())}"
+        return msg
+
     async def _handle_rollout_finish(self, rollout: Rollout) -> None:
+        """Log attempt metadata and emit adapted traces when a rollout ends."""
         store = self.get_store()
 
         rollout_id = rollout.rollout_id
@@ -76,35 +120,37 @@ class Baseline(FastAlgorithm):
         attempts = await store.query_attempts(rollout_id)
         for attempt in attempts:
             logger.info(
-                f"[Rollout {rollout_id} | Attempt {attempt.sequence_id}] ID: {attempt.attempt_id}. Status: {attempt.status}. Worker: {attempt.worker_id}"
+                "[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s",
+                rollout_id,
+                attempt.sequence_id,
+                attempt.attempt_id,
+                attempt.status,
+                attempt.worker_id,
             )
             spans = await store.query_spans(rollout_id=rollout_id)
             for span in spans:
-                prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
-                elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
-                logger.info(
-                    prefix_msg
-                    + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
-                    + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
-                    + f"{elapsed} seconds. Attributes: {span.attributes}"
-                )
+                if self.span_verbosity != "none":
+                    logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
 
         # Attempts to adapt the spans using the adapter if provided
         try:
             adapter = self.get_adapter()
+        except ValueError:
+            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
+            adapter = None
+        if adapter is not None:
             spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
             transformed_data = adapter.adapt(spans)
             logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
-        except ValueError:
-            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
 
     async def _enqueue_rollouts(
         self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
     ) -> None:
+        """Submit rollouts while respecting the maximum queue length."""
         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]
@@ -114,6 +160,7 @@ class Baseline(FastAlgorithm):
             await asyncio.sleep(self.polling_interval)
 
     async def _harvest_rollout_spans(self, rollout_id: str):
+        """Poll rollout status updates until completion and log transitions."""
         store = self.get_store()
         last_status: Optional[RolloutStatus] = None
 
@@ -140,16 +187,21 @@ 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:
+        """Execute the baseline loop across the provided datasets."""
         train_dataset_length = len(train_dataset) if train_dataset is not None else 0
         val_dataset_length = len(val_dataset) if val_dataset is not None else 0
         if train_dataset_length == 0 and val_dataset_length == 0:
             logger.error(
-                "MockAlgorithm requires at least a train_dataset or val_dataset to run. No train_dataset or val_dataset is provided. Exiting."
+                "MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running."
             )
             return
 
@@ -158,8 +210,8 @@ class Baseline(FastAlgorithm):
         ]
         train_indices = list(range(0, train_dataset_length))
         val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
-
-        store = self.get_store()
+        logger.debug(f"Train indices: {train_indices}")
+        logger.debug(f"Val indices: {val_indices}")
 
         # Currently we only supports a single resource update at the start.
         initial_resources = self.get_initial_resources()
@@ -175,18 +227,24 @@ class Baseline(FastAlgorithm):
             harvest_tasks: List[asyncio.Task[None]] = []
             logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
             for index in train_indices + val_indices:
-                queuing_rollouts = await store.query_rollouts(status=["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]
-                    mode = "train" if index in train_indices else "val"
-                    rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
-                    harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
-                    logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
-                else:
-                    # Sleep a bit and try again later.
-                    await asyncio.sleep(self.polling_interval)
+                logger.info(
+                    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_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]
+                        mode = "train" if index in train_indices else "val"
+                        rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                        harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
+                        logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
+                        break
+                    else:
+                        # Sleep a bit and try again later.
+                        await asyncio.sleep(self.polling_interval)
 
             # Wait for all harvest tasks to complete
+            logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
             if len(harvest_tasks) > 0:
                 await asyncio.gather(*harvest_tasks)

agentlightning/algorithm/utils.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import logging
+import random
+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]]:
+    """
+    Create an infinite iterator that yields batches from the dataset.
+
+    When batch_size >= dataset size, yields the entire shuffled dataset repeatedly.
+    When batch_size < dataset size, yields batches of the specified size, reshuffling
+    after each complete pass through the dataset.
+
+    Args:
+        dataset: The dataset to iterate over.
+        batch_size: The desired batch size.
+
+    Yields:
+        Sequences of tasks from the dataset. Each task appears at most once per epoch.
+    """
+    if batch_size >= len(dataset):
+        while True:
+            dataset_copy = [dataset[i] for i in range(len(dataset))]
+            random.shuffle(dataset_copy)
+            yield dataset_copy
+
+    else:
+        current_batch: List[int] = []
+        while True:
+            indices = list(range(len(dataset)))
+            random.shuffle(indices)
+            for index in indices:
+                if index in current_batch:
+                    continue
+                current_batch.append(index)
+                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

agentlightning/algorithm/verl/interface.py

@@ -1,30 +1,133 @@
 # 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
 
-from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.algorithm.base import Algorithm
 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(BaseAlgorithm):
-    """Algorithm leveraging VERL as the backend framework.
 
-    **Note on Customization:**
+class VERL(Algorithm):
+    """VERL-powered algorithm that delegates training to the VERL PPO runner.
 
-    At present, we recommend copying the source code from VERL and modifying it as needed to suit your requirements.
-    Native support for customizing training logic will be provided in future releases.
+    !!! warning
+        Advanced customisation currently requires copying the VERL source and
+        modifying it directly. Native hooks for overriding training behaviour
+        will land in a future release.
 
     Args:
-        config: The VERL configuration, matching what is typically provided when running VERL via the command line.
-            This config will be merged with VERL's base configuration and processed by Hydra.
+        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
+        from agentlightning.algorithm.verl import VERL
+
+        algorithm = VERL(
+            config={
+                "algorithm": {
+                    "adv_estimator": "grpo",
+                    "use_kl_in_reward": False,
+                },
+                "data": {
+                    "train_batch_size": 32,
+                    "max_prompt_length": 4096,
+                    "max_response_length": 2048,
+                },
+                "actor_rollout_ref": {
+                    "rollout": {
+                        "tensor_model_parallel_size": 1,
+                        "n": 4,
+                        "log_prob_micro_batch_size_per_gpu": 4,
+                        "multi_turn": {"format": "hermes"},
+                        "name": "vllm",
+                        "gpu_memory_utilization": 0.6,
+                    },
+                    "actor": {
+                        "ppo_mini_batch_size": 32,
+                        "ppo_micro_batch_size_per_gpu": 4,
+                        "optim": {"lr": 1e-6},
+                        "use_kl_loss": False,
+                        "kl_loss_coef": 0.0,
+                        "entropy_coeff": 0,
+                        "clip_ratio_low": 0.2,
+                        "clip_ratio_high": 0.3,
+                        "fsdp_config": {
+                            "param_offload": True,
+                            "optimizer_offload": True,
+                        },
+                    },
+                    "ref": {
+                        "log_prob_micro_batch_size_per_gpu": 8,
+                        "fsdp_config": {"param_offload": True},
+                    },
+                    "model": {
+                        "path": "Qwen/Qwen2.5-1.5B-Instruct",
+                        "use_remove_padding": True,
+                        "enable_gradient_checkpointing": True,
+                    },
+                },
+                "trainer": {
+                    "n_gpus_per_node": 1,
+                    "val_before_train": True,
+                    "critic_warmup": 0,
+                    "logger": ["console", "wandb"],
+                    "project_name": "AgentLightning",
+                    "experiment_name": "calc_x",
+                    "nnodes": 1,
+                    "save_freq": 64,
+                    "test_freq": 32,
+                    "total_epochs": 2,
+                },
+            }
+        )
+        trainer.fit(algorithm, train_dataset=my_train_dataset)
+        ```
     """
 
-    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:
@@ -33,13 +136,33 @@ class VERL(BaseAlgorithm):
 
         # Merge your dict overrides
         override_conf = OmegaConf.create(config)
+        # 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,
         train_dataset: Optional[Dataset[Any]] = None,
         val_dataset: Optional[Dataset[Any]] = None,
     ) -> None:
+        """Launch the VERL PPO entrypoint with the configured runtime context.
+
+        Args:
+            train_dataset: Optional dataset forwarded to VERL for training.
+            val_dataset: Optional dataset forwarded to VERL for evaluation.
+
+        Raises:
+            ValueError: If required dependencies such as the store, LLM proxy, or
+                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:
@@ -51,6 +174,8 @@ class VERL(BaseAlgorithm):
                 store=None,
                 llm_proxy=None,
                 adapter=None,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
             )
         else:
             print("Store is set. Assuming v1 execution mode.")
@@ -63,8 +188,15 @@ class VERL(BaseAlgorithm):
                 store=store,
                 llm_proxy=llm_proxy,
                 adapter=adapter,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
             )
 
     def get_client(self) -> AgentLightningClient:
+        """Create a client bound to the VERL-managed Agent Lightning server.
+
+        Deprecated:
+            Since v0.2.
+        """
         port = self.config.agentlightning.port
         return AgentLightningClient(endpoint=f"http://localhost:{port}")

agentlightning/cli/__init__.py

@@ -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."),
 }
 

agentlightning/cli/agentops_server.py

@@ -1,30 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import argparse
-import time
-from typing import Iterable
-
-from agentlightning.instrumentation.agentops import AgentOpsServerManager
-
-
-def main(argv: Iterable[str] | None = None) -> int:
-    parser = argparse.ArgumentParser(description="Start AgentOps server")
-    parser.add_argument("--daemon", action="store_true", help="Run server as a daemon")
-    parser.add_argument("--port", type=int, default=8002, help="Port to run the server on")
-    args = parser.parse_args(list(argv) if argv is not None else None)
-
-    manager = AgentOpsServerManager(daemon=args.daemon, port=args.port)
-    try:
-        manager.start()
-        # Wait forever
-        while True:
-            time.sleep(1)
-    except KeyboardInterrupt:
-        manager.stop()
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

agentlightning/cli/prometheus.py

@@ -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())

agentlightning/cli/store.py

@@ -6,23 +6,124 @@ from __future__ import annotations
 
 import argparse
 import asyncio
-from typing import Iterable
+import logging
+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",
+        dest="cors_origins",
+        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()
-    server = LightningStoreServer(store, host="0.0.0.0", port=args.port)
-    asyncio.run(server.run_forever())
+    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=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())
+    except RuntimeError as exc:
+        logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
+        return 1
     return 0
 
 

agentlightning/client.py

@@ -1,6 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Legacy client for interacting with a legacy Agent Lightning server."""
+"""Utilities for interacting with legacy Agent Lightning servers.
+
+This module contains compatibility shims that speak the deprecated HTTP
+interface used by older Agent Lightning deployments. Modern code should prefer
+the store-based APIs exposed by `agentlightning.store`, but keeping these
+clients available makes it easier to migrate existing workflows incrementally.
+"""
 
 import asyncio
 import logging
@@ -18,13 +24,24 @@ logger = logging.getLogger(__name__)
 
 
 class AgentLightningClient:
-    """
-    Client for interacting with a version-aware Agent Lightning Server.
+    """Client wrapper for the legacy version-aware Agent Lightning server.
 
-    This client handles polling for tasks, fetching specific versions of resources
-    (like model configurations), and posting completed rollouts back to the server.
-    It provides both synchronous and asynchronous methods for these operations and
-    includes a cache for resources.
+    The client exposes synchronous and asynchronous helpers for polling tasks,
+    retrieving resource bundles, and submitting rollouts. It also maintains a
+    simple in-memory cache keyed by the server-provided resource identifier to
+    avoid redundant network requests.
+
+    !!! warning "Deprecated"
+        [`AgentLightningClient`][agentlightning.client.AgentLightningClient] is part of
+        the legacy client/server stack. New code should rely on the store-based APIs
+        implemented in `agentlightning.store`.
+
+    Attributes:
+        endpoint: Base URL of the Agent Lightning server.
+        poll_interval: Delay in seconds between polling attempts when no task is
+            available.
+        timeout: Timeout in seconds applied to HTTP requests.
+        task_count: Number of tasks claimed during the lifetime of this client.
     """
 
     _next_task_uri = "/task"
@@ -33,12 +50,12 @@ class AgentLightningClient:
     _report_rollout_uri = "/rollout"
 
     def __init__(self, endpoint: str, poll_interval: float = 5.0, timeout: float = 10.0):
-        """Initializes the AgentLightningClient.
+        """Initialize the client.
 
         Args:
-            endpoint: The root URL of the Agent Lightning server.
-            poll_interval: The interval in seconds to wait between polling for new tasks.
-            timeout: The timeout in seconds for HTTP requests.
+            endpoint: Root URL of the Agent Lightning server.
+            poll_interval: Seconds to wait between polling attempts.
+            timeout: Seconds before a request to the server is considered timed out.
         """
         warnings.warn(
             "AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
@@ -51,13 +68,13 @@ class AgentLightningClient:
         self._default_headers = {"X-AgentLightning-Client": "true"}
 
     async def _request_json_async(self, url: str) -> Optional[Dict[str, Any]]:
-        """Makes an async GET request to the specified URL and returns the JSON response.
+        """Perform an asynchronous ``GET`` request and parse the JSON payload.
 
         Args:
-            url: The URL to request.
+            url: Fully qualified URL to query.
 
         Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
         """
         timeout = aiohttp.ClientTimeout(total=self.timeout)
         async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -70,14 +87,14 @@ class AgentLightningClient:
                 return None
 
     async def _post_json_async(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        """Makes an async POST request with a JSON payload.
+        """Perform an asynchronous ``POST`` request with a JSON body.
 
         Args:
-            url: The URL to post to.
-            payload: The dictionary data to send as JSON.
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.
 
         Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
         """
         timeout = aiohttp.ClientTimeout(total=self.timeout)
         async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -90,10 +107,11 @@ class AgentLightningClient:
                 return None
 
     async def poll_next_task_async(self) -> Optional[Task]:
-        """Polls the server asynchronously for the next task until one is available.
+        """Poll the server asynchronously until a task becomes available.
 
         Returns:
-            A Task object containing the task details.
+            The next [`Task`][agentlightning.Task] exposed by the server,
+            or ``None`` if polling fails.
         """
         url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
         while True:
@@ -108,13 +126,15 @@ class AgentLightningClient:
             await asyncio.sleep(self.poll_interval)
 
     async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
-        """Fetches a specific version of resources by its ID, using a cache.
+        """Fetch a specific resource bundle by identifier.
 
         Args:
-            resource_id: The ID of the resources to fetch, usually from a Task's metadata.
+            resource_id: Identifier sourced from the task metadata.
 
         Returns:
-            A ResourcesUpdate object containing the versioned resources, or None if not found.
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
         """
         if resource_id in self._resource_cache:
             logger.debug(f"Found resources '{resource_id}' in cache.")
@@ -130,10 +150,11 @@ class AgentLightningClient:
         return None
 
     async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
-        """Fetches the latest available resources from the server.
+        """Fetch the most recent resource bundle advertised by the server.
 
         Returns:
-            A ResourcesUpdate object containing the latest resources.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
         """
         url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
         response = await self._request_json_async(url)
@@ -145,26 +166,26 @@ class AgentLightningClient:
         return None
 
     async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
-        """Posts a completed rollout to the server asynchronously.
+        """Submit a completed rollout back to the server.
 
         Args:
-            rollout: A Rollout object containing the results of a task.
+            rollout: Legacy rollout payload produced by the executor.
 
         Returns:
-            The server's JSON response as a dictionary.
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
         """
         url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
         payload = rollout.model_dump(mode="json")
         return await self._post_json_async(url, payload)
 
     def _request_json(self, url: str) -> Optional[Dict[str, Any]]:
-        """Makes a sync GET request to the specified URL and returns the JSON response.
+        """Perform a blocking ``GET`` request and parse the JSON payload.
 
         Args:
-            url: The URL to request.
+            url: Fully qualified URL to query.
 
         Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
         """
         try:
             response = requests.get(url, timeout=self.timeout, headers=self._default_headers)
@@ -175,14 +196,14 @@ class AgentLightningClient:
             return None
 
     def _post_json(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        """Makes a sync POST request with a JSON payload.
+        """Perform a blocking ``POST`` request with a JSON payload.
 
         Args:
-            url: The URL to post to.
-            payload: The dictionary data to send as JSON.
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.
 
         Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
         """
         try:
             response = requests.post(url, json=payload, timeout=self.timeout, headers=self._default_headers)
@@ -193,10 +214,11 @@ class AgentLightningClient:
             return None
 
     def poll_next_task(self) -> Optional[Task]:
-        """Polls the server synchronously for the next task until one is available.
+        """Poll the server synchronously until a task becomes available.
 
         Returns:
-            A Task object containing the task details, including the required `resources_id`.
+            The next [`Task`][agentlightning.Task] available for execution, or
+            ``None`` if polling fails.
         """
         url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
         while True:
@@ -211,13 +233,15 @@ class AgentLightningClient:
             time.sleep(self.poll_interval)
 
     def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
-        """Fetches a specific version of resources by its ID synchronously, using a cache.
+        """Fetch a specific resource bundle by identifier.
 
         Args:
-            resource_id: The ID of the resources to fetch, usually from a Task's metadata.
+            resource_id: Identifier sourced from the task metadata.
 
         Returns:
-            A ResourcesUpdate object containing the versioned resources, or None if not found.
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
         """
         if resource_id in self._resource_cache:
             logger.debug(f"Found resources '{resource_id}' in cache.")
@@ -233,10 +257,11 @@ class AgentLightningClient:
         return None
 
     def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """Fetches the latest available resources from the server synchronously.
+        """Fetch the most recent resource bundle advertised by the server.
 
         Returns:
-            A ResourcesUpdate object containing the latest resources.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
         """
         url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
         response = self._request_json(url)
@@ -247,13 +272,13 @@ class AgentLightningClient:
         return None
 
     def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
-        """Posts a completed rollout to the server synchronously.
+        """Submit a completed rollout back to the server.
 
         Args:
-            rollout: A Rollout object containing the results of a task.
+            rollout: Legacy rollout payload produced by the executor.
 
         Returns:
-            The server's JSON response as a dictionary.
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
         """
         url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
         payload = rollout.model_dump(mode="json")
@@ -261,14 +286,16 @@ class AgentLightningClient:
 
 
 class DevTaskLoader(AgentLightningClient):
-    """A local task manager for development that provides sample tasks and resources.
+    """In-memory task loader used for development and integration tests.
 
-    This client mocks the server APIs by maintaining a local queue of tasks and resources
-    within the same process. It's designed for development, testing, and scenarios where
-    a full Agent Lightning server is not needed.
+    The loader mimics the behavior of the legacy HTTP server by storing tasks and
+    resources locally. Polling methods simply iterate over the provided collection,
+    allowing rapid iteration without provisioning any external infrastructure.
 
-    The DevTaskLoader overrides the polling and resource fetching methods to return data
-    from local collections instead of making HTTP requests to a remote server.
+    !!! warning "Deprecated"
+
+        [`DevTaskLoader`][agentlightning.client.DevTaskLoader] is a compatibility shim.
+        Prefer [`Trainer.dev`][agentlightning.Trainer.dev] for new code.
     """
 
     def __init__(
@@ -277,12 +304,17 @@ class DevTaskLoader(AgentLightningClient):
         resources: Union[NamedResources, ResourcesUpdate],
         **kwargs: Any,
     ):
-        """Initializes the DevTaskLoader with pre-defined tasks and resources.
+        """Initialize the loader with predefined tasks and resources.
 
         Args:
-            tasks: Either a List of TaskInput objects or a List of Task objects.
-            resources: Either NamedResources or ResourcesUpdate object.
-            **kwargs: Additional arguments passed to the parent AgentLightningClient.
+            tasks: Sequence of task inputs or preconstructed tasks that will be served in
+                order.
+            resources: Static resources returned for any `resources_id` query.
+            **kwargs: Additional keyword arguments forwarded to the parent client.
+
+        Raises:
+            ValueError: If no tasks are provided or both [`Task`][agentlightning.Task]
+                and [`TaskInput`][agentlightning.TaskInput] instances are mixed.
         """
         warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
         super().__init__(endpoint="local://", **kwargs)
@@ -300,24 +332,27 @@ class DevTaskLoader(AgentLightningClient):
         if isinstance(resources, ResourcesUpdate):
             self._resources_update = resources
         else:
-            self._resources_update = ResourcesUpdate(resources_id="local", resources=resources)
+            self._resources_update = ResourcesUpdate(
+                resources_id="local", resources=resources, create_time=time.time(), update_time=time.time(), version=1
+            )
 
         # Store rollouts posted back to the loader for easy debugging of local runs
         self._rollouts: List[RolloutLegacy] = []
 
     @property
     def rollouts(self) -> List[RolloutLegacy]:
-        """Return rollouts that have been posted back to the loader."""
+        """Return the rollouts posted back to the loader during development runs."""
         return self._rollouts
 
     def poll_next_task(self) -> Optional[Task]:
-        """Returns the next task from the local queue.
+        """Return the next task from the local queue.
 
-        If tasks are TaskInput objects, assembles them into Task objects.
-        If tasks are already Task objects, returns them directly.
+        If [`TaskInput`][agentlightning.TaskInput] instances were provided,
+        they are converted into [`Task`][agentlightning.Task] objects on the
+        fly. Otherwise, the preconstructed tasks are returned in sequence.
 
         Returns:
-            The next Task object from the local task list.
+            Next task to execute.
         """
         if self._task_index >= len(self._tasks):
             self._task_index = 0

agentlightning/config.py

@@ -83,12 +83,17 @@ def _str_to_bool(v: str) -> bool:
 
 
 def _get_param_type_details(param_annotation: Any) -> Tuple[Any, bool, bool]:
-    """
-    Determines the core type, if it's Optional, and if it's a List.
-    Returns: (core_type, is_optional, is_list)
-    - For Optional[T]: (T, True, is_list_status_of_T)
-    - For List[T]: (List[T], is_optional_status_of_List, True)
-    - For Optional[List[T]]: (List[T], True, True)
+    """Normalize an annotation into its core type, optionality, and list status.
+
+    Args:
+        param_annotation: The annotation to inspect.
+
+    Returns:
+        A tuple ``(core_type, is_optional, is_list)`` describing the normalized type.
+
+        - For ``Optional[T]`` → ``(T, True, is_list_status_of_T)``
+        - For ``List[T]`` → ``(List[T], is_optional_status_of_List, True)``
+        - For ``Optional[List[T]]`` → ``(List[T], True, True)``
     """
     is_optional = False
     is_list = False
@@ -309,6 +314,9 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
 
 
+# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
+
+
 def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
     """
     Parses command-line arguments to configure and instantiate provided CliConfigurable classes.

agentlightning/emitter/__init__.py

@@ -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",
 ]

agentlightning/emitter/annotation.py

@@ -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)

agentlightning/emitter/exception.py

@@ -1,38 +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:
-    """Emit an exception as a span."""
+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. 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"),
+    )

agentlightning/emitter/message.py

@@ -1,29 +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:
-    """Emit a string message as a span.
+def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
+    """Emit a textual message as an OpenTelemetry span.
 
-    OpenTelemetry has a dedicated design of logs by design, but we can also use spans to emit messages.
-    So that it can all be unified in the data store and analyzed together.
+    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)}.")

agentlightning/emitter/object.py

@@ -1,29 +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:
-    """Emit any object as a span. Make sure the object is JSON serializable."""
-    try:
-        serialized = json.dumps(object)
-    except (TypeError, ValueError):
-        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
-        return
+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.
 
-    tracer = get_tracer()
-    span = tracer.start_span(
-        SpanNames.OBJECT.value,
-        attributes={SpanAttributeNames.OBJECT.value: serialized},
+    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 will lead to a RuntimeError.
+    """
+    span_attributes = encode_object(object)
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
+
+    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

agentlightning/emitter/reward.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""Helpers for emitting reward spans and integrating with AgentOps telemetry."""
+
 import asyncio
 import inspect
 import json
@@ -18,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__)
 
@@ -32,35 +34,56 @@ __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:
-    """Check if AgentOps is initialized in the current context."""
+    """Return `True` when the AgentOps client has been configured."""
+    import agentops
+
     return agentops.get_client().initialized
 
 
-def reward(fn: FnType) -> FnType:
-    """
-    A decorator to wrap a function that computes rewards.
-    It will automatically handle the input and output of the function.
+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
+    the built-in telemetry otherwise. Both synchronous and asynchronous functions
+    are supported transparently.
+
+    Deprecated:
+        This decorator is deprecated. Use [`emit_reward`][agentlightning.emit_reward] instead.
+
+    Args:
+        fn: Callable that produces a numeric reward.
+
+    Returns:
+        Wrapped callable that preserves the original signature.
     """
 
-    def wrap_result(result: Optional[float]) -> RewardSpanData:
-        """
-        Wrap the result of the function in a dict.
-        """
+    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}
         if not isinstance(result, (float, int)):  # type: ignore
@@ -83,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
@@ -107,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)
@@ -118,30 +141,86 @@ def reward(fn: FnType) -> FnType:
         return wrapper  # type: ignore
 
 
-def emit_reward(reward: float) -> ReadableSpan:
-    """
-    Record a new reward as a new span.
+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:
+        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))
 
-    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]:
+    """Extract the reward value from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or Agent Lightning emitters.
+
+    Returns:
+        The primary reward encoded in the span or `None` when the span does not represent a reward.
     """
-    Get the reward value from a span.
-    """
+    # 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",
@@ -164,49 +243,71 @@ 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:
-    """
-    Check if a span is a reward span.
-    """
+    """Return ``True`` when the provided span encodes a reward value."""
     maybe_reward = get_reward_value(span)
     return maybe_reward is not None
 
 
 def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
-    """
-    Find all reward spans in the given list of spans.
+    """Return all reward spans in the provided sequence.
 
     Args:
-        spans: A list of spans (either ReadableSpan or Span).
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
 
     Returns:
-        A list of spans whose name matches the reward span name.
+        List of spans that could be parsed as rewards.
     """
     return [span for span in spans if is_reward_span(span)]
 
 
 def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
-    """
-    Get the last reward value from a list of spans.
+    """Return the last reward value present in the provided spans.
 
     Args:
-        spans: A list of spans (either ReadableSpan or Span).
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
 
     Returns:
-        The reward value from the last reward span, or None if not found.
+        Reward value from the latest reward span, or `None` when none are found.
     """
     for span in reversed(spans):
         reward = get_reward_value(span)

agentlightning/emitter/utils.py

@@ -1,22 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Common utilities for the emitter module."""
-
-import opentelemetry.trace as trace_api
-from opentelemetry.trace import get_tracer_provider
-
-
-def get_tracer() -> trace_api.Tracer:
-    """Return the tracer used for AgentLightning spans.
-
-    Raises:
-        RuntimeError: If the tracer is not initialized.
-
-    Returns:
-        The AgentLightning tracer instance.
-    """
-    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")

agentlightning/env_var.py

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

agentlightning/execution/base.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import logging
 from typing import Protocol
 
@@ -11,27 +13,52 @@ logger = logging.getLogger(__name__)
 
 
 class AlgorithmBundle(Protocol):
+    """Callable bundle produced by [`Trainer`][agentlightning.Trainer].
+
+    Execution strategies treat the returned coroutine as opaque, only providing
+    the shared store instance and cooperative stop event. Bundles typically
+    encapsulate algorithm setup plus adapter and LLM proxy, etc.
+    """
+
     async def __call__(self, store: LightningStore, event: ExecutionEvent) -> None:
-        """Initalization and execution logic."""
+        """Execute algorithm logic using ``store`` until completion or stop."""
 
 
 class RunnerBundle(Protocol):
+    """Callable bundle wrapping runner setup and the worker loop, as opposed to the
+    [`AlgorithmBundle`][agentlightning.AlgorithmBundle]."""
+
     async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
-        """Initalization and execution logic."""
+        """Execute runner logic for ``worker_id`` using ``store`` and ``event``."""
 
 
 class ExecutionStrategy:
-    """When trainer has created the executable of algorithm and runner in two bundles,
-    the execution strategy defines how to run them together, and how many parallel runners to run.
+    """Coordinate algorithm and runner bundles within a single process abstraction.
 
-    The store is the centric place for the two bundles to communicate.
+    Strategies decide how many worker bundles to launch, whether to communicate
+    through shared memory or an HTTP boundary, and how to react to shutdown
+    signals. They intentionally avoid inspecting the bundle internals; instead,
+    each bundle remains responsible for its own scheduling semantics.
 
-    The algorithm and runner's behavior (whether runner should perform one step or run forever,
-    whether the algo would send out the tasks or not) are defined inside the bundle,
-    and does not belong to the execution strategy.
-
-    The execute should support Ctrl+C to exit gracefully.
+    !!! note
+        Implementations must honor the [execute()][agentlightning.ExecutionStrategy.execute]
+        contract by propagating `KeyboardInterrupt` and ensuring resources are
+        released when an error occurs on either side of the algorithm/runner
+        pair.
     """
 
     def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        """Run the provided bundles using the configured orchestration model.
+
+        Args:
+            algorithm: Callable bundle responsible for algorithm execution.
+            runner: Callable bundle for runner workers.
+            store: Concrete [`LightningStore`][agentlightning.LightningStore]
+                shared across bundles.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the orchestration
+                implementation.
+        """
+
         raise NotImplementedError()

agentlightning/execution/client_server.py

@@ -9,6 +9,7 @@ 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
 
@@ -19,45 +20,44 @@ logger = logging.getLogger(__name__)
 
 
 class ClientServerExecutionStrategy(ExecutionStrategy):
-    """Run algorithm (server) and runners (clients) as separate processes over HTTP.
+    """Run algorithm and runner bundles as separate processes over HTTP.
 
-    **Execution Roles:**
+    Execution Roles:
 
-    - "algorithm": Start the HTTP server (`LightningStoreServer`) in-process and run the
-      algorithm bundle against it.
-    - "runner": Connect to an already running server via `LightningStoreClient` and
-      execute runner bundles (optionally in multiple processes).
-    - "both": Spawn the runner processes first, then launch the algorithm/server
-      bundle on the main process. This mode orchestrates the full loop locally.
+    - `"algorithm"`: Start [`LightningStoreServer`][agentlightning.LightningStoreServer]
+      in-process and execute the algorithm bundle against it.
+    - `"runner"`: Connect to an existing server with
+      [`LightningStoreClient`][agentlightning.LightningStoreClient] and run the
+      runner bundle locally (spawning multiple processes when requested).
+    - `"both"`: Spawn runner processes first, then execute the algorithm and
+      server on the same machine. This mode orchestrates the full loop locally.
 
-    When role == "both", you may choose which side runs on the main process via
-    `main_process` (debug helper). Running the runner bundle on the main process
-    is only supported with `n_runners == 1`.
+    When `role == "both"` you may choose which side runs on the main process
+    via `main_process`. The runner-on-main option is limited to
+    `n_runners == 1` because each additional runner requires its own event
+    loop and process.
 
-    Important: When `main_process == "runner"`, the algorithm runs in a subprocess
-    with the LightningStore server. This means any state modifications made during
-    execution remain in that subprocess and are NOT reflected in the original store
-    object passed to `execute()`. The main process runner accesses the store only
-    through the HTTP client interface.
+    !!! warning
+        When `main_process == "runner"` the algorithm and HTTP server execute
+        in a child process. Store mutations remain isolated inside that process,
+        so the original store instance passed to
+        [execute()][agentlightning.ExecutionStrategy.execute] is not updated.
 
-    **Abort / Stop Model (four-step escalation):**
+    Abort Model (four-step escalation):
 
-    1. Cooperative stop:
-       A shared :class:`~agentlightning.execution.events.MultiprocessingEvent`
-       (`stop_evt`) is passed to *all* bundles. Bundles should check it to exit.
-       Any crash (algorithm or runner) sets `stop_evt` so the other side can
-       stop cooperatively. Ctrl+C on the main process also flips the event.
-    2. KeyboardInterrupt synth:
-       Remaining subprocesses receive `SIGINT` to trigger `KeyboardInterrupt`
-       handlers.
-    3. Termination:
-       Stubborn subprocesses get `terminate()` (SIGTERM on POSIX).
-    4. Kill:
-        As a last resort we call `kill()` (SIGKILL on POSIX).
+    1. Cooperative stop. Every bundle receives a shared
+       [`MultiprocessingEvent`][agentlightning.MultiprocessingEvent] (`stop_evt`).
+       Any failure flips the event so peers can exit cleanly. Ctrl+C on the main
+       process also sets the flag.
+    2. KeyboardInterrupt synthesis. Remaining subprocesses receive ``SIGINT`` to
+       trigger `KeyboardInterrupt` handlers.
+    3. Termination. Stubborn processes are asked to ``terminate()``
+       (`SIGTERM` on POSIX).
+    4. Kill. As a last resort `kill()` is invoked (`SIGKILL` on POSIX).
 
-    Notes:
-        This mirrors the semantics implemented in :mod:`shared_memory`, but adapted
-        to multiple processes and the HTTP client/server boundary.
+    This mirrors the semantics implemented in
+    [`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy]
+    but adapts them to multiple processes and the HTTP client/server boundary.
     """
 
     alias: str = "cs"
@@ -68,20 +68,22 @@ 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.
 
         Args:
             role: Which side(s) to run in this process. When omitted, the
-                :envvar:`AGL_CURRENT_ROLE` environment variable is used.
+                `AGL_CURRENT_ROLE` environment variable is used.
             server_host: Interface the HTTP server binds to when running the
-                algorithm bundle locally. Defaults to :envvar:`AGL_SERVER_HOST`
-                or ``"localhost"`` if unset.
+                algorithm bundle locally. Defaults to `AGL_SERVER_HOST`
+                or `"localhost"` if unset.
             server_port: Port for the HTTP server in "algorithm"/"both" modes.
-                Defaults to :envvar:`AGL_SERVER_PORT` or ``4747`` if unset.
+                Defaults to `AGL_SERVER_PORT` or `4747` if unset.
             n_runners: Number of runner processes to spawn in "runner"/"both".
             graceful_timeout: How long to wait (seconds) after setting the stop
                 event before escalating to signals.
@@ -90,56 +92,61 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
             main_process: Which bundle runs on the main process when
                 `role == "both"`. `"runner"` requires `n_runners == 1` and is
                 primarily intended for debugging.
+            managed_store: When `True` (default) the strategy constructs
+                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:
-                raise ValueError("role must be provided via argument or AGL_CURRENT_ROLE env var")
-            if role_env not in ("algorithm", "runner", "both"):
-                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
-            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_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
     ) -> None:
-        logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
-        server_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
-        server_started = False
+        wrapper_store: LightningStore | None = None
+        if self.managed_store:
+            logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
+            wrapper_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
+            server_started = False
+        else:
+            wrapper_store = store
+            server_started = False
 
         try:
-            await server_store.start()
-            server_started = True
-            logger.debug("Algorithm bundle starting against endpoint %s", server_store.endpoint)
-            await algorithm(server_store, stop_evt)
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer):
+                await wrapper_store.start()
+                server_started = True
+                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()
@@ -149,20 +156,37 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
             stop_evt.set()
             raise
         finally:
-            if server_started:
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer) and server_started:
                 try:
-                    await server_store.stop()
+                    await wrapper_store.stop()
                 except Exception:
                     logger.exception("Error stopping LightningStore server")
                 else:
                     logger.debug("LightningStore server shutdown completed")
 
-    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
-        client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+    async def _execute_runner(
+        self,
+        runner: RunnerBundle,
+        worker_id: int,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+    ) -> None:
+        if self.managed_store:
+            # If managed, we actually do not use the provided store
+            client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+        else:
+            client_store = store
         try:
-            logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            if self.managed_store:
+                logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            else:
+                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()
@@ -172,16 +196,18 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
             stop_evt.set()
             raise
         finally:
-            try:
-                await client_store.close()
-            except Exception:
-                logger.exception("Error closing LightningStore client for runner %s", worker_id)
-            else:
-                logger.debug("Runner %s closed LightningStore client", worker_id)
+            if self.managed_store and isinstance(client_store, LightningStoreClient):
+                try:
+                    await client_store.close()
+                except Exception:
+                    logger.exception("Error closing LightningStore client for runner %s", worker_id)
+                else:
+                    logger.debug("Runner %s closed LightningStore client", worker_id)
 
     def _spawn_runners(
         self,
         runner: RunnerBundle,
+        store: LightningStore,
         stop_evt: ExecutionEvent,
         *,
         ctx: BaseContext,
@@ -189,15 +215,21 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
         """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
         processes: list[multiprocessing.Process] = []
 
-        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
+        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, 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(
                 multiprocessing.Process,
-                ctx.Process(target=_runner_sync, args=(runner, i, stop_evt), name=f"runner-{i}"),  # type: ignore
+                ctx.Process(target=_runner_sync, args=(runner, i, store, stop_evt), name=f"runner-{i}"),  # type: ignore
             )
             process.start()
             logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
@@ -216,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,
@@ -310,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(
@@ -341,10 +379,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
             elif self.role == "runner":
                 if self.n_runners == 1:
                     logger.info("Running runner solely...")
-                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
                 else:
                     logger.info("Spawning runner processes...")
-                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
                     # Wait for the processes to finish naturally.
                     for process in processes:
                         process.join()
@@ -352,7 +390,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
             elif self.role == "both":
                 if self.main_process == "algorithm":
                     logger.info("Spawning runner processes...")
-                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
                     try:
                         logger.info("Running algorithm...")
                         asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
@@ -373,7 +411,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                     # the background process spawned above (the provided
                     # store must therefore be picklable when using spawn).
                     logger.info("Running runner...")
-                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
 
                     # Wait for the algorithm process to finish.
                     algorithm_process.join()

agentlightning/execution/events.py

@@ -7,15 +7,18 @@ from typing import Optional, Protocol
 
 
 class ExecutionEvent(Protocol):
-    """
-    A minimal protocol similar to threading.Event.
+    """Protocol capturing the cooperative stop contract shared by strategies.
+
+    Implementations mirror the API of ``threading.Event`` and
+    ``multiprocessing.Event`` so the rest of the execution layer can remain
+    agnostic to the underlying concurrency primitive.
 
     Methods:
-        set(): Signal event like a cancellation (idempotent).
-        clear(): Reset to the non-set state.
-        is_set() -> bool: True if event has been signaled.
-        wait(timeout: Optional[float] = None) -> bool:
-            Block until event is set or timeout. Returns True if event has signaled.
+
+        set: Signal cancellation. The call must be idempotent.
+        clear: Reset the event to the unsignaled state.
+        is_set: Return ``True`` when cancellation has been requested.
+        wait: Block until the event is signaled or an optional timeout elapses.
     """
 
     def set(self) -> None: ...
@@ -25,11 +28,7 @@ class ExecutionEvent(Protocol):
 
 
 class ThreadingEvent:
-    """
-    An Event implementation using threading.Event.
-
-    Provides a thread-safe event object for signaling between threads.
-    """
+    """Thread-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
 
     __slots__ = ("_evt",)
 
@@ -50,12 +49,7 @@ class ThreadingEvent:
 
 
 class MultiprocessingEvent:
-    """
-    An Event implementation using multiprocessing.Event.
-
-    Provides a process-safe event object for signaling between processes.
-    Optionally accepts a multiprocessing context for custom process start methods.
-    """
+    """Process-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
 
     __slots__ = ("_evt",)
 

agentlightning/execution/inter_process.py

@@ -4,6 +4,12 @@ from .base import ExecutionStrategy
 
 
 class InterProcessExecutionStrategy(ExecutionStrategy):
+    """Placeholder strategy for future inter-process primitives.
+
+    The class exists to reserve the `ipc` alias and make the planned
+    implementation discoverable. Attempting to use it today will raise
+    `NotImplementedError` once the execution contract is finalized.
+    """
 
     alias: str = "ipc"
 

agentlightning/execution/shared_memory.py

@@ -7,6 +7,7 @@ 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
 
@@ -17,21 +18,26 @@ logger = logging.getLogger(__name__)
 
 
 class SharedMemoryExecutionStrategy(ExecutionStrategy):
-    """Run algorithm and runners in a single process with threads sharing memory.
+    """Execute bundles in a single process with cooperative worker threads.
 
-    Termination & abort model:
+    Stop Model:
 
-    - One shared ThreadingEvent (`stop_evt`) is passed to *all* bundles.
-    - The main thread (only) receives KeyboardInterrupt on Ctrl+C; we set `stop_evt` there.
-    - If any bundle raises, we set `stop_evt` from that thread to stop the rest.
-    - After the main-thread bundle finishes normally:
-      - If main_thread is "algorithm", we also set `stop_evt` to stop the runners.
-      - If main_thread is "runner", we do not set `stop_evt` to stop the algorithm.
-        We instead wait for the algorithm to finish naturally.
-    - Background threads are daemons; we join briefly and log any stragglers.
+    - All bundles share one [`ThreadingEvent`][agentlightning.ThreadingEvent]
+      named `stop_evt`.
+    - Only the main thread receives `KeyboardInterrupt`. When Ctrl+C occurs we
+      set `stop_evt`.
+    - Any exception raised inside a bundle sets `stop_evt` so other threads can
+      unwind cooperatively.
+    - Once the bundle running on the main thread exits successfully the
+      treatment depends on `main_thread`:
+        - `"algorithm"`: the runners are asked to stop by setting `stop_evt`.
+        - `"runner"`: the algorithm keeps running until it exits naturally.
+    - Background threads are marked as daemons. We join them briefly and log any
+      stragglers before shutting down.
 
-    Notes: Signals other than SIGINT (e.g., SIGTERM) are not intercepted; we respect
-    Python's default behavior for them.
+    !!! note
+        Signals other than `SIGINT` (such as `SIGTERM`) are not intercepted;
+        Python's default behavior for those signals is preserved.
     """
 
     alias: str = "shm"
@@ -43,32 +49,41 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         join_timeout: float = 15.0,
         graceful_delay: float = 5.0,
         poll_interval: float = 0.05,
+        managed_store: bool | None = None,
     ) -> None:
         if main_thread not in ("algorithm", "runner"):
             raise ValueError("main_thread must be 'algorithm' or 'runner'")
         if main_thread == "runner" and n_runners != 1:
-            raise ValueError("When main_thread is 'runner', n_runners must be 1")
+            raise ValueError(
+                "When main_thread is 'runner', n_runners must be 1. "
+                "Either use 'algorithm' on the main thread or set n_runners to 1."
+            )
         self.n_runners = n_runners
         self.main_thread = main_thread
         self.join_timeout = join_timeout
         self.graceful_delay = graceful_delay
         self.poll_interval = poll_interval
+        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.
 
         Control flow:
-          1) Start the bundle coroutine as `task`.
-          2) Start a watcher task that waits for `stop_evt` *without blocking* the loop
-             by periodically polling the threading event.
-          3) When the stop event flips:
-               a) Give the bundle *graceful_delay* seconds to finish on its own,
-                  because well-behaved bundles will check the event and return.
-               b) If still running after the grace period, cancel the bundle task.
-          4) Ensure both tasks are awaited; swallow `CancelledError` where appropriate.
+
+        1. Start the bundle coroutine as `task`.
+        2. Launch a watcher that polls `stop_evt` without blocking the loop.
+        3. When the stop event flips:
+            a. Give the bundle `graceful_delay` seconds to finish on its own,
+               because well-behaved bundles will check the event and return.
+            b. Cancel the bundle task if it is still running after the grace
+               period.
+        4. Await both tasks and swallow `CancelledError` where appropriate.
 
         This is a *backup* mechanism for bundles that might not poll the event
-        frequently; cooperative shutdown (checking `stop_evt` yourself) is still preferred.
+        frequently; cooperative shutdown (checking `stop_evt` inside the
+        bundle) remains the preferred approach.
         """
         task: asyncio.Task[Any] = asyncio.create_task(coro)  # type: ignore
         task_exception: Optional[BaseException] = None
@@ -191,7 +206,10 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
 
         # Create stop event and thread-safe store.
         stop_evt = ThreadingEvent()
-        thread_safe_store = LightningStoreThreaded(store)
+        if self.managed_store:
+            thread_safe_store = LightningStoreThreaded(store)
+        else:
+            thread_safe_store = store
 
         thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
         raised_from_thread: Optional[BaseException] = None

agentlightning/instrumentation/__init__.py

@@ -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

agentlightning/instrumentation/agentops.py

@@ -2,28 +2,78 @@
 
 from __future__ import annotations
 
+import json
 import logging
-import multiprocessing
-import signal
-import socket
-import time
-from typing import Any, Callable
+from typing import Any, Callable, no_type_check
 
-import flask
-import setproctitle
+import requests
+from agentops.client.api import V3Client, V4Client
+from agentops.client.api.types import AuthTokenResponse
+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 agentlightning.utils.otlp import LightningStoreOTLPExporter
 
 logger = logging.getLogger(__name__)
 
 __all__ = [
     "instrument_agentops",
     "uninstrument_agentops",
-    "agentops_local_server",
-    "AgentOpsServerManager",
 ]
 
 # Module-level storage for originals
 _original_handle_chat_attributes: Callable[..., Any] | None = None
 _original_handle_response: Callable[..., Any] | None = None
+_agentops_service_enabled = False
+
+
+def enable_agentops_service(enabled: bool = True) -> None:
+    """
+    Enable or disable communication with the AgentOps service.
+
+    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"AgentOps service enabled is set to {enabled}.")
+
+
+def _patch_exporters():
+    import agentops.client.api
+    import agentops.sdk.core
+
+    agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter  # type: ignore
+    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
+
+
+def _unpatch_exporters():
+    import agentops.client.api
+    import agentops.sdk.core
+
+    agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter  # type: ignore
+    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
+
+
+def _unwrap_legacy_response(response: Any) -> Any:
+    if hasattr(response, "parse") and callable(response.parse):
+        return response.parse()
+    return response
 
 
 def _patch_new_agentops():
@@ -39,41 +89,58 @@ def _patch_new_agentops():
 
     _original_handle_chat_attributes = handle_chat_attributes  # type: ignore
 
+    @no_type_check
     def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
-        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)  # type: ignore
-        if return_value is not None and hasattr(return_value, "prompt_token_ids"):  # type: ignore
-            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)  # type: ignore
-        if return_value is not None and hasattr(return_value, "response_token_ids"):  # type: ignore
-            attributes["response_token_ids"] = list(return_value.response_token_ids[0])  # type: ignore
+        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
+
+        # In some cases, response is a openai._legacy_response.LegacyAPIResponse (e.g., LiteLLM, or LangChain),
+        # This is created by client.with_raw_response.create()
+        return_value = _unwrap_legacy_response(return_value)
+
+        if (
+            return_value is not None
+            and hasattr(return_value, "prompt_token_ids")
+            and return_value.prompt_token_ids is not None
+        ):
+            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
+        if (
+            return_value is not None
+            and hasattr(return_value, "response_token_ids")
+            and return_value.response_token_ids is not None
+        ):
+            attributes["response_token_ids"] = list(return_value.response_token_ids[0])
 
         # For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
-        if (
-            not attributes.get("response_token_ids")
-            and return_value is not None
-            and hasattr(return_value, "choices")  # type: ignore
-            and return_value.choices  # type: ignore
-            and isinstance(return_value.choices, list)  # type: ignore
-        ):
-            first_choice = return_value.choices[0]  # type: ignore
-            if hasattr(first_choice, "token_ids"):  # type: ignore
-                attributes["response_token_ids"] = list(first_choice.token_ids)  # type: ignore
-            # newer versions of OpenAI client SDK
-            elif hasattr(first_choice, "provider_specific_fields") and "token_ids" in first_choice.provider_specific_fields:  # type: ignore
-                attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])  # type: ignore
-
-        # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
         if (
             return_value is not None
-            and hasattr(return_value, "http_response")  # type: ignore
-            and return_value.http_response is not None  # type: ignore
-            and hasattr(return_value.http_response, "json")  # type: ignore
+            and hasattr(return_value, "choices")
+            and return_value.choices
+            and isinstance(return_value.choices, list)
+            and len(return_value.choices) > 0
         ):
-            json_data = return_value.http_response.json()  # type: ignore
-            if isinstance(json_data, dict):
-                if "prompt_token_ids" in json_data:
-                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])  # type: ignore
-                if "response_token_ids" in json_data:
-                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])  # type: ignore
+            first_choice = return_value.choices[0]
+            # Token IDs from "choices[0].token_ids"
+            if "response_token_ids" not in attributes:
+                if hasattr(first_choice, "token_ids") and first_choice.token_ids is not None:
+                    attributes["response_token_ids"] = list(first_choice.token_ids)
+                # newer versions of OpenAI client SDK
+                elif (
+                    hasattr(first_choice, "provider_specific_fields")
+                    and first_choice.provider_specific_fields.get("token_ids") is not None
+                ):
+                    attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])
+
+            # log probability
+            # This is temporary. We need a unified convention for classifying and naming logprobs.
+            if hasattr(first_choice, "logprobs") and first_choice.logprobs is not None:
+                if hasattr(first_choice.logprobs, "content") and first_choice.logprobs.content is not None:
+                    attributes["logprobs.content"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.content]
+                    )
+                if hasattr(first_choice.logprobs, "refusal") and first_choice.logprobs.refusal is not None:
+                    attributes["logprobs.refusal"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.refusal]
+                    )
 
         return attributes
 
@@ -145,6 +212,8 @@ def instrument_agentops():
     Instrument agentops to capture token IDs.
     Automatically detects and uses the appropriate patching method based on the installed agentops version.
     """
+    _patch_exporters()
+
     # Try newest version first (tested for 0.4.16)
     try:
         return _patch_new_agentops()
@@ -164,6 +233,8 @@ def instrument_agentops():
 
 def uninstrument_agentops():
     """Uninstrument agentops to stop capturing token IDs."""
+    _unpatch_exporters()
+
     try:
         _unpatch_new_agentops()
     except Exception:
@@ -174,102 +245,70 @@ def uninstrument_agentops():
         pass
 
 
-def agentops_local_server():
+class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, AuthenticatedOTLPExporter):
     """
-    Returns a Flask app that can be used to test agentops integration.
-    This server provides endpoints for token fetching and a catch-all endpoint.
+    AuthenticatedOTLPExporter with switchable service control.
+
+    When `_agentops_service_enabled` is False, skip export and return success.
     """
-    app = flask.Flask(__name__)
 
-    @app.route("/v3/auth/token", methods=["POST"])
-    def fetch_token():  # type: ignore
-        return {"token": "dummy", "project_id": "dummy"}
-
-    @app.route("/", defaults={"path": ""}, methods=["GET", "POST"])
-    @app.route("/<path:path>", methods=["GET", "POST"])
-    def catch_all(path: str):  # type: ignore
-        return {"path": path}
-
-    return app
+    def should_bypass(self) -> bool:
+        return not _agentops_service_enabled
 
 
-def _run_server(**kwargs: Any):  # type: ignore
+class BypassableOTLPMetricExporter(OTLPMetricExporter):
     """
-    Internal function to run the Flask server.
-    This is used to avoid issues with multiprocessing and Flask's reloader.
+    OTLPMetricExporter with switchable service control.
+    When `_agentops_service_enabled` is False, skip export and return success.
     """
-    signal.signal(signal.SIGINT, signal.SIG_IGN)  # Ignore SIGINT in worker processes
-    setproctitle.setproctitle(multiprocessing.current_process().name)
-    app = agentops_local_server()
-    app.run(**kwargs)
 
-
-class AgentOpsServerManager:
-    """Manages a AgentOps local server to bypass the online service of AgentOps."""
-
-    def __init__(self, daemon: bool = True, port: int | None = None):
-        self.server_process: multiprocessing.Process | None = None
-        self.server_port = port
-        self.daemon = daemon
-        logger.info("AgentOpsServerManager initialized.")
-
-    def _find_available_port(self) -> int:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            return s.getsockname()[1]
-
-    def start(self):
-        if self.server_process and self.server_process.is_alive():
-            logger.warning("AgentOps server process appears to be already running.")
-            return
-
-        if self.server_port is None:
-            self.server_port = self._find_available_port()
-
-        logger.info(f"Starting AgentOps local server on port {self.server_port}...")
-
-        self.server_process = multiprocessing.Process(
-            target=_run_server,
-            kwargs={"host": "127.0.0.1", "port": self.server_port, "use_reloader": False, "debug": False},
-            daemon=self.daemon,
-            name="AgentLightning-AgentOpsServer",
-        )
-        self.server_process.start()
-        logger.info(
-            f"AgentOps local server process (PID: {self.server_process.pid}) started, targeting port {self.server_port}."
-        )
-        time.sleep(0.5)  # Brief wait for server to start up
-        if not self.server_process.is_alive():
-            logger.error(f"AgentOps local server failed to start or exited prematurely.")
-
-    def is_alive(self) -> bool:
-        if self.server_process and self.server_process.is_alive():
-            return True
-        return False
-
-    def stop(self):
-        if self.server_process is not None and self.server_process.is_alive():
-            logger.info(f"Stopping AgentOps local server (PID: {self.server_process.pid})...")
-            self.server_process.terminate()  # Send SIGTERM
-            self.server_process.join(timeout=5)  # Wait for clean exit
-            if self.server_process.is_alive():
-                logger.warning(
-                    f"AgentOps server (PID: {self.server_process.pid}) did not terminate gracefully, killing..."
-                )
-                self.server_process.kill()  # Force kill
-                self.server_process.join(timeout=10)  # Wait for kill
-            self.server_process = None
-            logger.info(f"AgentOps local server stopped.")
+    def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
+        if _agentops_service_enabled:
+            return super().export(*args, **kwargs)  # type: ignore[reportUnknownMemberType]
         else:
-            logger.info("AgentOps local server was not running or already stopped.")
+            logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
+            return MetricExportResult.SUCCESS
 
-    def get_port(self) -> int | None:
-        # Check liveness again in case it died since start()
-        if self.is_alive() and self.server_port is not None:
-            return self.server_port
-        # If called after server stopped or failed, port might be stale or None
-        if self.server_port is not None and (self.server_process is None or not self.server_process.is_alive()):
-            logger.warning(
-                f"AgentOps server port {self.server_port} is stored, but server process is not alive. Returning stored port."
-            )
-        return self.server_port
+
+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 should_bypass(self) -> bool:
+        return not _agentops_service_enabled
+
+
+class BypassableV3Client(V3Client):
+    """
+    V3Client with toggleable authentication calls.
+    Returns dummy auth response when `_agentops_service_enabled` is False.
+    """
+
+    # Temporary synchronous override of fetch_auth_token for mock purposes.
+    def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse:  # type: ignore[override]
+        if _agentops_service_enabled:
+            return super().fetch_auth_token(*args, **kwargs)  # type: ignore[override]
+        else:
+            logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
+            return AuthTokenResponse(token="dummy", project_id="dummy")
+
+
+class BypassableV4Client(V4Client):
+    """
+    V4Client with toggleable post requests.
+    Returns dummy response when `_agentops_service_enabled` is False.
+    """
+
+    def post(self, *args: Any, **kwargs: Any) -> requests.Response:
+        if _agentops_service_enabled:
+            return super().post(*args, **kwargs)
+        else:
+            logger.debug("SwitchableV4Client is switched off, skipping post request.")
+            response = requests.Response()
+            response.status_code = 200
+            response._content = b"{}"
+            return response

agentlightning/instrumentation/litellm.py

@@ -4,7 +4,8 @@
 
 It's unclear whether or not this file is useful.
 It seems that LiteLLM owns its own telemetry from their own entrance
-https://docs.litellm.ai/docs/observability/agentops_integration
+
+[Related documentation](https://docs.litellm.ai/docs/observability/agentops_integration).
 """
 
 from typing import Any, Optional

agentlightning/instrumentation/weave.py

@@ -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.")

agentlightning/litagent/decorator.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""Convenience decorators for building lightweight `LitAgent` implementations."""
+
 from __future__ import annotations
 
 import functools
@@ -90,24 +92,25 @@ class FunctionalLitAgentFunc(Protocol[T_contra]):
 
 
 class FunctionalLitAgent(LitAgent[T]):
-    """A specialized LitAgent that wraps a function-based rollout that accepts
-    dynamically a task input and a configured resource (LLM / prompt template / ...).
+    """Adapter that turns plain rollout functions into [`LitAgent`][agentlightning.LitAgent] instances.
 
-    This class allows users to define agent behavior using a simple function
-    that takes task input and a resource, rather than implementing a full
-    LitAgent subclass.
+    The helper inspects the wrapped function to determine which resources to
+    inject, allowing both synchronous and asynchronous callables to participate
+    in the training loop without writing a dedicated subclass.
     """
 
     def __init__(self, rollout_func: FunctionalLitAgentFunc[T], *, strip_proxy: bool = True) -> None:
-        """
-        Initialize the FunctionalLitAgent with a functional rollout function.
+        """Initialize the wrapper around a rollout function.
 
         Args:
-            rollout_func: A function that defines the agent's behavior.
-                          Can be sync or async, and can optionally accept a Rollout parameter.
-                          The function signature determines which resources are injected (llm, prompt_template, etc.).
-            strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource when the function accepts an llm parameter.
-                         Defaults to True.
+            rollout_func: Callable that implements the rollout. It may be synchronous
+                or asynchronous and can optionally receive a
+                [`Rollout`][agentlightning.Rollout] alongside resources such as
+                `llm` or `prompt_template`.
+            strip_proxy: When ``True``, convert
+                [`ProxyLLM`][agentlightning.ProxyLLM] inputs into
+                [`LLM`][agentlightning.LLM] instances before calling the
+                rollout function. Defaults to `True`.
         """
         super().__init__()
         self._rollout_func = rollout_func
@@ -138,12 +141,15 @@ class FunctionalLitAgent(LitAgent[T]):
         """Execute a synchronous rollout using the wrapped function.
 
         Args:
-            task: The task input data.
-            resources: Dictionary of named resources including LLMs.
-            rollout: The rollout object with metadata.
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.
 
         Returns:
-            The result from the wrapped rollout function.
+            Result produced by the wrapped rollout function.
+
+        Raises:
+            RuntimeError: If the wrapped function is asynchronous.
         """
         if self._is_async:
             raise RuntimeError(f"{self._rollout_func} is asynchronous. Use rollout_async instead.")
@@ -155,12 +161,15 @@ class FunctionalLitAgent(LitAgent[T]):
         """Execute an asynchronous rollout using the wrapped function.
 
         Args:
-            task: The task input data.
-            resources: Dictionary of named resources including LLMs.
-            rollout: The rollout object with metadata.
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.
 
         Returns:
-            The result from the wrapped rollout function.
+            Result produced by the wrapped rollout coroutine.
+
+        Raises:
+            RuntimeError: If the wrapped function is synchronous.
         """
         if not self._is_async:
             raise RuntimeError(f"{self._rollout_func} is synchronous. Use rollout instead.")
@@ -169,18 +178,19 @@ class FunctionalLitAgent(LitAgent[T]):
         return await self._rollout_func(task, **kwargs)  # type: ignore
 
     def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
-        """Extract the kwargs needed for the rollout function based on its signature.
+        """Prepare keyword arguments expected by the wrapped rollout function.
 
-        Dynamically builds the kwargs dictionary by inspecting the function signature and
+
+        It dynamically builds the `kwargs` dictionary by inspecting the function signature and
         including only the parameters the function accepts. This allows flexible function
         signatures that can request any combination of: rollout, llm, and/or prompt_template.
 
         Args:
-            resources: Dictionary of named resources available for the rollout.
-            rollout: The rollout object with metadata.
+            resources: Mapping of named resources available for the rollout.
+            rollout: Rollout metadata provided by the runtime.
 
         Returns:
-            A dictionary of kwargs to pass to the rollout function.
+            Dictionary of keyword arguments to forward to the rollout function.
         """
 
         kwargs: Dict[str, Any] = {}
@@ -194,19 +204,19 @@ class FunctionalLitAgent(LitAgent[T]):
         return kwargs
 
     def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
-        """Extract the first LLM resource from the resources dictionary.
+        """Retrieve the first LLM resource from the available resources.
 
         Strip the ProxyLLM resource into a LLM resource if needed.
 
         Args:
-            resources: Dictionary of named resources.
-            rollout: The rollout object with metadata.
+            resources: Mapping of named resources.
+            rollout: Rollout metadata used when stripping proxy endpoints.
 
         Returns:
-            The first LLM resource found.
+            First [`LLM`][agentlightning.LLM] resource encountered.
 
         Raises:
-            ValueError: If no LLM resource is found.
+            ValueError: If no LLM resource is present.
         """
         resource_found: LLM | None = None
         for name, resource in resources.items():
@@ -225,17 +235,17 @@ class FunctionalLitAgent(LitAgent[T]):
         return resource_found
 
     def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
-        """Extract the first PromptTemplate resource from the resources dictionary.
+        """Retrieve the first prompt template resource from the available resources.
 
         Args:
-            resources: Dictionary of named resources.
-            rollout: The rollout object with metadata. Not used in this method.
+            resources: Mapping of named resources.
+            rollout: Rollout metadata (unused).
 
         Returns:
-            The first PromptTemplate resource found.
+            First [`PromptTemplate`][agentlightning.PromptTemplate] resource encountered.
 
         Raises:
-            ValueError: If no PromptTemplate resource is found.
+            ValueError: If no prompt template resource is present.
         """
         resource_found: PromptTemplate | None = None
         for name, resource in resources.items():
@@ -253,21 +263,22 @@ class FunctionalLitAgent(LitAgent[T]):
         return resource_found
 
     def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
-        """Strip the ProxyLLM resource into a concrete LLM resource.
+        """Convert [`ProxyLLM`][agentlightning.ProxyLLM] instances into concrete LLMs.
 
-        This method resolves ProxyLLM instances to their concrete LLM implementation
+        It resolves ProxyLLM instances to their concrete LLM implementation
         by attaching the attempted rollout context. This is only used when the function
-        signature accepts an 'llm' parameter and strip_proxy is True.
+        signature accepts an `llm` parameter and strip_proxy is True.
 
         Args:
-            proxy_llm: The LLM resource, which may be a ProxyLLM.
-            rollout: The rollout object with metadata.
+            proxy_llm: Candidate LLM resource.
+            rollout: Rollout metadata that provides rollout and attempt identifiers.
 
         Returns:
-            The concrete LLM resource.
+            [`LLM`][agentlightning.LLM] with rollout context baked into the endpoint.
 
         Raises:
-            ValueError: If the rollout is not an AttemptedRollout (required for stripping ProxyLLM).
+            ValueError: If the rollout is not an
+                [`AttemptedRollout`][agentlightning.AttemptedRollout].
         """
 
         if not isinstance(proxy_llm, ProxyLLM):
@@ -293,41 +304,37 @@ def llm_rollout(*, strip_proxy: bool = True) -> Callable[[LlmRolloutFunc[T]], Fu
 def llm_rollout(
     func: LlmRolloutFunc[T] | None = None, *, strip_proxy: bool = True
 ) -> FunctionalLitAgent[T] | Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]:
-    """Create a FunctionalLitAgent from a function that takes (task, llm[, rollout]).
-
-    This decorator allows you to define an agent using a simple function
-    instead of creating a full LitAgent subclass. The returned FunctionalLitAgent
-    instance is callable, preserving the original function's behavior.
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for LLM-based rollouts.
 
     Args:
-        func: A function that defines the agent's behavior. Can be:
-              - sync: (task, llm) -> result
-              - sync with rollout: (task, llm, rollout) -> result
-              - async: async (task, llm) -> result
-              - async with rollout: async (task, llm, rollout) -> result
-        strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource.
-                     Defaults to True.
+        func: Callable defining the agent's behaviour. Supported signatures include:
+
+            * `(task, llm) -> result`
+            * `(task, llm, rollout) -> result`
+            * `async (task, llm) -> result`
+            * `async (task, llm, rollout) -> result`
+
+        strip_proxy: When `True`, convert proxy resources into concrete
+            [`LLM`][agentlightning.LLM] instances before calling the
+            function. Defaults to `True`.
 
     Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
 
-    Example:
+    Examples:
+        ```python
         @llm_rollout
         def my_agent(task, llm):
-            # Agent logic here
-            return response
+            return llm.endpoint
 
         @llm_rollout(strip_proxy=False)
         def my_agent_no_strip(task, llm):
-            # Agent logic here
-            return response
+            return llm.model
 
-        # Function is still callable with original behavior
         result = my_agent(task, llm)
-
-        # Agent methods are also available
         result = my_agent.rollout(task, resources, rollout)
+        ```
     """
 
     def decorator(f: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]:
@@ -343,19 +350,20 @@ def llm_rollout(
 
 
 def _validate_llm_rollout_func(func: Any) -> TypeGuard[LlmRolloutFunc[Any]]:
-    """Validate the function signature of a LLM rollout function.
+    """Validate the function signature of an LLM rollout function.
 
     Ensures the function follows the expected pattern for LLM-based rollouts:
+
     - Must have at least 2 parameters
     - First parameter must be named 'task'
     - Must have a parameter named 'llm'
     - Optionally can have a 'rollout' parameter
 
     Args:
-        func: The function to validate.
+        func: Function to inspect.
 
     Returns:
-        True if the function signature is valid.
+        `True` when the signature matches the supported patterns.
 
     Raises:
         ValueError: If the function signature does not match the expected pattern.
@@ -383,36 +391,34 @@ def prompt_rollout() -> Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
 def prompt_rollout(
     func: PromptRolloutFunc[T] | None = None,
 ) -> FunctionalLitAgent[T] | Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
-    """Create a FunctionalLitAgent from a function that takes (task, prompt_template[, rollout]).
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for prompt-based rollouts.
 
     This decorator is designed for agents that work with tunable prompt templates. It enables
     a workflow where algorithms manage and optimize the prompt template, while agents consume
     the template to perform rollouts. This is particularly useful for prompt optimization scenarios.
 
     Args:
-        func: A function that defines the agent's behavior. Can be:
-              - sync: (task, prompt_template) -> result
-              - sync with rollout: (task, prompt_template, rollout) -> result
-              - async: async (task, prompt_template) -> result
-              - async with rollout: async (task, prompt_template, rollout) -> result
+        func: Callable defining the agent's behavior. Supported signatures include:
+
+            * `(task, prompt_template) -> result`
+            * `(task, prompt_template, rollout) -> result`
+            * `async (task, prompt_template) -> result`
+            * `async (task, prompt_template, rollout) -> result`
 
     Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
 
-    Example:
+    Examples:
+        ```python
         @prompt_rollout
         def my_agent(task, prompt_template):
-            # Use the prompt template to generate a response
             messages = prompt_template.format(task=task.input)
-            # ... perform rollout with the formatted prompt
-            return response
+            return messages
 
-        # Function is still callable with original behavior
         result = my_agent(task, prompt_template)
-
-        # Agent methods are also available
         result = my_agent.rollout(task, resources, rollout)
+        ```
     """
 
     def decorator(f: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]:
@@ -429,16 +435,17 @@ def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]
     """Validate the function signature of a prompt rollout function.
 
     Ensures the function follows the expected pattern for prompt-template-based rollouts:
+
     - Must have at least 2 parameters
     - First parameter must be named 'task'
     - Must have a parameter named 'prompt_template'
     - Optionally can have a 'rollout' parameter
 
     Args:
-        func: The function to validate.
+        func: Function to inspect.
 
     Returns:
-        True if the function signature is valid.
+        `True` when the signature matches the supported patterns.
 
     Raises:
         ValueError: If the function signature does not match the expected pattern.
@@ -456,23 +463,30 @@ def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]
 
 
 def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., Any]]) -> FunctionalLitAgent[T]:
-    """Create a LitAgent from a function, automatically detecting the appropriate type.
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] from an arbitrary rollout function.
 
     This function inspects the provided callable and creates the appropriate
     agent type based on its signature. It supports both LLM-based and prompt-template-based
     agents. The returned agent instance is callable, preserving the original function's
     behavior and type hints.
 
+    See [`llm_rollout`][agentlightning.litagent.decorator.llm_rollout] and
+    [`prompt_rollout`][agentlightning.litagent.decorator.prompt_rollout] for more details.
+
     Args:
-        func: A function that defines the agent's behavior. Supported signatures:
-              - (task, llm[, rollout]) for LLM-based agents
-              - (task, prompt_template[, rollout]) for prompt-template-based agents
+        func: Callable that implements the rollout. Supported signatures:
+
+            - `[async ](task, llm[, rollout])` for LLM-based agents
+            - `[async ](task, prompt_template[, rollout])` for prompt-template-based agents
+
+            The supported output types of `func` is same as the return type of [`rollout`][agentlightning.LitAgent.rollout].
 
     Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
 
-    Example:
+    Examples:
+        ```python
         # LLM-based agent
         @rollout
         def my_llm_agent(task, llm):
@@ -495,6 +509,7 @@ def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., A
 
         # Agent methods are also available
         result = my_llm_agent.rollout(task, resources, rollout)
+        ```
 
     Raises:
         NotImplementedError: If the function signature doesn't match any known patterns.

agentlightning/litagent/litagent.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""Base abstractions for building agents that plug into Agent Lightning."""
+
 from __future__ import annotations
 
 import inspect
@@ -11,8 +13,8 @@ from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
 from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task
 
 if TYPE_CHECKING:
-    from agentlightning.runner import BaseRunner
-    from agentlightning.tracer import BaseTracer
+    from agentlightning.runner import Runner
+    from agentlightning.tracer import Tracer
     from agentlightning.trainer import Trainer
 
 
@@ -26,32 +28,38 @@ __all__ = [
 
 
 def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
-    """Check if the rollout API is v0.1.
-    Inspect the function signature to see if it has a rollout_id parameter.
+    """Return `True` when the rollout function uses the deprecated v0.1 signature.
+
+    The helper inspects the callable's signature to detect whether a `rollout_id`
+    parameter is present, which indicates the legacy API.
 
     Args:
-        func: The function to check.
+        func: Function to analyze.
+
+    Returns:
+        `True` if the callable exposes a `rollout_id` parameter.
     """
     return "rollout_id" in inspect.signature(func).parameters
 
 
 class LitAgent(Generic[T]):
-    """Base class for the training and validation logic of an agent.
+    """Base class for implementing agent rollouts.
 
-    Developers should subclass this class and implement the rollout methods
-    to define the agent's behavior for a single task. The agent's logic
-    is completely decoupled from the server communication and training
-    infrastructure.
+    Subclasses override the rollout methods to process tasks while the trainer and
+    runner infrastructure manages orchestration, tracing, and persistence.
     """
 
     def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
-        """
-        Initialize the LitAgent.
+        """Initialize the agent instance.
 
         Args:
-            trained_agents: Optional string representing the trained agents.
-                            This can be used to track which agents have been trained by this instance.
-                            Deprecated. Configure `agent_match` in adapter instead.
+            trained_agents: Optional identifier used by legacy tooling to mark trained
+                agents.
+
+        !!! warning "Deprecated"
+            The `trained_agents` flag is deprecated. Configure `agent_match` in the adapter
+            layer instead. See [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet]
+            for more details.
         """
         if trained_agents is not None:
             warnings.warn(
@@ -62,15 +70,12 @@ class LitAgent(Generic[T]):
         self.trained_agents = trained_agents
 
         self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
-        self._runner_ref: weakref.ReferenceType[BaseRunner[T]] | None = None
+        self._runner_ref: weakref.ReferenceType[Runner[T]] | None = None
 
     def is_async(self) -> bool:
-        """
-        Check if the agent implements asynchronous rollout methods.
-        Override this property for customized async detection logic.
+        """Return `True` when the agent overrides any asynchronous rollout methods.
 
-        Returns:
-            True if the agent has custom async rollout methods, False otherwise.
+        Override this method for customized async detection logic.
         """
         return (
             (
@@ -85,21 +90,15 @@ class LitAgent(Generic[T]):
         )
 
     def set_trainer(self, trainer: Trainer) -> None:
-        """
-        Set the trainer for this agent.
+        """Attach the trainer responsible for orchestration.
 
         Args:
-            trainer: The Trainer instance that will handle training and validation.
+            trainer: [`Trainer`][agentlightning.Trainer] that manages the agent.
         """
         self._trainer_ref = weakref.ref(trainer)
 
     def get_trainer(self) -> Trainer:
-        """
-        Get the trainer for this agent.
-
-        Returns:
-            The Trainer instance associated with this agent.
-        """
+        """Return the trainer associated with this agent."""
         if self._trainer_ref is None:
             raise ValueError("Trainer has not been set for this agent.")
         trainer = self._trainer_ref()
@@ -109,39 +108,31 @@ class LitAgent(Generic[T]):
 
     @property
     def trainer(self) -> Trainer:
-        """Convenient shortcut of self.get_trainer()."""
+        """Return the trainer associated with this agent."""
         return self.get_trainer()
 
-    def get_tracer(self) -> BaseTracer:
-        """
-        Get the tracer for this agent.
-
-        Returns:
-            The BaseTracer instance associated with this agent.
-        """
-        return self.trainer.tracer
+    def get_tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
+        if hasattr(self.runner, "tracer"):
+            return self.runner.tracer  # type: ignore
+        else:
+            return self.trainer.tracer
 
     @property
-    def tracer(self) -> BaseTracer:
-        """Convenient shortcut of self.get_tracer()."""
+    def tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
         return self.get_tracer()
 
-    def set_runner(self, runner: BaseRunner[T]) -> None:
-        """
-        Set the runner for this agent.
+    def set_runner(self, runner: Runner[T]) -> None:
+        """Attach the runner responsible for executing rollouts.
 
         Args:
-            runner: The runner instance that will handle the execution of rollouts.
+            runner: [`Runner`][agentlightning.Runner] coordinating execution.
         """
         self._runner_ref = weakref.ref(runner)
 
-    def get_runner(self) -> BaseRunner[T]:
-        """
-        Get the runner for this agent.
-
-        Returns:
-            The runner instance associated with this agent.
-        """
+    def get_runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
         if self._runner_ref is None:
             raise ValueError("Runner has not been set for this agent.")
         runner = self._runner_ref()
@@ -150,159 +141,112 @@ class LitAgent(Generic[T]):
         return runner
 
     @property
-    def runner(self) -> BaseRunner[T]:
-        """Convenient shortcut of self.get_runner()."""
+    def runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
         return self.get_runner()
 
-    def on_rollout_start(self, task: Task, runner: BaseRunner[T], tracer: BaseTracer) -> None:
-        """Hook called immediately before a rollout begins.
+    def on_rollout_start(self, task: Task, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked immediately before a rollout begins.
 
-        Deprecated in favor of `on_rollout_start` in the `Hook` interface.
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. The default implementation is a no-op.
 
         Args:
-            task: The :class:`Task` object that will be processed.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
+            task: [`Task`][agentlightning.Task] that will be processed.
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
 
-        Subclasses can override this method to implement custom logic such as
-        logging, metric collection, or resource setup. By default, this is a
-        no-op.
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_start`][agentlightning.Hook.on_rollout_start]
+            instead of this method when extending agents.
         """
 
-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: BaseRunner[T], tracer: BaseTracer) -> None:
-        """Hook called after a rollout completes.
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked after a rollout completes.
 
-        Deprecated in favor of `on_rollout_end` in the `Hook` interface.
+        Subclasses can override this method for cleanup or additional logging. The default
+        implementation is a no-op.
 
         Args:
-            task: The :class:`Task` object that was processed.
-            rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
+            task: [`Task`][agentlightning.Task] that was processed.
+            rollout: Resulting [`Rollout`][agentlightning.Rollout].
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
 
-        Subclasses can override this method for cleanup or additional
-        logging. By default, this is a no-op.
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_end`][agentlightning.Hook.on_rollout_end]
+            instead of this method when extending agents.
         """
 
     def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Main entry point for executing a rollout.
+        """Execute a rollout synchronously.
 
-        This method determines whether to call the synchronous or
-        asynchronous rollout method based on the agent's implementation.
 
         If you don't wish to implement both training rollout and validation
         rollout separately, you can just implement `rollout` which will work for both.
 
         Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
-                     Most agents should only use `task` and `resources`. Use `rollout`
-                     only if you need to access metadata like `rollout_id`.
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.
 
         Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
+            One of the following values:
+
+            * `None` when tracing is handled by the runner.
+            * `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.")
 
     async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of the main rollout method.
-
-        This method determines whether to call the synchronous or
-        asynchronous rollout method based on the agent's implementation.
+        """Execute a rollout asynchronously.
 
         Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
-                     Most agents should only use `task` and `resources`. Use `rollout`
-                     only if you need to access metadata like `rollout_id`.
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.
 
         Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
+            Same possible return values as
+            [`rollout`][agentlightning.LitAgent.rollout].
         """
         raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
 
     def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Defines the agent's behavior for a single training task.
+        """Process a single training task synchronously.
 
-        This method should contain the logic for how the agent processes an
-        input, uses the provided resources (like LLMs or prompts), and
-        produces a result.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
+        By default, this method delegates to
+        [`rollout`][agentlightning.LitAgent.rollout].
         """
         return self.rollout(task, resources, rollout)
 
     def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Defines the agent's behavior for a single validation task.
+        """Process a single validation task synchronously.
 
-        By default, this method redirects to `training_rollout`. Override it
-        if the agent should behave differently during validation.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the validation rollout. See `rollout` for
-            possible return types.
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout`][agentlightning.LitAgent.training_rollout].
         """
         return self.rollout(task, resources, rollout)
 
     async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of `training_rollout`.
+        """Process a single training task asynchronously.
 
-        This method should be implemented by agents that perform asynchronous
-        operations (e.g., non-blocking I/O, concurrent API calls).
-
-        Args:
-            task: The task object received from the server.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the asynchronous training rollout. See `rollout` for
-            possible return types.
+        By default, this method delegates to
+        [`rollout_async`][agentlightning.LitAgent.rollout_async].
         """
         return await self.rollout_async(task, resources, rollout)
 
     async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of `validation_rollout`.
+        """Process a single validation task asynchronously.
 
-        By default, this method redirects to `training_rollout_async`.
-        Override it for different asynchronous validation behavior.
-
-        Args:
-            task: The task object received from the server.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the asynchronous validation rollout. See `rollout` for
-            possible return types.
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout_async`][agentlightning.LitAgent.training_rollout_async].
         """
         return await self.rollout_async(task, resources, rollout)

agentlightning/logging.py

@@ -1,20 +1,370 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import logging
+from __future__ import annotations
 
-__all__ = ["configure_logger"]
+import logging
+import os
+import platform
+import sys
+import warnings
+from logging.config import dictConfig
+from typing import Any, Dict, Optional
+
+from rich.console import Console
+
+__all__ = ["setup", "configure_logger", "setup_module"]
 
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
-    logger = logging.getLogger(name)
-    logger.handlers.clear()  # clear existing handlers
+    """Create or reset a namespaced logger with a consistent console format.
 
-    # 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
+    This helper clears any previously attached handlers before binding a single
+    `StreamHandler` that writes to standard output. The resulting logger does
+    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`.
+        name: Dotted path for the logger instance. Defaults to
+            `"agentlightning"`.
+
+    Returns:
+        Configured logger instance ready for immediate use.
+
+    Examples:
+        ```python
+        from agentlightning import configure_logger
+
+        logger = configure_logger(level=logging.INFO)
+        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 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"
+
+    base_logger = setup_module(
+        level,
+        name="agentlightning",
+        console=console,
+        color=color,
+        propagate=propagate,
+        disable_existing_loggers=disable_existing_loggers,
+    )
+
+    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()

agentlightning/runner/__init__.py

@@ -1,11 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from .agent import LitAgentRunner
-from .base import BaseRunner
+from .base import Runner
 from .legacy import LegacyAgentRunner
 
 __all__ = [
-    "BaseRunner",
+    "Runner",
     "LegacyAgentRunner",
     "LitAgentRunner",
 ]

agentlightning/runner/agent.py

@@ -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 BaseTracer
+from agentlightning.tracer.base import Tracer
+from agentlightning.tracer.otel import OtelTracer
 from agentlightning.types import (
     AttemptedRollout,
     Hook,
@@ -29,43 +43,67 @@ 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
 
-from .base import BaseRunner
+from .base import Runner
 
 T_task = TypeVar("T_task")
 
 logger = logging.getLogger(__name__)
 
 
-class LitAgentRunner(BaseRunner[T_task]):
-    """Runner implementation for executing agent tasks with distributed support.
+class LitAgentRunner(Runner[T_task]):
+    """Execute [`LitAgent`][agentlightning.LitAgent] tasks with tracing support.
 
     This runner manages the complete lifecycle of agent rollout execution,
     including task polling, resource management, tracing, and hooks. It supports
     both continuous iteration over tasks from the store and single-step execution.
 
     Attributes:
-        worker_id: The unique identifier for this worker process.
+        worker_id: Identifier for the active worker process, if any.
     """
 
-    def __init__(self, tracer: BaseTracer, 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:
-            tracer: The tracer instance for recording execution traces and spans.
-            max_rollouts: Maximum number of tasks to process in iter() mode. If None,
-                the runner will continue indefinitely until interrupted.
-            poll_interval: Time in seconds to wait between polling attempts when
-                no tasks are available in the store.
+            tracer: [`Tracer`][agentlightning.Tracer] used for rollout spans.
+            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
@@ -80,10 +118,9 @@ class LitAgentRunner(BaseRunner[T_task]):
         initializes the tracer.
 
         Args:
-            agent: The LitAgent instance to be managed by this runner.
-            hooks: Optional sequence of Hook objects to be called at various
-                lifecycle stages (on_trace_start, on_trace_end, on_rollout_start,
-                on_rollout_end).
+            agent: [`LitAgent`][agentlightning.LitAgent] instance executed by the runner.
+            hooks: Optional sequence of [`Hook`][agentlightning.Hook]
+                callbacks invoked around tracing and rollout boundaries.
             **kwargs: Additional initialization arguments (currently unused).
         """
         self._agent = agent
@@ -100,13 +137,14 @@ class LitAgentRunner(BaseRunner[T_task]):
 
         Args:
             worker_id: Unique identifier for this worker process.
-            store: The LightningStore instance for task coordination and data persistence.
+            store: [`LightningStore`][agentlightning.LightningStore]
+                used for task coordination and persistence.
             **kwargs: Additional worker-specific initialization arguments (currently unused).
         """
         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.
@@ -131,7 +169,7 @@ class LitAgentRunner(BaseRunner[T_task]):
         This method cleans up worker-specific resources and resets the worker ID.
 
         Args:
-            worker_id: The unique identifier of the worker being torn down.
+            worker_id: Unique identifier of the worker being torn down.
             *args: Additional teardown arguments (currently unused).
             **kwargs: Additional teardown keyword arguments (currently unused).
         """
@@ -139,6 +177,15 @@ class LitAgentRunner(BaseRunner[T_task]):
 
         self._tracer.teardown_worker(worker_id)
 
+    @property
+    def tracer(self) -> Tracer:
+        """Get the tracer instance.
+
+        Returns:
+            The Tracer instance used by this runner.
+        """
+        return self._tracer
+
     def get_agent(self) -> LitAgent[T_task]:
         """Get the agent instance.
 
@@ -146,7 +193,7 @@ class LitAgentRunner(BaseRunner[T_task]):
             The LitAgent instance managed by this runner.
 
         Raises:
-            ValueError: If the agent has not been initialized via init().
+            ValueError: If the agent has not been initialized via [`init`][agentlightning.LitAgentRunner.init].
         """
         if self._agent is None:
             raise ValueError("Agent not initialized. Call init() first.")
@@ -159,7 +206,7 @@ class LitAgentRunner(BaseRunner[T_task]):
             The LightningStore instance for this worker.
 
         Raises:
-            ValueError: If the store has not been initialized via init_worker().
+            ValueError: If the store has not been initialized via [`init_worker`][agentlightning.LitAgentRunner.init_worker].
         """
         if self._store is None:
             raise ValueError("Store not initialized. Call init_worker() first.")
@@ -235,49 +282,84 @@ class LitAgentRunner(BaseRunner[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:
@@ -285,7 +367,8 @@ class LitAgentRunner(BaseRunner[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]
@@ -294,8 +377,225 @@ class LitAgentRunner(BaseRunner[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.
 
@@ -303,14 +603,16 @@ class LitAgentRunner(BaseRunner[T_task]):
         and return early if the event is set.
 
         Args:
-            event: Optional ExecutionEvent object that can be used to interrupt the sleep.
+            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():
@@ -348,6 +650,8 @@ class LitAgentRunner(BaseRunner[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
 
@@ -355,9 +659,11 @@ class LitAgentRunner(BaseRunner[T_task]):
             await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
 
             start_time = time.time()
-            with self._tracer.trace_context(
-                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+            logger.debug(f"{self._log_prefix(rollout_id)} Prepared for trace context.")
+            async with self._tracer.trace_context(
+                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
                 )
@@ -369,21 +675,27 @@ class LitAgentRunner(BaseRunner[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)
@@ -426,6 +738,7 @@ class LitAgentRunner(BaseRunner[T_task]):
         """Run the runner, continuously iterating over tasks in the store.
 
         This method polls the store for new rollouts and executes them until:
+
         - The event is set (if provided)
         - The max_rollouts limit is reached (if configured)
         - No more tasks are available
@@ -441,39 +754,40 @@ class LitAgentRunner(BaseRunner[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.")
 
@@ -488,7 +802,8 @@ class LitAgentRunner(BaseRunner[T_task]):
         """Execute a single task directly, bypassing the task queue.
 
         This method creates a new rollout for the given input and executes it
-        immediately. Unlike iter(), exceptions are propagated to the caller.
+        immediately. Unlike [`iter()`][agentlightning.LitAgentRunner.iter],
+        exceptions are propagated to the caller.
 
         Args:
             input: The task input to be processed by the agent.
@@ -515,7 +830,9 @@ class LitAgentRunner(BaseRunner[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)

agentlightning/runner/base.py

@@ -1,11 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Base runner interface for executing agent tasks.
-
-This module defines the abstract base class for all runner implementations
-in the agent-lightning framework. Runners are responsible for managing the
-execution lifecycle of agents and coordinating with the store.
-"""
+"""Abstract runner interface for executing agent tasks."""
 
 from __future__ import annotations
 
@@ -27,90 +22,75 @@ T_task = TypeVar("T_task")
 logger = logging.getLogger(__name__)
 
 
-class BaseRunner(ParallelWorkerBase, Generic[T_task]):
-    """Base class for all runners.
+class Runner(ParallelWorkerBase, Generic[T_task]):
+    """Abstract base class for long-running agent executors.
 
-    This abstract base class defines the interface that all runner implementations
-    must follow. Runners are responsible for executing agent tasks, managing the
-    execution lifecycle, and coordinating with the store.
+    Runner implementations coordinate [`LitAgent`][agentlightning.LitAgent]
+    instances, acquire work from a [`LightningStore`][agentlightning.LightningStore],
+    and emit [`Rollout`][agentlightning.Rollout] objects. Subclasses decide how
+    to schedule work (polling, streaming, etc.) while this base class provides a
+    minimal lifecycle contract.
     """
 
     def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
-        """Initialize the runner with the agent.
+        """Prepare the runner to execute tasks for `agent`.
 
-        This method is called once during setup to configure the runner with
-        the agent it will execute.
+        This method is called only once during the setup for all workers, not for each worker.
 
         Args:
-            agent: The LitAgent instance to be managed by this runner.
-            **kwargs: Additional initialization arguments specific to the runner implementation.
+            agent: Agent instance providing task-specific logic.
+            **kwargs: Optional runner-specific configuration.
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must supply the initialization
+                routine.
         """
         raise NotImplementedError()
 
     def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
-        """Initialize the runner for each worker with worker_id and store.
+        """Configure worker-local state before processing tasks.
 
-        This method is called once per worker process in a distributed setup.
-        It provides the worker with its unique ID and the store instance for
-        task coordination.
+        This method is called for **each** worker during the setup.
 
         Args:
-            worker_id: Unique identifier for this worker process.
-            store: The LightningStore instance for task coordination and data persistence.
-            **kwargs: Additional worker-specific initialization arguments.
+            worker_id: Unique identifier for this worker process or thread.
+            store: Shared [`LightningStore`][agentlightning.LightningStore]
+                backing task coordination.
+            **kwargs: Optional worker-specific configuration.
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must prepare per-worker resources.
         """
         raise NotImplementedError()
 
     def run(self, *args: Any, **kwargs: Any) -> None:
-        """Undefined method - use iter() or step() instead.
+        """Deprecated synchronous entry point.
 
-        This method is intentionally not implemented as the execution behavior
-        should be defined through iter() for continuous execution or step()
-        for single-task execution.
-
-        Args:
-            *args: Unused positional arguments.
-            **kwargs: Unused keyword arguments.
+        Use [`iter()`][agentlightning.Runner.iter] or [`step()`][agentlightning.Runner.step] instead.
 
         Raises:
-            RuntimeError: Always raised to indicate this method should not be used.
+            RuntimeError: Always raised to direct callers to
+                [iter()][agentlightning.Runner.iter] or
+                [step()][agentlightning.Runner.step].
         """
         raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
 
     def teardown(self, *args: Any, **kwargs: Any) -> None:
-        """Clean up runner resources and reset state.
-
-        This method is called once during shutdown to clean up any resources
-        allocated during initialization and reset the runner state.
-
-        Args:
-            *args: Additional teardown arguments.
-            **kwargs: Additional teardown keyword arguments.
+        """Release resources acquired during [`init()`][agentlightning.Runner.init].
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must implement the shutdown routine.
         """
         raise NotImplementedError()
 
     def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
-        """Clean up worker-specific resources.
-
-        This method is called once per worker during shutdown to clean up
-        any resources specific to that worker.
+        """Release per-worker resources allocated by [`init_worker()`][agentlightning.Runner.init_worker].
 
         Args:
-            worker_id: The unique identifier of the worker being torn down.
-            *args: Additional teardown arguments.
-            **kwargs: Additional teardown keyword arguments.
+            worker_id: Identifier of the worker being torn down.
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must implement the shutdown routine.
         """
         raise NotImplementedError()
 
@@ -122,18 +102,21 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
         store: LightningStore,
         hooks: Optional[Sequence[Hook]] = None,
         worker_id: Optional[int] = None,
-    ) -> Iterator[BaseRunner[T_task]]:
-        """Context manager for quickly init and teardown the runner,
-        so that you can debug the runner without a trainer environment.
+    ) -> Iterator[Runner[T_task]]:
+        """Initialize and tear down a runner within a simple context manager.
+
+        The helper is primarily intended for debugging runner implementations
+        outside of a full [`Trainer`][agentlightning.Trainer] stack.
 
         Args:
-            agent: The LitAgent instance to be managed by this runner.
-                   It should be the same agent that is to be run within the context.
-            store: The LightningStore instance for task coordination and data persistence.
-                   If you don't have one, you can easily create one with `InMemoryLightningStore()`.
-            hooks: Optional sequence of Hook instances to be used by the runner.
-                   Only some runners support hooks.
-            worker_id: Optional worker ID to be used by the runner.
+            agent: Agent executed by this runner.
+            store: Backing [`LightningStore`][agentlightning.LightningStore].
+                If you don't have one, you can easily create one with
+                [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+            hooks: Optional sequence of hooks recognised by the runner.
+                Not all runners support hooks.
+            worker_id: Override the worker identifier used during setup. Defaults
+                to `0`.
         """
         _initialized: bool = False
         _worker_initialized: bool = False
@@ -163,12 +146,11 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
         them until interrupted by the event or when no more tasks are available.
 
         Args:
-            event: Optional ExecutionEvent object that can be used to signal the runner
-                to stop gracefully. When set, the runner should finish its current
-                task and exit the iteration loop.
+            event: Cooperative stop signal. When set, the runner should complete
+                the current unit of work and exit the loop.
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses provide the iteration behavior.
         """
         raise NotImplementedError()
 
@@ -186,18 +168,15 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
         directly, bypassing the store's task queue.
 
         Args:
-            input: The task input to be processed by the agent.
-            resources: Optional named resources to be used for this specific task.
-                If not provided, the latest resources from the store will be used.
-            mode: Optional rollout mode (e.g., "train", "test"). If not provided,
-                the default mode will be used.
-            event: Optional ExecutionEvent object to signal interruption. When set, the
-                runner may abort the current execution.
+            input: Task payload consumed by the agent.
+            resources: Optional named resources scoped to this invocation.
+            mode: Optional rollout mode such as `"train"` or `"eval"`.
+            event: Cooperative stop signal for long-running tasks.
 
         Returns:
-            The completed rollout.
+            Completed rollout produced by the agent.
 
         Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses provide the execution behavior.
         """
         raise NotImplementedError()

agentlightning/runner/legacy.py

@@ -11,10 +11,10 @@ from agentlightning.adapter import TracerTraceToTriplet
 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 BaseTracer
-from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
+from agentlightning.tracer.base import Tracer
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet
 
-from .base import BaseRunner
+from .base import Runner
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +23,7 @@ __all__ = [
 ]
 
 
-class LegacyAgentRunner(BaseRunner[Any]):
+class LegacyAgentRunner(Runner[Any]):
     """Manages the agent's execution loop and integrates with AgentOps.
 
     This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -43,7 +43,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
         self,
         agent: LitAgent[Any],
         client: AgentLightningClient,
-        tracer: BaseTracer,
+        tracer: Tracer,
         triplet_exporter: TracerTraceToTriplet,
         worker_id: Optional[int] = None,
         max_tasks: Optional[int] = None,
@@ -58,7 +58,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
         self.worker_id = worker_id
         self.max_tasks = max_tasks
 
-    # These methods are overridden by BaseRunner, getting them back to old behavior.
+    # These methods are overridden by Runner, getting them back to old behavior.
     def init(self, *args: Any, **kwargs: Any) -> None:
         pass
 
@@ -99,7 +99,7 @@ class LegacyAgentRunner(BaseRunner[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(BaseRunner[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(BaseRunner[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:
@@ -180,7 +183,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
 
-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"):  # pyright: ignore[reportPrivateUsage]
                 start_time = time.time()
                 rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
                 # Pass the task input, not the whole task object
@@ -257,7 +260,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
 
-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            async with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                 start_time = time.time()
                 rollout_method = (
                     self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async

agentlightning/semconv.py

@@ -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."""

agentlightning/server.py

@@ -1,6 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Legacy server for the Agent Lightning framework. Deprecated in favor of agentlightning.store."""
+"""Legacy HTTP server compatible with the original Agent Lightning protocol.
+
+The implementation in this module predates the modern store-powered runtime and
+is kept for backwards compatibility with older deployments. New applications
+should migrate to the store architecture where possible.
+"""
 
 from __future__ import annotations
 
@@ -29,9 +34,15 @@ logger = logging.getLogger(__name__)
 
 
 class ServerDataStore:
-    """
-    A centralized, thread-safe, async, in-memory data store for the server's state.
-    This holds the task queue, versioned resources, and completed rollouts.
+    """Async-safe container for in-memory server state.
+
+    The store tracks queued tasks, claimed tasks, uploaded rollouts, and the
+    currently published resources. All interactions are guarded by asyncio locks
+    so that the FastAPI handlers can safely run in parallel.
+
+    !!! warning "Deprecated"
+        [`ServerDataStore`][agentlightning.server.ServerDataStore] is part of
+        the legacy client/server stack. Use [`LightningStore`][agentlightning.LightningStore] instead.
     """
 
     def __init__(self):
@@ -54,8 +65,18 @@ class ServerDataStore:
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> str:
-        """
-        Adds a new task to the queue with specific metadata and returns its unique ID.
+        """Enqueue a new task and return the generated rollout identifier.
+
+        Args:
+            sample: Payload that describes the task input.
+            mode: Phase in which the sample should be executed (`"train"`, `"val"`, or
+                `"test"`).
+            resources_id: Identifier of a resource bundle that the executor should
+                load before running the task.
+            metadata: Optional metadata forwarded to the executor.
+
+        Returns:
+            Unique rollout identifier assigned to the task.
         """
         rollout_id = f"rollout-{uuid.uuid4()}"
         task = Task(
@@ -72,9 +93,11 @@ class ServerDataStore:
         return rollout_id
 
     async def get_next_task(self) -> Optional[Task]:
-        """
-        Retrieves the next task from the queue without blocking.
-        Returns None if the queue is empty.
+        """Retrieve the next task from the queue without blocking.
+
+        Returns:
+            Next [`Task`][agentlightning.Task] ready to execute, or ``None``
+            when the queue is empty.
         """
         try:
             async with self._results_lock:
@@ -95,8 +118,10 @@ class ServerDataStore:
             return None
 
     async def update_resources(self, update: ResourcesUpdate):
-        """
-        Safely stores a new version of named resources and sets it as the latest.
+        """Persist a new resource bundle and mark it as the latest version.
+
+        Args:
+            update: Resource payload received from a client.
         """
         # TODO: evict old resources if necessary.
         async with self._resources_lock:
@@ -105,26 +130,38 @@ class ServerDataStore:
             logger.info(f"Resources updated. New version '{update.resources_id}' is now latest.")
 
     async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves a specific version of named resources by its ID.
+        """Retrieve a specific resource bundle by identifier.
+
+        Args:
+            resources_id: Identifier that was previously published to the store.
+
+        Returns:
+            Matching [`ResourcesUpdate`][agentlightning.ResourcesUpdate]
+            instance, or ``None`` when the identifier is unknown.
         """
         async with self._resources_lock:
             resources = self._resource_versions.get(resources_id)
             if resources:
-                return ResourcesUpdate(resources_id=resources_id, resources=resources)
+                return ResourcesUpdate(
+                    resources_id=resources_id,
+                    resources=resources,
+                    create_time=time.time(),
+                    update_time=time.time(),
+                    version=1,
+                )
             return None
 
     async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves the latest version of named resources.
-        """
+        """Return the most recent resource bundle, if one exists."""
         if self._latest_resources_id:
             return await self.get_resources_by_id(self._latest_resources_id)
         return None
 
     async def store_rollout(self, rollout: RolloutLegacy):
-        """
-        Safely stores a completed rollout from a client.
+        """Persist a completed rollout for later inspection.
+
+        Args:
+            rollout: Rollout returned by a client.
         """
         async with self._results_lock:
             self._processing_tasks.pop(rollout.rollout_id, None)
@@ -132,27 +169,31 @@ class ServerDataStore:
             logger.info(f"Rollout received and stored: {rollout.rollout_id}")
 
     async def retrieve_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
-        """
-        Safely retrieves a single rollout by its ID, removing it from the store.
+        """Retrieve and remove a stored rollout by identifier.
+
+        Args:
+            rollout_id: Identifier of the rollout to fetch.
+
+        Returns:
+            Stored [`RolloutLegacy`][agentlightning.RolloutLegacy], or ``None``
+            when the identifier is unknown.
         """
         async with self._results_lock:
             return self._completed_rollouts.pop(rollout_id, None)
 
     async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
-        """
-        Retrieves all completed rollouts and clears the store.
-        """
+        """Return all completed rollouts and clear the internal buffer."""
         async with self._results_lock:
             rollouts = list(self._completed_rollouts.values())
             self._completed_rollouts.clear()
             return rollouts
 
     def get_processing_tasks(self) -> Dict[str, Task]:
-        """Returns a copy of currently processing tasks for timeout checking."""
+        """Return a copy of currently processing tasks for timeout checking."""
         return self._processing_tasks.copy()
 
     async def requeue_task(self, task: Task):
-        """Requeues a task that has timed out and removes it from processing."""
+        """Requeue a task that timed out while being processed."""
         logger.warning(f"Requeuing task {task.rollout_id} after timeout (attempt {task.num_claims})")
         async with self._results_lock:
             # Remove from processing tasks
@@ -161,21 +202,26 @@ class ServerDataStore:
 
 
 class AgentLightningServer:
-    """
-    The main SDK class for developers to control the Agent Lightning Server.
+    """High-level controller for the legacy Agent Lightning FastAPI server.
 
-    This class manages the server lifecycle, task queueing, resources updates,
-    and retrieval of results, providing a simple interface for the optimization logic.
+    The controller orchestrates server start-up, task queueing, resource updates,
+    and retrieval of client rollouts. It is primarily used by existing systems that
+    still rely on the HTTP-based workflow.
+
+    !!! warning "Deprecated"
+        [`AgentLightningServer`][agentlightning.server.AgentLightningServer] is part of
+        the legacy client/server stack. Prefer the store-based runtime for new
+        integrations.
     """
 
     def __init__(self, host: str = "127.0.0.1", port: int = 8000, task_timeout_seconds: float = 300.0):
-        """
-        Initializes the server controller.
+        """Initialize the controller.
 
         Args:
-            host: The host to bind the server to.
-            port: The port to bind the server to.
-            task_timeout_seconds: Time in seconds after which a claimed task is considered stale and requeued.
+            host: Hostname or IP address to bind the HTTP server to.
+            port: TCP port exposed by the server.
+            task_timeout_seconds: Seconds before a claimed task is considered stale and
+                re-queued.
         """
         warnings.warn(
             "AgentLightningServer is deprecated. Please use LightningStoreServer instead.", DeprecationWarning
@@ -200,9 +246,7 @@ class AgentLightningServer:
     # --- ADDED: Lifespan context manager ---
     @asynccontextmanager
     async def _lifespan(self, app: FastAPI):
-        """
-        Manages server startup and shutdown. This runs inside the server's event loop.
-        """
+        """Manage server start-up and shutdown within the event loop."""
         logger.info("Server is starting up...")
         self.loop = asyncio.get_running_loop()
         self._store = ServerDataStore()  # Initialize data store here
@@ -216,9 +260,7 @@ class AgentLightningServer:
         self.loop = None
 
     async def _check_and_requeue_stale_tasks(self):
-        """
-        Check for stale tasks and requeue them. Called reactively during get_next_task.
-        """
+        """Check for stale tasks and requeue them when they exceed the timeout."""
         current_time = time.time()
         # Ensure store is initialized before checking
         if not self._store:
@@ -233,11 +275,11 @@ class AgentLightningServer:
                 )
 
     def _setup_routes(self):
-        """Setup FastAPI routes."""
+        """Configure the FastAPI routes that make up the legacy HTTP API."""
 
         @self._app.get("/task", response_model=TaskIfAny)
         async def next_task() -> TaskIfAny:  # type: ignore
-            """Endpoint for clients to poll for the next available task."""
+            """Provide the next available task to a client."""
             await self._check_and_requeue_stale_tasks()
 
             if not self._store:
@@ -253,7 +295,7 @@ class AgentLightningServer:
 
         @self._app.get("/resources/latest", response_model=ResourcesUpdate)
         async def fetch_latest_resources() -> ResourcesUpdate:  # type: ignore
-            """Endpoint for clients to poll for the latest available resources."""
+            """Return the most recent resource bundle published to the server."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
             resources_update = await self._store.get_latest_resources()
@@ -266,7 +308,7 @@ class AgentLightningServer:
         async def fetch_resources_by_id(  # type: ignore
             resource_id: str = Path(..., description="The unique identifier for the resource version.")
         ) -> ResourcesUpdate:
-            """Endpoint for clients to fetch a specific version of resources."""
+            """Return a specific version of resources by identifier."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
             resources_update = await self._store.get_resources_by_id(resource_id)
@@ -277,7 +319,7 @@ class AgentLightningServer:
 
         @self._app.post("/rollout", response_model=GenericResponse)
         async def post_rollout(payload: RolloutLegacy) -> GenericResponse:  # type: ignore
-            """Endpoint for clients to report a completed rollout."""
+            """Persist the rollout reported by a client."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
             await self._store.store_rollout(payload)
@@ -287,13 +329,13 @@ class AgentLightningServer:
             )
 
     async def start(self):
-        """Starts the FastAPI server in the background."""
+        """Start the FastAPI server in the background."""
         logger.info(f"Starting server at {self.endpoint}")
         asyncio.create_task(self._uvicorn_server.serve())
         await asyncio.sleep(1)  # Allow time for server to start up.
 
     async def stop(self):
-        """Gracefully stops the running FastAPI server."""
+        """Stop the FastAPI server and wait for a graceful shutdown."""
         if self._uvicorn_server.started:
             logger.info("Stopping server...")
             self._uvicorn_server.should_exit = True
@@ -301,10 +343,7 @@ class AgentLightningServer:
             logger.info("Server stopped.")
 
     async def run_forever(self):
-        """
-        Runs the server indefinitely until stopped.
-        This is useful when async start and stop methods do not work.
-        """
+        """Run the server indefinitely until `stop()` is invoked."""
         await self._uvicorn_server.serve()
 
     async def queue_task(
@@ -314,35 +353,37 @@ class AgentLightningServer:
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> str:
-        """
-        Adds a task to the queue for a client to process.
-        """
+        """Add a task to the queue for a client to process."""
         if not self._store:
             raise RuntimeError("Store not initialized. The server may not be running.")
         return await self._store.add_task(sample, mode=mode, resources_id=resources_id, metadata=metadata)
 
     async def update_resources(self, resources: NamedResources) -> str:
-        """
-        Updates the resources, creating a new version and setting it as the latest.
-        """
+        """Publish a new resource bundle and return its generated identifier."""
         if not self._store:
             raise RuntimeError("Store not initialized. The server may not be running.")
         resources_id = f"res-{uuid.uuid4()}"
-        update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+        update = ResourcesUpdate(
+            resources_id=resources_id, resources=resources, create_time=time.time(), update_time=time.time(), version=1
+        )
         await self._store.update_resources(update)
         return resources_id
 
     async def get_completed_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
-        """
-        Retrieves a specific completed rollout by its ID.
-        """
+        """Retrieve a specific completed rollout by identifier."""
         if not self._store:
             raise RuntimeError("Store not initialized. The server may not be running.")
         return await self._store.retrieve_rollout(rollout_id)
 
     async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[RolloutLegacy]:
-        """
-        Polls for a completed rollout by its ID, waiting up to `timeout` seconds.
+        """Poll for a completed rollout until it becomes available or a timeout expires.
+
+        Args:
+            rollout_id: Identifier of the rollout to wait for.
+            timeout: Maximum number of seconds to wait. ``None`` waits indefinitely.
+
+        Returns:
+            Retrieved rollout, or ``None`` when the timeout is reached without success.
         """
         start_time = time.time()
         while True:
@@ -354,9 +395,7 @@ class AgentLightningServer:
             await asyncio.sleep(1)
 
     async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
-        """
-        Retrieves all available completed trajectories and clears the internal store.
-        """
+        """Return every completed rollout and clear the internal buffer."""
         if not self._store:
             raise RuntimeError("Store not initialized. The server may not be running.")
         return await self._store.retrieve_completed_rollouts()

agentlightning/store/__init__.py

@@ -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",
 ]

agentlightning/store/base.py

@@ -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,33 +56,144 @@ UNSET = _UnsetType()
 Unset = _UnsetType  # Alias for convenience
 
 
-class LightningStore:
-    """
-    A centralized, thread-safe, async, data store for the lightning's state.
-    This holds the task queue, versioned resources, and completed rollouts.
+class LightningStoreCapabilities(TypedDict, total=False):
+    """Capability of a LightningStore implementation.
 
-    The store has a built-in clock and it should be responsible for tracking the times.
-    All the time-based operations like retry, timeout, etc. should be handled by the store.
+    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.
+
+    A `LightningStore` mediates every interaction between algorithms and runners:
+
+    - **Rollout lifecycle:** accept new rollouts, queue them for execution, create attempts,
+      and drive the rollout status machine (`"queuing"` → `"preparing"` → `"running"` →
+      `{"succeeded","failed","cancelled"}` or `"requeuing"` when a retry is justified).
+    - **Attempt tracking:** record each execution attempt, including progress heartbeats,
+      retry sequencing, and terminal states such as `"timeout"` or `"unresponsive"`.
+    - **Span ingest:** capture structured telemetry emitted by runners (either as native
+      [`Span`][agentlightning.Span] objects or as `opentelemetry.sdk.trace.ReadableSpan`
+      instances) so that algorithms can reconstruct trajectories and rewards.
+    - **Resource versioning:** manage immutable snapshots of named resources
+      (prompt templates, model checkpoints, proxy endpoints, …) and expose a single
+      "latest" snapshot that runners can fetch just after claiming work.
+
+    Implementations must provide thread-safe/async-safe semantics: each coroutine should
+    appear atomic to callers even when multiple algorithms or runners call the API concurrently.
+    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:
-        """
-        Add one incomplete rollout to the store, and get an attempt created for it.
-        This will immediately sets the rollout to a preparing state, and should be
-        used by whoever is going to execute the rollout.
+        """Register a rollout and immediately create its first attempt.
 
-        Return a special rollout with attempt object. Do not update it directly.
+        !!! note
+            Use [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout] when the
+            caller only wants to submit work for later scheduling.
 
-        But if the rollout fails or timeouts, it's still possible that the watchdog
-        sends it back to the queue for retry.
+        The rollout must be persisted with `status="preparing"` and an initial attempt
+        with `sequence_id == 1` so the caller can begin execution without visiting the
+        public queue. Implementations are expected to:
 
-        To enqueue a rollout to the task queue, use `enqueue_rollout` instead.
+        1. Generate a unique `rollout_id` and `attempt_id`.
+        2. Record `start_time` for both rollout and attempt based on the current clock.
+        3. Copy `config` and `metadata` so later mutations do not leak shared references.
+        4. Resolve `resources_id` to the latest resource snapshot when `None` is supplied.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode for downstream analytics (`"train"`, `"val"`, `"test"`).
+            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
+            the just-created attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide durable storage for the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
         """
         raise NotImplementedError()
 
@@ -87,34 +202,153 @@ class LightningStore:
         input: TaskInput,
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
+        config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> Rollout:
-        """
-        Adds a new task to the queue with specific metadata and
-        returns the rollout object with its unique ID.
+        """Persist a rollout in `queuing` state so runners can claim it later.
+
+        !!! note
+            Different from [`start_rollout()`][agentlightning.LightningStore.start_rollout],
+            this method is called when the caller only wants to submit work for later scheduling.
+
+        Implementations must generate a unique `rollout_id`, stamp `start_time` with
+        the current time, default `config` to a fresh [`RolloutConfig`][agentlightning.RolloutConfig],
+        and insert the rollout at the tail of the scheduling queue. No attempt is created yet.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode indicator (`"train"`, `"val"`, `"test"`).
+            resources_id: Resource snapshot used when a runner eventually executes the rollout.
+            config: Fine-grained retry/timeout parameters to persist with the rollout.
+            metadata: Free-form metadata stored verbatim with the rollout record.
+
+        Returns:
+            The stored [`Rollout`][agentlightning.Rollout] in `queuing` status.
+
+        Raises:
+            NotImplementedError: Subclasses must persist the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
         """
         raise NotImplementedError()
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        """
-        Retrieves the next task from the queue without blocking.
-        Returns None if the queue is empty.
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        """Persist multiple rollouts in `queuing` state.
 
-        Will set the rollout status to preparing.
+        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 start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        """
-        Create a new attempt for a given rollout ID and return the attempt details.
+    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.
+
+        Retrieval must be FIFO across rollouts that remain in `queuing` or `requeuing`
+        state. When a rollout is claimed, implementations must:
+
+        * Transition its status to `"preparing"`.
+        * Create a new attempt with `status="preparing"` and `sequence_id` equal to
+          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.
+
+        Raises:
+            NotImplementedError: Subclasses must implement queue retrieval.
         """
         raise NotImplementedError()
 
-    async def add_span(self, span: Span) -> Span:
-        """
-        Add a span to the store.
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        """Claim up to `limit` queued rollouts without blocking.
 
-        This method is responsible for updating the rollout/attempt status to "running" if needed.
+        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
+        normal queue flow (for example in an online RL setup).
+        Implementations must validate that the rollout exists, allocate a fresh `attempt_id`,
+        increment the `sequence_id` monotonically, stamp the new attempt with `status="preparing"`,
+        and return an up-to-date [`AttemptedRollout`][agentlightning.AttemptedRollout].
+
+        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.
+
+        Raises:
+            NotImplementedError: Subclasses must implement attempt creation.
+            ValueError: Implementations must raise when `rollout_id` is unknown.
+        """
+        raise NotImplementedError()
+
+    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`,
+        `attempt_id`, and `sequence_id`. Implementations must:
+
+        * Verify that both rollout and attempt exist.
+        * Ensure span ordering remains strictly increasing per attempt (rejecting or keeping duplicates).
+        * Treat the span arrival as a heartbeat: update the attempt's `last_heartbeat_time`
+          and transition both attempt and rollout to `"running"` if they were still
+          `"preparing"` or `"requeuing"`.
+
+        Args:
+            span: Fully populated span to persist.
+
+        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.
+            ValueError: Implementations must raise when the referenced rollout or attempt is missing.
         """
         raise NotImplementedError()
 
@@ -124,89 +358,360 @@ class LightningStore:
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Span:
-        """
-        Add an opentelemetry span to the store.
+    ) -> Optional[Span]:
+        """Convert and persist an OpenTelemetry span for a particular attempt.
 
-        If sequence_id is not provided, it will be fetched from `get_next_span_sequence_id` and assigned automatically.
+        Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
+        (typically via [`Span.from_opentelemetry()`][agentlightning.Span.from_opentelemetry]),
+        assign a strictly increasing `sequence_id` when one is not provided, and persist it
+        using the same semantics as [`add_span()`][agentlightning.LightningStore.add_span].
+
+        Args:
+            rollout_id: Identifier of the rollout that produced the span.
+            attempt_id: Attempt identifier the span belongs to.
+            readable_span: OpenTelemetry span in SDK form.
+            sequence_id: Optional explicit ordering hint. When omitted, call
+                [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
+                automatically.
+
+        Returns:
+            The stored span record. Return `None` if the span was not added due to a duplicate.
+
+        Raises:
+            NotImplementedError: Subclasses must implement span persistence.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
         """
         raise NotImplementedError()
 
     async def query_rollouts(
-        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Rollout]:
-        """
-        Query and retrieve rollouts filtered by their status.
-        If no status is provided, returns all 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,
+        # 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_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 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]:
-        """
-        Query and retrieve all attempts associated with a specific rollout ID.
-        Returns an empty list if no attempts are found.
+    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:
+            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.
+            ValueError: Implementations must raise when the rollout does not exist.
         """
         raise NotImplementedError()
 
     async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
-        """
-        Safely retrieves a specific rollout by its ID.
+        """Fetch a rollout by identifier without mutating its state.
+
+        Args:
+            rollout_id: Identifier to retrieve.
+
+        Returns:
+            The rollout when found, otherwise `None`.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
         """
         raise NotImplementedError()
 
     async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """Fetch the attempt with the highest `sequence_id` for `rollout_id`.
+
+        Args:
+            rollout_id: Identifier to inspect.
+
+        Returns:
+            The most recent attempt or `None` when no attempts exist yet.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+            ValueError: Implementations must raise when the rollout does not exist.
         """
-        Safely retrieves the latest attempt for a given rollout ID.
+        raise NotImplementedError()
+
+    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:
+            [`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.
         """
         raise NotImplementedError()
 
     async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves a specific version of named resources by its ID.
+        """Return a specific named resource snapshot by identifier.
+
+        Args:
+            resources_id: Identifier of the snapshot.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when missing.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
         """
         raise NotImplementedError()
 
     async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves the latest version of named resources.
+        """Fetch the latest resource snapshot marked as the global default.
+
+        Returns:
+            The current latest [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when
+            no resources have been registered yet.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
         """
         raise NotImplementedError()
 
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        """
-        Get the next span sequence ID for a given rollout and attempt.
-        This should be used to assign a unique sequence ID to each span within an attempt.
+        """Allocate the next strictly increasing sequence number used to order spans.
 
-        Recommend getting the ID before the operation even begins to avoid racing conditions.
+        Implementations must retain counters so repeated calls return `1, 2, ...` without
+        gaps unless spans were explicitly inserted with a custom `sequence_id`. The
+        counter may be scoped per rollout or per attempt, but the sequence must be
+        strictly increasing for spans emitted by the specified attempt so traces remain
+        totally ordered.
+
+        See [Distributed Tracing][distributed-tracing] for detailed motivations.
+
+        Args:
+            rollout_id: Identifier of the rollout emitting spans.
+            attempt_id: Attempt identifier for the upcoming span.
+
+        Returns:
+            The next integer sequence identifier, unique within the attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the allocator.
+            ValueError: Implementations must raise when the rollout or attempt does not exist.
+        """
+        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]:
-        """
-        Wait for specified rollouts to complete with a timeout.
-        Returns the completed rollouts, potentially incomplete if timeout is reached.
+        """Block until the targeted rollouts reach a terminal status or the timeout expires.
 
-        TODO: Add support for waiting for 20 new rollouts, or wait until 80% of the pending ids are completed.
+        Terminal statuses are `"succeeded"`, `"failed"`, and `"cancelled"`. When the timeout
+        elapses, implementations should return the subset of rollouts that are already terminal
+        and omit the rest.
+
+        !!! warning
+            It's dangerous and might be event-loop blocking to call this function
+            with a long timeout. It's a good idea to poll for the method to check
+            if new completed rollouts can coming. Be careful in implementing the sleep logic
+            to avoid busy-waiting.
+
+        Args:
+            rollout_ids: Identifiers of rollouts to watch.
+            timeout: Maximum time in seconds to wait. `None` waits indefinitely.
+
+        Returns:
+            Rollouts that finished before the deadline, in arbitrary order.
+
+        Raises:
+            NotImplementedError: Subclasses must implement waiting semantics.
+            ValueError: Implementations must raise when a rollout identifier is unknown.
         """
         raise NotImplementedError()
 
-    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
-        """
-        Query and retrieve all spans associated with a specific rollout ID.
-        Returns an empty list if no spans are found.
+    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.
+
+        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.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
         """
         raise NotImplementedError()
 
     async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
-        """
-        Safely stores a new version of named resources and sets it as the latest.
-        Not implemented by many stores yet.
+        """Persist a new immutable snapshot of named resources and mark it as latest.
+
+        Implementations must assign a fresh `resources_id` and ensure subsequent calls to
+        [`get_latest_resources()`][agentlightning.LightningStore.get_latest_resources] return the
+        snapshot produced here.
+
+        Args:
+            resources: Mapping of resource names to their serialized payloads.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate] including its generated id.
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
         """
         raise NotImplementedError()
 
     async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        """
-        Safely stores a new version or updates an existing version of named resources and sets it as the latest.
+        """Overwrite or extend an existing resource snapshot and mark it as latest.
+
+        This API is typically used by algorithms that maintain mutable resources (e.g., model
+        checkpoints) under a stable identifier.
+
+        Args:
+            resources_id: Identifier of the snapshot to replace.
+            resources: Updated mapping of resource names to payloads.
+
+        Returns:
+            The persisted [`ResourcesUpdate`][agentlightning.ResourcesUpdate].
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
+            ValueError: Implementations must raise when `resources_id` does not exist.
         """
         raise NotImplementedError()
 
@@ -220,22 +725,31 @@ class LightningStore:
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
     ) -> Rollout:
-        """
-        Update the rollout status and related metadata.
+        """Update rollout metadata and, when provided, drive status transitions.
 
-        Not-listed fields here either cannot be updated, or should be auto-updated (e.g., end_time).
+        Parameters default to the sentinel [`UNSET`][agentlightning.store.base.UNSET] to
+        distinguish omitted fields from explicit `None` assignments. Implementations must:
 
-        When status is updated to a finished / problematic state, other states like task
-        queues will be updated accordingly.
+        * Validate the rollout exists before mutating it.
+        * Replace each property when a concrete value (including `None`) is supplied.
+        * When the status switches into a terminal state, set `end_time` and signal any waiters.
+        * When the status re-enters a queueing state, ensure the rollout is enqueued exactly once.
 
         Args:
-            rollout_id: Unique identifier for the rollout to update
-            input: New input data for the rollout. If set, will be updated. Can be updated to None
-            mode: New mode for the rollout. If set, will be updated. Can be updated to None
-            resources_id: New resources ID for the rollout. If set, will be updated. Can be updated to None
-            status: New status for the rollout. If set, will be updated
-            config: New config for the rollout. If set, will be updated
-            metadata: Dictionary of additional metadata to update. If set, will replace the existing metadata
+            rollout_id: Identifier of the rollout to update.
+            input: Replacement task payload; pass `None` to explicitly clear the input.
+            mode: Replacement rollout mode.
+            resources_id: Replacement resources snapshot reference.
+            status: Target rollout status.
+            config: Replacement retry/timeout configuration.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated rollout record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            ValueError: Implementations must raise when the rollout is unknown or the update is invalid.
         """
         raise NotImplementedError()
 
@@ -248,18 +762,97 @@ class LightningStore:
         last_heartbeat_time: float | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
     ) -> Attempt:
-        """
-        Update a specific or latest attempt for a given rollout.
+        """Update attempt bookkeeping such as status, worker ownership, and heartbeats.
 
-        Update the latest attempt will NOT affect the corresponding rollout status.
+        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 [`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: Unique identifier for the rollout
-            attempt_id: Unique identifier for the attempt
-            status: Status to set for the attempt, update if provided
-            worker_id: Worker identifier, update if provided
-            last_heartbeat_time: Timestamp of the last heartbeat from the worker
-            metadata: Dictionary of additional metadata to update, will replace the existing metadata
+            rollout_id: Identifier of the rollout whose attempt will be updated.
+            attempt_id: Attempt identifier or `"latest"` as a convenience.
+            status: Replacement attempt status. Terminal statuses must set `end_time`.
+            worker_id: Identifier for the worker currently processing the attempt.
+            last_heartbeat_time: Wall-clock timestamp (seconds) of the latest heartbeat/span.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated attempt record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            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()

agentlightning/store/collection/__init__.py

@@ -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",
+]

agentlightning/store/collection/base.py

@@ -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

agentlightning/store/collection/memory.py

@@ -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()

agentlightning/store/mongo.py

@@ -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
+        ]

agentlightning/store/threading.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import threading
-from typing import Any, Dict, List, Literal, Optional, Sequence
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -11,6 +11,7 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -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,46 +38,117 @@ 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,
         mode: Literal["train", "val", "test"] | None = None,
         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, metadata)
+            return await self.store.start_rollout(
+                input,
+                mode,
+                resources_id,
+                config,
+                metadata,
+                worker_id,
+            )
 
     async def enqueue_rollout(
         self,
         input: TaskInput,
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
+        config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> Rollout:
         with self._lock:
-            return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+            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:
@@ -84,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)
@@ -100,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)
 
@@ -110,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)
 
@@ -122,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,
@@ -169,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,
+            )

agentlightning/store/utils.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import time
-from typing import Awaitable, Callable, List, cast
+from typing import Awaitable, Callable, Dict, List, Tuple
 
 from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
 
@@ -9,65 +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
@@ -75,52 +112,31 @@ async def healthcheck(
         # Get the latest attempt for this rollout
         latest_attempt = rollout.attempt
         if not latest_attempt:
-            continue
-
-        # Check if the attempt has already failed or succeeded
-        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
-            await propagate_status(update_rollout_status, latest_attempt, config)
+            # This should not happen
             continue
 
         # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
         if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "timeout",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
             continue
 
         # Check for unresponsive condition (based on last heartbeat)
-        if latest_attempt.last_heartbeat_time:
-            if latest_attempt.status == "preparing":
-                # If still preparing, mark it as running
-                latest_attempt = await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "running",
-                )
+        # (1) Haven't received heartbeat for a while
+        if (
+            latest_attempt.last_heartbeat_time
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
+        ):
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
 
-            # Haven't received heartbeat for a while
-            if (
-                config.unresponsive_seconds is not None
-                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
-            ):
-                await update_attempt_status(
-                    latest_attempt.rollout_id,
-                    latest_attempt.attempt_id,
-                    "unresponsive",
-                )
-                continue
-
-        # Check if there's no last heartbeat (no spans) at all
+        # (2) Check if there's no last heartbeat (no spans) at all
         if (
             latest_attempt.last_heartbeat_time is None
             and config.unresponsive_seconds is not None
             and current_time - latest_attempt.start_time > config.unresponsive_seconds
         ):
-            await update_attempt_status(
-                latest_attempt.rollout_id,
-                latest_attempt.attempt_id,
-                "unresponsive",
-            )
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
+
+    return updates

agentlightning/tracer/__init__.py

@@ -1,7 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from .agentops import AgentOpsTracer
-from .base import BaseTracer
+from .base import Tracer, clear_active_tracer, get_active_tracer, set_active_tracer
+from .dummy import DummyTracer
 from .otel import OtelTracer
 
-__all__ = ["AgentOpsTracer", "BaseTracer", "OtelTracer"]
+__all__ = [
+    "AgentOpsTracer",
+    "Tracer",
+    "OtelTracer",
+    "DummyTracer",
+    "get_active_tracer",
+    "set_active_tracer",
+    "clear_active_tracer",
+]

agentlightning/tracer/agentops.py

@@ -2,24 +2,25 @@
 
 from __future__ import annotations
 
-import asyncio
 import logging
 import os
-import threading
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Awaitable, Iterator, List, Optional
+import warnings
+from contextlib import asynccontextmanager, contextmanager
+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.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.instrumentation.agentops import AgentOpsServerManager
 from agentlightning.store.base import LightningStore
+from agentlightning.utils.otel import get_span_processors, get_tracer_provider
 
-from .base import BaseTracer
+from .base import with_active_tracer_context
+from .otel import LightningSpanProcessor, OtelTracer
 
 if TYPE_CHECKING:
     from agentops.integration.callbacks.langchain import LangchainCallbackHandler
@@ -28,7 +29,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class AgentOpsTracer(BaseTracer):
+class AgentOpsTracer(OtelTracer):
     """Traces agent execution using AgentOps.
 
     This tracer provides functionality to capture execution details using the
@@ -55,95 +56,44 @@ class AgentOpsTracer(BaseTracer):
         self.instrument_managed = instrument_managed
         self.daemon = daemon
 
-        self._agentops_server_manager = AgentOpsServerManager(self.daemon)
-        self._agentops_server_port_val: Optional[int] = None
-
         if not self.agentops_managed:
             logger.warning("agentops_managed=False. You are responsible for AgentOps setup.")
         if not self.instrument_managed:
             logger.warning("instrument_managed=False. You are responsible for all instrumentation.")
 
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        state["_agentops_server_manager"] = None  # Exclude the unpicklable server manager
-        # _agentops_server_port_val (int) is inherently picklable and will be included.
-        logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
-        return state
-
-    def __setstate__(self, state: Any):
-        self.__dict__.update(state)
-        # In child process, self._agentops_server_manager will be None.
-        logger.debug(f"Setting state for unpickled Trainer (PID {os.getpid()}). _agentops_server_manager is None.")
-
-    def init(self, *args: Any, **kwargs: Any):
-        if self.agentops_managed and self._agentops_server_manager:
-            self._agentops_server_manager.start()
-            self._agentops_server_port_val = self._agentops_server_manager.get_port()
-            if self._agentops_server_port_val is None:
-                if (
-                    self._agentops_server_manager.server_process is not None
-                    and self._agentops_server_manager.server_process.is_alive()
-                ):
-                    raise RuntimeError("AgentOps server started but port is None. Check server manager logic.")
-                elif (
-                    self._agentops_server_port_val is None and self._agentops_server_manager.server_process is None
-                ):  # Server failed to start
-                    raise RuntimeError("AgentOps server manager indicates server is not running and port is None.")
-
-    def teardown(self):
-        if self.agentops_managed:
-            self._agentops_server_manager.stop()
-            logger.info("AgentOps server stopped.")
-
     def instrument(self, worker_id: int):
         instrument_all()
 
     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)
             logger.info(f"[Worker {worker_id}] Instrumentation applied.")
 
         if self.agentops_managed:
-            if self._agentops_server_port_val:  # Use the stored, picklable port value
-                base_url = f"http://localhost:{self._agentops_server_port_val}"
-                env_vars_to_set = {
-                    "AGENTOPS_API_KEY": "dummy",
-                    "AGENTOPS_API_ENDPOINT": base_url,
-                    "AGENTOPS_APP_URL": f"{base_url}/notavailable",
-                    "AGENTOPS_EXPORTER_ENDPOINT": f"{base_url}/traces",
-                }
-                for key, value in env_vars_to_set.items():
-                    os.environ[key] = value
-                    logger.info(f"[Worker {worker_id}] Env var set: {key}={value}")
-            else:
-                logger.warning(
-                    f"[Worker {worker_id}] AgentOps managed, but local server port is not available. Client may not connect as expected."
-                )
-
+            os.environ.setdefault("AGENTOPS_API_KEY", "dummy")
             if not agentops.get_client().initialized:
-                agentops.init()  # type: ignore
+                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)
@@ -152,15 +102,19 @@ class AgentOpsTracer(BaseTracer):
             self.uninstrument(worker_id)
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
-    @contextmanager
-    def trace_context(
+        # 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,
         name: Optional[str] = None,
         *,
         store: Optional[LightningStore] = None,
         rollout_id: Optional[str] = None,
         attempt_id: Optional[str] = None,
-    ) -> Iterator[LightningSpanProcessor]:
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
@@ -171,33 +125,73 @@ class AgentOpsTracer(BaseTracer):
             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 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(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> 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 as processor:
-                yield processor
+            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:
-                yield 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")
 
-    def get_last_trace(self) -> List[ReadableSpan]:
-        """
-        Retrieves the raw list of captured spans from the most recent trace.
+    @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:
+            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}: {e}")
+            raise  # should reraise the error here so that runner can handle it
+        finally:
+            agentops.end_trace(trace, end_state=status)  # type: ignore
 
-        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_callback_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
         """
         Get the Langchain callback handler for integrating with Langchain.
 
@@ -221,104 +215,28 @@ class AgentOpsTracer(BaseTracer):
             )
         return LangchainCallbackHandler(api_key=api_key, tags=tags)
 
+    get_langchain_callback_handler = get_langchain_handler  # alias
 
-class LightningSpanProcessor(SpanProcessor):
-    def __init__(self):
-        self._spans: List[ReadableSpan] = []
+    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.")
 
-        # 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.")
-        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
-                self._await_in_loop(
-                    self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                    timeout=5.0,
+            if get_tracer_provider() is not instance.provider:
+                logger.error(
+                    "Mismatch between global singleton TracerProvider and AgentOps TracerProvider. "
+                    "AgentOps might not work properly."
                 )
-            except Exception:
-                # log; on_end MUST NOT raise
-                logger.exception(f"Error adding span to store: {span.name}")
 
-        self._spans.append(span)
+            if not isinstance(instance.provider, TracerProviderImpl):  # type: ignore
+                raise RuntimeError("Unsupported TracerProvider type for AgentOps instrumentation.")
+
+            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